Add cloudwatch_logs_options to customize `configure_cloudwatch_logs…

…` behavior (#116)

* Add `cloudwatch_logs_options` to customize `configure_cloudwatch_logs` behavior

* Add changelog entry

* Feedback from review

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
 

prefect_aws/ecs.py

@@ -212,6 +212,15 @@ class ECSTask(Infrastructure):
             "unless `stream_output` is set. "
         ),
     )
+    cloudwatch_logs_options: Dict[str, str] = Field(
+        default_factory=dict,
+        description=(
+            "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(
         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,
                 },
             }
 

tests/test_ecs.py

@@ -943,14 +943,31 @@ 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"],
             configure_cloudwatch_logs=True,
         )
 
 
+@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):