task(RHOAIENG-33283): Add runtime env e2e

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -18,7 +18,7 @@
 
 import logging
 import warnings
-from typing import Dict, Any, Optional, Tuple
+from typing import Dict, Any, Optional, Tuple, Union
 
 from ray.runtime_env import RuntimeEnv
 from codeflare_sdk.common.kueue.kueue import get_default_kueue_name
@@ -63,7 +63,7 @@ def __init__(
         cluster_name: Optional[str] = None,
         cluster_config: Optional[ManagedClusterConfig] = None,
         namespace: Optional[str] = None,
-        runtime_env: Optional[RuntimeEnv] = None,
+        runtime_env: Optional[Union[RuntimeEnv, Dict[str, Any]]] = None,
         ttl_seconds_after_finished: int = 0,
         active_deadline_seconds: Optional[int] = None,
         local_queue: Optional[str] = None,
@@ -77,7 +77,10 @@ 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 as RuntimeEnv object (optional)
+            runtime_env: Ray runtime environment configuration. Can be:
+                - RuntimeEnv object from ray.runtime_env
+                - Dict with keys like 'working_dir', 'pip', 'env_vars', etc.
+                Example: {"working_dir": "./my-scripts", "pip": ["requests"]}
             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)
@@ -109,7 +112,13 @@ def __init__(
 
         self.name = job_name
         self.entrypoint = entrypoint
-        self.runtime_env = runtime_env
+
+        # Convert dict to RuntimeEnv if needed for user convenience
+        if isinstance(runtime_env, dict):
+            self.runtime_env = RuntimeEnv(**runtime_env)
+        else:
+            self.runtime_env = runtime_env
+
         self.ttl_seconds_after_finished = ttl_seconds_after_finished
         self.active_deadline_seconds = active_deadline_seconds
         self.local_queue = local_queue
@@ -232,6 +241,9 @@ def _build_rayjob_cr(self) -> Dict[str, Any]:
             },
         }
 
+        # Extract files once and use for both runtime_env and submitter pod
+        files = extract_all_local_files(self)
+
         labels = {}
         # If cluster_config is provided, use the local_queue from the cluster_config
         if self._cluster_config is not None:
@@ -262,9 +274,6 @@ def _build_rayjob_cr(self) -> Dict[str, Any]:
         if self.active_deadline_seconds:
             rayjob_cr["spec"]["activeDeadlineSeconds"] = self.active_deadline_seconds
 
-        # Extract files once and use for both runtime_env and submitter pod
-        files = extract_all_local_files(self)
-
         # Add runtime environment (can be inferred even if not explicitly specified)
         processed_runtime_env = process_runtime_env(self, files)
         if processed_runtime_env:
