coiled 1.117.2.dev3__tar.gz → 1.129.3.dev13__tar.gz

{coiled-1.117.2.dev3 → coiled-1.129.3.dev13}/PKG-INFO

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

{coiled-1.117.2.dev3 → coiled-1.129.3.dev13}/coiled/batch.py

@@ -17,6 +17,8 @@ def run(
     workspace: str | None = None,
     software: str | None = None,
     container: str | None = None,
+    run_on_host: bool | None = None,
+    cluster_kwargs: dict | None = None,
     env: list | dict | None = None,
     secret_env: list | dict | None = None,
     tag: list | dict | None = None,
@@ -46,7 +48,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 +67,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 +86,8 @@ def run(
         workspace=workspace,
         software=software,
         container=container,
+        run_on_host=run_on_host,
+        cluster_kwargs=cluster_kwargs,
         env=env,
         secret_env=secret_env,
         tag=tag,
@@ -106,7 +118,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.117.2.dev3 → coiled-1.129.3.dev13}/coiled/capture_environment.py

@@ -1,6 +1,4 @@
-import asyncio
 import contextlib
-import logging
 import platform
 import sys
 import typing
@@ -69,49 +67,52 @@ async def approximate_packages(
     architecture: ArchitectureTypesEnum = ArchitectureTypesEnum.X86_64,
     pip_check_errors: Optional[Dict[str, List[str]]] = None,
     gpu_enabled: bool = False,
+    use_uv_installer: bool = True,
 ) -> typing.List[ResolvedPackageInfo]:
     user_conda_installed_python = next((p for p in packages if p["name"] == "python"), None)
-    user_conda_installed_pip = next(
-        (i for i, p in enumerate(packages) if p["name"] == "pip" and p["source"] == "conda"),
-        None,
-    )
-    if not user_conda_installed_pip:
-        # This means pip was installed by pip, or the system
-        # package manager
-        # Insert a conda version of pip to be installed first, it will
-        # then be used to install the users version of pip
-        pip = next(
-            (p for p in packages if p["name"] == "pip" and p["source"] == "pip"),
+    # Only add pip if we need it
+    if not use_uv_installer:
+        user_conda_installed_pip = next(
+            (i for i, p in enumerate(packages) if p["name"] == "pip" and p["source"] == "conda"),
             None,
         )
-        if not pip:
-            # insert a modern version and hope it does not introduce conflicts
-            packages.append({
-                "name": "pip",
-                "path": None,
-                "source": "conda",
-                "channel_url": "https://conda.anaconda.org/conda-forge/",
-                "channel": "conda-forge",
-                "subdir": "noarch",
-                "conda_name": "pip",
-                "version": "22.3.1",
-                "wheel_target": None,
-                "requested": False,
-            })
-        else:
-            # insert the users pip version and hope it exists on conda-forge
-            packages.append({
-                "name": "pip",
-                "path": None,
-                "source": "conda",
-                "channel_url": "https://conda.anaconda.org/conda-forge/",
-                "channel": "conda-forge",
-                "subdir": "noarch",
-                "conda_name": "pip",
-                "version": pip["version"],
-                "wheel_target": None,
-                "requested": True,
-            })
+        if not user_conda_installed_pip:
+            # This means pip was installed by pip, or the system
+            # package manager
+            # Insert a conda version of pip to be installed first, it will
+            # then be used to install the users version of pip
+            pip = next(
+                (p for p in packages if p["name"] == "pip" and p["source"] == "pip"),
+                None,
+            )
+            if not pip:
+                # insert a modern version and hope it does not introduce conflicts
+                packages.append({
+                    "name": "pip",
+                    "path": None,
+                    "source": "conda",
+                    "channel_url": "https://conda.anaconda.org/conda-forge/",
+                    "channel": "conda-forge",
+                    "subdir": "noarch",
+                    "conda_name": "pip",
+                    "version": "22.3.1",
+                    "wheel_target": None,
+                    "requested": False,
+                })
+            else:
+                # insert the users pip version and hope it exists on conda-forge
+                packages.append({
+                    "name": "pip",
+                    "path": None,
+                    "source": "conda",
+                    "channel_url": "https://conda.anaconda.org/conda-forge/",
+                    "channel": "conda-forge",
+                    "subdir": "noarch",
+                    "conda_name": "pip",
+                    "version": pip["version"],
+                    "wheel_target": None,
+                    "requested": True,
+                })
     coiled_selected_python = None
     if not user_conda_installed_python:
         # insert a special python package
@@ -208,6 +209,7 @@ async def create_environment_approximation(
     progress: Optional[Progress] = None,
     architecture: ArchitectureTypesEnum = ArchitectureTypesEnum.X86_64,
     gpu_enabled: bool = False,
+    use_uv_installer: bool = True,
 ) -> typing.List[ResolvedPackageInfo]:
     packages = await scan_prefix(progress=progress)
     pip_check_errors = await check_pip_happy(progress)
@@ -237,6 +239,7 @@ async def create_environment_approximation(
         architecture=architecture,
         pip_check_errors=pip_check_errors,
         gpu_enabled=gpu_enabled,
+        use_uv_installer=use_uv_installer,
     )
     return result
 
@@ -306,6 +309,7 @@ async def scan_and_create(
                 architecture=architecture,
                 gpu_enabled=gpu_enabled,
                 conda_extras=package_sync_conda_extras,
+                use_uv_installer=use_uv_installer,
             )
 
         if not package_sync_only:
@@ -427,39 +431,3 @@ If you use pip, venv, uv, pixi, etc. create a new environment and then:
 
 See https://docs.coiled.io/user_guide/software/package_sync_best_practices.html
 for more best practices. If that doesn't solve your issue, please contact support@coiled.io.""")
-
-
-if __name__ == "__main__":
-    from logging import basicConfig
-
-    basicConfig(level=logging.INFO)
-
-    from rich.console import Console
-    from rich.table import Table
-
-    async def run():
-        async with CloudV2(asynchronous=True) as cloud:
-            return await create_environment_approximation(
-                cloud=cloud,
-                priorities={
-                    ("dask", "conda"): PackageLevelEnum.CRITICAL,
-                    ("twisted", "conda"): PackageLevelEnum.IGNORE,
-                    ("graphviz", "conda"): PackageLevelEnum.LOOSE,
-                    ("icu", "conda"): PackageLevelEnum.LOOSE,
-                },
-            )
-
-    result = asyncio.run(run())
-
-    table = Table(title="Packages")
-    keys = ("name", "source", "include", "client_version", "specifier", "error", "note")
-
-    for key in keys:
-        table.add_column(key)
-
-    for pkg in result:
-        row_values = [str(pkg.get(key, "")) for key in keys]
-        table.add_row(*row_values)
-    console = Console()
-    console.print(table)
-    console.print(table)

{coiled-1.117.2.dev3 → coiled-1.129.3.dev13}/coiled/cli/batch/run.py

@@ -15,10 +15,13 @@ 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
 from coiled.cli.utils import CONTEXT_SETTINGS, fix_path_for_upload
 from coiled.credentials.aws import get_aws_local_session_token
+from coiled.filestore import FilestoreManager, upload_to_filestore_with_ui
 from coiled.utils import COILED_LOGGER_NAME, error_info_for_tracking, supress_logs
 
 console = Console(width=80)
@@ -103,7 +106,7 @@ def handle_possible_implicit_file(implicit_file):
             "remote_path": f"/scratch/{remote_rel_dir}{remote_base}",
             "content": file_content,
         }
-    elif any(implicit_file.endswith(t) for t in UPLOAD_FILE_TYPES):
+    elif any(implicit_file.endswith(t) and "$" not in implicit_file for t in UPLOAD_FILE_TYPES):
         console.print(
             f"[orange1]WARNING:[/orange1] {implicit_file} appears to be a filename, "
             "but this file not found locally and will not be copied to VMs",
@@ -193,6 +196,7 @@ def get_kwargs_from_header(f: dict, click_params: list):
         "default is to use the entrypoint (if any) set on the image."
     ),
 )
