coiled 1.126.1.dev19__tar.gz → 1.126.1.dev20__tar.gz

{coiled-1.126.1.dev19 → coiled-1.126.1.dev20}/PKG-INFO

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

{coiled-1.126.1.dev19 → coiled-1.126.1.dev20}/coiled/cli/batch/run.py

@@ -10,12 +10,12 @@ import shlex
 
 import click
 import dask.config
-import yaml
 from dask.utils import format_bytes, format_time, parse_timedelta
 from rich.console import Console
 from rich.panel import Panel
 
 import coiled
+from coiled.cli.batch.util import load_sidecar_spec
 from coiled.cli.batch.wait import batch_job_wait
 from coiled.cli.curl import sync_request
 from coiled.cli.run import dict_from_key_val_list
@@ -755,25 +755,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
 
     batch_job_container = f"{kwargs['container']}!" if kwargs["ignore_container_entrypoint"] else kwargs["container"]
 
-    scheduler_sidecars = []
-    if kwargs.get("scheduler_sidecar_spec"):
-        with open(kwargs["scheduler_sidecar_spec"]) as f:
-            if kwargs["scheduler_sidecar_spec"].endswith((".yaml", ".yml")):
-                sidecar_spec = yaml.safe_load(f)
-            elif kwargs["scheduler_sidecar_spec"].endswith(".json"):
-                sidecar_spec = json.load(f)
-            else:
-                raise ValueError(f"Unknown format for {kwargs['scheduler_sidecar_spec']}, json or yaml expected.")
-
-            # support either list-like or dict-like
-            if isinstance(sidecar_spec, list):
-                scheduler_sidecars = sidecar_spec
-            if isinstance(sidecar_spec, dict):
-                scheduler_sidecars = [{"name": key, **val} for key, val in sidecar_spec.items()]
-
-            for sidecar in scheduler_sidecars:
-                # allow `image` as the key, to match docker compose spec
-                sidecar["container"] = sidecar.get("container") or sidecar.get("image")
+    scheduler_sidecars = load_sidecar_spec(kwargs.get("scheduler_sidecar_spec"))
 
     dask_container = (
         kwargs.get("dask_container") or dask.config.get("coiled.batch.dask-container", None) or "ghcr.io/dask/dask"

coiled-1.126.1.dev20/coiled/cli/batch/util.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import json
+
+import yaml
+
+
+def load_sidecar_spec(spec_path: str | None):
+    scheduler_sidecars = []
+    if spec_path:
+        with open(spec_path) as f:
+            if spec_path.endswith((".yaml", ".yml")):
+                sidecar_spec = yaml.safe_load(f)
+            elif spec_path.endswith(".json"):
+                sidecar_spec = json.load(f)
+            else:
+                raise ValueError(f"Unknown format for {spec_path}, json or yaml expected.")
+
+            # support either list-like or dict-like
+            if isinstance(sidecar_spec, list):
+                scheduler_sidecars = sidecar_spec
+            if isinstance(sidecar_spec, dict):
+                scheduler_sidecars = [{"name": key, **val} for key, val in sidecar_spec.items()]
+
+            for sidecar in scheduler_sidecars:
+                # allow `image` as the key, to match docker compose spec
+                sidecar["container"] = sidecar.get("container") or sidecar.get("image")
+    return scheduler_sidecars

{coiled-1.126.1.dev19 → coiled-1.126.1.dev20}/coiled/cli/run.py

@@ -31,6 +31,7 @@ from coiled.utils import (
 from coiled.v2.cluster import ClusterKwargs
 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
 
@@ -238,7 +239,7 @@ 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) "
+            "(apt update -y && apt upgrade -y && apt install -y --no-install-recommends ca-certificates) "
             "2>&1 > /dev/null\n"
             f"{command_string}"
         )
@@ -550,6 +551,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,
@@ -580,6 +588,7 @@ def run(
     package_sync_strict,
     package_sync_conda_extras,
     docker_shm_size,
+    filestore_names,
     command,
 ):
     """
@@ -615,6 +624,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"])
 
@@ -652,6 +662,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)
@@ -694,9 +705,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]",
@@ -737,6 +759,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

{coiled-1.126.1.dev19 → coiled-1.126.1.dev20}/coiled/v2/cluster.py

@@ -210,6 +210,7 @@ class ClusterKwargs(TypedDict, total=False):
     scheduler_sidecars: list[dict] | None
     worker_sidecars: list[dict] | None
     pause_on_exit: bool | None
+    filestores_to_attach: list[dict] | None
 
 
 class Cluster(DistributedCluster, Generic[IsAsynchronous]):
@@ -478,6 +479,8 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
         Like ``scheduler_sidecars``, but run on worker VMs instead of scheduler.
     pause_on_exit
         Pause the cluster instead of shutting it down when exiting.
+    filestores_to_attach
+        List of filestores to attach (specified as ``{"id": id, "input": True, "output": True}``, not name).
     """
 
     _instances = weakref.WeakSet()
@@ -563,6 +566,7 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
         scheduler_sidecars: list[dict] | None = None,
         worker_sidecars: list[dict] | None = None,
         pause_on_exit: bool | None = None,
+        filestores_to_attach: list[dict] | None = None,
     ):
         self.pause_on_exit = pause_on_exit
         self.init_time = datetime.datetime.now(tz=datetime.timezone.utc)
@@ -1007,6 +1011,8 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
             no_client_timeout if no_client_timeout != NO_CLIENT_DEFAULT else (idle_timeout or "2 minutes")
         )
 
+        self.filestores_to_attach = filestores_to_attach
+
         if not self.asynchronous:
             # If we don't close the cluster, the user's ipython session gets spammed with
             # messages from distributed.
@@ -1617,6 +1623,7 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
                     worker_sidecars=self.worker_sidecars,
                     host_setup_script_content=self.host_setup_script_content,
                     pause_on_exit=self.pause_on_exit,
+                    filestores_to_attach=self.filestores_to_attach,
                     cluster_timeout_seconds=self.cluster_timeout_seconds,
                 )
                 cluster_created = not cluster_existed

{coiled-1.126.1.dev19 → coiled-1.126.1.dev20}/coiled/v2/core.py

@@ -643,6 +643,7 @@ class CloudV2(OldCloud, Generic[IsAsynchronous]):
         host_setup_script_content: str | None = None,
         pause_on_exit: bool | None = None,
         cluster_timeout_seconds: int | None = None,
+        filestores_to_attach: list[dict] | None = None,
     ) -> Tuple[int, bool]:
         # TODO (Declarative): support these args, or decide not to
         # https://gitlab.com/coiled/cloud/-/issues/4305
@@ -684,6 +685,7 @@ class CloudV2(OldCloud, Generic[IsAsynchronous]):
             "pause_on_exit": pause_on_exit,
             "cluster_timeout_seconds": cluster_timeout_seconds,
             "coiled_cloud_env_image": dask.config.get("coiled.cloud-env-image", None),
+            "filestores_to_attach": filestores_to_attach,
         }
 
         try: