"""Launcher tests are offline; no scheduler, SSH or GPU calls are made.""" from __future__ import annotations import importlib.util import json from pathlib import Path import shlex import signal import tarfile from types import SimpleNamespace import pytest SCRIPT = Path(__file__).resolve().parents[1] / "scripts/run-nemo-long-queued.py" spec = importlib.util.spec_from_file_location("nemo_queue_launcher", SCRIPT) launcher = importlib.util.module_from_spec(spec) spec.loader.exec_module(launcher) @pytest.fixture(autouse=True) def explicit_fixture_allocation(monkeypatch): monkeypatch.setattr(launcher, "OWNED_SHELL", "80818081-8081-4081-8081-808180818081") def test_scheduler_filter_never_returns_secret_keys(): raw = [{"id": launcher.OWNED_SHELL, "state": "STATE_RUNNING", "username": "owner", "resourcePool": "42_Compute", "workspaceId": 123, "privateKey": "DO_NOT_PRINT_PRIVATE_KEY", "container": {"token": "DO_NOT_PRINT_TOKEN"}}, {"id": "somebody-else", "privateKey": "ANOTHER_PRIVATE_KEY"}] snapshot = launcher.owned_snapshot(raw) assert snapshot["shell_id"] == launcher.OWNED_SHELL and snapshot["state"] == "RUNNING" assert "PRIVATE_KEY" not in json.dumps(snapshot) and "TOKEN" not in json.dumps(snapshot) with pytest.raises(launcher.LaunchError): launcher.owned_snapshot(raw[1:]) def test_scheduler_query_is_owned_list_only_and_captured(monkeypatch): calls = [] def run(command, **kwargs): calls.append(command) assert kwargs["capture_output"] and kwargs["timeout"] <= 30 return SimpleNamespace(returncode=0, stdout=json.dumps([{ "id": launcher.OWNED_SHELL, "state": "QUEUED", "username": "owner", "resourcePool": "42_Compute", "workspaceId": 123, "privateKey": "SECRET"}])) monkeypatch.setattr(launcher.subprocess, "run", run) assert launcher.scheduler_snapshot({})["state"] == "QUEUED" assert calls == [["det", "shell", "list", "--json"]] def ssh_command(cache, shell=None, proxy_extra=""): cache.mkdir(exist_ok=True) key = cache / "key" key.write_text("PRIVATE_KEY_FIXTURE") key.chmod(0o600) proxy = "/usr/bin/python3 -m determined.cli.tunnel https://cluster.example %h" + proxy_extra return shlex.join(["ssh", "-o", "ProxyCommand=" + proxy, "-o", "StrictHostKeyChecking=no", "-tt", "-o", "IdentitiesOnly=yes", "-i", str(key), "root@" + (shell or launcher.OWNED_SHELL)]) def test_connection_has_private_cache_paths_and_no_key_content(tmp_path): command = ssh_command(tmp_path) configuration = launcher.connection_config(command, tmp_path) assert launcher.OWNED_SHELL in configuration and "RequestTTY no" in configuration assert "PRIVATE_KEY_FIXTURE" not in configuration assert str(tmp_path / "known_hosts") in configuration @pytest.mark.parametrize("variant", ["other_shell", "external_key", "proxy_injection", "world_readable_key"]) def test_connection_rejects_unexpected_targets_keys_and_proxy(tmp_path, variant): cache = tmp_path / "cache" command = ssh_command(cache, shell="different-shell" if variant == "other_shell" else None, proxy_extra=" ; touch /tmp/unwanted" if variant == "proxy_injection" else "") if variant == "external_key": command = command.replace(str(cache / "key"), str(tmp_path / "elsewhere")) if variant == "world_readable_key": (cache / "key").chmod(0o644) with pytest.raises(launcher.LaunchError): launcher.connection_config(command, cache) @pytest.mark.parametrize("name,kind", [("../escape", tarfile.REGTYPE), ("/escape", tarfile.REGTYPE), ("link", tarfile.SYMTYPE), ("device", tarfile.CHRTYPE)]) def test_tar_staging_rejects_traversal_and_non_files(name, kind): member = tarfile.TarInfo(name) member.type = kind with pytest.raises(ValueError, match="unsafe staging"): launcher.validate_tar_member(member) def test_node_payload_compiles_and_preflights_before_experiment(): configuration = {"plan": "path with ' quote", "plan_sha256": "a" * 64} payload = launcher.node_program(configuration) compile(payload, "", "exec") prefix = payload.split("import hashlib,", 1)[0] namespace = {} exec(prefix, namespace) assert namespace["configuration"] == configuration assert payload.index("scripts/gpu-preflight.py") < payload.index('"hessian_recipes.gym_experiment", "run"') assert 'if (run / "experiment-status.json").exists()' in payload assert 'command.append("--resume")' in payload assert '"torch/lib", "nvidia", "triton"' in payload @pytest.fixture def launch_setup(tmp_path, monkeypatch): monkeypatch.setattr(launcher, "ROOT", tmp_path) root = tmp_path / "runs/experiment" root.mkdir(parents=True) configuration = {"plan": str(root / "experiment-plan.json"), "plan_sha256": "a" * 64, "launcher_sha256": "b" * 64, "budget_seconds": 60, "gym_python": "/fixture/gym/python"} calls = SimpleNamespace(released=0, ssh=0, payload=None) monkeypatch.setattr(launcher, "verify_inputs", lambda *args: configuration) monkeypatch.setattr(launcher, "prepare_connection", lambda environment, cache: cache / "fixture.conf") def release(environment): calls.released += 1 return True monkeypatch.setattr(launcher, "release_owned", release) class SSH: returncode = 0 def __init__(self, command, **kwargs): calls.ssh += 1 assert command[0] == "ssh" and launcher.OWNED_SHELL not in command[-1] assert kwargs["start_new_session"] def communicate(self, payload, timeout): assert 0 < timeout <= 60 calls.payload = payload def poll(self): return self.returncode monkeypatch.setattr(launcher.subprocess, "Popen", SSH) return configuration, calls, root def test_success_releases_exact_existing_allocation(launch_setup, monkeypatch, capsys): configuration, calls, root = launch_setup monkeypatch.setattr(launcher, "scheduler_snapshot", lambda env: { "shell_id": launcher.OWNED_SHELL, "state": "RUNNING", "username": "owner"}) assert launcher.launch(configuration) == 0 assert calls.released == calls.ssh == 1 status = json.loads((root / "queue-launcher-status.json").read_text()) assert status["result"] == "completed" and status["release_acknowledged"] assert "PRIVATE_KEY" not in capsys.readouterr().out @pytest.mark.parametrize("state,expected_releases", [("QUEUED", 0), ("PULLING", 1)]) def test_interrupt_preserves_queue_but_releases_active_allocation(launch_setup, monkeypatch, state, expected_releases): configuration, calls, root = launch_setup monkeypatch.setattr(launcher, "scheduler_snapshot", lambda env: { "shell_id": launcher.OWNED_SHELL, "state": state, "username": "owner"}) monkeypatch.setattr(launcher.time, "sleep", lambda delay: (_ for _ in ()).throw(launcher.Interrupted(signal.SIGTERM))) assert launcher.launch(configuration) == 128 + signal.SIGTERM assert calls.released == expected_releases and calls.ssh == 0 status = json.loads((root / "queue-launcher-status.json").read_text()) assert status["result"] == "interrupted" def test_source_change_after_allocation_releases_without_ssh(launch_setup, monkeypatch): configuration, calls, root = launch_setup monkeypatch.setattr(launcher, "scheduler_snapshot", lambda env: { "shell_id": launcher.OWNED_SHELL, "state": "RUNNING", "username": "owner"}) monkeypatch.setattr(launcher, "verify_inputs", lambda *args: (_ for _ in ()).throw(ValueError("PRIVATE_KEY_SENTINEL"))) assert launcher.launch(configuration) == 1 assert calls.released == 1 and calls.ssh == 0 assert "PRIVATE_KEY_SENTINEL" not in (root / "queue-launcher-status.json").read_text() def test_gpu_budget_includes_image_pulling_before_ssh(launch_setup, monkeypatch): configuration, calls, root = launch_setup monkeypatch.setattr(launcher, "scheduler_snapshot", lambda env: { "shell_id": launcher.OWNED_SHELL, "state": "PULLING", "username": "owner"}) ticks = iter([0, 61]) monkeypatch.setattr(launcher.time, "monotonic", lambda: next(ticks)) assert launcher.launch(configuration) == 1 assert calls.released == 1 and calls.ssh == 0 status = json.loads((root / "queue-launcher-status.json").read_text()) assert status["error_code"] == "active_gpu_budget_exhausted" def test_release_only_uses_the_fixed_owned_shell(monkeypatch): calls = [] def run(command, **kwargs): calls.append(command) return SimpleNamespace(returncode=0) monkeypatch.setattr(launcher.subprocess, "run", run) assert launcher.release_owned({}) assert calls == [["det", "shell", "kill", launcher.OWNED_SHELL]] def test_unset_controller_cannot_contact_or_release_any_allocation(monkeypatch): monkeypatch.setattr(launcher, "OWNED_SHELL", None) def forbidden(*args, **kwargs): raise AssertionError("Scheduler/SSH must not be contacted without an explicit allocation") monkeypatch.setattr(launcher.subprocess, "run", forbidden) with pytest.raises(launcher.LaunchError, match="explicit_owned_shell_id_required"): launcher.launch({}) with pytest.raises(launcher.LaunchError, match="explicit_owned_shell_id_required"): launcher.release_owned({}) def test_explicit_shell_from_another_workspace_is_rejected_before_ssh(): with pytest.raises(launcher.LaunchError, match="jaif"): launcher.owned_snapshot([{"id": launcher.OWNED_SHELL, "state": "RUNNING", "username": "owner", "resourcePool": "42_Compute", "workspaceId": 66}])