+@click.option("--run-on-host", default=None, help="Run code directly on host, not inside docker container.")
 @click.option(
     "--env",
     "-e",
@@ -335,6 +339,50 @@ def get_kwargs_from_header(f: dict, click_params: list):
         "into individual values. By default this is ',' for ``--map-over-values`` and newline for ``--map-over-file``."
     ),
 )
+@click.option("--wait", default=None, is_flag=True)
+@click.option(
+    "--upload",
+    "local_upload_path",
+    default=None,
+    type=str,
+    help=(
+        "File or directory to upload to cloud storage and download onto the VM(s). "
+        "By default files will be copied into the working directory on VM where your batch script runs."
+    ),
+)
+@click.option(
+    "--download",
+    "local_download_path",
+    default=None,
+    type=str,
+    help=(
+        "When used with ``--wait``, output files from job will be downloaded into this local directory "
+        "when job is complete. When used without ``--wait``, files won't be automatically downloaded, "
+        "but job will be configured to store result files in cloud storage for later download."
+    ),
+)
+@click.option(
+    "--sync",
+    "local_sync_path",
+    default=None,
+    type=str,
+    help="Equivalent to specifying both ``--upload`` and ``--download`` with the same local directory.",
+)
+@click.option(
+    "--pipe-to-files",
+    default=None,
+    is_flag=True,
+    help=(
+        "Write stdout and stderr from each task to files which can be downloaded when job is complete. "
+        "This is in addition to sending stdout and stderr to logs, and is more convenient than logs for when "
+        "you want to use outputs from tasks as inputs to further processing)."
+    ),
+)
+@click.option("--input-filestore", default=None, type=str, help="Name of input filestore")
+@click.option("--output-filestore", default=None, type=str, help="Name of output filestore")
+@click.option(
+    "--scheduler-sidecar-spec", default=None, type=str, help="Filename for scheduler sidecar spec (yaml or json)"
+)
 @click.option(
     "--ntasks",
     "--n-tasks",
@@ -425,7 +473,8 @@ def get_kwargs_from_header(f: dict, click_params: list):
         "For example, you can specify '30 minutes' or '1 hour'. Default is no timeout."
     ),
 )
