coiled 1.118.4.dev6__py3-none-any.whl → 1.129.3.dev10__py3-none-any.whl

coiled/cli/curl.py

@@ -1,3 +1,4 @@
+from json import dumps as json_dumps
 from json import loads as json_loads
 
 import click
@@ -24,16 +25,20 @@ def curl(url: str, request, data, json, json_output):
             url = f"{cloud.server}{url}"
         response = sync_request(cloud, url, method=request, data=all_data, json=json, json_output=json_output)
 
-    print(response)
+    if json_output:
+        print(json_dumps(response, indent=4))
+    else:
+        print(response)
 
 
-def sync_request(cloud, url, method, data, json: bool = False, json_output: bool = False):
+def sync_request(cloud, url, method, data=None, json: bool = False, json_output: bool = False):
     kwargs = {"method": method, "url": url}
 
     if json:
         kwargs["json"] = json_loads(data) if isinstance(data, str) else data
     else:
         kwargs["data"] = data
+
     response = cloud._sync(cloud._do_request, **kwargs)
     if response.status >= 400:
         print(f"{url} returned {response.status}")

coiled/cli/file.py

@@ -0,0 +1,116 @@
+import click
+
+import coiled
+from coiled.filestore import FilestoreManager
+
+from .cluster.utils import find_cluster
+from .utils import CONTEXT_SETTINGS
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.argument("cluster", default="", required=False)
+@click.option(
+    "--workspace",
+    default=None,
+    help="Coiled workspace (uses default workspace if not specified).",
+)
+@click.option(
+    "--filestore",
+    default=None,
+    help="Name of filestore (optional).",
+)
+@click.option(
+    "--filter",
+    "name_includes",
+    default=None,
+    help="Filter on file paths and/or names to download (optional).",
+)
+@click.option("--into", default=".")
+def download(cluster, workspace, filestore, name_includes, into):
+    if filestore:
+        filestores = FilestoreManager.get_filestore(name=filestore) or []
+        if not filestores:
+            print(f"{filestore} filestore not found")
+
+        for fs in filestores:
+            coiled.filestore.download_from_filestore_with_ui(
+                fs=fs,
+                into=into,
+                name_includes=name_includes,
+            )
+
+    else:
+        with coiled.Cloud(workspace=workspace) as cloud:
+            cluster_info = find_cluster(cloud, cluster)
+            cluster_id = cluster_info["id"]
+            attachments = FilestoreManager.get_cluster_attachments(cluster_id)
+        if not attachments:
+            print(f"No filestore found for {cluster_info['name']} ({cluster_info['id']})")
+
+        # TODO (possible enhancement) if there are multiple output filestores, let user pick which to download
+        for attachment in attachments:
+            if attachment["output"]:
+                coiled.filestore.download_from_filestore_with_ui(
+                    fs=attachment["filestore"],
+                    into=into,
+                    name_includes=name_includes,
+                )
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.argument("cluster", default="", required=False)
+@click.option(
+    "--workspace",
+    default=None,
+    help="Coiled workspace (uses default workspace if not specified).",
+)
+@click.option(
+    "--filestore",
+    default=None,
+    help="Name of filestore (optional).",
+)
+@click.option(
+    "--filter",
+    "name_includes",
+    default=None,
+    help="Filter on file paths and/or names to download (optional).",
+)
+def list_files(cluster, workspace, filestore, name_includes):
+    if filestore:
+        filestores = FilestoreManager.get_filestore(name=filestore) or []
+        if not filestores:
+            print(f"{filestore} filestore not found")
+
+        for fs in filestores:
+            coiled.filestore.list_files_ui(
+                fs=fs,
+                name_includes=name_includes,
+            )
+
+    else:
+        with coiled.Cloud(workspace=workspace) as cloud:
+            cluster_info = find_cluster(cloud, cluster)
+            cluster_id = cluster_info["id"]
+            attachments = FilestoreManager.get_cluster_attachments(cluster_id)
+        if not attachments:
+            print(f"No filestore found for {cluster_info['name']} ({cluster_info['id']})")
+
+        # TODO (possible enhancement) if there are multiple output filestores, let user pick which to download
+        for attachment in attachments:
+            if attachment["output"]:
+                coiled.filestore.list_files_ui(
+                    fs=attachment["filestore"],
+                    name_includes=name_includes,
+                )
+
+
+@click.group(name="file", context_settings=CONTEXT_SETTINGS)
+def file_group(): ...
+
+
+file_group.add_command(download)
+file_group.add_command(list_files, "list")

coiled/cli/hello/hello.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import json
-import subprocess
 import sys
 import time
 
