ECS Executor - Overriding Additional ECS Task Properties

[o-nikolas]

Description

From feedback here:

The executor config is scoped to overrides.containerOverrides. However there are relevant properties outside of overrides.containerOverrides that users may want to change.

For example, our ECS Cluster is actually composed of 3 capacity providers: A General-Purpose Capacity Provider (which is our cluster's default provider and runs on M7g instances), Memory-Optimized (R7g instances) and Compute-Optimized (C7g instances). My version of the ECS Executor allows users to set the appropriate Capacity Provider via the operator's executor_config param so that we can run our jobs in the most cost-efficient environment.

There are several other properties which airflow uses may want to set on a task-level, such as:

• overrides.cpu
• overrides.memory
• networkConfiguration
• tags
• placementConstraints
• placementStrategy

Use case/motivation

No response

Related issues

#34381

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[o-nikolas]

I started working on this one and there is a bit of nuance that I didn't expect.

Many of the config options are lists of dicts. Take tags for example, the type is:

 tags=[
        {
            'key': 'string',
            'value': 'string'
        },
    ],

So if you have a base set of run_task_kwargs configured in your airflow config such as:

AIRFLOW__AWS_ECS_EXECUTOR__RUN_TASK_KWARGS='{
"tags": [{"key": "fish", "value": "banana"}, {"key": "peanut", "value": "sauce"}],
"startedBy": "Niko",
"overrides": {"containerOverrides": [{"memory": 500, "environment":[{"name": "FOO", "value": "BAR"}]}]}}'

And in your executor_config on a particular task you specify:

@task(
    executor_config={
        "tags": [{"key": "FOO", "value": "BAR"}],
        "overrides": {"containerOverrides": [{"memory": 740}]},
    }
)
def hello_world():
    print("hello_world")

Should the final kwargs sent to ECS for tags be the sum of the two?: tags=[{'key': 'fish', 'value': 'banana'}, {'key': 'peanut', 'value': 'sauce'}, {"key": "FOO", "value": "BAR"}] or a zipping of the two where we update each dict in the list by index (copying/leaving the remainder if one list is smaller than the other): tags=[{'key': 'FOO', 'value': 'BAR'}, {'key': 'peanut', 'value': 'sauce'}] or should the executor_config completely override from?: "tags": [{'key': 'FOO', 'value': 'BAR'}],

There are many other configs that are lists of dicts in the ECS run_task API, I'm not sure one answer fits all of them. containerOverrides for example cannot be additive, we must update each dict with what is provided in the executor_config since each index in the list maps to a container, so a zipping update to each corresponding dict by index makes the most sense. However, something like tags above seems like it could be a different approach. We must evaluate many of the other configs as well (the full list is here).

I'd very much appreciate some input from @ferruzzi @shubham22 @syedahsn

As well as from @mshober who has an implementation of this, so I'd be interested to hear what approach they took and how it's working out for them:

My version of the ECS Executor allows users to set the appropriate Capacity Provider via the operator's executor_config param

Thanks for starting this @o-nikolas!

It is important to note that you can use propagateTags=TASK_DEFINITION as a way of setting default tags to all the tasks that Airflow launches. I haven't implemented tagging for my own environment yet, but that would be my approach to global tagging. That would render settings tags via run_task_kwargs obsolete and then tags set with executor_config could always be the complete set of tags that are passed to the RunTask API.

I'm also personally not a huge fan of the RUN_TASK_KWARGS config option. We already have explicit config options for several attributes (cluster, networkConfiguration, taskDefinition, platformVersion, launchType). I'd rather have more of the same. If we dictate what attributes can be set on a config level then it makes handling the executor_config much simpler. For example, I don't see the need of having overrides supported as a global config option when all of those properties can be set in the task definition.

As well as from @mshober who has an implementation of this, so I'd be interested to hear what approach they took and how it's working out for them:

My implementation is very specific and abstracted to my use case:

