dagster-cloud 1.8.2__py3-none-any.whl → 1.12.6__py3-none-any.whl

dagster_cloud/workspace/ecs/client.py

@@ -3,7 +3,7 @@ import json
 import logging
 import os
 import time
-from typing import List, Optional
+from typing import Optional
 
 import boto3
 import botocore
@@ -13,7 +13,7 @@ from botocore.exceptions import ClientError
 from dagster._utils.backoff import backoff
 from dagster._utils.cached_method import cached_method
 from dagster_aws.ecs.tasks import DagsterEcsTaskDefinitionConfig
-from dagster_aws.ecs.utils import task_definitions_match
+from dagster_aws.ecs.utils import is_transient_task_stopped_reason, task_definitions_match
 
 from dagster_cloud.workspace.ecs.service import Service
 
@@ -57,13 +57,14 @@ class Client:
         cluster_name: str,
         service_discovery_namespace_id: str,
         log_group: str,
-        subnet_ids: Optional[List[str]] = None,
-        security_group_ids: Optional[List[str]] = None,
+        subnet_ids: Optional[list[str]] = None,
+        security_group_ids: Optional[list[str]] = None,
         ecs_client=None,
         timeout: int = DEFAULT_ECS_TIMEOUT,
         grace_period: int = DEFAULT_ECS_GRACE_PERIOD,
         launch_type: str = "FARGATE",
         show_debug_cluster_info: bool = True,
+        assign_public_ip: Optional[bool] = None,
     ):
         self.ecs = ecs_client if ecs_client else boto3.client("ecs", config=config)
         self.logs = boto3.client("logs", config=config)
@@ -85,6 +86,7 @@ class Client:
         self.grace_period = check.int_param(grace_period, "grace_period")
         self.launch_type = check.str_param(launch_type, "launch_type")
         self._namespace: Optional[str] = None
+        self._assign_public_ip_override = assign_public_ip
 
     @property
     def ec2(self):
@@ -115,14 +117,17 @@ class Client:
     @property
     @cached_method
     def network_configuration(self):
+        if self.launch_type != "FARGATE":
+            assign_public_ip = None
+        elif self._assign_public_ip_override is not None:
+            assign_public_ip = "ENABLED" if self._assign_public_ip_override else "DISABLED"
+        else:
+            assign_public_ip = self._infer_assign_public_ip()
+
         network_configuration = {
             "awsvpcConfiguration": {
                 "subnets": self.subnet_ids,
-                **(
-                    {"assignPublicIp": self._assign_public_ip()}
-                    if self.launch_type == "FARGATE"
-                    else {}
-                ),
+                **({"assignPublicIp": assign_public_ip} if assign_public_ip else {}),
             },
         }
 
@@ -162,7 +167,7 @@ class Client:
             )
             logger.info(f"Created new task definition {task_definition_arn}")
         else:
-            task_definition_arn = existing_task_definition["taskDefinitionArn"]
+            task_definition_arn = check.not_none(existing_task_definition.get("taskDefinitionArn"))
             logger.info(f"Re-using existing task definition {task_definition_arn}")
 
         return task_definition_arn
@@ -333,7 +338,7 @@ class Client:
                 service=service.name,
                 desiredCount=0,
             )
