Add ability to specify default resource requests/limits for tasks via `pyflyte run` and `pyflyte register` #3229

redartera · 2025-04-14T13:14:43Z

Why are the changes needed?

When registering tasks with Flyte, users often need flexibility to specify resource requirements at registration time without modifying code. This is particularly useful in scenarios where:

Users want to test workflows with varied resource configurations
Teams need to quickly adjust resources for existing tasks without code changes during rapid development phases.

Currently, the only options are to specify resources explicitly in task definitions or rely on cluster-level defaults, which lacks flexibility for selective overrides.

What changes were proposed in this pull request?

This PR introduces the ability to specify default resource requests and limits for tasks that don't have resources explicitly defined. The implementation:

Adds a default_resources parameter to SerializationSettings that, if set, applies to tasks with unspecified resources
Implements CLI options --resource-requests and --resource-limits for both pyflyte run and pyflyte register commands
Adds a resource specification parser (resource_callback) that handles comma-separated key-value pairs
Ensures tasks with explicitly defined resources are not overridden
Updates the serialization pipeline to apply default resources during task registration

Example CLI usage:

# Specify resource requests only
pyflyte run --remote --resource-requests 'cpu=1,mem=2Gi,gpu=1' workflow.py my_workflow

# Specify both requests and limits
pyflyte register --resource-requests 'cpu=500m,mem=1Gi' --resource-limits 'cpu=1,mem=2Gi' package

# Register with resource overrides
pyflyte register --resource-requests 'cpu=1300m,mem=1100Mi' --image myimage --project myproject --domain development workflow.py

How was this patch tested?

Added comprehensive unit and integration tests:

Tests for the resource specification parser in tests/flytekit/unit/interaction/test_click_types.py
Integration tests in tests/flytekit/integration/remote/test_remote.py for both:
- pyflyte register with resource overrides
- pyflyte run --remote with resource overrides
Tests verify that the registered tasks have the correct resource specifications applied
Tests ensure proper error handling for invalid resource specifications (wrong separators, invalid keys, duplicate keys)

The integration tests verify end-to-end functionality by:

Registering workflows with resource overrides
Fetching the created tasks from the remote
Validating that the task templates contain the expected resource specifications

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Summary by Bito

This pull request enhances the Flyte framework by allowing users to specify default resource requests and limits for tasks during registration and execution. It introduces new CLI options for `pyflyte run` and `pyflyte register`, along with improved help text for clarity. Comprehensive tests ensure the new functionality operates correctly and handles errors effectively.

flyte-bot · 2025-04-14T13:14:55Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Automatic Review Skipped - Draft PR

Bito didn't auto-review because this pull request is in draft status.
To trigger review, mark the PR as ready or type /review in the comment and save.
You can change draft PR review settings here, or contact the agent instance creator at [email protected].

codecov · 2025-04-14T14:36:15Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 87.67%. Comparing base (c085352) to head (79035a0).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3229      +/-   ##
==========================================
+ Coverage   83.58%   87.67%   +4.08%     
==========================================
  Files           3       88      +85     
  Lines         195     4267    +4072     
==========================================
+ Hits          163     3741    +3578     
- Misses         32      526     +494

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

flyte-bot · 2025-04-14T14:42:27Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Automatic Review Skipped - Draft PR

Bito didn't auto-review because this pull request is in draft status.
To trigger review, mark the PR as ready or type /review in the comment and save.
You can change draft PR review settings here, or contact the agent instance creator at [email protected].

flyte-bot · 2025-04-14T16:12:15Z

Code Review Agent Run #85e95a

Actionable Suggestions - 0

Review Details

Files reviewed - 10 · Commit Range: 0eb6200..74afde8
- flytekit/clis/sdk_in_container/register.py
- flytekit/clis/sdk_in_container/run.py
- flytekit/configuration/__init__.py
- flytekit/core/python_auto_container.py
- flytekit/interaction/click_types.py
- flytekit/remote/remote.py
- flytekit/tools/repo.py
- tests/flytekit/unit/core/test_context_manager.py
- tests/flytekit/unit/core/test_serialization.py
- tests/flytekit/unit/interaction/test_click_types.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

