fix: test

Author: pawelpaszki

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -54,6 +54,8 @@ def __init__(
         shutdown_after_job_finishes: Optional[bool] = None,
         ttl_seconds_after_finished: int = 0,
         active_deadline_seconds: Optional[int] = None,
+        entrypoint_num_cpus: Optional[int] = None,
+        entrypoint_num_gpus: Optional[int] = None,
     ):
         """
         Initialize a RayJob instance.
@@ -100,6 +102,8 @@ def __init__(
         self.runtime_env = runtime_env
         self.ttl_seconds_after_finished = ttl_seconds_after_finished
         self.active_deadline_seconds = active_deadline_seconds
+        self.entrypoint_num_cpus = entrypoint_num_cpus
+        self.entrypoint_num_gpus = entrypoint_num_gpus
 
         # 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
@@ -189,6 +193,16 @@ def _build_rayjob_cr(self) -> Dict[str, Any]:
         if self.active_deadline_seconds:
             rayjob_cr["spec"]["activeDeadlineSeconds"] = self.active_deadline_seconds
 
+        # Add entrypoint resource requirements if specified
+        entrypoint_resources = {}
+        if self.entrypoint_num_cpus is not None:
+            entrypoint_resources["cpu"] = str(self.entrypoint_num_cpus)
+        if self.entrypoint_num_gpus is not None:
+            entrypoint_resources["gpu"] = str(self.entrypoint_num_gpus)
+
+        if entrypoint_resources:
+            rayjob_cr["spec"]["entrypointResources"] = entrypoint_resources
+
         # Add runtime environment if specified
         if self.runtime_env:
             rayjob_cr["spec"]["runtimeEnvYAML"] = str(self.runtime_env)

tests/e2e/rayjob_existing_cluster_kind_test.py

@@ -96,7 +96,7 @@ def assert_rayjob_submit_against_existing_cluster(
                 "env_vars": get_setup_env_variables(ACCELERATOR=accelerator),
             },
             shutdown_after_job_finishes=False,
-            entrypoint_num_gpus=number_of_gpus,
+            entrypoint_num_gpus=number_of_gpus if number_of_gpus > 0 else None,
         )
 
         # Submit the job