artist.util.environment_setup

Attributes

log

A logger for the environment.

Functions

`initialize_ddp_environment`(→ tuple[torch.device, bool, ...)	Set up the distributed environment.
`create_subgroups_for_nested_ddp`(→ tuple[int, int, ...)	Assign the current process (rank) to a subgroup based on a predefined group assignment map.
`setup_distributed_environment`(→ Generator[dict[str, ...)	Set up the distributed environment.
`distribute_groups_among_ranks`(→ tuple[dict[int, ...)	Distribute groups among ranks in round-robin fashion.
`get_device`(→ torch.device)	Get the correct GPU device type for common operating systems, default to CPU if none is found.

Module Contents

artist.util.environment_setup.log: A logger for the environment.

artist.util.environment_setup.initialize_ddp_environment(device: torch.device | None = None) → tuple[torch.device, bool, int, int]

Set up the distributed environment.

Based on the available devices, the outer process group is initialized with the appropriate backend. For computation on GPUs the nccl backend optimized for NVIDIA GPUs is chosen. For computation on CPUs gloo is used as backend. If the program is run without the intention of being distributed, the world_size will be set to 1, accordingly the only rank is 0.

Parameters

devicetorch.device | None: The device on which to perform computations or load tensors and models (default is None). If None, ARTIST will automatically select the most appropriate device (CUDA or CPU) based on availability and OS.

Yields

torch.device: The device for each rank.
bool: Distributed mode enabled or disabled.
int: The rank of the current process.
int: The world size or total number of processes.

artist.util.environment_setup.create_subgroups_for_nested_ddp(rank: int, groups_to_ranks_mapping: dict[int, list[int]]) → tuple[int, int, torch.distributed.ProcessGroup | None, dict[int, list[int]]]

Assign the current process (rank) to a subgroup based on a predefined group assignment map.

Parameters

rankint: The current process.
groups_to_ranks_mappingdict[int, list[int]]: The mapping from heliostat group to rank.

Returns

int: The rank within the heliostat group.
int: The world size of the heliostat group.
torch.distributed.ProcessGroup | None: The distributed process group.
dict[int, list[int]]: The mapping from ranks to heliostat groups.

artist.util.environment_setup.setup_distributed_environment(number_of_heliostat_groups: int, device: torch.device | None = None) → Generator[dict[str, Any], None, None]

Set up the distributed environment.

Parameters

deviceOptional[torch.device]: The device on which to perform computations or load tensors and models (default is None). If None, ARTIST will automatically select the most appropriate device (CUDA or CPU) based on availability and OS.
heliostat_group_assignmentsOptional[dict[str, list[int]]]: The mapping from rank to heliostat group.

Yields

dict[str, torch.device | bool | int | torch.distributed.ProcessGroup | dict[int, list[int]] | None]: The distributed setup including the torch device assigned to this rank, whether the environment is running in distributed and or nested mode, the global rank of the current process, the total number of processes in the global process group, the ProcessGroup object representing the subgroup, the mapping from rank to heliostat group, the rank of the current process within its assigned subgroup and the number of processes in the current process subgroup.

artist.util.environment_setup.distribute_groups_among_ranks(world_size: int, number_of_heliostat_groups: int) → tuple[dict[int, list[int]], bool]

Distribute groups among ranks in round-robin fashion.

If there are fewer ranks than groups, some ranks receive multiple groups. If there are more ranks than groups, some groups are handled by multiple ranks, enabling nested distribution.

Parameters

world_sizeint: Total number of processes in the global process group.
number_of_heliostat_groupsint: The number of heliostat groups.

Returns

dict[int, list[int]]: The dictionary mapping heliostat groups to ranks.
bool: Indicates whether the distributed setup is nested or not.

artist.util.environment_setup.get_device(device: torch.device | None = None) → torch.device

Get the correct GPU device type for common operating systems, default to CPU if none is found.

Parameters

devicetorch.device | None: The device on which to perform computations or load tensors and models (default is None). If None, ARTIST will automatically select the most appropriate device (CUDA or CPU) based on availability and OS. MPS (for Mac) is not supported due to limitations in torch.

Returns

torch.device: The device.