flyte-bot · 2025-04-14T16:27:05Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
New Feature - CLI Enhancements for Default Resources	- `register.py` - Added '--default-resources' CLI option and updated resource spec parsing; adjusted import structure. - `run.py` - Introduced support for '--default-resources' in run command and ensured proper parameter passing. - `click_types.py` - Implemented resource_spec_callback to parse default resource strings and updated associated test cases.
Feature Improvement - Resource Default Settings Implementation	- `__init__.py` - Integrated default_resources into configuration settings with updated registration parameters. - `python_auto_container.py` - Added logic to override task resource specs when not explicitly set. - `remote.py` - Propagated default_resources parameter for remote task and workflow registrations. - `repo.py` - Updated repository functions to include default resource specifications during task registration.
Testing - Enhanced Test Coverage for Default Resources	- `test_remote.py` - Introduced integration tests validating default resource overrides in both register and run commands. - `test_context_manager.py` - Incorporated tests to verify the context manager handles default_resources correctly. - `test_serialization.py` - Added tests to check that default resource settings apply to tasks without explicit resource declarations while preserving explicit settings. - `test_click_types.py` - Enhanced tests for resource_spec_callback to ensure robust parsing of various resource specification formats.

flyte-bot · 2025-04-14T17:32:19Z

Code Review Agent Run #f176ed

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 74afde8..4bb75b7
- tests/flytekit/integration/remote/test_remote.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

Signed-off-by: redartera <[email protected]>

flyte-bot · 2025-04-15T13:42:19Z

Code Review Agent Run #a4a15b

Actionable Suggestions - 0

Review Details

Files reviewed - 11 · Commit Range: 4bb75b7..59d0aec
- flytekit/clis/sdk_in_container/register.py
- flytekit/clis/sdk_in_container/run.py
- flytekit/configuration/__init__.py
- flytekit/core/python_auto_container.py
- flytekit/interaction/click_types.py
- flytekit/remote/remote.py
- flytekit/tools/repo.py
- tests/flytekit/integration/remote/test_remote.py
- tests/flytekit/unit/core/test_context_manager.py
- tests/flytekit/unit/core/test_serialization.py
- tests/flytekit/unit/interaction/test_click_types.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

flyte-bot · 2025-04-18T04:02:15Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Automatic Review Skipped - Branch Excluded

Bito didn't auto-review because the source or target branch is excluded from automatic reviews.
To trigger review, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

flyte-bot · 2025-04-29T12:27:55Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Skipped - No Changes

Bito didn't review because we did not see any changes in the PR to review.
To trigger review again, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

flyte-bot · 2025-05-02T13:43:15Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Skipped - No Changes

Bito didn't review because we did not see any changes in the PR to review.
To trigger review again, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

pingsutw

Thanks, it looks great. I think it's useful, and I'd like to get other reviewers' opinions as well

Signed-off-by: redartera <[email protected]>

flyte-bot · 2025-05-04T12:33:13Z

Code Review Agent Run #7dbdeb

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 59d0aec..e25d8d4
- flytekit/core/python_auto_container.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

flyte-bot · 2025-05-07T20:33:14Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Skipped - No Changes

Bito didn't review because we did not see any changes in the PR to review.
To trigger review again, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: redartera <[email protected]>

flyte-bot · 2025-05-12T21:11:32Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Code Review Skipped - No Changes

Bito didn't review because we did not see any changes in the PR to review.
To trigger review again, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

machichima · 2025-05-25T14:48:01Z

Hi @redartera ,

Could we use separate option instead? Like:

❯ pyflyte run file.py --cpu=1 --mem=2Gi

❯ pyflyte run file.py --cpu-limit=1 --mem-limit=2Gi

This has follow benefits:

Prevent typo when using, like using "cpi" instead of "cpu", click will directly gives error
We do not need to split string into parts
Easier to read (in my perspective)

redartera · 2025-05-26T17:47:04Z

Hi @redartera ,

Could we use separate option instead? Like:
❯ pyflyte run file.py --cpu=1 --mem=2Gi
❯ pyflyte run file.py --cpu-limit=1 --mem-limit=2Gi
This has follow benefits:

Prevent typo when using, like using "cpi" instead of "cpu", click will directly gives error

