unilab.base.process_device

Training-worker device routing for UniSim backends.

The learner device and the device consumed by a simulator are related, but they do not always use the same index namespace. In particular, the off-policy launcher keeps the host-visible CUDA namespace while the PPO torchrun launcher remaps CUDA_VISIBLE_DEVICES and therefore exposes a rank-local index to child processes. The helpers in this module keep that translation on the cold path, next to the backend process binding contract.

Functions

apply_backend_env_device_override(...[, ...])

Return an env override carrying the rank-selected simulator device.

apply_collector_tensor_device_override(...)

Return an env override selecting optional HOST_BRIDGE Torch carriers.

bind_backend_process_device(backend_type, ...)

Bind one backend's process-global accelerator device.

bind_genesis_process_device(resolved)

Select the CUDA device used by an in-process Genesis worker.

canonicalize_collector_tensor_device(...)

Return the concrete public-manager collector tensor device for a request.

configure_backend_process_device(...[, ...])

is_host_bridge_backend(backend_type)

Return whether the backend executes through the packed HOST_BRIDGE plane.

pin_genesis_device_before_cuda_init(...[, ...])

Pin Genesis to its rank device before the first torch CUDA call.

resolve_backend_env_device_id(backend_type, *)

Resolve the simulator payload id from the rank-local CUDA namespace.

resolve_backend_process_device(backend_type, ...)

unilab.base.process_device.canonicalize_collector_tensor_device(requested_device)[source]

Return the concrete public-manager collector tensor device for a request.

Parameters:

requested_device (str | None)

Return type:

str

unilab.base.process_device.apply_backend_env_device_override(env_cfg_override, backend_type, *, learner_device=None)[source]

Return an env override carrying the rank-selected simulator device.

The input mapping is never mutated. If no topology/device can be resolved, the owner-configured value is preserved. This lets one helper serve training, playback, and custom entrypoints while retaining the historical default (device zero) for single-process calls.

Parameters:
Return type:

dict[str, Any]

unilab.base.process_device.apply_collector_tensor_device_override(env_cfg_override, backend_type, *, collector_tensor_device=None, learner_device=None)[source]

Return an env override selecting optional HOST_BRIDGE Torch carriers.

The input mapping is never mutated. training.collector_tensor_device is the explicit training contract: cpu remains on CPU and unindexed cuda is resolved to the current device by the environment. auto follows the learner device: a CUDA learner routes the rank-local (indexed) learner device string, while a CPU learner routes CPU and a torch.mps learner retains the owner/default CPU placement. Eligibility comes from the public SDK-free backend capability inventory: only a HOST_BRIDGE backend may route this key, and the requested device family/index must be in its declared torch_devices. auto is a no-op for backends outside that scope and when no learner device is known. The learner_device argument is the fallback for callers that do not compose a training configuration.

An explicit cpu request with a known CUDA learner fails closed (issue #2483) unless the build is ROCm/HIP or the run carries the UNILAB_ALLOW_CUDA_LEARNER_CPU_CARRIER attribution-baseline/debug opt-in.

Parameters:
Return type:

dict[str, Any]

unilab.base.process_device.bind_backend_process_device(backend_type, resolved)[source]

Bind one backend’s process-global accelerator device.

This top-level callable is intentionally backend-aware and lazy. It can be wrapped with functools.partial() and injected into the RL runtime’s spawn-based collectors while remaining pickleable by module reference.

Parameters:

  • backend_type (str)

  • resolved (str)

Return type:

str | None

unilab.base.process_device.bind_genesis_process_device(resolved)[source]

Select the CUDA device used by an in-process Genesis worker.

Genesis initializes a process-wide session whose Quadrants CUDA runtime always binds the first entry of CUDA_VISIBLE_DEVICES; torch’s current device alone is not honored (issue #1508). A non-zero request is therefore honored by pinning visibility to the target GPU before any CUDA context exists, after which the in-process device is cuda:0. The returned string is the device the rest of this process must actually use — callers that computed a pre-pin device (learner, probes, collectors) have to adopt the returned value. Binding must happen before constructing the backend (and before gs.init), including in spawn-based collector processes, and remains in effect for the lifetime of the process.

Parameters:

resolved (str)

Return type:

str

unilab.base.process_device.configure_backend_process_device(backend_type, learner_device, *, backend_device_id=None)[source]
Parameters:
Return type:

str | None

unilab.base.process_device.is_host_bridge_backend(backend_type)[source]

Return whether the backend executes through the packed HOST_BRIDGE plane.

Only HOST_BRIDGE backends route training.collector_tensor_device; callers that record or display the resolved device use this predicate instead of naming individual backends.

Parameters:

backend_type (str)

Return type:

bool

unilab.base.process_device.pin_genesis_device_before_cuda_init(backend_type, *, learner_device=None)[source]

Pin Genesis to its rank device before the first torch CUDA call.

torch.cuda.is_available() already latches CUDA_VISIBLE_DEVICES in the CUDA runtime, so the pin must run ahead of any torch CUDA query — entrypoints should call this before registry/bootstrap/device detection. An explicit non-Genesis cuda:N learner device resolves without touching torch. Returns the in-process device the caller must use when a pin happened, else None.

Parameters:
Return type:

str | None

unilab.base.process_device.resolve_backend_env_device_id(backend_type, *, learner_device=None)[source]

Resolve the simulator payload id from the rank-local CUDA namespace.

None means the backend has no payload device field. With CUDA visible, exactly one rank-local entry is required and the payload id is always 0; otherwise the explicit non-Genesis learner device fallback remains valid.

Parameters:
Return type:

int | None

unilab.base.process_device.resolve_backend_process_device(backend_type, learner_device)[source]
Parameters:
Return type:

str | None