You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "o-nikolas (via GitHub)" <gi...@apache.org> on 2025/03/28 21:35:56 UTC
[PR] Update ECS executor to support Task SDK [airflow]
o-nikolas opened a new pull request, #48513:
URL: https://github.com/apache/airflow/pull/48513
Follows the paradigm used in the k8s executor except the serialized workload is passed along as string instead of a file, since there is no (easy) way to plumb a file into the ECS container at run time.
In the future using a secret to pass along the workload could be another possible approach, but this requires the use of another cloud service (i.e. Secret Manager or Parameter Store).
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2031565043
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -114,6 +123,32 @@ def __init__(self, *args, **kwargs):
self.run_task_kwargs = self._load_run_kwargs()
+ def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
+ from airflow.executors import workloads
+
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+
+ def _process_workloads(self, workloads: list[workloads.All]) -> None:
+ from airflow.executors.workloads import ExecuteTask
+
+ # Airflow V3 version
+ for w in workloads:
+ if not isinstance(w, ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ # TODO: AIP-72 handle populating tokens once https://github.com/apache/airflow/issues/45107 is handled.
+ command = [w]
+ key = w.ti.key
Review Comment:
People are busy and I haven't heard back on this one. I can see the tokens being passed through so I think no action is needed here other than dropping that comment. Resolving this conversation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on PR #48513:
URL: https://github.com/apache/airflow/pull/48513#issuecomment-2766207330
Closing and re-opening to run test (they somehow did not run)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #48513:
URL: https://github.com/apache/airflow/pull/48513#issuecomment-2770825768
Follow up on this:
> > Overall looks good! Good job with figuring out all these things. Two questions:
> > ```
> > * Can we add unit tests? ECS executor changes are not tested
> > ```
>
> Sure, I can look at adding some. But most all of the actual behaviour of the executor itself hasn't changed. Just the interface to it (which should honestly be tested by the base executor once all this stuff moves into that class). But I'll look at adding a few.
Unit test added! It pretty comprehensively tests the task SDK path beginning to end. Let me know what you think!
>
> > ```
> > * It looks like it is but just double checking. With these changes, is ECS executor still compatible with AF2?
> > ```
>
> Yupp it should be still working with AF2, I can do a manual test to verify 👍
Manually setup a repro with AF2 and the current executor, everything works as expected!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #48513:
URL: https://github.com/apache/airflow/pull/48513#issuecomment-2770872016
CC @ashb @amoghrajesh
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas merged PR #48513:
URL: https://github.com/apache/airflow/pull/48513
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2028240448
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -74,16 +71,44 @@ def main():
parser = argparse.ArgumentParser(
description="Execute a workload in a Containerised executor using the task SDK."
)
- parser.add_argument(
- "input_file", help="Path to the input JSON file containing the execution workload payload."
+
+ # Create a mutually exclusive group to ensure that only one of the flags is set
+ group = parser.add_mutually_exclusive_group(required=True)
+ group.add_argument(
+ "--json-path",
+ help="Path to the input JSON file containing the execution workload payload.",
+ type=str,
+ )
+ group.add_argument(
+ "--json-string",
+ help="The JSON string itself containing the execution workload payload.",
+ type=str,
Review Comment:
Yeah sounds fair
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -462,6 +497,24 @@ def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None,
"""Save the task to be executed in the next sync by inserting the commands into a queue."""
if executor_config and ("name" in executor_config or "command" in executor_config):
raise ValueError('Executor Config should never override "name" or "command"')
+ if len(command) == 1:
+ from airflow.executors.workloads import ExecuteTask
+
+ if isinstance(command[0], ExecuteTask):
+ workload = command[0]
+ ser_input = workload.model_dump_json()
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
Review Comment:
I think this will show up in the logs? Are you ok with that?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py:
##########
@@ -61,7 +61,7 @@ class EcsQueuedTask:
key: TaskInstanceKey
command: CommandType
- queue: str
+ queue: str | None
Review Comment:
Can there be a None queue?
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -35,9 +35,7 @@
log = structlog.get_logger(logger_name=__name__)
-def execute_workload(input: str) -> None:
- from pydantic import TypeAdapter
-
+def execute_workload(workload) -> None:
Review Comment:
Can we add typing for workload?
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -74,16 +71,44 @@ def main():
parser = argparse.ArgumentParser(
description="Execute a workload in a Containerised executor using the task SDK."
)
- parser.add_argument(
- "input_file", help="Path to the input JSON file containing the execution workload payload."
+
+ # Create a mutually exclusive group to ensure that only one of the flags is set
+ group = parser.add_mutually_exclusive_group(required=True)
+ group.add_argument(
+ "--json-path",
+ help="Path to the input JSON file containing the execution workload payload.",
+ type=str,
+ )
+ group.add_argument(
+ "--json-string",
+ help="The JSON string itself containing the execution workload payload.",
+ type=str,
)
args = parser.parse_args()
- with open(args.input_file) as file:
- input_data = file.read()
+ from pydantic import TypeAdapter
- execute_workload(input_data)
+ from airflow.executors import workloads
+
+ decoder = TypeAdapter[workloads.All](workloads.All)
+ if args.json_path:
+ try:
+ with open(args.json_path) as file:
+ input_data = file.read()
+ workload = decoder.validate_json(input_data)
+ except Exception as e:
+ log.error("Failed to read file", error=str(e))
+ sys.exit(1)
+
+ elif args.json_string:
+ try:
+ workload = decoder.validate_json(args.json_string)
+ except Exception as e:
+ log.error("Failed to parse input JSON string", error=str(e))
+ sys.exit(1)
Review Comment:
A check if none are set?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029460358
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -35,9 +35,7 @@
log = structlog.get_logger(logger_name=__name__)
-def execute_workload(input: str) -> None:
- from pydantic import TypeAdapter
-
+def execute_workload(workload) -> None:
Review Comment:
Sure, will do
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029480952
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/utils.py:
##########
@@ -61,7 +61,7 @@ class EcsQueuedTask:
key: TaskInstanceKey
command: CommandType
- queue: str
+ queue: str | None
Review Comment:
I swear I was getting type error before, but I don't seem to now. I removed it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029458311
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -74,16 +71,44 @@ def main():
parser = argparse.ArgumentParser(
description="Execute a workload in a Containerised executor using the task SDK."
)
- parser.add_argument(
- "input_file", help="Path to the input JSON file containing the execution workload payload."
+
+ # Create a mutually exclusive group to ensure that only one of the flags is set
+ group = parser.add_mutually_exclusive_group(required=True)
+ group.add_argument(
+ "--json-path",
+ help="Path to the input JSON file containing the execution workload payload.",
+ type=str,
+ )
+ group.add_argument(
+ "--json-string",
+ help="The JSON string itself containing the execution workload payload.",
+ type=str,
)
args = parser.parse_args()
- with open(args.input_file) as file:
- input_data = file.read()
+ from pydantic import TypeAdapter
- execute_workload(input_data)
+ from airflow.executors import workloads
+
+ decoder = TypeAdapter[workloads.All](workloads.All)
+ if args.json_path:
+ try:
+ with open(args.json_path) as file:
+ input_data = file.read()
+ workload = decoder.validate_json(input_data)
+ except Exception as e:
+ log.error("Failed to read file", error=str(e))
+ sys.exit(1)
+
+ elif args.json_string:
+ try:
+ workload = decoder.validate_json(args.json_string)
+ except Exception as e:
+ log.error("Failed to parse input JSON string", error=str(e))
+ sys.exit(1)
Review Comment:
Argparse will do that for us, the mutually exclusive group is set to be required (see `required=True` on line 76), so if one isn't provided argparse will throw an exception with a message saying so.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck closed pull request #48513: Update ECS executor to support Task SDK
URL: https://github.com/apache/airflow/pull/48513
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #48513:
URL: https://github.com/apache/airflow/pull/48513#issuecomment-2766822038
> Overall looks good! Good job with figuring out all these things. Two questions:
>
> * Can we add unit tests? ECS executor changes are not tested
Sure, I can look at adding some. But most all of the actual behaviour of the executor itself hasn't changed. Just the interface to it (which should honestly be tested by the base executor once all this stuff moves into that class). But I'll look at adding a few.
>
> * It looks like it is but just double checking. With these changes, is ECS executor still compatible with AF2?
Yupp it should be still working with AF2, I can do a manual test to verify :+1:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2027602348
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -114,6 +123,32 @@ def __init__(self, *args, **kwargs):
self.run_task_kwargs = self._load_run_kwargs()
+ def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
+ from airflow.executors import workloads
+
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+
+ def _process_workloads(self, workloads: list[workloads.All]) -> None:
+ from airflow.executors.workloads import ExecuteTask
+
+ # Airflow V3 version
+ for w in workloads:
+ if not isinstance(w, ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ # TODO: AIP-72 handle populating tokens once https://github.com/apache/airflow/issues/45107 is handled.
+ command = [w]
+ key = w.ti.key
Review Comment:
Tokens is done now btw.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2027752631
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -114,6 +123,32 @@ def __init__(self, *args, **kwargs):
self.run_task_kwargs = self._load_run_kwargs()
+ def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
+ from airflow.executors import workloads
+
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+
+ def _process_workloads(self, workloads: list[workloads.All]) -> None:
+ from airflow.executors.workloads import ExecuteTask
+
+ # Airflow V3 version
+ for w in workloads:
+ if not isinstance(w, ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ # TODO: AIP-72 handle populating tokens once https://github.com/apache/airflow/issues/45107 is handled.
+ command = [w]
+ key = w.ti.key
Review Comment:
Fair, must have happened in main while I was branched. I'll drop that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029452199
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -114,6 +123,32 @@ def __init__(self, *args, **kwargs):
self.run_task_kwargs = self._load_run_kwargs()
+ def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
+ from airflow.executors import workloads
+
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+
+ def _process_workloads(self, workloads: list[workloads.All]) -> None:
+ from airflow.executors.workloads import ExecuteTask
+
+ # Airflow V3 version
+ for w in workloads:
+ if not isinstance(w, ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ # TODO: AIP-72 handle populating tokens once https://github.com/apache/airflow/issues/45107 is handled.
+ command = [w]
+ key = w.ti.key
Review Comment:
Revisiting this, I've removed the comment, but I'm not actually sure what work needs to be done. That TODO was just lifted from @amoghrajesh's work on the kubernetes executor (and I still see the TODO there). I assume the token is just passed through with the workload to the supervisor?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029452617
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py:
##########
@@ -462,6 +497,24 @@ def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None,
"""Save the task to be executed in the next sync by inserting the commands into a queue."""
if executor_config and ("name" in executor_config or "command" in executor_config):
raise ValueError('Executor Config should never override "name" or "command"')
+ if len(command) == 1:
+ from airflow.executors.workloads import ExecuteTask
+
+ if isinstance(command[0], ExecuteTask):
+ workload = command[0]
+ ser_input = workload.model_dump_json()
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
Review Comment:
For now yes, there isn't really an easier way to get this into the containers without other external services. So as a basic implementation, we're okay with this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update ECS executor to support Task SDK [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2023783804
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
Review Comment:
CC: @amoghrajesh
The changes here will be relevant for you to review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org