Resilient Training with Ray and Jax

This guide demonstrates how resilient training of models written in libraries built on Jax can be achieved using Ray. More specifically, we will delve into the following:

• What a Ray cluster is and how to bring a Ray cluster up on underlying hardware infrastructure
• How to leverage the brought up Ray cluster to implement fault tolerant finetuning of a pretrained LLM; this includes:
  • automatic recovery from failures and hangs by leveraging Ray’s machinery for failure detection and recovery from a model checkpoint
  • fast recovery from failures and hangs by leveraging model compilation caching

For a full runnable example of resilient JAX training using Ray see the simple toy example, which includes a Docker image you can try out on a single node with a least two GPUs. The guide below gives an overview of the example code and how you might want to extend it to your own model.

Why Ray?

Ray’s runtime is based on a “head”/“worker” architecture where a “head” process called the coordinator, spawns “worker” processes called actors. This architecture allows computation to be run across the workers and the coordinator to act as a watchdog and manage the control plane. As a result, the coordinator can detect any crash or hang experienced by an actor and start the recovery process. This can be viewed as a single-controller runtime where the coordinator is the single-controller. We are going to build our resilient training system using Ray’s single-controller runtime, with the actors being responsible for training the model and the coordinator being responsible for failure detection and recovery.

The workflow for implementing resilient training with Ray and JAX can be broadly split up into the following steps:

• Launching a Ray cluster
• Implementing the coordinator and actors
• Launching a job with the coordinator and actors on the Ray cluster

We will delve into each of these steps in detail in that order in the rest of this guide.

Starting the Ray Cluster

The first step to launching any workload using Ray is to start a Ray Cluster. The Ray cluster consists of a single “head node” and a number of “worker nodes”. There is a slight abuse of terminology that can be a little confusing here with the usage of the word “node” but both Ray head nodes and Ray worker nodes can in general be thought of as sub-allocations within a physical node. This means that on a physical node with 8 GPUs and 128 CPUs, a Ray worker node could pertain to a sub-allocation of up to 8 GPUs and 128 CPUs. And similarly for the Ray head node. The coordinator process always runs on the Ray head node and actors always run on Ray worker nodes. In this guide we will assume that each actor gets 1 GPU, 16 CPUs and that we operate in a 1 GPU per process setting.

Starting a Ray Cluster manually

We will begin with a simple example of how to manually start a Ray cluster on 2 physical nodes. This will involve a single Ray head node and 2 Ray worker nodes, where each Ray worker node is allocated all GPUs of the node it runs on. Moreover, each worker node will have as many actors scheduled on it as it has GPUs available, with each actor being able to address only a single GPU. This ensures that we are indeed operating in a 1 GPU per process setting. We will assume the IP addresses of the physical nodes are IP_ADDR_1 and IP_ADDR_2 and that the head node will be allocated on the physical node with IP_ADDR_1.

First, run the following script on one physical node:

1
#!/bin/bash
2
3
# Number of GPUs per node
4
gpus_per_node=8
5
min_worker_port=10001
6
max_worker_port=10257
7
8
# On the physical node with `IP_ADDR_1`:
9
CUDA_VISIBLE_DEVICES="" ray start --head --node-ip-address=IP_ADDR_1 --port=<HEAD_NODE_PORT>
10
ray start --address "IP_ADDR_1:<HEAD_NODE_PORT>" \
11
                    --resources="{\"worker_units\": $gpus_per_node}" \
12
                    --min-worker-port=$min_worker_port \
13
                    --max-worker-port=$max_worker_port

The script above will allocate the Ray head node on the physical node with IP_ADDR_1 and also start 1 ray worker node on the same physical node. We note the following about the script above:

• The address argument to ray start refers to the address of the Ray head node.
• We specify CUDA_VISIBLE_DEVICES="" when allocating the Ray head node. This is because the coordinator will not be performing any GPU computation and thus the Ray head node required no GPUs.
• Resources like GPUs, CPUs and memory are physically afforded to Ray worker nodes by the underlying physical node. But Ray allows users to define custom virtual resources to give users more finegrained control of actor scheduling. In this case we define a physical resource called a worker_unit. Since we are going to be assigning one process per GPU we give each ray worker node 8 units of the custom worker_unit resource. This allows us to “tag” the ray worker nodes that will be used as the primary workers for our finetuning job. To ensure that actors only ever gets scheduled on these ray worker nodes we will request a single worker_unit for each actor during scheduling.

