Repository: PrefectHQ/prefect

Issue:
ECSRun should accept task_definition_arn
## Current behavior
The new ECSRun config offers great flexibility and fine-grained control over how a task is defined and executed, especially when custom values for specific parameters have to provided to override the default configuration.
However, in a CI/CD environment, you may want to create the task definition outside Prefect and then provide the task_definition_arn as the only parameter.




## Proposed behavior
Extend ECSRun config to accept a new parameter `task_definition_arn` that, if provided, take precedence over `task_definition` and `task_definition_path`.




## Example
In a CI/CD environment, you may want to manage task definition outside Prefect.
Moreover, being able to just pass the `task_definition_arn`, would also simplify Prefect code, because the definition of the infrastructure/environment is completely externalized.


Comments:

Patch:
diff --git a/changes/pr3681.yaml b/changes/pr3681.yaml
new file mode 100644
index 0000000000..e707ece5c9
--- /dev/null
+++ b/changes/pr3681.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Add `task_definition_arn` to `ECSRun` run-config - [#3681](https://github.com/PrefectHQ/prefect/pull/3681)"
diff --git a/src/prefect/agent/ecs/agent.py b/src/prefect/agent/ecs/agent.py
index 7d9df496a4..1052add1dc 100644
--- a/src/prefect/agent/ecs/agent.py
+++ b/src/prefect/agent/ecs/agent.py
@@ -310,8 +310,7 @@ class ECSAgent(Agent):
             )
             raise ValueError("Flow is missing a `run_config`")
 
-        # Check if a task definition already exists
-        taskdef_arn = self.lookup_task_definition_arn(flow_run)
+        taskdef_arn = self.get_task_definition_arn(flow_run, run_config)
         if taskdef_arn is None:
             # Register a new task definition
             self.logger.debug(
@@ -359,16 +358,22 @@ class ECSAgent(Agent):
             "prefect:flow-version": str(flow_run.flow.version),
         }
 
