Skip to content

[Pools] Add heterogeneity support.#8315

Merged
lloyd-brown merged 5 commits into
masterfrom
lloyd/pools-heterogeneous
Jan 10, 2026
Merged

[Pools] Add heterogeneity support.#8315
lloyd-brown merged 5 commits into
masterfrom
lloyd/pools-heterogeneous

Conversation

@lloyd-brown

@lloyd-brown lloyd-brown commented Dec 16, 2025

Copy link
Copy Markdown
Collaborator

Tested (run the relevant ones):

  • Code formatting: install pre-commit (auto-check on commit) or bash format.sh
  • Any manual or new tests for this PR (please specify below)
  • All smoke tests: /smoke-test (CI) or pytest tests/test_smoke.py (local)
  • Relevant individual tests: /smoke-test -k test_name (CI) or pytest tests/test_smoke.py::test_name (local)
  • Backward compatibility: /quicktest-core (CI) or pytest tests/smoke_tests/test_backward_compat.py (local)

@gemini-code-assist

gemini-code-assist Bot commented Dec 16, 2025

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @lloyd-brown, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the 'Pools' feature by introducing full support for heterogeneous resource configurations. This allows users to define jobs with flexible resource requirements (e.g., 'any_of' T4 or A100 GPUs), and the system will intelligently select and record the specific resource used from the available pool workers. The changes involve core logic updates for resource handling, database integration for tracking resolved resources, and the removal of prior limitations on heterogeneous pools, all validated by new, dedicated smoke tests.

Highlights

  • Heterogeneous Resource Support: The system now correctly handles tasks with 'any_of' or 'ordered' resource configurations, allowing for more flexible resource allocation in pools.
  • Dynamic Resource Resolution: When a job with heterogeneous resource requirements is scheduled, the specific resource chosen from the 'any_of'/'ordered' list is now recorded in the job's 'full_resources' in the database.
  • Removed Heterogeneity Restrictions: Previous validation that prevented the creation of pools with heterogeneous accelerators has been removed, enabling this feature.
  • Enhanced Resource-Aware Scheduling: The scheduling logic has been updated to account for heterogeneous resources, including a fallback mechanism for unresolved configurations and specific updates when a resource is selected.
  • New Smoke Tests: Comprehensive smoke tests have been added to validate the new heterogeneous pool functionality, including scenarios for 'any_of' accelerators and resource-aware scheduling.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@lloyd-brown

lloyd-brown commented Dec 16, 2025

Copy link
Copy Markdown
Collaborator Author

/smoke-test -k test_pools_heterogeneous

gemini-code-assist[bot]
gemini-code-assist Bot reviewed Dec 16, 2025
View reviewed changes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces heterogeneity support for resource pools, a significant enhancement. The changes are well-structured and include updates to resource configuration handling, scheduling logic for heterogeneous resources, and database state management. The new tests for heterogeneity and resource-aware scheduling are comprehensive. However, I've identified a critical issue in the _task_fits function, which is now used by the new scheduling logic. This function incorrectly compares CPU and memory requirements, which could lead to incorrect scheduling decisions.

Comment thread sky/serve/serve_utils.py
@lloyd-brown

lloyd-brown commented Dec 16, 2025

Copy link
Copy Markdown
Collaborator Author

/smoke-test -k test_pools_heterogeneous

@lloyd-brown

lloyd-brown commented Dec 16, 2025

Copy link
Copy Markdown
Collaborator Author

/smoke-test -k test_pools

@lloyd-brown lloyd-brown marked this pull request as ready for review December 16, 2025 18:34
SeungjinYang
SeungjinYang reviewed Jan 7, 2026
View reviewed changes
Comment thread sky/jobs/state.py Outdated
SeungjinYang
SeungjinYang reviewed Jan 7, 2026
View reviewed changes
Comment thread sky/serve/serve_utils.py
cg505
cg505 approved these changes Jan 10, 2026
View reviewed changes
Comment thread sky/serve/serve_utils.py
Comment thread sky/serve/serve_utils.py
@lloyd-brown lloyd-brown force-pushed the lloyd/pools-heterogeneous branch from 922166a to 732df90 Compare January 10, 2026 02:16
@lloyd-brown

lloyd-brown commented Jan 10, 2026

Copy link
Copy Markdown
Collaborator Author

/smoke-test --kubernetes

@lloyd-brown lloyd-brown merged commit 717a4cd into master Jan 10, 2026
21 checks passed
@lloyd-brown lloyd-brown deleted the lloyd/pools-heterogeneous branch January 10, 2026 03:50
@kevinmingtarja kevinmingtarja mentioned this pull request Feb 11, 2026
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@SeungjinYang SeungjinYang SeungjinYang left review comments
@cg505 cg505 cg505 approved these changes
+1 more reviewer
@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@lloyd-brown @cg505 @SeungjinYang