coiled 1.127.1.dev3__tar.gz → 1.127.1.dev4__tar.gz

{coiled-1.127.1.dev3 → coiled-1.127.1.dev4}/PKG-INFO

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

{coiled-1.127.1.dev3 → coiled-1.127.1.dev4}/coiled/batch.py

@@ -17,6 +17,7 @@ def run(
     workspace: str | None = None,
     software: str | None = None,
     container: str | None = None,
+    run_on_host: bool | None = None,
     env: list | dict | None = None,
     secret_env: list | dict | None = None,
     tag: list | dict | None = None,
@@ -46,7 +47,11 @@ def run(
     package_sync_strict: bool = False,
     package_sync_conda_extras: list | None = None,
     package_sync_ignore: list[str] | None = None,
+    local_upload_path: str | None = None,
+    buffers_to_upload: list[dict] | None = None,
     host_setup_script: str | None = None,
+    host_setup_script_content: str | None = None,
+    command_as_script: bool | None = None,
     ignore_container_entrypoint: bool | None = None,
     job_timeout: str | None = None,
     logger=None,
@@ -61,8 +66,12 @@ def run(
         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
         ``FOO=3 BAR=4`` to another.
+    buffers_to_upload
+        takes a list of dictionaries, each should have path where file should be written on VM(s)
+        relative to working directory, and ``io.BytesIO`` which provides content of file,
+        for example ``[{"relative_path": "hello.txt", "buffer": io.BytesIO(b"hello")}]``.
     """
-    if isinstance(command, str):
+    if isinstance(command, str) and not command.startswith("#!") and not command_as_script:
         command = shlex.split(command)
 
     env = dict_to_key_val_list(env)
@@ -76,6 +85,7 @@ def run(
         workspace=workspace,
         software=software,
         container=container,
+        run_on_host=run_on_host,
         env=env,
         secret_env=secret_env,
         tag=tag,
@@ -106,7 +116,11 @@ def run(
         package_sync_strict=package_sync_strict,
         package_sync_conda_extras=package_sync_conda_extras,
         package_sync_ignore=package_sync_ignore,
+        local_upload_path=local_upload_path,
+        buffers_to_upload=buffers_to_upload,
         host_setup_script=host_setup_script,
+        host_setup_script_content=host_setup_script_content,
+        command_as_script=command_as_script,
         ignore_container_entrypoint=ignore_container_entrypoint,
         job_timeout=job_timeout,
         logger=logger,

{coiled-1.127.1.dev3 → coiled-1.127.1.dev4}/coiled/cli/batch/run.py

@@ -503,6 +503,14 @@ def batch_run_cli(ctx, **kwargs):
 
 def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
     command = kwargs["command"]
+    user_files = []
+
+    if isinstance(command, str) and (command.startswith("#!") or kwargs.get("command_as_script")):
+        user_files.append({
+            "path": "script",
+            "content": command,
+        })
+        command = ["script"]
 
     # Handle command as string case (e.g. `coiled batch run "python myscript.py"`)
     if len(command) == 1:
@@ -524,7 +532,6 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
     # unescape escaped COILED env vars in command
     command = [part.replace("\\$COILED", "$COILED") for part in command]
 
-    user_files = []
     kwargs_from_header = None
 
     # identify implicit files referenced in commands like "python foo.py" or "foo.sh"
@@ -715,8 +722,8 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
     if user_files_from_content:
         user_files.extend(user_files_from_content)
 
-    host_setup_content = None
-    if kwargs["host_setup_script"]:
+    host_setup_content = kwargs.get("host_setup_script_content")
+    if not host_setup_content and kwargs["host_setup_script"]:
         with open(kwargs["host_setup_script"]) as f:
             host_setup_content = f.read()
 
@@ -870,6 +877,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
             or kwargs.get("pipe_to_files")
             or kwargs.get("input_filestore")
             or kwargs.get("output_filestore")
+            or kwargs.get("buffers_to_upload")
         ):
             fs_base_name = kwargs["name"] or f"batch-job-{job_id}"
 
@@ -890,9 +898,11 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
                 {"id": out_fs["id"], "output": True, "path": "/scratch/batch/", "primary": True},
             ])
 
-            if kwargs.get("local_upload_path") or kwargs.get("local_sync_path"):
+            if kwargs.get("local_upload_path") or kwargs.get("local_sync_path") or kwargs.get("buffers_to_upload"):
                 upload_to_filestore_with_ui(
-                    fs=in_fs, local_dir=kwargs.get("local_upload_path") or kwargs.get("local_sync_path")
+                    fs=in_fs,
+                    local_dir=kwargs.get("local_upload_path") or kwargs.get("local_sync_path"),
+                    file_buffers=kwargs.get("buffers_to_upload"),
                 )
 
         # Run the job on a cluster

{coiled-1.127.1.dev3 → coiled-1.127.1.dev4}/coiled/filestore.py

@@ -111,7 +111,7 @@ def download_from_filestore_with_ui(fs, into=".", name_includes=None):
             )
 
 
-def upload_to_filestore_with_ui(fs, local_dir):
+def upload_to_filestore_with_ui(fs, local_dir, file_buffers=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
@@ -123,16 +123,24 @@ def upload_to_filestore_with_ui(fs, local_dir):
             Align.left(f"Currently uploading: [blue]{f or ''}[/blue]"),
         )
 
-    if fs and local_dir:
-        files, total_bytes = FilestoreManager.get_files_for_upload(local_dir)
+    files = []
+    total_bytes = None
 
+    if fs:
+        if local_dir:
+            files_from_path, total_bytes = FilestoreManager.get_files_for_upload(local_dir)
+            files.extend(files_from_path)
+        if file_buffers:
+            files.extend(file_buffers)
+
+    if files:
         size_label = "Bytes"
         size_scale = 1
 
-        if total_bytes > 10_000_000:
+        if total_bytes and total_bytes > 10_000_000:
             size_label = "Mb"
             size_scale = 1_000_000
-        elif total_bytes > 10_000:
+        elif total_bytes and total_bytes > 10_000:
             size_label = "Kb"
             size_scale = 1_000
 
@@ -146,7 +154,75 @@ def upload_to_filestore_with_ui(fs, local_dir):
                     "label": size_label,
                     "total": total_bytes / size_scale if size_scale > 1 else total_bytes,
                     "completed": done_bytes / size_scale if size_scale > 1 else done_bytes,
-                },
+                }
+                if total_bytes
+                else {},
+            ])
+
+            # 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_urls = upload_info.get("urls")
+            existing_blobs = upload_info.get("existing")
+
+            for file in files:
+                relative_path = file.get("relative_path")
+                local_path = file.get("local_path")
+                buffer = file.get("buffer")
+                if local_path:
+                    size = file.get("size")
+                    skip_upload = False
+                    existing_blob_info = existing_blobs.get(relative_path)
+                    if existing_blob_info:
+                        modified = os.path.getmtime(local_path)
+                        if size == existing_blob_info["size"] and modified < existing_blob_info["modified"]:
+                            skip_upload = True
+
+                    if not skip_upload:
+                        progress.batch_title = progress_title(local_path)
+                        progress.refresh()
+
+                        FilestoreManager.upload_to_signed_url(local_path, upload_urls[relative_path])
+
+                    done_bytes += size
+
+                elif buffer:
+                    FilestoreManager.upload_bytes_to_signed_url(buffer, upload_urls[relative_path])
+
+                done_files += 1
+
+                progress.update_progress([
+                    {"label": "Files", "total": len(files), "completed": done_files},
+                    {
+                        "label": size_label,
+                        "total": total_bytes / size_scale if size_scale > 1 else total_bytes,
+                        "completed": done_bytes / size_scale if size_scale > 1 else done_bytes,
+                    }
+                    if total_bytes
+                    else {},
+                ])
+
+            progress.update_title(Align.left(f"Uploaded to cloud storage: [green]{fs['name']}[green]"))
+
+        # TODO (future enhancement) send write status
+        # FilestoreManager.post_fs_write_status(fs["id"], "finish", {"complete": True, "file_count": len(files)})
+
+        return len(files)
+
+
+def upload_bytes_to_fs(fs, files):
+    def progress_title(f=None):
+        return Group(
+            Align.left(Status(f"Uploading to cloud storage: [green]{fs['name']}[green]", spinner="dots")),
+            Align.left(f"Currently uploading: [blue]{f or ''}[/blue]"),
+        )
+
+    if fs and files:
+        with coiled.utils.SimpleRichProgressPanel.from_defaults(title=progress_title()) as progress:
+            done_files = 0
+
+            progress.update_progress([
+                {"label": "Files", "total": len(files), "completed": done_files},
             ])
 
             upload_info = FilestoreManager.get_signed_upload_urls(fs["id"], files_for_upload=files)
@@ -154,7 +230,10 @@ def upload_to_filestore_with_ui(fs, local_dir):
             upload_urls = upload_info.get("urls")
             existing_blobs = upload_info.get("existing")
 
-            for local_path, relative_path, size in files:
+            for file in files:
+                local_path = file.get("local_path")
+                relative_path = file.get("relative_path")
+                size = file.get("size")
                 skip_upload = False
                 existing_blob_info = existing_blobs.get(relative_path)
                 if existing_blob_info:
@@ -169,15 +248,9 @@ def upload_to_filestore_with_ui(fs, local_dir):
                     FilestoreManager.upload_to_signed_url(local_path, upload_urls[relative_path])
 
                 done_files += 1
-                done_bytes += size
 
                 progress.update_progress([
                     {"label": "Files", "total": len(files), "completed": done_files},
-                    {
-                        "label": size_label,
-                        "total": total_bytes / size_scale if size_scale > 1 else total_bytes,
-                        "completed": done_bytes / size_scale if size_scale > 1 else done_bytes,
-                    },
                 ])
 
             progress.update_title(Align.left(f"Uploaded to cloud storage: [green]{fs['name']}[green]"))
@@ -230,7 +303,7 @@ class FilestoreManagerWithoutHttp:
 
     @classmethod
     def get_signed_upload_urls(cls, fs_id, files_for_upload):
-        paths = [p for _, p, _ in files_for_upload]  # relative paths
+        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})
 
     @classmethod
@@ -269,7 +342,7 @@ class FilestoreManagerWithoutHttp:
             relative_path = Path(os.path.relpath(local_path, local_dir)).as_posix()
             size = os.path.getsize(local_path)
 
-            files.append((local_path, relative_path, size))
+            files.append({"local_path": local_path, "relative_path": relative_path, "size": size})
             total_bytes += size
 
             return files, total_bytes
@@ -301,32 +374,36 @@ class FilestoreManagerWithoutHttp:
                 relative_path = Path(os.path.relpath(local_path, local_dir)).as_posix()
                 size = os.path.getsize(local_path)
 
-                files.append((local_path, relative_path, size))
+                files.append({"local_path": local_path, "relative_path": relative_path, "size": size})
                 total_bytes += size
         return files, total_bytes
 
     @classmethod
-    def upload_to_signed_url(cls, local_path, url):
+    def upload_to_signed_url(cls, local_path: str, url: str):
         with open(local_path, "rb") as f:
             buffer = io.BytesIO(f.read())
-            buffer.seek(0)
-            num_bytes = len(buffer.getvalue())
-            with httpx.Client(http2=cls.http2) as client:
-                headers = {"Content-Type": "binary/octet-stream", "Content-Length": str(num_bytes)}
-                if "blob.core.windows.net" in url:
-                    headers["x-ms-blob-type"] = "BlockBlob"
-                # TODO error handling
-                client.put(
-                    url,
-                    # content must be set to an iterable of bytes, rather than a
-                    # bytes object (like file.read()) because files >2GB need
-                    # to be sent in chunks to avoid an OverflowError in the
-                    # Python stdlib ssl module, and httpx will not chunk up a
-                    # bytes object automatically.
-                    content=buffer,
-                    timeout=60,
-                    headers=headers,
-                )
+            cls.upload_bytes_to_signed_url(buffer=buffer, url=url)
+
+    @classmethod
+    def upload_bytes_to_signed_url(cls, buffer: io.BytesIO, url: str):
+        buffer.seek(0)
+        num_bytes = len(buffer.getvalue())
+        with httpx.Client(http2=cls.http2) as client:
+            headers = {"Content-Type": "binary/octet-stream", "Content-Length": str(num_bytes)}
+            if "blob.core.windows.net" in url:
+                headers["x-ms-blob-type"] = "BlockBlob"
+            # TODO error handling
+            client.put(
+                url,
+                # content must be set to an iterable of bytes, rather than a
+                # bytes object (like file.read()) because files >2GB need
+                # to be sent in chunks to avoid an OverflowError in the
+                # Python stdlib ssl module, and httpx will not chunk up a
+                # bytes object automatically.
+                content=buffer,
+                timeout=60,
+                headers=headers,
+            )
 
     @classmethod
     def download_from_signed_url(cls, local_path, url, max_retries=3, verbose=False):

{coiled-1.127.1.dev3 → coiled-1.127.1.dev4}/coiled/utils.py

@@ -2176,6 +2176,8 @@ class SimpleRichProgressPanel(Progress):
 
     def update_progress(self, tasks: list[dict]):
         for task in tasks:
+            if not task:
+                continue
             if task["label"] not in self._tasks_from_dicts:
                 self._tasks_from_dicts[task["label"]] = self.add_task(task["label"])