-    def lookup_task_definition_arn(self, flow_run: GraphQLResult) -> Optional[str]:
-        """Lookup an existing task definition ARN for a flow run.
+    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
diff --git a/src/prefect/run_configs/ecs.py b/src/prefect/run_configs/ecs.py
index 6e2f3d9183..f75f55c80f 100644
--- a/src/prefect/run_configs/ecs.py
+++ b/src/prefect/run_configs/ecs.py
@@ -11,11 +11,11 @@ class ECSRun(RunConfig):
 
     ECS Tasks are composed of task definitions and runtime parameters.
 
-    Task definitions can be configured using either the `task_definition` or
-    `task_definition_path` parameters. If neither is specified, the default
-    configured on the agent will be used. At runtime this task definition will
-    be registered once per flow version - subsequent runs of the same flow
-    version will reuse the existing definition.
+    Task definitions can be configured using either the `task_definition`,
+    `task_definition_path`, or `task_definition_arn` parameters. If neither is
+    specified, the default configured on the agent will be used. At runtime
+    this task definition will be registered once per flow version - subsequent
+    runs of the same flow version will reuse the existing definition.
 
     Runtime parameters can be specified via `run_task_kwargs`. These will be
     merged with any runtime parameters configured on the agent when starting
@@ -34,6 +34,9 @@ class ECSRun(RunConfig):
             Otherwise the task definition will be loaded at runtime on the
             agent.  Supported runtime file schemes include (`s3`, `gcs`, and
             `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).
         - 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.
@@ -99,6 +102,7 @@ ecs.html#ECS.Client.run_task
         *,
         task_definition: dict = None,
         task_definition_path: str = None,
+        task_definition_arn: str = None,
         image: str = None,
         env: dict = None,
         cpu: Union[int, str] = None,
@@ -109,9 +113,19 @@ ecs.html#ECS.Client.run_task
     ) -> None:
         super().__init__(labels=labels)
 
-        if task_definition is not None and task_definition_path is not None:
+        if (
+            sum(
+                [
+                    task_definition is not None,
+                    task_definition_path is not None,
+                    task_definition_arn is not None,
+                ]
+            )
+            > 1
+        ):
             raise ValueError(
-                "Cannot provide both `task_definition` and `task_definition_path`"
+                "Can only provide one of `task_definition`, `task_definition_path`, "
+                "or `task_definition_arn`"
             )
         if task_definition_path is not None:
             parsed = parse_path(task_definition_path)
@@ -127,6 +141,7 @@ ecs.html#ECS.Client.run_task
 
         self.task_definition = task_definition
         self.task_definition_path = task_definition_path
+        self.task_definition_arn = task_definition_arn
         self.image = image
         self.env = env
         self.cpu = cpu
diff --git a/src/prefect/serialization/run_config.py b/src/prefect/serialization/run_config.py
index 96129353cf..ca94f386ee 100644
--- a/src/prefect/serialization/run_config.py
+++ b/src/prefect/serialization/run_config.py
@@ -28,6 +28,7 @@ class ECSRunSchema(RunConfigSchemaBase):
 
     task_definition_path = fields.String(allow_none=True)
     task_definition = JSONCompatible(allow_none=True)
+    task_definition_arn = fields.String(allow_none=True)
     image = fields.String(allow_none=True)
     env = fields.Dict(keys=fields.String(), allow_none=True)
     cpu = fields.String(allow_none=True)


Test patch:
diff --git a/tests/agent/test_ecs_agent.py b/tests/agent/test_ecs_agent.py
index f6de39f5d0..79925e4d43 100644
--- a/tests/agent/test_ecs_agent.py
+++ b/tests/agent/test_ecs_agent.py
@@ -506,7 +506,7 @@ class TestGetRunTaskKwargs:
 
 
 @pytest.mark.parametrize("kind", ["exists", "missing", "error"])
-def test_lookup_task_definition_arn(aws, kind):
+def test_get_task_definition_arn(aws, kind):
     if kind == "exists":
         aws.resourcegroupstaggingapi.get_resources.return_value = {
             "ResourceTagMappingList": [{"ResourceARN": "my-taskdef-arn"}]
@@ -525,10 +525,11 @@ def test_lookup_task_definition_arn(aws, kind):
         )
         expected = None
 
+    run_config = ECSRun()
     flow_run = GraphQLResult({"flow": GraphQLResult({"id": "flow-id", "version": 1})})
     agent = ECSAgent()
 
-    res = agent.lookup_task_definition_arn(flow_run)
+    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"]) == [
@@ -538,6 +539,15 @@ def test_lookup_task_definition_arn(aws, kind):
     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, **kwargs):
         agent = ECSAgent(**kwargs)
@@ -599,6 +609,15 @@ class TestDeployFlow:
         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"}]}
+
+        res = self.deploy_flow(ECSRun(task_definition_arn="my-taskdef-arn"))
+        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 "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"}]
diff --git a/tests/run_configs/test_ecs.py b/tests/run_configs/test_ecs.py
index 64c04c303b..ffa2251573 100644
--- a/tests/run_configs/test_ecs.py
+++ b/tests/run_configs/test_ecs.py
@@ -10,6 +10,7 @@ def test_no_args():
     config = ECSRun()
     assert config.task_definition is None
     assert config.task_definition_path is None
+    assert config.task_definition_arn is None
     assert config.image is None
     assert config.env is None
     assert config.cpu is None
@@ -45,15 +46,29 @@ def test_labels():
     assert config.labels == {"a", "b"}
 
 
-def test_cant_specify_both_task_definition_and_task_definition_path():
-    with pytest.raises(ValueError, match="Cannot provide both"):
-        ECSRun(task_definition={}, task_definition_path="/some/path")
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        dict(task_definition={}, task_definition_path="/some/path"),
+        dict(task_definition={}, task_definition_arn="some_arn"),
+        dict(task_definition_path="/some/path", task_definition_arn="some_arn"),
+        dict(
+            task_definition={},
+            task_definition_path="/some/path",
+            task_definition_arn="some_arn",
+        ),
+    ],
+)
+def test_can_only_specify_task_definition_one_way(kwargs):
+    with pytest.raises(ValueError, match="Can only provide one of"):
+        ECSRun(**kwargs)
 
 
 def test_remote_task_definition_path():
     config = ECSRun(task_definition_path="s3://bucket/example.yaml")
     assert config.task_definition_path == "s3://bucket/example.yaml"
     assert config.task_definition is None
+    assert config.task_definition_arn is None
 
 
 @pytest.mark.parametrize("scheme", ["local", "file", None])
@@ -78,9 +93,17 @@ def test_local_task_definition_path(tmpdir, scheme):
     config = ECSRun(task_definition_path=task_definition_path)
 
     assert config.task_definition_path is None
+    assert config.task_definition_arn is None
     assert config.task_definition == task_definition
 
 
+def test_task_definition_arn():
+    config = ECSRun(task_definition_arn="my-task-definition")
+    assert config.task_definition_arn == "my-task-definition"
+    assert config.task_definition is None
+    assert config.task_definition_path is None
+
+
 def test_task_definition():
     task_definition = {
         "containerDefinitions": [
@@ -90,6 +113,7 @@ def test_task_definition():
     config = ECSRun(task_definition=task_definition)
 
     assert config.task_definition_path is None
+    assert config.task_definition_arn is None
     assert config.task_definition == task_definition
 
 
diff --git a/tests/serialization/test_run_configs.py b/tests/serialization/test_run_configs.py
index f4cb6c00ad..38d9c23c28 100644
--- a/tests/serialization/test_run_configs.py
+++ b/tests/serialization/test_run_configs.py
@@ -109,6 +109,7 @@ def test_serialize_docker_run(config):
                 ]
             }
         ),
+        ECSRun(task_definition_arn="my-task-definition"),
     ],
 )
 def test_serialize_ecs_run(config):
@@ -118,6 +119,7 @@ def test_serialize_ecs_run(config):
     fields = [
         "task_definition",
         "task_definition_path",
+        "task_definition_arn",
         "image",
         "env",
         "cpu",