At this point the Ray cluster has 1 Ray head node and 1 Ray worker node. We will allocate the second Ray worker node on the physical node with ip address IP_ADDR_2 by running the following script:

1
#!/bin/bash
2
3
# Number of GPUs per node
4
gpus_per_node=8
5
min_worker_port=10001
6
max_worker_port=10257
7
8
ray start --address "IP_ADDR_1:<HEAD_NODE_PORT>" \
9
                    --resources="{\"worker_units\": $gpus_per_node}" \
10
                    --min-worker-port=$min_worker_port \
11
                    --max-worker-port=$max_worker_port

Now the Ray cluster consists of 1 Ray head node and 2 Ray worker nodes, all allocated across the 2 physical nodes with IP addresses IP_ADDR_1 and IP_ADDR_2.

Implementing the coordinator and actors

Now that we have an understanding of how to bring up a ray cluster, we are almost in a position to be able to launch a job on this cluster. Before we do so we need to define the actions of the coordinator and the actors.

Designing the coordinator

Let us start defining a class called RayClusterCoordinator, that encapsulates the functionality of the coordinator. We will start simple and add functionality as we go:

1
import ray
2
import random
3
import asyncio
4
import redis
5
        
6
class RayClusterCoordinator:
7
    def __init__(self, worker_cls) -> None:
8
        self.worker_cls = worker_cls
9
        self.num_workers = int(os.environ.get('NGPUS'))
10
11
        self.workers = [worker_cls.options(num_gpus=1, 
12
                                            num_cpus=16, 
13
                                            resources={"worker_units": 1}).remote() for _ in range(self.num_workers)]
14
   
15
    def initialize_workers(self, **kwargs):
16
        self.worker_init_kwargs = kwargs
17
        coordinator_ip = ray.get(self.workers[0].get_host_ip.remote())
18
        coordinator_port = random.randint(1, 100000)  % 2**12 + (65535 - 2**12 + 1)
19
        self.jax_coordinator_addr = f"{coordinator_ip}:{coordinator_port}"
20
        ray.get([w.initialize.remote(i, self.jax_coordinator_addr, self.num_workers, **kwargs) for i, w in enumerate(self.workers)])
21
        self.workers_initialized = True
22
    
23
    async def run(self, **kwargs):
24
        if not self.workers_initialized:
25
            raise Exception("""Cannot run workers without initializing them first. 
26
                               Please call the initialize_workers method of your cluster coordinator first.""")
27
        # Launch computation on actors...
28
    

The first job of the coordinator is to spawn actors and have them be scheduled on Ray worker nodes. This happens in the constructor of the RayClusterCoordinator. worker_cls is the class that encapsulates the functionality of each actor. Notice how we request 1 GPU, 16 CPUs and 1 worker_unit per actor as mentioned earlier. There are two additional mandatory methods that RayClusterCoordinator implements:

• initialize_workers: the purpose of this method is to trigger the initialize method of the worker_cls representing an actor
• run: the purpose of this method is to launch computation on all the scheduled actors. We haven’t included the details of the run method quite yet, nor have we discussed why it is an async method. We will get to both after defining some of the functionality of the class representing an actor. In the class definition above NGPUS is an environment variable that represents the total number of GPUs available to the job across all nodes.

Designing the actor

Our design of the actor involves defining two classes: ResilientWorker and ModelTrainer. ResilientWorker will define a high level interface for the role an actor plays in this system and ModelTrainer is a subclass of ResilientWorker that will implement specific functionality to train a model in a fault-tolerant manner, including model checkpointing and compilation caching. Let’s start with ResilientWorker.

1
import jax
2
import socket
3
import os
4
import redis
5
6
class ResilientWorker:
7
    def __init__(self):
8
        self.host_ip = socket.gethostbyname(socket.gethostname())
9
        self.logical_gpu_id = int(os.environ.get('CUDA_VISIBLE_DEVICES'))
10
        
11
    def get_process_id(self):
12
        return self.process_id
13
    
14
    def get_host_ip(self):
15
        return self.host_ip
16
     
17
    def get_logical_gpu_id(self):
18
        return self.logical_gpu_id
19
        
20
    def initialize(self, process_id, coordinator_addr, num_processes):
21
        self.process_id = process_id
22
        jax.distributed.initialize(coordinator_address=coordinator_addr, num_processes=num_processes, process_id=self.process_id, local_device_ids=0)
23
    
24
    def run(self, *args, **kwargs):
25
        raise NotImplementedError

The ResilientWorker interface defines a number of methods, the most important of which are the initialize and run methods:

• initialize: is responsible for initializing jax.distributed across all the spawned actors. Every class that inherits from ResilientWorker can initialize any task specific state in its initialize method, but must call the base class’ initialize method to setup the jax.distributed runtime.
• run: every class that inherits from ResilientWorker must implement its own run method. This is the entry point to the computation that the actors will perform.

To understand how this class and inheritance structure works, let’s start implementing the ModelTrainer class.

1
import ray
2
import jax
3
import orbax.checkpoint as ocp
4
from optax._src.transform import ScaleByAdamState
5
from optax._src.base import EmptyState
6
7
@ray.remote
8
class ModelTrainer(ResilientWorker):
9
    def __init__(self):
10
        super().__init__()
11
    
12
    def initialize(self, process_id, coordinator_addr, num_processes, **kwargs) -> None:
13
        super().initialize(process_id, coordinator_addr, num_processes)
14
        self.config = self._get_training_config(kwargs['training_config_dict'])
15
        self.ckpt_dir = kwargs['ckpt_dir']
16
        self.ckpt_freq = kwargs['ckpt_freq']
17
        ckpt_mgr_options = ocp.CheckpointManagerOptions(save_interval_steps=self.ckpt_freq, max_to_keep=5)
18
        self.ckpt_mgr = ocp.CheckpointManager(self.ckpt_dir, 
19
                                              item_names=('params', 'opt_state', 'stop_point'), 
20
                                              options=ckpt_mgr_options)
21
        jax.config.update("jax_compilation_cache_dir", kwargs['jax_compilation_cache'])
22
                                              
23
    def run(self, num_physical_nodes, restore=False) -> List[Any] | None:
24
        # Method implementing the training loop

In this guide, the objective of the actors is to finetune a pretrained LLM. Additionally, they have to do so in a fault tolerant manner. ModelTrainer is the class that encapsulates this behavior and inherits from ResilientWorker. Its responsibility to finetune an LLM while being able to recover quickly and automatically when met with a failure. The class is decorated with ray.remote since instances of this class are to be scheduled as actors.

As a subclass of ResilientWorker, ModelTrainer must implement its own run method. But before we look at the details of the run method let’s look at the details of the overridden initialize method in ModelTrainer. It is not mandatory for subclasses of ResilientWorker to override initialize, but it is often useful to initialize child class specific state by doing so. In this example, given our objective, in addition to calling the base class’ initialize method to initialize jax.distributed, ModelTrainer initializes the following in its initialize method:

1. model checkpointing state: initialization of an Orbax checkpoint manager. One of the important aspects of fault tolerant training is to be able to recover from the latest checkpoint. Hence our model is checkpointed at fixed intervals during the finetuning process
2. jax compilation cache: to avoid the penalty of recompilation when restoring the training run from the latest checkpoint after a crash or a hang

We will now look at the details of the run method of ModelTrainer with the following caveat: we will merely gloss over parts of the run method that don’t pertain directly to fault tolerant training, these include:

• Dataset creation and tokenizing
• Sharding model parameters
• Dataloading details
• Implementation of a single training step

The methods below provide signatures for some of these functions (that will be used in run):

1
def _load_model_params(self):
2
    ### Load model parameters from a pretrained checkpoint ###
3
    ...
4
    ...
5
    return pretrained_checkpoint
6
7
def _shard_parameters(self, params, mesh):
8
    ### Shard parameters across mesh ###
9
    ...
10
    ...
11
    return sharded_parameters
12
13
def _shard_inputs(self, *inputs, mesh):
14
    ### Shard inputs across mesh ###
15
    ...
16
    ...
17
    return sharded_inputs
18
19
@staticmethod
20
def train_step(
21
    model, 
22
    params, 
23
    optimizer: optax.GradientTransformation, 
24
    opt_state: optax.OptState, 
25
    tokens: jax.Array, 
26
    input_mask: jax.Array,
27
) -> Tuple[jax.Array, ParameterPytree, optax.OptState]:
28
    ### Single Forward + Backward pass ###

