Repository: PrefectHQ/prefect

Issue:
Docker Agent: Passing custom parameters when creating docker.APIClient
## Current behavior
prefect docker agent is using default parameters to create an instance of docker.APIClient. This makes it impossible to customize the docker client. Things such as to increase the timeout of docker client cannot be done.

https://github.com/PrefectHQ/prefect/blob/master/src/prefect/agent/docker/agent.py#L186


## Proposed behavior
We should have a way to pass custom parameters to when creating an instance of docker.APIClient. 
Either as a dictionary or as a json config file




## Example
When the somehow get overloaded, docker daemon failed to create a container within 60 seconds, docker.APIClient will then return a timeout error similar to https://github.com/docker/compose/issues/3927. In my case increase the timeout from 60 to 120 seconds would mitigate this problem.

Discussion:

Patch:
diff --git a/changes/pr4211.yaml b/changes/pr4211.yaml
deleted file mode 100644
index cdeb4fc989..0000000000
--- a/changes/pr4211.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-fix:
-  - "Move `command`, `environment`, `cpu`, `memory`, `execution_role_arn`, and `task_role_arn` configuration for ECS tasks from definition time to run time in the ECS agent - [#4211](https://github.com/PrefectHQ/prefect/pull/4211)"
-  - "Register (and deregister) a new task definition for every flow run in ECS agent - [#4211](https://github.com/PrefectHQ/prefect/pull/4211)"
diff --git a/changes/pr4232.yaml b/changes/pr4232.yaml
new file mode 100644
index 0000000000..a9b08be9a5
--- /dev/null
+++ b/changes/pr4232.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Add `--docker-client-timeout` flag to docker agent, for configuring the timeout for all docker API requests - [#4232](https://github.com/PrefectHQ/prefect/pull/4232)"
diff --git a/docs/orchestration/agents/ecs.md b/docs/orchestration/agents/ecs.md
index 81c77ad912..6cfc7f57d1 100644
--- a/docs/orchestration/agents/ecs.md
+++ b/docs/orchestration/agents/ecs.md
@@ -133,17 +133,18 @@ their respective [ECSRun](/orchestration/flow_config/run_configs.md#ecsrun) `run
 
 ### Execution Role ARN
 
-ECS tasks use [execution
-roles](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html)
-to grant permissions to the ECS infrastructure to make AWS API calls on your
-behalf. If actions taken to *start* your task require external AWS services
-(e.g. pulling an image from ECR), you'll need to configure an execution role.
-Permissions used by your code once your task starts are granted via [task
-roles](#task-role-arn) instead (see above).
-
-ECS provides a builtin policy `AmazonECSTaskExecutionPolicy` that provides
-common settings.  This supports pulling images from ECR and enables using
-CloudWatch logs. The full policy is below:
+The following policy is the AmazonECSTaskExecutionPolicy. The execution-role-arn 
+can be used to pull the image from ECR or enable logs in CloudWatch. More information for creating this role
+can be found [here](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html). 
+You can configure a default execution role for tasks started by the agent using the `--execution-role-arn` option:
+
+```bash
+prefect agent ecs start --execution-role-arn my-execution-role-arn
+```
+
+Flows can override this agent default by passing the `execution_role_arn` option to
+their respective [ECSRun](/orchestration/flow_config/run_configs.md#ecsrun) `run_config`.
+
 
 ```json
 {
@@ -165,21 +166,6 @@ CloudWatch logs. The full policy is below:
 }
 ```
 
-Usually AWS will automatically create an IAM role with this policy named
-`ecsTaskExecutionRole` (if not, you may need to create one yourself, see [the
-AWS docs for more
-info](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html)).
-
-You can configure a default execution role for tasks started by the agent using
-the `--execution-role-arn` option:
-
-```bash
-prefect agent ecs start --execution-role-arn my-execution-role-arn
-```
-
-Flows can override this agent default by passing the `execution_role_arn` option to
-their respective [ECSRun](/orchestration/flow_config/run_configs.md#ecsrun) `run_config`.
-
 ### Custom Task Definition Template
 
 For deeper customization of the Tasks created by the agent, you may want to
@@ -188,10 +174,7 @@ configured per-flow (on the
 [ECSRun](/orchestration/flow_config/run_configs.md#ecsrun) `run_config`), or on
 the Agent as a default for flows that don't provide their own template.
 
-The flow will be executed in a container named `flow` - if a container named
-`flow` isn't part of the task definition template Prefect will add a new
-container with that name (this allows adding sidecar containers without
-requiring the user to define a `flow` container as well). Any option available to
+Any option available to
 [`register_task_definition`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.register_task_definition)
 may be specified here. For reference, the default template packaged with
 Prefect can be found
diff --git a/src/prefect/agent/docker/agent.py b/src/prefect/agent/docker/agent.py
index 218fd74449..766bf32337 100644
--- a/src/prefect/agent/docker/agent.py
+++ b/src/prefect/agent/docker/agent.py
@@ -18,17 +18,18 @@ if TYPE_CHECKING:
     import docker
 
 
-def _stream_container_logs(base_url: str, container_id: str) -> None:
+def _stream_container_logs(base_url: str, timeout: int, container_id: str) -> None:
     """
     Stream container logs back to stdout
 
     Args:
         - base_url (str): URL for a Docker daemon server
+        - timeout (int): timeout for docker api requests
         - container_id (str): ID of a container to stream logs
     """
     import docker
 
-    client = docker.APIClient(base_url=base_url, version="auto")
+    client = docker.APIClient(base_url=base_url, timeout=timeout, version="auto")
     for log in client.logs(container=container_id, stream=True, follow=True):
         print(str(log, "utf-8").rstrip())
 
@@ -82,6 +83,8 @@ class DockerAgent(Agent):
             some Docker-in-Docker setups that users may be running their agent with.
         - reg_allow_list (List[str], optional): Limits Docker Agent to only pull images
             from the listed registries.
+        - docker_client_timeout (int, optional): The timeout to use for docker
+            API calls, defaults to 60 seconds.
     """
 
     def __init__(
@@ -101,6 +104,7 @@ class DockerAgent(Agent):
         networks: List[str] = None,
         docker_interface: bool = True,
         reg_allow_list: List[str] = None,
+        docker_client_timeout: int = None,
     ) -> None:
         super().__init__(
             agent_config_id=agent_config_id,
@@ -150,6 +154,7 @@ class DockerAgent(Agent):
         self.networks = networks
         self.logger.debug(f"Docker networks set to {self.networks}")
 
+        self.docker_client_timeout = docker_client_timeout or 60
         self.docker_interface = docker_interface
         self.logger.debug(
             "Docker interface toggle set to {}".format(self.docker_interface)
@@ -183,7 +188,11 @@ class DockerAgent(Agent):
         # the 'import prefect' time low
         import docker
 
-        return docker.APIClient(base_url=self.base_url, version="auto")
+        return docker.APIClient(
+            base_url=self.base_url,
+            version="auto",
+            timeout=self.docker_client_timeout,
+        )
 
     def heartbeat(self) -> None:
         try:
@@ -483,6 +492,7 @@ class DockerAgent(Agent):
             target=_stream_container_logs,
             kwargs={
                 "base_url": self.base_url,
+                "timeout": self.docker_client_timeout,
                 "container_id": container_id,
             },
         )
diff --git a/src/prefect/agent/ecs/agent.py b/src/prefect/agent/ecs/agent.py
index 1542685431..9caf518b3a 100644
--- a/src/prefect/agent/ecs/agent.py
+++ b/src/prefect/agent/ecs/agent.py
@@ -1,6 +1,6 @@
 import os
 from copy import deepcopy
-from typing import Iterable, Dict, Any
+from typing import Iterable, Dict, Optional, Any
 
 import slugify
 import yaml
@@ -192,6 +192,9 @@ class ECSAgent(Agent):
         )  # type: Dict[str, Any]
 
         self.ecs_client = get_boto_client("ecs", **self.boto_kwargs)
+        self.rgtag_client = get_boto_client(
+            "resourcegroupstaggingapi", **self.boto_kwargs
+        )
 
         # Load default task definition
         if not task_definition_path:
@@ -224,6 +227,17 @@ class ECSAgent(Agent):
         else:
             self.run_task_kwargs = {}
 
+        # If `task_role_arn` is configured on the agent, add it to the default
+        # template. The agent default `task_role_arn` is only applied if using
+        # the agent's default template.
+        if self.task_role_arn:
+            self.task_definition["taskRoleArn"] = self.task_role_arn
+
+        # If `execution_role_arn` is configured on the agent, add it to the
+        # default task definition template.
+        if self.execution_role_arn:
+            self.task_definition["executionRoleArn"] = self.execution_role_arn
+
         # If running on fargate, auto-configure `networkConfiguration` for the
         # user if they didn't configure it themselves.
         if self.launch_type == "FARGATE" and not self.run_task_kwargs.get(
@@ -288,7 +302,8 @@ class ECSAgent(Agent):
         run_config = self._get_run_config(flow_run, ECSRun)
         assert isinstance(run_config, ECSRun)  # mypy
 
-        if run_config.task_definition_arn is None:
+        taskdef_arn = self.get_task_definition_arn(flow_run, run_config)
+        if taskdef_arn is None:
             # Register a new task definition
             self.logger.debug(
                 "Registering new task definition for flow %s", flow_run.flow.id
@@ -296,22 +311,12 @@ class ECSAgent(Agent):
             taskdef = self.generate_task_definition(flow_run, run_config)
             resp = self.ecs_client.register_task_definition(**taskdef)
             taskdef_arn = resp["taskDefinition"]["taskDefinitionArn"]
-            new_taskdef_arn = True
             self.logger.debug(
                 "Registered task definition %s for flow %s",
                 taskdef_arn,
                 flow_run.flow.id,
             )
         else:
-            from prefect.serialization.storage import StorageSchema
-            from prefect.storage import Docker
-
-            if isinstance(StorageSchema().load(flow_run.flow.storage), Docker):
-                raise ValueError(
-                    "Cannot provide `task_definition_arn` when using `Docker` storage"
-                )
-            taskdef_arn = run_config.task_definition_arn
-            new_taskdef_arn = False
             self.logger.debug(
                 "Using task definition %s for flow %s", taskdef_arn, flow_run.flow.id
             )
@@ -320,12 +325,6 @@ class ECSAgent(Agent):
         kwargs = self.get_run_task_kwargs(flow_run, run_config)
 
         resp = self.ecs_client.run_task(taskDefinition=taskdef_arn, **kwargs)
-
-        # Always deregister the task definition if a new one was registered
-        if new_taskdef_arn:
-            self.logger.debug("Deregistering task definition %s", taskdef_arn)
-            self.ecs_client.deregister_task_definition(taskDefinition=taskdef_arn)
-
         if resp.get("tasks"):
             task_arn = resp["tasks"][0]["taskArn"]
             self.logger.debug("Started task %r for flow run %r", task_arn, flow_run.id)
@@ -337,6 +336,51 @@ class ECSAgent(Agent):
             )
         )
 
+    def get_task_definition_tags(self, flow_run: GraphQLResult) -> dict:
+        """Get required task definition tags from a flow run.
+
+        Args:
+            - flow_run (GraphQLResult): the flow run
+
+        Returns:
+            - dict: a dict of tags to use
+        """
+        return {
+            "prefect:flow-id": flow_run.flow.id,
+            "prefect:flow-version": str(flow_run.flow.version),
+        }
+
+    def get_task_definition_arn(
+        self, flow_run: GraphQLResult, run_config: ECSRun
+    ) -> Optional[str]:
+        """Get an existing task definition ARN for a flow run.
+
+        Args:
+            - flow_run (GraphQLResult): the flow run
+            - run_config (ECSRun): The flow's run config
+
+        Returns:
+            - Optional[str]: the task definition ARN. Returns `None` if no
+                existing definition is found.
+        """
+        if run_config.task_definition_arn is not None:
+            return run_config.task_definition_arn
+
+        tags = self.get_task_definition_tags(flow_run)
+
+        from botocore.exceptions import ClientError
+
+        try:
+            res = self.rgtag_client.get_resources(
+                TagFilters=[{"Key": k, "Values": [v]} for k, v in tags.items()],
+                ResourceTypeFilters=["ecs:task-definition"],
+            )
+            if res["ResourceTagMappingList"]:
+                return res["ResourceTagMappingList"][0]["ResourceARN"]
+            return None
+        except ClientError:
+            return None
+
     def generate_task_definition(
         self, flow_run: GraphQLResult, run_config: ECSRun
     ) -> Dict[str, Any]:
@@ -367,15 +411,17 @@ class ECSAgent(Agent):
             word_boundary=True,
             save_order=True,
         )
-        taskdef["family"] = f"prefect-{slug}"
-
-        # Add some metadata tags for easier tracking by users
-        taskdef.setdefault("tags", []).extend(
-            [
-                {"key": "prefect:flow-id", "value": flow_run.flow.id},
-                {"key": "prefect:flow-version", "value": str(flow_run.flow.version)},
-            ]
-        )
+        family = f"prefect-{slug}"
+
+        tags = self.get_task_definition_tags(flow_run)
+
+        taskdef["family"] = family
+
+        taskdef_tags = [{"key": k, "value": v} for k, v in tags.items()]
+        for entry in taskdef.get("tags", []):
+            if entry["key"] not in tags:
+                taskdef_tags.append(entry)
+        taskdef["tags"] = taskdef_tags
 
         # Get the flow container (creating one if it doesn't already exist)
         containers = taskdef.setdefault("containerDefinitions", [])
@@ -389,18 +435,42 @@ class ECSAgent(Agent):
         # Set flow image
         container["image"] = image = get_flow_image(flow_run)
 
-        # Add `PREFECT__CONTEXT__IMAGE` environment variable
-        env = {"PREFECT__CONTEXT__IMAGE": image}
+        # Set flow run command
+        container["command"] = ["/bin/sh", "-c", get_flow_run_command(flow_run)]
+
+        # Set taskRoleArn if configured
+        if run_config.task_role_arn:
+            taskdef["taskRoleArn"] = run_config.task_role_arn
+
+        # Set executionRoleArn if configured
+        if run_config.execution_role_arn:
+            taskdef["executionRoleArn"] = run_config.execution_role_arn
+
+        # Populate static environment variables from the following sources,
+        # with precedence:
+        # - Static environment variables, hardcoded below
+        # - Values in the task definition template
+        env = {
+            "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
+            "PREFECT__CONTEXT__IMAGE": image,
+            "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
+            "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
+        }
         container_env = [{"name": k, "value": v} for k, v in env.items()]
         for entry in container.get("environment", []):
             if entry["name"] not in env:
                 container_env.append(entry)
         container["environment"] = container_env
 
-        # Ensure that cpu/memory are strings not integers
-        if "cpu" in taskdef:
+        # Set resource requirements, if provided
+        # Also ensure that cpu/memory are strings not integers
+        if run_config.cpu:
+            taskdef["cpu"] = str(run_config.cpu)
+        elif "cpu" in taskdef:
             taskdef["cpu"] = str(taskdef["cpu"])
-        if "memory" in taskdef:
+        if run_config.memory:
+            taskdef["memory"] = str(run_config.memory)
+        elif "memory" in taskdef:
             taskdef["memory"] = str(taskdef["memory"])
 
         return taskdef
@@ -438,35 +508,9 @@ class ECSAgent(Agent):
             container = {"name": "flow"}
             container_overrides.append(container)
 
-        # Set taskRoleArn if configured
-        if run_config.task_role_arn:
-            overrides["taskRoleArn"] = run_config.task_role_arn
-        elif self.task_role_arn:
-            overrides["taskRoleArn"] = self.task_role_arn
-
-        # Set executionRoleArn if configured
-        if run_config.execution_role_arn:
-            overrides["executionRoleArn"] = run_config.execution_role_arn
-        elif self.execution_role_arn:
-            overrides["executionRoleArn"] = self.execution_role_arn
-
-        # Set resource requirements, if provided
-        # Also ensure that cpu/memory are strings not integers
-        if run_config.cpu:
-            overrides["cpu"] = str(run_config.cpu)
-        elif "cpu" in overrides:
-            overrides["cpu"] = str(overrides["cpu"])
-        if run_config.memory:
-            overrides["memory"] = str(run_config.memory)
-        elif "memory" in overrides:
-            overrides["memory"] = str(overrides["memory"])
-
-        # Set flow run command
-        container["command"] = ["/bin/sh", "-c", get_flow_run_command(flow_run)]
-
         # Populate environment variables from the following sources,
         # with precedence:
-        # - Values required for flow execution, hardcoded below
+        # - Dynamic values required for flow execution, hardcoded below
         # - Values set on the ECSRun object
         # - Values set using the `--env` CLI flag on the agent
         env = self.env_vars.copy()
@@ -474,9 +518,6 @@ class ECSAgent(Agent):
             env.update(run_config.env)
         env.update(
             {
-                "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
-                "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
-                "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
                 "PREFECT__BACKEND": config.backend,
                 "PREFECT__CLOUD__API": config.cloud.api,
                 "PREFECT__CONTEXT__FLOW_RUN_ID": flow_run.id,
diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py
index b90fa8b6aa..f26b4f91c5 100644
--- a/src/prefect/cli/agent.py
+++ b/src/prefect/cli/agent.py
@@ -223,6 +223,12 @@ def docker():
         "setups that users may be running their agent with."
     ),
 )
+@click.option(
+    "--docker-client-timeout",
+    default=None,
+    type=int,
+    help="The timeout to use for docker API calls, defaults to 60 seconds.",
+)
 def start(volumes, no_docker_interface, **kwargs):
     """Start a docker agent"""
     from prefect.agent.docker import DockerAgent
@@ -524,6 +530,12 @@ _agents = {
     help="Disable presence of a Docker interface.",
     hidden=True,
 )
+@click.option(
+    "--docker-client-timeout",
+    default=None,
+    type=int,
+    hidden=True,
+)
 @click.pass_context
 def start(
     ctx,
@@ -549,6 +561,7 @@ def start(
     agent_address,
     hostname_label,
     storage_labels,
+    docker_client_timeout,
 ):
     """
     Start an agent.
@@ -611,6 +624,7 @@ def start(
         --no-docker-interface       Disable the check of a Docker interface on this machine.
                                     Note: This is mostly relevant for some Docker-in-Docker
                                     setups that users may be running their agent with.
+        --docker-client-timeout     Timeout for docker api requests
 
     \b
     Kubernetes Agent:
@@ -692,6 +706,7 @@ def start(
                 volumes=list(volume),
                 networks=tuple(network),
                 docker_interface=not no_docker_interface,
+                docker_client_timeout=docker_client_timeout,
             ).start()
         elif agent_option == "fargate":
             from_qualified_name(retrieved_agent)(
diff --git a/src/prefect/run_configs/ecs.py b/src/prefect/run_configs/ecs.py
index ed87a3d251..f3077bedad 100644
--- a/src/prefect/run_configs/ecs.py
+++ b/src/prefect/run_configs/ecs.py
@@ -23,13 +23,10 @@ class ECSRun(RunConfig):
 
     Args:
         - task_definition (dict, optional): An in-memory task definition spec
-            to use. The flow will be executed in a container named `flow` - if
-            a container named `flow` isn't part of the task definition, Prefect
-            will add a new container with that name.  See the
-            [ECS.Client.register_task_definition][1] docs for more information
-            on task definitions. Note that this definition will be stored
-            directly in Prefect Cloud/Server - use `task_definition_path`
-            instead if you wish to avoid this.
+            to use. See the [ECS.Client.register_task_definition][1] docs for
+            more information on task definitions. Note that this definition
+            will be stored directly in Prefect Cloud/Server - use
+            `task_definition_path` instead if you wish to avoid this.
         - task_definition_path (str, optional): Path to a task definition spec
             to use. If a local path (no file scheme, or a `file`/`local`
             scheme), the task definition will be loaded on initialization and
@@ -39,8 +36,7 @@ class ECSRun(RunConfig):
             `agent` (for paths local to the runtime agent)).
         - task_definition_arn (str, optional): A pre-registered task definition
             ARN to use (either `family`, `family:version`, or a full task
-            definition ARN). This task definition must include a container
-            named `flow` (which will be used to run the flow).
+            definition ARN).
         - image (str, optional): The image to use for this task. If not
             provided, will be either inferred from the flow's storage (if using
             `Docker` storage), or use the default configured on the agent.
@@ -135,12 +131,6 @@ ecs.html#ECS.Client.run_task
                 "Can only provide one of `task_definition`, `task_definition_path`, "
                 "or `task_definition_arn`"
             )
-        if task_definition_arn is not None and image is not None:
-            raise ValueError(
-                "Cannot provide both `task_definition_arn` and `image`, since an ECS "
-                "task's `image` can only be configured as part of the task definition"
-            )
-
         if task_definition_path is not None:
             parsed = parse_path(task_definition_path)
             if parsed.scheme == "file":


Test patch:
diff --git a/tests/agent/test_docker_agent.py b/tests/agent/test_docker_agent.py
index 92bdad2372..a671c7dd70 100644
--- a/tests/agent/test_docker_agent.py
+++ b/tests/agent/test_docker_agent.py
@@ -68,11 +68,12 @@ def test_docker_agent_config_options_populated(monkeypatch):
     api = MagicMock()
     monkeypatch.setattr("docker.APIClient", api)
 
-    agent = DockerAgent(base_url="url", no_pull=True)
+    agent = DockerAgent(base_url="url", no_pull=True, docker_client_timeout=123)
     assert agent.client.get_auth_token() == "TEST_TOKEN"
     assert agent.logger
     assert agent.no_pull
     assert api.call_args[1]["base_url"] == "url"
+    assert api.call_args[1]["timeout"] == 123
 
 
 def test_docker_agent_no_pull(api):
@@ -515,7 +516,11 @@ def test_docker_agent_deploy_flow_show_flow_logs(api, monkeypatch):
 
     process_kwargs = dict(
         target=_stream_container_logs,
-        kwargs={"base_url": agent.base_url, "container_id": "container_id"},
+        kwargs={
+            "base_url": agent.base_url,
+            "container_id": "container_id",
+            "timeout": 60,
+        },
     )
     process.assert_called_with(**process_kwargs)
     # Check all arguments to `multiprocessing.Process` are pickleable
diff --git a/tests/agent/test_ecs_agent.py b/tests/agent/test_ecs_agent.py
index 629fb44998..c94a7c6709 100644
--- a/tests/agent/test_ecs_agent.py
+++ b/tests/agent/test_ecs_agent.py
@@ -48,7 +48,9 @@ def aws(monkeypatch):
         "Subnets": [{"SubnetId": "test-subnet-id-1"}, {"SubnetId": "test-subnet-id-2"}]
     }
 
-    clients = dict(ecs=MagicMock(), ec2=ec2, s3=MagicMock())
+    clients = dict(
+        ecs=MagicMock(), ec2=ec2, s3=MagicMock(), resourcegroupstaggingapi=MagicMock()
+    )
 
     def get_client(key, **kwargs):
         if key not in clients:
@@ -343,6 +345,53 @@ class TestGenerateTaskDefinition:
         taskdef = self.generate_task_definition(run_config, storage)
         assert taskdef["containerDefinitions"][0]["image"] == expected
 
+    def test_generate_task_definition_command(self):
+        taskdef = self.generate_task_definition(ECSRun())
+        assert taskdef["containerDefinitions"][0]["command"] == [
+            "/bin/sh",
+            "-c",
+            "prefect execute flow-run",
+        ]
+
+    def test_generate_task_definition_resources(self):
+        taskdef = self.generate_task_definition(ECSRun(cpu="2048", memory="4096"))
+        assert taskdef["cpu"] == "2048"
+        assert taskdef["memory"] == "4096"
+
+    @pytest.mark.parametrize(
+        "on_run_config, on_agent, expected",
+        [
+            (None, None, None),
+            ("task-role-1", None, "task-role-1"),
+            (None, "task-role-2", "task-role-2"),
+            ("task-role-1", "task-role-2", "task-role-1"),
+        ],
+    )
+    def test_generate_task_definition_task_role_arn(
+        self, on_run_config, on_agent, expected
+    ):
+        taskdef = self.generate_task_definition(
+            ECSRun(task_role_arn=on_run_config), task_role_arn=on_agent
+        )
+        assert taskdef.get("taskRoleArn") == expected
+
+    @pytest.mark.parametrize(
+        "on_run_config, on_agent, expected",
+        [
+            (None, None, None),
+            ("execution-role-1", None, "execution-role-1"),
+            (None, "execution-role-2", "execution-role-2"),
+            ("execution-role-1", "execution-role-2", "execution-role-1"),
+        ],
+    )
+    def test_generate_task_definition_execution_role_arn(
+        self, on_run_config, on_agent, expected
+    ):
+        taskdef = self.generate_task_definition(
+            ECSRun(execution_role_arn=on_run_config), execution_role_arn=on_agent
+        )
+        assert taskdef.get("executionRoleArn") == expected
+
     def test_generate_task_definition_environment(self):
         run_config = ECSRun(
             image="test-image",
@@ -367,6 +416,9 @@ class TestGenerateTaskDefinition:
         env = {item["name"]: item["value"] for item in env_list}
         # Agent and run-config level envs are only set at runtime
         assert env == {
+            "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
+            "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
+            "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
             "PREFECT__CONTEXT__IMAGE": "test-image",
             "CUSTOM1": "VALUE1",
             "CUSTOM2": "VALUE2",
@@ -439,51 +491,6 @@ class TestGetRunTaskKwargs:
             "overrides": {"cpu": "2048", "memory": "2048", "taskRoleArn": "testing"},
         }
 
-    def test_get_run_task_kwargs_command(self):
-        kwargs = self.get_run_task_kwargs(ECSRun())
-        assert kwargs["overrides"]["containerOverrides"][0]["command"] == [
-            "/bin/sh",
-            "-c",
-            "prefect execute flow-run",
-        ]
-
-    def test_get_run_task_kwargs_resources(self):
-        kwargs = self.get_run_task_kwargs(ECSRun(cpu="2048", memory="4096"))
-        assert kwargs["overrides"]["cpu"] == "2048"
-        assert kwargs["overrides"]["memory"] == "4096"
-
-    @pytest.mark.parametrize(
-        "on_run_config, on_agent, expected",
-        [
-            (None, None, None),
-            ("task-role-1", None, "task-role-1"),
-            (None, "task-role-2", "task-role-2"),
-            ("task-role-1", "task-role-2", "task-role-1"),
-        ],
-    )
-    def test_get_task_run_kwargs_task_role_arn(self, on_run_config, on_agent, expected):
-        kwargs = self.get_run_task_kwargs(
-            ECSRun(task_role_arn=on_run_config), task_role_arn=on_agent
-        )
-        assert kwargs["overrides"].get("taskRoleArn") == expected
-
-    @pytest.mark.parametrize(
-        "on_run_config, on_agent, expected",
-        [
-            (None, None, None),
-            ("execution-role-1", None, "execution-role-1"),
-            (None, "execution-role-2", "execution-role-2"),
-            ("execution-role-1", "execution-role-2", "execution-role-1"),
-        ],
-    )
-    def test_get_task_run_kwargs_execution_role_arn(
-        self, on_run_config, on_agent, expected
-    ):
-        kwargs = self.get_run_task_kwargs(
-            ECSRun(execution_role_arn=on_run_config), execution_role_arn=on_agent
-        )
-        assert kwargs["overrides"].get("executionRoleArn") == expected
-
     def test_get_run_task_kwargs_environment(self, tmpdir, backend):
         path = str(tmpdir.join("kwargs.yaml"))
         with open(path, "w") as f:
@@ -512,9 +519,6 @@ class TestGetRunTaskKwargs:
         env_list = kwargs["overrides"]["containerOverrides"][0]["environment"]
         env = {item["name"]: item["value"] for item in env_list}
         assert env == {
-            "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
-            "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
-            "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
             "PREFECT__BACKEND": backend,
             "PREFECT__CLOUD__API": prefect.config.cloud.api,
             "PREFECT__CLOUD__AUTH_TOKEN": "TEST_TOKEN",
@@ -529,14 +533,57 @@ class TestGetRunTaskKwargs:
         }
 
 
+@pytest.mark.parametrize("kind", ["exists", "missing", "error"])
+def test_get_task_definition_arn(aws, kind):
+    if kind == "exists":
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}]
+        }
+        expected = "my-taskdef-arn"
+    elif kind == "missing":
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": []
+        }
+        expected = None
+    else:
+        from botocore.exceptions import ClientError
+
+        aws.resourcegroupstaggingapi.get_resources.side_effect = ClientError(
+            {}, "GetResources"
+        )
+        expected = None
+
+    run_config = ECSRun()
+    flow_run = GraphQLResult({"flow": GraphQLResult({"id": "flow-id", "version": 1})})
+    agent = ECSAgent()
+
+    res = agent.get_task_definition_arn(flow_run, run_config)
+    assert res == expected
+    kwargs = aws.resourcegroupstaggingapi.get_resources.call_args[1]
+    assert sorted(kwargs["TagFilters"], key=lambda x: x["Key"]) == [
+        {"Key": "prefect:flow-id", "Values": ["flow-id"]},
+        {"Key": "prefect:flow-version", "Values": ["1"]},
+    ]
+    assert kwargs["ResourceTypeFilters"] == ["ecs:task-definition"]
+
+
+def test_get_task_definition_arn_provided_task_definition_arn():
+    run_config = ECSRun(task_definition_arn="my-taskdef-arn")
+    flow_run = GraphQLResult({"flow": GraphQLResult({"id": "flow-id", "version": 1})})
+    agent = ECSAgent()
+
+    res = agent.get_task_definition_arn(flow_run, run_config)
+    assert res == "my-taskdef-arn"
+
+
 class TestDeployFlow:
-    def deploy_flow(self, run_config, storage=None, **kwargs):
+    def deploy_flow(self, run_config, **kwargs):
         agent = ECSAgent(**kwargs)
         flow_run = GraphQLResult(
             {
                 "flow": GraphQLResult(
                     {
-                        "storage": (storage or Local()).serialize(),
+                        "storage": Local().serialize(),
                         "id": "flow-id",
                         "version": 1,
                         "name": "Test Flow",
@@ -556,18 +603,11 @@ class TestDeployFlow:
         ):
             self.deploy_flow(LocalRun())
 
-    def test_deploy_flow_errors_if_mix_task_definition_arn_and_docker_storage(self):
-        with pytest.raises(
-            ValueError,
-            match="Cannot provide `task_definition_arn` when using `Docker` storage",
-        ):
-            self.deploy_flow(
-                ECSRun(task_definition_arn="my-taskdef-arn"),
-                storage=Docker(registry_url="test", image_name="name", image_tag="tag"),
-            )
-
     @pytest.mark.parametrize("run_config", [ECSRun(), UniversalRun(), None])
-    def test_deploy_flow_registers_new_task_definition(self, run_config, aws):
+    def test_deploy_flow_registers_taskdef_if_not_found(self, run_config, aws):
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": []
+        }
         aws.ecs.register_task_definition.return_value = {
             "taskDefinition": {"taskDefinitionArn": "my-taskdef-arn"}
         }
@@ -582,11 +622,19 @@ class TestDeployFlow:
         assert aws.ecs.run_task.called
         assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn"
         assert "my-task-arn" in res
-        assert aws.ecs.deregister_task_definition.called
-        assert (
-            aws.ecs.deregister_task_definition.call_args[1]["taskDefinition"]
-            == "my-taskdef-arn"
-        )
+
+    def test_deploy_flow_does_not_register_taskdef_if_found(self, aws):
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}]
+        }
+        aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]}
+
+        res = self.deploy_flow(ECSRun(run_task_kwargs={"enableECSManagedTags": True}))
+        assert not aws.ecs.register_task_definition.called
+        assert aws.ecs.run_task.called
+        assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn"
+        assert aws.ecs.run_task.call_args[1]["enableECSManagedTags"] is True
+        assert "my-task-arn" in res
 
     def test_deploy_flow_uses_provided_task_definition_arn(self, aws):
         aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]}
