"""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")
