[Ray] Add elasticity and fault tolerance features to jobs launched on ray cluster

[ntlm1686]

Two features for Elastic Distributed Training are added to job launched by TorchX on Ray Cluster in this PR:

1. Fault Tolerance - Node failure throws RayActorError which can be captured. Placement Groups have built-in fault tolerance, and can recover from node failure automatically.
2. Elasticity - the execution of placement groups are pending tasks that will be scheduled by GCS when resources become availiable.

The logic of the new ray_driver.py is in the plot below:

Test plan:

[Note]: This PR is related to the previous PR #559. All future changes will be submitted to this PR.

[d4l3k]

getting into a nicer shape :) I like how you used different return values as state machine inputs

[d4l3k]

nit: min_replicas: Optional[int] = None

[d4l3k]

nit: would be nice to get rid of this parsing by passing in min_replicas instead of rep

[d4l3k]

not a big fan of this but not sure if there's any better approach

[ntlm1686]

It seems actor_id is the only thing useful we get from the exception without changing ray's code. Only error messages. https://github.com/ray-project/ray/blob/4c5c5763efff50bf9f76ae73a8c0073183dbb0cc/python/ray/exceptions.py#L233

[d4l3k]

instead of parsing can we use a map between the ObjectRef to the actor info? Seems like the ObjectRefs returned from wait should be equal to the ones passed in

https://github.com/ray-project/ray/blob/43aa2299e6623c8f8c7c4a1b80133459d0aa68b0/python/ray/includes/object_ref.pxi#L38

[ntlm1686]

I checked, but I think there isn't actor related information. https://github.com/ray-project/ray/blob/4c5c5763efff50bf9f76ae73a8c0073183dbb0cc/python/ray/includes/object_ref.pxi#L38

But we can create some issues about this one and the exception one.

[d4l3k]

nit spelling

[d4l3k]

it might be better to explicitly make the calls here. + we need to wait for the first placement group before creating the remainders otherwise the smaller groups might get scheduled first (unless ray does fifo queuing, which if it does we should document that)

initial = create_placement_group_async(self.replicas[0: self.min_nodes])
pg.wait(itimeout_seconds=...)
self.groups = []
for range(self.min_nodes, self.max_nodes):
    self.groups.append(create_placement_group_async(self.replicas[i:i+1]))

[ntlm1686]

Ray doesn't do FIFO, will change.

[d4l3k]

this only breaks the inner loop -- is there a case where active_tasks >0 but command_actors_count == 0?

[ntlm1686]

Yes, active_tasks may contain some actors which haven't been scheduled successfully in a placement group, even after the job has finished. I guess I have to use return here.

[d4l3k]

we should check the role max_retries/retry_policy here and throw an error if any of the workers exits more than N times

[ntlm1686]

Actually, here if we add a max_retries, it actually means how many node failures we can endure, so it's different from the max_retries in the ray actor's context.

[ntlm1686]

As long as the number of working nodes is bigger than min_nnodes, we should allow infinite node failures, unless nodes are failing too frequently. That can be a retry_policy.

[d4l3k]

would be good to add some "state machine" mock tests that runs through that driver loop step by step

[d4l3k]

can we split loop out from run so it has a "_step()" and "run()" method so we can manually step through the behavior from a mocked unit test?

[atinsood]

@ljjsalt @d4l3k

I have been thinking more about this and I am wondering if ray.util.queue is a better way of implementing this.

you basically create 2 actors, PlacementGroupManager actor and CommandManager actor and exchange information between them using ray queue.

PGManager group actor is responsible for creating the placement group and then putting a message in the queue for the CommandManager actor to process which can then create the command actor. this helps us keep the knowledge of pg creation and command actor creation compartmentalized in these actors.

additionally, we can @Remote these fns so both of these can be run in parallel

https://docs.ray.io/en/releases-1.2.0/advanced.html#message-passing-using-ray-queue

[ntlm1686]

@ljjsalt @d4l3k

I have been thinking more about this and I am wondering if ray.util.queue is a better way of implementing this.

you basically create 2 actors, PlacementGroupManager actor and CommandManager actor and exchange information between them using ray queue.

PGManager group actor is responsible for creating the placement group and then putting a message in the queue for the CommandManager actor to process which can then create the command actor. this helps us keep the knowledge of pg creation and command actor creation compartmentalized in these actors.

additionally, we can @Remote these fns so both of these can be run in parallel

https://docs.ray.io/en/releases-1.2.0/advanced.html#message-passing-using-ray-queue

@atinsood

The reason that we should use Placement Group is: if we only use ray actors here, even though we can rerun the failed actor after it throws a RayActorError. But this job could potentially lose the computation resource it used to have, then fail the job even the node is recovered.

For example, there are 3 nodes each with 1 cpu, job 1 requires minimum nodes of 3, and it’s running with all 3 nodes, another job 2 launched later and it requires 1 cpu, and becomes a pending task. Once a node failure happens, we have to rerun the actor which becomes a pending task after job 2, and job 2 will take that node, job 1 will fail since there aren't enough nodes to restart.

Since placement groups rescheduling have the highest priority, this situation won't happen.

I didn't get it why to use queue, since we didn't use any inter process communication here.

I understand your concern in using ray.wait here. So I wrapped the return value of command(remote) actors with RayResult class. If anything unexpected happen, it throws an error, we will know right away. But based on the logic that ray.wait supposed to work, I think it's right to use it here.

