From 08b487fc2068cfb8812602a029560e56f8707744 Mon Sep 17 00:00:00 2001
From: Amanda Richardson <amanda162000@gmail.com>
Date: Tue, 16 Jul 2024 19:03:21 -0500
Subject: [PATCH] updates to adapt to ensemble

---
 smartsim/_core/generation/generator.py | 35 ++++++++++++++++----------
 smartsim/entity/ensemble.py            |  2 +-
 smartsim/launchable/job.py             |  5 ++++
 tests/test_generator/test_generator.py | 13 +++++++++-
 4 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/smartsim/_core/generation/generator.py b/smartsim/_core/generation/generator.py
index 1c8d4043c6..749657ba4d 100644
--- a/smartsim/_core/generation/generator.py
+++ b/smartsim/_core/generation/generator.py
@@ -66,17 +66,29 @@ def __init__(self, gen_path: str, run_ID: str, job: Job) -> None:
         """
         self.job = job
         # TODO revisit this check
-        if isinstance(job, (Job, JobGroup)):
+        if job._ensemble_name is None:
             job_type = f"{job.__class__.__name__.lower()}s"
-        entity_type = f"{job.entity.__class__.__name__.lower()}-{create_short_id_str()}"
-        self.path = os.path.join(
-            gen_path,
-            run_ID,
-            job_type,
-            f"{job.name}-{create_short_id_str()}",
-            entity_type,
-            "run",
-        )
+            entity_type = f"{job.entity.__class__.__name__.lower()}-{create_short_id_str()}"
+            self.path = os.path.join(
+                gen_path,
+                run_ID,
+                job_type,
+                f"{job.name}-{create_short_id_str()}",
+                entity_type,
+                "run",
+            )
+        else:
+            job_type = "ensembles"
+            entity_type = f"{job.entity.__class__.__name__.lower()}-{create_short_id_str()}"
+            self.path = os.path.join(
+                gen_path,
+                run_ID,
+                job_type,
+                job._ensemble_name,
+                f"{job.name}",
+                entity_type,
+                "run",
+            )
 
     @property
     def log_level(self) -> int:
@@ -130,9 +142,6 @@ def generate_experiment(self) -> str:
 
         """
         pathlib.Path(self.path).mkdir(exist_ok=True, parents=True)
-        # logger.log(
-        #     level=self.log_level, msg="Working in experiment "
-        # )
 
         # The log_file only keeps track of the last generation
         # this is to avoid gigantic files in case the user repeats
diff --git a/smartsim/entity/ensemble.py b/smartsim/entity/ensemble.py
index 711cfdee34..883b69f9a7 100644
--- a/smartsim/entity/ensemble.py
+++ b/smartsim/entity/ensemble.py
@@ -110,4 +110,4 @@ def as_jobs(self, settings: LaunchSettings) -> tuple[Job, ...]:
         apps = self._create_applications()
         if not apps:
             raise ValueError("There are no members as part of this ensemble")
-        return tuple(Job(app, settings) for app in apps)
+        return tuple(Job(app, settings, f"job_{i}", ensemble_name=self.name) for i, app in enumerate(apps, 1))
diff --git a/smartsim/launchable/job.py b/smartsim/launchable/job.py
index 2e441c9163..da73b8317d 100644
--- a/smartsim/launchable/job.py
+++ b/smartsim/launchable/job.py
@@ -32,6 +32,7 @@
 from smartsim._core.commands.launchCommands import LaunchCommands
 from smartsim.launchable.basejob import BaseJob
 from smartsim.settings import LaunchSettings
+from smartsim._core.utils.helpers import create_short_id_str
 
 if t.TYPE_CHECKING:
     from smartsim.entity.entity import SmartSimEntity
@@ -51,11 +52,15 @@ def __init__(
         entity: SmartSimEntity,
         launch_settings: LaunchSettings,
         name: str = "job",
+        **kwargs: t.Any,
     ):
         super().__init__()
         self._entity = deepcopy(entity)
         self._launch_settings = deepcopy(launch_settings)
         self._name = deepcopy(name)
+        self._ensemble_name = kwargs.get('ensemble_name', None)
+        if self._ensemble_name is not None:
+            self._ensemble_name += f"-{create_short_id_str()}"
         # TODO: self.warehouse_runner = JobWarehouseRunner
 
     # TODO do we want the user to be allowed to reset the Job name? Therefore, add setter
diff --git a/tests/test_generator/test_generator.py b/tests/test_generator/test_generator.py
index 3feab6f4a3..5558f2ff90 100644
--- a/tests/test_generator/test_generator.py
+++ b/tests/test_generator/test_generator.py
@@ -7,7 +7,7 @@
 
 from smartsim import Experiment
 from smartsim._core.generation.generator import Generator
-from smartsim.entity.model import Application
+from smartsim.entity import Application, Ensemble
 from smartsim.launchable import Job, JobGroup
 from smartsim.settings.builders.launch import SlurmArgBuilder
 from smartsim.settings.dispatch import Dispatcher
@@ -97,3 +97,14 @@ def test_full_exp_generate_job_directory(test_dir, job_instance):
     )
     job_execution_path = no_op_exp._generate(job_instance)
     assert osp.isdir(job_execution_path)
+
+def test_generate_ensemble_directory(test_dir, wlmutils):
+    ensemble = Ensemble("ensemble-name", "echo", replicas=2)
+    launch_settings = LaunchSettings(wlmutils.get_test_launcher())
+    job_list = ensemble.as_jobs(launch_settings)
+    for job in job_list:
+        run_ID = "temp_run"
+        gen = Generator(gen_path=test_dir, run_ID=run_ID, job=job)
+        print(gen.path)
+        
+    
\ No newline at end of file