Details on how to implement these aspects of model training can be found in many excellent guides online, one of which is this tutorial on finetuning a Gemma model. We will instead focus on the aspects that demonstrate how Ray can be used to achieve fault tolerant training. With the helper methods defined as above, the run method looks as follows:

1
def run(self, num_nodes, restore=False) -> List[Any] | None:
2
        # Init / load model parameters
3
        if not restore:
4
            # Load from pre-trained checkpoint
5
            model, params = self._load_model_params()
6
            start_epoch, start_iter, ckpt_step = 0, 0, 0
7
        else:
8
            # Load from checkpoint saved during training, includes saved optimizer state
9
            model, params, opt_state, stop_points = self.restore()
10
            start_epoch, start_iter, ckpt_step = stop_points["epoch"], stop_points["iter"], stop_points["ckpt_step"]
11
12
        # Shard model parameters
13
        dp_dim = num_nodes
14
        tp_dim = len(jax.devices()) // dp_dim
15
        devices = np.asarray(jax.devices()).reshape(dp_dim, tp_dim)
16
        mesh = Mesh(devices, ('DP', 'TP'))
17
18
        if not restore:
19
            params = self._shard_parameters(params, mesh)
20
21
        # Create data loader
22
        train_dataloader = DataLoader(...)
23
24
        # Compile the function that performs a single step of forward + backward
25
        jitted_train_step = jax.jit(ModelTrainer.train_step, static_argnames=['model','optimizer'])
26
        
27
        # Initialize optimizer
28
        optimizer = optax.adam(learning_rate=1e-3)
29
        if not restore:
30
            # First time initializing the optimizer state
31
            opt_state = optimizer.init(params)
32
        
33
        # Training loop
34
        for epoch in range(start_epoch, NUM_EPOCHS):
35
            for i, (tokens, mask) in enumerate(train_dataloader):
36
                # Skip computation till we reach the position
37
                # in the dataloader where the failure occurred
38
                if epoch == start_epoch and i <= start_iter:
39
                    continue
40
                
41
42
                global_tokens, global_mask = self._shard_inputs(tokens, mask, mesh)      
43
                train_loss, params, opt_state = jitted_train_step(model, 
44
                                                                  params, 
45
                                                                  optimizer, 
46
                                                                  opt_state, 
47
                                                                  global_tokens, 
48
                                                                  global_mask, 
49
                )
50
                self.checkpoint(params, opt_state, epoch, i, ckpt_step)             

Let’s examine the details of the run method above. It takes two arguments num_nodes and restore. num_nodes is the number of physical nodes that all the actors are spawned across and is only used to create the mesh that computation is going to be spread out across. restore is a parameter that lets run know if run has been called during the recovery process after a failure or if it is at the start of a training job. If restore is false, the method looks very familiar to a standard training loop: a model is created, a data loader is created, parameters are sharded, the optimizer state is initialized and the model is optimized in the training loop with checkpointing being facilitated by the following method:

1
def checkpoint(self, params, opt_state, epoch, iter, ckpt_step):
2
    self.ckpt_mgr.wait_until_finished()
3
    stop_point = {"epoch" : epoch, "iter" : iter, 'ckpt_step' : ckpt_step}
4
    opt_state_pytree = {"count" : opt_state[0].count, "mu" : opt_state[0].mu, "nu" : opt_state[0].nu}
5
    self.ckpt_mgr.save(ckpt_step, args=ocp.args.Composite(params=ocp.args.StandardSave(params), 
6
                                                  opt_state=ocp.args.StandardSave(opt_state_pytree), 
7
                                                  stop_point=ocp.args.JsonSave(stop_point))
8
    )

When restore is true, there are a number of differences. The first is that parameters are not loaded from the pretrained checkpoint, but are instead loaded from the latest saved checkpoint. The latest saved checkpoint also contains the latest optimizer state and some additional parameters (epoch, iteration anhd ckpt_step) that help pick up training from the point at which the failure took place. The restore method looks as follows:

1
def restore(self):
2
        ckpt =  self.ckpt_mgr.restore(self.ckpt_mgr.latest_step(), 
3
                                      args=ocp.args.Composite(params=ocp.args.StandardRestore(), 
4
                                                              opt_state=ocp.args.StandardRestore(), 
5
                                                              stop_point=ocp.args.JsonRestore())
6
        )
