artist.util.environment_setup
=============================

.. py:module:: artist.util.environment_setup


Attributes
----------

.. autoapisummary::

   artist.util.environment_setup.log


Functions
---------

.. autoapisummary::

   artist.util.environment_setup.initialize_ddp_environment
   artist.util.environment_setup.create_subgroups_for_nested_ddp
   artist.util.environment_setup.setup_distributed_environment
   artist.util.environment_setup.distribute_groups_among_ranks
   artist.util.environment_setup.get_device


Module Contents
---------------

.. py:data:: log

   A logger for the environment.


.. py:function:: 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
   ----------
   device : torch.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.


.. py:function:: 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
   ----------
   rank : int
       The current process.
   groups_to_ranks_mapping : dict[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.


.. py:function:: 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
   ----------
   device : Optional[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_assignments : Optional[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.


.. py:function:: 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_size : int
       Total number of processes in the global process group.
   number_of_heliostat_groups : int
       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.


.. py:function:: 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
   ----------
   device : torch.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.