Use a low-resource pod to move data to Lepton

Running a low-resource pod to copy experiment data to the Lepton cluster
is more reliable and broadly compatible with various cluster types
versus the storage API.

Signed-Off-By: Robert Clark <[email protected]>

Author: roclark

nemo_run/core/execution/lepton.py

@@ -1,19 +1,24 @@
+import base64
 import logging
 import os
 import re
 import subprocess
+import tempfile
 import time
 from dataclasses import dataclass, field
+from datetime import datetime
 from pathlib import Path
-from typing import Any, Optional, Set, Type
+from typing import Any, List, Optional, Set, Type
 
 from invoke.context import Context
 from leptonai.api.v1.client import APIClient
 from leptonai.api.v1.types.affinity import LeptonResourceAffinity
-from leptonai.api.v1.types.common import Metadata
+from leptonai.api.v1.types.common import LeptonVisibility, Metadata
 from leptonai.api.v1.types.dedicated_node_group import DedicatedNodeGroup
 from leptonai.api.v1.types.deployment import (EnvVar, LeptonContainer,
-                                              LeptonDeployment, Mount)
+                                              LeptonDeployment,
+                                              LeptonDeploymentUserSpec, Mount,
+                                              ResourceRequirement)
 from leptonai.api.v1.types.job import (LeptonJob, LeptonJobState,
                                        LeptonJobUserSpec)
 from leptonai.api.v1.types.replica import Replica
@@ -25,8 +30,6 @@
 
 logger = logging.getLogger(__name__)
 
-LEPTON_STORAGE_DEPLOYMENT = "storage-rsync-by-lepton"
-
 
 @dataclass(kw_only=True)
 class LeptonExecutor(Executor):
@@ -78,54 +81,56 @@ def _replica_public_ip(self, name: str, client: APIClient) -> str | None:
             raise RuntimeError(f"Unable to find status of pod {name}.")
         return replicas[0].status.public_ip
 
-    def _rsync_password(self, deployment: LeptonDeployment) -> str:
-        """
-        Find the rsync password on the cluster.
-        """
-        for env in deployment.spec.envs:
-            if env.name == "PASSWORD":
-                return env.value
-        raise RuntimeError("Unable to find rsync password on cluster.")
+    def copy_directory_data_command(self, local_dir_path: str, dest_path: str) -> List:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tarball_path = os.path.join(temp_dir, "archive.tar.gz")
+            subprocess.run(f"tar -czf {tarball_path} -C {local_dir_path} .", shell=True, check=True)
+            with open(tarball_path, "rb") as file:
+                file_data = file.read()
+            encoded_data = base64.b64encode(file_data).decode("utf-8")
 
+            # Delete and recreate directory if it already exists, command to decode base64 data, save to a file, and extract inside the pod
+            cmd = f"rm -rf {dest_path} && mkdir -p {dest_path} && echo {encoded_data} | base64 -d > {dest_path}/archive.tar.gz && tar -xzf {dest_path}/archive.tar.gz -C {dest_path} && rm {dest_path}/archive.tar.gz"
+            full_command = ["sh", "-c", cmd]
+            return full_command
 
     def move_data(self, sleep: float = 10) -> None:
         """
-        Moves job directory into PVC and deletes the workload after completion.
-
-        Referenced from the LeptonAI Python SDK: https://github.com/leptonai/leptonai/blob/main/leptonai/cli/storage.py
+        Moves job directory into remote storage and deletes the workload after completion.
         """
         client = APIClient()
-        client.storage.create_dir(additional_path=self.lepton_job_dir)
-
-        # Lepton comes with a hidden storage deplyment specifically for rsync'ing
-        # data to storage with a deterministic name. This deployment can be used
-        # for transferring data to the storage.
-        storage_deployment = client.deployment.get(LEPTON_STORAGE_DEPLOYMENT)
-
-        if not storage_deployment:
-            raise RuntimeError("Rsync deployment not found on cluster. Ensure your cluster supports rsync.")
-        port = storage_deployment.spec.container.ports[0].host_port
-        ip = self._replica_public_ip(LEPTON_STORAGE_DEPLOYMENT, client)
-        workspace_id = client.get_workspace_id()
-
-        password = self._rsync_password(storage_deployment)
+        cmd = self.copy_directory_data_command(self.job_dir, self.lepton_job_dir)
+        node_group_id = self._node_group_id(client)
+        valid_node_ids = self._valid_node_ids(node_group_id, client)
 
-        command = (
-            f"rsync -a {self.job_dir}/ rsync://{workspace_id}@{ip}:{port}/volume{self.lepton_job_dir}"
+        spec = LeptonDeploymentUserSpec(
+            container=LeptonContainer(
+                image="busybox:1.37.0",  # Use a very low resource container
+                command=cmd,
+            ),
+            mounts=[Mount(path=mount["path"], mount_path=mount["mount_path"]) for mount in self.mounts],
+        )
+        spec.resource_requirement = ResourceRequirement(
+            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,
         )
+        custom_name = f"data-mover-{int(datetime.now().timestamp())}"
 
-        process = subprocess.Popen(
-            command,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            env={"RSYNC_PASSWORD": password},
-            shell=True,
-            universal_newlines=True
+        deployment = LeptonDeployment(
+            metadata=Metadata(
+                id=custom_name,
+                name=custom_name,
+                visibility=LeptonVisibility("private"),
+            ),
+            spec=spec,
         )
 
-        for line in process.stdout:
-            print(line, end="")
-        process.wait()
+        client.deployment.create(deployment)
 
     def _node_group_id(self, client: APIClient) -> DedicatedNodeGroup:
         """