coiled 1.135.4.dev17__tar.gz → 1.135.4.dev18__tar.gz

{coiled-1.135.4.dev17 → coiled-1.135.4.dev18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: coiled
-Version: 1.135.4.dev17
+Version: 1.135.4.dev18
 Summary: Python client for coiled.io dask clusters
 Project-URL: Homepage, https://coiled.io
 Maintainer-email: Coiled <info@coiled.io>

{coiled-1.135.4.dev17 → coiled-1.135.4.dev18}/coiled/batch.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import shlex
 import time
+from contextlib import nullcontext
 
 from aiohttp import ServerDisconnectedError
+from typing_extensions import Literal
 
 import coiled
 from coiled.cli.curl import sync_request
@@ -15,6 +17,7 @@ def run(
     *,
     name: str | None = None,
     workspace: str | None = None,
+    cloud: coiled.Cloud[Literal[False]] | None = None,
     software: str | None = None,
     container: str | None = None,
     run_on_host: bool | None = None,
@@ -64,6 +67,11 @@ def run(
 
     Additional Parameters
     ---------------------
+    cloud
+        Authenticated Coiled Cloud client to use for this submission. If not
+        provided, a client is created from the current Coiled configuration.
+        The client's default workspace is used. ``cloud`` and ``workspace``
+        are mutually exclusive.
     map_over_task_var_dicts
         takes a list of dictionaries, so you can specify multiple environment variables for each task.
         For example, ``[{"FOO": 1, "BAR": 2}, {"FOO": 3, "BAR": 4}]`` will pass ``FOO=1 BAR=2`` to one task and
@@ -81,6 +89,9 @@ def run(
     tag = dict_to_key_val_list(tag)
     vm_type = [vm_type] if isinstance(vm_type, str) else vm_type
 
+    if cloud is not None and workspace is not None:
+        raise ValueError("cloud and workspace are mutually exclusive")
+
     kwargs = dict(
         name=name,
         command=command,
@@ -144,25 +155,32 @@ def run(
         **default_kwargs,
     }
 
-    success = True
-    exception = None
-    try:
-        return _batch_run(default_kwargs, **kwargs)
-    except Exception as e:
-        success = False
-        exception = e
-        raise
-    finally:
-        coiled.add_interaction(
-            "coiled-batch-python",
-            success=success,
-            **error_info_for_tracking(exception),
-        )
-
-
-def wait_for_job_done(job_id: int, timeout: int | None = None) -> str | None:
+    with coiled.Cloud(workspace=workspace) if cloud is None else nullcontext(cloud) as cloud:
+        success = True
+        exception = None
+        try:
+            return _batch_run(default_kwargs, cloud=cloud, **kwargs)
+        except Exception as e:
+            success = False
+            exception = e
+            raise
+        finally:
+            tracking_info = error_info_for_tracking(exception)
+            _ = cloud.add_interaction(
+                action="coiled-batch-python",
+                success=success,
+                error_message=tracking_info.pop("error_message", None),
+                additional_data=tracking_info or None,
+            )
+
+
+def wait_for_job_done(
+    job_id: int,
+    timeout: int | None = None,
+    cloud: coiled.Cloud[Literal[False]] | None = None,
+) -> str | None:
     timeout_at = time.monotonic() + timeout if timeout is not None else None
-    with coiled.Cloud() as cloud:
+    with coiled.Cloud() if cloud is None else nullcontext(cloud) as cloud:
         url = f"{cloud.server}/api/v2/jobs/{job_id}"
         while timeout_at is None or time.monotonic() < timeout_at:
             try:

{coiled-1.135.4.dev17 → coiled-1.135.4.dev18}/coiled/cli/batch/run.py

@@ -7,12 +7,14 @@ import logging
 import os
 import re
 import shlex
+from contextlib import nullcontext
 
 import click
 import dask.config
 from dask.utils import format_bytes, format_time, parse_timedelta
 from rich.console import Console
 from rich.panel import Panel
+from typing_extensions import Literal
 
 import coiled
 from coiled.cli.batch.util import load_sidecar_spec
@@ -524,7 +526,13 @@ def batch_run_cli(ctx, **kwargs):
         )
 
 
-def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
+def _batch_run(
+    default_kwargs,
+    logger=None,
+    from_cli=False,
+    cloud: coiled.Cloud[Literal[False]] | None = None,
+    **kwargs,
+) -> dict:
     command = kwargs["command"]
     user_files = []
 
@@ -868,8 +876,9 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
         "max_retries": kwargs.get("max_retries") or 0,
     }
 
-    with coiled.Cloud(workspace=kwargs["workspace"]) as cloud:
+    with coiled.Cloud(workspace=kwargs["workspace"]) if cloud is None else nullcontext(cloud) as cloud:
         job_spec["workspace"] = cloud.default_workspace
+        cluster_kwargs["workspace"] = cloud.default_workspace
 
         compressed_data = gzip.compress(json.dumps(job_spec).encode())
         if len(compressed_data) > 2_400_000:
@@ -914,6 +923,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
                 names=[in_fs_name, out_fs_name],
                 workspace=job_spec["workspace"],
                 region=kwargs["region"],
+                cloud=cloud,
             )
 
             in_fs = filestores[0]
@@ -929,6 +939,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
                     fs=in_fs,
                     local_dir=kwargs.get("local_upload_path") or kwargs.get("local_sync_path"),
                     file_buffers=kwargs.get("buffers_to_upload"),
+                    cloud=cloud,
                 )
 
         # Run the job on a cluster
@@ -944,6 +955,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
             FilestoreManager.attach_filestores_to_cluster(
                 cluster_id=cluster.cluster_id,
                 attachments=filestores_to_attach,
+                cloud=cloud,
             )
 
         if logger:

{coiled-1.135.4.dev17 → coiled-1.135.4.dev18}/coiled/filestore.py

@@ -3,13 +3,16 @@ from __future__ import annotations
 import io
 import os
 import time
+from contextlib import nullcontext
 from pathlib import Path
+from typing import Any
 
 import httpx
 from rich.align import Align
 from rich.console import Group
 from rich.prompt import Confirm
 from rich.status import Status
+from typing_extensions import Literal
 
 import coiled
 from coiled.cli.curl import sync_request
@@ -111,7 +114,7 @@ def download_from_filestore_with_ui(fs, into=".", name_includes=None):
             )
 
 
