"""ssh 引数とサーバ上で実行するコマンドの組み立てのテスト（SPEC_03「接続」「取得」「run」）.""" from __future__ import annotations import pytest from tias_judge_model.remote.commands import ( build_job_env, build_job_script, build_probe_command, build_ssh_argv, format_git_commit, validate_job_name, ) def _argv(jump: str | None, timeout: int = 15) -> list[str]: return build_ssh_argv( "ssh", user="hasegawa", address="192.168.23.14", jump_host=jump, connect_timeout_sec=timeout, remote_command="echo hi", ) def _has_option(argv: list[str], value: str) -> bool: return any(argv[i] == "-o" and argv[i + 1] == value for i in range(len(argv) - 1)) class TestBuildSshArgv: def test_with_jump_host(self) -> None: argv = _argv("labpc") assert argv[0] == "ssh" j = argv.index("-J") assert argv[j + 1] == "labpc" assert "hasegawa@192.168.23.14" in argv assert argv[-1] == "echo hi" # 接続先より前にオプションがある assert j < argv.index("hasegawa@192.168.23.14") def test_without_jump_host(self) -> None: argv = _argv(None) assert "-J" not in argv assert "hasegawa@192.168.23.14" in argv @pytest.mark.parametrize("jump", ["labpc", None]) def test_batch_mode_and_timeout_always(self, jump: str | None) -> None: argv = _argv(jump, timeout=7) assert _has_option(argv, "BatchMode=yes") assert _has_option(argv, "ConnectTimeout=7") def test_remote_command_is_single_argument(self) -> None: argv = build_ssh_argv( "ssh", user="u", address="h", jump_host=None, connect_timeout_sec=5, remote_command="cd ~/x; echo a b", ) assert argv[-1] == "cd ~/x; echo a b" @pytest.mark.parametrize("bad", ["-oProxyCommand=evil", "a b", ""]) def test_rejects_option_like_destination(self, bad: str) -> None: with pytest.raises(ValueError): build_ssh_argv( "ssh", user="u", address=bad, jump_host=None, connect_timeout_sec=5, remote_command="true", ) class TestBuildProbeCommand: def test_queries_listed_in_spec(self) -> None: command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) assert "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu" in command assert "--query-compute-apps=gpu_uuid,pid,used_memory" in command assert "--format=csv,noheader,nounits" in command def test_observation_uses_interval_and_duration(self) -> None: command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) assert "-lms 500" in command assert "timeout 10" in command assert "pmon -s u" in command def test_observe_zero_does_not_observe(self) -> None: command = build_probe_command("~/work/x", observe_sec=0, sample_interval_ms=500) assert "-lms" not in command assert "pmon" not in command assert "--query-gpu=index,uuid,name" in command def test_is_single_line(self) -> None: command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) assert "\n" not in command class TestJobSettings: @pytest.mark.parametrize("name", ["a", "exp1", "Exp_1-b", "x" * 64]) def test_valid_job_names(self, name: str) -> None: assert validate_job_name(name) == name @pytest.mark.parametrize("name", ["", "_a", "-a", "a b", "a.b", "a:b", "x" * 65, "a/b"]) def test_invalid_job_names(self, name: str) -> None: with pytest.raises(ValueError): validate_job_name(name) def test_job_env_contains_selected_gpu_and_extra(self) -> None: env = build_job_env(2, ["FOO=bar", "EMPTY="]) assert env["CUDA_VISIBLE_DEVICES"] == "2" assert env["FOO"] == "bar" assert env["EMPTY"] == "" class TestGitCommit: """`TIAS_GIT_COMMIT`（SPEC_03「run」・SPEC_01「実験の出力」）.""" def test_clean_tree_passes_short_hash(self) -> None: # git rev-parse --short HEAD の出力（末尾改行付き） assert format_git_commit("abc1234\n", dirty=False) == "abc1234" def test_dirty_tree_appends_suffix(self) -> None: assert format_git_commit("abc1234\n", dirty=True) == "abc1234-dirty" def test_job_env_contains_git_commit(self) -> None: env = build_job_env(0, [], git_commit="abc1234-dirty") assert env["TIAS_GIT_COMMIT"] == "abc1234-dirty" def test_job_env_without_git_commit_does_not_set_it(self) -> None: env = build_job_env(0, []) assert "TIAS_GIT_COMMIT" not in env @pytest.mark.parametrize("git_commit", ["abc1234", None]) def test_env_option_cannot_override_git_commit(self, git_commit: str | None) -> None: # git が使えず渡さないときも --env で指定させない（実験の記録を迂回させないため） with pytest.raises(ValueError): build_job_env(0, ["TIAS_GIT_COMMIT=fake"], git_commit=git_commit) class TestReservedEnv: @pytest.mark.parametrize( "key", [ "CUDA_VISIBLE_DEVICES", "TIAS_DATA_ROOT", "CUDA_DEVICE_ORDER", "PYTHONUNBUFFERED", "TIAS_GIT_COMMIT", ], ) @pytest.mark.parametrize("gpu_index", [0, None]) def test_reserved_keys_cannot_be_overridden(self, key: str, gpu_index: int | None) -> None: with pytest.raises(ValueError): build_job_env(gpu_index, [f"{key}=x"], git_commit="abc1234") def test_fixed_values(self) -> None: env = build_job_env(3, [], git_commit="abc1234") assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" assert env["PYTHONUNBUFFERED"] == "1" assert env["CUDA_VISIBLE_DEVICES"] == "3" class TestCpuJob: """`run --cpu`（SPEC_03「run」）. CUDA_VISIBLE_DEVICES を空文字で渡し，ログの開始行の GPU は「なし（CPU）」とする． """ def test_cpu_job_env_hides_gpus(self) -> None: env = build_job_env(None, ["FOO=bar"], git_commit="abc1234") assert "CUDA_VISIBLE_DEVICES" in env assert env["CUDA_VISIBLE_DEVICES"] == "" assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" assert env["TIAS_GIT_COMMIT"] == "abc1234" assert env["FOO"] == "bar" def _script(self, gpu_index: int | None) -> str: return build_job_script( "~/work/TIASJudgeModel", "survey", ["tias-dataset", "survey", "--config", "configs/newton.yaml"], uv_path="~/.local/bin/uv", server_name="Newton", gpu_index=gpu_index, env=build_job_env(gpu_index, [], git_commit="abc1234"), path_env={"TIAS_DATA_ROOT": "~/data/TIASJudgeModel"}, ) def test_cpu_job_log_header_says_no_gpu(self) -> None: script = self._script(None) assert "gpu=なし（CPU）" in script assert "gpu=None" not in script # 空文字として export する（未設定ではなく空にして GPU を見えなくする） assert "export CUDA_VISIBLE_DEVICES=''" in script def test_gpu_job_log_header_has_index(self) -> None: script = self._script(1) assert "gpu=1" in script assert "なし（CPU）" not in script assert "export CUDA_VISIBLE_DEVICES=1" in script assert "export TIAS_GIT_COMMIT=abc1234" in script