@@ -595,10 +643,12 @@ class TestDeployFlow:
         assert not aws.ecs.register_task_definition.called
         assert aws.ecs.run_task.called
         assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn"
-        assert not aws.ecs.deregister_task_definition.called
         assert "my-task-arn" in res
 
     def test_deploy_flow_run_task_fails(self, aws):
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}]
+        }
         aws.ecs.run_task.return_value = {
             "tasks": [],
             "failures": [{"reason": "my-reason"}],
@@ -606,16 +656,16 @@ class TestDeployFlow:
         with pytest.raises(ValueError) as exc:
             self.deploy_flow(ECSRun())
         assert aws.ecs.run_task.called
-        assert aws.ecs.deregister_task_definition.called
         assert "my-reason" in str(exc.value)
 
     def test_deploy_flow_forwards_run_task_kwargs(self, aws):
-        aws.ecs.register_task_definition.return_value = {
-            "taskDefinition": {"taskDefinitionArn": "my-taskdef-arn"}
+        aws.resourcegroupstaggingapi.get_resources.return_value = {
+            "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}]
         }
         aws.ecs.run_task.return_value = {"tasks": [{"taskArn": "my-task-arn"}]}
 
         res = self.deploy_flow(ECSRun(run_task_kwargs={"enableECSManagedTags": True}))
+        assert not aws.ecs.register_task_definition.called
         assert aws.ecs.run_task.called
         assert aws.ecs.run_task.call_args[1]["taskDefinition"] == "my-taskdef-arn"
         assert aws.ecs.run_task.call_args[1]["enableECSManagedTags"] is True
diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py
index d0d42981e2..4a283e6df9 100644
--- a/tests/cli/test_agent.py
+++ b/tests/cli/test_agent.py
@@ -59,7 +59,7 @@ def test_help(cmd):
             (
                 "--base-url testurl --no-pull --show-flow-logs --volume volume1 "
                 "--volume volume2 --network testnetwork1 --network testnetwork2 "
-                "--no-docker-interface"
+                "--no-docker-interface --docker-client-timeout 123"
             ),
             {
                 "base_url": "testurl",
@@ -68,6 +68,7 @@ def test_help(cmd):
                 "no_pull": True,
                 "show_flow_logs": True,
                 "docker_interface": False,
+                "docker_client_timeout": 123,
             },
         ),
         (
diff --git a/tests/run_configs/test_ecs.py b/tests/run_configs/test_ecs.py
index e04b60f99f..3a48557bcc 100644
--- a/tests/run_configs/test_ecs.py
+++ b/tests/run_configs/test_ecs.py
@@ -106,10 +106,6 @@ def test_task_definition_arn():
     assert config.task_definition is None
     assert config.task_definition_path is None
 
-    # Can't mix `image` and `task_definition_arn`
-    with pytest.raises(ValueError, match="task_definition_arn"):
-        ECSRun(task_definition_arn="my-task-definition", image="my-image")
-
 
 def test_task_definition():
     task_definition = {