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

[redartera]

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:

1. Users want to test workflows with varied resource configurations
2. 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 resources for tasks that don't have resources explicitly defined. The implementation:

1. Adds a default_resources parameter to SerializationSettings that, if set, applies to tasks with unspecified resources
2. Implements a CLI option --default-resources for both pyflyte run and pyflyte register commands
3. Adds a resource specification parser that handles both request-only and request-limit specifications
4. Ensures tasks with explicitly defined resources are not overridden

Example CLI usage:

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

# Specify both requests and limits
pyflyte register --default-resources 'cpu=(0.5,1);mem=(2Gi,4Gi);gpu=1' package

How was this patch tested?

Added comprehensive unit tests:

1. Tests for the resource specification parser in tests/flytekit/unit/interaction/test_click_types.py
2. Tests for default resource overrides in tests/flytekit/unit/core/test_serialization.py covering:
   • Regular @task decorators
   • @dynamic workflow tasks
   • @eager tasks
   • map_task functions
3. Integration tests for pyflyte register and pyflyte run --remote

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 PR adds a new CLI feature with '--default-resources' option to 'pyflyte run' and 'pyflyte register' commands, enabling users to specify default resource requests and limits for tasks. The implementation spans CLI commands, configuration settings, and repository functions to ensure proper propagation of default settings while preserving explicitly defined resources. These changes enhance resource management flexibility without requiring task code modifications.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 4

[flyte-bot]

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]

Codecov Report

Attention: Patch coverage is 12.12121% with 29 lines in your changes missing coverage. Please review.

Project coverage is 42.83%. Comparing base (8abecfe) to head (4bb75b7).

Files with missing lines	Patch %	Lines
flytekit/interaction/click_types.py	0.00%	24 Missing ⚠️
flytekit/core/python_auto_container.py	0.00%	4 Missing ⚠️
flytekit/tools/repo.py	0.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (8abecfe) and HEAD (4bb75b7). Click for more details.

HEAD has 11 uploads less than BASE

Flag	BASE (8abecfe)	HEAD (4bb75b7)
	28	17

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3229       +/-   ##
===========================================
- Coverage   89.81%   42.83%   -46.98%     
===========================================
  Files         136      289      +153     
  Lines        6599    25599    +19000     
  Branches        0     2932     +2932     
===========================================
+ Hits         5927    10966     +5039     
- Misses        672    14514    +13842     
- Partials        0      119      +119     

☔ 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]

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]

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

• Review rules
• General documentation
• FAQ

AI Code Review powered by

[flyte-bot]

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Enhanced Default Resource Handling	- register.py - Modified CLI command to accept and process default resource parameters. - run.py - Added CLI option for default resources and integrated resource specification callback. - __init__.py - Introduced a default_resources parameter and updated configuration settings to support it. - python_auto_container.py - Implemented logic to apply default resources when no explicit resource configuration is set. - click_types.py - Developed resource_spec_callback for parsing and validating default resource input strings. - remote.py - Integrated the default_resources parameter into remote task registration and workflow execution functions. - repo.py - Extended repository functions to include default resource override support during task registration.
Testing - Added and Updated Tests	- test_remote.py - Added integration tests for verifying default resource overrides in both registration and execution flows. - test_context_manager.py - Updated tests to account for default resource settings in serialization and context management. - test_serialization.py - Introduced unit tests that validate application of default resources and ensure explicit resource settings remain unmodified. - test_click_types.py - Included tests for resource_spec_callback to verify correct parsing and error handling of resource specifications.

[flyte-bot]

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

• Review rules
• General documentation
• FAQ

AI Code Review powered by

[flyte-bot]

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

• Review rules
• General documentation
• FAQ

AI Code Review powered by