class_docstr: Coordinator of n (possibly remote) Learner workers.
  
  Each Learner worker has a copy of the RLModule, the loss function(s), and
  one or more optimizers.
docstr: Load the checkpoints of the modules being trained by this LearnerGroup.
  
     This method only needs to be called if the LearnerGroup is training
     distributed learners (e.g num_learner_workers > 0).
  
  Args:
      marl_module_ckpt_dir: The path to the checkpoint for the
          MultiAgentRLModule.
      modules_to_load: A set of module ids to load from the checkpoint.
      rl_module_ckpt_dirs: A mapping from module ids to the path to a
          checkpoint for a single agent RLModule.
docstr_tok: ["Load","the","checkpoints","of","the","modules","being","trained","by","this","LearnerGroup",".","This","method","only","needs","to","be","called","if","the","LearnerGroup","is","training","distributed","learners","(","e.g","num_learner_workers",">","0",")",".","Args",":","marl_module_ckpt_dir",":","The","path","to","the","checkpoint","for","the","MultiAgentRLModule",".","modules_to_load",":","A","set","of","module","ids","to","load","from","the","checkpoint",".","rl_module_ckpt_dirs",":","A","mapping","from","module","ids","to","the","path","to","a","checkpoint","for","a","single","agent","RLModule","."]
code: def _distributed_load_module_state(
      self,
      *,
      marl_module_ckpt_dir: Optional[str] = None,
      modules_to_load: Optional[Set[str]] = None,
      rl_module_ckpt_dirs: Optional[Dict[ModuleID, str]] = None,
  ):
      """Load the checkpoints of the modules being trained by this LearnerGroup.
  
         This method only needs to be called if the LearnerGroup is training
         distributed learners (e.g num_learner_workers > 0).
  
      Args:
          marl_module_ckpt_dir: The path to the checkpoint for the
              MultiAgentRLModule.
          modules_to_load: A set of module ids to load from the checkpoint.
          rl_module_ckpt_dirs: A mapping from module ids to the path to a
              checkpoint for a single agent RLModule.
  
      """
  
      assert (
          len(self._workers)
          == self._worker_manager.num_healthy_actors()
      )
      workers = self._worker_manager.healthy_actor_ids()
      head_node_ip = ray.util.get_node_ip_address()
  
      def _load_module_state(w):
          # doing imports here since they might not be imported on the worker
          import ray
          import tempfile
          import shutil
  
          worker_node_ip = ray.util.get_node_ip_address()
          # sync the checkpoints from the head to the worker if the worker is not
          # on the same node as the head
          tmp_marl_module_ckpt_dir = marl_module_ckpt_dir
          tmp_rl_module_ckpt_dirs = rl_module_ckpt_dirs
          if worker_node_ip != head_node_ip:
              if marl_module_ckpt_dir:
                  tmp_marl_module_ckpt_dir = tempfile.mkdtemp()
                  sync_dir_between_nodes(
                      source_ip=head_node_ip,
                      source_path=marl_module_ckpt_dir,
                      target_ip=worker_node_ip,
                      target_path=tmp_marl_module_ckpt_dir,
                  )
              if rl_module_ckpt_dirs:
                  tmp_rl_module_ckpt_dirs = {}
                  for module_id, path in rl_module_ckpt_dirs.items():
                      tmp_rl_module_ckpt_dirs[
                          module_id
                      ] = tempfile.mkdtemp()
                      sync_dir_between_nodes(
                          source_ip=head_node_ip,
                          source_path=path,
                          target_ip=worker_node_ip,
                          target_path=tmp_rl_module_ckpt_dirs[
                              module_id
                          ],
                      )
                      tmp_rl_module_ckpt_dirs[module_id] = pathlib.Path(
                          tmp_rl_module_ckpt_dirs[module_id]
                      )
          if marl_module_ckpt_dir:
              # load the MARLModule checkpoint if they were specified
              w.module.load_state(
                  tmp_marl_module_ckpt_dir,
                  modules_to_load=modules_to_load,
              )
          if rl_module_ckpt_dirs:
              # load the RLModule if they were specified
              for module_id, path in tmp_rl_module_ckpt_dirs.items():
                  w.module[module_id].load_state(
                      path / RLMODULE_STATE_DIR_NAME
                  )
  
          # remove the temporary directories on the worker if any were created
          if worker_node_ip != head_node_ip:
              if marl_module_ckpt_dir:
                  shutil.rmtree(tmp_marl_module_ckpt_dir)
              if rl_module_ckpt_dirs:
                  for (
                      module_id,
                      path,
                  ) in tmp_rl_module_ckpt_dirs.items():
                      shutil.rmtree(path)
  
      self._worker_manager.foreach_actor(
          _load_module_state, remote_actor_ids=workers
      )
  
imports: from collections import defaultdict, Counter
  from functools import partial
  import pathlib
  from typing import Any, Callable, Dict, List, Optional, Set, Type, Union
  import uuid
  import ray
  from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
  from ray.rllib.core.learner.learner import Learner
  from ray.rllib.core.learner.reduce_result_dict_fn import _reduce_mean_results
  from ray.rllib.core.rl_module.rl_module import SingleAgentRLModuleSpec, RLMODULE_STATE_DIR_NAME
  from ray.rllib.policy.sample_batch import MultiAgentBatch
  from ray.rllib.utils.actor_manager import FaultTolerantActorManager
  from ray.rllib.utils.deprecation import Deprecated, deprecation_warning
  from ray.rllib.utils.minibatch_utils import ShardBatchIterator, ShardEpisodesIterator
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.typing import EpisodeType, ModuleID, ResultDict, RLModuleSpec, ShouldModuleBeUpdatedFn
  from ray.train._internal.backend_executor import BackendExecutor
  from ray.tune.utils.file_transfer import sync_dir_between_nodes
  from ray.util.annotations import PublicAPI
label_desc: Use image node_id 22 for calling the LearnerGroup obj's underlying member method code with example usage: obj._distributed_load_module_state() without return types