Distributed training with Ray

[dragosmc]

Hi Ludwig team,

I'm trying to setup Ray to work with Ludwig and I'm facing some config issues. The docs here only mention how to train from the CLI https://ludwig.ai/latest/user_guide/distributed_training/, but our use case is training from Python code.

The setup I'm trying this on is:

• a Ray cluster on AWS
• Ludwig with programatic call to model.train()
• the dataset is passed as pandas dataframe to train

I'm wondering what are Ludiwg's assumptions about distributed training or, in other words, how does Ludwig handle the setup on Ray:

• does it expect to have the code on the worker nodes at a specific location?
• any other assumption and perhaps any "escape hatches" that allows us to configure even more Ray?

An example of such config would help a lot as well.

Last but not least, I see that the Ludwig-provided docker images for Ray use Python 3.8 w/ Ray 2.2.0. Would be great to build various images with Python 3.9 and 3.10 and Ray 2.3.0 (I can help with this with a PR). Additionally, whilst on the same topic of "FAQ" and "guides" I can collate the discussion here into a richer guide on how to use Ray with Ludwig without much friction.

Dragos.

[tgaddair]

Hey @dragosmc, good questions:

• The main assumptions are that (1) all the workers have identical Ludwig versions / Ludwig dependencies installed, and (2) the dataset can be read from all of the workers in the cluster (preferably: from an object storage location like S3). User code doesn't need to live on all the workers unless you have custom encoders etc. that Ludwig needs to use.
• Ray training can be configured pretty extensively via the backend section of the config (see docs here). We also have a section in the docs here about distributed training that covers various gotchas. Let me know if you additional questions after looking at the docs, happy to put together a more complete example if needed.
• We're working on upgrading our images to Ray 2.3 now. This will also bump to PyTorch 2.0 and CUDA 11.8. We can additionally bump Python to 3.9 or 3.10. Maybe 3.9 just to play it safe for users if that's fine for you.

[tgaddair]

And would definitely love any help on improving docs as you work through some of this. Thanks for offering to help with putting together richer guides and FAQ!

[dragosmc]

Thank you for the bulletpoints.
Regarding the first one:

• (1) I have a setup_commands in my Ray configuration which installs Ludwig and al other dependencies that are required.
• (2) since I'm calling model.train() is there an extra step required to ensure that the dataset is available to the workers? I (wrongly) assumed that the chunks are being transferred over the wire from the head to the workers
• (3) I do have a 3rd point here regarding custom resources for Ludwig. We do indeed have a custom F1 metric for our classifier which gets registered with Ludwig as part of a Callback. We're now trying to create a custom encoder of a sequence of text sequences which we added to the same callback but it's hard to ensure the order of adding and registering new components to Ludwig is the correct one (input definitions, encoder, etc.) - (I'm went a bit offtopic with this point but the aim is to see if the way we're doing things is what you had in mind when architecting the plugin/extensibility functionality for Ludwig)

On the second one:

• One gotcha that I'd have liked to know is "how to tell Ludwig that I have a cluster at this IP address". The setup we're trying get working with Ray is a notebook running in Sagemaker Studio with a Ray cluster sitting in the same VPC which Ludwig is accessing. Not sure if this is a common pattern, but having ray installed without setting the RAY_ADDREESS environment variable would only start a new cluster on the notebook instance. If you think this is a scenario other teams might find themselves in the we can perhaps update the docs to point this (in hindsight) obvious setting - or even add such a setting to the ray backend config as a property
• On the same note of documentation I'm thinking that an example notebook, or python file, that exemplifies the assumptions and that "just works" once someone has a connectable cluster of Ray would be great
• Last but not least I was caught by the fact that my version of Python that I'm using to run the model was different than the one inside the ray's default docker image and ludwig's default docker image - it'd also be a nice thing to have a note on why one should use the ludwig image provided vs. the ray recommended one

On the third one:

• I don't really mind, it's just that I wasn't expecting the images in the example (or the main branch) to use a Python 3.8. Maybe this is because I'm used with Sagemaker and I always pick the latest image which comes with Python 3.10. Having a matrix of images for various combinations will definitely help

Once more, thank you for your time and for answering my questions.

Note: My background is Java and Node so some of what I've written above might sound dumb from a pythonic point of view :D

[dragosmc]

Update on the entire process: I have followed the tutorial in the docs and wasn't able to get it to work.

Here's the setup:

• a Ray cluster with the latest ludwig-ray image ludwigai/ludwig-ray:latest
• a simple Ludwig config that expects 3 columns in a dataframe as features and a label column
• a fake dataset available from S3 at s3://dragos-test/data.csv (https://dragos-test.s3.eu-west-2.amazonaws.com/data.csv)
  The entire setup can be found here https://github.com/dragosmc/ludwig-ray

The commands I've ran are:

• first setup the cluster

ray up cluster.yaml

• try running ludwig on the cluster

ray submit cluster.yaml --verbose \
    ludwig train --config config.yaml --dataset s3://dragos-test/data.csv

On the last command I get the following error:

(wp38) ➜  test-ray git:(main) ✗ ray submit cluster.yaml --verbose \
    ludwig train --config config.yaml --dataset s3://dragos-test/data.csv
2023-04-04 14:29:00,215 INFO util.py:372 -- setting max workers for head node type to 0
Loaded cached provider configuration from /var/folders/76/2dtcyszd1p3g1yzmvtfsq_zh0000gn/T/ray-config-44e79340569cb51c7caef67b6e7893999dcf3b54
If you experience issues with the cloud provider, try re-running the command with --no-config-cache.
Creating AWS resource `ec2` in `eu-west-2`
Creating AWS resource `ec2` in `eu-west-2`
Fetched IP: x.x.x.x
Running `mkdir -p /tmp/ray_tmp_mount/koble-cluster/~ && chown -R ubuntu /tmp/ray_tmp_mount/koble-cluster/~`
Warning: Permanently added 'x.x.x.x' (ED25519) to the list of known hosts.
Shared connection to x.x.x.x closed.
Running `rsync --rsh ssh -i /Users/dragos/.ssh/ray-autoscaler_eu-west-2.pem -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o ExitOnForwardFailure=yes -o ServerAliveInterval=5 -o ServerAliveCountMax=3 -o ControlMaster=auto -o ControlPath=/tmp/ray_ssh_d18928faba/36c97cbe08/%C -o ControlPersist=10s -o ConnectTimeout=120s -avz ludwig ubuntu@x.x.x.x:/tmp/ray_tmp_mount/koble-cluster/~/ludwig`
building file list ... rsync: link_stat "/Users/dragos/work/test-ray/ludwig" failed: No such file or directory (2)
done

sent 29 bytes  received 20 bytes  98.00 bytes/sec
total size is 0  speedup is 0.00
rsync error: some files could not be transferred (code 23) at /AppleInternal/Library/BuildRoots/97f6331a-ba75-11ed-a4bc-863efbbaf80d/Library/Caches/com.apple.xbs/Sources/rsync/rsync/main.c(996) [sender=2.6.9]
Error: SSH command failed.
(wp38) ➜  test-ray git:(main) ✗ 

From the error I see that Ray expects a folder/file named ludwig inside the working directory and this is where I kinda got stuck... Anything obviously wrong here?

[tgaddair]

Sorry for the late response here, let me take a stab at these questions:

since I'm calling model.train() is there an extra step required to ensure that the dataset is available to the workers? I (wrongly) assumed that the chunks are being transferred over the wire from the head to the workers

If you're providing the dataset as a Pandas DF, then it will be sent over the wire to the workers. No worries there. The only issue arises if the dataset is a file. Also, you likely want to configure a remote cache for the intermediate (preprocessed) data so it is available to the workers. See this section of the config.

I do have a 3rd point here regarding custom resources for Ludwig. We do indeed have a custom F1 metric for our classifier which gets registered with Ludwig as part of a Callback.

Would be great to have a separate discussion on this. We can definitely add F1 as a metric pretty easily, but would love feedback on the custom Metric API, as I wouldn't be surprised if there were some rough edges from user experience point of view.

One gotcha that I'd have liked to know is "how to tell Ludwig that I have a cluster at this IP address". The setup we're trying get working with Ray is a notebook running in Sagemaker Studio with a Ray cluster sitting in the same VPC which Ludwig is accessing. Not sure if this is a common pattern, but having ray installed without setting the RAY_ADDREESS environment variable would only start a new cluster on the notebook instance. If you think this is a scenario other teams might find themselves in the we can perhaps update the docs to point this (in hindsight) obvious setting - or even add such a setting to the ray backend config as a property

We can definitely add the remote address to the backend config. That would be a very reasonable change.

Last but not least I was caught by the fact that my version of Python that I'm using to run the model was different than the one inside the ray's default docker image and ludwig's default docker image - it'd also be a nice thing to have a note on why one should use the ludwig image provided vs. the ray recommended one

Good callout, we definitely could use better examples with Ray, and explain the differences between the images better.

Update on the entire process: I have followed the tutorial in the docs and wasn't able to get it to work.

Ah, this may be something that changed in newer versions of the Ray CLI. It should probably be ray exec ... instead of ray submit ..., as you want to execute a command on the remote machine, rather than upload a Python script to run. Can you give that a shot and let us know how it goes? Should hopefully be a quick doc update if so.

Going to loop in some folks from our side to take a look at this thread and help with addressing some of these issues.

cc @justinxzhao @arnavgarg1 @connor-mccorm @ksbrar @abidwael