Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
changed field formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
jeanluciano committed Apr 10, 2024
1 parent daa9a24 commit bbdae80
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 6 deletions.
11 changes: 8 additions & 3 deletions prefect_aws/workers/ecs_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@
taskRoleArn: "{{ task_role_arn }}"
tags: "{{ labels }}"
taskDefinition: "{{ task_definition_arn }}"
capacityProviderStrategy: "{{ capacity_provider_strategy }}"
"""

# Create task run retry settings
Expand Down Expand Up @@ -372,7 +373,7 @@ class CapacityProvider(BaseModel):
The capacity provider strategy to use when running the task.
"""

capacity_provider: str
capacityProvider: str
weight: int
base: int

Expand Down Expand Up @@ -436,7 +437,7 @@ class ECSVariables(BaseVariables):
)
)
capacity_provider_strategy: Optional[List[CapacityProvider]] = Field(
default=None,
default_factory=List[CapacityProvider],
description=(
"The capacity provider strategy to use when running the task. This is only"
"If a capacityProviderStrategy is specified, we will omit the launchType"
Expand Down Expand Up @@ -1466,8 +1467,12 @@ def _prepare_task_run_request(

task_run_request.setdefault("taskDefinition", task_definition_arn)
assert task_run_request["taskDefinition"] == task_definition_arn
capacityProviderStrategy = (
task_run_request.get("capacityProviderStrategy")
or configuration.capacity_provider_strategy
)

if "capacityProviderStrategy" in task_run_request:
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(
"Removing launchType from task run request. Due to finding"
Expand Down
3 changes: 0 additions & 3 deletions tests/workers/test_ecs_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -2026,9 +2026,6 @@ async def test_user_defined_capacity_provider_strategy(
{"base": 0, "weight": 1, "capacityProvider": "r6i.large"}
],
)

assert "launchType" not in configuration.task_run_request

session = aws_credentials.get_boto3_session()
ecs_client = session.client("ecs")

Expand Down

0 comments on commit bbdae80

Please sign in to comment.