From fe15a2723110ccc88d3b1cb25ac0343748cb541b Mon Sep 17 00:00:00 2001 From: Julia Putko Date: Wed, 25 Sep 2024 10:32:00 -0700 Subject: [PATCH] type fixes, param tests --- smartsim/_core/utils/helpers.py | 17 +++ smartsim/experiment.py | 24 +--- tests/test_experiment.py | 222 +++++++++++++++----------------- 3 files changed, 124 insertions(+), 139 deletions(-) diff --git a/smartsim/_core/utils/helpers.py b/smartsim/_core/utils/helpers.py index 1133358a6..06ceb28df 100644 --- a/smartsim/_core/utils/helpers.py +++ b/smartsim/_core/utils/helpers.py @@ -43,11 +43,15 @@ 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") @@ -55,6 +59,19 @@ _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 index 6d435206a..0861cc1ed 100644 --- 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 index 0863151ab..8034a1738 100644 --- 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)