Repository: PrefectHQ/prefect-aws (python)

Issue:
Accept `None` as an argument to Launch Type for ECSTask with `publish_as_work_pool`
<!-- Thanks for submitting an issue! Please describe the issue. -->
ECS doesn't let you use both a Capacity Provider and a Launch Type.  We use Capacity Providers in order to specify specific Instance Types (`r6i.large` etc) for our jobs.  This works fine if you pass `None` to the `launch_type` argument, but when I tried to use one as a foundation for `publish_as_work_pool`, it gave me this error `"Validation failed for field 'launch_type'. Failure reason: None is not one of ['FARGATE', 'EC2', 'EXTERNAL', 'FARGATE_SPOT']"`.  This also works with `ECSTask` normally at the moment.

Relevant Issue from Prefect 1 that lead to support for Capacity Providers in ECS: https://github.com/PrefectHQ/prefect/issues/5210

# Expectation / Proposal
Change the validation so that `launch_type=None` is okay when using `publish_as_work_pool`.  


# Traceback / Example
```
ecs_task_block = ECSTask(**full_ecs_args)

ecs_task_block.publish_as_work_pool("ecs_task_block_test")

"Validation failed for field 'launch_type'. Failure reason: None is not one of ['FARGATE', 'EC2', 'EXTERNAL', 'FARGATE_SPOT']"
```


Pull request:
Support for capacity provider
<!-- Thanks for contributing 🎉! Please ensure the title neatly summarizes the proposed changes. -->

<!-- Overview -->
Enables passing down  `capacityProviderStrategy` to `task_run_request`.

Closes #362 



### 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 d216efa..d862fdd 100644
--- a/prefect_aws/workers/ecs_worker.py
+++ b/prefect_aws/workers/ecs_worker.py
@@ -70,9 +70,9 @@ from prefect.workers.base import (
 from pydantic import VERSION as PYDANTIC_VERSION
 
 if PYDANTIC_VERSION.startswith("2."):
-    from pydantic.v1 import Field, root_validator
+    from pydantic.v1 import BaseModel, Field, root_validator
 else:
-    from pydantic import Field, root_validator
+    from pydantic import Field, root_validator, BaseModel
 
 from slugify import slugify
 from tenacity import retry, stop_after_attempt, wait_fixed, wait_random
@@ -126,6 +126,7 @@ overrides:
   taskRoleArn: "{{ task_role_arn }}"
 tags: "{{ labels }}"
 taskDefinition: "{{ task_definition_arn }}"
+capacityProviderStrategy: "{{ capacity_provider_strategy }}"
 """
 
 # Create task run retry settings
@@ -245,6 +246,16 @@ def mask_api_key(task_run_request):
     )
 
 
+class CapacityProvider(BaseModel):
+    """
+    The capacity provider strategy to use when running the task.
+    """
+
+    capacityProvider: str
+    weight: int
+    base: int
+
+
 class ECSJobConfiguration(BaseJobConfiguration):
     """
     Job configuration for an ECS worker.
@@ -425,6 +436,14 @@ class ECSVariables(BaseVariables):
             ),
         )
     )
+    capacity_provider_strategy: Optional[List[CapacityProvider]] = Field(
+        default_factory=list,
+        description=(
+            "The capacity provider strategy to use when running the task. "
+            "If a capacity provider strategy is specified, the selected launch"
+            " type will be ignored."
+        ),
+    )
     image: Optional[str] = Field(
         default=None,
         description=(
@@ -1449,17 +1468,24 @@ class ECSWorker(BaseWorker):
 
         task_run_request.setdefault("taskDefinition", task_definition_arn)
         assert task_run_request["taskDefinition"] == task_definition_arn
+        capacityProviderStrategy = task_run_request.get("capacityProviderStrategy")
 
-        if task_run_request.get("launchType") == "FARGATE_SPOT":
+        if capacityProviderStrategy:
+            # Should not be provided at all if capacityProviderStrategy is set, see https://docs.aws.amazon.com/AmazonECS/latest/APIReference/API_RunTask.html#ECS-RunTask-request-capacityProviderStrategy  # noqa
+            self._logger.warning(
+                "Found capacityProviderStrategy. "
+                "Removing launchType from task run request."
+            )
+            task_run_request.pop("launchType", None)
+
+        elif task_run_request.get("launchType") == "FARGATE_SPOT":
             # Should not be provided at all for FARGATE SPOT
             task_run_request.pop("launchType", None)
 
             # A capacity provider strategy is required for FARGATE SPOT
-            task_run_request.setdefault(
-                "capacityProviderStrategy",
-                [{"capacityProvider": "FARGATE_SPOT", "weight": 1}],
-            )
-
+            task_run_request["capacityProviderStrategy"] = [
+                {"capacityProvider": "FARGATE_SPOT", "weight": 1}
+            ]
         overrides = task_run_request.get("overrides", {})
         container_overrides = overrides.get("containerOverrides", [])
 


Test patch:
diff --git a/tests/workers/test_ecs_worker.py b/tests/workers/test_ecs_worker.py
index dc32931..c5c7817 100644
--- a/tests/workers/test_ecs_worker.py
+++ b/tests/workers/test_ecs_worker.py
@@ -506,6 +506,7 @@ async def test_launch_types(
         # Instead, it requires a capacity provider strategy but this is not supported
         # by moto and is not present on the task even when provided so we assert on the
         # mock call to ensure it is sent
+
         assert mock_run_task.call_args[0][1].get("capacityProviderStrategy") == [
             {"capacityProvider": "FARGATE_SPOT", "weight": 1}
         ]
@@ -2050,6 +2051,41 @@ async def test_user_defined_environment_variables_in_task_definition_template(
     ]
 
 
+@pytest.mark.usefixtures("ecs_mocks")
+async def test_user_defined_capacity_provider_strategy(
+    aws_credentials: AwsCredentials, flow_run: FlowRun
+):
+    configuration = await construct_configuration(
+        aws_credentials=aws_credentials,
+        capacity_provider_strategy=[
+            {"base": 0, "weight": 1, "capacityProvider": "r6i.large"}
+        ],
+    )
+    session = aws_credentials.get_boto3_session()
+    ecs_client = session.client("ecs")
+
+    async with ECSWorker(work_pool_name="test") as worker:
+        # Capture the task run call because moto does not track
+        # 'capacityProviderStrategy'
+        original_run_task = worker._create_task_run
+        mock_run_task = MagicMock(side_effect=original_run_task)
+        worker._create_task_run = mock_run_task
+
+        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)
+    assert not task.get("launchType")
+    # Instead, it requires a capacity provider strategy but this is not supported
+    # by moto and is not present on the task even when provided so we assert on the
+    # mock call to ensure it is sent
+    assert mock_run_task.call_args[0][1].get("capacityProviderStrategy") == [
+        {"base": 0, "weight": 1, "capacityProvider": "r6i.large"},
+    ]
+
+
 @pytest.mark.usefixtures("ecs_mocks")
 async def test_user_defined_environment_variables_in_task_run_request_template(
     aws_credentials: AwsCredentials, flow_run: FlowRun