From 8a058021e870177172e4a56d43db178253f10deb Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 8 Oct 2026 22:47:39 +0800 Subject: [PATCH 1/2] feat(training): select recorded CUDA MPS daemon --- .../1-training/7-tensor_runtime_production.md | 16 ++- .../1-training/7-tensor_runtime_production.md | 15 +- src/unilab/scripts/train_offpolicy.py | 52 +++++++ src/unilab/training/cuda_mps_cli.py | 82 +++++++++++ src/unilab/training/cuda_process_sharing.py | 2 +- .../test_offpolicy_double_buffer_runner.py | 65 +++++++++ tests/training/test_cuda_mps_cli.py | 130 +++++++++++++++++- tests/training/test_cuda_process_sharing.py | 2 +- 8 files changed, 357 insertions(+), 7 deletions(-) diff --git a/docs/sphinx/source/en/2-user_guide/1-training/7-tensor_runtime_production.md b/docs/sphinx/source/en/2-user_guide/1-training/7-tensor_runtime_production.md index abf075e1e..b3579128b 100644 --- a/docs/sphinx/source/en/2-user_guide/1-training/7-tensor_runtime_production.md +++ b/docs/sphinx/source/en/2-user_guide/1-training/7-tensor_runtime_production.md @@ -107,11 +107,25 @@ On a single-GPU host, the defaults are already complete: ```bash uv run uni-cumps start -eval "$(uv run uni-cumps env)" uv run uni-cumps doctor uv run uni-cumps stop ``` +After `start`, a trainer request using `training.cuda_process_sharing=mps` +selects the sole live UniLab-recorded daemon covering its GPU. Shell +integration is optional: + +```bash +uv run uni-cumps start +uv run --extra mjwarp train --algo flashsac \ + --task g1_motion_tracking --sim mjwarp \ + training.cuda_process_sharing=mps +``` + +An explicit `CUDA_MPS_PIPE_DIRECTORY` in the caller environment still takes +precedence. If no matching live daemon exists, training fails closed before +environment/learner/collector construction and points to `uni-cumps start`. + Default values: | Setting | Default | diff --git a/docs/sphinx/source/zh_CN/2-user_guide/1-training/7-tensor_runtime_production.md b/docs/sphinx/source/zh_CN/2-user_guide/1-training/7-tensor_runtime_production.md index 12cf5b778..70cb487cc 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/1-training/7-tensor_runtime_production.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/1-training/7-tensor_runtime_production.md @@ -102,11 +102,24 @@ UniLab 提供显式的用户自有生命周期命令，但训练路径仍不会 ```bash uv run uni-cumps start -eval "$(uv run uni-cumps env)" uv run uni-cumps doctor uv run uni-cumps stop ``` +`start` 后，设置 `training.cuda_process_sharing=mps` 的 trainer 会自动选择 +覆盖该 GPU 的唯一 live UniLab-recorded daemon。Shell 集成变为可选： + +```bash +uv run uni-cumps start +uv run --extra mjwarp train --algo flashsac \ + --task g1_motion_tracking --sim mjwarp \ + training.cuda_process_sharing=mps +``` + +调用方显式提供的 `CUDA_MPS_PIPE_DIRECTORY` 仍然优先。若没有匹配的 live +daemon，训练会在 environment/learner/collector 构造前 fail closed，并指向 +`uni-cumps start`。 + 默认参数： | 设置 | 默认值 | diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py index efa9f69d1..f073f35ea 100644 --- a/src/unilab/scripts/train_offpolicy.py +++ b/src/unilab/scripts/train_offpolicy.py @@ -51,6 +51,10 @@ resolve_nan_guard_cfg, should_run_playback, ) +from unilab.training.cuda_mps_cli import ( + CudaMpsCliError, + select_environment_for_gpu, +) from unilab.training.cuda_process_sharing import ( CudaProcessSharingEvidence, probe_cuda_process_sharing, @@ -249,6 +253,11 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): str(cfg.training.sim_backend), rank_device, ) + _select_cuda_mps_daemon_environment( + getattr(cfg.training, "cuda_process_sharing", None), + str(cfg.training.sim_backend), + rank_device, + ) cuda_process_sharing = probe_cuda_process_sharing( getattr(cfg.training, "cuda_process_sharing", None), rank_device, @@ -365,6 +374,49 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): return runner +def _select_cuda_mps_daemon_environment( + requested: Any, + backend: str, + learner_device: str, +) -> None: + """Bind the sole live UniLab-recorded daemon in this trainer process. + + Explicit deployment environment wins. Otherwise an ``mps`` request resolves + the user-owned daemon record for the rank-local GPU before the fail-closed + probe. This never starts or stops a daemon and never mutates a parent shell. + """ + + del backend + if requested != "mps" or os.environ.get("CUDA_MPS_PIPE_DIRECTORY"): + return + try: + import torch + except ImportError: + return + if not torch.cuda.is_available(): + return + normalized_device = learner_device.strip().lower() + if normalized_device == "cuda": + index = 0 + else: + base, separator, index_text = normalized_device.partition(":") + if base != "cuda" or not separator or not index_text.isdigit(): + return + index = int(index_text) + if index >= int(torch.cuda.device_count()): + return + properties = torch.cuda.get_device_properties(index) + uuid = str(getattr(properties, "uuid", "")).strip() + if not uuid: + return + try: + environment = select_environment_for_gpu(uuid) + except CudaMpsCliError as exc: + raise ValueError(str(exc)) from exc + os.environ.setdefault("CUDA_MPS_PIPE_DIRECTORY", environment["CUDA_MPS_PIPE_DIRECTORY"]) + os.environ.setdefault("CUDA_MPS_LOG_DIRECTORY", environment["CUDA_MPS_LOG_DIRECTORY"]) + + def _attach_cuda_process_sharing_manifest( runner: Any, evidence: CudaProcessSharingEvidence, diff --git a/src/unilab/training/cuda_mps_cli.py b/src/unilab/training/cuda_mps_cli.py index 907c4d4ac..9daacc51a 100644 --- a/src/unilab/training/cuda_mps_cli.py +++ b/src/unilab/training/cuda_mps_cli.py @@ -42,6 +42,7 @@ _NAME_PATTERN = re.compile(r"^[a-z0-9][a-z0-9-]{0,63}$") _DEFAULT_CONTROL_COMMAND = ("nvidia-cuda-mps-control", "-d") _QUIT_COMMAND = ("nvidia-cuda-mps-control",) +_SERVER_LIST_QUERY_COMMAND = ("nvidia-cuda-mps-control", "get_server_list") _MAX_CONTROL_SOCKET_BYTES = 95 _INERT_CONTROL_ARTIFACTS = ("control_lock", "log") _GPU_QUERY_COMMAND = ( @@ -189,6 +190,9 @@ def _canonical_uuid(value: str, *, field: str = "uuid") -> str: raise CudaMpsCliError( f"MIG device {value!r} is not supported by uni-cumps yet; select a physical GPU UUID." ) + # Torch properties commonly expose the bare UUID without NVIDIA's GPU- tag. + if not uuid.upper().startswith("GPU-") and len(uuid) == 36 and uuid.count("-") == 4: + uuid = f"GPU-{uuid}" if not uuid.upper().startswith("GPU-"): raise CudaMpsCliError(f"Invalid NVIDIA GPU {field}: {value!r}; expected a GPU- UUID.") raw_uuid = uuid[4:] @@ -492,6 +496,40 @@ def _wait_for_control(pipe_directory: Path, timeout: float = 5.0) -> bool: return False +def _wait_for_control_query( + environment: Mapping[str, str], + run_command: RunCommand, + *, + timeout: float = 5.0, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + remaining = max(0.0, deadline - time.monotonic()) + try: + _run_command( + run_command, + _SERVER_LIST_QUERY_COMMAND, + timeout=remaining, + env=environment, + ) + except CudaMpsCliError: + time.sleep(min(0.05, max(0.0, deadline - time.monotonic()))) + continue + return True + return False + + +def _stop_unrecorded_control( + environment: Mapping[str, str], + run_command: RunCommand, +) -> bool: + try: + _run_command(run_command, _QUIT_COMMAND, input="quit\n", env=environment) + except CudaMpsCliError: + return False + return True + + def _remove_inert_control_artifacts(pipe_directory: Path) -> None: """Remove only artifacts left by a failed control launch; preserve sockets.""" @@ -837,6 +875,18 @@ def start_daemon( f"MPS control did not create a usable control pipe at {control}. " "Inspect CUDA_MPS_LOG_DIRECTORY and the host deployment." ) + if not _wait_for_control_query(env, run_command): + stopped = _stop_unrecorded_control(env, run_command) + if using_default_runtime: + _remove_orphan_default_control_sockets(pipe_directory) + tail = _control_log_tail(log_directory) + detail = f" Control log tail:\n{tail}" if tail is not None else "" + if not stopped: + detail += " The launched control daemon may still be running." + raise CudaMpsCliError( + "MPS control created a pipe but did not answer a readiness query. " + f"Inspect CUDA_MPS_LOG_DIRECTORY and the host deployment.{detail}" + ) process_identity = _control_daemon_identity(pipe_directory) if process_identity is None: raise CudaMpsCliError( @@ -880,6 +930,38 @@ def environment_for_daemon(daemon: DaemonRecord) -> dict[str, str]: } +def select_environment_for_gpu( + gpu_uuid: str, + *, + root: Path | None = None, + proc_stat: Callable[[int], tuple[int, int] | None] = _proc_stat, +) -> dict[str, str]: + """Select the sole live user-owned daemon environment for one GPU. + + This is launcher-facing data only. It does not mutate the current process, + start a daemon, stop a daemon, or provide stop authority. + """ + + canonical = _canonical_uuid(gpu_uuid) + matching = [ + daemon + for daemon in list_daemon_records(root) + if canonical in daemon.gpu_uuids and daemon_process_is_live(daemon, proc_stat=proc_stat) + ] + if not matching: + raise CudaMpsCliError( + f"No live UniLab-recorded CUDA MPS daemon covers GPU {canonical}. " + "Run `uni-cumps start` first or set CUDA_MPS_PIPE_DIRECTORY explicitly." + ) + if len(matching) > 1: + names = ", ".join(daemon.name for daemon in matching) + raise CudaMpsCliError( + f"Multiple live UniLab-recorded daemons cover GPU {canonical}: {names}. " + "Set CUDA_MPS_PIPE_DIRECTORY to the intended daemon explicitly." + ) + return environment_for_daemon(matching[0]) + + def stop_daemon( name: str, *, diff --git a/src/unilab/training/cuda_process_sharing.py b/src/unilab/training/cuda_process_sharing.py index 16daa2adf..e77d859be 100644 --- a/src/unilab/training/cuda_process_sharing.py +++ b/src/unilab/training/cuda_process_sharing.py @@ -179,7 +179,7 @@ def _server_evidence( environment["CUDA_MPS_PIPE_DIRECTORY"] = str(control_pipe.parent) try: result = run_command( - ["nvidia-cuda-mps-control", "get-server-list"], + ["nvidia-cuda-mps-control", "get_server_list"], text=True, capture_output=True, timeout=_MPS_CONTROL_QUERY_TIMEOUT_SEC, diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 0a930dbbb..6edcc1f54 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util +import os import platform import queue import socket @@ -215,6 +216,65 @@ def reject_factory(*args, **kwargs): module.build_runner(algo, cfg) +def test_cuda_process_sharing_selects_sole_recorded_daemon_before_probe( + monkeypatch: pytest.MonkeyPatch, +): + module = _offpolicy() + _cuda_torch_module(monkeypatch) + cfg = _offpolicy_cfg(["task=g1_walk_flat/mjwarp", "training.cuda_process_sharing=mps"]) + monkeypatch.delenv("CUDA_MPS_PIPE_DIRECTORY", raising=False) + monkeypatch.delenv("CUDA_MPS_LOG_DIRECTORY", raising=False) + selected: dict[str, str] = {} + + def fake_selector(requested, backend, learner_device): + selected.update( + { + "requested": requested, + "backend": backend, + "learner_device": learner_device, + "pipe": "/tmp/unilab-recorded/pipe", + "log": "/tmp/unilab-recorded/log", + } + ) + os.environ["CUDA_MPS_PIPE_DIRECTORY"] = selected["pipe"] + os.environ["CUDA_MPS_LOG_DIRECTORY"] = selected["log"] + + def reject_selector(*args, **kwargs): + del args, kwargs + raise AssertionError("explicit environment must skip daemon selection") + + monkeypatch.setattr(module, "_select_cuda_mps_daemon_environment", fake_selector) + monkeypatch.setattr( + module, + "probe_cuda_process_sharing", + lambda *_args, **kwargs: _FakeCudaProcessSharingEvidence(), + ) + monkeypatch.setattr(module, "registry_env_factory", lambda *args, **kwargs: _fake_env_factory) + monkeypatch.setattr( + module, + "configure_backend_process_device", + lambda _backend, device: device, + ) + import uni_rl.algos.sac.double_buffer as owner_module + + monkeypatch.setattr(owner_module, "SACLearner", _FakeLearner) + monkeypatch.setattr(owner_module, "DoubleBufferOffPolicyRunner", _FakeRunner) + module.build_runner("sac", cfg) + + assert selected == { + "requested": "mps", + "backend": "mjwarp", + "learner_device": "cuda:0", + "pipe": "/tmp/unilab-recorded/pipe", + "log": "/tmp/unilab-recorded/log", + } + + # Direct helper test documents explicit-environment precedence: deployment + # provided a pipe, so a selector that would otherwise reject is never used. + monkeypatch.setattr(module, "select_environment_for_gpu", reject_selector) + module._select_cuda_mps_daemon_environment("mps", "mjwarp", "cuda:0") + + def test_valid_cuda_process_sharing_evidence_enters_runner_manifest( monkeypatch: pytest.MonkeyPatch, ): @@ -233,6 +293,11 @@ def test_valid_cuda_process_sharing_evidence_enters_runner_manifest( "probe_cuda_process_sharing", lambda *args, **kwargs: evidence, ) + monkeypatch.setattr( + module, + "_select_cuda_mps_daemon_environment", + lambda *_args, **_kwargs: None, + ) import uni_rl.algos.sac.double_buffer as owner_module diff --git a/tests/training/test_cuda_mps_cli.py b/tests/training/test_cuda_mps_cli.py index b1b54aa2f..d325d75b7 100644 --- a/tests/training/test_cuda_mps_cli.py +++ b/tests/training/test_cuda_mps_cli.py @@ -103,6 +103,10 @@ def test_visible_gpus_canonicalizes_uuids_and_visibility() -> None: assert [(gpu.index, gpu.uuid) for gpu in gpus] == [(0, GPU_A)] +def test_bare_torch_gpu_uuid_is_canonicalized() -> None: + assert mps._canonical_uuid("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") == GPU_A + + def test_mig_selectors_are_explicitly_rejected() -> None: with pytest.raises(mps.CudaMpsCliError, match="MIG device"): mps.resolve_topology("MIG-abc", run_command=fake_run()) @@ -244,6 +248,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) return Completed("") + if command == list(mps._SERVER_LIST_QUERY_COMMAND): + return Completed("") raise AssertionError(f"unexpected command {command}") record = mps.start_daemon( @@ -264,6 +270,91 @@ def run(command: list[str], **kwargs: Any) -> Completed: assert (linux_host / "prod" / "daemon.json").is_file() +def test_start_daemon_waits_until_control_answers_query( + linux_host: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(mps.shutil, "which", lambda _name: "/fake/bin/nvidia-cuda-mps-control") + pipe_dir = linux_host / "ready-pipe" + log_dir = linux_host / "ready-log" + queries: list[list[str]] = [] + + def run(command: list[str], **kwargs: Any) -> Completed: + if command[0] == "nvidia-smi": + return Completed(f"0, {GPU_A}, Default\n") + if command == list(mps._DEFAULT_CONTROL_COMMAND): + make_control(pipe_dir) + (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") + monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) + return Completed("") + if command == list(mps._SERVER_LIST_QUERY_COMMAND): + queries.append(command) + assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == str(pipe_dir) + assert kwargs["timeout"] > 0.25 + if len(queries) == 1: + return Completed("control is starting", returncode=1) + return Completed("") + raise AssertionError(f"unexpected command {command}") + + record = mps.start_daemon( + GPU_A, + name="ready", + pipe_dir=str(pipe_dir), + log_dir=str(log_dir), + root=linux_host, + run_command=run, + ) + + assert len(queries) == 2 + assert record.pid == 4242 + + +def test_start_daemon_cleans_up_unrecorded_control_after_readiness_failure( + linux_host: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(mps.shutil, "which", lambda _name: "/fake/bin/nvidia-cuda-mps-control") + pipe_dir = linux_host / "failed-pipe" + log_dir = linux_host / "failed-log" + quit_commands = 0 + + def run(command: list[str], **kwargs: Any) -> Completed: + nonlocal quit_commands + if command[0] == "nvidia-smi": + return Completed(f"0, {GPU_A}, Default\n") + if command == list(mps._DEFAULT_CONTROL_COMMAND): + make_control(pipe_dir) + (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") + monkeypatch.setattr(mps, "_proc_stat", lambda _pid: None) + return Completed("") + if command == list(mps._SERVER_LIST_QUERY_COMMAND): + assert kwargs["timeout"] > 0.25 + return Completed("control is not ready", returncode=1) + if command == list(mps._QUIT_COMMAND): + quit_commands += 1 + assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == str(pipe_dir) + return Completed("") + raise AssertionError(f"unexpected command {command}") + + monkeypatch.setattr(mps, "_wait_for_control_query", lambda *_args, **_kwargs: False) + monkeypatch.setattr( + mps, + "_remove_orphan_default_control_sockets", + lambda pipe_directory: pipe_directory.mkdir(parents=True, exist_ok=True), + ) + + with pytest.raises(mps.CudaMpsCliError, match="did not answer a readiness query"): + mps.start_daemon( + GPU_A, + name="failed", + pipe_dir=str(pipe_dir), + log_dir=str(log_dir), + root=linux_host, + run_command=run, + ) + + assert quit_commands == 1 + assert not (linux_host / "failed" / "daemon.json").exists() + + def test_default_daemon_name_uses_complete_uuid() -> None: name = mps._default_name((mps.GpuIdentity(0, GPU_A),)) @@ -460,6 +551,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) return Completed("") + if command == list(mps._SERVER_LIST_QUERY_COMMAND): + return Completed("") raise AssertionError(f"unexpected command {command}") record = mps.start_daemon( @@ -475,6 +568,37 @@ def run(command: list[str], **kwargs: Any) -> Completed: assert list((linux_host / ".stale").iterdir()) +def test_select_environment_for_gpu_requires_sole_live_matching_record( + linux_host: Path, +) -> None: + record = daemon_record(linux_host, name="prod") + write_record(linux_host, record) + + with pytest.raises(mps.CudaMpsCliError, match="No live UniLab-recorded CUDA MPS daemon"): + mps.select_environment_for_gpu(GPU_A, root=linux_host, proc_stat=lambda _pid: None) + + values = mps.select_environment_for_gpu( + GPU_A, root=linux_host, proc_stat=live_proc_stat(record.pid) + ) + assert values["CUDA_MPS_PIPE_DIRECTORY"] == record.pipe_directory + assert values["CUDA_MPS_LOG_DIRECTORY"] == record.log_directory + + other = daemon_record( + linux_host, + name="other", + pid=4243, + process_start_ticks=778, + gpu_uuids=(GPU_A,), + ) + write_record(linux_host, other) + with pytest.raises(mps.CudaMpsCliError, match="Multiple live UniLab-recorded daemons"): + mps.select_environment_for_gpu( + GPU_A, + root=linux_host, + proc_stat=lambda pid: (pid, 777 if pid == record.pid else 778), + ) + + def test_environment_for_daemon_does_not_mutate_process_environment( linux_host: Path, ) -> None: @@ -562,7 +686,7 @@ def run(command: list[str], **kwargs: Any) -> Completed: if "--query-gpu=index,uuid" in command: return Completed(f"0, {GPU_A}\n") return Completed(f"0, {GPU_A}, Default\n") - if command == ["nvidia-cuda-mps-control", "get-server-list"]: + if command == ["nvidia-cuda-mps-control", "get_server_list"]: return Completed("4242\n") raise AssertionError(f"unexpected command {command}") @@ -581,7 +705,7 @@ def run(command: list[str], **kwargs: Any) -> Completed: Path(record.pipe_directory) / "control" ) daemon_query = commands[-1] - assert daemon_query[0] == ["nvidia-cuda-mps-control", "get-server-list"] + assert daemon_query[0] == ["nvidia-cuda-mps-control", "get_server_list"] assert daemon_query[1]["env"]["CUDA_MPS_PIPE_DIRECTORY"] == record.pipe_directory @@ -601,7 +725,7 @@ def run(command: list[str], **kwargs: Any) -> Completed: if "--query-gpu=index,uuid" in command: return Completed(f"0, {GPU_A}\n") return Completed(f"0, {GPU_A}, Default\n") - if command == ["nvidia-cuda-mps-control", "get-server-list"]: + if command == ["nvidia-cuda-mps-control", "get_server_list"]: return Completed("Cannot find MPS control daemon process", returncode=1) raise AssertionError(f"unexpected command {command}") diff --git a/tests/training/test_cuda_process_sharing.py b/tests/training/test_cuda_process_sharing.py index 04664ed27..8d215633a 100644 --- a/tests/training/test_cuda_process_sharing.py +++ b/tests/training/test_cuda_process_sharing.py @@ -37,7 +37,7 @@ def run_command(command: list[str], **kwargs: Any) -> _Completed: del kwargs if command[0] == "nvidia-smi": return _Completed(f"0, {gpu_uuid}\n") - if command == ["nvidia-cuda-mps-control", "get-server-list"]: + if command == ["nvidia-cuda-mps-control", "get_server_list"]: return _Completed(f"{server_pid}\n") raise AssertionError(f"unexpected command: {command}") From b41fdac52d2b50d1c46db3231cbd1fd8785695b2 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 8 Oct 2026 23:53:15 +0800 Subject: [PATCH 2/2] fix(training): send CUDA MPS queries through stdin --- src/unilab/training/cuda_mps_cli.py | 3 ++- src/unilab/training/cuda_process_sharing.py | 3 ++- tests/training/test_cuda_mps_cli.py | 22 ++++++++++++--------- tests/training/test_cuda_process_sharing.py | 4 ++-- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/unilab/training/cuda_mps_cli.py b/src/unilab/training/cuda_mps_cli.py index 9daacc51a..ba3152ab1 100644 --- a/src/unilab/training/cuda_mps_cli.py +++ b/src/unilab/training/cuda_mps_cli.py @@ -42,7 +42,7 @@ _NAME_PATTERN = re.compile(r"^[a-z0-9][a-z0-9-]{0,63}$") _DEFAULT_CONTROL_COMMAND = ("nvidia-cuda-mps-control", "-d") _QUIT_COMMAND = ("nvidia-cuda-mps-control",) -_SERVER_LIST_QUERY_COMMAND = ("nvidia-cuda-mps-control", "get_server_list") +_SERVER_LIST_QUERY_COMMAND = ("nvidia-cuda-mps-control",) _MAX_CONTROL_SOCKET_BYTES = 95 _INERT_CONTROL_ARTIFACTS = ("control_lock", "log") _GPU_QUERY_COMMAND = ( @@ -509,6 +509,7 @@ def _wait_for_control_query( _run_command( run_command, _SERVER_LIST_QUERY_COMMAND, + input="get_server_list\n", timeout=remaining, env=environment, ) diff --git a/src/unilab/training/cuda_process_sharing.py b/src/unilab/training/cuda_process_sharing.py index e77d859be..c48fd7c4d 100644 --- a/src/unilab/training/cuda_process_sharing.py +++ b/src/unilab/training/cuda_process_sharing.py @@ -179,7 +179,8 @@ def _server_evidence( environment["CUDA_MPS_PIPE_DIRECTORY"] = str(control_pipe.parent) try: result = run_command( - ["nvidia-cuda-mps-control", "get_server_list"], + ["nvidia-cuda-mps-control"], + input="get_server_list\n", text=True, capture_output=True, timeout=_MPS_CONTROL_QUERY_TIMEOUT_SEC, diff --git a/tests/training/test_cuda_mps_cli.py b/tests/training/test_cuda_mps_cli.py index d325d75b7..69b3939d1 100644 --- a/tests/training/test_cuda_mps_cli.py +++ b/tests/training/test_cuda_mps_cli.py @@ -248,7 +248,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) return Completed("") - if command == list(mps._SERVER_LIST_QUERY_COMMAND): + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" return Completed("") raise AssertionError(f"unexpected command {command}") @@ -286,7 +287,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) return Completed("") - if command == list(mps._SERVER_LIST_QUERY_COMMAND): + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" queries.append(command) assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == str(pipe_dir) assert kwargs["timeout"] > 0.25 @@ -325,10 +327,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", lambda _pid: None) return Completed("") - if command == list(mps._SERVER_LIST_QUERY_COMMAND): - assert kwargs["timeout"] > 0.25 - return Completed("control is not ready", returncode=1) if command == list(mps._QUIT_COMMAND): + assert kwargs["input"] == "quit\n" quit_commands += 1 assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == str(pipe_dir) return Completed("") @@ -551,7 +551,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: (pipe_dir / "nvidia-cuda-mps-control.pid").write_text("4242\n") monkeypatch.setattr(mps, "_proc_stat", live_proc_stat()) return Completed("") - if command == list(mps._SERVER_LIST_QUERY_COMMAND): + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" return Completed("") raise AssertionError(f"unexpected command {command}") @@ -686,7 +687,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: if "--query-gpu=index,uuid" in command: return Completed(f"0, {GPU_A}\n") return Completed(f"0, {GPU_A}, Default\n") - if command == ["nvidia-cuda-mps-control", "get_server_list"]: + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" return Completed("4242\n") raise AssertionError(f"unexpected command {command}") @@ -705,7 +707,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: Path(record.pipe_directory) / "control" ) daemon_query = commands[-1] - assert daemon_query[0] == ["nvidia-cuda-mps-control", "get_server_list"] + assert daemon_query[0] == ["nvidia-cuda-mps-control"] + assert daemon_query[1]["input"] == "get_server_list\n" assert daemon_query[1]["env"]["CUDA_MPS_PIPE_DIRECTORY"] == record.pipe_directory @@ -725,7 +728,8 @@ def run(command: list[str], **kwargs: Any) -> Completed: if "--query-gpu=index,uuid" in command: return Completed(f"0, {GPU_A}\n") return Completed(f"0, {GPU_A}, Default\n") - if command == ["nvidia-cuda-mps-control", "get_server_list"]: + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" return Completed("Cannot find MPS control daemon process", returncode=1) raise AssertionError(f"unexpected command {command}") diff --git a/tests/training/test_cuda_process_sharing.py b/tests/training/test_cuda_process_sharing.py index 8d215633a..fc39b71c3 100644 --- a/tests/training/test_cuda_process_sharing.py +++ b/tests/training/test_cuda_process_sharing.py @@ -34,10 +34,10 @@ def _torch(uuid: str) -> Any: def _fake_commands(gpu_uuid: str = "GPU-a", *, server_pid: int = 2768293) -> Any: def run_command(command: list[str], **kwargs: Any) -> _Completed: - del kwargs if command[0] == "nvidia-smi": return _Completed(f"0, {gpu_uuid}\n") - if command == ["nvidia-cuda-mps-control", "get_server_list"]: + if command == ["nvidia-cuda-mps-control"]: + assert kwargs["input"] == "get_server_list\n" return _Completed(f"{server_pid}\n") raise AssertionError(f"unexpected command: {command}")