itwinai.distributed

class itwinai.distributed.DistributedStrategy[source]

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)[source]

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_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {}

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

model_fields: ClassVar[dict[str, FieldInfo]] = {'global_rank': FieldInfo(annotation=int, required=False, default=0), 'global_world_size': FieldInfo(annotation=int, required=False, default=1), 'local_rank': FieldInfo(annotation=int, required=False, default=0), 'local_world_size': FieldInfo(annotation=int, required=False, default=1)}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

itwinai.distributed.detect_distributed_environment() ClusterEnvironment[source]

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) None[source]

Disable print() when not in main worker.

Parameters:

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