Add target_job_queue to SubmitAnyscaleJob operator

anyscale_provider/operators/anyscale.py

@@ -9,7 +9,7 @@
 from airflow.models import BaseOperator
 from airflow.utils.context import Context
 from anyscale.compute_config.models import ComputeConfig
-from anyscale.job.models import JobConfig, JobState
+from anyscale.job.models import JobConfig, JobQueueConfig, JobState
 from anyscale.service.models import RayGCSExternalStorageConfig, ServiceConfig, ServiceState
 
 from anyscale_provider.hooks.anyscale import AnyscaleHook
@@ -81,6 +81,7 @@ def __init__(
         wait_for_completion: bool = True,
         job_timeout_seconds: float = 3600,
         poll_interval: float = 60,
+        job_queue_config: JobQueueConfig | None = None,
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -103,6 +104,7 @@ def __init__(
         self.wait_for_completion = wait_for_completion
         self.job_timeout_seconds = job_timeout_seconds
         self.poll_interval = poll_interval
+        self.job_queue_config = job_queue_config
 
         self.job_id: str | None = None
 
@@ -148,6 +150,7 @@ def execute(self, context: Context) -> str:
             "cloud": self.cloud,
             "project": self.project,
             "max_retries": self.max_retries,
+            "job_queue_config": self.job_queue_config,
         }
 
         self.log.info(f"Using Anyscale version {anyscale.__version__}")

docs/index.rst

@@ -22,7 +22,7 @@ Hook
 
 Operators
 ~~~~~~~~~
-- **SubmitAnyscaleJob**: This operator submits a job to Anyscale. It takes configuration parameters for the job, such as the entrypoint, image URI, and compute configuration. The operator uses ``AnyscaleHook`` to handle the submission process.
+- **SubmitAnyscaleJob**: This operator submits a job to Anyscale. It takes configuration parameters for the job, such as the entrypoint, image URI, compute configuration, and job queue configuration. The operator uses ``AnyscaleHook`` to handle the submission process.
 - **RolloutAnyscaleService**: Similar to the job submission operator, this operator is designed to manage services on Anyscale. It can be used to deploy new services or update existing ones, leveraging ``AnyscaleHook`` for all interactions with the Anyscale API.
 
 Triggers

example_dags/anyscale_job.py

@@ -4,6 +4,8 @@
 from airflow import DAG
 
 from anyscale_provider.operators.anyscale import SubmitAnyscaleJob
+from anyscale.job.models import JobQueueConfig, JobQueueSpec, JobQueueExecutionMode
+
 
 default_args = {
     "owner": "airflow",
@@ -44,6 +46,77 @@
     dag=dag,
 )
 
+JOB_QUEUE_NAME = "test-job-queue-180s-idle-timeout"
+
+submit_anyscale_job_with_new_job_queue = SubmitAnyscaleJob(
+    task_id="submit_anyscale_job_with_new_job_queue",
+    conn_id=ANYSCALE_CONN_ID,
+    name="AstroJobWithJobQueue",
+    image_uri="anyscale/image/airflow-integration-testing:1",
+    compute_config="airflow-integration-testing:1",
+    working_dir=str(FOLDER_PATH),
+    entrypoint="python ray-job.py",
+    max_retries=1,
+    job_timeout_seconds=3000,
+    poll_interval=30,
+    dag=dag,
+    job_queue_config=JobQueueConfig(
+        priority=100,
+        job_queue_spec=JobQueueSpec(
+            name=JOB_QUEUE_NAME,
+            execution_mode=JobQueueExecutionMode.PRIORITY,
+            max_concurrency=5,
+            idle_timeout_s=180,
+        ),
+    ),
+)
+
+submit_anyscale_job_with_existing_job_queue = SubmitAnyscaleJob(
+    task_id="submit_anyscale_job_with_existing_job_queue",
+    conn_id=ANYSCALE_CONN_ID,
+    name="AstroJobWithJobQueue",
+    image_uri="anyscale/image/airflow-integration-testing:1",
+    compute_config="airflow-integration-testing:1",
+    working_dir=str(FOLDER_PATH),
+    entrypoint="python ray-job.py",
+    max_retries=1,
+    job_timeout_seconds=3000,
+    poll_interval=30,
+    dag=dag,
+    job_queue_config=JobQueueConfig(
+        priority=100,
+        # NOTE: This will reuse the existing job queue given it has the same spec
+        job_queue_spec=JobQueueSpec(
+            name=JOB_QUEUE_NAME,
+            execution_mode=JobQueueExecutionMode.PRIORITY,
+            max_concurrency=5,
+            idle_timeout_s=180,
+        ),
+    ),
+)
+
+submit_another_anyscale_job_with_existing_job_queue = SubmitAnyscaleJob(
+    task_id="submit_another_anyscale_job_with_existing_job_queue",
+    conn_id=ANYSCALE_CONN_ID,
+    name="AstroJobWithJobQueue",
+    image_uri="anyscale/image/airflow-integration-testing:1",
+    compute_config="airflow-integration-testing:1",
+    working_dir=str(FOLDER_PATH),
+    entrypoint="python ray-job.py",
+    max_retries=1,
+    job_timeout_seconds=3000,
+    poll_interval=30,
+    dag=dag,
+    job_queue_config=JobQueueConfig(
+        priority=100,
+        target_job_queue_name=JOB_QUEUE_NAME,
+    ),
+)
+
 
 # Defining the task sequence
 submit_anyscale_job
+submit_anyscale_job_with_new_job_queue >> [
+    submit_anyscale_job_with_existing_job_queue,
+    submit_another_anyscale_job_with_existing_job_queue,
+]

pyproject.toml

@@ -26,7 +26,7 @@ requires-python = ">=3.8, <3.13"
 dependencies = [
     "apache-airflow>=2.7",
     "pyyaml",
-    "anyscale==0.24.44",
+    "anyscale>=0.24.54",
 ]
 
 [project.urls]