ocp.v1.multihost module#
Multihost functionalities.
- orbax.checkpoint.experimental.v1.multihost.is_pathways_backend()[source][source]#
Checks if a Pathways backend is detected.
- Return type:
bool- Returns:
True if a Pathways (native or proxy) backend is detected, False otherwise.
- async orbax.checkpoint.experimental.v1.multihost.sync_global_processes(key, *, operation_id, timeout=None, processes=None, record_event_name='/jax/checkpoint/sync_global_devices_duration_sec')[source][source]#
Barrier to sync concurrent processes.
NOTE: The barrier name must be unique, i.e. no process should wait on the same barrier name multiple times.
- Parameters:
key (
str) – barrier name. Must be unique.operation_id (
str) – The barrier name will be prefixed with the operation id.timeout (
UnionType[int,None]) – timeout in seconds.processes (
Optional[Collection[int],None]) – If None, expects to wait across all processes and devices. Otherwise, creates a barrier only across devices associated with the given processes.record_event_name (
str) – The name of the event to record the duration of the synchronization.
- Raises:
ValueError – If timeout is less than or equal to 0, or if processes is provided and the current process index (process_index()) is not in processes.
- orbax.checkpoint.experimental.v1.multihost.is_primary_host(primary_host)[source][source]#
Checks whether the current host is the primary host.
- Parameters:
primary_host (
UnionType[int,None]) – The host id of the primary host. If None, all hosts are considered as primary (True is returned).- Return type:
bool- Returns:
True if primary_host is None or equals process_index(), False otherwise.
- orbax.checkpoint.experimental.v1.multihost.process_count()[source][source]#
Returns the number of JAX processes associated with the distributed cluster.
- Return type:
int- Returns:
The total number of JAX processes (jax.process_count()).
- orbax.checkpoint.experimental.v1.multihost.process_index()[source][source]#
Returns the unique index of the current JAX process in the cluster.
When running on distributed backends (non-Pathways), this retrieves the distributed process_id (global_state.process_id) rather than jax.process_index(), which ensures correct barrier coordination across a subset of processes.
- Return type:
int- Returns:
The integer index of the current process.