-@click.argument("command", nargs=-1)
+@click.option("--dask-container", default=None, type=str)
+@click.argument("command", nargs=-1, required=True)
 def batch_run_cli(ctx, **kwargs):
     """
     Submit a batch job to run on Coiled.
@@ -454,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:
@@ -475,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"
@@ -666,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()
 
@@ -707,6 +763,12 @@ 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 = 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"
+    )
+
     cluster_kwargs = {
         "name": kwargs["name"],
         "workspace": kwargs["workspace"],
@@ -719,7 +781,9 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
         # if batch job is running in extra container, then we just need a pretty minimal dask container
         # so for now switch the default in that case to basic dask container
         # TODO would it be better to use a pre-built senv with our `cloud-env-run` container instead?
-        "container": "daskdev/dask:latest" if kwargs["container"] and not kwargs["software"] else None,
+        "container": dask_container
+        if (kwargs["container"] or kwargs.get("run_on_host")) and not kwargs["software"]
+        else None,
         "region": kwargs["region"],
         "scheduler_options": {
             "idle_timeout": "520 weeks",  # TODO allow job timeout?
@@ -739,6 +803,8 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
         "package_sync_conda_extras": kwargs.get("package_sync_conda_extras"),
         "package_sync_ignore": kwargs.get("package_sync_ignore"),
         "allow_cross_zone": True if kwargs["allow_cross_zone"] is None else kwargs["allow_cross_zone"],
+        "scheduler_sidecars": scheduler_sidecars,
+        **(kwargs.get("cluster_kwargs") or {}),
     }
 
     # when task will run on scheduler, give it the same VM specs as worker node
@@ -758,7 +824,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
 
     # Create a job
     job_spec = {
-        "user_command": " ".join(command),  # don't use shlex because we don't things like `$foo` or `&&` escaped
+        "user_command": coiled.utils.join_command_parts(command),
         "user_files": user_files,
         **job_array_kwargs,
         "scheduler_task_array": scheduler_task_ids,
@@ -771,8 +837,10 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
         # Avoid possibly breaking prefect batch jobs
         # https://github.com/coiled/platform/pull/8655#pullrequestreview-2826448869
         "workdir": None if "flow-run-id" in tags else "/scratch/batch",
+        "pipe_to_files": bool(kwargs.get("pipe_to_files")),
         "host_setup": host_setup_content,
         "job_timeout_seconds": parse_timedelta(kwargs["job_timeout"]) if kwargs["job_timeout"] else None,
+        "run_in_container": not kwargs.get("run_on_host"),
     }
 
     with coiled.Cloud(workspace=kwargs["workspace"]) as cloud:
@@ -796,6 +864,48 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
 
         job_id = response["id"]
 
+        filestores_to_attach = []
+
+        for sidecar in scheduler_sidecars:
+            for attachment in sidecar.get("filestores") or []:
+                filestores_to_attach.append({"worker": False, "input": True, "output": True, **attachment})
+
+        # only create and attach filestores if using upload or indicate desire to store results
+        if (
+            kwargs.get("local_upload_path")
+            or kwargs.get("local_sync_path")
+            or kwargs.get("local_download_path")
+            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}"
+
+            in_fs_name = kwargs.get("input_filestore") or f"{fs_base_name}-input"
+            out_fs_name = kwargs.get("output_filestore") or f"{fs_base_name}-output"
+
+            filestores = FilestoreManager.get_or_create_filestores(
+                names=[in_fs_name, out_fs_name],
+                workspace=job_spec["workspace"],
+                region=kwargs["region"],
+            )
+
+            in_fs = filestores[0]
+            out_fs = filestores[1]
+
+            filestores_to_attach.extend([
+                {"id": in_fs["id"], "input": True, "path": "/scratch/batch/", "primary": True},
+                {"id": out_fs["id"], "output": True, "path": "/scratch/batch/", "primary": True},
+            ])
+
+            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"),
+                    file_buffers=kwargs.get("buffers_to_upload"),
+                )
+
         # Run the job on a cluster
         with supress_logs([COILED_LOGGER_NAME], level=logging.WARNING):
             cluster = coiled.Cluster(
@@ -804,9 +914,16 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
                 **cluster_kwargs,
             )
 
+        # TODO support for attaching as part of create request
+        if filestores_to_attach:
+            FilestoreManager.attach_filestores_to_cluster(
+                cluster_id=cluster.cluster_id,
+                attachments=filestores_to_attach,
+            )
+
         if logger:
             message = f"""
