From 450896b0d7a5bb55d1dd0dc56d772f69e7bc4318 Mon Sep 17 00:00:00 2001 From: Michael Adkins Date: Wed, 21 Sep 2022 13:21:32 -0500 Subject: [PATCH 1/3] Add `cloudwatch_logs_options` to customize `configure_cloudwatch_logs` behavior --- prefect_aws/ecs.py | 27 +++++++++++++++++++++ tests/test_ecs.py | 60 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/prefect_aws/ecs.py b/prefect_aws/ecs.py index cac8cc66..b2725070 100644 --- a/prefect_aws/ecs.py +++ b/prefect_aws/ecs.py @@ -212,6 +212,15 @@ class ECSTask(Infrastructure): "unless `stream_output` is set. " ), ) + cloudwatch_logs_options: Optional[Dict[str, str]] = Field( + default=None, + description=( + "When `configure_cloudwatch_logs`, this setting may be used to pass " + "additional options to the CloudWatch logs configuration or override " + "the default options. See [the AWS documentation](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/using_awslogs.html#create_awslogs_logdriver_options) " # noqa + "for available options." + ), + ) stream_output: bool = Field( default=None, description=( @@ -321,6 +330,23 @@ def configure_cloudwatch_logs_requires_execution_role_arn( ) return values + @root_validator + def cloudwatch_logs_options_requires_configure_cloudwatch_logs( + cls, values: dict + ) -> dict: + """ + Enforces that an execution role arn is provided (or could be provided by a + runtime task definition) when configuring logging. + """ + if values.get("cloudwatch_logs_options") and not values.get( + "configure_cloudwatch_logs" + ): + raise ValueError( + "`configure_cloudwatch_log` must be enabled to use " + "`cloudwatch_logs_options`." + ) + return values + @root_validator def image_is_required(cls, values: dict) -> dict: """ @@ -848,6 +874,7 @@ def _prepare_task_definition(self, task_definition: dict, region: str) -> dict: "awslogs-group": "prefect", "awslogs-region": region, "awslogs-stream-prefix": self.name or "prefect", + **(self.cloudwatch_logs_options or {}), }, } diff --git a/tests/test_ecs.py b/tests/test_ecs.py index 132be7bd..69775b1b 100644 --- a/tests/test_ecs.py +++ b/tests/test_ecs.py @@ -943,7 +943,10 @@ async def test_network_config_from_vpc_with_no_subnets(aws_credentials): @pytest.mark.usefixtures("ecs_mocks") async def test_logging_requires_execution_role_arn(aws_credentials): - with pytest.raises(ValidationError, match="`execution_role_arn` must be provided"): + with pytest.raises( + ValidationError, + match="`execution_role_arn` must be provided", + ): ECSTask( aws_credentials=aws_credentials, command=["prefect", "version"], @@ -951,6 +954,20 @@ async def test_logging_requires_execution_role_arn(aws_credentials): ) +@pytest.mark.usefixtures("ecs_mocks") +async def test_log_options_requires_logging(aws_credentials): + with pytest.raises( + ValidationError, + match="`configure_cloudwatch_log` must be enabled to use `cloudwatch_logs_options`", # noqa + ): + ECSTask( + aws_credentials=aws_credentials, + command=["prefect", "version"], + configure_cloudwatch_logs=False, + cloudwatch_logs_options={"foo": " bar"}, + ) + + @pytest.mark.usefixtures("ecs_mocks") async def test_logging_requires_execution_role_arn_at_runtime(aws_credentials): # In constrast to `test_logging_requires_execution_role_arn`, a task definition @@ -1012,6 +1029,47 @@ async def test_configure_cloudwatch_logging(aws_credentials): assert "logConfiguration" not in container +@pytest.mark.usefixtures("ecs_mocks") +async def test_cloudwatch_log_options(aws_credentials): + session = aws_credentials.get_boto3_session() + ecs_client = session.client("ecs") + + with mock_logs(): + task = ECSTask( + aws_credentials=aws_credentials, + auto_deregister_task_definition=False, + command=["prefect", "version"], + configure_cloudwatch_logs=True, + execution_role_arn="test", + cloudwatch_logs_options={ + "awslogs-stream-prefix": "override-prefix", + "max-buffer-size": "2m", + }, + ) + + task_arn = await run_then_stop_task(task) + task = describe_task(ecs_client, task_arn) + task_definition = describe_task_definition(ecs_client, task) + + for container in task_definition["containerDefinitions"]: + if container["name"] == "prefect": + # Assert that the 'prefect' container has logging configured with user + # provided options + assert container["logConfiguration"] == { + "logDriver": "awslogs", + "options": { + "awslogs-create-group": "true", + "awslogs-group": "prefect", + "awslogs-region": "us-east-1", + "awslogs-stream-prefix": "override-prefix", + "max-buffer-size": "2m", + }, + } + else: + # Other containers should not be modifed + assert "logConfiguration" not in container + + @pytest.mark.usefixtures("ecs_mocks") @pytest.mark.parametrize("launch_type", ["FARGATE", "FARGATE_SPOT"]) async def test_bridge_network_mode_warns_on_fargate(aws_credentials, launch_type: str): From a334699529d20ac352863f9d25d910d2356683e3 Mon Sep 17 00:00:00 2001 From: Michael Adkins Date: Wed, 21 Sep 2022 13:23:26 -0500 Subject: [PATCH 2/3] Add changelog entry --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e759274f..a4411905 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased - Fix configuration to submit doc edits via GitHub - [#110](https://github.com/PrefectHQ/prefect-aws/pull/110) +- Add `ECSTask.cloudwatch_logs_options` for customization of CloudWatch logging — [#116](https://github.com/PrefectHQ/prefect-aws/pull/116) ### Added From 945bf852136020c57de8a573f1a156fd2412f75e Mon Sep 17 00:00:00 2001 From: Michael Adkins Date: Wed, 21 Sep 2022 14:03:32 -0500 Subject: [PATCH 3/3] Feedback from review --- prefect_aws/ecs.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/prefect_aws/ecs.py b/prefect_aws/ecs.py index b2725070..51bc1e3f 100644 --- a/prefect_aws/ecs.py +++ b/prefect_aws/ecs.py @@ -212,13 +212,13 @@ class ECSTask(Infrastructure): "unless `stream_output` is set. " ), ) - cloudwatch_logs_options: Optional[Dict[str, str]] = Field( - default=None, + cloudwatch_logs_options: Dict[str, str] = Field( + default_factory=dict, description=( - "When `configure_cloudwatch_logs`, this setting may be used to pass " - "additional options to the CloudWatch logs configuration or override " - "the default options. See [the AWS documentation](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/using_awslogs.html#create_awslogs_logdriver_options) " # noqa - "for available options." + "When `configure_cloudwatch_logs` is enabled, this setting may be used to " + "pass additional options to the CloudWatch logs configuration or override " + "the default options. See the AWS documentation for available options. " + "https://docs.aws.amazon.com/AmazonECS/latest/developerguide/using_awslogs.html#create_awslogs_logdriver_options" # noqa ), ) stream_output: bool = Field( @@ -874,7 +874,7 @@ def _prepare_task_definition(self, task_definition: dict, region: str) -> dict: "awslogs-group": "prefect", "awslogs-region": region, "awslogs-stream-prefix": self.name or "prefect", - **(self.cloudwatch_logs_options or {}), + **self.cloudwatch_logs_options, }, }