7
8
        params, opt_state_pytree, stop_point = ckpt.params, ckpt.opt_state, ckpt.stop_point
9
        model = Model(...) # Flax implementation of the language model being finetuned
10
        opt_state = (ScaleByAdamState(count=opt_state_pytree["count"], mu=opt_state_pytree["mu"], nu=opt_state_pytree["nu"]), EmptyState())
11
        return model, params, opt_state, stop_point

So now that we have the run method for our actors in place, how exactly do we launch them? And more importantly, how do we detect and recover from failures? To do so, we have to go back to the coordinator and flesh out the run method of the RayClusterCoordinator class.

Revisiting the coordinator

As mentioned earlier, the purpose of the run method of the RayClusterCoordinator is to launch computation on actors. Now that we have the computation that each actor is going to perform defined in the run method of ModelTrainer, we will have the RayClusterCoordinator call it to launch computation on the actors. Additionally, we will also include logic to detect failures and start recovery in the run method of the RayClusterCoordinator. The first version of this method looks as follows:

1
# Run method of RayClusterCoordinator
2
async def run(self, **kwargs):
3
    if not self.workers_initialized:
4
        print(f"Cannot run workers without initializing them first. Please call the initialize_workers method of your cluster coordinator first.")
5
        return []
6
7
    task = asyncio.create_task(self._run_workers_async(**kwargs))
8
    while True:
9
        try:
10
            task_results = await asyncio.wait(task)
11
        except Exception as e:
12
            # Failure in one or more of the launched Ray actors, start recovery
13
            
14
            # Step 1: Kill the current task
15
            task.cancel()
16
            
17
            # Step 2: Kill all actors to free resources (all actors besides the failed actors will be in hanged state)
18
            for w in self.workers:
19
                ray.kill(w)
20
            self.workers_initialized = False
21
22
            # Step 3: Recreate actor handles and assign them to worker nodes
23
            self.workers = [self.worker_cls.options(num_gpus=1, num_cpus=16, resources={"worker_units": 1}).remote() 
24
                            for _ in range(self.num_workers)]
25
            
26
            # Step 4: Initialize the recreated actors and set the restore argument of the Actor run method to True
27
            self.initialize_workers(**self.worker_init_kwargs)
28
            kwargs["restore"] = True
29
30
            # Step 5: Restart task to launch actor computation
31
            task = asyncio.create_task(self._run_workers_async(**kwargs))
32
        else:
33
            return task_results[0]

Let’s examine this method. The first detail to note (as mentioned when RayClusterCoordinator was introduced) is that the method is an async method. The reason for this is that we want the coordinator to be able to perform multiple non-blocking tasks. One of which is launching actors and being able detect when any of them fail. This is what the _run_workers_async method does:

1
async def _run_workers_async(self, **kwargs):
2
    worker_run_futures = [w.run.remote(self.num_physical_nodes, **kwargs) for w in self.workers]
3
    while True:
4
        for i, wf in enumerate(worker_run_futures):
5
            try:
6
                ray.get(wf, timeout=0)
7
            except ray.exceptions.GetTimeoutError:
8
                pass
9
10
        # Give control back to the asyncio event loop
11
        await asyncio.sleep(30)

In this method each w is a reference to a ModelTrainer actor and [w.run.remote(self.num_physical_nodes, **kwargs) for w in self.workers] calls the run method of each actor to start the finetuning job. When a Ray actor fails, it sends a signal back to the coordinator that can be handled. This happens whether the failure is “graceful”, in that the Actor is able to throw an exception; or if the failure is less “graceful” like a segmentation fault. Calling ray.get(wf, timeout=0) on any future from an actor that has not failed, will always result in a ray.exceptions.GetTimeoutError, if the actor’s run method has not completed execution. However, if the actor has failed in any manner, graceful or not, ray.get(wf, timeout=0) will always throw one of ray.exceptions.RayTaskError or ray.exception.ActorDiedError, which propagates back to the coordinator.

When the coordinator receives a failure signal from the _run_workers_async task, it can start the recovery process. This is what the code in the except block of the run method in RayClusterCoordinator does.

Adding hang detection capabilities

In its current state, this system can recover from crashes in any of the actors. But another important, and trickier class of failures to be able to detect and recover from is hangs. Hangs can occur for a variety of reasons during a run and so it’s important for the system to be able to detect them in a general sense.