We do not need to split string into parts

Easier to read (in my perspective)

Makes sense @machichima - how about just --cpu=1 and --cpu=(1,3) to be consistent with the syntax of Resources and avoid the need for a separate click option for requests and limits?

Also, this would probably require we remove --default-resources - users would need to read the --help message to understand that this only overrides default resources and tasks with pre-defined resource specifications will remain intact

flyte-bot · 2025-05-26T18:21:22Z

Bito Review Skipped - No Changes Detected

Bito didn't review this pull request because we did not detect any changes in the pull request to review.

machichima · 2025-06-01T17:59:29Z

Hi @redartera

Makes sense @machichima - how about just --cpu=1 and --cpu=(1,3) to be consistent with the syntax of Resources and avoid the need for a separate click option for requests and limits?

I still prefer separating request and limit to two, like what majority of CLI tools did (e.g. kubectl), which provide a more explicit way of setting.

kubectl ... --limits=cpu=200m,memory=512Mi --requests=cpu=0,memory=0

It's unfortunate that click cannot set a flexible nargs, otherwise being able to use the following will be the best...

--default-resources cpu=1 mem=2Gi gpu=1 --limit-resources cpu=2 mem=4Gi
--default-resources cpu=1 mem=2Gi --limit-resources cpu=2

Also, this would probably require we remove --default-resources - users would need to read the --help message to understand that this only overrides default resources and tasks with pre-defined resource specifications will remain intact

I think it's fine as long as we document it well

Signed-off-by: redartera <[email protected]>

redartera · 2025-06-02T14:14:12Z

Hi @redartera

Makes sense @machichima - how about just --cpu=1 and --cpu=(1,3) to be consistent with the syntax of Resources and avoid the need for a separate click option for requests and limits?

I still prefer separating request and limit to two, like what majority of CLI tools did (e.g. kubectl), which provide a more explicit way of setting.

kubectl ... --limits=cpu=200m,memory=512Mi --requests=cpu=0,memory=0

It's unfortunate that click cannot set a flexible nargs, otherwise being able to use the following will be the best...
--default-resources cpu=1 mem=2Gi gpu=1 --limit-resources cpu=2 mem=4Gi
--default-resources cpu=1 mem=2Gi --limit-resources cpu=2
Also, this would probably require we remove --default-resources - users would need to read the --help message to understand that this only overrides default resources and tasks with pre-defined resource specifications will remain intact

I think it's fine as long as we document it well

@machichima I separated requests and limits into two options as recommended, PTAL. I chose --resource-limits as opposed to just --limits to avoid confusion with another --limit CLI option that exists in multiple places.

machichima

Just did a quick walk through. Overall looks good, just some nit.

Will look into detail and try it out later

flytekit/clis/sdk_in_container/run.py

Signed-off-by: redartera <[email protected]>

machichima · 2025-06-03T22:18:41Z

Verified that it works:

machichima

Thank you for the complete tests!!

Generally, my comment is about putting the resource settings into a variable to use when setting and asserting, making it easier if we want to extend or modify the test in the future

flytekit/interaction/click_types.py

tests/flytekit/integration/remote/test_remote.py

tests/flytekit/unit/core/test_context_manager.py

tests/flytekit/unit/core/test_serialization.py

Signed-off-by: redartera <[email protected]>

redartera · 2025-06-05T13:12:33Z

@machichima PTAL at 968a041

machichima · 2025-06-05T14:45:20Z

Hi @redartera ,
Just rerun the integration test as it failed. Please ping me again when the test passed

machichima

LGTM. Thanks!!

machichima · 2025-06-05T16:24:30Z

cc @pingsutw

redartera marked this pull request as ready for review April 14, 2025 15:06

redartera requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan, Future-Outlier and davidmirror-ops as code owners April 14, 2025 15:06

redartera force-pushed the pyflyte-override-resources branch from 4bb75b7 to 8291b77 Compare April 15, 2025 12:40

add support for default task resource overrides

59d0aec

Signed-off-by: redartera <[email protected]>

redartera force-pushed the pyflyte-override-resources branch from 8291b77 to 59d0aec Compare April 15, 2025 12:41

