Upload files with rsync

Uploading files with rsync allows the files to retain their existing
permissions. This means files that are executable locally will also be
executable on the remote storage.

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

Author: roclark

nemo_run/core/execution/lepton.py

@@ -12,7 +12,8 @@
 from leptonai.api.v1.types.affinity import LeptonResourceAffinity
 from leptonai.api.v1.types.common import Metadata
 from leptonai.api.v1.types.dedicated_node_group import DedicatedNodeGroup
-from leptonai.api.v1.types.deployment import EnvVar, LeptonContainer, Mount
+from leptonai.api.v1.types.deployment import (EnvVar, LeptonContainer,
+                                              LeptonDeployment, Mount)
 from leptonai.api.v1.types.job import (LeptonJob, LeptonJobState,
                                        LeptonJobUserSpec)
 from leptonai.api.v1.types.replica import Replica
@@ -24,6 +25,8 @@
 
 logger = logging.getLogger(__name__)
 
+LEPTON_STORAGE_DEPLOYMENT = "storage-rsync-by-lepton"
+
 
 @dataclass(kw_only=True)
 class LeptonExecutor(Executor):
@@ -63,29 +66,66 @@ def stop_job(self, job_id: str):
         client.job.update(job_id, spec={"spec": {"stopped": True}})
         logger.info(f"Job {job_id} stopped successfully.")
 
+    def _replica_public_ip(self, name: str, client: APIClient) -> str | None:
+        """
+        Find the public IP address for a given replica.
+        """
+        replicas = client.deployment.get_replicas(name)
+
+        if len(replicas) != 1:
+            raise RuntimeError(f"Pod {name} has more than one replica. This is not supported.")
+        if not replicas[0].status:
+            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 move_data(self, sleep: float = 10) -> None:
         """
-        Moves job directory into PVC and deletes the workload after completion
+        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
         """
         client = APIClient()
         client.storage.create_dir(additional_path=self.lepton_job_dir)
 
-        # Create all sub-directories in the directory tree
-        # Then, copy all files to the storage
-        for root, dirs, files in os.walk(self.job_dir):
-            # Create the sub-directories
-            for dir in dirs:
-                abs_path = os.path.join(root, dir)
-                relative_path = os.path.join(self.lepton_job_dir, abs_path.replace(self.job_dir, "").lstrip("/"))
-                client.storage.create_dir(additional_path=relative_path)
-            # Copy the files in each sub-directory to the remote filesystem
-            for file in files:
-                abs_path = os.path.join(root, file)
-                relative_path = os.path.join(self.lepton_job_dir, abs_path.replace(self.job_dir, "").lstrip("/"))
-                client.storage.create_file(
-                    local_path=abs_path,
-                    remote_path=relative_path
-                )
+        # 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)
+
+        command = (
+            f"rsync -a {self.job_dir}/ rsync://{workspace_id}@{ip}:{port}/volume{self.lepton_job_dir}"
+        )
+
+        process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            env={"RSYNC_PASSWORD": password},
+            shell=True,
+            universal_newlines=True
+        )
+
+        for line in process.stdout:
+            print(line, end="")
+        process.wait()
 
     def _node_group_id(self, client: APIClient) -> DedicatedNodeGroup:
         """
@@ -127,7 +167,7 @@ def create_lepton_job(self, name: str):
         cmd = [
             "/bin/bash",
             "-c",
-            f"chmod +x {self.lepton_job_dir}/launch_script.sh && bash {self.lepton_job_dir}/launch_script.sh"
+            f"bash {self.lepton_job_dir}/launch_script.sh"
         ]
 
         # Get ID of requested node group