def run(self, func: Callable, *args, **kwargs):
"""Execute ``func`` with provided arguments in distributed context.
Example
.. code-block:: python
def training(local_rank, config, **kwargs):
# ...
print(idist.get_rank(), ": run with config:", config, "- backend=", idist.backend())
# ...
Args:
func (Callable): function to execute. First argument of the function should be `local_rank` - local process
index.
*args: positional arguments of ``func`` (without `local_rank`).
**kwargs: keyword arguments of ``func``.
"""
if self._spawn_params is not None:
self.logger.info("Spawn function '{}' in {} processes".format(
func, self._spawn_params["nproc_per_node"]))
idist.spawn(self.backend,
func,
args=args,
kwargs_dict=kwargs,
**self._spawn_params)
else:
self.logger.info("- Run '{}' in {} processes".format(
func, idist.get_world_size()))
local_rank = idist.get_local_rank()
func(local_rank, *args, **kwargs)
self.logger.info("End of run")
def run(self, func: Callable, *args: Any, **kwargs: Any) -> None:
"""Execute ``func`` with provided arguments in distributed context.
Example
.. code-block:: python
def training(local_rank, config, **kwargs):
# ...
print(idist.get_rank(), ": run with config:", config, "- backend=", idist.backend())
# ...
with idist.Parallel(backend=backend) as parallel:
parallel.run(training, config, a=1, b=2)
Args:
func: function to execute. First argument of the function should be `local_rank` - local process
index.
args: positional arguments of ``func`` (without `local_rank`).
kwargs: keyword arguments of ``func``.
"""
if self._spawn_params is not None and self.backend is not None:
self._logger.info( # type: ignore[attr-defined]
f"Spawn function '{func}' in {self._spawn_params['nproc_per_node']} processes"
)
idist.spawn(self.backend,
func,
args=args,
kwargs_dict=kwargs,
**self._spawn_params)
else:
self._logger.info(
f"- Run '{func}' in {idist.get_world_size()} processes"
) # type: ignore[attr-defined]
local_rank = idist.get_local_rank()
func(local_rank, *args, **kwargs)
self._logger.info("End of run") # type: ignore[attr-defined]
