coiled 1.128.0__py3-none-any.whl → 1.128.1__py3-none-any.whl

coiled/cli/mpi.py

@@ -6,8 +6,10 @@ import fabric.connection
 
 import coiled
 
+from .. import AWSOptions
+from ..v2.cluster_comms import get_cluster_connection_info, get_comm_from_connection_info
 from .cluster.utils import find_cluster
-from .run import get_ssh_connection, write_via_ssh
+from .run import KeepaliveSession, get_ssh_connection, write_via_ssh
 from .utils import CONTEXT_SETTINGS
 
 
@@ -15,32 +17,73 @@ from .utils import CONTEXT_SETTINGS
     context_settings=CONTEXT_SETTINGS,
 )
 @click.option("--worker-nodes", default=1, type=int)
+@click.option(
+    "--include-head/--exclude-head",
+    default=True,
+    type=bool,
+    help="Include head (scheduler) node in placement group with EFA",
+)
 @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)
+@click.option("--head-vm-type", default=None, type=str)
+@click.option("--worker-vm-type", default=None, type=str)
+@click.option("--pip", multiple=True, type=str, help="Packages to install with pip")
+@click.option("--apt", multiple=True, type=str, help="Packages to install with apt install")
+@click.option("--setup-script", default=None, type=str, help="Path to additional host setup script")
+@click.option("--cluster-timeout", default="1h", type=str, help="Maximum lifetime for cluster")
+@click.option(
+    "--idle-timeout",
+    default="10m",
+    type=str,
+    help="How long to wait when nothing is running before shutting down cluster",
+)
+def setup(
+    worker_nodes,
+    include_head,
+    vm_type,
+    head_vm_type,
+    worker_vm_type,
+    pip,
+    apt,
+    setup_script,
+    cluster_timeout,
+    idle_timeout,
+):
+    additional_host_setup = None
+    if setup_script:
+        with open(setup_script) as f:
+            additional_host_setup = f.read()
+    setup_script = get_host_setup_script(pip_install=pip, apt_install=apt, additional_setup=additional_host_setup)
+
+    efa_settings = (
+        {"use_placement_group": True, "use_efa": True}
+        if include_head
+        else {"use_worker_placement_group": True, "use_worker_efa": True}
+    )
 
     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],
+        backend_options=AWSOptions(**{**efa_settings, "ami_version": "DL"}),
+        scheduler_vm_types=[head_vm_type or vm_type],
+        worker_vm_types=[worker_vm_type or vm_type],
         worker_disk_size="100GB",
         scheduler_disk_size="100GB",
         shutdown_on_close=False,
-        idle_timeout=idle_timeout,
+        idle_timeout="520 weeks",  # don't use idle timeout
+        cluster_timeout=cluster_timeout,
     )
 
+    cluster.set_keepalive(keepalive=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)
+    with KeepaliveSession(cluster=cluster, prefix="mpi-ssh-setup"):
+        setup_mpi_ssh(connection)
 
     print("MPI is ready")
 
@@ -88,6 +131,8 @@ legate \
         cluster_info = find_cluster(cloud, cluster)
         cluster_id = cluster_info["id"]
         connection = get_ssh_connection(cloud, cluster_id)
+        address, security = get_cluster_connection_info(cluster_id, cloud)
+        comms = get_comm_from_connection_info(address, security)
 
     setup_mpi_ssh(connection, include_scheduler=include_head)
 
@@ -116,8 +161,10 @@ legate \
 
     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"})
+    with KeepaliveSession(cluster=None, comms=comms, prefix="mpi-ssh"):
+        connection.run(wrapped_command, hide=False, pty=True, warn=True, env={"PATH": "/tmp/host-user-venv/bin:$PATH"})
+
+    comms.close_comms()
 
 
 def setup_mpi_ssh(connection, include_scheduler=True):
@@ -138,11 +185,11 @@ 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 []
+def get_host_setup_script(venv_path="/tmp/host-user-venv", apt_install=None, pip_install=None, additional_setup=None):
+    apt_install = list(apt_install or [])
     apt_install.extend(["openmpi-bin", "python3-pip", "python3-venv"])
 
-    pip_install = pip_install or []
+    pip_install = list(pip_install or [])
 
     pip_install_line = f"{venv_path}/bin/python -m pip install {' '.join(pip_install)}" if pip_install else ""
 
@@ -153,6 +200,7 @@ mkdir {venv_path}
 python3 -m venv {venv_path}
 
 {pip_install_line}
+{additional_setup or ""}
 
 echo 'done' > /tmp/host-setup-done
     """

coiled/cli/run.py

@@ -29,6 +29,8 @@ from coiled.utils import (
     unset_single_thread_defaults,
 )
 from coiled.v2.cluster import ClusterKwargs
+from coiled.v2.cluster_comms import use_comm_rpc
+from coiled.v2.core import Cloud
 from coiled.v2.widgets.rich import LightRichClusterWidget
 
 from ..filestore import FilestoreManager
@@ -46,20 +48,37 @@ USER_CONTAINER_NAME = "tmp-user-1"
 
 
 class KeepaliveSession:
-    def __init__(self, cluster, prefix="", monitor_proc_activity=False):
+    def __init__(self, cluster, comms=None, prefix="", monitor_proc_activity=False):
         self.cluster = cluster
+        self.comms = comms
         self.monitor_proc_activity = monitor_proc_activity
         rand_uuid = short_random_string()
         self.session_id = f"{prefix}-{rand_uuid}" if prefix else rand_uuid
 
+        if self.comms:
+            self.cloud = Cloud.current(asynchronous=False)
+
     def __enter__(self):
         # keepalive session lets us use keepalive without dask client
-        self.cluster._call_scheduler_comm(
-            "coiled_add_keepalive_session", name=self.session_id, monitor_proc_activity=self.monitor_proc_activity
-        )
+        if self.cluster:
+            self.cluster._call_scheduler_comm(
+                "coiled_add_keepalive_session", name=self.session_id, monitor_proc_activity=self.monitor_proc_activity
+            )
+        elif self.comms:
+            use_comm_rpc(
+                self.cloud,
+                self.comms,
+                "coiled_add_keepalive_session",
+                name=self.session_id,
+                monitor_proc_activity=self.monitor_proc_activity,
+            )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.cluster._call_scheduler_comm("coiled_end_keepalive_session", name=self.session_id)
+        if self.cluster:
+            self.cluster._call_scheduler_comm("coiled_end_keepalive_session", name=self.session_id)
+        elif self.comms:
+            use_comm_rpc(self.cloud, self.comms, "coiled_end_keepalive_session", name=self.session_id)
+            self.cloud.close()
 
 
 def get_ssh_connection(cloud, cluster_id) -> fabric.connection.Connection:

coiled/v2/cluster_comms.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from coiled.v2.states import ProcessStateEnum
+
+
+def get_cluster_connection_info(
+    cluster_id: int,
+    cloud,
+    *,
+    use_scheduler_public_ip: bool = True,
+) -> tuple[str, dict]:
+    """
+    Get the comms info we need to connect to Dask in a running cluster.
+
+    (This is a bit of a hack. It would be nicer to have a way to tell coiled.Cluster not to
+    create, just retrieve. But Cluster is a bit hard to deal with... )
+    """
+
+    cluster_info = cloud._get_cluster_details_synced(cluster_id=cluster_id)  # type: ignore
+
+    if ProcessStateEnum(cluster_info["scheduler"]["current_state"]["state"]) != ProcessStateEnum.started:
+        scheduler_state = cluster_info["scheduler"]["current_state"]["state"]
+        raise RuntimeError(f"Cannot get security info for cluster {cluster_id}, scheduler state is {scheduler_state}")
+
+    public_ip = cluster_info["scheduler"]["instance"]["public_ip_address"]
+    private_ip = cluster_info["scheduler"]["instance"]["private_ip_address"]
+    tls_cert = cluster_info["cluster_options"]["tls_cert"]
+    tls_key = cluster_info["cluster_options"]["tls_key"]
+    scheduler_port = cluster_info["scheduler_port"]
+    dashboard_address = cluster_info["scheduler"]["dashboard_address"]
+    give_scheduler_public_ip = cluster_info["cluster_infra"]["give_scheduler_public_ip"]
+
+    private_address = f"tls://{private_ip}:{scheduler_port}"
+    public_address = f"tls://{public_ip}:{scheduler_port}"
+
+    use_public_address = give_scheduler_public_ip and use_scheduler_public_ip
+    if use_public_address:
+        if not public_ip:
+            raise RuntimeError(
+                "Your Coiled client is configured to use the public IP address, but the scheduler VM does not "
+                "have a public IP address."
+            )
+        address_to_use = public_address
+    else:
+        address_to_use = private_address
+
+    security_info = {
+        "tls_key": tls_key,
+        "tls_cert": tls_cert,
+        "dashboard_address": dashboard_address,
+        "public_address": public_address,
+        "private_address": private_address,
+        "address_to_use": address_to_use,
+    }
+
+    return address_to_use, security_info
+
+
+def get_comm_from_connection_info(address, security):
+    from distributed import rpc
+
+    from coiled.utils import GatewaySecurity
+
+    security_obj = GatewaySecurity(security["tls_key"], security["tls_cert"])
+    return rpc(address, connection_args=security_obj.get_connection_args("client"))
+
+
+def use_comm_rpc(cloud, comm, function, **kwargs):
+    async def foo():
+        await getattr(comm, function)(**kwargs)
+
+    cloud._sync(foo)

{coiled-1.128.0.dist-info → coiled-1.128.1.dist-info}/METADATA

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

{coiled-1.128.0.dist-info → coiled-1.128.1.dist-info}/RECORD

@@ -33,11 +33,11 @@ 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/mpi.py,sha256=_i5GPTu4kE1kc-DQ1JLixIEpVVNJ0suwyLUgjp-iUww,6793
 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
-coiled/cli/run.py,sha256=IHy-7a9RCUkNQauugqZYij_B-qMiMEk1gCJSf_EHBdc,30898
+coiled/cli/run.py,sha256=CtUmIarGMBC4n2xpbplStTrGGgGZiP5mTTKXOoEN_yo,31600
 coiled/cli/sync.py,sha256=S5PzB9GSPJn3HvviOMLKVbo4ET46FlPwLYK_7sRyonQ,9726
 coiled/cli/utils.py,sha256=cp7ToFGRpUKi6iNL6BbLjzgrgeTYSX_C55lYhaKWHHA,3479
 coiled/cli/batch/__init__.py,sha256=539CnfnqqcW7ndSufTS-Ie5FGZiElMYxE0Ptu70wo8M,660
@@ -89,6 +89,7 @@ coiled/extensions/prefect/runners.py,sha256=AcaGS1637TnqFPKnjmmLHpdzjwAsxBLDKrOF
 coiled/extensions/prefect/workers.py,sha256=Z2VuAhTm5AjWEKyCniMZrTxqtkn3uJp3sO3bFeR2Rr0,1642
 coiled/v2/__init__.py,sha256=KaCULaAqatcsYbTbj_SQtTLocbSKZa-uQXiyCICKFRM,805
 coiled/v2/cluster.py,sha256=hGs5_SVdrh2zKWpd_8RNPC7LbNd1F6ggEgAj9fCEDm8,148134
+coiled/v2/cluster_comms.py,sha256=UcJWLeZlc68S0uaNd9lLKbF5uaDhYqqkdTsA0CBXYRI,2643
 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
@@ -96,8 +97,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.128.0.dist-info/METADATA,sha256=3Vk38O2A51tSoz3OfVKlWQArARO9aRTTLn_rvGNBBS8,2176
-coiled-1.128.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-coiled-1.128.0.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
-coiled-1.128.0.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
-coiled-1.128.0.dist-info/RECORD,,
+coiled-1.128.1.dist-info/METADATA,sha256=ju4vAo9MSxF8AdwdzeX8wZX2IdsCIrARVEafamEoUbo,2176
+coiled-1.128.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+coiled-1.128.1.dist-info/entry_points.txt,sha256=C8dz1ST_bTlTO-kNvuHBJQma9PyJPotg0S4xpPt5aHY,47
+coiled-1.128.1.dist-info/licenses/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
+coiled-1.128.1.dist-info/RECORD,,