Being able to detect hangs requires more proactivity on the part of the coordinator, for the simple reason that if any actor is in a hanged state, it cannot send any signals to the coordinator to communicate that it is in such a state. So in addition to the crash detection logic outlined in _run_workers_async we will have the coordinator perform its second non-blocking task via another async function called _detect_worker_hang_async, into which we will add the logic to detect hangs. At a high level, the logic to detect hangs depends on the actors “checking in” periodically with the coordinator via a distributed key-value store. This is what we will use Redis for. To start the redis-server simply run the following on the same physical node that the ray head node is running on.

1
redis-server --bind $head_node_ip --port 6380 --protected-mode no --daemonize yes &&
2
export REDIS_ADDR=$head_node_ip:6380

With the redis server running we need to have both the coordinator and each actor connect to the server, which we can do by adding the following two lines to the constructor of RayClusterCoordinator and ResilientWorker:

1
self.redis_addr = os.environ.get('REDIS_ADDR').split(':')
2
self.redis = redis.Redis(host=self.redis_addr[0], port=int(self.redis_addr[1]), decode_responses=True, password=None)

As mentioned, the the hang detection logic we are yet to implement in _detect_worker_hang_async depends on actors “checking in” with the coordinator. This is facilitated through a heartbeat in the form of a timestamp from each actor being written to the redis KV store. They key is the process ID and the value is the timestamp of the last time the actor was alive. This is implemented through a context manager that we add to the ResilientWorker class:

1
@contextmanager
2
def EnableHeartbeat(self):
3
    try:
4
        yield
5
    finally:
6
        self.redis.set(self.process_id, time.time())

and then wrapping the computation in the the main training loop in the run method of ModelTrainer with this context manager:

1
def run(self, num_nodes, restore=False) -> List[Any] | None:
2
    
3
    # All the code before the training loop is the same ...
4
5
    for epoch in range(start_epoch, NUM_EPOCHS):
6
        for i, (tokens, mask) in enumerate(train_dataloader):
7
            if epoch == start_epoch and i <= start_iter:
8
                continue
9
10
            with self.EnableHeartbeat():
11
                global_tokens, global_mask = self._shard_inputs(tokens, mask, mesh)      
12
                train_loss, params, opt_state = jitted_train_step(model, 
13
                                                                  params, 
14
                                                                  optimizer, 
15
                                                                  opt_state, 
16
                                                                  global_tokens, 
17
                                                                  global_mask, 
18
                )
19
                self.checkpoint(params, opt_state, epoch, i, ckpt_step)   

That’s it for the actors. Now we can implement _detect_worker_hang_async to complete the hang detection feature. The logic here is straightforward, wait a certain amount of time t seconds and then check the timestamp of each actor’s heartbeat in the redis KV store. If an actor hasn’t checked in in t seconds, it’s likely hanged. It looks as follows:

1
async def _detect_worker_hang_async(self):       
2
    # Check if processes are hanging
3
    while True:
4
        # Wait for 5 minutes
5
        await asyncio.sleep(300)
6
        for pid in range(self.num_workers):
7
            elapsed = time.time() - float(self.redis.get(pid))
8
            if elapsed > 300:
9
                raise Exception(f"Worker {self._worker_process_ids[pid]} appears to have hanged")

A couple of points to note about the method above:

• The first is that the asyncio.sleep serves two purposes. One if to wait for t as part of the hang detection logic. The other is to give control back to the asyncio event loop so that it can switch contexts to the _run_workers_async method to check if any actors have failed.
• The chosen time of t=300 does mean that it will take at least 5 minutes to detect a hang after its occurence. This is the best case and occurs if and only if an actor hangs at the exact moment the asyncio.sleep(300) command is issued. In the worst case, the time from occurrence to detection is 10 minutes and happen if the hang occurs at the exact moment the asyncio.sleep(300) command finishes. However, t is configurable and the chosen value here is just meant to illustrate functionality.

With _detect_worker_hang_async implemented, we just need to change the run method of RayClusterCoordinator to launch the task. The final updated method looks as follows:

1
async def run(self, **kwargs):
2
    if not self.workers_initialized:
3
        print(f"Cannot run workers without initializing them first. Please call the initialize_workers method of your cluster coordinator first.")
4
        return []
