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.

bind_backend_process_device(resolved)

Bind a resolved backend process device in the current process.

bind_backend_process_device_for_backend(...)

Bind one backend's process-global accelerator device.

bind_genesis_process_device(resolved)

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

configure_backend_process_device(...)

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 integer simulator device id for a training rank.

resolve_backend_process_device(backend_type, ...)

warn_if_backend_device_collision(...[, source])

Warn when a multi-rank simulator still resolves to device zero.
unilab.base.process_device.apply_backend_env_device_override(env_cfg_override, backend_type, *, devices=None, rank=0, local_rank=None, world_size=1, 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.bind_backend_process_device(resolved)[source]

Bind a resolved backend process device in the current process.

Top-level on purpose: uni_rl’s off-policy collectors receive this as the injected backend_device_binder and pickle it by reference into spawn-based subprocesses. The mjwarp import stays lazy so the binder is importable without the mjwarp extra installed.

Parameters:

resolved (str)

Return type:

str | None

unilab.base.process_device.bind_backend_process_device_for_backend(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 uni_rl’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)[source]
Parameters:
Return type:

str | None

unilab.base.process_device.pin_genesis_device_before_cuda_init(backend_type, *, devices=None, rank=0, local_rank=None, world_size=1, 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. Pure config topology (training.devices / LOCAL_RANK / an explicit 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, *, devices=None, rank=0, local_rank=None, world_size=1, learner_device=None)[source]

Resolve the integer simulator device id for a training rank.

Parameters:

  • backend_type (str) – Selected UniSim backend.

  • devices (Sequence[int] | None) – training.devices in the host-visible namespace. This is used by the off-policy launcher and by single-process PPO.

  • rank (int) – Off-policy data-parallel rank (rank zero by default).

  • local_rank (int | None) – LOCAL_RANK from torchrun. In a distributed PPO worker this is the logical index inside the launcher’s remapped CUDA_VISIBLE_DEVICES list.

  • world_size (int) – Torchrun world size. Values greater than one select the local_rank namespace; values of one select devices[rank].

  • learner_device (str | None) – Explicit learner device fallback when no topology was configured (for example APPO or a single-device play command).

Return type:

int | None

Returns None for backends without an explicit simulator device field. For a torchrun worker the returned id is intentionally local (rather than the host index in devices), because the worker subprocess inherits the remapped CUDA_VISIBLE_DEVICES environment.

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

str | None

unilab.base.process_device.warn_if_backend_device_collision(backend_type, *, devices, rank, device_id, source='environment')[source]

Warn when a multi-rank simulator still resolves to device zero.

This is a transition guard for older adapters/configuration paths. Rank zero legitimately owns device zero; only a non-zero rank resolving to zero is a collision. The warning is intentionally emitted at construction time, never from a hot simulation path.

Parameters:
Return type:

None