Repository: PrefectHQ/prefect-aws (python)

Issue:
Worker does not template ECSTask family
PR #152 changed the behavior of the `family` top-level settings on the `ECSTask` infrastructure block. Rather than just being set to `prefect`, it instead templates using the flow and deployment names: `prefect__{flow-name}__{deployment-name}`. This change does not seem to have made it into the [worker](https://github.com/PrefectHQ/prefect-aws/blob/main/prefect_aws/workers/ecs_worker.py#L1185-L1190). When users have multiple heterogeneous deployments, using a single family (`prefect`) can lead to excessive task definition revisions and may trigger service rate limits.

Pull request:
Generates family field
<!-- Thanks for contributing 🎉! Please ensure the title neatly summarizes the proposed changes. -->

<!-- Overview -->

Closes #289 

### Example
<!-- A code blurb is best. Changes to features should include an example that is executable by a new user. -->

### Screenshots
<!--
Any relevant screenshots
  - The updated docs page from `mkdocs serve`.
  - Output from running the example.
  - Service integration test results.
-->

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [ x ] References any related issue by including "Closes #<Issue Number>" or "Closes <Issue URL>".
  - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect-aws/issues/new/choose) first.
- [ x ] Includes tests or only affects documentation.
- [ x ] Passes `pre-commit` checks.
  - Run `pre-commit install && pre-commit run --all` locally for formatting and linting.
- [ ] Includes screenshots of documentation updates.
  - Run `mkdocs serve` view documentation locally.


Patch:
diff --git a/prefect_aws/workers/ecs_worker.py b/prefect_aws/workers/ecs_worker.py
index edbe843..1a5c3d2 100644
--- a/prefect_aws/workers/ecs_worker.py
+++ b/prefect_aws/workers/ecs_worker.py
@@ -673,7 +673,7 @@ class ECSWorker(BaseWorker):
 
         if not task_definition_arn:
             task_definition = self._prepare_task_definition(
-                configuration, region=ecs_client.meta.region_name
+                configuration, region=ecs_client.meta.region_name, flow_run=flow_run
             )
             (
                 task_definition_arn,
@@ -1205,10 +1205,28 @@ class ECSWorker(BaseWorker):
                 )
             time.sleep(configuration.task_watch_poll_interval)
 
+    def _get_or_generate_family(self, task_definition: dict, flow_run: FlowRun) -> str:
+        """
+        Gets or generate a family for the task definition.
+        """
+        family = task_definition.get("family")
+        if not family:
+            assert self._work_pool_name and flow_run.deployment_id
+            family = (
+                f"{ECS_DEFAULT_FAMILY}_{self._work_pool_name}_{flow_run.deployment_id}"
+            )
+        slugify(
+            family,
+            max_length=255,
+            regex_pattern=r"[^a-zA-Z0-9-_]+",
+        )
+        return family
+
     def _prepare_task_definition(
         self,
         configuration: ECSJobConfiguration,
         region: str,
+        flow_run: FlowRun,
     ) -> dict:
         """
         Prepare a task definition by inferring any defaults and merging overrides.
@@ -1269,13 +1287,9 @@ class ECSWorker(BaseWorker):
                 },
             }
 
-        family = task_definition.get("family") or ECS_DEFAULT_FAMILY
-        task_definition["family"] = slugify(
-            family,
-            max_length=255,
-            regex_pattern=r"[^a-zA-Z0-9-_]+",
+        task_definition["family"] = self._get_or_generate_family(
+            task_definition, flow_run
         )
-
         # CPU and memory are required in some cases, retrieve the value to use
         cpu = task_definition.get("cpu") or ECS_DEFAULT_CPU
         memory = task_definition.get("memory") or ECS_DEFAULT_MEMORY


Test patch:
diff --git a/tests/workers/test_ecs_worker.py b/tests/workers/test_ecs_worker.py
index 702f2dd..e4dab38 100644
--- a/tests/workers/test_ecs_worker.py
+++ b/tests/workers/test_ecs_worker.py
@@ -26,6 +26,7 @@ from prefect_aws.workers.ecs_worker import (
     _TASK_DEFINITION_CACHE,
     ECS_DEFAULT_CONTAINER_NAME,
     ECS_DEFAULT_CPU,
+    ECS_DEFAULT_FAMILY,
     ECS_DEFAULT_MEMORY,
     AwsCredentials,
     ECSJobConfiguration,
@@ -648,6 +649,7 @@ async def test_task_definition_arn(aws_credentials: AwsCredentials, flow_run: Fl
     _, task_arn = parse_identifier(result.identifier)
 
     task = describe_task(ecs_client, task_arn)
+    print(task)
     assert (
         task["taskDefinitionArn"] == task_definition_arn
     ), "The task definition should be used without registering a new one"
@@ -2316,3 +2318,27 @@ async def test_mask_sensitive_env_values():
         res["overrides"]["containerOverrides"][0]["environment"][1]["value"]
         == "NORMAL_VALUE"
     )
+
+
+@pytest.mark.usefixtures("ecs_mocks")
+async def test_get_or_generate_family(
+    aws_credentials: AwsCredentials, flow_run: FlowRun
+):
+    configuration = await construct_configuration(
+        aws_credentials=aws_credentials,
+    )
+
+    work_pool_name = "test"
+    session = aws_credentials.get_boto3_session()
+    ecs_client = session.client("ecs")
+    family = f"{ECS_DEFAULT_FAMILY}_{work_pool_name}_{flow_run.deployment_id}"
+
+    async with ECSWorker(work_pool_name=work_pool_name) as worker:
+        result = await run_then_stop_task(worker, configuration, flow_run)
+
+    assert result.status_code == 0
+    _, task_arn = parse_identifier(result.identifier)
+
+    task = describe_task(ecs_client, task_arn)
+    task_definition = describe_task_definition(ecs_client, task)
+    assert task_definition["family"] == family