Adding ray as a distributor

[Vaishaal]

For running large jobs on AWS I found a couple problems with the spark backend.

1. Spark runs things stage-wise by default so either you need to have a huge subjob_size and wait for the reader to finish before the job starts or you run stage by stage but then you are vulnerable to stragglers in each stage. In my experience I saw the the stragglers slowed down the pipeline by 2-10x depending how unlucky you got with a shard
2. when you are running on a non dedicated cluster spark doesn't handle autoscaling well so you basically need to pay for num_cpus x total runtime, whereas natively ray auto scales number of cpu nodes based on queue size and download speed. This is also nice when you get an error (as you often do for these large DL jobs), as you didn't waste money keeping cluster alive.
3. As of right now the ray aws utilities are much nicer to work with than the native spark ones.

This PR just adds a ray distributor (with same interface as spark and multiprocessing distributors), and adds an example launch script, and a cluster_config.yaml file if people want to spin up their own AWS cluster.

Using this I was able to get over 200k/images second on a cluster of 100 m5.24xlarges consistently for 24 hours.

[Vaishaal]

@rom1504 can you take a look and merge this :)

[rom1504]

Looks pretty good. Any way to add a test for this distributor ?
It would guarantee this would keep working with future changes

[jelech]

Hi, is there any plan to merge this PR?

[rom1504]

@jelech I want to but I would prefer we add some testing as otherwise it'll eventually get broken over time

[Vaishaal]

Let me do it!

…

________________________________ From: Romain Beaumont ***@***.***> Sent: Saturday, July 15, 2023 2:41:14 PM To: rom1504/img2dataset ***@***.***> Cc: Vaishaal ***@***.***>; Author ***@***.***> Subject: Re: [rom1504/img2dataset] Adding ray as a distributor (PR #272) @jelech<https://github.com/jelech> I want to but I would prefer we add some testing as otherwise it'll eventually get broken over time — Reply to this email directly, view it on GitHub<#272 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AAKMJEEKH4TQTX72XPRCWP3XQLP4VANCNFSM6AAAAAAT5LYA3Y>. You are receiving this because you authored the thread.Message ID: ***@***.***>

[rom1504]

I think we need this https://docs.ray.io/en/latest/cluster/vms/user-guides/launching-clusters/on-premises.html

eg in gh action run

pip install ray[default]
ray start --head --port=6379
ray start --address=127.0.0.1:6379

[rom1504]

then adapt https://github.com/rom1504/img2dataset/blob/main/tests/test_main.py#L363

plus either automatically start a ray cluster if there is none in the main code

or do it only in tests

[rom1504]

ok so I added tests

I'll rework the doc here a little bit and then we're good for merging

[Vaishaal]

Nice thanks so much!

[rom1504]

thanks for the PR!