-Command:     {" ".join(command)}
+Command:     {coiled.utils.join_command_parts(command)}
 Cluster ID:  {cluster.cluster_id}
 URL:         {cluster.details_url}
 Tasks:       {n_tasks}
@@ -822,17 +939,31 @@ Tasks:       {n_tasks}
                     status_command = f"{status_command} --workspace {kwargs['workspace']}"
             else:
                 status_command = f"coiled.batch.status({cluster.cluster_id})"
-            message = f"""
-[bold]Command[/]:     [bright_blue]{" ".join(command)}[/]
-[bold]Cluster ID[/]:  [bright_blue]{cluster.cluster_id}[/]
-[bold]URL[/]:         [link][bright_blue]{cluster.details_url}[/bright_blue][/link]
-[bold]Tasks[/]:       [bright_blue]{n_tasks}[/]
+            track_status_message = (
+                f"""
 
 To track progress run:
 
   [green]{status_command}[/]
-{extra_message}"""
+"""
+                if not kwargs.get("wait")
+                else ""
+            )
+
+            message = f"""
+[bold]Command[/]:     [bright_blue]{coiled.utils.join_command_parts(command)}[/]
+[bold]Cluster ID[/]:  [bright_blue]{cluster.cluster_id}[/]
+[bold]URL[/]:         [link][bright_blue]{cluster.details_url}[/bright_blue][/link]
+[bold]Tasks[/]:       [bright_blue]{n_tasks}[/]
+{track_status_message}{extra_message}"""
 
             console.print(Panel(message, title="Coiled Batch"))
 
+        if kwargs.get("wait") and cluster.cluster_id:
+            batch_job_wait(
+                cluster_id=cluster.cluster_id,
+                workspace=job_spec["workspace"],
+                download=kwargs.get("local_download_path") or kwargs.get("local_sync_path"),
+            )
+
         return {"cluster_id": cluster.cluster_id, "cluster_name": cluster.name, "job_id": job_id}