def _run_task_kwargs(self, task_id: TaskInstanceKeyType, cmd: CommandType, queue: str, exec_config: ExecutorConfigType) -> dict:
    capacity_provider = exec_config.get("capacity_provider")
    memory_reservation_mib: int | None = exec_config.get("memory_reservation_mib")
    memory_limit_mib: int | None = exec_config.get("memory_limit_mib")
    cpu_reservation: int | None = exec_config.get("cpu_reservation")

    run_task_api = deepcopy(self.run_task_kwargs)
    container_override = self.get_container(run_task_api["overrides"])
    container_override["command"] = cmd
    if cpu_reservation is not None:
        container_override["cpu"] = int(cpu_reservation)
    if memory_reservation_mib is not None:
        container_override["memoryReservation"] = int(memory_reservation_mib)
    if memory_limit_mib is not None:
        container_override["memory"] = int(memory_limit_mib)
    if capacity_provider:
        run_task_api["capacityProviderStrategy"] = [{"capacityProvider": capacity_provider}]
    return run_task_api

so my code is likely not too helpful. I'll spend some time thinking about how it can be improved to fit all use cases.

[o-nikolas]

Thanks for starting this @o-nikolas!

It is important to note that you can use propagateTags=TASK_DEFINITION as a way of setting default tags to all the tasks that Airflow launches. I haven't implemented tagging for my own environment yet, but that would be my approach to global tagging. That would render settings tags via run_task_kwargs obsolete and then tags set with executor_config could always be the complete set of tags that are passed to the RunTask API.

It's not really so much that tags are the specific case that needs a work around, just the example I chose, there are many others.

I'm also personally not a huge fan of the RUN_TASK_KWARGS config option. We already have explicit config options for several attributes (cluster, networkConfiguration, taskDefinition, platformVersion, launchType). I'd rather have more of the same. If we dictate what attributes can be set on a config level then it makes handling the executor_config much simpler. For example, I don't see the need of having overrides supported as a global config option when all of those properties can be set in the task definition.

I'm not in favour of this approach. There are many config options for ECS RunTask, creating a first class Airflow config option for each one is untenable. Especially because that couples us very tightly to the ECS API, every change it makes we must reflect in our config options. The few important configs we have now plus the catch-all of the run_task_kwargs gives users full control without overburdening the maintenance and quality of the underlying code.
We have the same situation with Airflow operators (which usually wrap an underlying Boto API or two), those operators cannot possibly take every parameter that the underlying Boto API can, so we pull out some of the mandatory and important ones, and then often have a <boto_api>_kwargs arg that allows the user to specify anything else.

As well as from @mshober who has an implementation of this, so I'd be interested to hear what approach they took and how it's working out for them:

My implementation is very specific and abstracted to my use case:

...

so my code is likely not too helpful. I'll spend some time thinking about how it can be improved to fit all use cases.

Mmm, yeah, that is very specific to your usecase it seems. Thanks for sharing it either way!

[ferruzzi]

I think tags should be additive.... AFAIK it will always be a list of len(1) and we should have an order of precedence; just tags[0]update() it so they get appended or overridden as expected seems reasonable to me. I think that's a pretty straightforward solution and aligns with how we handle tags elsewhere.

I'm pretty sure the container_overrides is the same and there can only be one container (for now?) so it would only ever be a len(1) list? In which case I like your solution. If I'm mistaken, then the issue becomes how can we identify the container? In your example above, if it is possible to have multiple containers, there is no indication about which container is being referenced so I don't think we can make that call. We can't assume that list[0] is always the same container if there is no container id stored to verify it. If there can only be one container then I don't see the issue with how we handle it everywhere else; set an order of precedence and update() up the chain to override or append as necessary.

I definitely agree that I wouldn't want an explicit config option for every possible field. Since we don't have any control over the boto API and don't get any kind of prior warning to changes they might make, we'll be stuck chasing it any time they decide to make a change. I suppose the flip side of that is that we're just passing the buck to the users instead, but as Niko mentioned it's how we handle it elsewhere so there's precedent.

[o-nikolas]

Thanks for taking the time to read and consider this issue @ferruzzi, it's a sticky one and I appreciate the brain cycles.

I think tags should be additive.... AFAIK it will always be a list of len(1) and we should have an order of precedence; just tags[0]update() it so they get appended or overridden as expected seems reasonable to me. I think that's a pretty straightforward solution and aligns with how we handle tags elsewhere.

If we go the purely additive approach, it means that there's no way for a task to use executor_config to override the tags provided by the run_task_kwargs config template. Which seems like a sad compromise, but if everyone else agrees then that's fair.