-        except botocore.exceptions.ClientError as error:
+        except botocore.exceptions.ClientError as error:  # pyright: ignore[reportAttributeAccessIssue]
             if error.response["Error"]["Code"] in [
                 "ServiceNotFoundException",
                 "ServiceNotActiveException",
@@ -416,7 +421,7 @@ class Client:
                             if resource_arn in actual_services:
                                 services.append(Service(client=self, arn=resource_arn))
 
-                except botocore.exceptions.ClientError as error:
+                except botocore.exceptions.ClientError as error:  # pyright: ignore[reportAttributeAccessIssue]
                     if error.response["Error"]["Code"] == "AccessDeniedException":
                         self._use_legacy_tag_filtering = True
                         logger.warning(
@@ -487,7 +492,7 @@ class Client:
         )
 
         if exit_code:
-            raise Exception(self.get_task_logs(task_arn))
+            raise Exception(self.get_task_logs(task_arn))  # pyright: ignore[reportCallIssue]
 
         return True
 
@@ -608,7 +613,7 @@ class Client:
 
     async def check_service_has_running_tasks(
         self, service_name, container_name, logger=None
-    ) -> List[str]:
+    ) -> list[str]:
         # return the ARN of the task if it starts
         logger = logger or logging.getLogger("dagster_cloud.EcsClient")
         start_time = time.time()
@@ -622,33 +627,38 @@ class Client:
                     cluster=self.cluster_name,
                     services=[service_name],
                 )
-                if not services:
+                if not services or not services.get("services"):
                     raise Exception(
                         f"Service description not found for {self.cluster_name}/{service_name}"
                     )
 
-                service = services.get("services")[0]
+                service = services["services"][0]
                 desired_count = service.get("desiredCount")
                 running_count = service.get("runningCount")
 
                 # If the service has reached the desired count, we can start tracking the tasks
-                if desired_count == running_count:
+                if desired_count and (desired_count > 0) and (desired_count == running_count):
                     running_tasks = self.ecs.list_tasks(
                         cluster=self.cluster_name,
                         serviceName=service_name,
                         desiredStatus="RUNNING",
                     ).get("taskArns")
 
-                    if not running_tasks:
-                        raise Exception(
-                            f"Unexpected error obtaining tasks for {service_name} in {self.cluster_name}"
-                        )
+                    if running_tasks:
+                        tasks_to_track = running_tasks
+
+            if not tasks_to_track and time.time() > start_time + STOPPED_TASK_GRACE_PERIOD:
+                # If there are still no running_tasks tasks after a certain grace period, check for stopped tasks
+                stopped_tasks = self._check_for_stopped_tasks(service_name)
+                if stopped_tasks:
+                    latest_stopped_task = stopped_tasks[0]
+                    stopped_reason = latest_stopped_task.get("stoppedReason", "")
 
-                    tasks_to_track = running_tasks
-                elif time.time() > start_time + STOPPED_TASK_GRACE_PERIOD:
-                    # If there are still no running_tasks tasks after a certain grace period, check for stopped tasks
-                    stopped_tasks = self._check_for_stopped_tasks(service_name)
-                    if stopped_tasks:
+                    if is_transient_task_stopped_reason(stopped_reason):
+                        logger.warning(
+                            f"Task stopped with a transient stoppedReason: {stopped_reason} - waiting for the service to launch a new task"
+                        )
+                    else:
                         self._raise_failed_task(stopped_tasks[0], container_name, logger)
 
             if tasks_to_track:
@@ -663,7 +673,15 @@ class Client:
                         if not self._check_all_essential_containers_are_running(task):
                             all_tasks_running = False
                     elif task.get("lastStatus") == "STOPPED":
-                        self._raise_failed_task(task, container_name, logger)
+                        stopped_reason = task.get("stoppedReason", "")
+                        if is_transient_task_stopped_reason(stopped_reason):
+                            logger.warning(
+                                f"Running task stopped with a transient stoppedReason: {stopped_reason} - waiting for the service to launch a new task"
+                            )
+                            tasks_to_track = []
+                            all_tasks_running = False
+                        else:
+                            self._raise_failed_task(task, container_name, logger)
 
                 if all_tasks_running:
                     return tasks_to_track
@@ -680,7 +698,7 @@ class Client:
             )
             if response.get("services"):
                 service = response["services"][0]
-                service_events = [event.get("message") for event in service.get("events")]
+                service_events = [str(event.get("message")) for event in service.get("events", [])]
                 service_events_str = "Service events:\n" + "\n".join(service_events)
         except:
             logger.exception(f"Error trying to get service event logs from service {service_name}")
