Loadtest job length setting

Author: burakemirsezen

src/swiss_ai_model_launch/cli/main.py

@@ -47,6 +47,7 @@
 _DEFAULT_LOADTEST_READY_POLL_SECONDS = 10
 _LOADTEST_READY_PROGRESS_SECONDS = 300
 _DEFAULT_LOADTEST_METRICS_REMOTE_WRITE_URL = "https://prometheus-dev.swissai.svc.cscs.ch/api/v1/write"
+_DEFAULT_LOADTEST_JOB_TIME = "00:30:00"
 
 
 def _make_firecrest_launcher_config(
@@ -178,6 +179,13 @@ def _add_loadtest_arguments(
         default=True,
         help="Wait for the cluster loadtest job and download/copy the summary (default: true).",
     )
+    parser.add_argument(
+        "--loadtest-job-time",
+        dest="loadtest_job_time",
+        default=_DEFAULT_LOADTEST_JOB_TIME,
+        metavar="HH:MM:SS",
+        help=f"SLURM time limit for the cluster k6 loadtest job (default: {_DEFAULT_LOADTEST_JOB_TIME}).",
+    )
     parser.add_argument(
         "--loadtest-server-url",
         dest="loadtest_server_url",
@@ -870,8 +878,11 @@ def _make_cluster_loadtest_config(
         raise ValueError("--cancel-after-loadtest requires --wait-for-loadtest.")
     if args.loadtest_ready_timeout <= 0:
         raise ValueError("--loadtest-ready-timeout must be greater than 0.")
+    if not re.fullmatch(r"[0-9]{1,2}:[0-5][0-9]:[0-5][0-9]", args.loadtest_job_time):
+        raise ValueError("--loadtest-job-time must be in HH:MM:SS format.")
     return ClusterLoadtestConfig(
         container_image=str(DEFAULT_CLUSTER_CONTAINER_IMAGE),
+        time=args.loadtest_job_time,
         wait=args.wait_for_loadtest,
         reservation=reservation or getattr(args, "reservation", None),
         metrics_remote_write_url=(
@@ -899,7 +910,7 @@ async def _run_k6_on_cluster(
     print(f"Loadtest prompts file: {prompts_file}")
     if cluster_config.metrics_remote_write_url:
         print(f"Loadtest metrics remote write: {cluster_config.metrics_remote_write_url}")
-    job_id = await submit_cluster_loadtest(
+    submission = await submit_cluster_loadtest(
         launcher=launcher,
         server=server,
         bench=loadtest_config,
@@ -908,7 +919,8 @@ async def _run_k6_on_cluster(
         summary_path=summary_path,
         cluster=cluster_config,
     )
-    print(f"Cluster loadtest job submitted: {job_id}")
+    print(f"Cluster loadtest job submitted: {submission.job_id}")
+    print(f"Loadtest run label: {submission.run_label}")
     if cluster_config.wait:
         print(f"Loadtest summary: {summary_path}")
 

src/swiss_ai_model_launch/loadtest/cluster.py

@@ -28,6 +28,12 @@ class ClusterLoadtestConfig:
     metrics_remote_write_url: str | None = None
 
 
+@dataclass(frozen=True)
+class ClusterLoadtestSubmission:
+    job_id: int
+    run_label: str
+
+
 def build_cluster_loadtest_script(
     *,
     bench: LoadtestConfig,
@@ -205,7 +211,7 @@ async def submit_cluster_loadtest(
     prompts_file: Path,
     summary_path: Path,
     cluster: ClusterLoadtestConfig,
-) -> int:
+) -> ClusterLoadtestSubmission:
     run_label = f"loadtest_{bench.scenario}_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{create_salt(6)}"
     prompts_path, container_mounts = _container_mounts_for_external_prompts(run_label, prompts_file)
     script = build_cluster_loadtest_script(
@@ -242,7 +248,7 @@ async def submit_cluster_loadtest(
             if remote_summary.exists():
                 summary_path.parent.mkdir(parents=True, exist_ok=True)
                 shutil.copyfile(remote_summary, summary_path)
-        return job_id
+        return ClusterLoadtestSubmission(job_id=job_id, run_label=run_label)
 
     if isinstance(launcher, FirecRESTLauncher):
         firecrest_working_dir = str(launcher._get_working_dir())
@@ -294,6 +300,6 @@ async def submit_cluster_loadtest(
                 )
             except f7t.FirecrestException as e:
                 raise RuntimeError(f"Could not download cluster loadtest summary for job {job_id}: {e}") from e
-        return job_id
+        return ClusterLoadtestSubmission(job_id=job_id, run_label=run_label)
 
     raise TypeError(f"Cluster loadtests are not supported for launcher type {type(launcher).__name__}")

tests/unit/test_cluster_loadtest.py

@@ -3,12 +3,15 @@
 
 import pytest
 
+import swiss_ai_model_launch.loadtest.cluster as cluster_module
+from swiss_ai_model_launch.launchers import SlurmLauncher
 from swiss_ai_model_launch.loadtest.cluster import (
     ClusterLoadtestConfig,
     _container_mounts_for_external_prompts,
     build_cluster_loadtest_script,
+    submit_cluster_loadtest,
 )
-from swiss_ai_model_launch.loadtest.models import LoadtestConfig
+from swiss_ai_model_launch.loadtest.models import LoadtestConfig, ServerConfig
 
 
 @pytest.fixture
@@ -123,9 +126,9 @@ def test_script_k6_run_command_shell_quotes_dynamic_values(
     tokens = shlex.split(inner_command)
 
     assert tokens[tokens.index("--tag") + 1] == "scenario=throughput"
-    assert f"run_label=run-$USER's-label" in tokens
-    assert f"model=model-$USER's-name" in tokens
-    assert f"PROMPTS_FILE=/capstor/prompt files/$USER/prompts.json" in tokens
+    assert "run_label=run-$USER's-label" in tokens
+    assert "model=model-$USER's-name" in tokens
+    assert "PROMPTS_FILE=/capstor/prompt files/$USER/prompts.json" in tokens
     assert 'sh -lc "' not in script
 
 
@@ -147,3 +150,42 @@ def test_container_mounts_includes_top_level_dir() -> None:
 def test_container_mounts_prompts_path_returned() -> None:
     prompts_path, _ = _container_mounts_for_external_prompts("my_run", Path("/scratch/data/prompts.jsonl"))
     assert prompts_path == "/scratch/data/prompts.jsonl"
+
+
+@pytest.mark.asyncio
+async def test_submit_cluster_loadtest_returns_job_id_and_run_label(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    bench: LoadtestConfig,
+) -> None:
+    async def fake_run_checked(*cmd: str) -> str:
+        assert cmd[0] == "sbatch"
+        return "Submitted batch job 123\n"
+
+    monkeypatch.setenv("HOME", str(tmp_path))
+    monkeypatch.setattr(cluster_module, "_run_checked", fake_run_checked)
+    monkeypatch.setattr(cluster_module, "create_salt", lambda length: "X" * length)
+
+    k6_script = tmp_path / "script.js"
+    k6_script.write_text("export default function() {}\n")
+    prompts_file = tmp_path / "prompts.jsonl"
+    prompts_file.write_text("{}\n")
+
+    submission = await submit_cluster_loadtest(
+        launcher=SlurmLauncher(
+            system_name="test-system",
+            username="test-user",
+            account="test-account",
+            partition="test-partition",
+        ),
+        server=ServerConfig(url="https://example.test", api_key="secret", model="test-model", is_swissai=True),
+        bench=bench,
+        k6_script=k6_script,
+        prompts_file=prompts_file,
+        summary_path=tmp_path / "summary.json",
+        cluster=ClusterLoadtestConfig(container_image="/images/k6.sqsh", wait=False),
+    )
+
+    assert submission.job_id == 123
+    assert submission.run_label.startswith("loadtest_throughput_")
+    assert submission.run_label.endswith("_XXXXXX")

tests/unit/test_loadtest_cli.py

@@ -36,11 +36,11 @@ def test_loadtest_run_help_excludes_removed_scenario_owned_flags(capsys: pytest.
     assert "--loadtest-ignore-eos" in help_text
     assert "--wait-until-healthy" in help_text
     assert "--loadtest-metrics-remote-write" in help_text
+    assert "--loadtest-job-time" in help_text
     assert "--job-id" not in help_text
     assert "--loadtest-chat-mode" not in help_text
     assert "--loadtest-cpus-per-task" not in help_text
     assert "--loadtest-k6-script" not in help_text
-    assert "--loadtest-job-time" not in help_text
     assert "--loadtest-ready-poll-interval" not in help_text
     assert "--loadtest-think-time" not in help_text
     assert "--loadtest-request-timeout" not in help_text
@@ -143,7 +143,6 @@ def test_loadtest_parser_does_not_expose_api_key_override() -> None:
         ("--no-loadtest-chat-mode", None),
         ("--job-id", "123"),
         ("--loadtest-cpus-per-task", "8"),
-        ("--loadtest-job-time", "01:00:00"),
         ("--loadtest-k6-script", "script.js"),
         ("--loadtest-ready-poll-interval", "30"),
         ("--loadtest-think-time", "0"),
@@ -174,6 +173,28 @@ def test_loadtest_metrics_remote_write_enabled_by_default() -> None:
     )
 
 
+def test_loadtest_job_time_defaults_to_cluster_default() -> None:
+    parser = _build_parser()
+    args = parser.parse_args(["loadtest", "run"])
+
+    assert _make_cluster_loadtest_config(args).time == "00:30:00"
+
+
+def test_loadtest_job_time_can_be_overridden() -> None:
+    parser = _build_parser()
+    args = parser.parse_args(["loadtest", "run", "--loadtest-job-time", "01:00:00"])
+
+    assert _make_cluster_loadtest_config(args).time == "01:00:00"
+
+
+def test_loadtest_job_time_rejects_invalid_format() -> None:
+    parser = _build_parser()
+    args = parser.parse_args(["loadtest", "run", "--loadtest-job-time", "1h"])
+
+    with pytest.raises(ValueError, match="--loadtest-job-time"):
+        _make_cluster_loadtest_config(args)
+
+
 def test_loadtest_metrics_remote_write_can_be_disabled() -> None:
     parser = _build_parser()
     args = parser.parse_args(["loadtest", "run", "--no-loadtest-metrics-remote-write"])