Add Lepton RayCluster tests

Signed-Off-By: Robert Clark <roclark@nvidia.com>

Author: roclark

nemo_run/core/execution/lepton.py

@@ -80,9 +80,9 @@ class LeptonExecutor(Executor):
         default_factory=list
     )  # Image pull secrets for container registry authentication
     custom_spec: dict[str, Any] = field(default_factory=dict)
-    pre_launch_commands: list[str] = field(default_factory=list)
+    pre_launch_commands: list[str] = field(default_factory=list)  # Custom commands before launch
     head_resource_shape: Optional[str] = ""  # Only used for LeptonRayCluster
-    ray_version: Optional[str] = None  # Only used for LeptonRayCluster  # Custom commands before launch
+    ray_version: Optional[str] = None  # Only used for LeptonRayCluster
 
     def stop_job(self, job_id: str):
         """

nemo_run/run/ray/cluster.py

@@ -40,8 +40,6 @@ class RayCluster:
     name: str
     executor: Executor
     log_level: str = "INFO"
-    ray_version: Optional[str] = None  # Only used for LeptonRayCluster
-    head_resource_shape: str = None  # Only used for LeptonRayCluster
 
     def __post_init__(self):
         configure_logging(level=self.log_level)
@@ -61,10 +59,6 @@ def __post_init__(self):
 
         self._port_forward_map = {}
 
-        if isinstance(self.backend, LeptonRayCluster):
-            self.backend.ray_version = self.ray_version
-            self.backend.head_resource_shape = self.head_resource_shape
-
     def start(
         self,
         wait_until_ready: bool = True,

nemo_run/run/ray/job.py

@@ -44,7 +44,6 @@ class RayJob:
     pre_ray_start_commands: Optional[list[str]] = None
     log_level: str = "INFO"
     cluster_name: Optional[str] = None  # Used to connect to existing RayCluster
-    ray_version: Optional[str] = None  # Only used for LeptonRayJob
     cluster_ready_timeout: Optional[int] = 1800  # Only used for LeptonRayJob
 
     def __post_init__(self) -> None:  # noqa: D401 – simple implementation
@@ -65,7 +64,6 @@ def __post_init__(self) -> None:  # noqa: D401 – simple implementation
 
         if isinstance(self.executor, LeptonExecutor):
             self.backend.cluster_name = self.cluster_name
-            self.backend.ray_version = self.ray_version
             self.backend.cluster_ready_timeout = self.cluster_ready_timeout
 
     # ------------------------------------------------------------------

nemo_run/run/ray/lepton.py

@@ -22,7 +22,7 @@
 import warnings
 from dataclasses import dataclass
 from ray.job_submission import JobSubmissionClient
-from rich.pretty import Pretty
+from rich.pretty import pretty_repr
 from typing import Any, Optional, TypeAlias
 
 from leptonai.api.v1.types.affinity import LeptonResourceAffinity
@@ -134,7 +134,7 @@ def status(
 
         if display and status_dict:
             cluster = client.raycluster.get(self.name)
-            logger.info(Pretty(cluster))
+            logger.info(pretty_repr(vars(cluster), expand_all=True))
 
         return status_dict
 
@@ -176,6 +176,7 @@ def create(
         if len(name) > 35:
             logger.warning("length of name exceeds 35 characters. Shortening...")
             name = name[:34]
+            self.name = name
 
         executor = self.executor
         client = APIClient()
@@ -305,15 +306,15 @@ def delete(
         Returns
         -------
         bool
-            *True* if the job was successfully cancelled (or already gone), *False* otherwise.
+            *True* if the cluster was confirmed deleted (or already gone), *False* otherwise.
         """
         name = self.name
         logger.debug(f"Deleting RayCluster '{name}'")
 
         client = APIClient()
         status = self._status(client)
 
-        if status["cluster_name"] is None:
+        if status is None or status["cluster_name"] is None:
             logger.warning(f"RayCluster '{name}' does not exist or is already deleted")
             return True
 
@@ -331,6 +332,10 @@ def delete(
 
                 if not status or status["cluster_name"] is None:
                     logger.info(f"RayCluster '{name}' successfully deleted")
+
+                    if name in self.cluster_map:
+                        del self.cluster_map[name]
+
                     return True
 
                 logger.debug(f"RayCluster '{name}' is still being deleted, waiting...")
@@ -443,8 +448,6 @@ def _ray_client(self, create_if_not_exists: bool = False) -> APIClient:
                 cluster = LeptonRayCluster(
                     name=name,
                     executor=self.executor,
-                    ray_version=self.executor.ray_version,
-                    head_resource_shape=self.executor.head_resource_shape,
                 )
                 cluster.create()
                 logger.info(f"Waiting for RayCluster '{name}' to be ready...")