itwinai.distributed

class itwinai.distributed.DistributedStrategy

Bases: ABC

Abstract class to define the distributed backend methods.

class itwinai.distributed.ClusterEnvironment(*, global_rank: int = 0, local_rank: int = 0, global_world_size: int = 1, local_world_size: int = 1)

Bases: BaseModel

Stores information about distributed environment.

global_rank: int

Global rank of current worker, in a distributed environment. global_rank==0 identifies the main worker. Defaults to 0.

local_rank: int

Local rank of current worker, in a distributed environment. Defaults to 0.

global_world_size: int

Total number of workers in a distributed environment. Defaults to 1.

local_world_size: int

Number of workers on the same node in a distributed environment. Defaults to 1.

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

itwinai.distributed.detect_distributed_environment() → ClusterEnvironment

Detects distributed environment, extracting information like global ans local ranks, and world size.

itwinai.distributed.builtin_print()

Save original builtin print before patching it in distributed environments

itwinai.distributed.distributed_patch_print(is_main: bool) → Callable

Disable print() when not in main worker.

Parameters:

is_main (bool) – whether it is called from main worker.

Returns:

patched print().

Return type:

Callable

itwinai.distributed.suppress_workers_print(func: Callable) → Callable

Decorator to suppress print() calls in workers having global rank different from 0. To force printing on all workers you need to use print(..., force=True).

itwinai.distributed.suppress_workers_output(func)

Decorator to suppress stadout and stderr in workers having global rank different from 0.