coiled 1.131.1.dev2__tar.gz → 1.133.0__tar.gz

{coiled-1.131.1.dev2 → coiled-1.133.0}/PKG-INFO

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

{coiled-1.131.1.dev2 → coiled-1.133.0}/coiled/cli/notebook/notebook.py

@@ -155,11 +155,36 @@ def check_jupyter() -> bool:
         "allows you to use an io2 EBS volume on AWS."
     ),
 )
+@click.option(
+    "--capacity-reservation-id",
+    default=None,
+    help="For AWS, allows you to specify existing Capacity Reservation to use.",
+)
+@click.option(
+    "--capacity-reservation-group-arn",
+    default=None,
+    help="For AWS, allows you to specify existing Capacity Reservation group to use.",
+)
+@click.option(
+    "--docker-shm-size",
+    default=None,
+    help="Non-default value for shm_size (for example, '3 GiB').",
+)
+@click.option(
+    "--ami-version",
+    default=None,
+    help="Use non-default Coiled AMI, e.g. 'DL' for the Deep Learning Base OSS Nvidia Driver GPU AMI.",
+)
 @click.option(
     "--region",
     default=None,
     help="The cloud provider region in which to run the notebook.",
 )
+@click.option(
+    "--zone",
+    default=None,
+    help="The cloud provider zone in which to run the notebook.",
+)
 @click.option(
     "--open",
     default=True,
@@ -256,7 +281,12 @@ def start_notebook(
     gpu: bool,
     disk_size: int | None,
     disk_config: str | None,
+    capacity_reservation_id: str | None,
+    capacity_reservation_group_arn: str | None,
+    docker_shm_size: str | None,
+    ami_version: str | None,
     region: str | None,
+    zone: str | None,
     open: bool,
     block: bool,
     include_vcs: bool,
@@ -294,7 +324,12 @@ def start_notebook(
         gpu=gpu,
         disk_size=disk_size,
         disk_config=disk_config_dict,
+        capacity_reservation_id=capacity_reservation_id,
+        capacity_reservation_group_arn=capacity_reservation_group_arn,
+        docker_shm_size=docker_shm_size,
+        ami_version=ami_version,
         region=region,
+        zone=zone,
         open=open,
         block=block,
         include_vcs=include_vcs,
@@ -324,7 +359,12 @@ def _start_notebook(
     gpu: bool = False,
     disk_size: int | None = None,
     disk_config: dict | None = None,
+    capacity_reservation_id: str | None = None,
+    capacity_reservation_group_arn: str | None = None,
+    docker_shm_size: str | None = None,
+    ami_version: str | None = None,
     region: str | None = None,
+    zone: str | None = None,
     open: bool = False,
     block: bool = True,
     include_vcs: bool = False,
@@ -399,6 +439,18 @@ def _start_notebook(
             except FileNotFoundError:
                 package_sync_ignore = []
 
+            backend_options = dict()
+            if zone:
+                backend_options["zone_name"] = zone
+            if capacity_reservation_id:
+                backend_options["capacity_reservation_id"] = capacity_reservation_id
+            if capacity_reservation_group_arn:
+                backend_options["capacity_reservation_group_arn"] = capacity_reservation_group_arn
+            if ami_version:
+                backend_options["ami_version"] = ami_version
+            if docker_shm_size:
+                backend_options["docker_shm_size"] = docker_shm_size
+
             cluster = coiled.Cluster(
                 name=name,
                 workspace=workspace,
@@ -428,6 +480,7 @@ def _start_notebook(
                 unset_single_threading_variables=True,
                 pause_on_exit=resumable,
                 private_to_creator=private,
+                backend_options=backend_options or None,  # type: ignore
             )
             info["cluster_id"] = cluster.cluster_id
 

{coiled-1.131.1.dev2 → coiled-1.133.0}/coiled/cli/run.py

@@ -576,6 +576,11 @@ def get_entrypoint(connection, container_name) -> str:
     default=None,
     help="Don't automatically mount SSD/NVMe (if any) to ``/scratch``.",
 )
+@click.option(
+    "--ami-version",
+    default=None,
+    help="Use non-default Coiled AMI, e.g. 'DL' for the Deep Learning Base OSS Nvidia Driver GPU AMI.",
+)
 @click.option(
     "--filestore",
     "filestore_names",
@@ -614,6 +619,7 @@ def run(
     package_sync_conda_extras,
     docker_shm_size,
     skip_ssd_mount,
+    ami_version,
     filestore_names,
     command,
 ):
@@ -650,6 +656,7 @@ def run(
         package_sync_strict=package_sync_strict,
         package_sync_conda_extras=package_sync_conda_extras,
         docker_shm_size=docker_shm_size,
+        ami_version=ami_version,
         filestore_names=filestore_names,
         skip_ssd_mount=skip_ssd_mount,
     )
@@ -691,6 +698,7 @@ def start_run(
     docker_shm_size: str | None = None,
     filestore_names: list[str] | None = None,
     skip_ssd_mount: bool | None = None,
+    ami_version: str | None = None,
 ):
     runtime_env_dict = dict_from_key_val_list(env)
     tags = dict_from_key_val_list(tag)
@@ -761,6 +769,13 @@ def start_run(
                     workspace=workspace,
                 )
                 info["workspace"] = workspace
+
+                backend_options = {}
+                if ami_version:
+                    backend_options["ami_version"] = ami_version
+                if docker_shm_size:
+                    backend_options["docker_shm_size"] = docker_shm_size
+
                 cluster_kwargs: ClusterKwargs = {
                     "workspace": workspace,
                     "n_workers": 0,
@@ -785,7 +800,7 @@ def start_run(
                     "host_setup_script": host_setup_script,
                     "package_sync_strict": package_sync_strict,
                     "package_sync_conda_extras": package_sync_conda_extras,
-                    "backend_options": {"docker_shm_size": docker_shm_size} if docker_shm_size else None,
+                    "backend_options": backend_options or None,  # type: ignore
                     "unset_single_threading_variables": True,
                     "skip_ssd_mount": skip_ssd_mount,
                     "filestores_to_attach": filestores_to_attach,

{coiled-1.131.1.dev2 → coiled-1.133.0}/coiled/types.py

@@ -447,6 +447,10 @@ class AWSOptions(BackendOptions, total=False):
     ami_version
         Use non-default type of AMI.
         Supported options include "DL" for the Deep Learning Base OSS Nvidia Driver GPU AMI.
+    capacity_reservation_id
+        Use specified existing Capacity Reservation.
+    capacity_reservation_group_arn
+        Use specified existing Capacity Reservation group.
     """
 
     keypair_name: Optional[str]
@@ -455,6 +459,8 @@ class AWSOptions(BackendOptions, total=False):
     use_efa: Optional[bool]
     use_worker_efa: Optional[bool]
     ami_version: Optional[str]
+    capacity_reservation_id: Optional[str]
+    capacity_reservation_group_arn: Optional[str]
 
 
 class GCPOptions(BackendOptions, total=False):