Add doc and tests for mpi (#162)

* add doc and UTs for mpi

* fixes

* update

* fixes lint

* fixes

* fixes

* fixes import

Author: ConeyLiu

.github/workflows/raydp.yml

@@ -46,7 +46,12 @@ jobs:
       - name: Install extra dependencies for macOS
         if: matrix.os == 'macos-latest'
         run: |
-          brew install libuv libomp
+          brew install libuv libomp mpich
+      - name: Install extra dependencies for Ubuntu
+        if: matrix.os == 'ubuntu-latest'
+        run: |
+          sudo apt-get install -y mpich
+
       - name: Cache pip - Ubuntu
         if: matrix.os == 'ubuntu-latest'
         uses: actions/cache@v2

README.md

@@ -1,28 +1,18 @@
 # RayDP
 
-RayDP is a distributed data processing library that provides simple APIs for running Spark on [Ray](https://github.com/ray-project/ray) and integrating Spark with distributed deep learning and machine learning frameworks. RayDP makes it simple to build distributed end-to-end data analytics and AI pipeline. Instead of using lots of glue code or an orchestration framework to stitch multiple distributed programs, RayDP allows you to write Spark, PyTorch, Tensorflow, XGBoost code in a single python program with increased productivity and performance. You can build an end-to-end pipeline on a single Ray cluster by using Spark for data preprocessing, RaySGD or Horovod for distributed deep learning, RayTune for hyperparameter tuning and RayServe for model serving.
-
-### Spark on Ray
-
-RayDP provides an API for starting a Spark job on Ray in your python program without a need to setup a Spark cluster manually. RayDP supports Ray as a Spark resource manger and runs Spark executors in Ray actors. RayDP utilizes Ray's in-memory object store to efficiently exchange data between Spark and other Ray libraries. You can use Spark to read the input data, process the data using SQL, Spark DataFrame, or Pandas (via [Koalas](https://github.com/databricks/koalas)) API, extract and transform features using Spark MLLib, and feed the output to deep learning and machine learning frameworks.
-
-### Integrating Spark with Deep Learning and Machine Learning Frameworks
-
-#### MLDataset API
-RayDP provides an API for creating a Ray MLDataset from a Spark dataframe. MLDataset represents a distributed dataset stored in Ray's in-memory object store. It supports transformation on each shard and can be converted to a PyTorch or Tensorflow dataset for distributed training. If you prefer to using Horovod on Ray or RaySGD for distributed training, you can use MLDataset to seamlessly integrate Spark with them.
-
-#### Estimator API
-RayDP also provides high level scikit-learn style Estimator APIs for distributed training. The Estimator APIs allow you to train a deep neural network directly on a Spark DataFrame, leveraging Ray’s ability to scale out across the cluster. The Estimator APIs are wrappers of RaySGD and hide the complexity of converting a Spark DataFrame to a PyTorch/Tensorflow dataset and distributing the training.
+RayDP is a distributed data processing library that provides simple APIs for running Spark/MPI on [Ray](https://github.com/ray-project/ray) and integrating Spark with distributed deep learning and machine learning frameworks. RayDP makes it simple to build distributed end-to-end data analytics and AI pipeline. Instead of using lots of glue code or an orchestration framework to stitch multiple distributed programs, RayDP allows you to write Spark, PyTorch, Tensorflow, XGBoost code in a single python program with increased productivity and performance. You can build an end-to-end pipeline on a single Ray cluster by using Spark for data preprocessing, RaySGD or Horovod for distributed deep learning, RayTune for hyperparameter tuning and RayServe for model serving.
 
 ## Installation
 
 
 You can install latest RayDP using pip. RayDP requires Ray (>=1.3.0) and PySpark (>=3.0.0). Please also make sure java is installed and JAVA_HOME is set properly.
+
 ```shell
 pip install raydp
 ```
 
 Or you can install our nightly build:
+
 ```shell
 pip install raydp-nightly
 ```
@@ -34,38 +24,75 @@ If you'd like to build and install the latest master, use the following command:
 pip install dist/raydp*.whl
 ```
 
-## Getting Started
-To start a Spark job on Ray, you can use the `raydp.init_spark` API. You can write Spark, PyTorch/Tensorflow, Ray code in the same python program to easily implement an end-to-end pipeline.
+## Spark on Ray
+
+RayDP provides an API for starting a Spark job on Ray in your python program without a need to setup a Spark cluster manually. RayDP supports Ray as a Spark resource manger and runs Spark executors in Ray actors. RayDP utilizes Ray's in-memory object store to efficiently exchange data between Spark and other Ray libraries. You can use Spark to read the input data, process the data using SQL, Spark DataFrame, or Pandas (via [Koalas](https://github.com/databricks/koalas)) API, extract and transform features using Spark MLLib, and feed the output to deep learning and machine learning frameworks.
 
 ### Classic Spark Word Count Example
-After we use RayDP to initialize a Spark cluster, we can use Spark as usual. 
+
+To start a Spark job on Ray, you can use the `raydp.init_spark` API. After we use RayDP to initialize a Spark cluster, we can use Spark as usual. 
+
 ```python
 import ray
 import raydp
 
+# connect to ray cluster
 ray.init(address='auto')
 
+# create a Spark cluster with specified resource requirements
 spark = raydp.init_spark('word_count',
                          num_executors=2,
                          executor_cores=2,
                          executor_memory='1G')
 
+# normal data processesing with Spark
 df = spark.createDataFrame([('look',), ('spark',), ('tutorial',), ('spark',), ('look', ), ('python', )], ['word'])
 df.show()
 word_count = df.groupBy('word').count()
 word_count.show()
 
+# stop the spark cluster
 raydp.stop_spark()
 ```
 
-### Integration with PyTorch
-However, combined with other ray components, such as RaySGD and RayServe, we can easily build an end-to-end deep learning pipeline. In this example. we show how to use our estimator API, which is a wrapper around RaySGD, to perform data preprocessing using Spark, and train a model using PyTorch.
+### Dynamic Resource Allocation
+
+RayDP now supports External Shuffle Serivce. To enable it, you can either set `spark.shuffle.service.enabled` to `true` in `spark-defaults.conf`, or you can provide a config to `raydp.init_spark`, as shown below:
+
+```python
+raydp.init_spark(..., configs={"spark.shuffle.service.enabled": "true"})
+```
+
+The user-provided config will overwrite those specified in `spark-defaults.conf`. By default Spark will load `spark-defaults.conf` from `$SPARK_HOME/conf`, you can also modify this location by setting `SPARK_CONF_DIR`.
+
+Similarly, you can also enable Dynamic Executor Allocation this way. However, because Ray does not support object ownership tranferring now(1.3.0), you must use Dynamic Executor Allocation with data persistence. You can write the data frame in spark to HDFS as a parquet as shown below:
+
+```python
+ds = RayMLDataset.from_spark(..., fs_directory="hdfs://host:port/your/directory")
+```
+
+### Spark Submit
+
+RayDP provides a substitute for spark-submit in Apache Spark. You can run your java or scala application on RayDP cluster by using `bin/raydp-submit`. You can add it to `PATH` for convenience. When using `raydp-submit`, you should specify number of executors, number of cores and memory each executor by Spark properties, such as `--conf spark.executor.cores=1`, `--conf spark.executor.instances=1` and `--conf spark.executor.memory=500m`. `raydp-submit` only supports Ray cluster. Spark standalone, Apache Mesos, Apache Yarn are not supported, please use traditional `spark-submit` in that case. For the same reason, you do not need to specify `--master` in the command. Besides, RayDP does not support cluster as deploy-mode.
+
+### Integrating Spark with Deep Learning and Machine Learning Frameworks
+
+Combined with other ray components, such as RaySGD and RayServe, we can easily build an end-to-end deep learning pipeline.
+
+***MLDataset API***
+
+RayDP provides an API for creating a Ray MLDataset from a Spark dataframe. MLDataset represents a distributed dataset stored in Ray's in-memory object store. It supports transformation on each shard and can be converted to a PyTorch or Tensorflow dataset for distributed training. If you prefer to using Horovod on Ray or RaySGD for distributed training, you can use MLDataset to seamlessly integrate Spark with them.
+
+***Estimator API***
+
+RayDP also provides high level scikit-learn style Estimator APIs for distributed training. The Estimator APIs allow you to train a deep neural network directly on a Spark DataFrame, leveraging Ray’s ability to scale out across the cluster. The Estimator APIs are wrappers of RaySGD and hide the complexity of converting a Spark DataFrame to a PyTorch/Tensorflow dataset and distributing the training.
+
 ```python
 import ray
 import raydp
 from raydp.torch import TorchEstimator
 
-ray.init()
+ray.init(address="auto")
 spark = raydp.init_spark(app_name="RayDP example",
                          num_executors=2,
                          executor_cores=2,
@@ -86,20 +113,9 @@ estimator.fit_on_spark(train_df)
 raydp.stop_spark()
 ```
 
-## Dynamic Executor Allocation
-RayDP now supports External Shuffle Serivce. To enable it, you can either set `spark.shuffle.service.enabled` to `true` in `spark-defaults.conf`, or you can provide a config to `raydp.init_spark`, as shown below:
-```python
-raydp.init_spark(..., configs={"spark.shuffle.service.enabled": "true"})
-```
-The user-provided config will overwrite those specified in `spark-defaults.conf`. By default Spark will load `spark-defaults.conf` from `$SPARK_HOME/conf`, you can also modify this location by setting `SPARK_CONF_DIR`.
+## MPI on Ray
 
-Similarly, you can also enable Dynamic Executor Allocation this way. However, because Ray does not support object ownership tranferring now(1.3.0), you must use Dynamic Executor Allocation with data persistence. You can write the data frame in spark to HDFS as a parquet as shown below:
-```python
-ds = RayMLDataset.from_spark(..., fs_directory="hdfs://host:port/your/directory")
-```
+RayDP also provides a simple API to running MPI job on top of Ray. Currently, we support three types of MPI: `intel_mpi`, `openmpi` and `MPICH`. You can refer [doc/mpi.md](./doc/mpi.md) for more details.
 
 ## More Examples
 Not sure how to use RayDP? Check the `examples` folder. We have added many examples showing how RayDP works together with PyTorch, TensorFlow, XGBoost, Horovod, and so on. If you still cannot find what you want, feel free to post an issue to ask us!
-
-## Spark Submit
-RayDP provides a substitute for spark-submit in Apache Spark. You can run your java or scala application on RayDP cluster by using `bin/raydp-submit`. You can add it to `PATH` for convenience. When using `raydp-submit`, you should specify number of executors, number of cores and memory each executor by Spark properties, such as `--conf spark.executor.cores=1`, `--conf spark.executor.instances=1` and `--conf spark.executor.memory=500m`. `raydp-submit` only supports Ray cluster. Spark standalone, Apache Mesos, Apache Yarn are not supported, please use traditional `spark-submit` in that case. For the same reason, you do not need to specify `--master` in the command. Besides, RayDP does not support cluster as deploy-mode.

doc/mpi.md

@@ -0,0 +1,98 @@
+# MPI on Ray
+
+RayDP also provides a simple API to running MPI job on top of Ray. Currently, we support three types of MPI: `intel_mpi`, `openmpi` and `MPICH`. To use the following API, make sure you have installed the given type of MPI on each of Ray worker node.
+
+### API
+
+```python
+def create_mpi_job(job_name: str,
+                   world_size: int,
+                   num_cpus_per_process: int,
+                   num_processes_per_node: int,
+                   mpi_script_prepare_fn: Callable = None,
+                   timeout: int = 1,
+                   mpi_type: str = "intel_mpi",
+                   placement_group=None,
+                   placement_group_bundle_indexes: List[int] = None) -> MPIJob:
+    """ Create a MPI Job
+
+    :param job_name: the job name
+    :param world_size: the world size
+    :param num_cpus_per_process: num cpus per process, this used to request resource from Ray
+    :param num_processes_per_node: num processes per node
+    :param mpi_script_prepare_fn: a function used to create mpi script, it will pass in a
+        MPIJobcontext instance. It will use the default script if not provides.
+    :param timeout: the timeout used to wait for job creation
+    :param mpi_type: the mpi type, now only support openmpi, intel_mpi and mpich
+    :param placement_group: the placement_group for request mpi resources
+    :param placement_group_bundle_indexes: this should be equal with
+        world_size / num_processes_per_node if provides.
+    """
+```
+
+### Create a simple MPI Job
+
+```python
+from raydp.mpi import create_mpi_job, MPIJobContext, WorkerContext
+
+# Define the MPI JOb. We want to create a 4 world_size MPIJob, and each process requires 2 cpus.
+# We have set the num_processes_per_node to 2, so the processes will be strictly spread into two nodes.
+
+# You could also to specify the placement group to reserve the resources for MPI job. The num_cpus_per_process 
+# will be ignored if the placement group is provided. And the size of
+# placement_group_bundle_indexes should be equal with world_size // num_processes_per_node.
+job = create_mpi_job(job_name="example",
+                     world_size=4,
+                     num_cpus_per_process=2,
+                     num_processes_per_node=2,
+                     timeout=5,
+                     mpi_type="intel_mpi",
+                     placement_group=None,
+                     placement_group_bundle_indexes: List[int] = None)
+
+# Start the MPI Job, this will start up the MPI processes and connect to the ray cluster
+job.start()
+
+# define the MPI task function
+def func(context: WorkerContext):
+    return context.job_id
+
+# run the MPI task, this is a blocking operation. And the results is a world_size array.
+results = job.run(func)
+
+# stop the MPI job
+job.stop()
+```
+
+### Use `with` auto start/stop MPIJob
+```python
+with create_mpi_job(job_name="example",
+                    world_size=4,
+                    num_cpus_per_process=2,
+                    num_processes_per_node=2,
+                    timeout=5,
+                    mpi_type="intel_mpi") as job:
+    def f(context: WorkerContext):
+        return context.job_id
+    results = job.run(f)
+```
+
+### Specify the MPI script and environments
+
+You could customize the MPI job environments and MPI scritps with `mpi_script_prepare_fn` argument.
+
+```python
+def script_prepare_fn(context: MPIJobContext):
+    context.add_env("OMP_NUM_THREADS", "2")
+    default_script = ["mpirun", "--allow-run-as-root", "--tag-output", "-H",
+                      ",".join(context.hosts), "-N", f"{context.num_procs_per_node}"]
+    return default_script
+  
+job = create_mpi_job(job_name="example",
+                     world_size=4,
+                     num_cpus_per_process=2,
+                     num_processes_per_node=2,
+                     timeout=5,
+                     mpi_type="intel_mpi",
+                     mpi_script_prepare_fn=script_prepare_fn)
+```

python/raydp/mpi/__init__.py

@@ -18,14 +18,17 @@
 
 from typing import Callable, List
 
-from raydp.mpi.mpi_job import MPIJob, MPIType, IntelMPIJob, OpenMPIJob, MPIJobContext
+from .mpi_job import MPIJob, MPIType, IntelMPIJob, OpenMPIJob, MPICHJob, MPIJobContext
+from .mpi_worker import WorkerContext
 
 
 def _get_mpi_type(mpi_type: str) -> MPIType:
     if mpi_type.strip().lower() == "openmpi":
         return MPIType.OPEN_MPI
     elif mpi_type.strip().lower() == "intel_mpi":
         return MPIType.INTEL_MPI
+    elif mpi_type.strip().lower() == "mpich":
+        return MPIType.MPICH
     else:
         return None
 
@@ -39,16 +42,16 @@ def create_mpi_job(job_name: str,
                    mpi_type: str = "intel_mpi",
                    placement_group=None,
                    placement_group_bundle_indexes: List[int] = None) -> MPIJob:
-    """
-    Create a MPI Job
+    """Create a MPI Job
+
     :param job_name: the job name
     :param world_size: the world size
     :param num_cpus_per_process: num cpus per process, this used to request resource from Ray
     :param num_processes_per_node: num processes per node
     :param mpi_script_prepare_fn: a function used to create mpi script, it will pass in a
-        MPIJobcontext instance. It will use the default script if not provides.
+        MPIJobContext instance. It will use the default script if not provides.
     :param timeout: the timeout used to wait for job creation
-    :param mpi_type: the mpi type, now only support openmpi and intel_mpi
+    :param mpi_type: the mpi type, now only support openmpi, intel_mpi and MPICH
     :param placement_group: the placement_group for request mpi resources
     :param placement_group_bundle_indexes: this should be equal with
         world_size / num_processes_per_node if provides.
@@ -74,8 +77,18 @@ def create_mpi_job(job_name: str,
                            timeout=timeout,
                            placement_group=placement_group,
                            placement_group_bundle_indexes=placement_group_bundle_indexes)
+    elif mpi_type == MPIType.MPICH:
+        return MPICHJob(mpi_type=MPIType.MPICH,
+                        job_name=job_name,
+                        world_size=world_size,
+                        num_cpus_per_process=num_cpus_per_process,
+                        num_processes_per_node=num_processes_per_node,
+                        mpi_script_prepare_fn=mpi_script_prepare_fn,
+                        timeout=timeout,
+                        placement_group=placement_group,
+                        placement_group_bundle_indexes=placement_group_bundle_indexes)
     else:
         raise Exception(f"MPI type: {mpi_type} not supported now")
 
 
-__all__ = ["create_mpi_job", "MPIJobContext"]
+__all__ = ["create_mpi_job", "MPIJobContext", "WorkerContext"]