Orchestrate Ray jobs on Anyscale with Apache Airflow®

​

Time to complete

​

Assumed knowledge

• Ray basics. See the Getting Started section of the Ray documentation.
• Anyscale basics. See Get started section of the Anyscale documentation.
• Airflow operators. See Airflow operators.

​

Prerequisites

• The Astro CLI.
• An Anyscale account with AI platform features enabled. You also need to have at least one suitable image and compute config available in your Anyscale account.

​

Step 1: Configure your Astro project

1. Create a new Astro project:
   $ mkdir astro-anyscale-tutorial && cd astro-anyscale-tutorial
   $ astro dev init
   
2. In the requirements.txt file, add the Anyscale provider.
   astro-provider-anyscale==1.0.1
   
3. Run the following command to start your Airflow project:
   astro dev start
   

​

Step 2: Configure a Ray connection

1. In the Airflow UI, go to Admin -> Connections and click +.
2. Create a new connection and choose the Anyscale connection type. Enter the following information:
   • Connection ID: anyscale_conn
   • API Key: Your Anyscale API key
3. Click Save.

​

Step 3: Write a DAG to orchestrate Anyscale jobs

1. Create a new file in your dags directory called anyscale_script.py and add the following code:
   # anyscale_script.py
   import numpy as np
   import ray
   import argparse
   
   @ray.remote
   def square(x):
       return x**2
   
   def main(data):
       ray.init()
       data = np.array(data)
       futures = [square.remote(x) for x in data]
       results = ray.get(futures)
       mean = np.mean(results)
       print(f"Mean squared value: {mean}")
       return mean
   
   if __name__ == "__main__":
       parser = argparse.ArgumentParser(description="Process some integers.")
       parser.add_argument(
           "data", nargs="+", type=float, help="List of numbers to process"
       )
       args = parser.parse_args()
       data = args.data
       main(data)
   
   
2. Create a new file in your dags directory called anyscale_tutorial.py.
3. Copy and paste the code below into the file:
   """
   ## Anyscale Tutorial
   
   This tutorial demonstrates how to use the Anyscale provider in Airflow to
   parallelize a task using Ray on Anyscale.
   """
   
   from airflow.decorators import dag
   from airflow.operators.python import PythonOperator
   from anyscale_provider.operators.anyscale import SubmitAnyscaleJob
   from airflow.models.baseoperator import chain
   from pathlib import Path
   
   CONN_ID = "anyscale_conn"
   FOLDER_PATH = Path(__file__).parent
   
   
   def _generate_data() -> list:
       """
       Generate sample data
       Returns:
           list: List of integers
       """
       import random
   
       return [random.randint(1, 100) for _ in range(10)]
   
   
   @dag(
       start_date=None,
       schedule=None,
       catchup=False,
       tags=["ray", "example"],
       doc_md=__doc__,
   )
   def anyscale_tutorial():
   
       data = PythonOperator(
           task_id="generate_data",
           python_callable=_generate_data,
       )
   
       get_mean_squared_value = SubmitAnyscaleJob(
           task_id="SubmitRayJob",
           conn_id=CONN_ID,
           name="AstroJob",
           image_uri="< your image uri >",  # e.g. "anyscale/ray:2.35.0-slim-py312-cpu"
           compute_config="< your compute config >",  # e.g. airflow-integration-testing:1
           entrypoint="python anyscale_script.py {{ ti.xcom_pull(task_ids='generate_data') | join(' ') }}",
           working_dir=str(FOLDER_PATH),  # the folder containing the script
           requirements=["requests", "pandas", "numpy", "torch"],
           max_retries=1,
           job_timeout_seconds=3000,
           poll_interval=30,
       )
   
       chain(data, get_mean_squared_value)
   
   
   anyscale_tutorial()
   

   See all 60 lines
   • The generate_data task randomly generates a list of 10 integers.
   • The get_mean_squared_value task submits a Ray job on Anyscale to calculate the mean squared value of the list of integers.

​

Step 4: Run the DAG

1. In the Airflow UI, click the play button to manually run your DAG.
2. After the DAG runs successfully, check your Anyscale account to see the job submitted by Airflow.
   

​

Conclusion