kueue-tests

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/config.py

@@ -131,10 +131,8 @@ class ManagedClusterConfig:
         accelerator_configs:
             A dictionary of custom resource mappings to map extended resource requests to RayCluster resource names.
             Defaults to DEFAULT_ACCELERATORS but can be overridden with custom mappings.
-        local_queue:
-            The name of the queue to use for the cluster.
         annotations:
-            A dictionary of annotations to apply to the cluster.
+            A dictionary of annotations to apply to the Job.
         volumes:
             A list of V1Volume objects to add to the Cluster
         volume_mounts:
@@ -161,7 +159,6 @@ class ManagedClusterConfig:
     accelerator_configs: Dict[str, str] = field(
         default_factory=lambda: DEFAULT_ACCELERATORS.copy()
     )
-    local_queue: Optional[str] = None
     annotations: Dict[str, str] = field(default_factory=dict)
     volumes: list[V1Volume] = field(default_factory=list)
     volume_mounts: list[V1VolumeMount] = field(default_factory=list)
@@ -248,7 +245,6 @@ def build_ray_cluster_spec(self, cluster_name: str) -> Dict[str, Any]:
         """
         ray_cluster_spec = {
             "rayVersion": RAY_VERSION,
-            "enableInTreeAutoscaling": False,
             "headGroupSpec": self._build_head_group_spec(),
             "workerGroupSpecs": [self._build_worker_group_spec(cluster_name)],
         }
@@ -346,12 +342,9 @@ def _build_head_container(self) -> V1Container:
                 self.head_accelerators,
             ),
             volume_mounts=self._generate_volume_mounts(),
+            env=self._build_env_vars() if hasattr(self, "envs") and self.envs else None,
         )
 
-        # Add environment variables if specified
-        if hasattr(self, "envs") and self.envs:
-            container.env = self._build_env_vars()
-
         return container
 
     def _build_worker_container(self) -> V1Container:
@@ -373,12 +366,9 @@ def _build_worker_container(self) -> V1Container:
                 self.worker_accelerators,
             ),
             volume_mounts=self._generate_volume_mounts(),
+            env=self._build_env_vars() if hasattr(self, "envs") and self.envs else None,
         )
 
-        # Add environment variables if specified
-        if hasattr(self, "envs") and self.envs:
-            container.env = self._build_env_vars()
-
         return container
 
     def _build_resource_requirements(

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -22,6 +22,7 @@
 import re
 import ast
 from typing import Dict, Any, Optional, Tuple
+from codeflare_sdk.common.kueue.kueue import get_default_kueue_name
 from codeflare_sdk.common.utils.constants import MOUNT_PATH
 from kubernetes import client
 from ...common.kubernetes_cluster.auth import get_api_client
@@ -59,9 +60,9 @@ def __init__(
         cluster_config: Optional[ManagedClusterConfig] = None,
         namespace: Optional[str] = None,
         runtime_env: Optional[Dict[str, Any]] = None,
-        shutdown_after_job_finishes: Optional[bool] = None,
         ttl_seconds_after_finished: int = 0,
         active_deadline_seconds: Optional[int] = None,
+        local_queue: Optional[str] = None,
     ):
         """
         Initialize a RayJob instance.
@@ -73,12 +74,11 @@ def __init__(
             cluster_config: Configuration for creating a new cluster (optional if cluster_name provided)
             namespace: The Kubernetes namespace (auto-detected if not specified)
             runtime_env: Ray runtime environment configuration (optional)
-            shutdown_after_job_finishes: Whether to shut down cluster after job finishes (optional)
             ttl_seconds_after_finished: Seconds to wait before cleanup after job finishes (default: 0)
             active_deadline_seconds: Maximum time the job can run before being terminated (optional)
+            local_queue: The Kueue LocalQueue to submit the job to (optional)
 
         Note:
-            shutdown_after_job_finishes is automatically detected but can be overridden:
             - True if cluster_config is provided (new cluster will be cleaned up)
             - False if cluster_name is provided (existing cluster will not be shut down)
             - User can explicitly set this value to override auto-detection
@@ -108,17 +108,7 @@ def __init__(
         self.runtime_env = runtime_env
         self.ttl_seconds_after_finished = ttl_seconds_after_finished
         self.active_deadline_seconds = active_deadline_seconds
-
-        # Auto-set shutdown_after_job_finishes based on cluster_config presence
-        # If cluster_config is provided, we want to clean up the cluster after job finishes
-        # If using existing cluster, we don't want to shut it down
-        # User can override this behavior by explicitly setting shutdown_after_job_finishes
-        if shutdown_after_job_finishes is not None:
-            self.shutdown_after_job_finishes = shutdown_after_job_finishes
-        elif cluster_config is not None:
-            self.shutdown_after_job_finishes = True
-        else:
-            self.shutdown_after_job_finishes = False
+        self.local_queue = local_queue
 
         if namespace is None:
             detected_namespace = get_current_namespace()
@@ -177,10 +167,6 @@ def submit(self) -> str:
                 if scripts:
                     self._handle_script_volumes_for_existing_cluster(scripts, result)
 
-            if self.shutdown_after_job_finishes:
-                logger.info(
-                    f"Cluster will be automatically cleaned up {self.ttl_seconds_after_finished}s after job completion"
-                )
             return self.name
         else:
             raise RuntimeError(f"Failed to submit RayJob {self.name}")
@@ -230,11 +216,20 @@ def _build_rayjob_cr(self) -> Dict[str, Any]:
             },
             "spec": {
                 "entrypoint": self.entrypoint,
-                "shutdownAfterJobFinishes": self.shutdown_after_job_finishes,
                 "ttlSecondsAfterFinished": self.ttl_seconds_after_finished,
+                "shutdownAfterJobFinishes": True,
+                "enableInTreeAutoscaling": False,
             },
         }
 
+        labels = {}
+        if self.local_queue:
+            labels["kueue.x-k8s.io/queue-name"] = self.local_queue
+        else:
+            labels["kueue.x-k8s.io/queue-name"] = get_default_kueue_name(self.namespace)
+
+        rayjob_cr["metadata"]["labels"] = labels
+
         # Add active deadline if specified
         if self.active_deadline_seconds:
             rayjob_cr["spec"]["activeDeadlineSeconds"] = self.active_deadline_seconds

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -498,40 +498,6 @@ def test_namespace_explicit_override(mocker):
     assert rayjob.namespace == "explicit-ns"
 
 
-def test_shutdown_behavior_with_cluster_config(mocker):
-    """Test that shutdown_after_job_finishes is True when cluster_config is provided."""
-    mocker.patch("kubernetes.config.load_kube_config")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
-
-    cluster_config = ManagedClusterConfig()
-
-    rayjob = RayJob(
-        job_name="test-job",
-        entrypoint="python script.py",
-        cluster_config=cluster_config,
-        namespace="test-namespace",
-    )
-
-    assert rayjob.shutdown_after_job_finishes is True
-
-
-def test_shutdown_behavior_with_existing_cluster(mocker):
-    """Test that shutdown_after_job_finishes is False when using existing cluster."""
-    mocker.patch("kubernetes.config.load_kube_config")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
-
-    rayjob = RayJob(
-        job_name="test-job",
-        entrypoint="python script.py",
-        cluster_name="existing-cluster",
-        namespace="test-namespace",
-    )
-
-    assert rayjob.shutdown_after_job_finishes is False
-
-
 def test_rayjob_with_rayjob_cluster_config(mocker):
     """Test RayJob with the new ManagedClusterConfig."""
     mocker.patch("kubernetes.config.load_kube_config")
@@ -935,48 +901,6 @@ def test_build_ray_cluster_spec_with_image_pull_secrets(mocker):
     assert worker_secrets[1].name == "another-secret"
 
 
-def test_rayjob_user_override_shutdown_behavior(mocker):
-    """Test that user can override the auto-detected shutdown behavior."""
-    mocker.patch("kubernetes.config.load_kube_config")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
-    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
-
-    # Test 1: User overrides shutdown to True even when using existing cluster
-    rayjob_existing_override = RayJob(
-        job_name="test-job",
-        entrypoint="python script.py",
-        cluster_name="existing-cluster",
-        shutdown_after_job_finishes=True,  # User override
-        namespace="test-namespace",  # Explicitly specify namespace
-    )
-
-    assert rayjob_existing_override.shutdown_after_job_finishes is True
-
-    # Test 2: User overrides shutdown to False even when creating new cluster
-    cluster_config = ManagedClusterConfig()
-
-    rayjob_new_override = RayJob(
-        job_name="test-job",
-        entrypoint="python script.py",
-        cluster_config=cluster_config,
-        shutdown_after_job_finishes=False,  # User override
-        namespace="test-namespace",  # Explicitly specify namespace
-    )
-
-    assert rayjob_new_override.shutdown_after_job_finishes is False
-
-    # Test 3: User override takes precedence over auto-detection
-    rayjob_override_priority = RayJob(
-        job_name="test-job",
-        entrypoint="python script.py",
-        cluster_config=cluster_config,
-        shutdown_after_job_finishes=True,  # Should override auto-detection
-        namespace="test-namespace",  # Explicitly specify namespace
-    )
-
-    assert rayjob_override_priority.shutdown_after_job_finishes is True
-
-
 class TestRayVersionValidation:
     """Test Ray version validation in RayJob."""
 
@@ -2353,7 +2277,7 @@ def test_rayjob_init_missing_cluster_name_with_no_config(mocker):
         rayjob.runtime_env = None
         rayjob.ttl_seconds_after_finished = 0
         rayjob.active_deadline_seconds = None
-        rayjob.shutdown_after_job_finishes = False
+        rayjob.shutdown_after_job_finishes = True
         rayjob.namespace = "test-namespace"
         rayjob._cluster_name = None
         rayjob._cluster_config = None