I'm pretty sure the container_overrides is the same and there can only be one container (for now?) so it would only ever be a len(1) list? In which case I like your solution. If I'm mistaken, then the issue becomes how can we identify the container? In your example above, if it is possible to have multiple containers, there is no indication about which container is being referenced so I don't think we can make that call. We can't assume that list[0] is always the same container if there is no container id stored to verify it. If there can only be one container then I don't see the issue with how we handle it everywhere else; set an order of precedence and update() up the chain to override or append as necessary.

Perhaps we can assert that the number of items in the overrides from the executor_config matches that we already have from the airflow config options. Also the override config itself contains sub configs that are themselves lists, so you run into the problem for each of those as well (should we do additive, override, zip, etc) sigh.

Plus there are other config options that are lists of dicts as well:

• capacityProviderStrategy
• placementConstraints
• placementStrategy

[ferruzzi]

If we go the purely additive approach, it means that there's no way for a task to use executor_config to override the tags provided by the run_task_kwargs config template.

Overwrite, yes. Remove, no. So maybe the solution is to make tags additive with values getting overridden if the name already exists, and some way to explicitly remove a tag? "If tag name starts with ~ then pop it; else add/update it"? I'd have to double check, but I think a valid tag name can start with a - so we can't use that, but I don't think tag names can include ~.

Plus there are other config options that are lists of dicts as well:

It's entirely possible I haven't thought this through down to the bone, but it seems to me that all of them you listed should be override-with-hierarchy. It's possible I'm missing a usecase, of course, but when I don't think it's unreasonable to expect the finer-grained placementStrategy to override the more general one, for example. It seems like the tags are the only one that jumps out at me as really needing add/merge, where you might want to add a tag for which environment launched the task, and another tag at the DAG level stating something, and keep building up that list.

[o-nikolas]

If we go the purely additive approach, it means that there's no way for a task to use executor_config to override the tags provided by the run_task_kwargs config template.

Mmm, ya that's an interesting idea. I'm a little concerned with coming up with something too bespoke for each config type. That's a lot of code/complexity to maintain. Interested to hear what others think.

Overwrite, yes. Remove, no. So maybe the solution is to make tags additive with values getting overridden if the name already exists, and some way to explicitly remove a tag? "If tag name starts with ~ then pop it; else add/update it"? I'd have to double check, but I think a valid tag name can start with a - so we can't use that, but I don't think tag names can include ~.

Plus there are other config options that are lists of dicts as well:

It's entirely possible I haven't thought this through down to the bone, but it seems to me that all of them you listed should be override-with-hierarchy. It's possible I'm missing a usecase, of course, but when I don't think it's unreasonable to expect the finer-grained placementStrategy to override the more general one, for example. It seems like the tags are the only one that jumps out at me as really needing add/merge, where you might want to add a tag for which environment launched the task, and another tag at the DAG level stating something, and keep building up that list.

Interesting, I'm not sure of all the edge cases there either. I don't use those configs enough to say for sure, maybe I'll do some more testing with those to see what's up.

[shubham22]

I'd vote for simplicity and keeping things consistent irrespective of the type of the field, which means executor_config should always override the run_task_kwargs for any field specified in both. If the user really cares about passing the tags that they've already defined in run_task_kwargs, then they can mention them again. If we go with additive approach, it is another documentation that a user must read and they wouldn't have a way to remove certain tags. Yes, you can design the system with ~ or something, but that's unnecessary logic on top of the logic (additive) which is already inconsistent.

[o-nikolas]

I'd vote for simplicity and keeping things consistent irrespective of the type of the field, which means executor_config should always override the run_task_kwargs for any field specified in both. If the user really cares about passing the tags that they've already defined in run_task_kwargs, then they can mention them again. If we go with additive approach, it is another documentation that a user must read and they wouldn't have a way to remove certain tags. Yes, you can design the system with ~ or something, but that's unnecessary logic on top of the logic (additive) which is already inconsistent.

This approach is definitely enticing. Thoughts @syedahsn @ferruzzi?

@shubham22, regarding your approach. In my original example above I have overrides specified in the run_task_kwargs template and in the executor_config. Do you think the smaller executor_config should overwrite the run_task kwargs like option 1 or 2?

1. {"containerOverrides": [{"memory": 740, "environment":[{"name": "FOO", "value": "BAR"}]}]}

2. {"containerOverrides": [{"memory": 740}]}