@@ -326,8 +335,19 @@ def _build_submitter_pod_template(
 
         # Build ConfigMap items for each file
         config_map_items = []
+        entrypoint_path = files.get(
+            "__entrypoint_path__"
+        )  # Metadata for single file case
+
         for file_name in files.keys():
-            config_map_items.append({"key": file_name, "path": file_name})
+            if file_name == "__entrypoint_path__":
+                continue  # Skip metadata key
+
+            # For single file case, use the preserved path structure
+            if entrypoint_path:
+                config_map_items.append({"key": file_name, "path": entrypoint_path})
+            else:
+                config_map_items.append({"key": file_name, "path": file_name})
 
         # Check if we need to unzip working_dir
         has_working_dir_zip = "working_dir.zip" in files

src/codeflare_sdk/ray/rayjobs/runtime_env.py

@@ -22,7 +22,8 @@
 logger = logging.getLogger(__name__)
 
 # Regex pattern for finding Python files in entrypoint commands
-PYTHON_FILE_PATTERN = r"(?:python\s+)?([./\w/]+\.py)"
+# Matches paths like: test.py, ./test.py, dir/test.py, my-dir/test.py
+PYTHON_FILE_PATTERN = r"(?:python\s+)?([./\w/-]+\.py)"
 
 # Path where working_dir will be unzipped on submitter pod
 UNZIP_PATH = "/tmp/rayjob-working-dir"
@@ -126,11 +127,15 @@ def _extract_single_entrypoint_file(job: RayJob) -> Optional[Dict[str, str]]:
     """
     Extract single Python file from entrypoint if no working_dir specified.
 
+    Returns a dict with metadata about the file path structure so we can
+    preserve it when mounting via ConfigMap.
+
     Args:
         job: RayJob instance
 
     Returns:
-        Dict of {filename: content} or None
+        Dict with special format: {"__entrypoint_path__": path, "filename": content}
+        This allows us to preserve directory structure when mounting
     """
     if not job.entrypoint:
         return None
@@ -145,9 +150,15 @@ def _extract_single_entrypoint_file(job: RayJob) -> Optional[Dict[str, str]]:
                 with open(file_path, "r") as f:
                     content = f.read()
 
+                # Use basename as key (ConfigMap keys can't have slashes)
+                # But store the full path for later use in ConfigMap item.path
                 filename = os.path.basename(file_path)
+                relative_path = file_path.lstrip("./")
+
                 logger.info(f"Extracted single entrypoint file: {file_path}")
-                return {filename: content}
+
+                # Return special format with metadata
+                return {"__entrypoint_path__": relative_path, filename: content}
 
             except (IOError, OSError) as e:
                 logger.warning(f"Could not read entrypoint file {file_path}: {e}")
@@ -348,10 +359,13 @@ def create_file_configmap(
     # Use a basic config builder for ConfigMap creation
     config_builder = ManagedClusterConfig()
 
+    # Filter out metadata keys (like __entrypoint_path__) from ConfigMap data
+    configmap_files = {k: v for k, v in files.items() if not k.startswith("__")}
+
     # Validate and build ConfigMap spec
-    config_builder.validate_configmap_size(files)
+    config_builder.validate_configmap_size(configmap_files)
     configmap_spec = config_builder.build_file_configmap_spec(
-        job_name=job.name, namespace=job.namespace, files=files
+        job_name=job.name, namespace=job.namespace, files=configmap_files
     )
 
     # Create ConfigMap with owner reference

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

@@ -491,6 +491,38 @@ def test_rayjob_with_runtime_env(auto_mock_setup):
     assert rayjob_cr["spec"]["runtimeEnvYAML"] == "pip:\n- numpy\n- pandas\n"
 
 
+def test_rayjob_with_runtime_env_dict(auto_mock_setup):
+    """
+    Test RayJob with runtime environment as dict (user convenience).
+    Users can pass a dict instead of having to import RuntimeEnv.
+    """
+    # User can pass dict instead of RuntimeEnv object
+    runtime_env_dict = {
+        "pip": ["numpy", "pandas"],
+        "env_vars": {"TEST_VAR": "test_value"},
+    }
+
+    rayjob = RayJob(
+        job_name="test-job",
+        entrypoint="python test.py",
+        cluster_name="test-cluster",
+        runtime_env=runtime_env_dict,
+        namespace="test-namespace",
+    )
+
+    # Should be converted to RuntimeEnv internally
+    assert isinstance(rayjob.runtime_env, RuntimeEnv)
+    assert rayjob.runtime_env["env_vars"] == {"TEST_VAR": "test_value"}
+
+    # Verify it generates proper YAML output
+    rayjob_cr = rayjob._build_rayjob_cr()
+    assert "runtimeEnvYAML" in rayjob_cr["spec"]
+    runtime_yaml = rayjob_cr["spec"]["runtimeEnvYAML"]
+    assert "pip:" in runtime_yaml or "pip_packages:" in runtime_yaml
+    assert "env_vars:" in runtime_yaml
+    assert "TEST_VAR" in runtime_yaml
+
+
 def test_rayjob_with_active_deadline_and_ttl(auto_mock_setup):
     """
     Test RayJob with both active deadline and TTL settings.

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

@@ -431,9 +431,13 @@ def track_create_cm(*args, **kwargs):
     mock_create_cm.assert_called_once()
 
     # Verify create_file_configmap was called with: (job, files, rayjob_result)
-    mock_create_cm.assert_called_with(
-        rayjob, {"test.py": "print('test')"}, submit_result
-    )
+    # Files dict includes metadata key __entrypoint_path__ for single file case
+    call_args = mock_create_cm.call_args[0]
+    assert call_args[0] == rayjob
+    assert call_args[2] == submit_result
+    # Check that the actual file content is present
+    assert "test.py" in call_args[1]
+    assert call_args[1]["test.py"] == "print('test')"
 
 
 def test_rayjob_submit_with_files_new_cluster(auto_mock_setup, tmp_path):