Merge branch 'flyteorg:master' into pyflyte-override-resources

412fdbb

Merge branch 'master' into pyflyte-override-resources

0a98881

Merge branch 'flyteorg:master' into pyflyte-override-resources

718b6ab

pingsutw reviewed May 3, 2025

View reviewed changes

apply lint

e25d8d4

Signed-off-by: redartera <[email protected]>

Merge branch 'flyteorg:master' into pyflyte-override-resources

23968a4

merge

46425d3

Signed-off-by: redartera <[email protected]>

Merge branch 'flyteorg:master' into pyflyte-override-resources

79035a0

redartera added 2 commits June 2, 2025 13:40

separate requests and limits per review comment

622dcdd

Signed-off-by: redartera <[email protected]>

merge

4dc3d9a

Signed-off-by: redartera <[email protected]>

machichima reviewed Jun 2, 2025

View reviewed changes

flytekit/clis/sdk_in_container/run.py Outdated Show resolved Hide resolved

flytekit/clis/sdk_in_container/run.py Outdated Show resolved Hide resolved

address comment

a5ca772

Signed-off-by: redartera <[email protected]>

redartera force-pushed the pyflyte-override-resources branch from 8d76fd5 to a5ca772 Compare June 3, 2025 21:20

one more

5c3cbe4

Signed-off-by: redartera <[email protected]>

machichima reviewed Jun 3, 2025

View reviewed changes

shorter docstring - put resources in variables for tests

968a041

Signed-off-by: redartera <[email protected]>

machichima approved these changes Jun 5, 2025

View reviewed changes

pingsutw approved these changes Jun 5, 2025

View reviewed changes

pingsutw merged commit 1a7133a into flyteorg:master Jun 5, 2025
227 of 228 checks passed

Add ability to specify default resource requests/limits for tasks via pyflyte run and pyflyte register #3229

Add ability to specify default resource requests/limits for tasks via pyflyte run and pyflyte register #3229

Conversation

redartera commented Apr 14, 2025 • edited by flyte-bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Check all the applicable boxes

Summary by Bito

Uh oh!

flyte-bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

codecov bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

flyte-bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

flyte-bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #85e95a

Uh oh!

flyte-bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changelist by Bito

Uh oh!

flyte-bot commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #f176ed

Uh oh!

flyte-bot commented Apr 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #a4a15b

Uh oh!

flyte-bot commented Apr 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

flyte-bot commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

flyte-bot commented May 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

pingsutw left a comment

Choose a reason for hiding this comment

Uh oh!

flyte-bot commented May 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #7dbdeb

Uh oh!

flyte-bot commented May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

flyte-bot commented May 12, 2025

Code Review Agent Run Status

Uh oh!

machichima commented May 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

redartera commented May 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

flyte-bot commented May 26, 2025

Uh oh!

machichima commented Jun 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

redartera commented Jun 2, 2025

Uh oh!

machichima left a comment

Choose a reason for hiding this comment

Add ability to specify default resource requests/limits for tasks via `pyflyte run` and `pyflyte register` #3229

Add ability to specify default resource requests/limits for tasks via `pyflyte run` and `pyflyte register` #3229

redartera commented Apr 14, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Apr 14, 2025 •

edited

Loading

codecov bot commented Apr 14, 2025 •

edited

Loading

flyte-bot commented Apr 14, 2025 •

edited

Loading

flyte-bot commented Apr 14, 2025 •

edited

Loading

flyte-bot commented Apr 14, 2025 •

edited

Loading

flyte-bot commented Apr 14, 2025 •

edited

Loading

flyte-bot commented Apr 15, 2025 •

edited

Loading

flyte-bot commented Apr 18, 2025 •

edited

Loading

flyte-bot commented Apr 29, 2025 •

edited

Loading

flyte-bot commented May 2, 2025 •

edited

Loading

flyte-bot commented May 4, 2025 •

edited

Loading

flyte-bot commented May 7, 2025 •

edited

Loading

machichima commented May 25, 2025 •

edited

Loading

redartera commented May 26, 2025 •

edited

Loading

machichima commented Jun 1, 2025 •

edited

Loading

machichima commented Jun 5, 2025 •

edited

Loading