5
6
    tasks = [asyncio.create_task(self._run_workers_async(**kwargs)), 
7
             asyncio.create_task(self._detect_worker_hang_async())]
8
    while True:
9
        try:
10
            task_results = await asyncio.gather(*tasks)
11
        except Exception as e:
12
        
13
            for task in tasks:
14
                task.cancel()
15
16
            for w in self.workers:
17
                ray.kill(w)
18
            self.workers_initialized = False
19
20
            self.workers = [self.worker_cls.options(num_gpus=1, num_cpus=16, resources={"worker_units": 1}).remote(self.log_dir) 
21
                            for _ in range(self.num_workers)]
22
            self.initialize_workers(**self.worker_init_kwargs)
23
            kwargs["restore"] = True
24
25
            tasks = [asyncio.create_task(self._run_workers_async(**kwargs)), 
26
                     asyncio.create_task(self._detect_worker_hang_async())]
27
        else:
28
            return task_results[0]

When the coordinator determines an actor has hanged, it raises an exception that is handled exactly like a crash is handled, except during the recovery process it needs to kill two tasks instead of one.

Launching a job Ray job

With that we’ve described every important aspect of the coordinator and the actors that help them work together to achieve failure resilient training. The final piece of the puzzle is to launch the training job that leverages all the logic implemented in RayClusterCoordinator, ResilientWorker and ModelTrainer, on the Ray cluster. This is achieved through the following entrypoint (which could be in its own main.py file or all the code above as well as the entrypoint could be implemented in a large main.py file):

1
# main.py
2
import os
3
import ray
4
5
ray.init(address='auto')
6
# Create the Ray cluster coordinator object
7
cluster_coorindator = RayClusterCoordinator(ModelTrainer)
8
# Get the job configuration set during launch.
9
# This is automatically set by Ray
10
job_runtime_env = json.loads(os.environ.get('RAY_JOB_CONFIG_JSON_ENV_VAR'))['runtime_env']
11
12
# Initialize workers
13
cluster_coordinator.initialize_workers(jax_compilation_cache=job_runtime_env['jax_compilation_cache'],
14
                                       ckpt_dir=job_runtime_env['ckpt_dir'],
15
                                       ckpt_freq=job_runtime_env['ckpt_freq'])
16
17
# Run the workers
18
run_results = asyncio.run(cluster_coordinator.run(restore=False))

and a script called the driver as shown below:

1
# launch_ray_cluster_job.py
2
3
from ray.job_submission import JobSubmissionClient, JobStatus
4
5
# Jobs are scheduled on a Ray cluster through a JobSubmissionClient
6
# The JobSubmissionClient connect to the Ray cluster via localhost
7
# Create the JobSubmissionClient
8
client = JobSubmissionClient("http://127.0.0.1:8265")
9
# Submit main.py to run on the Ray cluster whose head node's IP
10
# address is localhost
11
job_id = client.submit_job(
12
            entrypoint="python3 main.py",
13
            runtime_env={"working_dir" : ".",
14
                        "jax_compilation_cache" : "/path/to/compilation_cache",
15
                        "ckpt_dir" : "/path/to/checkpoint/directory",
16
                        "ckpt_freq" : 100,
17
            }
18
)
19
20
prev_logs = ''
21
while True:
22
    status = client.get_job_status(job_id)
23
    if status in {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}:
24
        if status in {JobStatus.STOPPED, JobStatus.FAILED}:
25
            logs = client.get_job_logs(job_id)
26
            print(logs, flush=True)
27
        break
28
    time.sleep(5)
29
    if status == JobStatus.RUNNING:
30
        logs = client.get_job_logs(job_id)
31
        print(logs[len(prev_logs):], flush=True)
32
        prev_logs = logs

The driver script is run on the same physical node as the Ray head node and is the reason the address to connect the JobSubmissionClient to the Ray cluster is http://127.0.0.1 or localhost.

Starting a Ray Cluster with a SLURM allocation

Sometimes it is likely that physical resources are allocated to an application by a scheduler like SLURM or Kubernetes. The script below is a template of a SLURM batch script that sets up a ray cluster with the specification as above (1 Ray worker node per physical node / 8 GPUs per Ray worker node), but over <NUM_NODES> physical nodes.