@@ -715,9 +733,9 @@ class Client:
         ).get("taskDefinition")
 
         essential_containers = {
-            container["name"]
-            for container in task_definition["containerDefinitions"]
-            if container["essential"]
+            check.not_none(container.get("name"))
+            for container in task_definition.get("containerDefinitions", [])
+            if container.get("essential") and container.get("name")
         }
 
         # Just because the task is RUNNING doesn't mean everything has started up correctly -
@@ -729,7 +747,7 @@ class Client:
         )
 
     def _get_service_discovery_id(self, hostname):
-        service_name = hostname.split("." + self.namespace)[0]
+        service_name = hostname.split("." + self.namespace)[0]  # pyright: ignore[reportOperatorIssue]
 
         paginator = self.service_discovery.get_paginator("list_services")
         for page in paginator.paginate(
@@ -747,7 +765,7 @@ class Client:
                 if service["Name"] == service_name:
                     return service["Id"]
 
-    def _assign_public_ip(self):
+    def _infer_assign_public_ip(self):
         # https://docs.aws.amazon.com/AmazonECS/latest/userguide/fargate-task-networking.html
         # Assign a public IP if any of the subnets are public
         route_tables = self.ec2.route_tables.filter(

dagster_cloud/workspace/ecs/launcher.py

@@ -1,7 +1,8 @@
 import asyncio
 import os
+from collections.abc import Collection, Mapping, Sequence
 from pathlib import Path
-from typing import Any, Collection, Dict, List, Mapping, Optional, Sequence, cast
+from typing import Any, Optional, cast
 
 import boto3
 import grpc
@@ -30,9 +31,15 @@ from dagster_cloud.workspace.ecs.client import (
     DEFAULT_ECS_TIMEOUT,
     ECS_EXEC_LINUX_PARAMETERS,
     Client,
+    get_debug_ecs_prompt,
 )
+from dagster_cloud.workspace.ecs.run_launcher import CloudEcsRunLauncher
 from dagster_cloud.workspace.ecs.service import Service
-from dagster_cloud.workspace.ecs.utils import get_ecs_human_readable_label, unique_ecs_resource_name
+from dagster_cloud.workspace.ecs.utils import (
+    get_ecs_human_readable_label,
+    get_server_task_definition_family,
+    unique_ecs_resource_name,
+)
 from dagster_cloud.workspace.user_code_launcher import (
     DEFAULT_SERVER_PROCESS_STARTUP_TIMEOUT,
     SHARED_USER_CODE_LAUNCHER_CONFIG,
@@ -44,11 +51,10 @@ from dagster_cloud.workspace.user_code_launcher.user_code_launcher import (
     UserCodeLauncherEntry,
     async_serialize_exceptions,
 )
-from dagster_cloud.workspace.user_code_launcher.utils import deterministic_label_for_location
-
-from .client import get_debug_ecs_prompt
-from .run_launcher import CloudEcsRunLauncher
-from .utils import get_server_task_definition_family
+from dagster_cloud.workspace.user_code_launcher.utils import (
+    deterministic_label_for_location,
+    get_grpc_server_env,
+)
 
 EcsServerHandleType = Service
 
@@ -60,12 +66,12 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
     def __init__(
         self,
         cluster: str,
-        subnets: List[str],
+        subnets: list[str],
         execution_role_arn: str,
         log_group: str,
         service_discovery_namespace_id: str,
         task_role_arn: Optional[str] = None,
-        security_group_ids: Optional[List[str]] = None,
+        security_group_ids: Optional[list[str]] = None,
         inst_data: Optional[ConfigurableClassData] = None,
         secrets=None,
         secrets_tag=None,
@@ -83,6 +89,10 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
         server_ecs_tags: Optional[Sequence[Mapping[str, Optional[str]]]] = None,
         run_ecs_tags: Optional[Sequence[Mapping[str, Optional[str]]]] = None,
         server_health_check: Optional[Mapping[str, Any]] = None,
+        enable_ecs_exec=False,
+        server_task_definition_prefix: str = "server",
+        run_task_definition_prefix: str = "run",
+        assign_public_ip: Optional[bool] = None,
         **kwargs,
     ):
         self.ecs = boto3.client("ecs")
@@ -140,6 +150,22 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             run_sidecar_containers, "run_sidecar_containers"
         )
 
+        self.server_task_definition_prefix = check.str_param(
+            server_task_definition_prefix, "server_task_definition_prefix"
+        )
+        check.invariant(
+            len(self.server_task_definition_prefix) <= 16,
+            "server_task_definition_prefix must be at most 16 characters",
+        )
+        self.run_task_definition_prefix = check.str_param(
+            run_task_definition_prefix, "run_task_definition_prefix"
+        )
+
+        check.invariant(
+            len(self.run_task_definition_prefix) <= 16,
+            "run_task_definition_prefix must be at most 16 characters",
+        )
+
         self.server_ecs_tags = check.opt_sequence_param(server_ecs_tags, "server_ecs_tags")
         self.run_ecs_tags = check.opt_sequence_param(run_ecs_tags, "run_ecs_tags")
 
@@ -147,6 +173,8 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             server_health_check, "server_health_check"
         )
 
+        self._enable_ecs_exec = enable_ecs_exec
+
         self.client = Client(
             cluster_name=self.cluster,
             subnet_ids=self.subnets,
@@ -157,8 +185,9 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             timeout=self._ecs_timeout,
             grace_period=self._ecs_grace_period,
             launch_type=self.launch_type,
+            assign_public_ip=assign_public_ip,
         )
-        super(EcsUserCodeLauncher, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
     @property
     def show_debug_cluster_info(self) -> bool:
@@ -261,13 +290,34 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
                     {"enabled": Field(bool, is_required=False, default_value=False)},
                     is_required=False,
                 ),
+                "enable_ecs_exec": Field(
+                    bool,
+                    is_required=False,
+                    default_value=False,
+                ),
+                "server_task_definition_prefix": Field(
+                    str, is_required=False, default_value="server"
+                ),
+                "run_task_definition_prefix": Field(
+                    str, is_required=False, default_value="dagsterrun"
+                ),
+                "assign_public_ip": Field(
+                    Noneable(bool),
+                    is_required=False,
+                    default_value=None,
+                    description=(
+                        "When using the FARGATE launch type, the launcher will attempt to automatically determine if it is "
+                        "necessary to assign a public IP to the ECS task. In complex network topologies, this automatic "
+                        "determination may not be accurate. In this case, you can explicitly set this value to True or False."
+                    ),
+                ),
             },
             SHARED_ECS_CONFIG,
             SHARED_USER_CODE_LAUNCHER_CONFIG,
         )
 
     @classmethod
-    def from_config_value(cls, inst_data: ConfigurableClassData, config_value: Dict[str, Any]):  # pyright: ignore[reportIncompatibleMethodOverride], fix me!
+    def from_config_value(cls, inst_data: ConfigurableClassData, config_value: dict[str, Any]):  # pyright: ignore[reportIncompatibleMethodOverride], fix me!
         return EcsUserCodeLauncher(inst_data=inst_data, **config_value)
 
     @property
@@ -314,12 +364,12 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
         return container_context.repository_credentials
 
     def _get_enable_ecs_exec(self) -> bool:
-        return False
+        return self._enable_ecs_exec
 
-    def _get_additional_grpc_server_env(self) -> Dict[str, str]:
+    def _get_additional_grpc_server_env(self) -> dict[str, str]:
         return {}
 
-    def _get_dagster_tags(self, deployment_name: str, location_name: str) -> Dict[str, str]:
+    def _get_dagster_tags(self, deployment_name: str, location_name: str) -> dict[str, str]:
         return {
             "dagster/deployment_name": get_ecs_human_readable_label(deployment_name),
             "dagster/location_name": get_ecs_human_readable_label(
@@ -337,7 +387,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             f"Getting resource limits for {deployment_name}:{location_name}. resources: {self.server_resources}"
         )
         metadata = self._actual_entries[(deployment_name, location_name)].code_location_deploy_data
-        resources = metadata.container_context.get("ecs", {}).get("server_resources")
+        resources = metadata.container_context.get("ecs", {}).get("server_resources", {})
         return {
             "ecs": {
                 "cpu_limit": resources.get("cpu"),
@@ -367,8 +417,8 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             command = metadata.get_grpc_server_command(
                 metrics_enabled=self._instance.user_code_launcher.code_server_metrics_enabled
             )
-            additional_env = metadata.get_grpc_server_env(
-                PORT, location_name, self._instance.ref_for_deployment(deployment_name)
+            additional_env = get_grpc_server_env(
+                metadata, PORT, location_name, self._instance.ref_for_deployment(deployment_name)
             )
             tags = {
                 "dagster/grpc_server": "1",
@@ -420,7 +470,10 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
         self._logger.info(f"Creating a new service for {deployment_name}:{location_name}...")
 
         family = get_server_task_definition_family(
-            self._instance.organization_name, deployment_name, location_name
+            self.server_task_definition_prefix,
+            self._instance.organization_name,
+            deployment_name,
+            location_name,
         )
 
         system_tags = {**self._get_dagster_tags(deployment_name, location_name), **tags}
@@ -524,7 +577,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             task_logs = "Task logs:\n" + "\n".join(logs) if logs else "No logs in task."
             sections.append(task_logs)
         except:
-            self._logger.exception("Error trying to get logs for failed task", task_arn=task_arn)
+            self._logger.exception("Error trying to get logs for failed task", task_arn=task_arn)  # pyright: ignore[reportCallIssue]
 
         if self.show_debug_cluster_info:
             sections.append(get_debug_ecs_prompt(self.cluster, task_arn))
@@ -575,7 +628,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             except Exception as e:
                 if (
                     isinstance(e.__cause__, grpc.RpcError)
-                    and cast(grpc.RpcError, e.__cause__).code() == grpc.StatusCode.UNIMPLEMENTED
+                    and cast("grpc.RpcError", e.__cause__).code() == grpc.StatusCode.UNIMPLEMENTED
                 ):
                     # New gRPC method not implemented on old multipex server versions
                     pass
@@ -640,7 +693,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
         }
         return self.client.list_services(tags)
 
-    def _list_server_handles(self) -> List[EcsServerHandleType]:
+    def _list_server_handles(self) -> list[EcsServerHandleType]:
         return [
             service
             for service in self.client.list_services()
@@ -654,7 +707,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
     def get_server_create_timestamp(self, handle: EcsServerHandleType) -> Optional[float]:
         return handle.create_timestamp
 
-    def _run_launcher_kwargs(self) -> Dict[str, Any]:
+    def _run_launcher_kwargs(self) -> dict[str, Any]:
         return dict(
             task_definition={
                 "log_group": self.log_group,
@@ -688,6 +741,7 @@ class EcsUserCodeLauncher(DagsterCloudUserCodeLauncher[EcsServerHandleType], Con
             run_ecs_tags=self.run_ecs_tags,
             container_name=CONTAINER_NAME,
             run_resources=self.run_resources,
+            task_definition_prefix=self.run_task_definition_prefix,
         )
 
     def run_launcher(self) -> CloudEcsRunLauncher:  # pyright: ignore[reportIncompatibleMethodOverride], fix me!

dagster_cloud/workspace/ecs/run_launcher.py

@@ -2,14 +2,14 @@ import dagster._check as check
 from dagster_aws.ecs import EcsRunLauncher
 
 from dagster_cloud.instance import DagsterCloudAgentInstance
-
-from .utils import get_run_task_definition_family
+from dagster_cloud.workspace.ecs.utils import get_run_task_definition_family
 
 
 class CloudEcsRunLauncher(EcsRunLauncher[DagsterCloudAgentInstance]):
     def _get_run_task_definition_family(self, run) -> str:
         return get_run_task_definition_family(
+            self._task_definition_prefix,
             self._instance.organization_name,
             check.not_none(self._instance.deployment_name),
-            check.not_none(run.external_job_origin),
+            check.not_none(run.remote_job_origin),
         )

dagster_cloud/workspace/ecs/utils.py

@@ -2,10 +2,13 @@ import hashlib
 import re
 from typing import Optional
 
-from dagster._core.remote_representation.origin import RemoteJobOrigin
+from dagster._core.remote_origin import RemoteJobOrigin
 from dagster_aws.ecs.utils import sanitize_family
 
-from ..user_code_launcher.utils import get_human_readable_label, unique_resource_name
+from dagster_cloud.workspace.user_code_launcher.utils import (
+    get_human_readable_label,
+    unique_resource_name,
+)
 
 
 def unique_ecs_resource_name(deployment_name, location_name):
@@ -29,10 +32,11 @@ def _get_family_hash(name, max_length=32, hash_size=8):
     m = hashlib.sha1()
     m.update(name.encode("utf-8"))
     name_hash = m.hexdigest()[:hash_size]
-    return f"{name[:(max_length-hash_size-1)]}_{name_hash}"
+    return f"{name[: (max_length - hash_size - 1)]}_{name_hash}"
 
 
 def get_server_task_definition_family(
+    task_definition_prefix: str,
     organization_name: Optional[str],
     deployment_name: str,
     location_name: str,
@@ -43,9 +47,12 @@ def get_server_task_definition_family(
     m = hashlib.sha1()
     m.update(location_name.encode("utf-8"))
 
+    # '{16}_{64}_{64}_{64}': max 211 characters
     truncated_location_name = _get_family_hash(location_name, max_length=64)
 
-    final_family = f"server_{organization_name}_{deployment_name}_{truncated_location_name}"
+    final_family: str = (
+        f"{task_definition_prefix}_{organization_name}_{deployment_name}_{truncated_location_name}"
+    )
 
     assert len(final_family) <= 255
 
@@ -53,6 +60,7 @@ def get_server_task_definition_family(
 
 
 def get_run_task_definition_family(
+    task_definition_prefix: str,
     organization_name: Optional[str],
     deployment_name: str,
     job_origin: RemoteJobOrigin,
@@ -64,12 +72,13 @@ def get_run_task_definition_family(
     repo_name = job_origin.repository_origin.repository_name
     location_name = job_origin.repository_origin.code_location_origin.location_name
 
+    assert len(task_definition_prefix) <= 16
     assert len(str(organization_name)) <= 64
     assert len(deployment_name) <= 64
 
     # '{16}_{64}_{64}_{32}_{32}_{32}': max 245 characters
 
-    final_family = f"run_{organization_name}_{deployment_name}_{_get_family_hash(location_name)}_{_get_family_hash(repo_name)}_{_get_family_hash(job_name)}"
+    final_family = f"{task_definition_prefix}_{organization_name}_{deployment_name}_{_get_family_hash(location_name)}_{_get_family_hash(repo_name)}_{_get_family_hash(job_name)}"
 
     assert len(final_family) <= 255
 

dagster_cloud/workspace/kubernetes/__init__.py

@@ -1 +1 @@
-from .launcher import K8sUserCodeLauncher as K8sUserCodeLauncher
+from dagster_cloud.workspace.kubernetes.launcher import K8sUserCodeLauncher as K8sUserCodeLauncher