type fixes, param tests

CrayLabs · Sep 25, 2024 · fe15a27 · fe15a27
1 parent 6b54a9a
commit fe15a27
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 139 deletions.
diff --git a/smartsim/_core/utils/helpers.py b/smartsim/_core/utils/helpers.py
@@ -43,18 +43,35 @@
 from datetime import datetime
 from shutil import which
 
+from typing_extensions import TypeAlias
+
 if t.TYPE_CHECKING:
     from types import FrameType
 
     from typing_extensions import TypeVarTuple, Unpack
 
+    from smartsim.launchable.job import Job
+
     _Ts = TypeVarTuple("_Ts")
 
 
 _T = t.TypeVar("_T")
 _HashableT = t.TypeVar("_HashableT", bound=t.Hashable)
 _TSignalHandlerFn = t.Callable[[int, t.Optional["FrameType"]], object]
 
+_nested_seq: TypeAlias = "t.Sequence[Job | _nested_seq]"
+
+
+def unpack(value: _nested_seq) -> t.Generator[Job, None, None]:
+    """Unpack any iterable input in order to obtain a
+    single sequence of values"""
+
+    for item in value:
+        if isinstance(item, t.Iterable):
+            yield from unpack(item)
+        else:
+            yield item
+
 
 def check_name(name: str) -> None:
     """

diff --git a/smartsim/experiment.py b/smartsim/experiment.py
@@ -55,17 +55,6 @@
 logger = get_logger(__name__)
 
 
-def _unpack(jobs: t.Sequence[t.Tuple[Job]] | t.Sequence[str]) -> t.Iterator[str]:
-    """Unpack any iterable input into exp.start in order to obtain a
-    single sequence of jobs that can be launched"""
-    for item in jobs:
-
-        if isinstance(item, t.Iterable):
-            yield from _unpack(item)
-        else:
-            yield item
-
-
 def _exp_path_map(exp: "Experiment") -> str:
     """Mapping function for use by method contextualizer to place the path of
     the currently-executing experiment into context for log enrichment"""
@@ -162,9 +151,7 @@ def __init__(self, name: str, exp_path: str | None = None):
         experiment
         """
 
-    def start(
-        self, *jobs: Job | t.Sequence[t.Tuple[Job]] | t.Sequence[Job]
-    ) -> tuple[LaunchedJobID, ...]:
+    def start(self, *jobs: Job | t.Sequence[Job]) -> tuple[LaunchedJobID, ...]:
         """Execute a collection of `Job` instances.
 
         :param jobs: A collection of other job instances to start
@@ -173,20 +160,21 @@ def start(
             particular execution of the job.
         """
         # If item is instance iterable then unpack and extend the list
-        if isinstance(jobs, t.Iterable):
-            jobs = list(tuple(_unpack(jobs)))
+
+        jobs_ = list(tuple(_helpers.unpack(jobs)))
+        # jobs = list(tuple(_helpers.unpack(jobs)))
         # Create the run id
         run_id = datetime.datetime.now().replace(microsecond=0).isoformat()
         # Generate the root path
         root = pathlib.Path(self.exp_path, run_id)
-        return self._dispatch(Generator(root), dispatch.DEFAULT_DISPATCHER, *jobs)
+        return self._dispatch(Generator(root), dispatch.DEFAULT_DISPATCHER, *jobs_)
 
     def _dispatch(
         self,
         generator: Generator,
         dispatcher: dispatch.Dispatcher,
         job: Job,
-        *jobs: Job | t.Sequence[t.Tuple[Job]],
+        *jobs: Job,
     ) -> tuple[LaunchedJobID, ...]:
         """Dispatch a series of jobs with a particular dispatcher
 

diff --git a/tests/test_experiment.py b/tests/test_experiment.py
@@ -56,11 +56,6 @@
 _ID_GENERATOR = (str(i) for i in itertools.count())
 
 
-@pytest.fixture
-def get_gen_symlink_dir(fileutils):
-    yield fileutils.get_test_conf_path(osp.join("generator_files", "to_symlink_dir"))
-
-
 def random_id():
     return next(_ID_GENERATOR)
 
@@ -626,141 +621,126 @@ def test_experiment_stop_does_not_raise_on_unknown_job_id(
     assert before_cancel == after_cancel
 
 
-def test_start_sequence_5(
-    test_dir,
-    wlmutils,
-    monkeypatch,
-):
-    monkeypatch.setattr(
-        "smartsim._core.dispatch._LauncherAdapter.start",
-        lambda launch, exe, job_execution_path, env, out, err: random_id(),
-    )
-    """Test the unpacking of a tuple of tuples 
-            exp.start(job1, (job2, job3))
+def test_start_sequences_of_jobs_good(test_dir, wlmutils, monkeypatch, job_list):
+    """Test the unpacking of a tuple of tuples
+    exp.start(job1, (job2, job3))
     """
-    ensemble = Ensemble("ensemble-name", "echo", replicas=2)
-
-    launch_settings = launchSettings.LaunchSettings(wlmutils.get_test_launcher())
-    job_list = ensemble.as_jobs(launch_settings)
-
-    exp = Experiment(name="exp_name", exp_path=test_dir)
-
-    exp.start(job_list)
-
-
-def test_start_with_two_element_tuple(
-    test_dir,
-    wlmutils,
-    monkeypatch,
-):
-    """Test unpacking a tuple of two jobs"""
-
     monkeypatch.setattr(
         "smartsim._core.dispatch._LauncherAdapter.start",
         lambda launch, exe, job_execution_path, env, out, err: random_id(),
     )
