[BUG] Resource overrides doesn't works if the task is not defined in the same workflow file.

[xshen8888]

Describe the bug

Resource overrides works if the task foo() is defined in the same workflow file. However, if I move the task foo() to another package (so that it’s shared and re-usable task), with_overrides() like below will not be effective. Do you experience this, and what is the solution? In another workflow file a.py:

@task(
    requests=Resources(
        cpu="3",
        mem="3Gi",
    ),
)
def foo(...) -> None:
    ...

In my workflow file:
from a import foo
...
@workflow
def my_wf():
    foo(...).with_overrides(
        requests=Resources(
            cpu="5",
            mem="5Gi",
        ),
    )

Result: Task foo has been allocated resources cpu 3 & member 3Gi.

Expected behavior

Task foo should have been allocated resources cpu 5 & member 5Gi.

Additional context to reproduce

No response

Screenshots

No response

Are you sure this issue hasn't been raised already?

• [X] Yes

Have you read the Code of Conduct?

• [X] Yes

[dosubot[bot]]

🤖

Hey @xshen8888, nice to see you around here again! 🚀

Based on the information you've provided, it seems like the resource override might not be working when the task is moved to a different package because the GetTaskResources function in Flyte checks for resource overrides based on the id of the task. This id is a combination of the project, domain, and workflow name. If the task is moved to a different package, its id might change, and the function might not be able to find the correct resource override.

Here is the relevant code from the flyteadmin/pkg/manager/impl/util/resources.go file:

request := interfaces.ResourceRequest{
	ResourceType: admin.MatchableResource_TASK_RESOURCE,
}
if id != nil && len(id.Project) > 0 {
	request.Project = id.Project
}
if id != nil && len(id.Domain) > 0 {
	request.Domain = id.Domain
}
if id != nil && id.ResourceType == core.ResourceType_WORKFLOW && len(id.Name) > 0 {
	request.Workflow = id.Name
}

resource, err := resourceManager.GetResource(ctx, request)

In this code, the GetResource function is called with a ResourceRequest that includes the project, domain, and workflow name from the task's id. If the task is moved to a different package, these values might change, and the function might not be able to find the correct resource override.

As for the with_overrides() function, the context provided does not contain any specific information about it. Instead, it discusses a similar function called with_runtime_override(name: str). This function is used to register a "hook" on tasks, which allows the task to be identified by name and overridden with new config at runtime. The overrides can be specified in a simple map and the structure of the workflow graph does not matter/does not have to be shown. The overrides can be applied through multiple ways including UI, workflow decorator, launch plan, etc.

If you are looking for information specifically about with_overrides(), I'm afraid the provided context does not contain sufficient information to answer your question. Could you please provide more details or context about the with_overrides() function?

Sources

• flyteadmin/pkg/manager/impl/util/resources.go
• rfc/system/3553-config-override.md

About Dosu

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[xshen8888]

Hi @pingsutw , Will with_runtime_override(name: str) solve my problem? If so, could you point me to example code using with_runtime_override(name: str)? Thanks! Frank

[arbaobao]

#take

[Future-Outlier]

Doing it

[xshen8888]

Hi @Future-Outlier , @arbaobao , @dosu-bot, I cannot figure out how function with_overrides() can be applied to task, like in this code example: https://github.com/flyteorg/flytekit/blob/551924aa1c803da209591aa786c9f64d876874f9/tests/flytekit/unit/core/test_composition.py#L27 Could you explain where this function is defined? Thanks! Frank

[eapolinario]

@xshen8888 , can you double-check if you're not being hit by the issue described in https://github.com/flyteorg/flytekit/pull/2151?