HyperPod elastic agent

The HyperPod elastic agent is an extension of PyTorch’s ElasticAgent. It orchestrates lifecycles of training workers on each container and communicates with the HyperPod training operator. To use the HyperPod training operator, you must first install the HyperPod elastic agent into your training image before you can submit and run jobs using the operator. The following is a docker file that installs elastic agent and uses hyperpodrun to create the job launcher. For more information about the HyperPod training operator, installation of the operator, and how to use it, see the official AWS documentation.

RUN pip install hyperpod-elastic-agent

ENTRYPOINT ["entrypoint.sh"]
# entrypoint.sh
...
hyperpodrun --nnodes=node_count --nproc-per-node=proc_count \
            --rdzv-backend hyperpod \ # Optional
            ... # Other torchrun args
            # pre-traing arg_group
            --pre-train-script pre.sh --pre-train-args "pre_1 pre_2 pre_3" \
            # post-train arg_group
            --post-train-script post.sh --post-train-args "post_1 post_2 post_3" \
            training.py --script-args

You can now submit jobs with kubectl.

Note: When invoking hyperpodrun in a container entrypoint, ensure that hyperpodrun is PID 1 or alternatively is able to receive signals(specifically SIGTERM) forwared from PID 1 in the container to allow graceful termination during container/pod teardown. This can be done in the following ways:

• Use exec when specifying ENTRYPOINT in shell mode, for ex: ENTRYPOINT exec hyperpodrun .....
• Use exec format to speicify the container entrypoint, for ex: ENTRYPOINT ["hyperpodrun", .....]
• Use an init process that can forward signals to its child processes, for ex: tini

HyperPod elastic agent arguments

The HyperPod elastic agent supports all of the original arguments from the PyTorch Elastic Agent official documentation. The following is a list of additional arguments available in the HyperPod elastic agent: