shrinath-suresh/blogs-docs-splitted huggingface dataset:
the metrics pushed to a metric handling service in your
infrastructure, implement a
torch.distributed.elastic.metrics.MetricHandler and configure it
in your custom launcher.
# my_launcher.py
import torch.distributed.elastic.metrics as metrics
class MyMetricHandler(metrics.MetricHandler):
       def emit(self, metric_data: metrics.MetricData):
           # push metric_data to your metric sink
def main():
     metrics.configure(MyMetricHandler())
 spec = WorkerSpec(...)
 agent = LocalElasticAgent(spec)
 agent.run()

Events Handler
TorchElastic supports events recording (see Events). The events module
defines API that allows you to record events and implement custom
EventHandler. EventHandler is used for publishing events produced
during torchelastic execution to different sources, e.g.  AWS
CloudWatch. By default it uses
torch.distributed.elastic.events.NullEventHandler that ignores
events. To configure custom events handler you need to implement
torch.distributed.elastic.events.EventHandler interface and
configure it in your custom launcher.
# my_launcher.py
import torch.distributed.elastic.events as events
class MyEventHandler(events.EventHandler):
       def record(self, event: events.Event):
           # process event
def main():
     events.configure(MyEventHandler())
 spec = WorkerSpec(...)
 agent = LocalElasticAgent(spec)
 agent.run()
Multiprocessing
Library that launches and manages "n" copies of worker subprocesses
either specified by a function or a binary.
For functions, it uses "torch.multiprocessing" (and therefore python
"multiprocessing") to spawn/fork worker processes. For binaries it
uses python "subprocessing.Popen" to create worker processes.
Usage 1: Launching two trainers as a function
from torch.distributed.elastic.multiprocessing import Std, start_processes
def trainer(a, b, c):
       pass # train
# runs two trainers
   # LOCAL_RANK=0 trainer(1,2,3)
   # LOCAL_RANK=1 trainer(4,5,6)
   ctx = start_processes(
           name="trainer",
           entrypoint=trainer,
           args={0: (1,2,3), 1: (4,5,6)},
           envs={0: {"LOCAL_RANK": 0}, 1: {"LOCAL_RANK": 1}},
           log_dir="/tmp/foobar",
           redirects=Std.ALL, # write all worker stdout/stderr to a log file
           tee={0: Std.ERR}, # tee only local rank 0's stderr to console
         )
)
# waits for all copies of trainer to finish
   ctx.wait()
Usage 2: Launching 2 echo workers as a binary
# same as invoking
   # echo hello
   # echo world > stdout.log
   ctx = start_processes(
           name="echo"
           entrypoint="echo",
           log_dir="/tmp/foobar",
           args={0: "hello", 1: "world"},
           redirects={1: Std.OUT},
          )
Just like "torch.multiprocessing", the return value of the function
"start_processes()" is a process context ("api.PContext"). If a
function was launched, a "api.MultiprocessContext" is returned and if
a binary was launched a "api.SubprocessContext" is returned. Both are
specific implementations of the parent "api.PContext" class.
Starting Multiple Workers
torch.distributed.elastic.multiprocessing.start_processes(name, entrypoint, args, envs, log_dir, start_method='spawn', redirects=Std.NONE, tee=Std.NONE)
Starts "n" copies of "entrypoint" processes with the provided
options. "entrypoint" is either a "Callable" (function) or a "str"
   (binary). The number of copies is determined by the number of
   entries for "args" and "envs" arguments, which need to have the
   same key set.
"args" and "env" parameters are the arguments and environment
   variables to pass down to the entrypoint mapped by the replica
   index (local rank). All local ranks must be accounted for. That is,
   the keyset should be "{0,1,...,(nprocs-1)}".
Note:
 When the "entrypoint" is a binary ("str"), "args" can only be
 strings. If any other type is given, then it is casted to a
 string representation (e.g. "str(arg1)"). Furthermore, a binary
 failure will only write an "error.json" error file if the main
 function is annotated with
 "torch.distributed.elastic.multiprocessing.errors.record". For
 function launches, this is done by default and there is no need
 to manually annotate with the "@record" annotation.
"redirects" and "tee" are bitmasks specifying which std stream(s)
   to redirect to a log file in the "log_dir". Valid mask values are
   defined in "Std". To redirect/tee only certain local ranks, pass
   "redirects" as a map with the key as the local rank to specify the
   redirect behavior for. Any missing local ranks will default to
   "Std.NONE".
"tee" acts like the unix "tee" command in that it redirects +
   prints to console. To avoid worker stdout/stderr from printing to
   console, use the "redirects" parameter.
