dagster-aws 0.28.0__py3-none-any.whl → 0.28.2__py3-none-any.whl

dagster_aws/ecs/launcher.py

@@ -5,7 +5,7 @@ import uuid
 import warnings
 from collections import namedtuple
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import boto3
 from botocore.exceptions import ClientError
@@ -30,11 +30,12 @@ from dagster._core.launcher.base import (
     WorkerStatus,
 )
 from dagster._core.storage.dagster_run import DagsterRun
-from dagster._core.storage.tags import RUN_WORKER_ID_TAG
+from dagster._core.storage.tags import HIDDEN_TAG_PREFIX, RUN_WORKER_ID_TAG
 from dagster._grpc.types import ExecuteRunArgs
 from dagster._serdes import ConfigurableClass
 from dagster._serdes.config_class import ConfigurableClassData
 from dagster._utils.backoff import backoff
+from dagster._utils.tags import get_boolean_tag_value
 from typing_extensions import Self
 
 from dagster_aws.ecs.container_context import (
@@ -81,6 +82,9 @@ TAGS_TO_EXCLUDE_FROM_PROPAGATION = {"dagster/op_selection", "dagster/solid_selec
 DEFAULT_REGISTER_TASK_DEFINITION_RETRIES = 5
 DEFAULT_RUN_TASK_RETRIES = 5
 
+if TYPE_CHECKING:
+    from botocore.config import Config
+
 
 class EcsRunLauncher(RunLauncher[T_DagsterInstance], ConfigurableClass):
     """RunLauncher that starts a task in ECS for each Dagster job run.
@@ -119,10 +123,13 @@ class EcsRunLauncher(RunLauncher[T_DagsterInstance], ConfigurableClass):
         task_definition_prefix: str = "run",
     ):
         self._inst_data = inst_data
-        self.ecs = boto3.client("ecs")
+
+        boto_client_config = self.get_boto_client_config()
+
+        self.ecs = boto3.client("ecs", config=boto_client_config)
         self.ec2 = boto3.resource("ec2")
-        self.secrets_manager = boto3.client("secretsmanager")
-        self.logs = boto3.client("logs")
+        self.secrets_manager = boto3.client("secretsmanager", config=boto_client_config)
+        self.logs = boto3.client("logs", config=boto_client_config)
 
         self._task_definition_prefix = task_definition_prefix
 
@@ -236,6 +243,9 @@ class EcsRunLauncher(RunLauncher[T_DagsterInstance], ConfigurableClass):
         self._current_task_metadata = None
         self._current_task = None
 
+    def get_boto_client_config(self) -> Optional["Config"]:
+        return None
+
     @property
     def inst_data(self):
         return self._inst_data
@@ -855,6 +865,28 @@ class EcsRunLauncher(RunLauncher[T_DagsterInstance], ConfigurableClass):
             task.get("stoppedReason", "")
         )
 
+    def _add_eni_id_tags(self, run: DagsterRun, task: dict[str, Any]):
+        attachments = task.get("attachments", [])
+        eni_ids = {}
+        eni_count = 0
+        for attachment in attachments:
+            if attachment.get("type") == "ElasticNetworkInterface":
+                details = {d["name"]: d["value"] for d in attachment.get("details", [])}
+
+                if "networkInterfaceId" in details:
+                    if eni_count == 0:
+                        eni_ids[f"{HIDDEN_TAG_PREFIX}eni_id"] = details["networkInterfaceId"]
+                    else:
+                        eni_ids[f"{HIDDEN_TAG_PREFIX}eni_id_{eni_count}"] = details[
+                            "networkInterfaceId"
+                        ]
+                    eni_count += 1
+        self._instance.add_run_tags(run.run_id, eni_ids)
+        if eni_count > 0:
+            logging.info(f"Added {eni_count} ENI ID tags for run {run.run_id}: {eni_ids}")
+        else:
+            logging.warning(f"No ENI IDs found for run {run.run_id}")
+
     def check_run_worker_health(self, run: DagsterRun):
         run_worker_id = run.tags.get(RUN_WORKER_ID_TAG)
 
@@ -870,6 +902,14 @@ class EcsRunLauncher(RunLauncher[T_DagsterInstance], ConfigurableClass):
 
         t = tasks[0]
 
+        if get_boolean_tag_value(os.getenv("DAGSTER_AWS_ENI_TAGGING_ENABLED")) and not run.tags.get(
+            f"{HIDDEN_TAG_PREFIX}eni_id"
+        ):
+            try:
+                self._add_eni_id_tags(run, t)
+            except Exception:
+                logging.exception(f"Error adding ENI ID tags for run {run.run_id}")
+
         if t.get("lastStatus") in RUNNING_STATUSES:
             return CheckRunHealthResult(WorkerStatus.RUNNING, run_worker_id=run_worker_id)
         elif t.get("lastStatus") in STOPPED_STATUSES:

dagster_aws/pipes/message_readers.py

@@ -4,11 +4,11 @@ import os
 import random
 import string
 import sys
-from collections.abc import Generator, Iterator, Sequence
+from collections.abc import Callable, Generator, Iterator, Sequence
 from contextlib import contextmanager
 from datetime import datetime
 from threading import Event, Thread
-from typing import IO, TYPE_CHECKING, Any, Callable, Optional, cast
+from typing import IO, TYPE_CHECKING, Any, Optional, cast
 
 import boto3
 import dagster._check as check

dagster_aws/version.py

@@ -1 +1 @@
-__version__ = "0.28.0"
+__version__ = "0.28.2"

{dagster_aws-0.28.0.dist-info → dagster_aws-0.28.2.dist-info}/METADATA

@@ -1,22 +1,21 @@
 Metadata-Version: 2.4
 Name: dagster-aws
-Version: 0.28.0
+Version: 0.28.2
 Summary: Package for AWS-specific Dagster framework solid and resource components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-aws
 Author: Dagster Labs
 Author-email: hello@dagsterlabs.com
 License: Apache-2.0
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.9,<3.14
+Requires-Python: >=3.10,<3.14
 License-File: LICENSE
 Requires-Dist: boto3
-Requires-Dist: dagster==1.12.0
+Requires-Dist: dagster==1.12.2
 Requires-Dist: packaging
 Requires-Dist: requests
 Provides-Extra: redshift

{dagster_aws-0.28.0.dist-info → dagster_aws-0.28.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
 dagster_aws/__init__.py,sha256=n5d1NvkQvkVnjWetqKFOLCZUkymTJ0v2WMxLGsbA2Hk,166
 dagster_aws/_stubs.py,sha256=3C6ve6WLtuT2HtCn50hBAMIz-eEAwcZo4fCVnNxGBQw,1112
 dagster_aws/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
-dagster_aws/version.py,sha256=MRQGtOXBhcDKeeNOL0LiB-cllo6kfd8_KGJOvaDp0XQ,23
+dagster_aws/version.py,sha256=K-TM2fq9AmH_Dk8Cadam72wILDZ_6qftLHvY9P1Fc3I,23
 dagster_aws/athena/__init__.py,sha256=J0-QydhSP3pUQk7CawztpaTm3Q1NFEpbubTM1eULbGo,413
 dagster_aws/athena/resources.py,sha256=h7Md_KQXHvtBHA019WZFzAFTClzEUDItDZNdhR2RrPY,11907
 dagster_aws/cloudwatch/__init__.py,sha256=iiyLEyTfNMThU4_SNM_Ns7IGXKR7wTKIvOUz9v_RQBA,82
@@ -11,7 +11,7 @@ dagster_aws/ecr/resources.py,sha256=3SiatAFl9vvkXFmxZjB6e2L20bM6gsKfaveiyTWSRdM,
 dagster_aws/ecs/__init__.py,sha256=E6SAG_YKJcdUIw57amZf3AqzqVUN-GLjrLCpPxSemMg,233
 dagster_aws/ecs/container_context.py,sha256=vB8C4d04iIhd0Gu9vRdYAfua9ib95N25uohEoev9yIY,18154
 dagster_aws/ecs/executor.py,sha256=Fv9pbtJyE0D1DZ2e2uguYS8JHinZFjvcIfwqguYhYDM,16969
-dagster_aws/ecs/launcher.py,sha256=dqNXv2Qm1NqTESywoV5vdC6X2y99mV7_SX4HOLymv6E,37508
+dagster_aws/ecs/launcher.py,sha256=FN5u1tFTwyRsHn1wwoqW7zrrej1JqoZKYBrUsWoG0jo,39234
 dagster_aws/ecs/tasks.py,sha256=n5zigCqVhTsFQJS9g09eVHVCerZYJjhN_Hd-EEZJS7c,18443
 dagster_aws/ecs/test_utils.py,sha256=VqzOIws-fLoDxGrOMzE9ZuAbOF4t9AURIwjyM2Jid-E,2041
 dagster_aws/ecs/utils.py,sha256=CGgMjdoprtS9golz-H3yIDkSXPHxvOTmy5LJc43VVSo,5389
@@ -26,7 +26,7 @@ dagster_aws/emr/types.py,sha256=-JQH3H2niodJuwiyrGW0HqgkV75bD-SRNeXdXuqYY9I,3698
 dagster_aws/emr/utils.py,sha256=AgxJ8L-uprSyjeJcxFgbMMGLsgqtjcuW11X7mIWi1tI,580
 dagster_aws/pipes/__init__.py,sha256=FEqfChhBYzJRVxqAV1jfaXecRUgxJaLq4xgy8AKjNkI,895
 dagster_aws/pipes/context_injectors.py,sha256=DV356jY7k53gB9G--iQjOxx5C7P-ZAxTySbeaMs9gv0,2282
-dagster_aws/pipes/message_readers.py,sha256=kMYdzDPdYxUqzo_Oms5iknA6ezGN-27oHYU6E3-yK1c,14690
+dagster_aws/pipes/message_readers.py,sha256=1qp47TcXK0cERctRK9KFzIdZ3v-X_RSl9nUaR0Cmobo,14690
 dagster_aws/pipes/clients/__init__.py,sha256=p9jEkLonG7gaI0vS6LA_dlZBSk_YSRfH-svT8eU7Myg,564
 dagster_aws/pipes/clients/ecs.py,sha256=eiEXUbB51cFbnQ5llL6fLjCJ8pFPA0UfuWIghBwnf1g,15482
 dagster_aws/pipes/clients/emr.py,sha256=KvJW70N9WLDEDIXjeLBlMzR29a79ZSmqBj55O8sAJfs,14101
@@ -59,8 +59,8 @@ dagster_aws/utils/mrjob/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG
 dagster_aws/utils/mrjob/log4j.py,sha256=Gj2D5USX8swfBXlWf3-NiqeFvrs3h1SGJkOPbP3mke8,4392
 dagster_aws/utils/mrjob/retry.py,sha256=o8Y1Sgu8ABNnsm9jMCfagJpXmgpjcwDw6ufQmDV_Dus,4805
 dagster_aws/utils/mrjob/utils.py,sha256=pFKobTkgwPxQoQNlzGzD3jFxwIP-flQOspQxNKGg2dc,3684
-dagster_aws-0.28.0.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
-dagster_aws-0.28.0.dist-info/METADATA,sha256=923U7KzM7U-7Fa2pb3rCbwI__9JUUc5fHd2liuDc2zA,1695
-dagster_aws-0.28.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dagster_aws-0.28.0.dist-info/top_level.txt,sha256=-hOwIi00tt_K9eXRhiyjst45X-Ffu5_QCdbKcm8UpDg,12
-dagster_aws-0.28.0.dist-info/RECORD,,
+dagster_aws-0.28.2.dist-info/licenses/LICENSE,sha256=4lsMW-RCvfVD4_F57wrmpe3vX1xwUk_OAKKmV_XT7Z0,11348
+dagster_aws-0.28.2.dist-info/METADATA,sha256=H6hZQHJ7uzLx9QgVCj06pBzFeM8GmPkzE316ELaA50M,1646
+dagster_aws-0.28.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dagster_aws-0.28.2.dist-info/top_level.txt,sha256=-hOwIi00tt_K9eXRhiyjst45X-Ffu5_QCdbKcm8UpDg,12
+dagster_aws-0.28.2.dist-info/RECORD,,