"""GPU の空き判定（4 段階）と選択のテスト（SPEC_03「GPU の空き判定と選択」）. - 観測中の最大値で閾値を判定する（瞬間値が 0% でも観測中に超えたら使用中） - 計算プロセスを所有者を問わず数える．一覧が取れなければ閾値内の GPU は「空き（利用者あり）」 - GPU の行が解釈できない・サンプル 0 件なら「状態不明」 - 選択では「未使用」をサーバの優先順位より優先する """ from __future__ import annotations import pytest from probe_helpers import PMON_HEADER, gpu_line, probe_text, sample_lines from tias_judge_model.remote.config import GpuFreeThreshold, ServerConfig from tias_judge_model.remote.gpu import ( GpuInfo, GpuSelection, GpuState, ProbeOutput, ServerProbe, check_cpu_options, check_selection_options, classify_gpu, gpu_users, parse_probe_output, select_cpu_server, select_gpu, ) from tias_judge_model.remote.observe import UsagePattern OBSERVE = GpuFreeThreshold( max_memory_used_mib=1024, max_utilization_percent=10, observe_sec=10, sample_interval_ms=500 ) INSTANT = GpuFreeThreshold( max_memory_used_mib=1024, max_utilization_percent=10, observe_sec=0, sample_interval_ms=500 ) UUID0 = "GPU-00000000-aaaa-bbbb-cccc-000000000000" UUID1 = "GPU-11111111-aaaa-bbbb-cccc-111111111111" def _states(output: ProbeOutput) -> dict[int | None, GpuState]: return {gpu.index: gpu.state for gpu in output.gpus} def _gpu(output: ProbeOutput, index: int) -> GpuInfo: return next(g for g in output.gpus if g.index == index) # --------------------------------------------------------------------------- # 判定（classify_gpu） # --------------------------------------------------------------------------- class TestClassifyGpu: def test_within_threshold_without_processes_is_unused(self) -> None: assert classify_gpu(0, 0, OBSERVE, 0) is GpuState.UNUSED def test_within_threshold_with_processes_is_shared(self) -> None: assert classify_gpu(500, 5, OBSERVE, 1) is GpuState.SHARED def test_unknown_process_list_is_shared_not_unused(self) -> None: assert classify_gpu(0, 0, OBSERVE, None) is GpuState.SHARED @pytest.mark.parametrize( ("memory", "util", "expected"), [ (1023, 9, GpuState.UNUSED), (1024, 9, GpuState.BUSY), # 「未満」なのでちょうどは超えた扱い (1023, 10, GpuState.BUSY), (5000, 0, GpuState.BUSY), (0, 100, GpuState.BUSY), ], ) def test_threshold_boundaries(self, memory: int, util: int, expected: GpuState) -> None: assert classify_gpu(memory, util, OBSERVE, 0) is expected def test_busy_even_if_process_list_unknown(self) -> None: assert classify_gpu(0, 50, OBSERVE, None) is GpuState.BUSY # --------------------------------------------------------------------------- # probe 出力の解析と判定（parse_probe_output） # --------------------------------------------------------------------------- class TestObservedJudgement: def test_max_during_observation_is_used_even_if_instant_is_zero(self) -> None: # 瞬間値は 0% だが，観測中に 74% まで上がった（2026-09-27 Newton の実測を模す） text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 0, 74, 0, 0, 30, 0, 0, 0, 0]), ) output = parse_probe_output(text, OBSERVE) assert _states(output) == {0: GpuState.BUSY} def test_memory_peak_during_observation_makes_busy(self) -> None: samples = ["0, 0, 10", "0, 0, 2000", "0, 0, 10"] text = probe_text([gpu_line(0, UUID0, used=10, util=0)], samples=samples) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.BUSY} def test_always_within_threshold_without_processes_is_unused(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 3, 9, 0, 1]) ) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNUSED} def test_single_sample_over_threshold_is_busy(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 0, 10, 0]) ) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.BUSY} def test_no_samples_for_gpu_is_unknown(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=0)], samples=sample_lines(0, [0, 0, 0]), ) assert _states(parse_probe_output(text, OBSERVE)) == { 0: GpuState.UNUSED, 1: GpuState.UNKNOWN, } def test_missing_samples_section_makes_all_unknown(self) -> None: text = probe_text([gpu_line(0, UUID0, used=5, util=0)], samples=None) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNKNOWN} def test_observe_zero_uses_instant_value(self) -> None: text = probe_text([gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=50)]) output = parse_probe_output(text, INSTANT) assert _states(output) == {0: GpuState.UNUSED, 1: GpuState.BUSY} def test_observe_zero_does_not_require_samples(self) -> None: text = probe_text([gpu_line(0, UUID0, used=2000, util=0)]) assert _states(parse_probe_output(text, INSTANT)) == {0: GpuState.BUSY} class TestProcessesAndUnknown: def test_processes_are_matched_by_gpu_uuid(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=500, util=0), gpu_line(1, UUID1, used=5, util=0)], samples=sample_lines(0, [0, 0]) + sample_lines(1, [0, 0]), apps=[f"{UUID0}, 4321, 480"], owners=["4321 alice 99720"], ) output = parse_probe_output(text, OBSERVE) assert _states(output) == {0: GpuState.SHARED, 1: GpuState.UNUSED} process = _gpu(output, 0).processes[0] assert process.pid == 4321 assert process.user == "alice" assert process.elapsed_sec == 99720 assert _gpu(output, 1).processes == () def test_own_process_also_counts(self) -> None: # 自分の別ジョブが載っている GPU にも次のジョブを重ねない（所有者を問わず数える） text = probe_text( [gpu_line(0, UUID0, used=300, util=0)], samples=sample_lines(0, [0, 0]), apps=[f"{UUID0}, 100, 300"], owners=["100 hasegawa 10"], ) output = parse_probe_output(text, OBSERVE) gpu = _gpu(output, 0) assert gpu.state is GpuState.SHARED assert gpu_users(gpu, "hasegawa") == ["自分"] assert gpu_users(gpu, "someone_else") == ["hasegawa"] def test_process_list_failure_makes_within_threshold_gpus_shared(self) -> None: text = probe_text( [ gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=0), ], samples=sample_lines(0, [0, 0]) + sample_lines(1, [0, 90]), apps=[], apps_rc=9, ) output = parse_probe_output(text, OBSERVE) assert not output.processes_query_ok # 閾値内は未使用と誤認せず「空き（利用者あり）」，閾値を超えたものは使用中のまま assert _states(output) == {0: GpuState.SHARED, 1: GpuState.BUSY} def test_unparsable_process_list_makes_gpus_shared(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=0)], samples=sample_lines(0, [0]) + sample_lines(1, [0]), apps=["this is not csv"], ) output = parse_probe_output(text, OBSERVE) assert _states(output) == {0: GpuState.SHARED, 1: GpuState.SHARED} def test_missing_process_section_makes_gpus_shared(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0]), apps=None ) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.SHARED} @pytest.mark.parametrize( "line", [ f"0, {UUID0}, NVIDIA RTX A6000, [N/A], 49140, 0", f"0, {UUID0}, NVIDIA RTX A6000, 5, 49140, [N/A]", "0, garbage", "garbage", ], ) def test_unparsable_gpu_line_is_unknown(self, line: str) -> None: text = probe_text( [line, gpu_line(1, UUID1, used=5, util=0)], samples=sample_lines(0, [0]) + sample_lines(1, [0]), ) output = parse_probe_output(text, OBSERVE) assert output.gpus[0].state is GpuState.UNKNOWN assert _gpu(output, 1).state is GpuState.UNUSED def test_gpu_query_failure_yields_no_candidates(self) -> None: text = probe_text([], gpu_rc=127, apps=None) output = parse_probe_output(text, OBSERVE) assert not output.gpu_query_ok assert output.gpus == () def test_graphics_process_is_not_counted_from_pmon(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=300, util=5)], samples=sample_lines(0, [5, 0]), pmon=[ *PMON_HEADER, " 0 1234 C 74 10 - - - - python", " 0 1234 C 3 1 - - - - python", " 0 999 G 50 2 - - - - Xorg", " 0 - - - - - - - - -", ], apps=[f"{UUID0}, 1234, 280"], owners=["1234 alice 45"], ) output = parse_probe_output(text, OBSERVE) gpu = _gpu(output, 0) assert [p.pid for p in gpu.processes] == [1234] assert gpu.processes[0].max_sm_percent == 74 def test_preamble_before_markers_is_ignored(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0]), preamble=["Welcome to Ubuntu", "Last login: Sat Sep 27 10:00:00 2026"], ) assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNUSED} def test_truncated_output_is_error(self) -> None: text = probe_text([gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0])) truncated = text.split("@@tias-end")[0] with pytest.raises(ValueError): parse_probe_output(truncated, OBSERVE) def test_jobs_and_other_sessions(self) -> None: text = probe_text([], jobs=["exp1"], others=["mywork"], samples=[]) output = parse_probe_output(text, OBSERVE) assert output.jobs == ("exp1",) assert output.other_sessions == ("mywork",) class TestPatternDoesNotAffectState: def test_intermittent_but_within_threshold_is_unused(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 5, 0, 0, 0]) ) gpu = _gpu(parse_probe_output(text, OBSERVE), 0) assert gpu.pattern is UsagePattern.INTERMITTENT assert gpu.state is GpuState.UNUSED def test_constant_but_within_threshold_is_unused(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=5, util=3)], samples=sample_lines(0, [3, 3, 3, 3, 3]) ) gpu = _gpu(parse_probe_output(text, OBSERVE), 0) assert gpu.pattern is UsagePattern.CONSTANT assert gpu.state is GpuState.UNUSED def test_memory_only_pattern_is_shared(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=500, util=0)], samples=sample_lines(0, [0, 0, 0], memory=500), apps=[f"{UUID0}, 1, 500"], owners=["1 bob 100"], ) gpu = _gpu(parse_probe_output(text, OBSERVE), 0) assert gpu.pattern is UsagePattern.MEMORY_ONLY assert gpu.state is GpuState.SHARED def test_idle_pattern_but_memory_over_threshold_is_busy(self) -> None: text = probe_text( [gpu_line(0, UUID0, used=2000, util=0)], samples=sample_lines(0, [0, 0, 0], memory=2000), ) gpu = _gpu(parse_probe_output(text, OBSERVE), 0) assert gpu.pattern is UsagePattern.IDLE assert gpu.state is GpuState.BUSY # --------------------------------------------------------------------------- # 選択（select_gpu） # --------------------------------------------------------------------------- def _server(name: str, priority: int, *, data_ready: bool = True) -> ServerConfig: return ServerConfig( name=name, address="192.168.0.1", user="hasegawa", priority=priority, remote_dir="~/work/TIASJudgeModel", data_dir="~/data/TIASJudgeModel", data_ready=data_ready, ) def _info(index: int, state: GpuState) -> GpuInfo: return GpuInfo( index=index, uuid=f"GPU-{index}", name="GPU", memory_used_mib=0, memory_total_mib=49140, utilization_percent=0, state=state, processes_known=True, ) def _probe(server: ServerConfig, *states: GpuState, gpu_query_ok: bool = True) -> ServerProbe: gpus = tuple(_info(i, s) for i, s in enumerate(states)) return ServerProbe( server, ProbeOutput( jobs=(), other_sessions=(), gpu_query_ok=gpu_query_ok, processes_query_ok=True, gpus=gpus if gpu_query_ok else (), ), ) def _unreachable(server: ServerConfig) -> ServerProbe: return ServerProbe(server, None, "Connection timed out") U, S, B, X = GpuState.UNUSED, GpuState.SHARED, GpuState.BUSY, GpuState.UNKNOWN NEWTON = _server("Newton", 1) ARCHIMEDES = _server("Archimedes", 2) def _picked(selection: GpuSelection | None) -> tuple[str, int, GpuState] | None: if selection is None: return None return selection.server.name, selection.gpu_index, selection.state class TestSelectGpuAutomatic: def test_unused_beats_server_priority(self) -> None: # Newton（優先）は利用者ありのみ，Archimedes に未使用 → Archimedes を選ぶ probes = [_probe(NEWTON, S, S, B), _probe(ARCHIMEDES, B, U)] assert _picked(select_gpu(probes)) == ("Archimedes", 1, U) def test_priority_order_among_unused(self) -> None: probes = [_probe(ARCHIMEDES, U), _probe(NEWTON, B, U)] assert _picked(select_gpu(probes)) == ("Newton", 1, U) def test_same_priority_uses_config_order(self) -> None: first = _server("Alpha", 1) second = _server("Beta", 1) assert _picked(select_gpu([_probe(first, U), _probe(second, U)])) == ("Alpha", 0, U) assert _picked(select_gpu([_probe(second, U), _probe(first, U)])) == ("Beta", 0, U) def test_lowest_index_within_server(self) -> None: probe = _probe(NEWTON, B, U, U) assert _picked(select_gpu([probe])) == ("Newton", 1, U) def test_index_order_even_if_listed_out_of_order(self) -> None: output = ProbeOutput( jobs=(), other_sessions=(), gpu_query_ok=True, processes_query_ok=True, gpus=(_info(3, U), _info(1, U), _info(2, U)), ) assert _picked(select_gpu([ServerProbe(NEWTON, output)])) == ("Newton", 1, U) def test_shared_when_no_unused_by_priority(self) -> None: probes = [_probe(NEWTON, B, S), _probe(ARCHIMEDES, S)] assert _picked(select_gpu(probes)) == ("Newton", 1, S) def test_none_when_only_busy_or_unknown(self) -> None: probes = [_probe(NEWTON, B, X), _probe(ARCHIMEDES, X, B)] assert select_gpu(probes) is None def test_data_not_ready_excluded_by_default(self) -> None: archimedes = _server("Archimedes", 2, data_ready=False) probes = [_probe(NEWTON, S), _probe(archimedes, U)] assert _picked(select_gpu(probes)) == ("Newton", 0, S) assert _picked(select_gpu(probes, include_not_ready=True)) == ("Archimedes", 0, U) def test_only_data_not_ready_has_free_gpu(self) -> None: archimedes = _server("Archimedes", 2, data_ready=False) probes = [_probe(NEWTON, B), _probe(archimedes, U)] assert select_gpu(probes) is None assert _picked(select_gpu(probes, include_not_ready=True)) == ("Archimedes", 0, U) def test_server_option_restricts_candidates(self) -> None: probes = [_probe(NEWTON, U), _probe(ARCHIMEDES, B, S)] assert _picked(select_gpu(probes, server_name="Archimedes")) == ("Archimedes", 1, S) def test_server_option_none_available(self) -> None: probes = [_probe(NEWTON, U), _probe(ARCHIMEDES, B)] assert select_gpu(probes, server_name="Archimedes") is None def test_unreachable_server_is_skipped(self) -> None: probes = [_unreachable(NEWTON), _probe(ARCHIMEDES, S)] assert _picked(select_gpu(probes)) == ("Archimedes", 0, S) def test_server_with_failed_gpu_query_is_skipped(self) -> None: probes = [_probe(NEWTON, gpu_query_ok=False), _probe(ARCHIMEDES, S)] assert _picked(select_gpu(probes)) == ("Archimedes", 0, S) def test_unknown_server_name_is_error(self) -> None: with pytest.raises(ValueError): select_gpu([_probe(NEWTON, U)], server_name="Nowhere") class TestSelectGpuExplicit: @pytest.mark.parametrize("state", [U, S]) def test_unused_or_shared_is_accepted(self, state: GpuState) -> None: probes = [_probe(NEWTON, B, state)] selection = select_gpu(probes, server_name="Newton", gpu_index=1) assert _picked(selection) == ("Newton", 1, state) assert selection is not None and not selection.forced @pytest.mark.parametrize("state", [B, X]) def test_busy_or_unknown_is_error_without_force(self, state: GpuState) -> None: probes = [_probe(NEWTON, U, state)] with pytest.raises(ValueError): select_gpu(probes, server_name="Newton", gpu_index=1) @pytest.mark.parametrize("state", [B, X]) def test_busy_or_unknown_is_allowed_with_force(self, state: GpuState) -> None: probes = [_probe(NEWTON, U, state)] selection = select_gpu(probes, server_name="Newton", gpu_index=1, force=True) assert _picked(selection) == ("Newton", 1, state) assert selection is not None and selection.forced def test_explicit_gpu_even_if_other_server_has_unused(self) -> None: probes = [_probe(NEWTON, S), _probe(ARCHIMEDES, U)] assert _picked(select_gpu(probes, server_name="Newton", gpu_index=0)) == ("Newton", 0, S) def test_gpu_without_server_is_error(self) -> None: with pytest.raises(ValueError): select_gpu([_probe(NEWTON, U)], gpu_index=0) with pytest.raises(ValueError): check_selection_options(server_name=None, gpu_index=0, force=False) def test_gpu_with_server_passes_option_check(self) -> None: check_selection_options(server_name="Newton", gpu_index=0, force=False) check_selection_options(server_name="Newton", gpu_index=0, force=True) check_selection_options(server_name=None, gpu_index=None, force=False) # --------------------------------------------------------------------------- # run --cpu（SPEC_03「run」） # --------------------------------------------------------------------------- class TestCheckCpuOptions: def test_cpu_alone_is_accepted(self) -> None: check_cpu_options(gpu_index=None, force=False, observe_sec=None) @pytest.mark.parametrize( ("gpu_index", "force", "observe_sec"), [(0, False, None), (None, True, None), (None, False, 0), (None, False, 10), (1, True, 5)], ) def test_gpu_related_options_conflict( self, gpu_index: int | None, force: bool, observe_sec: int | None ) -> None: with pytest.raises(ValueError): check_cpu_options(gpu_index=gpu_index, force=force, observe_sec=observe_sec) class TestSelectCpuServer: def test_gpu_state_is_ignored(self) -> None: # GPU がすべて使用中・状態不明・一覧取得失敗でも，接続できれば選ぶ probes = [_probe(NEWTON, B, X), _probe(ARCHIMEDES, U)] assert select_cpu_server(probes).name == "Newton" probes = [_probe(NEWTON, gpu_query_ok=False), _probe(ARCHIMEDES, U)] assert select_cpu_server(probes).name == "Newton" def test_priority_order(self) -> None: # 設定の記載順ではなく priority 順 probes = [_probe(ARCHIMEDES, U), _probe(NEWTON, B)] assert select_cpu_server(probes).name == "Newton" def test_same_priority_uses_config_order(self) -> None: a, b = _server("A", 1), _server("B", 1) assert select_cpu_server([_probe(b, U), _probe(a, U)]).name == "B" def test_unreachable_is_skipped(self) -> None: probes = [_unreachable(NEWTON), _probe(ARCHIMEDES, B)] assert select_cpu_server(probes).name == "Archimedes" def test_all_unreachable_is_none(self) -> None: assert select_cpu_server([_unreachable(NEWTON), _unreachable(ARCHIMEDES)]) is None def test_data_not_ready_is_excluded_by_default(self) -> None: newton = _server("Newton", 1, data_ready=False) probes = [_probe(newton, U), _probe(ARCHIMEDES, B)] assert select_cpu_server(probes).name == "Archimedes" assert select_cpu_server([_probe(newton, U)]) is None def test_no_data_includes_not_ready(self) -> None: newton = _server("Newton", 1, data_ready=False) probes = [_probe(newton, U), _probe(ARCHIMEDES, B)] assert select_cpu_server(probes, include_not_ready=True).name == "Newton" def test_explicit_server(self) -> None: probes = [_probe(NEWTON, U), _probe(ARCHIMEDES, B)] assert select_cpu_server(probes, server_name="Archimedes").name == "Archimedes" def test_explicit_unreachable_server_is_error(self) -> None: probes = [_probe(NEWTON, U), _unreachable(ARCHIMEDES)] with pytest.raises(ValueError): select_cpu_server(probes, server_name="Archimedes") def test_explicit_not_ready_server_requires_no_data(self) -> None: newton = _server("Newton", 1, data_ready=False) probes = [_probe(newton, U), _probe(ARCHIMEDES, U)] with pytest.raises(ValueError): select_cpu_server(probes, server_name="Newton") assert select_cpu_server(probes, server_name="Newton", include_not_ready=True) is newton def test_unknown_server_is_error(self) -> None: with pytest.raises(ValueError): select_cpu_server([_probe(NEWTON, U)], server_name="Euler")