refactor DGXC Lepton data mover: switch to BatchJob with auto cleanup and sleep after every run (#265)

* refactor DGXC Lepton data mover: switch to BatchJob with auto cleanup and sleep after every run

Signed-off-by: prekshivyas <[email protected]>

Author: prekshivyas

nemo_run/core/execution/lepton.py

@@ -13,15 +13,12 @@
 from invoke.context import Context
 from leptonai.api.v2.client import APIClient
 from leptonai.api.v1.types.affinity import LeptonResourceAffinity
-from leptonai.api.v1.types.common import LeptonVisibility, Metadata
+from leptonai.api.v1.types.common import Metadata, LeptonVisibility
 from leptonai.api.v1.types.dedicated_node_group import DedicatedNodeGroup
 from leptonai.api.v1.types.deployment import (
     EnvVar,
     LeptonContainer,
-    LeptonDeployment,
-    LeptonDeploymentUserSpec,
     Mount,
-    ResourceRequirement,
 )
 from leptonai.api.v1.types.job import LeptonJob, LeptonJobState, LeptonJobUserSpec
 from leptonai.api.v1.types.replica import Replica
@@ -85,7 +82,7 @@ def copy_directory_data_command(self, local_dir_path: str, dest_path: str) -> Li
             full_command = ["sh", "-c", cmd]
             return full_command
 
-    def move_data(self, sleep: float = 10) -> None:
+    def move_data(self, sleep: float = 10, timeout: int = 600, poll_interval: int = 5) -> None:
         """
         Moves job directory into remote storage and deletes the workload after completion.
         """
@@ -94,34 +91,60 @@ def move_data(self, sleep: float = 10) -> None:
         node_group_id = self._node_group_id(client)
         valid_node_ids = self._valid_node_ids(node_group_id, client)
 
-        spec = LeptonDeploymentUserSpec(
-            container=LeptonContainer(
-                image="busybox:1.37.0",  # Use a very low resource container
-                command=cmd,
-            ),
-            mounts=[Mount(**mount) for mount in self.mounts],
-        )
-        spec.resource_requirement = ResourceRequirement(
+        job_spec = LeptonJobUserSpec(
             resource_shape="cpu.small",
             affinity=LeptonResourceAffinity(
                 allowed_dedicated_node_groups=[node_group_id.metadata.id_],
                 allowed_nodes_in_node_group=valid_node_ids,
             ),
-            min_replicas=1,
-            max_replicas=1,
+            container=LeptonContainer(
+                image="busybox:1.37.0",
+                command=cmd,
+            ),
+            completions=1,
+            parallelism=1,
+            mounts=[Mount(**mount) for mount in self.mounts],
         )
+
         custom_name = f"data-mover-{int(datetime.now().timestamp())}"
 
-        deployment = LeptonDeployment(
+        job = LeptonJob(
             metadata=Metadata(
                 id=custom_name,
                 name=custom_name,
                 visibility=LeptonVisibility("private"),
             ),
-            spec=spec,
+            spec=job_spec,
         )
 
-        client.deployment.create(deployment)
+        response = client.job.create(job)
+        job_id = response.metadata.id_
+
+        start_time = time.time()
+        count = 0
+
+        while True:
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds.")
+            current_job = client.job.get(job_id)
+            current_job_status = current_job.status.state
+            if count > 0 and current_job_status in [
+                LeptonJobState.Completed,
+                LeptonJobState.Failed,
+                LeptonJobState.Unknown,
+            ]:
+                break
+            count += 1
+            time.sleep(poll_interval)
+
+        if current_job_status != LeptonJobState.Completed:
+            raise RuntimeError(f"Job {job_id} failed with status: {current_job_status}")
+
+        delete_success = client.job.delete(job_id)
+        if delete_success:
+            logging.info(f"Successfully deleted job {job_id}")
+        else:
+            logging.error(f"Failed to delete job {job_id}")
 
     def _node_group_id(self, client: APIClient) -> DedicatedNodeGroup:
         """

test/core/execution/test_lepton.py

@@ -22,13 +22,10 @@
 from leptonai.api.v1.types.common import LeptonVisibility, Metadata
 from leptonai.api.v1.types.deployment import (
     LeptonContainer,
-    LeptonDeployment,
-    LeptonDeploymentUserSpec,
     LeptonResourceAffinity,
     Mount,
-    ResourceRequirement,
 )
-from leptonai.api.v1.types.job import LeptonJob
+from leptonai.api.v1.types.job import LeptonJob, LeptonJobUserSpec
 
 from nemo_run.core.execution.lepton import LeptonExecutor, LeptonJobState
 from nemo_run.core.packaging.git import GitArchivePackager
@@ -173,26 +170,23 @@ def test_copy_directory_data_command_fails(self, mock_tempdir):
     @patch("nemo_run.core.execution.lepton.APIClient")
     def test_move_data_success(self, mock_APIClient, mock_datetime, mock_copy):
         mock_instance = MagicMock()
-        mock_deployment_api = MagicMock()
-        mock_instance.deployment = mock_deployment_api
-        mock_deployment_api.create = MagicMock()
+        mock_job_api = MagicMock()
+        mock_instance.job = mock_job_api
         mock_copy.return_value = ["sh", "-c", "echo 'hello world'"]
         mock_APIClient.return_value = mock_instance
         mock_client = mock_APIClient.return_value
         mock_nodegroup = MagicMock()
-        mock_datetime = MagicMock()
         mock_datetime_now = MagicMock()
-        mock_datetime.now = mock_datetime_now
-        mock_timestamp = MagicMock()
-        mock_datetime_now.timestamp = MagicMock()
-        mock_timestamp.return_value = "1"
+        mock_datetime.now.return_value = mock_datetime_now
+        mock_datetime_now.timestamp.return_value = 1
         mock_client.nodegroup = mock_nodegroup
         mock_nodegroup.list_all.return_value = [
             SimpleNamespace(metadata=SimpleNamespace(name="123456", id_="my-node-id"))
         ]
         mock_nodegroup.list_nodes.return_value = [
             SimpleNamespace(metadata=SimpleNamespace(id_="10-10-10-10"))
         ]
+        mock_job_api.get.return_value = SimpleNamespace(status=SimpleNamespace(state="Completed"))
 
         executor = LeptonExecutor(
             container_image="nvcr.io/nvidia/test:latest",
@@ -203,32 +197,35 @@ def test_move_data_success(self, mock_APIClient, mock_datetime, mock_copy):
 
         executor.move_data()
 
-        spec = LeptonDeploymentUserSpec(
-            container=LeptonContainer(
-                image="busybox:1.37.0",
-                command=["sh", "-c", "echo 'hello world'"],
-            ),
-            mounts=[Mount(path="/workspace", mount_path="/workspace")],
-        )
-        spec.resource_requirement = ResourceRequirement(
+        expected_cmd = ["sh", "-c", "echo 'hello world'"]
+        expected_spec = LeptonJobUserSpec(
             resource_shape="cpu.small",
             affinity=LeptonResourceAffinity(
                 allowed_dedicated_node_groups=["my-node-id"],
                 allowed_nodes_in_node_group=["10-10-10-10"],
             ),
-            min_replicas=1,
-            max_replicas=1,
+            container=LeptonContainer(
+                image="busybox:1.37.0",
+                command=expected_cmd,
+            ),
+            completions=1,
+            parallelism=1,
+            mounts=[Mount(path="/workspace", mount_path="/workspace")],
         )
+
         custom_name = "data-mover-1"
-        deployment = LeptonDeployment(
+        expected_job = LeptonJob(
             metadata=Metadata(
-                id=custom_name, name=custom_name, visibility=LeptonVisibility("private")
+                id=custom_name,
+                name=custom_name,
+                visibility=LeptonVisibility("private"),
             ),
-            spec=spec,
+            spec=expected_spec,
         )
 
-        mock_copy.assert_called_once_with("", "")
-        mock_deployment_api.create.assert_called_once_with(deployment)
+        mock_copy.assert_called_once_with(executor.job_dir, executor.lepton_job_dir)
+        mock_job_api.create.assert_called_once_with(expected_job)
+        mock_job_api.delete.assert_called_once_with(mock_job_api.create.return_value.metadata.id_)
 
     def test_node_group_id(self):
         mock_client = MagicMock(