coiled 1.127.1.dev3__py3-none-any.whl → 1.127.1.dev5__py3-none-any.whl

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/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/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/types.py

@@ -435,10 +435,20 @@ class AWSOptions(BackendOptions, total=False):
         If possible, this will attempt to put workers in the same cluster placement group (in theory this can
         result in better network between workers, since they'd be physically close to each other in datacenter,
         though we haven't seen this to have much benefit in practice).
+    use_worker_placement_group:
+        Cluster placement group for only the workers, not the scheduler.
+    use_efa
+        Attach Elastic Fabric Adaptor for faster inter-connect between instances.
+        Only some instance types are supported.
+    use_worker_efa
+        Attach Elastic Fabric Adaptor only on cluster workers, not the scheduler.
     """
 
     keypair_name: Optional[str]
     use_placement_group: Optional[bool]
+    use_worker_placement_group: Optional[bool]
+    use_efa: Optional[bool]
+    use_worker_efa: Optional[bool]
 
 
 class GCPOptions(BackendOptions, total=False):

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"])
 

{coiled-1.127.1.dev3.dist-info → coiled-1.127.1.dev5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coiled
-Version: 1.127.1.dev3
+Version: 1.127.1.dev5
 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.dist-info → coiled-1.127.1.dev5.dist-info}/RECORD

@@ -2,7 +2,7 @@ coiled/__init__.py,sha256=SslrfRlKfunoTJMCfopqezFePDDKS8LG_FhLkMMo_TE,2274
 coiled/__main__.py,sha256=4XILBmm4ChZYo7h3JzgslFU0tjQVzdX0XtYcQLhCv0w,171
 coiled/analytics.py,sha256=96CeL8KVnm3-76lvT4fNkgML0lHebaLea-YP3wW-KqM,7486
 coiled/auth.py,sha256=go7vWtCwBbwtWyNrNBxg28xBrdjrETbE-mn3KaN5Yl8,1867
-coiled/batch.py,sha256=1gsUboonjwOQjFa55VFL3pdBWhZ-Vh-fk8N8qkWNqa4,6214
+coiled/batch.py,sha256=LLrkq910sXNtrssPEqK-KU6-aWRS9XuQEaRQw8wPUP4,7023
 coiled/capture_environment.py,sha256=K5mNwUe8KM_l203h3oZvcZUJTrGozT-CH1GrtuPYv8U,18458
 coiled/cluster.py,sha256=wwK9-SefbFBUEHJjYHXlWN3YvPcvR6XD2J-RdPCGhgc,5049
 coiled/coiled.yaml,sha256=z70xzNUy0E8b8Yt12tYYmjJDDmp-U63oUD61ccuu5N0,1037
@@ -12,7 +12,7 @@ coiled/context.py,sha256=MXWsW0swdYU-x32U7NiM0xt-t65maiEO8rvsGGeScFw,4754
 coiled/core.py,sha256=Cu6hKBXRWSztbpF8huAyU_1glnt1gacnO9vExvG-Cwo,110796
 coiled/errors.py,sha256=5aXhNXgidMm0VgPYT3MZMwlHhRE57MeSmqAJFHYaa8Y,305
 coiled/exceptions.py,sha256=jUXgmfO0LitGe8ztSmAlzb9eQV3X5c0kNO2BwtEDTYg,3099
-coiled/filestore.py,sha256=GkRLsGAPR2l6_GPwysvfB0rC3Gt8HIkQaP87fTKN5kM,14912
+coiled/filestore.py,sha256=Tc594sTm2e_TGsuxYQNFO8jCyT2fxjRMVzEmNRAJdCM,17952
 coiled/function.py,sha256=pONtcTUDRr0dykhVV73AWXqU_pb-4-lvOA0tR3i_PlA,21654
 coiled/plugins.py,sha256=w03H2Sck54QmwrVOM1BVscNiVeQsHyGm1yWNzPPPWKs,3424
 coiled/prefect.py,sha256=j1EOg7Xuw82TNRonAGEoZ3ANlwN8GM5aDXRYSjC0lnA,1497
@@ -22,8 +22,8 @@ coiled/software.py,sha256=eh3kZ8QBuIt_SPvTy_x6TXEv87SGqOJkO4HW-LCSsas,8701
 coiled/software_utils.py,sha256=JqGO8nstm0Hi-UCIBhHa25reNeVO-XOnv5eLoIyRcBo,40367
 coiled/spans.py,sha256=Aq2MOX6JXaJ72XiEmymPcsefs-kID85MEw6t-kOdPWI,2078
 coiled/spark.py,sha256=kooZCZT4dLMG_AQEOlaf6gj86G3UdowDfbw-Eiq94MU,9059
-coiled/types.py,sha256=abXo79w_FraFFCqC-6kX6S9D6yEPJzCsWXwYbpo1W-c,13984
-coiled/utils.py,sha256=Urochb9rYx_lKu54Sp_CPyXxtJ7iyfPZeRgf6ld8uEk,78464
+coiled/types.py,sha256=mpYmhX9FGoe_pE_GU8Nx2nZXdwbn-DBSHYXl87ferjM,14442
+coiled/utils.py,sha256=WalMzNUbjVUJvAMgXaTTyDC0HPSM_zsHiYRHK7lmkkk,78514
 coiled/websockets.py,sha256=BaCNiOgPVtm55R_rf1TK78tzoFSKLp4z2UCW7S57iNA,5956
 coiled/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 coiled/cli/config.py,sha256=WKZkDpPYywYS665krQLqid2RgSchDOddZqBkkwBtyVk,2267
@@ -42,7 +42,7 @@ coiled/cli/utils.py,sha256=cp7ToFGRpUKi6iNL6BbLjzgrgeTYSX_C55lYhaKWHHA,3479
 coiled/cli/batch/__init__.py,sha256=539CnfnqqcW7ndSufTS-Ie5FGZiElMYxE0Ptu70wo8M,660
 coiled/cli/batch/list.py,sha256=lU3mXeKUHltzpdbo7Txcd64T1-XpHE-wxTFIBoTR28w,3300
 coiled/cli/batch/logs.py,sha256=CbGK5GRjtu0UKGetKY_yTcjB-3PdHy5Xf4MLksYTE8g,1947
-coiled/cli/batch/run.py,sha256=zdIVlX3nVNzqeVxaS3lrvZPd3vZaKSS6RauQKE0B_aw,36616
+coiled/cli/batch/run.py,sha256=_hFcSETJMhXCqcDXe06wb2Kly4FJanTNjkIPdxO65bg,37077
 coiled/cli/batch/status.py,sha256=cSpMRC0YHFQ18y-XEXFjTOza_OTd7mtT-NX7sRSLMhk,4193
 coiled/cli/batch/util.py,sha256=ztisJzDHpsYswrdk_hI7USUkfponI8oLhcKAa6DXJo4,1026
 coiled/cli/batch/wait.py,sha256=dEP1OH0IYteqaYU2UdrGm_vU7IDE7h3l3Cb3KBoaCCY,3879
@@ -95,8 +95,8 @@ coiled/v2/widgets/__init__.py,sha256=Bt3GHTTyri-kFUaqGRVydDM-sCg5NdNujDg2RyvgV8U
 coiled/v2/widgets/interface.py,sha256=YeMQ5qdRbbpM04x9qIg2LE1xwxyRxFbdDYnkrwHazPk,301
 coiled/v2/widgets/rich.py,sha256=3rU5-yso92NdeEh3uSvEE-GwPNyp6i0Nb5PE5czXCik,28974
 coiled/v2/widgets/util.py,sha256=Y8qpGqwNzqfCzgyRFRy7vcscBoXqop-Upi4HLPpXLgg,3120
-coiled-1.127.1.dev3.dist-info/METADATA,sha256=kiRJC5gmRLQevypb0ekf9beDPcXbsNnQQjSC_dh2rMI,2181
-coiled-1.127.1.dev3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-coiled-1.127.1.dev3.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
-coiled-1.127.1.dev3.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
-coiled-1.127.1.dev3.dist-info/RECORD,,
+coiled-1.127.1.dev5.dist-info/METADATA,sha256=n6zw1SHWBo_ZlfyKYVUTTjKxDnlnDFzg64BZpTtsW5Y,2181
+coiled-1.127.1.dev5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+coiled-1.127.1.dev5.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
+coiled-1.127.1.dev5.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
+coiled-1.127.1.dev5.dist-info/RECORD,,