[core] Add MPI support on Ray cluster #40917

fishbone · 2023-11-02T22:03:52Z

Why are these changes needed?

This PR adds the support to run MPI based code on top of Ray.

The support is done with runtime env plugin. To enable it, the following decorator needs to be added inside ray remote options:

@ray.remote(
        runtime_env={
            "mpi": {
                "args": ["-n", "4"],
                "worker_entry": "mpi_worker.run",
            }
        }
)
def f():
    pass

Here the mpi_worker.run is the function the process with rank > 0 will run. It'll run as import mpi_worker; mpi_worker.run(). The parameter needs to be passed with MPI comm.bcast.

Here the process with rank 0 sill will run the remote function f.

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

rkooo567

Generally lgtm. Some comments in input validation & API doc

python/ray/_private/runtime_env/mpi.py

rkooo567 · 2023-11-03T00:53:40Z

python/ray/_private/runtime_env/mpi.py

+                ["mpirun", "--version"], capture_output=True, check=True
+            )
+        except subprocess.CalledProcessError:
+            logger.error(


Maybe logger.exception to print stacktrace?

rkooo567 · 2023-11-03T00:54:18Z

python/ray/_private/runtime_env/mpi.py

+        except subprocess.CalledProcessError:
+            logger.error(
+                "Failed to run mpi run. Please make sure mpi has been installed"
+            )


should we kill proc here? Or does it guarantee the proc is killed? (if so can you comment here?)

Modify context is in runtime env agent I think. Exception should be good? I can test it.

rkooo567 · 2023-11-03T00:55:33Z

python/ray/_private/runtime_env/mpi.py

+        from pathlib import Path
+
+        # mpirun -n 10 python mpi.py worker_entry_func
+        worker_entry = mpi_config["worker_entry"]


maybe we need to handle a case where "worker_entry" doesn't exist because we don't have input validation iiuc

worker_Entry = mpi_config.get("worker_entry") if worker_entry is None: raise

rkooo567 · 2023-11-03T00:55:45Z

python/ray/_private/runtime_env/mpi.py

+
+        # mpirun -n 10 python mpi.py worker_entry_func
+        worker_entry = mpi_config["worker_entry"]
+        assert Path(worker_entry).is_file()


Add a error message?

assert Path(worker_entry).is_file(), "worker_entry is not a file but ..."

rkooo567 · 2023-11-03T00:57:10Z

python/ray/_private/runtime_env/mpi.py

+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Setup MPI worker")


is this intentional?

Do you mean the main function? or the parser?

This will will be used as the mpi entry point.

main function. I don't see any main function from other plugins though. Maybe it should be a part of mpi_worker.py not here? (this means the function is executed when you import it?)

MPIRUN is like a fork and the rest plugin doesn't have this.
The function won't execute when import since it checks main. If you import, __name__ won't be __main__.

This piece of code is part of the plugin, that's why I put it here and it's simple.

But open to move if you insist.

Hmm I see. I think this is a bit confusing when I read code first time. I prefer to move it to a separate file (something like mpi_start.py), but it is also okay if you add comments in details in the main block. E.g., "the plugin starts a subprocess that runs this main method. It is not executed as a part of normal plugin" or sth like that

I thought about moving it to python/ray/_private/workers/mpi_workers.py but feel it just moves the code to far away from the place where it's used. I think split it and move it to the other file is better.

rkooo567 · 2023-11-03T00:57:48Z

python/ray/runtime_env/runtime_env.py

@@ -287,6 +288,7 @@ def __init__(
        nsight: Optional[Union[str, Dict[str, str]]] = None,
        config: Optional[Union[Dict, RuntimeEnvConfig]] = None,
        _validate: bool = True,
+        mpi: Optional[Dict] = None,


We should update a doc, or consider _mpi!

we'll update the doc. i think it's a feature.

rkooo567 · 2023-11-03T00:58:10Z

python/ray/tests/mpi/test_mpi.py

+        runtime_env={
+            "mpi": {
+                "args": ["-n", "4"],
+                "worker_entry": "mpi_worker.py",


Where does it find the file? From the current directory?

I think we should be clear about it in the docstring

I think it's from the working dir. I'll add the doc.

python/ray/widgets/util.py

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

ci/docker/core.build.Dockerfile

rkooo567

Hmm having main inside runtime env plugin file seems wrong?

rkooo567 · 2023-11-06T09:06:14Z

python/ray/_private/runtime_env/mpi.py

+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Setup MPI worker")


main function. I don't see any main function from other plugins though. Maybe it should be a part of mpi_worker.py not here? (this means the function is executed when you import it?)

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

rkooo567

I think code lgtm now. Can you ping me after updating the docstring for mpi API? I think we also need api approval as it is a new plugin

rkooo567 · 2023-11-07T11:58:50Z

python/ray/_private/runtime_env/mpi.py

+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Setup MPI worker")


Hmm I see. I think this is a bit confusing when I read code first time. I prefer to move it to a separate file (something like mpi_start.py), but it is also okay if you add comments in details in the main block. E.g., "the plugin starts a subprocess that runs this main method. It is not executed as a part of normal plugin" or sth like that

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone · 2023-11-08T02:39:54Z

@rkooo567 I'll create another PR for the doc.

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

rkooo567 · 2023-11-08T11:06:25Z

ci/env/install-dependencies.sh

@@ -160,6 +160,9 @@ install_miniconda() {
    )
  fi

+  # Install mpi4py
+  "${WORKSPACE_DIR}"/ci/suppress_output conda install -c anaconda mpi4py -y


Hmm it is a bit weird we have it here? (it will be requested by every dev to download mpi4py)

Why don't we just make it test requirement?

The problem is that only conda can install it. It has system deps and doesn't work with py311 :(

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone · 2023-11-09T00:05:09Z

@rkooo567 I updated the API to avoid the extra worker file.

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

This PR adds the support to run MPI based code on top of Ray. The support is done with runtime env plugin. To enable it, the following decorator needs to be added inside ray remote options: @ray.remote( runtime_env={ "mpi": { "args": ["-n", "4"], "worker_entry": "mpi_worker.run", } } ) def f(): pass Here the mpi_worker.run is the function the process with rank > 0 will run. It'll run as import mpi_worker; mpi_worker.run(). The parameter needs to be passed with MPI comm.bcast. Here the process with rank 0 sill will run the remote function f. Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone added 7 commits November 2, 2023 19:38

draft

0da2538

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

40e2346

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

356fda7

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

c6806ad

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

3978df9

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

add mpi test

51d523d

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

14d03ae

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone assigned scv119 and rkooo567 Nov 2, 2023

rkooo567 reviewed Nov 3, 2023

View reviewed changes

rkooo567 added the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Nov 3, 2023

fix

0d92333

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone requested review from aslonnie and can-anyscale as code owners November 3, 2023 19:35

aslonnie reviewed Nov 3, 2023

View reviewed changes

ci/docker/core.build.Dockerfile Outdated Show resolved Hide resolved

rkooo567 reviewed Nov 6, 2023

View reviewed changes

fishbone added 3 commits November 6, 2023 23:54

fix build

789a4e2

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

43ae02a

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

update gpu docker

e471637

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

rkooo567 reviewed Nov 7, 2023

View reviewed changes

fishbone added 10 commits November 7, 2023 17:51

fix

7727229

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

format

183dfbb

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

032a5fb

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

2bdf42e

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

7005fa3

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

add

c9611da

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

97ccd24

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

a6831d1

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

02afd02

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

split

3852f9c

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone removed the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Nov 8, 2023

Merge branch 'master' into mpi-ray

271e49f

fishbone added 4 commits November 8, 2023 08:16

rev

f780c05

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

a784b57

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

5002815

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

a492234

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

rkooo567 approved these changes Nov 8, 2023

View reviewed changes

fishbone added 2 commits November 8, 2023 23:57

fix

2807165

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fix

7268695

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

fishbone and others added 4 commits November 9, 2023 00:07

remove dbg

8844056

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

Merge remote-tracking branch 'upstream/master' into mpi-ray

f4abd66

Signed-off-by: Yi Cheng <74173148+iycheng@users.noreply.github.com>

Merge branch 'master' into mpi-ray

9044563

Merge branch 'master' into mpi-ray

22eabbd

can-anyscale merged commit 99b1a2c into ray-project:master Nov 9, 2023
23 of 30 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[core] Add MPI support on Ray cluster #40917

[core] Add MPI support on Ray cluster #40917

fishbone commented Nov 2, 2023 •

edited

rkooo567 left a comment

rkooo567 Nov 3, 2023

rkooo567 Nov 3, 2023

fishbone Nov 3, 2023 •

edited by scv119

rkooo567 Nov 3, 2023

rkooo567 Nov 3, 2023

rkooo567 Nov 3, 2023

rkooo567 Nov 3, 2023

fishbone Nov 3, 2023

fishbone Nov 3, 2023

rkooo567 Nov 6, 2023

fishbone Nov 6, 2023

rkooo567 Nov 7, 2023

fishbone Nov 8, 2023

rkooo567 Nov 3, 2023

fishbone Nov 3, 2023

rkooo567 Nov 3, 2023

rkooo567 Nov 3, 2023

fishbone Nov 3, 2023

rkooo567 left a comment

rkooo567 Nov 6, 2023

rkooo567 left a comment •

edited

rkooo567 Nov 7, 2023

fishbone commented Nov 8, 2023

rkooo567 Nov 8, 2023

fishbone Nov 8, 2023

fishbone commented Nov 9, 2023



		if __name__ == "__main__":
		parser = argparse.ArgumentParser(description="Setup MPI worker")

[core] Add MPI support on Ray cluster #40917

[core] Add MPI support on Ray cluster #40917

Conversation

fishbone commented Nov 2, 2023 • edited

Why are these changes needed?

Related issue number

Checks

rkooo567 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fishbone Nov 3, 2023 • edited by scv119

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fishbone commented Nov 8, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fishbone commented Nov 9, 2023

fishbone commented Nov 2, 2023 •

edited

fishbone Nov 3, 2023 •

edited by scv119

rkooo567 left a comment •

edited