For each process, the "log_dir" will contain:


"{local_rank}/error.json": if the process failed, a file with
      the error info


"{local_rank}/stdout.json": if "redirect & STDOUT == STDOUT"


"{local_rank}/stderr.json": if "redirect & STDERR == STDERR"


Note:
 It is expected that the "log_dir" exists, is empty, and is a
 directory.

Example:
  log_dir = "/tmp/test"

  # ok; two copies of foo: foo("bar0"), foo("bar1")
start_processes(
         name="trainer",
         entrypoint=foo,
         args:{0:("bar0",), 1:("bar1",),
         envs:{0:{}, 1:{}},
         log_dir=log_dir
      )
  # invalid; envs missing for local rank 1
  start_processes(
     name="trainer",
     entrypoint=foo,
     args:{0:("bar0",), 1:("bar1",),
     envs:{0:{}},
     log_dir=log_dir
  )

  # ok; two copies of /usr/bin/touch: touch file1, touch file2
  start_processes(
     name="trainer",
     entrypoint="/usr/bin/touch",
     args:{0:("file1",), 1:("file2",),
     envs:{0:{}, 1:{}},
     log_dir=log_dir
   )

  # caution; arguments casted to string, runs:
  # echo "1" "2" "3" and echo "[1, 2, 3]"
  start_processes(
     name="trainer",
     entrypoint="/usr/bin/echo",
     args:{0:(1,2,3), 1:([1,2,3],),
     envs:{0:{}, 1:{}},
     log_dir=log_dir
   )

Parameters:
      * name (str) -- a human readable short name that describes
what the processes are (used as header when tee'ing
        stdout/stderr outputs)
  * **entrypoint** (*Union**[**Callable**, **str**]*) -- either a
    "Callable" (function) or "cmd" (binary)

  * **args** (*Dict**[**int**, **Tuple**]*) -- arguments to each
    replica

  * **envs** (*Dict**[**int**, **Dict**[**str**, **str**]**]*) --
    env vars to each replica

  * **log_dir** (*str*) -- directory used to write log files

  * **start_method** (*str*) -- multiprocessing start method
    (spawn, fork, forkserver) ignored for binaries

  * **redirects** (*Union**[**Std**, **Dict**[**int**,
    **Std**]**]*) -- which std streams to redirect to a log file

  * **tee** (*Union**[**Std**, **Dict**[**int**, **Std**]**]*) --
    which std streams to redirect + print to console

Return type:
      PContext
Process Context
Process Context
class torch.distributed.elastic.multiprocessing.api.PContext(name, entrypoint, args, envs, stdouts, stderrs, tee_stdouts, tee_stderrs, error_files)
The base class that standardizes operations over a set of processes
   that are launched via different mechanisms. The name "PContext" is
   intentional to disambiguate with
   "torch.multiprocessing.ProcessContext".
Warning:
 stdouts and stderrs should ALWAYS be a superset of tee_stdouts
 and tee_stderrs (respectively) this is b/c tee is implemented as
 a redirect + tail -f <stdout/stderr.log>

class torch.distributed.elastic.multiprocessing.api.MultiprocessContext(name, entrypoint, args, envs, stdouts, stderrs, tee_stdouts, tee_stderrs, error_files, start_method)
"PContext" holding worker processes invoked as a function.
class torch.distributed.elastic.multiprocessing.api.SubprocessContext(name, entrypoint, args, envs, stdouts, stderrs, tee_stdouts, tee_stderrs, error_files)
"PContext" holding worker processes invoked as a binary.
class torch.distributed.elastic.multiprocessing.api.RunProcsResult(return_values=, failures=, stdouts=, stderrs=)
Results of a completed run of processes started with
   "start_processes()". Returned by "PContext".
Note the following:


All fields are mapped by local rank


"return_values" - only populated for functions (not the
      binaries).


"stdouts" - path to stdout.log (empty string if no redirect)


"stderrs" - path to stderr.log (empty string if no redirect)
Elastic Agent
Server
The elastic agent is the control plane of torchelastic. It is a
process that launches and manages underlying worker processes. The
agent is responsible for:


Working with distributed torch: the workers are started with all
   the necessary information to successfully and trivially call
   "torch.distributed.init_process_group()".


Fault tolerance: monitors workers and upon detecting worker
   failures or unhealthiness, tears down all workers and restarts
   everyone.


Elasticity: Reacts to membership changes and restarts workers with
   the new members.


The simplest agents are deployed per node and works with local
processes. A more advanced agent can launch and manage workers
remotely. Agents can be completely decentralized, making decisions
based on the workers it manages. Or can be coordinated, communicating
to other agents (that manage workers in the same job) to make a
collective decision.
collective decision.
Below is a diagram of an agent that manages a local group of workers.
[image]
Concepts
This section describes the high-level classes and concepts that are
relevant to understanding the role of the "agent" in torchelastic.
class torch.distributed.elastic.agent.server.ElasticAgent
Agent process responsible for managing one or more worker
   processes. The worker processes are assumed to be regular
   distributed PyTorch scripts. When the worker process is created by
   the agent, the agent provides the necessary information for the
   worker processes to properly initialize a torch process group.
The exact deployment topology and ratio of agent-to-worker is
   dependent on the specific implementation of the agent and the
   user's job placement preferences. For instance, to run a
   distributed training job on GPU with 8 trainers (one per GPU) one
   can:

Use 8 x single GPU instances, place an agent per instance,
      managing 1 worker per agent.
managing 1 worker per agent.


Use 4 x double GPU instances, place an agent per instance,
      managing 2 workers per agent.


Use 2 x quad GPU instances, place an agent per instance,
      managing 4 workers per agent.


Use 1 x 8 GPU instance, place an agent per instance, managing 8
      workers per agent.


Usage
  group_result = agent.run()
   if group_result.is_failed():
     # workers failed
     failure = group_result.failures[0]
     log.exception(f"worker 0 failed with exit code : {failure.exit_code}")
   else:
     return group_result.return_values[0] # return rank 0's results

abstract get_worker_group(role='default')
  Returns:
     The "WorkerGroup" for the given "role". Note that the worker
     group is a mutable object and hence in a multi-
     threaded/process environment it may change state.
     Implementors are encouraged (but not required) to return a
     defensive read-only copy.

  Return type:
Return type:
         WorkerGroup
abstract run(role='default')
  Runs the agent, retrying the worker group on failures up to
  "max_restarts".

  Returns:
     The result of the execution, containing the return values or
     failure details for each worker mapped by the worker's global
     rank.

  Raises:
     **Exception - any other failures NOT related to worker
     process** --

  Return type:
     *RunResult*

class torch.distributed.elastic.agent.server.WorkerSpec(role, local_world_size, rdzv_handler, fn=None, entrypoint=None, args=(), max_restarts=3, monitor_interval=30.0, master_port=None, master_addr=None, local_addr=None, redirects=Std.NONE, tee=Std.NONE)
Contains blueprint information about a particular type of worker.
   For a given role, there must only exist a single worker spec.
   Worker spec is expected to be homogenous across all nodes
   (machine), that is each node runs the same number of workers for a
   particular spec.
particular spec.
Parameters:
      * role (str) -- user-defined role for the workers with
        this spec
  * **local_world_size** (*int*) -- number local workers to run

  * **fn** (*Optional**[**Callable**]*) -- (deprecated use
    entrypoint instead)

  * **entrypoint** (*Optional**[**Union**[**Callable**,
    **str**]**]*) -- worker function or command

  * **args** (*Tuple*) -- arguments to pass to "entrypoint"

  * **rdzv_handler** (*RendezvousHandler*) -- handles rdzv for
    this set of workers

  * **max_restarts** (*int*) -- number of max retries for the
    workers

  * **monitor_interval** (*float*) -- monitor status of workers
    every "n" seconds

  * **master_port** (*Optional**[**int**]*) -- fixed port to run
    the c10d store on rank 0 if not specified then will chose a
    random free port

  * **master_addr** (*Optional**[**str**]*) -- fixed master_addr
to run the c10d store on rank 0 if not specified then will
        chose hostname on agent rank 0
  * **redirects** (*Union**[**Std**, **Dict**[**int**,
    **Std**]**]*) -- redirect std streams to a file, selectively
    redirect for a particular local rank by passing a map

  * **tee** (*Union**[**Std**, **Dict**[**int**, **Std**]**]*) --
    tees the specified std stream(s) to console + file,
    selectively tee for a particular local rank by passing a map,
    takes precedence over "redirects" settings.

get_entrypoint_name()
  If the entrypoint is a function (e.g. "Callable") returns its
  "__qualname__", else if the entrypoint is a binary (e.g. "str"),
  returns the binary name.

class torch.distributed.elastic.agent.server.WorkerState(value)
State of the "WorkerGroup". Workers in a worker group change state
   as a unit. If a single worker in a worker group fails the entire
   set is considered failed:
set is considered failed:
  UNKNOWN - agent lost track of worker group state, unrecoverable
  INIT - worker group object created not yet started
  HEALTHY - workers running and healthy
  UNHEALTHY - workers running and unhealthy
  STOPPED - workers stopped (interrupted) by the agent
  SUCCEEDED - workers finished running (exit 0)
  FAILED - workers failed to successfully finish (exit !0)

A worker group starts from an initial "INIT" state, then progresses
   to "HEALTHY" or "UNHEALTHY" states, and finally reaches a terminal
   "SUCCEEDED" or "FAILED" state.
Worker groups can be interrupted and temporarily put into "STOPPED"
   state by the agent. Workers in "STOPPED" state are scheduled to be
   restarted in the near future by the agent. Some examples of workers
   being put into "STOPPED" state are:


Worker group failure|unhealthy observed


Membership change detected


When actions (start, stop, rdzv, retry, etc) on worker group fails
and results in the action being partially applied to the worker
   group the state will be "UNKNOWN". Typically this happens on
   uncaught/unhandled exceptions during state change events on the
   agent. The agent is not expected to recover worker groups in
   "UNKNOWN" state and is better off self terminating and allowing the
   job manager to retry the node.
static is_running(state)
  Returns:
     True if the worker state represents workers still running
     (e.g. that the process exists but not necessarily healthy).

  Return type:
     bool

class torch.distributed.elastic.agent.server.Worker(local_rank, global_rank=- 1, role_rank=- 1, world_size=- 1, role_world_size=- 1)
Represents a worker instance. Contrast this with "WorkerSpec" that
   represents the specifications of a worker. A "Worker" is created
   from a "WorkerSpec". A "Worker" is to a "WorkerSpec" as an object
   is to a class.
The "id" of the worker is interpreted by the specific
implementation of "ElasticAgent". For a local agent, it could be
   the "pid (int)" of the worker, for a remote agent it could be
   encoded as "host:port (string)".
Parameters:
      * id (Any) -- uniquely identifies a worker (interpreted by
        the agent)
  * **local_rank** (*int*) -- local rank of the worker

  * **global_rank** (*int*) -- global rank of the worker

  * **role_rank** (*int*) -- rank of the worker across all workers
    that have the same role

  * **world_size** (*int*) -- number of workers (globally)

  * **role_world_size** (*int*) -- number of workers that have the
    same role

class torch.distributed.elastic.agent.server.WorkerGroup(spec)
Represents the set of "Worker" instances for the given "WorkerSpec"
   managed by "ElasticAgent". Whether the worker group contains cross
   instance workers or not depends on the implementation of the agent.
Implementations
Implementations
Below are the agent implementations provided by torchelastic.
class torch.distributed.elastic.agent.server.local_elastic_agent.LocalElasticAgent(spec, start_method='spawn', exit_barrier_timeout=300, log_dir=None)
An implementation of "torchelastic.agent.server.ElasticAgent" that
   handles host-local workers. This agent is deployed per host and is
   configured to spawn "n" workers. When using GPUs, "n" maps to the
   number of GPUs available on the host.
The local agent does not communicate to other local agents deployed
   on other hosts, even if the workers may communicate inter-host. The
   worker id is interpreted to be a local process. The agent starts
   and stops all worker processes as a single unit.
The worker function and argument passed to the worker function must
   be python multiprocessing compatible. To pass multiprocessing data
   structures to the workers you may create the data structure in the
same multiprocessing context as the specified "start_method" and
   pass it as a function argument.
The "exit_barrier_timeout" specifies the amount of time (in
   seconds) to wait for other agents to finish. This acts as a safety
   net to handle cases where workers finish at different times, to
   prevent agents from viewing workers that finished early as a scale-
   down event. It is strongly advised that the user code deal with
   ensuring that workers are terminated in a synchronous manner rather
   than relying on the exit_barrier_timeout.
A named pipe based watchdog can be enabled in "LocalElasticAgent"
   if an environment variable "TORCHELASTIC_ENABLE_FILE_TIMER" with
   value 1 has been defined in the "LocalElasticAgent" process.
   Optionally, another environment variable
   "TORCHELASTIC_TIMER_FILE" can be set with a unique file name for
   the named pipe. If the environment variable
   "TORCHELASTIC_TIMER_FILE" is not set, "LocalElasticAgent" will
internally create a unique file name and set it to the environment
   variable "TORCHELASTIC_TIMER_FILE", and this environment variable
   will be propagated to the worker processes to allow them to connect
   to the same named pipe that "LocalElasticAgent" uses.
Example launching function
  def trainer(args) -> str:
      return "do train"

  def main():
      start_method="spawn"
      shared_queue= multiprocessing.get_context(start_method).Queue()
      spec = WorkerSpec(
                  role="trainer",
                  local_world_size=nproc_per_process,
                  entrypoint=trainer,
                  args=("foobar",),
                  ...<OTHER_PARAMS...>)
      agent = LocalElasticAgent(spec, start_method)
      results = agent.run()

      if results.is_failed():
          print("trainer failed")
      else:
          print(f"rank 0 return value: {results.return_values[0]}")
prints -> rank 0 return value: do train
Example launching binary
  def main():
      spec = WorkerSpec(
                  role="trainer",
                  local_world_size=nproc_per_process,
                  entrypoint="/usr/local/bin/trainer",
                  args=("--trainer_args", "foobar"),
                  ...<OTHER_PARAMS...>)
      agent = LocalElasticAgent(spec)
      results = agent.run()

      if not results.is_failed():
          print("binary launches do not have return values")

Extending the Agent
To extend the agent you can implement "`ElasticAgent" directly,
however we recommend you extend "SimpleElasticAgent" instead, which
provides most of the scaffolding and leaves you with a few specific
abstract methods to implement.
class torch.distributed.elastic.agent.server.SimpleElasticAgent(spec, exit_barrier_timeout=300)
An "ElasticAgent" that manages workers ("WorkerGroup") for a single
"WorkerSpec" (e.g. one particular type of worker role).
_assign_worker_ranks(store, group_rank, group_world_size, spec)
  Determines proper ranks for worker processes. The rank
  assignment is done according to the following algorithm:

  1. Each agent writes its configuration(group_rank,
     group_world_size , num_workers) to the common store.

  2. Each agent retrieves configuration for all agents and
     performs two level sort using role and rank.

  3. Determine the global rank: the global rank of the workers for
     the current agent is the offset of the infos array up to
     group_rank of the agent. The offset is computed as a sum of
     local_world_size of all agents that have rank less than the
     group_rank. The workers would have the ranks: [offset,
     offset+local_world_size)

  4. Determine the role rank: The role rank is determined using
     the algorithms in the point 3 with the exception that the
offset is done from the first agent that has the same role as
         current one and has the minimum group rank.
  Return type:
     *List*[*Worker*]

_exit_barrier()
  Wait for "exit_barrier_timeout" seconds for all agents to finish
  executing their local workers (either successfully or not). This
  acts as a safety guard against user scripts that terminate at
  different times. This barrier keeps the agent process alive
  until all workers finish.

_initialize_workers(worker_group)
  Starts a fresh set of workers for the worker_group. Essentially
  a rendezvous followed by a start_workers.

  The caller should first call "_stop_workers()" to stop running
  workers prior to calling this method.

  Optimistically sets the state of the worker group that just
  started as "HEALTHY" and delegates the actual monitoring of
  state to "_monitor_workers()" method

abstract _monitor_workers(worker_group)
abstract _monitor_workers(worker_group)
  Checks on the workers for the "worker_group" and returns the new
  state of the worker group.

  Return type:
     *RunResult*

_rendezvous(worker_group)
  Runs rendezvous for the workers specified by worker spec.
  Assigns workers a new global rank and world size. Updates the
  rendezvous store for the worker group.

_restart_workers(worker_group)
  Restarts (stops, rendezvous, starts) all local workers in the
  group.

abstract _shutdown(death_sig=Signals.SIGTERM)
  Cleans up any resources that were allocated during the agent's
  work.

  Parameters:
     **death_sig** (*Signals*) -- Signal to send to the child
     process, SIGTERM is default

abstract _start_workers(worker_group)
  Starts "worker_group.spec.local_world_size" number of workers
  according to worker spec for the worker group .

  Returns a map of "local_rank" to worker "id".

  Return type:
Return type:
         Dict[int, Any]
abstract _stop_workers(worker_group)
  Stops all workers in the given worker group. Implementors must
  deal with workers in all states defined by "WorkerState". That
  is, it must gracefully handle stopping non-existent workers,
  unhealthy (stuck) workers, etc.

class torch.distributed.elastic.agent.server.api.RunResult(state, return_values=, failures=)
Results returned by the worker executions. Run results follow an
   "all-or-nothing" policy where the run is successful if and only if
   ALL local workers managed by this agent complete successfully.
If the result is successful (e.g. "is_failed() = False") then the
   "return_values" field contains the outputs (return values) of the
   workers managed by THIS agent mapped by their GLOBAL ranks. That is
   "result.return_values[0]" is the return value of global rank 0.
Note:
 "return_values" are only meaningful for when the worker