coiled 1.118.4.dev6__py3-none-any.whl → 1.129.3.dev10__py3-none-any.whl

coiled/cli/setup/aws.py

@@ -240,7 +240,7 @@ setup_doc = """{
 }"""
 
 
-def get_ongoing_doc(ecr=True, custom_bucket_prefix=None) -> str:
+def get_ongoing_doc(ecr=True, package_sync_bucket_prefix=None) -> str:
     optional_permissions = []
 
     if ecr:
@@ -266,7 +266,7 @@ def get_ongoing_doc(ecr=True, custom_bucket_prefix=None) -> str:
             ],
         })
 
-    if custom_bucket_prefix:
+    if package_sync_bucket_prefix:
         optional_permissions.extend([
             {
                 "Sid": "OngoingPackageSyncBucketCreate",
@@ -277,7 +277,7 @@ def get_ongoing_doc(ecr=True, custom_bucket_prefix=None) -> str:
                     "s3:PutBucketOwnershipControls",
                     "s3:PutBucketPolicy",
                 ],
-                "Resource": [f"arn:*:s3:::{custom_bucket_prefix}-*"],
+                "Resource": [f"arn:*:s3:::{package_sync_bucket_prefix}-*"],
             },
             {
                 "Sid": "OngoingPackageSyncEnvUploadDownload",
@@ -286,10 +286,34 @@ def get_ongoing_doc(ecr=True, custom_bucket_prefix=None) -> str:
                     "s3:PutObject",
                     "s3:GetObject",
                 ],
-                "Resource": [f"arn:*:s3:::{custom_bucket_prefix}-*/*"],
+                "Resource": [f"arn:*:s3:::{package_sync_bucket_prefix}-*/*"],
             },
         ])
 
+    # TODO should this be configurable?
+    optional_permissions.extend([
+        {
+            "Sid": "OngoingPersistentDataBuckets",
+            "Effect": "Allow",
+            "Action": [
+                "s3:CreateBucket",
+                "s3:ListBucket",
+                "s3:PutBucketOwnershipControls",
+                "s3:PutBucketPolicy",
+            ],
+            "Resource": ["arn:*:s3:::coiled-data-*"],
+        },
+        {
+            "Sid": "OngoingPersistentDataObjects",
+            "Effect": "Allow",
+            "Action": [
+                "s3:PutObject",
+                "s3:GetObject",
+            ],
+            "Resource": ["arn:*:s3:::coiled-data-*/*"],
+        },
+    ])
+
     ongoing_policy_statement = {
         "Statement": [
             {
@@ -297,6 +321,7 @@ def get_ongoing_doc(ecr=True, custom_bucket_prefix=None) -> str:
                 "Effect": "Allow",
                 "Resource": "*",
                 "Action": [
+                    "ec2:AuthorizeSecurityGroupEgress",
                     "ec2:AuthorizeSecurityGroupIngress",
                     "ec2:CreateFleet",
                     "ec2:CreateLaunchTemplate",
@@ -361,6 +386,16 @@ def get_ongoing_doc(ecr=True, custom_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",
@@ -1393,7 +1428,7 @@ def do_setup(
         user_name = slug
         setup_name = setup_name or f"{slug}-setup"
         ongoing_name = ongoing_name or f"{slug}-ongoing"
-        ongoing_doc = get_ongoing_doc(custom_bucket_prefix=custom_s3_bucket_prefix)
+        ongoing_doc = get_ongoing_doc(package_sync_bucket_prefix=custom_s3_bucket_prefix)
 
         try:
             aws_account = do_intro(sts, iam, region=region, coiled_account=coiled_account)
@@ -1508,26 +1543,27 @@ def check_local_aws_creds():
         return False
 
 
-def update_instance_profile_policy(iam, aws_account, yes):
-    policy_name = "CoiledInstancePolicy"
-    policy_doc = """{
+INSTANCE_POLICY_DOCUMENT = """{
     "Version": "2012-10-17",
     "Statement": [
         {
-            "Sid": "CoiledEC2Policy",
+            "Sid": "CoiledEC2LogPolicy",
             "Effect": "Allow",
             "Action": [
                 "logs:CreateLogGroup",
                 "logs:CreateLogStream",
-                "logs:PutLogEvents",
-                "cloudwatch:PutMetricData",
-                "aps:RemoteWrite"
+                "logs:PutLogEvents"
             ],
             "Resource": "*"
         }
     ]
 }"""
 
+
+def update_instance_profile_policy(iam, aws_account, yes):
+    policy_name = "CoiledInstancePolicy"
+    policy_doc = INSTANCE_POLICY_DOCUMENT
+
     # Check for existing policy
     policy_arn, policy_diff = get_policy_diff(iam, aws_account, policy_name, policy_doc)
 

coiled/cli/setup/azure.py

@@ -175,8 +175,21 @@ coiled curl -X POST "${SETUP_ENDPOINT}" --json --data "{\\"credentials\\": {\\"t
         "service principal for Coiled to use."
     ),
 )
+@click.option(
+    "--refresh-for-app-id", default=None, help="Refresh the secret key used by Coiled for specified Application ID."
+)
 @click.command(context_settings=CONTEXT_SETTINGS)
-def azure_setup(subscription, resource_group, region, account, iam_user, keep_existing_access, save_script, ship_token):
+def azure_setup(
+    subscription,
+    resource_group,
+    region,
+    account,
+    iam_user,
+    keep_existing_access,
+    save_script,
+    ship_token,
+    refresh_for_app_id,
+):
     print(
         "Coiled on Azure is currently in [bold]public beta[/bold], "
         "please contact [link]support@coiled.io[/link] if you have any questions or problems."
@@ -264,6 +277,17 @@ def azure_setup(subscription, resource_group, region, account, iam_user, keep_ex
         f"with [green]{region}[/green] as the default region\n"
     )
 
+    if refresh_for_app_id:
+        refresh_app_creds(
+            app_id=refresh_for_app_id,
+            coiled_account=coiled_account,
+            sub_id=sub_id,
+            rg_name=rg_name,
+            region=region,
+            keep_existing_keys=True,
+        )
+        return
+
     if ship_token:
         enable_providers(creds, sub_id)
         ship_token_creds(
@@ -378,6 +402,31 @@ def setup_with_service_principal(
     return True
 
 
+def refresh_app_creds(app_id, keep_existing_keys, coiled_account, sub_id, rg_name, region):
+    print(f"  [bright_black]Resetting/retrieving credentials for {app_id}...")
+    cred_reset_opts = "--append" if keep_existing_keys else ""
+    app_creds_json = az_cli_wrapper(f"az ad app credential reset --id {app_id} {cred_reset_opts}")
+    app_creds = json.loads(app_creds_json)
+
+    creds_to_submit = {
+        "tenant_id": app_creds["tenant"],
+        "client_id": app_creds["appId"],
+        "client_secret": app_creds["password"],
+    }
+
+    print("Sending Azure credentials to Coiled... ", end="")
+    submit_azure_credentials(
+        coiled_account=coiled_account,
+        sub_id=sub_id,
+        rg_name=rg_name,
+        region=region,
+        creds_to_submit=creds_to_submit,
+    )
+    print(f"Azure credentials have been updated for {coiled_account} using Azure app {app_id} and {rg_name}!")
+
+    coiled.add_interaction(action="RefreshCloudCredentials", success=True)
+
+
 def submit_azure_credentials(coiled_account, sub_id, rg_name, region, creds_to_submit, check_after: bool = False):
     with coiled.Cloud(account=coiled_account) as cloud:
         setup_endpoint = f"/api/v2/cloud-credentials/{coiled_account}/azure"

coiled/context.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-import asyncio
 import functools
+import inspect
 import random
 import string
 from contextlib import contextmanager, nullcontext
@@ -103,7 +103,7 @@ def get_trace_context(func: Union[SyncFuncType, AsyncFuncType]):
 
 
 def track_context(func: F) -> F:
-    if asyncio.iscoroutinefunction(func):
+    if inspect.iscoroutinefunction(func):
 
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:

coiled/credentials/google.py

@@ -4,8 +4,6 @@ import logging
 import os
 from typing import Tuple, cast
 
-from distributed import Client
-
 from coiled.utils import supress_logs
 
 
@@ -128,25 +126,8 @@ def get_long_lived_adc_to_forward(scopes=None):
 def send_application_default_credentials(cluster, scopes=None, return_creds: bool = False):
     to_forward = get_long_lived_adc_to_forward(scopes)
 
-    def write_adf_files():
-        for path, content in to_forward["files"].items():
-            abs_path = os.path.expanduser(path)
-            os.makedirs(os.path.dirname(abs_path), exist_ok=True)
-            with open(abs_path, "w") as f:
-                f.write(content)
-
-        os.makedirs(os.path.expanduser("~/.config"), exist_ok=True)
-        container_adc_dir = os.path.expanduser("~/.config/gcloud")
-
-        # files are written to /gcloud-config, which is then symlinked to ~/.config/gcloud
-        if not os.path.exists(container_adc_dir):
-            os.symlink("/gcloud-config", container_adc_dir, target_is_directory=True)
-        # otherwise, maybe we should copy the files to both places?
-
     cluster.send_private_envs(to_forward["env"])
-    with Client(cluster, name="non-user-write-gcp-adf") as client:
-        client.run_on_scheduler(write_adf_files)
-        client.run(write_adf_files)
+    cluster.write_files_for_dask(files=to_forward["files"], symlink_dirs={"/gcloud-config": "~/.config/gcloud"})
 
     cluster._queue_cluster_event("credentials", "Google Application Default Credentials forwarded")
     print(