Besides, considering each command actor is a finite state machine that has four states(SCHEDULING, FAILING, RUNNING, TASK_COMPLETED).
For each command actor, it starts with SCHEDULING state, and actor.schedule.remote() becomes an asynchronous step function. When it's SCHEDULING, the next state is RUNNING. When it's RUNNING, the next state is TASK_COMPLETED or FAILING. When it's FAILING, the next state is SCHEDULING.

[atinsood]

The reason that we should use Placement Group is

yeah, not disagreeing on that. I am just thinking how do we manage the interaction between pg creation and command actor creation

@ray.remote
class PlacementGroupManager(object):
    def __init__(self, min_nodes, max_nodes, queue):
        self.queue = queue
        pass

    def run() -> PlacementGroup:
        # step 1. create a PG with min replicas

        #put the initial pg in the queue
        self.queue.put(ready)

        # step 2. while loop to keep create the rest of PGs incrementatlly,

        #go through the rest of the pgs one by one and keep adding them to the queue
        #we can deal with the logic of pg failure here or when to stop pg creation if needed


@ray.remote
class CommandActorManager(object):
    def __init__(self, queue, active_workers):
        self.queue = queue
        self.active_workers = []
        pass

    def run(self):
        #some logic on when to stop the while loop, either a poison pill or a signal actor
        while True:
            self.queue.get() #get notified that the pg was created
            # logic to create the command actor and goes here

            active_workers.append(command_actor.exec_module.remote()) # or may be use another queue, pretty sure there is a better way to deal with this

            #once you are done reading all the PGs from the queue or if the queue 





def main() -> None:  # pragma: no cover
    actors: List[RayActor] = load_actor_json("actors.json")
    # pyre-fixme[16]: Module `worker` has no attribute `init`.
    ray.init(address="auto", namespace="torchx-ray")
    q = Queue() # we can set the set of the queue upfront
    pg_manager = PlacementGroupManager(2,4,q) #min, max and queue
    pg_manager.run.remote()
    cmd_actor_manager = CommandActorManager(q, active_workers)
    cmd_actor_manager.run.remote()
    # Await return result of remote ray function
    while len(active_workers) > 0:
        _logger.info(f"running ray.wait on {active_workers}")

        # pyre-fixme[16]: Module `worker` has no attribute `wait`.
        completed_workers, active_workers = ray.wait(active_workers)
        # If a failure occurs the ObjectRef will be marked as completed.
        # Calling ray.get will expose the failure as a RayActorError.
        for object_ref in completed_workers:
            ray.get(object_ref)

I was thinking something like this, this is not the correct code, but more of a thought process on how to do interaction between the logic that is creating placement groups and the logic that is creating command actors and how to keep them compartmentalized

[ntlm1686]

@ljjsalt @d4l3k

I have been thinking more about this and I am wondering if ray.util.queue is a better way of implementing this.

you basically create 2 actors, PlacementGroupManager actor and CommandManager actor and exchange information between them using ray queue.

PGManager group actor is responsible for creating the placement group and then putting a message in the queue for the CommandManager actor to process which can then create the command actor. this helps us keep the knowledge of pg creation and command actor creation compartmentalized in these actors.

additionally, we can @Remote these fns so both of these can be run in parallel

https://docs.ray.io/en/releases-1.2.0/advanced.html#message-passing-using-ray-queue

        # step 2. while loop to keep create the rest of PGs incrementatlly,

        #go through the rest of the pgs one by one and keep adding them to the queue
        #we can deal with the logic of pg failure here or when to stop pg creation if needed

Actually, this step is not necessary. First, increasing PG one by one can be a problem once the number of nodes is large.

Please read the description of the X axis, it only makes the time that a PG creation event added to GCS pending queue longer.

The second change: instead of creating the command actor after a placement group has been scheduled, we create all the command actors at the beginning too(with SCHEDULING state).

[codecov]

Codecov Report

Merging #572 (3228b9c) into main (b051e3f) will increase coverage by 0.15%.
The diff coverage is 98.41%.

@@            Coverage Diff             @@
##             main     #572      +/-   ##
==========================================
+ Coverage   94.85%   95.00%   +0.15%     
==========================================
  Files          66       66              
  Lines        4042     4144     +102     
==========================================
+ Hits         3834     3937     +103     
+ Misses        208      207       -1     

Impacted Files	Coverage Δ
torchx/schedulers/ray_scheduler.py	95.26% <ø> (ø)
torchx/schedulers/ray/ray_driver.py	97.10% <98.07%> (+1.26%)	⬆️
torchx/components/dist.py	96.42% <100.00%> (+7.06%)	⬆️
torchx/schedulers/ray/ray_common.py	100.00% <100.00%> (ø)
torchx/specs/api.py	98.40% <100.00%> (+<0.01%)	⬆️
torchx/runner/api.py	96.85% <0.00%> (+0.01%)	⬆️

📣 Codecov can now indicate which changes are the most critical in Pull Requests. Learn more

[d4l3k]

If we did want to separate the concerns here with scheduling vs retries we could use the builtin task max_retries config though that has some other implications.

I'm not sure that using a queue with two actors would simplify this logic -- state machines that you can step through are nice from a testing perspective

[ntlm1686]

If we did want to separate the concerns here with scheduling vs retries we could use the builtin task max_retries config though that has some other implications.

I'm not sure that using a queue with two actors would simplify this logic -- state machines that you can step through are nice from a testing perspective

But we are just running schedule(which cannot go wrong) and exec_module functions(which cannot go wrong unless script fails), it it necessary to use max_retries here? Of course, the feature can be easily added.