-    ensemble = Ensemble("ensemble-name", "echo", replicas=2)
-
-    application = Application(
-        "test_name",
-        exe="echo",
-        exe_args=["spam", "eggs"],
-    )
-
-    launch_settings = launchSettings.LaunchSettings(wlmutils.get_test_launcher())
-    job_list = ensemble.as_jobs(launch_settings)
-
-    job2 = job.Job(application, launch_settings)
-    exp = Experiment(name="exp_name", exp_path=test_dir)
-
-    exp.start((job2, job_list))
 
-
-def test_start_with_nested_tuple(
-    test_dir,
-    wlmutils,
-    monkeypatch,
-):
-    """Test unpacking a nested tuple"""
-    monkeypatch.setattr(
-        "smartsim._core.dispatch._LauncherAdapter.start",
-        lambda launch, exe, job_execution_path, env, out, err: random_id(),
-    )
     ensemble = Ensemble("ensemble-name", "echo", replicas=2)
 
     launch_settings = launchSettings.LaunchSettings(wlmutils.get_test_launcher())
     job_list = ensemble.as_jobs(launch_settings)
 
-    application = Application(
-        "test_name",
-        exe="echo",
-        exe_args=["spam", "eggs"],
-    )
-
-    job2 = job.Job(application, launch_settings)
-
-    application_2 = Application(
-        "test_name_2",
-        exe="echo",
-        exe_args=["spam", "eggs"],
-    )
-    job_3 = job.Job(application_2, launch_settings)
-
     exp = Experiment(name="exp_name", exp_path=test_dir)
 
-    exp.start((job2, (job_list, job_3)))
-
-
-def test_start_with_one_element_tuple(
-    test_dir,
-    wlmutils,
-    monkeypatch,
-):
-    """Test unpacking a tuple of one job"""
-    monkeypatch.setattr(
-        "smartsim._core.dispatch._LauncherAdapter.start",
-        lambda launch, exe, job_execution_path, env, out, err: random_id(),
-    )
-    ensemble = Ensemble("ensemble-name", "echo", replicas=2)
-
-    launch_settings = launchSettings.LaunchSettings(wlmutils.get_test_launcher())
-    job_list = ensemble.as_jobs(launch_settings)
-
-    exp = Experiment(name="exp_name", exp_path=test_dir)
+    exp.start(job_list)
 
-    exp.start((job_list))
 
+@pytest.mark.parametrize(
+    "job_list",
+    (
+        pytest.param(
+            [
+                (
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launchSettings.LaunchSettings("local"),
+                    ),
+                    Ensemble("ensemble-name", "echo", replicas=2).as_jobs(
+                        launchSettings.LaunchSettings("local")
+                    ),
+                )
+            ],
+            id="(job1, job2)",
+        ),
+        pytest.param(
+            [
+                (
+                    Ensemble("ensemble-name", "echo", replicas=2).as_jobs(
+                        launchSettings.LaunchSettings("local")
+                    ),
+                    (
+                        job.Job(
+                            Application(
+                                "test_name",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launchSettings.LaunchSettings("local"),
+                        ),
+                        job.Job(
+                            Application(
+                                "test_name_2",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launchSettings.LaunchSettings("local"),
+                        ),
+                    ),
+                )
+            ],
+            id="(job1, (job2, job3))",
+        ),
+        pytest.param(
+            [
+                (
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launchSettings.LaunchSettings("local"),
+                    ),
+                )
+            ],
+            id="(job,)",
+        ),
+        pytest.param(
+            [
+                [
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launchSettings.LaunchSettings("local"),
+                    ),
+                    (
+                        Ensemble("ensemble-name", "echo", replicas=2).as_jobs(
+                            launchSettings.LaunchSettings("local")
+                        ),
+                        job.Job(
+                            Application(
+                                "test_name_2",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launchSettings.LaunchSettings("local"),
+                        ),
+                    ),
+                ]
+            ],
+            id="[job_1, (job_2, job_3)]",
+        ),
+    ),
+)
+def test_start_unpack(test_dir, wlmutils, monkeypatch, job_list):
+    """Test unpacking a sequences of jobs"""
 
-def test_start_list_of_jobs(
-    test_dir,
-    wlmutils,
-    monkeypatch,
-):
-    """Test unpacking a list of tuples"""
     monkeypatch.setattr(
         "smartsim._core.dispatch._LauncherAdapter.start",
         lambda launch, exe, job_execution_path, env, out, err: random_id(),
     )
-    ensemble = Ensemble("ensemble-name", "echo", replicas=2)
-
-    launch_settings = launchSettings.LaunchSettings(wlmutils.get_test_launcher())
-    job_list = ensemble.as_jobs(launch_settings)
-
-    application = Application(
-        "test_name",
-        exe="echo",
-        exe_args=["spam", "eggs"],
-    )
-
-    job2 = job.Job(application, launch_settings)
-
-    application_2 = Application(
-        "test_name_2",
-        exe="echo",
-        exe_args=["spam", "eggs"],
-    )
-    job_3 = job.Job(application_2, launch_settings)
 
     exp = Experiment(name="exp_name", exp_path=test_dir)
-
-    exp.start([job2, (job_list, job_3)])
+    exp.start(*job_list)