Repository navigation
perf(e2e): reorder tests by model affinity to minimize GPU swaps - #556
Conversation
📝 WalkthroughWalkthroughPer-test model affinity mapping added and persisted during collection; local backend set introduced. Collection reset clears affinity. Tests are reordered so cloud tests run first, then contiguous groups per model (preserving within-group order) to minimize model swaps. Changes
Sequence Diagram(s)sequenceDiagram
actor Pytest
participant Collector as Collection Logic
participant Affinity as Affinity Store
participant Reorder as Reordering Logic
participant Items as Test Items
Pytest->>Collector: pytest_collection_modifyitems(items)
Collector->>Items: iterate items to compute backend/model info
Collector->>Affinity: record _item_model_affinity[item.nodeid] = model_id or None
loop per item
Affinity->>Affinity: store mapping
end
Collector->>Reorder: call _reorder_to_minimize_model_swaps(items)
Reorder->>Items: move cloud (None) tests first
Reorder->>Items: then group items by model_id in first-appearance order, preserving intra-group order
Reorder->>Pytest: return reordered items
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
Summary of ChangesHello, 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 introduces significant optimizations to the test execution process by implementing a system to track and leverage GPU model affinity. The primary goal is to reduce the overhead associated with switching between different GPU models during test runs, thereby improving efficiency and potentially reducing overall test execution time. This is achieved by reordering tests to run contiguously based on their required GPU model, with cloud tests (which have no specific GPU model affinity) running first. Highlights
Changelog
Activity
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a test reordering mechanism to minimize GPU model swaps by grouping tests based on their required model. The implementation tracks model affinity for each test and reorders them in pytest_collection_modifyitems. While the overall approach is sound, I've identified a logic issue in how model affinity is assigned, which could cause tests using cloud backends to be incorrectly grouped with local GPU tests. My review includes a specific code suggestion to correct this, ensuring that only tests utilizing local GPUs are assigned model affinity.
8ae3e0d to
2884aa7
Compare
key4ng
changed the title
|
Hi @key4ng, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
… minimize GPU model swaps Signed-off-by: key4ng <rukeyang@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2884aa7e28
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@e2e_test/fixtures/hooks.py`:
- Around line 41-42: The hardcoded _LOCAL_BACKENDS frozenset can drift from the
authoritative source; change it to derive its values from the canonical
enum/constant (e.g., import and iterate over ConnectionMode or the infra
module's central constant) so the set is computed (filtering modes that require
local GPU workers) instead of being manually maintained, or at minimum add a
clear comment pointing to the exact definition of ConnectionMode/infra constant
(include module/class name) so maintainers know where to update when new
backends are added; update the reference in hooks.py to use that derived set
(symbol: _LOCAL_BACKENDS) and ensure tests still pass.
9880062 to
fb29df2
Compare
Signed-off-by: key4ng <rukeyang@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fb29df29ec
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Signed-off-by: key4ng <rukeyang@gmail.com>
fb29df2 to
c35fbe5
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@e2e_test/fixtures/hooks.py`:
- Around line 239-245: The affinity check is using the aggregated parametrize
marker values (backends) instead of the concrete backend for the current test
iteration; change the condition to look up the per-item backend from
item.callspec.params (e.g., backend = item.callspec.params.get("setup_backend")
or the real param name) and then check if that single backend is in
_LOCAL_BACKENDS, falling back to the existing `any(b in _LOCAL_BACKENDS for b in
backends)` only if callspec or the specific param is not present; update the
expression used in the if-block that references model_id, backends,
_LOCAL_BACKENDS and is_e2e to use this per-item backend variable.
4 participants
Description
Problem
E2E tests take ~16 minutes, with a significant portion spent on GPU model swaps rather than actual testing. Tests alternate between models (
openai/gpt-oss-20bfor Harmony tests,Qwen/Qwen2.5-14B-Instructfor Local tests) across different files, causing repeated MRU eviction + relaunch cycles (~35–100s each). Both models require 2 GPUs (tp=2), but only GPUs 0–1 are available for swapping.Solution
Reorder tests at collection time to group them by GPU model affinity — cloud tests first, then all tests for each model contiguously — reducing model swaps from ~7 to 1 and bringing E2E time from ~16 min down to ~12 min.
How it works
The reordering happens in
pytest_collection_modifyitems():Nonefor cloud-only tests_reorder_to_minimize_model_swaps()groups tests:Changes
e2e_test/fixtures/hooks.py: Added_item_model_affinitydict and_LOCAL_BACKENDSfrozenset to track each test's GPU model requirement. Populated affinity during the existing marker scan loop. Added_reorder_to_minimize_model_swaps()function that groups tests by model affinity (cloud first, then per-model). Updatedreset_collection_state()to clear the new dict.Test plan
Reordered tests to minimize model swaps: cloud=N, ModelA=N, ModelB=N🤖 Generated with Claude Code
Summary by CodeRabbit