diff --git a/smartsim/_core/launcher_/slurm/slurm_launcher.py b/smartsim/_core/launcher_/slurm/slurm_launcher.py index ee1dac5d7..24dcacaef 100644 --- a/smartsim/_core/launcher_/slurm/slurm_launcher.py +++ b/smartsim/_core/launcher_/slurm/slurm_launcher.py @@ -33,6 +33,8 @@ import time import typing as t +from typing_extensions import TypeAlias + from smartsim._core.config import CONFIG from smartsim._core.launcher_.slurm import slurm_commands as commands from smartsim._core.launcher_.slurm import slurm_parser as parser @@ -118,7 +120,7 @@ def start(self) -> None: ) -_SlurmCommandType: t.TypeAlias = SrunCommand +_SlurmCommandType: TypeAlias = SrunCommand """Types that are capable of being launched by the `SlurmLauncher`""" diff --git a/tests/test_slurm_launcher.py b/tests/test_slurm_launcher.py index a4152ac8b..70c8102c0 100644 --- a/tests/test_slurm_launcher.py +++ b/tests/test_slurm_launcher.py @@ -126,6 +126,7 @@ def test_slurm_launcher_can_start_a_command(monkeypatch, make_srun_command): ) launcher = SlurmLauncher() srun = make_srun_command(["-N", "1", "-n", "1"], ["echo", "spam", "eggs"]) + monkeypatch.setattr(srun, "start", lambda *_, **__: ...) id_ = launcher.start(srun) info = launcher._launched[id_] assert info.slurm_id == "mock-step-id" @@ -145,6 +146,7 @@ def test_slurm_launcher_errors_if_cannot_parse_id(monkeypatch, make_srun_command monkeypatch.setattr(time, "sleep", lambda *_, **__: ...) launcher = SlurmLauncher() srun = make_srun_command(["-N", "1", "-n", "1"], ["echo", "spam", "eggs"]) + monkeypatch.setattr(srun, "start", lambda *_, **__: ...) with pytest.raises( errors.LauncherError, match=r"Could not find id of launched job step" ):