feat(training): select recorded CUDA MPS daemon

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 |

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`。
+
 默认参数：
 
 | 设置 | 默认值 |

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,

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",)
 _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-<id> UUID.")
     raw_uuid = uuid[4:]
@@ -492,6 +496,41 @@ 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,
+                input="get_server_list\n",
+                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 +876,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 +931,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,
     *,

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,

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