test: add loadtest unit and integration coverage

Author: burakemirsezen

tests/integration/test_loadtest_e2e.py

@@ -0,0 +1,96 @@
+import os
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+import firecrest as f7t
+import pytest
+
+from swiss_ai_model_launch.launchers.firecrest_launcher import FirecRESTLauncher
+from swiss_ai_model_launch.loadtest.cluster import ClusterLoadtestConfig, submit_cluster_loadtest
+from swiss_ai_model_launch.loadtest.models import LoadtestConfig, ServerConfig
+from swiss_ai_model_launch.loadtest.setup import DEFAULT_CLUSTER_CONTAINER_IMAGE, K6_SCRIPT
+from tests.integration.utils import wait_for_job_running
+
+_LOADTEST_RUNNING_TIMEOUT_MIN = 20
+_LOADTEST_SERVER_URL = "https://api.swissai.svc.cscs.ch"
+
+_REQUIRED_ENV_VARS = [
+    "SML_CSCS_API_KEY",
+    "SML_FIRECREST_CLIENT_ID",
+    "SML_FIRECREST_CLIENT_SECRET",
+    "SML_FIRECREST_SYSTEM",
+    "SML_FIRECREST_TOKEN_URI",
+    "SML_FIRECREST_URL",
+    "SML_LOADTEST_MODEL",
+    "SML_LOADTEST_PROMPTS_FILE",
+    "SML_PARTITION",
+    "SML_RESERVATION",
+]
+
+
+@pytest.fixture(scope="function")  # type: ignore[misc]
+def env() -> dict[str, str]:
+    missing = [v for v in _REQUIRED_ENV_VARS if os.environ.get(v) is None]
+    if missing:
+        pytest.fail(
+            "Missing required environment variables: " + ", ".join(missing),
+            pytrace=False,
+        )
+    return {v: os.environ[v] for v in _REQUIRED_ENV_VARS}
+
+
+@pytest.fixture(scope="function")  # type: ignore[misc]
+async def launcher(env: dict[str, str]) -> AsyncIterator[FirecRESTLauncher]:
+    client = f7t.v2.AsyncFirecrest(
+        firecrest_url=env["SML_FIRECREST_URL"],
+        authorization=f7t.ClientCredentialsAuth(
+            client_id=env["SML_FIRECREST_CLIENT_ID"],
+            client_secret=env["SML_FIRECREST_CLIENT_SECRET"],
+            token_uri=env["SML_FIRECREST_TOKEN_URI"],
+            min_token_validity=90,
+        ),
+    )
+    try:
+        yield await FirecRESTLauncher.from_client(
+            client=client,
+            system_name=env["SML_FIRECREST_SYSTEM"],
+            partition=env["SML_PARTITION"],
+            reservation=env["SML_RESERVATION"] or None,
+        )
+    finally:
+        await client.close_session()
+
+
+@pytest.mark.comprehensive
+async def test_submit_cluster_loadtest_starts_cluster_job(
+    launcher: FirecRESTLauncher,
+    env: dict[str, str],
+    tmp_path: Path,
+) -> None:
+    submission = await submit_cluster_loadtest(
+        launcher=launcher,
+        server=ServerConfig(
+            url=os.environ.get("SML_LOADTEST_SERVER_URL", _LOADTEST_SERVER_URL),
+            api_key=env["SML_CSCS_API_KEY"],
+            model=env["SML_LOADTEST_MODEL"],
+            is_swissai=True,
+        ),
+        bench=LoadtestConfig(
+            scenario=os.environ.get("SML_LOADTEST_SCENARIO", "throughput"),
+            think_time="0",
+            max_tokens=os.environ.get("SML_LOADTEST_MAX_TOKENS", "16"),
+        ),
+        k6_script=K6_SCRIPT,
+        prompts_file=Path(env["SML_LOADTEST_PROMPTS_FILE"]),
+        summary_path=tmp_path / "summary.json",
+        cluster=ClusterLoadtestConfig(
+            container_image=str(DEFAULT_CLUSTER_CONTAINER_IMAGE),
+            wait=False,
+            reservation=env["SML_RESERVATION"] or None,
+        ),
+    )
+
+    try:
+        await wait_for_job_running(launcher, submission.job_id, _LOADTEST_RUNNING_TIMEOUT_MIN)
+    finally:
+        await launcher.cancel_job(submission.job_id)

tests/unit/test_cluster_loadtest.py

@@ -3,15 +3,12 @@
 
 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, ServerConfig
+from swiss_ai_model_launch.loadtest.models import LoadtestConfig
 
 
 @pytest.fixture
@@ -150,42 +147,3 @@ 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")