@@ -18,6 +16,7 @@ import coiled
 from coiled.scan import scan_prefix
 from coiled.utils import login_if_required
 
+from ..curl import sync_request
 from .examples import examples
 from .utils import PRIMARY_COLOR, Panel, console, has_macos_system_python, log_interactions
 
@@ -41,8 +40,10 @@ def needs_login():
 
 
 def get_interactions():
-    response = subprocess.check_output(["coiled", "curl", "/api/v2/interactions/user-interactions/hello"])
-    return json.loads(response, strict=False)
+    with coiled.Cloud() as cloud:
+        return sync_request(
+            cloud, url=f"{cloud.server}/api/v2/interactions/user-interactions/hello", method="get", json_output=True
+        )
 
 
 def get_already_run_examples():
@@ -294,7 +295,7 @@ Choose any computation you'd like to run:
 Yee-haw you've done all my examples 🎉  
 Now you can:
 - Try Coiled in your own use case
-- [Talk to us](https://calendly.com/d/cmph-386-cjt/coiled-help-getting-started)
+- [Ask us questions](mailto:support@coiled.io)
 - Explore the [docs](https://docs.coiled.io?utm_source=coiled-hello&utm_medium=finished) to see all the other things Coiled can do
 """),  # noqa
                     border_style="green",

coiled/cli/mpi.py

@@ -0,0 +1,252 @@
+import os.path
+import pathlib
+import shlex
+
+import click
+import fabric.connection
+
+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, upload_file
+from .utils import CONTEXT_SETTINGS
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.option("--worker-nodes", default=1, type=int)
+@click.option(
+    "--include-head/--exclude-head",
+    default=True,
+    type=bool,
+    help="Include head (scheduler) node in placement group with EFA",
+)
+@click.option("--vm-type", default="g6.8xlarge", type=str)
+@click.option("--head-vm-type", default=None, type=str)
+@click.option("--worker-vm-type", default=None, type=str)
+@click.option("--pip", multiple=True, type=str, help="Packages to install with pip")
+@click.option("--apt", multiple=True, type=str, help="Packages to install with apt install")
+@click.option("--setup-script", default=None, type=str, help="Path to additional host setup script")
+@click.option("--cluster-timeout", default="1h", type=str, help="Maximum lifetime for cluster")
+@click.option(
+    "--idle-timeout",
+    default="10m",
+    type=str,
+    help="How long to wait when nothing is running before shutting down cluster",
+)
+def setup(
+    worker_nodes,
+    include_head,
+    vm_type,
+    head_vm_type,
+    worker_vm_type,
+    pip,
+    apt,
+    setup_script,
+    cluster_timeout,
+    idle_timeout,
+):
+    additional_host_setup = None
+    if setup_script:
+        with open(setup_script) as f:
+            additional_host_setup = f.read()
+    setup_script = get_host_setup_script(pip_install=pip, apt_install=apt, additional_setup=additional_host_setup)
+
+    efa_settings = (
+        {"use_placement_group": True, "use_efa": True}
+        if include_head
+        else {"use_worker_placement_group": True, "use_worker_efa": True}
+    )
+
+    cluster = coiled.Cluster(
+        n_workers=worker_nodes,
+        container="daskdev/dask:latest",
+        allow_ssh_from="me",
+        host_setup_script=setup_script,
+        backend_options=AWSOptions(**{**efa_settings, "ami_version": "DL"}),
+        scheduler_vm_types=[head_vm_type or vm_type],
+        worker_vm_types=[worker_vm_type or vm_type],
+        worker_disk_size="100GB",
+        scheduler_disk_size="100GB",
+        shutdown_on_close=False,
+        idle_timeout="520 weeks",  # don't use idle timeout
+        cluster_timeout=cluster_timeout,
+    )
+
+    cluster.set_keepalive(keepalive=idle_timeout)
+
+    print("Cluster created, installing software for MPI...")
+
+    with coiled.Cloud() as cloud:
+        connection = get_ssh_connection(cloud, cluster.cluster_id)
+
+    with KeepaliveSession(cluster=cluster, prefix="mpi-ssh-setup"):
+        setup_mpi_ssh(connection)
+
+    print("MPI is ready")
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.option("--cluster", default=None)
+@click.option("--workspace", default=None, type=str)
+@click.option("--legate", is_flag=True, default=False, type=bool)
+@click.option(
+    "--include-head/--exclude-head",
+    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, upload, command):
+    nodes = "$(cat workers | wc -w)"
+
+    command = list(command)
+
+    with coiled.Cloud(workspace=workspace) as cloud:
+        cluster_info = find_cluster(cloud, cluster)
+        cluster_id = cluster_info["id"]
+        connection = get_ssh_connection(cloud, cluster_id)
+        address, security = get_cluster_connection_info(cluster_id, cloud)
+        comms = get_comm_from_connection_info(address, security)
+
+    setup_mpi_ssh(connection, include_scheduler=include_head)
+
+    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")
+                or not worker["instance"].get("current_state")
+                or worker["instance"]["current_state"]["state"] != "ready"
+            ):
+                continue
+            worker_address = worker["instance"]["private_ip_address"]
+
+            worker_connections.append(
+                fabric.connection.Connection(
+                    worker_address, gateway=connection, user=connection.user, connect_kwargs=connection.connect_kwargs
+                )
+            )
+
+        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:
+                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}")
+
+    with KeepaliveSession(cluster=None, comms=comms, prefix="mpi-ssh"):
+        connection.run(wrapped_command, hide=False, pty=True, warn=True, env={"PATH": "/tmp/host-user-venv/bin:$PATH"})
+
+    comms.close_comms()
+
+
+def setup_mpi_ssh(connection, include_scheduler=True):
+    add_scheduler_line = 'printf "\n127.0.0.1" >> workers' if include_scheduler else ""
+
+    setup_mpi = f"""
+/bin/coiled_agent list-worker-ips | sudo tee workers && sudo chown ubuntu workers
+ssh-keyscan -f workers -t ed25519 >> ~/.ssh/known_hosts
+{add_scheduler_line}
+
+# block until host setup script has finished, at least on schedule node
+until [ -f /tmp/host-setup-done ]
+do
+     sleep 5
+done
+"""
+
+    _ = connection.run(setup_mpi, hide=True, pty=False)
+
+
+def get_host_setup_script(venv_path="/tmp/host-user-venv", apt_install=None, pip_install=None, additional_setup=None):
+    apt_install = list(apt_install or [])
+    apt_install.extend(["openmpi-bin", "python3-pip", "python3-venv"])
+
+    pip_install = list(pip_install or [])
+
+    pip_install_line = f"{venv_path}/bin/python -m pip install {' '.join(pip_install)}" if pip_install else ""
+
+    return f"""
+sudo apt install {" ".join(apt_install)} -y
+
+mkdir {venv_path}
+python3 -m venv {venv_path}
+
+{pip_install_line}
+{additional_setup or ""}
+
+echo 'done' > /tmp/host-setup-done
+    """
+
+
+@click.group(name="mpi", context_settings=CONTEXT_SETTINGS)
+def mpi_group(): ...
+
+
+mpi_group.add_command(setup)
+mpi_group.add_command(run)

coiled/cli/notebook/notebook.py

@@ -237,6 +237,12 @@ def check_jupyter() -> bool:
     ),
     hidden=True,
 )
+@click.option(
+    "--private",
+    default=False,
+    is_flag=True,
+    help="Make this notebook private to you (other workspace members cannot access it).",
+)
 def start_notebook(
     name: str | None,
     account: str | None,
@@ -262,6 +268,7 @@ def start_notebook(
     mount_bucket: List[str] | None,
     host_setup_script: str | None,
     resumable: bool = False,
+    private: bool = False,
 ):
     """
     Launch or re-open a notebook session, with optional file syncing.
@@ -299,6 +306,7 @@ def start_notebook(
         mount_bucket=mount_bucket,
         host_setup_script=host_setup_script,
         resumable=resumable,
+        private=private,
     )
 
 
@@ -329,6 +337,7 @@ def _start_notebook(
     mount_bucket: List[str] | None = None,
     host_setup_script: str | None = None,
     resumable: bool = False,
+    private: bool = False,
 ) -> coiled.Cluster | None:
     """
     Launch or re-open a notebook session, with optional file syncing.
@@ -418,6 +427,7 @@ def _start_notebook(
                 shutdown_on_close=False,
                 unset_single_threading_variables=True,
                 pause_on_exit=resumable,
+                private_to_creator=private,
             )
             info["cluster_id"] = cluster.cluster_id
 

coiled/cli/run.py

@@ -29,8 +29,11 @@ from coiled.utils import (
     unset_single_thread_defaults,
 )
 from coiled.v2.cluster import ClusterKwargs
+from coiled.v2.cluster_comms import use_comm_rpc
+from coiled.v2.core import Cloud
 from coiled.v2.widgets.rich import LightRichClusterWidget
 
+from ..filestore import FilestoreManager
 from .sync import SYNC_TARGET, start_sync, stop_sync
 from .utils import CONTEXT_SETTINGS
 
@@ -45,20 +48,37 @@ USER_CONTAINER_NAME = "tmp-user-1"
 
 
 class KeepaliveSession:
-    def __init__(self, cluster, prefix="", monitor_proc_activity=False):
+    def __init__(self, cluster, comms=None, prefix="", monitor_proc_activity=False):
         self.cluster = cluster
+        self.comms = comms
         self.monitor_proc_activity = monitor_proc_activity
         rand_uuid = short_random_string()
         self.session_id = f"{prefix}-{rand_uuid}" if prefix else rand_uuid
 
+        if self.comms:
+            self.cloud = Cloud.current(asynchronous=False)
+
     def __enter__(self):
         # keepalive session lets us use keepalive without dask client
-        self.cluster._call_scheduler_comm(
-            "coiled_add_keepalive_session", name=self.session_id, monitor_proc_activity=self.monitor_proc_activity
-        )
+        if self.cluster:
+            self.cluster._call_scheduler_comm(
+                "coiled_add_keepalive_session", name=self.session_id, monitor_proc_activity=self.monitor_proc_activity
+            )
+        elif self.comms:
+            use_comm_rpc(
+                self.cloud,
+                self.comms,
+                "coiled_add_keepalive_session",
+                name=self.session_id,
+                monitor_proc_activity=self.monitor_proc_activity,
+            )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.cluster._call_scheduler_comm("coiled_end_keepalive_session", name=self.session_id)
+        if self.cluster:
+            self.cluster._call_scheduler_comm("coiled_end_keepalive_session", name=self.session_id)
+        elif self.comms:
+            use_comm_rpc(self.cloud, self.comms, "coiled_end_keepalive_session", name=self.session_id)
+            self.cloud.close()
 
 
 def get_ssh_connection(cloud, cluster_id) -> fabric.connection.Connection:
@@ -105,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
@@ -131,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(
@@ -238,7 +258,8 @@ def run_via_ssh(
 
     if container and "/uv:" in container and command_string.startswith("uv"):
         command_string = (
-            "(apt update && apt upgrade && apt install -y --no-install-recommends ca-certificates) 2>&1 > /dev/null\n"
+            "(apt update -y && apt upgrade -y && apt install -y --no-install-recommends ca-certificates) "
+            "2>&1 > /dev/null\n"
             f"{command_string}"
         )
 
@@ -549,6 +570,13 @@ def get_entrypoint(connection, container_name) -> str:
     default=None,
     help="Non-default value for shm_size (for example, '3 GiB').",
 )
+@click.option(
+    "--filestore",
+    "filestore_names",
+    default=None,
+    multiple=True,
+    help="Name of filestore to attach; can be specified multiple times for multiple filestores.",
+)
 @click.argument("command", nargs=-1)
 def run(
     name: str | None,
@@ -579,6 +607,7 @@ def run(
     package_sync_strict,
     package_sync_conda_extras,
     docker_shm_size,
+    filestore_names,
     command,
 ):
     """
@@ -614,6 +643,7 @@ def run(
         package_sync_strict=package_sync_strict,
         package_sync_conda_extras=package_sync_conda_extras,
         docker_shm_size=docker_shm_size,
+        filestore_names=filestore_names,
     )
     sys.exit(info["exit_code"])
 
@@ -651,6 +681,7 @@ def start_run(
     package_sync_strict: bool = False,
     package_sync_conda_extras: List[str] | None = None,
     docker_shm_size: str | None = None,
+    filestore_names: list[str] | None = None,
 ):
     runtime_env_dict = dict_from_key_val_list(env)
     tags = dict_from_key_val_list(tag)
@@ -693,9 +724,20 @@ def start_run(
     # fail early if user specified `--file` that doesn't exist
     check_explicit_files(file)
 
+    filestores_to_attach = []
+
     try:
         with coiled.Cloud(workspace=workspace or account) as cloud:
             workspace = workspace or cloud.default_workspace
+
+            if filestore_names:
+                filestores = FilestoreManager.get_or_create_filestores(
+                    names=filestore_names,
+                    workspace=workspace,
+                    region=region,
+                )
+                filestores_to_attach.extend([{"id": fs["id"], "input": True, "output": True} for fs in filestores])
+
             with LightRichClusterWidget(
                 workspace=workspace,
                 title=f"Running [bold]{coiled.utils.join_command_parts(command)}[/bold]",
@@ -736,6 +778,7 @@ def start_run(
                     "package_sync_conda_extras": package_sync_conda_extras,
                     "backend_options": {"docker_shm_size": docker_shm_size} if docker_shm_size else None,
                     "unset_single_threading_variables": True,
+                    "filestores_to_attach": filestores_to_attach,
                 }
                 cluster_kwargs["name"] = name or f"run-{short_random_string()}"
                 cluster_kwargs["cloud"] = cloud