1
#!/bin/bash
2
#SBATCH --nodes=<NUM_NODES>
3
#SBATCH --exclusive
4
#SBATCH --account=<SLURM_ACCOUNT>
5
#SBATCH --partition=<SLURM_PARTITION>
6
#SBATCH --time=<SLURM_ALLOCATION_TIME>
7
8
# Number of GPUs per node
9
gpus_per_node=8
10
11
# Getting the node names and IP addresses in the SLURM allocation
12
nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
13
nodes_array=($nodes)
14
ip_addresses_array=()
15
16
for node in $nodes; do
17
    ip_address=$(host $node | awk '/has address/ { print $4 }')
18
    # Add the IP address to the array
19
    ip_addresses_array+=("$ip_address")
20
done
21
22
head_node=${nodes_array[0]}
23
head_node_ip=${ip_addresses_array[0]}
24
25
port=<PORT>
26
ip_head=$head_node_ip:$port
27
28
# First we start the head of the ray cluster on one of the physical nodes
29
# In this case we are giving an entire physical node to the ray head node
30
# The ray head node is marked by including --head to the ray start command
31
srun --nodes=1 --ntasks=1 -w "$head_node" bash -c "redis-server --bind $head_node_ip --port 6380 --daemonize yes && \
32
                                                   ray start --head --node-ip-address="$head_node_ip" --port=$port --block" &
33
34
export REDIS_ADDR=$head_node_ip:6380
35
36
# We now need to start the Ray worker nodes
37
# We define the following variables to help
38
num_ray_worker_nodes=$SLURM_JOB_NUM_NODES 
39
export NGPUS=$((gpus_per_node * num_ray_worker_nodes)) # 8 actors per ray worker node (one for each GPU)
40
41
# Start Ray worker nodes
42
# We want 1 Ray worker node per physical node
43
# Worker nodes are started with ray start but without the --head flag
44
min_worker_port=10001
45
max_worker_port=10257
46
for ((i = 0; i < num_ray_worker_nodes; i++)); do
47
    node_i=${nodes_array[$i]}
48
    
49
    srun --exact --nodes=1 --ntasks=1 --cpus-per-task=$((16 * $gpus_per_node)) -w "$node_i" \
50
    ray start --address "$ip_head" \
51
              --resources="{\"worker_units\": $gpus_per_node}" \
52
              --min-worker-port=$min_worker_port \
53
              --max-worker-port=$max_worker_port --block &
54
    sleep 3
55
  done
56
done
57
58
# At this stage the Ray cluster bringup has started on the physical nodes in the allocation
59
# Before we launch a job on this cluster we need to make sure that the bringup is complete
60
# We do so by querying the number of worker_units in the ray cluster and asserting = NGPUS
61
extract_worker_units() {
62
  status_output=$(srun --overlap --nodes=1 --ntasks=1 -w "$head_node" ray status)
63
  if echo "$status_output" | grep -q "worker_units"; then
64
    worker_units=$(echo "$status_output" | grep "worker_units" | awk -F'[/. ]' '{print $4}')
65
    echo $worker_units
66
  else
67
    echo 0
68
  fi
69
}
70
71
# Poll to make sure that all Ray worker nodes have connected to the head.
72
# All workers have connected when number of GPUs in ray cluster
73
# is equal to NGPUS. We use the utility function above
74
# to check how many GPUs have come online in the ray cluster
75
while true; do
76
  worker_units=$(extract_worker_units)
77
  if [ "$worker_units" -eq "$NGPUS" ]; then
78
    break
79
  fi
80
  sleep 2
81
done
82
83
echo "All workers connected!"
84
85
# We can now launch a job on this cluster
86
# We do so by launching the driver process on the physical node that the head node is on
87
# This driver process is responsible for launching a job on the Ray cluster
88
srun --overlap --nodes=1 --ntasks=1 -w "$head_node" python3 launch_ray_cluster_job.py

Sharp bits and current limitations

This tutorial is meant to serve as a guide for how to use Ray to implement a fault tolerant training framework with JAX, but we’d like to note some limitations of the infrastructure explained in this tutorial:

1. While most aspects of the design are agnostic to the type of hardware being used, it has only been tested on GPUs. It is likely that the same principles can be applied to make the design applicable with minimal change to run on TPUs and other accelerators but that has not been tested yet.

2. The implicit assumption in the design is that the same number of physical nodes is available before and after failure; meaning that it doesn’t support recovery in the case of dead nodes unless hot spares are available.