Allow for sequences into exp.start(), and unpack iterables

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:
     """

smartsim/experiment.py

@@ -151,19 +151,23 @@
         experiment
         """
 
-    def start(self, *jobs: 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
         :returns: A sequence of ids with order corresponding to the sequence of
             jobs that can be used to query or alter the status of that
             particular execution of the job.
         """
+        # If item is instance iterable then unpack and extend the list
+
+        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,

tests/test_experiment.py

@@ -34,23 +34,32 @@
 import time
 import typing as t
 import uuid
+from os import path as osp
 
 import pytest
 
 from smartsim._core import dispatch
-from smartsim._core.control.interval import SynchronousTimeInterval
 from smartsim._core.control.launch_history import LaunchHistory
 from smartsim._core.utils.launcher import LauncherProtocol, create_job_id
+from smartsim.builders.ensemble import Ensemble
 from smartsim.entity import entity
+from smartsim.entity.application import Application
 from smartsim.error import errors
 from smartsim.experiment import Experiment
 from smartsim.launchable import job
 from smartsim.settings import launch_settings
 from smartsim.settings.arguments import launch_arguments
 from smartsim.status import InvalidJobStatus, JobStatus
+from smartsim.types import LaunchedJobID
 
 pytestmark = pytest.mark.group_a
 
+_ID_GENERATOR = (str(i) for i in itertools.count())
+
+
+def random_id():
+    return next(_ID_GENERATOR)
+
 
 @pytest.fixture
 def experiment(monkeypatch, test_dir, dispatcher):
@@ -611,3 +620,128 @@ def test_experiment_stop_does_not_raise_on_unknown_job_id(
     assert stat == InvalidJobStatus.NEVER_STARTED
     after_cancel = exp.get_status(*all_known_ids)
     assert before_cancel == after_cancel
+
+
+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))
+    """
+    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 = launch_settings.LaunchSettings(wlmutils.get_test_launcher())
+    job_list = ensemble.build_jobs(launch_settings)
+
+    exp = Experiment(name="exp_name", exp_path=test_dir)
+
+    exp.start(job_list)
+
+
+@pytest.mark.parametrize(
+    "job_list",
+    (
+        pytest.param(
+            [
+                (
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launch_settings.LaunchSettings("local"),
+                    ),
+                    Ensemble("ensemble-name", "echo", replicas=2).build_jobs(
+                        launch_settings.LaunchSettings("local")
+                    ),
+                )
+            ],
+            id="(job1, job2)",
+        ),
+        pytest.param(
+            [
+                (
+                    Ensemble("ensemble-name", "echo", replicas=2).build_jobs(
+                        launch_settings.LaunchSettings("local")
+                    ),
+                    (
+                        job.Job(
+                            Application(
+                                "test_name",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launch_settings.LaunchSettings("local"),
+                        ),
+                        job.Job(
+                            Application(
+                                "test_name_2",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launch_settings.LaunchSettings("local"),
+                        ),
+                    ),
+                )
+            ],
+            id="(job1, (job2, job3))",
+        ),
+        pytest.param(
+            [
+                (
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launch_settings.LaunchSettings("local"),
+                    ),
+                )
+            ],
+            id="(job,)",
+        ),
+        pytest.param(
+            [
+                [
+                    job.Job(
+                        Application(
+                            "test_name",
+                            exe="echo",
+                            exe_args=["spam", "eggs"],
+                        ),
+                        launch_settings.LaunchSettings("local"),
+                    ),
+                    (
+                        Ensemble("ensemble-name", "echo", replicas=2).build_jobs(
+                            launch_settings.LaunchSettings("local")
+                        ),
+                        job.Job(
+                            Application(
+                                "test_name_2",
+                                exe="echo",
+                                exe_args=["spam", "eggs"],
+                            ),
+                            launch_settings.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"""
+
+    monkeypatch.setattr(
+        "smartsim._core.dispatch._LauncherAdapter.start",
+        lambda launch, exe, job_execution_path, env, out, err: random_id(),
+    )
+
+    exp = Experiment(name="exp_name", exp_path=test_dir)
+    exp.start(*job_list)