https://outerbounds.com/ logo
Join Slack
Powered by
Hi team, I'd be interested in contributing to ope...
# ask-metaflow
a
Hi team, I'd be interested in contributing to open source to get Ray running on top of AWS Batch multi-node parallel jobs through Metaflow. I asked the Ray team, and they indicated that it is likely feasible to run Ray on AWS Batch. In my mind, I'm thinking the user could use the 
@parallel
decorator to initiate the gang-scheduled cluster, and then maybe do all the Ray configuration through the AWS Batch job definition? Would love a second perspective from Outerbounds on this. These are the docs for installing Ray on on-premise clusters: https://docs.ray.io/en/latest/cluster/vms/user-guides/launching-clusters/on-premises.html#manual-cluster-launcher
a
@acoustic-van-30942 a quick way to verify this would be to confirm if it is indeed feasible to run Ray on multi-node AWS Batch outside of Metaflow. Once there is a path, we can work out how to stitch it together with Metaflow. One question though - what would be the benefit of running Ray on AWS Batch?
a
Thanks @ancient-application-36103. I think the big thing is that it would allow users to adopt a module called "Ray Train" which is just a framework that more or less competes with Tensorflow, MPI, etc. For users who are familiar and comfortable with Ray Train, which abstracts away a lot of the boilerplate that comes with other distributed training libraries, that would be very valuable. But Ray Train is only compatible with "Ray Clusters", so essentially, if there's a way to transform the AWS Batch ephemeral clusters into Ray clusters, that would be a HUGE victory. In a nutshell, the benefit would be that our current managed training infra uses Metaflow + AWS Batch. The shortest path to getting Ray integrated would be to have it running on Batch. Otherwise, we will have to go down the K8s route, which will require much more engineering effort.
a
@acoustic-van-30942 were you able to make any progress? also happy to give you a deep dive on how the current 
@parallel
implementation works if that's helpful. let me know!
a
Hey @ancient-application-36103. I'm actually very familiar with the 
@parallel
decorator. It's been fantastic. Our team has been using it for a bunch of POCs and benchmarking results. I'm now trying to pair it with Ray
👍🏼 1
Hey @ancient-application-36103 Sorry for the dumb question, but if I create my own Ray decorator that looks something like this:
Copy code
import inspect
import subprocess
import pickle
import tempfile
import os
import sys
from metaflow import current
from metaflow.plugins.parallel_decorator import ParallelDecorator

class RayParallelDecorator(ParallelDecorator):
    name = "ray_parallel"
    defaults = {"master_port": None}
    IS_PARALLEL = True

    def task_decorate(
        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
    ):
        return super().task_decorate(
            step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
        )

    def setup_distributed_env(self, flow):
        setup_ray_distributed(self.attributes["master_port"])
        
def setup_ray_distributed(master_port=None):
    """
    Manually set up Ray cluster
    """
    # Choose port depending on run id to reduce probability of collisions, unless
    # provided by the user.
    try:
        master_port = master_port or (9001 + abs(int(current.run_id)) % 1000)
    except:
        # if `int()` fails, i.e. `run_id` is not an `int`, use just a constant port. Can't use `hash()`,
        # as that is not constant.
        master_port = 9001
    
    if current.parallel.node_index == 0:
        subprocess.run([sys.executable, "-m", "ray", "start", "--head", f"--port={master_port}"])
    else:
        address = f"{current.parallel.main_ip}:{master_port}"
        subprocess.run([sys.executable, "-m", "ray", "start", "--address", address])
How do I call it?
a
are you creating this decorator in a fork or within metaflow extensions?
if in a fork, just register it here and you can start using 
@ray_parallel
on top of the step
a
Ah OK I'll try it from a fork then. Thanks Savin!
I think I was able to get the Ray cluster set-up with this 
@ray_parallel
decorator (see here) But it's not quite clear how how to submit the Ray training job from Metaflow. Unfortunately, the workers die when I submit the job. I have a simple reproducible example here Node logs below
stderr_RayFlow_486_train_2214-0.txtstderr_RayFlow_486_train_2213-0 (1).txt
a
@acoustic-van-30942 following up here - is this the latest state that I should look at or do you have any more updates on this?
a
Thanks for your assistance Savin! One update I had was that I did a print statement of 
ray.cluster_resources()
And unfortunately, it only shows a cluster with 1 node, which is probably why it's failing.
Copy code
{'node:__internal_head__': 1.0, 'memory': 30534839706.0, 'GPU': 1.0, 'object_store_memory': 10000000000.0, 'node:10.14.52.26': 1.0, 'accelerator_type:A10G': 1.0, 'CPU': 16.0}
a
Do you happen to know how Ray does cluster discovery under the hood?
a
We are installing it on-premise (via the 
@ray_parallel
that I created). Essentially, you choose a node to be the head node, then for each of the worker nodes you provide the IP address of the head node. That, supposedly, should have created the Ray cluster. https://docs.ray.io/en/latest/cluster/vms/user-guides/launching-clusters/on-premises.htmlhttps://docs.ray.io/en/latest/cluster/vms/user-guides/launching-clusters/on-premises.html
Hi @hallowed-glass-14538, @ancient-application-36103 For the 
@parallel
decorator, is there a way to check the status of the control node from the worker nodes using either Metaflow client API or some other method? For example, I want to keep the worker nodes alive (infinite while loop), but if the control node is finished, then I want to break out of that while loop for the control nodes.
Asking because the Ray training job will work if I can keep the worker nodes alive, but now I'm stuck in an infinite loop where the control node task is finished, but the worker nodes don't seem to be completing their task(s)
Example:
Copy code
@parallel
    @step
    def train(self):
        if current.parallel.node_index == 0:  
            import ray
            import subprocess
            import time
            p = subprocess.Popen('ray start --head --port=6379',  
                         shell=True).wait()
            ray.init()
            result = subprocess.run(["python", "train.py", "--num_workers", str(self.num_parallel)], capture_output=True, text=True, check=True)
            print(result.stdout)
            ray.shutdown() 
        else:
            import ray
            import time
            import subprocess
            p = subprocess.Popen(f"ray start --address='{current.parallel.main_ip}:6379'", shell=True).wait()
            ray.init()
            print(ray.cluster_resources())       
            print(ray.nodes())
            try:
                while ray.is_initialized():
                    time.sleep(10)
            except:
                pass

        self.next(self.multinode_end)
11 Views
Open in Slack
PreviousNext
Powered by