-def upload_to_filestore_with_ui(fs, local_dir, file_buffers=None):
+def upload_to_filestore_with_ui(fs, local_dir, file_buffers=None, cloud=None):
     # TODO (future enhancement) send write status
     #   this is tricky because status is stored on the "attachment" object, which might not exist yet
     #   because we want to be able to upload files before cluster has been created
@@ -160,7 +163,7 @@ def upload_to_filestore_with_ui(fs, local_dir, file_buffers=None):
             ])
 
             # files_for_upload is type list[dict] where each dict has "relative_path" key
-            upload_info = FilestoreManager.get_signed_upload_urls(fs["id"], files_for_upload=files)
+            upload_info = FilestoreManager.get_signed_upload_urls(fs["id"], files_for_upload=files, cloud=cloud)
 
             upload_urls = upload_info.get("urls")
             existing_blobs = upload_info.get("existing")
@@ -440,10 +443,41 @@ class FilestoreManagerWithoutHttp:
 class FilestoreManager(FilestoreManagerWithoutHttp):
     http2 = True
 
+    @classmethod
+    def get_or_create_filestores(cls, names, workspace, region, cloud=None):
+        return cls.make_req(
+            "/api/v2/filestore/list",
+            post=True,
+            data={"names": names, "workspace": workspace, "region": region},
+            cloud=cloud,
+        ).get("filestores")
+
+    @classmethod
+    def get_signed_upload_urls(cls, fs_id, files_for_upload, cloud=None):
+        paths = [f["relative_path"] for f in files_for_upload]  # relative paths
+        return cls.make_req(
+            f"/api/v2/filestore/fs/{fs_id}/signed-urls/upload",
+            post=True,
+            data={"paths": paths},
+            cloud=cloud,
+        )
+
+    @classmethod
+    def attach_filestores_to_cluster(cls, cluster_id, attachments, cloud=None):
+        return cls.make_req(
+            "/api/v2/filestore/attach",
+            post=True,
+            data={
+                "cluster_id": cluster_id,
+                "attachments": attachments,
+            },
+            cloud=cloud,
+        )
+
     @staticmethod
-    def make_req(api_path, post=False, data=None):
+    def make_req(api_path, post=False, data=None, cloud: coiled.Cloud[Literal[False]] | None = None) -> Any:
         workspace = (data or {}).get("workspace")
-        with coiled.Cloud(workspace=workspace) as cloud:
+        with coiled.Cloud(workspace=workspace) if cloud is None else nullcontext(cloud) as cloud:
             url = f"{cloud.server}{api_path}"
             response = sync_request(
                 cloud=cloud,