task(RHOAIENG-33283): Utilize runtime_env from Ray Public API

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -19,6 +19,8 @@
 import logging
 import warnings
 from typing import Dict, Any, Optional, Tuple
+
+from ray.runtime_env import RuntimeEnv
 from codeflare_sdk.common.kueue.kueue import get_default_kueue_name
 from codeflare_sdk.common.utils.constants import MOUNT_PATH
 
@@ -61,7 +63,7 @@ def __init__(
         cluster_name: Optional[str] = None,
         cluster_config: Optional[ManagedClusterConfig] = None,
         namespace: Optional[str] = None,
-        runtime_env: Optional[Dict[str, Any]] = None,
+        runtime_env: Optional[RuntimeEnv] = None,
         ttl_seconds_after_finished: int = 0,
         active_deadline_seconds: Optional[int] = None,
         local_queue: Optional[str] = None,
@@ -75,7 +77,7 @@ def __init__(
             cluster_name: The name of an existing Ray cluster (optional if cluster_config provided)
             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)
+            runtime_env: Ray runtime environment configuration as RuntimeEnv object (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)

src/codeflare_sdk/ray/rayjobs/runtime_env.py

@@ -8,6 +8,7 @@
 from typing import Dict, Any, Optional, List, TYPE_CHECKING
 from codeflare_sdk.common.utils.constants import MOUNT_PATH
 from kubernetes import client
+from ray.runtime_env import RuntimeEnv
 
 from codeflare_sdk.ray.rayjobs.config import ManagedClusterConfig
 from ...common.kubernetes_cluster.auth import get_api_client
@@ -22,6 +23,14 @@
 PYTHON_FILE_PATTERN = r"(?:python\s+)?([./\w/]+\.py)"
 
 
+def _normalize_runtime_env(
+    runtime_env: Optional[RuntimeEnv],
+) -> Optional[Dict[str, Any]]:
+    if runtime_env is None:
+        return None
+    return runtime_env.to_dict()
+
+
 def extract_all_local_files(job: RayJob) -> Optional[Dict[str, str]]:
     """
     Extract all local files from both entrypoint and runtime_env working_dir.
@@ -32,14 +41,17 @@ def extract_all_local_files(job: RayJob) -> Optional[Dict[str, str]]:
     Returns:
         Dict of {file_name: file_content} if local files found, None otherwise
     """
+    # Convert RuntimeEnv to dict for processing
+    runtime_env_dict = _normalize_runtime_env(job.runtime_env)
+
     # If there's a remote working_dir, don't extract local files to avoid conflicts
     if (
-        job.runtime_env
-        and "working_dir" in job.runtime_env
-        and not os.path.isdir(job.runtime_env["working_dir"])
+        runtime_env_dict
+        and "working_dir" in runtime_env_dict
+        and not os.path.isdir(runtime_env_dict["working_dir"])
     ):
         logger.info(
-            f"Remote working_dir detected: {job.runtime_env['working_dir']}. "
+            f"Remote working_dir detected: {runtime_env_dict['working_dir']}. "
             "Skipping local file extraction - all files should come from remote source."
         )
         return None
@@ -55,18 +67,18 @@ def extract_all_local_files(job: RayJob) -> Optional[Dict[str, str]]:
 
     # Extract files from runtime_env working_dir if it's a local directory
     if (
-        job.runtime_env
-        and "working_dir" in job.runtime_env
-        and os.path.isdir(job.runtime_env["working_dir"])
+        runtime_env_dict
+        and "working_dir" in runtime_env_dict
+        and os.path.isdir(runtime_env_dict["working_dir"])
     ):
         working_dir_files = extract_working_dir_files(
-            job.runtime_env["working_dir"], processed_files
+            runtime_env_dict["working_dir"], processed_files
         )
         if working_dir_files:
             files.update(working_dir_files)
 
     # If no working_dir specified in runtime_env, try to infer and extract files from inferred directory
-    elif not job.runtime_env or "working_dir" not in job.runtime_env:
+    elif not runtime_env_dict or "working_dir" not in runtime_env_dict:
         inferred_working_dir = infer_working_dir_from_entrypoint(job)
         if inferred_working_dir:
             working_dir_files = extract_working_dir_files(
@@ -221,24 +233,27 @@ def process_runtime_env(
     Returns:
         Processed runtime environment as YAML string, or None if no processing needed
     """
+    # Convert RuntimeEnv to dict for processing
+    runtime_env_dict = _normalize_runtime_env(job.runtime_env)
+
     processed_env = {}
 
     # Handle env_vars
-    if job.runtime_env and "env_vars" in job.runtime_env:
-        processed_env["env_vars"] = job.runtime_env["env_vars"]
+    if runtime_env_dict and "env_vars" in runtime_env_dict:
+        processed_env["env_vars"] = runtime_env_dict["env_vars"]
         logger.info(
-            f"Added {len(job.runtime_env['env_vars'])} environment variables to runtime_env"
+            f"Added {len(runtime_env_dict['env_vars'])} environment variables to runtime_env"
         )
 
     # Handle pip dependencies
-    if job.runtime_env and "pip" in job.runtime_env:
-        pip_deps = process_pip_dependencies(job, job.runtime_env["pip"])
+    if runtime_env_dict and "pip" in runtime_env_dict:
+        pip_deps = process_pip_dependencies(job, runtime_env_dict["pip"])
         if pip_deps:
             processed_env["pip"] = pip_deps
 
     # Handle working_dir - if it's a local path, set it to mount path
-    if job.runtime_env and "working_dir" in job.runtime_env:
-        working_dir = job.runtime_env["working_dir"]
+    if runtime_env_dict and "working_dir" in runtime_env_dict:
+        working_dir = runtime_env_dict["working_dir"]
         if os.path.isdir(working_dir):
             # Local working directory - will be mounted at MOUNT_PATH
             processed_env["working_dir"] = MOUNT_PATH
@@ -252,7 +267,7 @@ def process_runtime_env(
             logger.info(f"Using remote working directory: {working_dir}")
 
     # If no working_dir specified but we have files, set working_dir to mount path
-    elif not job.runtime_env or "working_dir" not in job.runtime_env:
+    elif not runtime_env_dict or "working_dir" not in runtime_env_dict:
         if files:
             # Local files found - will be mounted at MOUNT_PATH
             processed_env["working_dir"] = MOUNT_PATH

src/codeflare_sdk/ray/rayjobs/test/test_rayjob.py

@@ -15,6 +15,7 @@
 import pytest
 from unittest.mock import MagicMock
 from codeflare_sdk.common.utils.constants import RAY_VERSION
+from ray.runtime_env import RuntimeEnv
 
 from codeflare_sdk.ray.rayjobs.rayjob import RayJob
 from codeflare_sdk.ray.cluster.config import ClusterConfiguration
@@ -35,7 +36,7 @@ def test_rayjob_submit_success(auto_mock_setup):
         cluster_name="test-ray-cluster",
         namespace="test-namespace",
         entrypoint="python -c 'print(\"hello world\")'",
-        runtime_env={"pip": ["requests"]},
+        runtime_env=RuntimeEnv(pip=["requests"]),
     )
 
     job_id = rayjob.submit()
@@ -68,7 +69,7 @@ def test_rayjob_submit_failure(auto_mock_setup):
         cluster_name="test-ray-cluster",
         namespace="default",
         entrypoint="python test.py",
-        runtime_env={"pip": ["numpy"]},
+        runtime_env=RuntimeEnv(pip=["numpy"]),
     )
 
     with pytest.raises(RuntimeError, match="Failed to submit RayJob test-rayjob"):
@@ -474,7 +475,7 @@ def test_rayjob_with_runtime_env(auto_mock_setup):
     """
     Test RayJob with runtime environment configuration.
     """
-    runtime_env = {"pip": ["numpy", "pandas"]}
+    runtime_env = RuntimeEnv(pip=["numpy", "pandas"])
 
     rayjob = RayJob(
         job_name="test-job",
@@ -567,7 +568,7 @@ def test_rayjob_constructor_parameter_validation(auto_mock_setup):
         entrypoint="python test.py",
         cluster_name="test-cluster",
         namespace="test-ns",
-        runtime_env={"pip": ["numpy"]},
+        runtime_env=RuntimeEnv(pip=["numpy"]),
         ttl_seconds_after_finished=300,
         active_deadline_seconds=600,
     )
@@ -576,7 +577,12 @@ def test_rayjob_constructor_parameter_validation(auto_mock_setup):
     assert rayjob.entrypoint == "python test.py"
     assert rayjob.cluster_name == "test-cluster"
     assert rayjob.namespace == "test-ns"
-    assert rayjob.runtime_env == {"pip": ["numpy"]}
+    # Check that runtime_env is a RuntimeEnv object and contains pip dependencies
+    assert isinstance(rayjob.runtime_env, RuntimeEnv)
+    runtime_env_dict = rayjob.runtime_env.to_dict()
+    assert "pip" in runtime_env_dict
+    # Ray transforms pip to dict format with 'packages' key
+    assert runtime_env_dict["pip"]["packages"] == ["numpy"]
     assert rayjob.ttl_seconds_after_finished == 300
     assert rayjob.active_deadline_seconds == 600
 

src/codeflare_sdk/ray/rayjobs/test/test_runtime_env.py

@@ -17,6 +17,7 @@
 import os
 from unittest.mock import MagicMock, patch
 from codeflare_sdk.common.utils.constants import MOUNT_PATH, RAY_VERSION
+from ray.runtime_env import RuntimeEnv
 
 from codeflare_sdk.ray.rayjobs.rayjob import RayJob
 from codeflare_sdk.ray.cluster.config import ClusterConfiguration
@@ -43,11 +44,11 @@ def test_rayjob_with_remote_working_dir(auto_mock_setup):
     Test RayJob with remote working directory in runtime_env.
     Should not extract local files and should pass through remote URL.
     """
-    runtime_env = {
-        "working_dir": "https://github.com/org/repo/archive/refs/heads/main.zip",
-        "pip": ["numpy", "pandas"],
-        "env_vars": {"TEST_VAR": "test_value"},
-    }
+    runtime_env = RuntimeEnv(
+        working_dir="https://github.com/org/repo/archive/refs/heads/main.zip",
+        pip=["numpy", "pandas"],
+        env_vars={"TEST_VAR": "test_value"},
+    )
 
     rayjob = RayJob(
         job_name="test-job",