coiled 1.127.1.dev16__py3-none-any.whl → 1.127.1.dev20__py3-none-any.whl

coiled/batch.py

@@ -18,6 +18,7 @@ def run(
     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,
@@ -86,6 +87,7 @@ def run(
         software=software,
         container=container,
         run_on_host=run_on_host,
+        cluster_kwargs=cluster_kwargs,
         env=env,
         secret_env=secret_env,
         tag=tag,

coiled/cli/batch/run.py

@@ -804,6 +804,7 @@ def _batch_run(default_kwargs, logger=None, from_cli=False, **kwargs) -> dict:
         "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

coiled/cli/core.py

@@ -10,6 +10,7 @@ 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
@@ -42,3 +43,4 @@ 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")

coiled/cli/mpi.py

@@ -0,0 +1,166 @@
+import os.path
+import shlex
+
+import click
+import fabric.connection
+
+import coiled
+
+from .cluster.utils import find_cluster
+from .run import get_ssh_connection, write_via_ssh
+from .utils import CONTEXT_SETTINGS
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.option("--worker-nodes", default=1, type=int)
+@click.option("--vm-type", default="g6.8xlarge", type=str)
+@click.option("--pip", multiple=True, type=str)
+@click.option("--idle-timeout", default=None, type=str)
+def setup(worker_nodes, vm_type, pip, idle_timeout):
+    setup_script = get_host_setup_script(pip_install=pip)
+
+    cluster = coiled.Cluster(
+        n_workers=worker_nodes,
+        container="daskdev/dask:latest",
+        allow_ssh_from="me",
+        host_setup_script=setup_script,
+        backend_options={"use_placement_group": True, "use_efa": True, "ami_version": "DL"},
+        scheduler_vm_types=[vm_type],
+        worker_vm_types=[vm_type],
+        worker_disk_size="100GB",
+        scheduler_disk_size="100GB",
+        shutdown_on_close=False,
+        idle_timeout=idle_timeout,
+    )
+
+    print("Cluster created, installing software for MPI...")
+
+    with coiled.Cloud() as cloud:
+        connection = get_ssh_connection(cloud, cluster.cluster_id)
+
+    setup_mpi_ssh(connection)
+
+    print("MPI is ready")
+
+
+@click.command(
+    context_settings=CONTEXT_SETTINGS,
+)
+@click.option("--cluster", default=None)
+@click.option("--workspace", default=None, type=str)
+@click.option("--legate", is_flag=True, default=False, type=bool)
+@click.option(
+    "--include-head/--exclude-head",
+    default=True,
+    type=bool,
+)
+@click.argument("command", nargs=-1, required=True)
+def run(cluster, workspace, legate, include_head, command):
+    nodes = "$(cat workers | wc -w)"
+
+    command = list(command)
+
+    files = {}
+    for i, c in enumerate(command):
+        if os.path.exists(c):
+            remote_path = f"/scratch/batch/{os.path.basename(c)}"
+            command[i] = remote_path
+            with open(c) as f:
+                content = f.read()
+            files[remote_path] = content
+
+    if legate:
+        # TODO make "--gpus 1 --sysmem 2000 --fbmem 20000" configurable
+        wrapped_command = f"""
+legate \
+  --gpus 1 --sysmem 2000 --fbmem 20000 \
+  --nodes {nodes} \
+  --launcher mpirun \
+  --launcher-extra ' --hostfile workers -x PATH ' \
+  {shlex.join(command)}
+"""
+    else:
+        wrapped_command = f"mpirun --hostfile workers -x PATH {shlex.join(command)}"
+
+    with coiled.Cloud(workspace=workspace) as cloud:
+        cluster_info = find_cluster(cloud, cluster)
+        cluster_id = cluster_info["id"]
+        connection = get_ssh_connection(cloud, cluster_id)
+
+    setup_mpi_ssh(connection, include_scheduler=include_head)
+
+    if files:
+        worker_connections = []
+
+        for worker in cluster_info["workers"]:
+            if (
+                not worker.get("instance")
+                or not worker["instance"].get("current_state")
+                or worker["instance"]["current_state"]["state"] != "ready"
+            ):
+                continue
+            worker_address = worker["instance"]["private_ip_address"]
+
+            worker_connections.append(
+                fabric.connection.Connection(
+                    worker_address, gateway=connection, user=connection.user, connect_kwargs=connection.connect_kwargs
+                )
+            )
+
+        for path, content in files.items():
+            write_via_ssh(connection, content=content, path=path)
+            for conn in worker_connections:
+                write_via_ssh(conn, content=content, path=path)  # , mode=0o555
+
+    print(f"Running command:\n{wrapped_command}")
+
+    # TODO keepalive session so this will interact correctly with idle timeout / keepalive
+    connection.run(wrapped_command, hide=False, pty=True, warn=True, env={"PATH": "/tmp/host-user-venv/bin:$PATH"})
+
+
+def setup_mpi_ssh(connection, include_scheduler=True):
+    add_scheduler_line = 'printf "\n127.0.0.1" >> workers' if include_scheduler else ""
+
+    setup_mpi = f"""
+/bin/coiled_agent list-worker-ips | sudo tee workers && sudo chown ubuntu workers
+ssh-keyscan -f workers -t ed25519 >> ~/.ssh/known_hosts
+{add_scheduler_line}
+
+# block until host setup script has finished, at least on schedule node
+until [ -f /tmp/host-setup-done ]
+do
+     sleep 5
+done
+"""
+
+    _ = connection.run(setup_mpi, hide=True, pty=False)
+
+
+def get_host_setup_script(venv_path="/tmp/host-user-venv", apt_install=None, pip_install=None):
+    apt_install = apt_install or []
+    apt_install.extend(["openmpi-bin", "python3-pip", "python3-venv"])
+
+    pip_install = pip_install or []
+
+    pip_install_line = f"{venv_path}/bin/python -m pip install {' '.join(pip_install)}" if pip_install else ""
+
+    return f"""
+sudo apt install {" ".join(apt_install)} -y
+
+mkdir {venv_path}
+python3 -m venv {venv_path}
+
+{pip_install_line}
+
+echo 'done' > /tmp/host-setup-done
+    """
+
+
+@click.group(name="mpi", context_settings=CONTEXT_SETTINGS)
+def mpi_group(): ...
+
+
+mpi_group.add_command(setup)
+mpi_group.add_command(run)

coiled/cli/setup/aws.py

@@ -321,6 +321,7 @@ def get_ongoing_doc(ecr=True, package_sync_bucket_prefix=None) -> str:
                 "Effect": "Allow",
                 "Resource": "*",
                 "Action": [
+                    "ec2:AuthorizeSecurityGroupEgress",
                     "ec2:AuthorizeSecurityGroupIngress",
                     "ec2:CreateFleet",
                     "ec2:CreateLaunchTemplate",
@@ -385,6 +386,16 @@ def get_ongoing_doc(ecr=True, package_sync_bucket_prefix=None) -> str:
                 ],
                 "Condition": {"StringEquals": {"ec2:ResourceTag/owner": "coiled"}},
             },
+            {
+                "Sid": "OngoingPlacementGroupPolicy",
+                "Effect": "Allow",
+                "Resource": "arn:*:ec2:*:*:placement-group/coiled-*",
+                "Action": [
+                    "ec2:CreatePlacementGroup",
+                    "ec2:DescribePlacementGroups",
+                    "ec2:DeletePlacementGroup",
+                ],
+            },
             {
                 "Sid": "OptionalLogPull",
                 "Effect": "Allow",

coiled/types.py

@@ -442,6 +442,9 @@ class AWSOptions(BackendOptions, total=False):
         Only some instance types are supported.
     use_worker_efa
         Attach Elastic Fabric Adaptor only on cluster workers, not the scheduler.
+    ami_version
+        Use non-default type of AMI.
+        Supported options include "DL" for the Deep Learning Base OSS Nvidia Driver GPU AMI.
     """
 
     keypair_name: Optional[str]
@@ -449,6 +452,7 @@ class AWSOptions(BackendOptions, total=False):
     use_worker_placement_group: Optional[bool]
     use_efa: Optional[bool]
     use_worker_efa: Optional[bool]
+    ami_version: Optional[str]
 
 
 class GCPOptions(BackendOptions, total=False):

coiled/v2/cluster.py

@@ -498,7 +498,7 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
         worker_vm_types: list | None = None,
         worker_cpu: Union[int, List[int]] | None = None,
         worker_memory: Union[str, List[str]] | None = None,
-        worker_disk_size: Union[int, str] | None = None,
+        worker_disk_size: int | str | None = None,
         worker_disk_throughput: int | None = None,
         worker_disk_config: dict | None = None,
         worker_gpu: Union[int, bool] | None = None,
@@ -507,7 +507,7 @@ class Cluster(DistributedCluster, Generic[IsAsynchronous]):
         scheduler_vm_types: list | None = None,
         scheduler_cpu: Union[int, List[int]] | None = None,
         scheduler_memory: Union[str, List[str]] | None = None,
-        scheduler_disk_size: int | None = None,
+        scheduler_disk_size: int | str | None = None,
         scheduler_disk_config: dict | None = None,
         scheduler_gpu: bool | None = None,
         asynchronous: bool = False,

{coiled-1.127.1.dev16.dist-info → coiled-1.127.1.dev20.dist-info}/METADATA

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

{coiled-1.127.1.dev16.dist-info → coiled-1.127.1.dev20.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=LLrkq910sXNtrssPEqK-KU6-aWRS9XuQEaRQw8wPUP4,7023
+coiled/batch.py,sha256=QH-BMlMKkjdToPbw6q0I1W1TTJIDHu24B363mUGDL2c,7102
 coiled/capture_environment.py,sha256=K5mNwUe8KM_l203h3oZvcZUJTrGozT-CH1GrtuPYv8U,18458
 coiled/cluster.py,sha256=wwK9-SefbFBUEHJjYHXlWN3YvPcvR6XD2J-RdPCGhgc,5049
 coiled/coiled.yaml,sha256=z70xzNUy0E8b8Yt12tYYmjJDDmp-U63oUD61ccuu5N0,1037
@@ -22,17 +22,18 @@ 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=mpYmhX9FGoe_pE_GU8Nx2nZXdwbn-DBSHYXl87ferjM,14442
+coiled/types.py,sha256=xJh5t_Kk7S-LeZnZ5C4oTtl1_el3mZuQeITz1QfPHjA,14619
 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
-coiled/cli/core.py,sha256=Yw5g-Y7p75a3k3tVa_jbVhzR6xDyOqwO2yOEB6npUzk,1125
+coiled/cli/core.py,sha256=iZC9v_LyhDOlQwvb5p4w6ul2sfzI7aNiZFHgIWIoNns,1186
 coiled/cli/curl.py,sha256=cYDQXvc1kZNDhFMkREhZoGYh0a9Ea06tEbobK8ZBCJ8,1589
 coiled/cli/diagnostics.py,sha256=1jIeue7xLOaf7LQFsNc6NmO5yU1jqmPFpKZSKjGN4rs,394
 coiled/cli/env.py,sha256=NHh7ZSq9yfongkpFqzon1eLhnH1FwToVvkKFIhqXRBE,6932
 coiled/cli/file.py,sha256=fJmOG3YhxpxXokGYu90wpjdwkJpp1XVqPJ_iveb5ShA,3623
 coiled/cli/login.py,sha256=cByVXmMsfGEuY2TkYU_Y8zq1zVTUHAxOe_wpw2uHsBs,2242
+coiled/cli/mpi.py,sha256=37yPngCYcAbYn35PNVO6-SA3YY6Ik8KVe5blmF3yK4g,5048
 coiled/cli/package_sync.py,sha256=lABDY20yjfLYGfPlQu8ugI-Q8doY4JtN8_0nb9PkcT4,4101
 coiled/cli/prefect.py,sha256=T-SSFey4jlA_jpEI0DqAhVIPwlt2GvBFogEqYCwwevI,302
 coiled/cli/prefect_serve.py,sha256=gemq6YOVbnBoq4k3tSaU2gFJR3aMSxXLNxH6jB8V3n8,4378
@@ -42,7 +43,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=_hFcSETJMhXCqcDXe06wb2Kly4FJanTNjkIPdxO65bg,37077
+coiled/cli/batch/run.py,sha256=jywwmcJbgwfNOZt1eqZRq4a-2n3BZa4zt4fylgz2Grw,37125
 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
@@ -73,7 +74,7 @@ coiled/cli/notebook/__init__.py,sha256=khKDheFK7ociEbb1ODp5oygNLBo_1K7PsL08OaEME
 coiled/cli/notebook/notebook.py,sha256=i_XD03RK2cYeYn_TVl20Uv-kJ_2x-0Oe5iRUTm6w1Tc,23293
 coiled/cli/setup/__init__.py,sha256=BiGnIH9vXGhCFOEPuSUkitcrwAA97wTsfcwMXC0DkYg,837
 coiled/cli/setup/amp.py,sha256=_zlZtqsd_LkSF5C_G8qDm0To-t30C0Z6XKMdDzrm7qg,5039
-coiled/cli/setup/aws.py,sha256=XWzj-YUhbb3KM_PAErnM81KYDRltspN0sxOyRPM7IM0,65447
+coiled/cli/setup/aws.py,sha256=MS4Au1AGoALeVO_VuTdq_RRzL3JzOGgpTgcM69avXU0,65885
 coiled/cli/setup/azure.py,sha256=JNBdKyY9NTcWfr6NAb4d-QM45y6bPBuFFCsCY21equI,25288
 coiled/cli/setup/entry.py,sha256=2PKtvH_ARWt5c5qjeb7dfmJOcFTqRGoskPidNoQTiOg,2425
 coiled/cli/setup/gcp.py,sha256=i67kFRJJpDORrqkVfDu1jFseN80iDbKe1vswk6jxRI8,38817
@@ -87,7 +88,7 @@ coiled/extensions/prefect/__init__.py,sha256=cZp1mqX29FrnINoQsuH6pz4z4uuOACs0mgi
 coiled/extensions/prefect/runners.py,sha256=AcaGS1637TnqFPKnjmmLHpdzjwAsxBLDKrOF7OpfEwM,987
 coiled/extensions/prefect/workers.py,sha256=Z2VuAhTm5AjWEKyCniMZrTxqtkn3uJp3sO3bFeR2Rr0,1642
 coiled/v2/__init__.py,sha256=KaCULaAqatcsYbTbj_SQtTLocbSKZa-uQXiyCICKFRM,805
-coiled/v2/cluster.py,sha256=DelsScilcQ97c7MDT1mBa-CFYHX-5JwdFvWhG2F2gLs,148134
+coiled/v2/cluster.py,sha256=hGs5_SVdrh2zKWpd_8RNPC7LbNd1F6ggEgAj9fCEDm8,148134
 coiled/v2/core.py,sha256=Bf5A_rzK3tuUqqMVAgN5vix-tX_F8AEWR2pICnG3YcA,71615
 coiled/v2/cwi_log_link.py,sha256=d4k6wRYhcdDVdhWYZIX6WL1g0lscXY0yq__H1sPUNWk,1883
 coiled/v2/states.py,sha256=VduyWuf6rByG_wg5AXTxZpe88cCTSdIa4HrPjk1jBcA,9031
@@ -95,8 +96,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.dev16.dist-info/METADATA,sha256=qMOXvhRHA3knlti4m8OYkLrJMiew0sPowR66YqZrhhg,2182
-coiled-1.127.1.dev16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-coiled-1.127.1.dev16.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
-coiled-1.127.1.dev16.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
-coiled-1.127.1.dev16.dist-info/RECORD,,
+coiled-1.127.1.dev20.dist-info/METADATA,sha256=O69puerNR6NRQsz4T2IVMwdxChm9ERigc3HyOH3KrdI,2182
+coiled-1.127.1.dev20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+coiled-1.127.1.dev20.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
+coiled-1.127.1.dev20.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
+coiled-1.127.1.dev20.dist-info/RECORD,,