Add target_job_queue to SubmitAnyscaleJob operator (#52)

Anyscale Job Queues is a new functionality that supports submitting
multiple workloads to the same cluster. An Anyscale Customer requested
this feature on the Airflow integration.

[Job Queues](https://docs.anyscale.com/1.0.0/jobs/job-queues)

And the [Job Queue
Spec](https://docs.anyscale.com/canary/reference/job-api#jobqueuespec)

Adds two optional parameters to submit to existing job queues. _Does
not_ implement creation of new job queues.

Closes: #44

---------

Co-authored-by: marwan116 <sarieddine.marwan@gmail.com>

Author: RCdeWit

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
@@ -26,6 +26,7 @@ class SubmitAnyscaleJob(BaseOperator):
 
     :param conn_id: Required. The connection ID for Anyscale.
     :param entrypoint: Required. Command that will be run to execute the job, e.g., `python main.py`.
+    :param job_queue_config: Optional. Allow users to specify JobQueueConfig (more information: https://docs.anyscale.com/platform/jobs/job-queues/ and https://docs.anyscale.com/reference/job-api/).
     :param name: Optional. Name of the job. Multiple jobs can be submitted with the same name.
     :param image_uri: Optional. URI of an existing image. Exclusive with `containerfile`.
     :param containerfile: Optional. The file path to a containerfile that will be built into an image before running the workload. Exclusive with `image_uri`.
@@ -81,6 +82,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 +105,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 +151,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]