coiled 1.128.1.dev2__tar.gz → 1.128.2__tar.gz

{coiled-1.128.1.dev2 → coiled-1.128.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coiled
-Version: 1.128.1.dev2
+Version: 1.128.2
 Summary: Python client for coiled.io dask clusters
 Project-URL: Homepage, https://coiled.io
 Maintainer-email: Coiled <info@coiled.io>

{coiled-1.128.1.dev2 → coiled-1.128.2}/coiled/cli/mpi.py

@@ -1,4 +1,5 @@
 import os.path
+import pathlib
 import shlex
 
 import click
@@ -9,7 +10,7 @@ import coiled
 from .. import AWSOptions
 from ..v2.cluster_comms import get_cluster_connection_info, get_comm_from_connection_info
 from .cluster.utils import find_cluster
-from .run import KeepaliveSession, get_ssh_connection, write_via_ssh
+from .run import KeepaliveSession, get_ssh_connection, upload_file
 from .utils import CONTEXT_SETTINGS
 
 
@@ -99,34 +100,13 @@ def setup(
     default=True,
     type=bool,
 )
+@click.option("--upload", type=str, multiple=True)
 @click.argument("command", nargs=-1, required=True)
-def run(cluster, workspace, legate, include_head, command):
+def run(cluster, workspace, legate, include_head, upload, command):
     nodes = "$(cat workers | wc -w)"
 
     command = list(command)
 
-    files = {}
-    for i, c in enumerate(command):
-        if os.path.exists(c):
-            remote_path = f"/scratch/batch/{os.path.basename(c)}"
-            command[i] = remote_path
-            with open(c) as f:
-                content = f.read()
-            files[remote_path] = content
-
-    if legate:
-        # TODO make "--gpus 1 --sysmem 2000 --fbmem 20000" configurable
-        wrapped_command = f"""
-legate \
-  --gpus 1 --sysmem 2000 --fbmem 20000 \
-  --nodes {nodes} \
-  --launcher mpirun \
-  --launcher-extra ' --hostfile workers -x PATH ' \
-  {shlex.join(command)}
-"""
-    else:
-        wrapped_command = f"mpirun --hostfile workers -x PATH {shlex.join(command)}"
-
     with coiled.Cloud(workspace=workspace) as cloud:
         cluster_info = find_cluster(cloud, cluster)
         cluster_id = cluster_info["id"]
@@ -136,9 +116,11 @@ legate \
 
     setup_mpi_ssh(connection, include_scheduler=include_head)
 
-    if files:
-        worker_connections = []
+    has_implicit_file = any(os.path.exists(c) for c in command)
 
+    if has_implicit_file or upload:
+        # get SSH connections to each of the workers (using scheduler as jump server)
+        worker_connections = []
         for worker in cluster_info["workers"]:
             if (
                 not worker.get("instance")
@@ -154,10 +136,66 @@ legate \
                 )
             )
 
-        for path, content in files.items():
-            write_via_ssh(connection, content=content, path=path)
+        for idx, implicit_file in enumerate(command):
+            if os.path.exists(implicit_file) and os.path.isfile(implicit_file):
+                # this will preserve path structure relative to cwd
+                # so `coiled run python ./subdir/foo.py` will go to `/scratch/subdir/foo.py`
+                remote_path = upload_file(connection, implicit_file, remote_root="/scratch/batch")
+                print(f"Uploaded {implicit_file} to {remote_path}")
+
+                for conn in worker_connections:
+                    upload_file(conn, implicit_file, remote_root="/scratch/batch")
+
+                # adjust command to reference path on VM
+                command[idx] = remote_path
+
+        files_to_upload = []
+        for f in upload:
+            path = pathlib.Path(f)
+            if not path.exists():
+                raise FileNotFoundError(f"Cannot find specified file {f}")
+
+            if path.is_file():
+                files_to_upload.append({"f": path})
+            elif path.is_dir():
+                # for paths outside cwd, parent_dir is used as the root so that path structure from there is preserved
+                parent_dir = pathlib.Path(path).parent
+                for subfile in path.rglob("*"):
+                    if subfile.is_file():
+                        files_to_upload.append({"f": subfile, "specified_root": parent_dir})
+
+        if files_to_upload:
+            print(
+                f"Uploading {len(files_to_upload)} file{'s' if len(files_to_upload) > 1 else ''} "
+                "from local machine to cloud VM..."
+            )
+        for i, file_to_upload in enumerate(files_to_upload):
+            try:
+                mb_size = file_to_upload["f"].stat().st_size / 1_000_000
+                if mb_size > 1:
+                    print(f"  {file_to_upload['f']} is {mb_size:.2f} MB, this may be slow to upload")
+            except Exception:
+                pass
+            remote_path = upload_file(connection, remote_root="/scratch/batch", **file_to_upload)
+            print("    ", remote_path)
             for conn in worker_connections:
-                write_via_ssh(conn, content=content, path=path)  # , mode=0o555
+                upload_file(conn, remote_root="/scratch/batch", **file_to_upload)
+            if i and (i % 20 == 0 or i + 1 == len(files_to_upload)):
+                print(f"  {i + 1}/{len(files_to_upload)} files uploaded")
+
+    if legate:
+        # TODO make "--gpus 1 --sysmem 2000 --fbmem 20000" configurable
+        #         wrapped_command = f"""
+        wrapped_command = f"""
+legate \
+  --gpus 1 --sysmem 2000 --fbmem 20000 \
+  --nodes {nodes} \
+  --launcher mpirun \
+  --launcher-extra ' --hostfile workers -x PATH ' \
+  {shlex.join(command)}
+"""
+    else:
+        wrapped_command = f"mpirun --hostfile workers -x PATH {shlex.join(command)}"
 
     print(f"Running command:\n{wrapped_command}")
 

{coiled-1.128.1.dev2 → coiled-1.128.2}/coiled/cli/run.py

@@ -125,7 +125,7 @@ def write_files_into_container(connection, container_name: str, files: Dict[str,
         )
 
 
-def upload_file(connection, f, specified_root=None) -> str:
+def upload_file(connection: fabric.Connection, f: str, specified_root=None, remote_root="/scratch") -> str:
     cwd = os.path.abspath(os.path.curdir)
     base = os.path.basename(f)
     is_under_cwd = os.path.commonpath((os.path.abspath(f), cwd)) == cwd
@@ -151,16 +151,16 @@ def upload_file(connection, f, specified_root=None) -> str:
         # For example, if user specified `--file /absolute/subdir/`, then preserve path structure relative
         # to `/absolute/subdir/`, so `/absolute/subdir/foo/bar.txt` would go to `/scratch/subdir/foo/bar.txt`.
         specified_path_dir = os.path.dirname(os.path.relpath(f, relative_to))
-        remote_dir = f"/scratch/{specified_path_dir}/"
+        remote_dir = f"{remote_root}/{specified_path_dir}/"
         make_remote_dir(connection, remote_dir)
     else:
-        remote_dir = "/scratch/"
+        remote_dir = f"{remote_root}/"
 
     connection.put(f, remote_dir)
 
     # we want path on Linux VM, which might not match os.path.join run client-side, so join path manually
     # remote_dir should already end in "/"
-    return f"{remote_dir}{base}"
+    return f"{remote_dir}{base}".replace("//", "/")
 
 
 def run_via_ssh(