coiled-1.129.3.dev13/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.129.3.dev13/coiled/cli/batch/wait.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import itertools
+import time
+from collections import Counter
+
+import click
+from rich.align import Align
+from rich.console import Console, Group
+from rich.progress import BarColumn, TextColumn
+from rich.status import Status
+
+import coiled
+
+from ..cluster.utils import find_cluster
+from ..utils import CONTEXT_SETTINGS
+
+console = Console(width=80)
+
+
+@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(
+    "--download",
+    default=None,
+)
+def batch_wait_cli(
+    cluster: str,
+    workspace: str | None,
+    download: str | None,
+):
+    """Monitor the progress of a Coiled Batch job."""
+    with coiled.Cloud(workspace=workspace) as cloud:
+        cluster_info = find_cluster(cloud, cluster)
+        cluster_id = cluster_info["id"]
+        jobs = coiled.batch.status(cluster=cluster_id, workspace=workspace)
+        if not jobs:
+            print(f"No batch jobs for cluster {cluster_id}")
+            return
+    batch_job_wait(cluster_id=cluster_id, workspace=workspace, download=download)
+
+
+def batch_job_wait(cluster_id: int, workspace: str | None, download: str | None):
+    cluster_errored = False
+    with coiled.utils.SimpleRichProgressPanel(
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(complete_style="progress.remaining"),
+        TextColumn("[progress.percentage]{task.completed}"),
+        console=console,
+        batch_title="",
+    ) as progress:
+        progress_tasks = {
+            "pending": progress.add_task("[yellow]Pending"),
+            "running": progress.add_task("[green]Processing"),
+            "done": progress.add_task("[blue]Done"),
+            "error": progress.add_task("[red]Error"),
+        }
+        done = False
+        while not done:
+            jobs = coiled.batch.status(cluster=cluster_id, workspace=workspace)
+            progress.batch_title = format_batch_title(jobs)  # type: ignore
+            tasks = list(itertools.chain.from_iterable(job["tasks"] for job in jobs))
+            states = [task["state"] if not task["exit_code"] else "error" for task in tasks]
+            counts = Counter(states)
+            for state, task in progress_tasks.items():
+                # show both "running" and "assigned" as "Processing..."
+                state_to_show = "running" if state == "assigned" else state
+
+                progress.update(task, total=len(tasks), completed=counts[state_to_show], refresh=True)
+
+            all_tasks_completed = (counts["done"] + counts["error"]) == len(tasks)
+            cluster_errored = jobs[0]["cluster_state"] == "error"
+            done = all_tasks_completed or cluster_errored
+            if not done:
+                time.sleep(2)
+
+    if download and not cluster_errored:
+        attachments = coiled.filestore.wait_until_complete(cluster_id)
+
+        if not attachments:
+            print("No filestores found for this job")
+            return
+
+        for attachment in attachments:
+            if attachment["output"] and attachment["primary"]:
+                coiled.filestore.download_from_filestore_with_ui(
+                    fs=attachment["filestore"],
+                    into=download,
+                )
+
+        coiled.filestore.clear_filestores_with_ui([a["filestore"] for a in attachments if a["primary"]])
+
+
+def format_batch_title(jobs):
+    cluster_id = jobs[0]["cluster_id"]
+    cluster_state = jobs[0]["cluster_state"]
+    user_command = jobs[0]["user_command"]
+    return Group(
+        Align.center(
+            Status(f"Monitoring jobs for cluster {cluster_id} ([bold]{cluster_state}[/bold])", spinner="dots")
+        ),
+        Align.center(f"[bold]Command:[/bold] [green]{user_command}[/green]"),
+    )

{coiled-1.117.2.dev3 → coiled-1.129.3.dev13}/coiled/cli/core.py

@@ -7,8 +7,10 @@ from .config import config
 from .curl import curl
 from .diagnostics import diagnostics
 from .env import env
+from .file import file_group
 from .hello import hello
 from .login import login
+from .mpi import mpi_group
 from .notebook import notebook_group
 from .package_sync import package_sync
 from .prefect import prefect
@@ -40,3 +42,5 @@ cli.add_command(batch_group)
 cli.add_command(better_logs, "logs")
 cli.add_command(hello)
 cli.add_command(hello, "quickstart")
+cli.add_command(file_group)
+cli.add_command(mpi_group, "mpi")