Handle NVIDIA PyTorch updates #1525
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughAdds configurable NVIDIA PyTorch update management: new types/constants and settings, VirtualEnvironment policy state and checks, InstallationManager flow to conditionally prompt/apply updates (auto/defer/pinned), mirror recommendation/normalization, and unit tests for these behaviors. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant InstallMgr as InstallationManager
participant VEnv as VirtualEnvironment
participant Config as DesktopConfig
participant System
User->>InstallMgr: trigger install/validate
InstallMgr->>InstallMgr: maybeUpdateNvidiaTorch()
InstallMgr->>VEnv: check device & policy
alt policy allows check
VEnv->>System: getInstalledTorchPackageVersions() (spawn uv)
System-->>VEnv: installed versions
VEnv-->>InstallMgr: versions
InstallMgr->>InstallMgr: determine if update needed
alt update needed
InstallMgr->>Config: read lastPromptedVersion & policy
InstallMgr->>User: prompt update decision (auto/defer/pin)
User-->>InstallMgr: decision
InstallMgr->>VEnv: updateTorchUpdatePolicy(decision, pinned, version)
InstallMgr->>VEnv: ensureRecommendedNvidiaTorch(mirrorOverride?)
VEnv->>System: install/update torch packages
System-->>VEnv: success / failure
alt failure
InstallMgr->>User: prompt failure handling (silence/retry)
User-->>InstallMgr: response
end
end
else policy prevents update
InstallMgr-->>User: skip update
end
Possibly related PRs
✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
dosubot
bot
added
size:L
Contributor
There was a problem hiding this comment.
Pull request overview
This PR implements a prompt-based approach for NVIDIA PyTorch updates, allowing users to control when and how PyTorch packages are updated. It adds persistence for user update preferences and introduces custom mirror detection capabilities.
Changes:
- Adds user prompts before updating NVIDIA PyTorch packages with options to update, defer, or permanently silence updates
- Persists user update policy preferences (auto/defer/pinned) and package versions in desktop settings
- Adds detection and unit tests for custom torch mirror usage
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/constants.ts | Defines new types for torch update policy and pinned packages, plus recommended version key |
| src/store/desktopSettings.ts | Adds fields to persist torch update policy, pinned packages, and prompt history |
| src/virtualEnvironment.ts | Implements policy checking, custom mirror detection, and makes torch version methods public |
| src/main-process/comfyInstallation.ts | Passes torch update configuration from settings to VirtualEnvironment |
| src/install/installationManager.ts | Implements the prompt flow and update decision logic in maybeUpdateNvidiaTorch |
| tests/unit/virtualEnvironment.test.ts | Adds tests for custom mirror detection and torch version retrieval |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| return this.torchMirror !== getDeviceDefaultTorchMirror(this.selectedDevice); | ||
| } | ||
|
|
||
| private shouldSkipNvidiaTorchUpgrade(): boolean { |
There was a problem hiding this comment.
The isUsingCustomTorchMirror method is defined but never used in the codebase. According to the PR description, updates should be skipped when using a custom mirror, but this check is not implemented in maybeUpdateNvidiaTorch or anywhere else. Either this method should be integrated into the update flow, or it should be removed if it's not part of the intended feature set.
Suggested change
| private shouldSkipNvidiaTorchUpgrade(): boolean { | |
| private shouldSkipNvidiaTorchUpgrade(): boolean { | |
| if (this.isUsingCustomTorchMirror()) return true; |
src/install/installationManager.ts
Outdated
| type: 'question', | ||
| title: 'Update PyTorch?', | ||
| message: | ||
| 'Your NVIDIA PyTorch build is out of date. We can update to the recommended build for performance improvements. This may change memory use and/or compatibility with custom nodes', |
There was a problem hiding this comment.
Missing period at the end of the sentence.
Suggested change
| 'Your NVIDIA PyTorch build is out of date. We can update to the recommended build for performance improvements. This may change memory use and/or compatibility with custom nodes', | |
| 'Your NVIDIA PyTorch build is out of date. We can update to the recommended build for performance improvements. This may change memory use and/or compatibility with custom nodes.', |
Comment on lines
+325
to
+326
| test('returns parsed torch versions from uv json output', async ({ virtualEnv }) => { | ||
| const output = `[{"name":"aiohappyeyeballs","version":"2.6.1"},{"name":"aiohttp","version":"3.13.3"},{"name":"aiosignal","version":"1.4.0"},{"name":"alembic","version":"1.18.1"},{"name":"annotated-types","version":"0.7.0"},{"name":"attrs","version":"25.4.0"},{"name":"av","version":"16.1.0"},{"name":"certifi","version":"2026.1.4"},{"name":"charset-normalizer","version":"3.4.4"},{"name":"comfy-kitchen","version":"0.2.6"},{"name":"comfyui-embedded-docs","version":"0.4.0"},{"name":"comfyui-frontend-package","version":"1.36.14"},{"name":"comfyui-workflow-templates","version":"0.8.4"},{"name":"comfyui-workflow-templates-core","version":"0.3.88"},{"name":"comfyui-workflow-templates-media-api","version":"0.3.39"},{"name":"comfyui-workflow-templates-media-image","version":"0.3.55"},{"name":"comfyui-workflow-templates-media-other","version":"0.3.80"},{"name":"comfyui-workflow-templates-media-video","version":"0.3.38"},{"name":"einops","version":"0.8.1"},{"name":"filelock","version":"3.20.3"},{"name":"frozenlist","version":"1.8.0"},{"name":"fsspec","version":"2026.1.0"},{"name":"greenlet","version":"3.3.0"},{"name":"hf-xet","version":"1.2.0"},{"name":"huggingface-hub","version":"0.36.0"},{"name":"idna","version":"3.11"},{"name":"jinja2","version":"3.1.6"},{"name":"kornia","version":"0.8.2"},{"name":"kornia-rs","version":"0.1.10"},{"name":"mako","version":"1.3.10"},{"name":"markupsafe","version":"3.0.3"},{"name":"mpmath","version":"1.3.0"},{"name":"multidict","version":"6.7.0"},{"name":"networkx","version":"3.6.1"},{"name":"numpy","version":"2.4.1"},{"name":"nvidia-cublas","version":"13.0.0.19"},{"name":"nvidia-cuda-cupti","version":"13.0.48"},{"name":"nvidia-cuda-nvrtc","version":"13.0.48"},{"name":"nvidia-cuda-runtime","version":"13.0.48"},{"name":"nvidia-cudnn-cu13","version":"9.13.0.50"},{"name":"nvidia-cufft","version":"12.0.0.15"},{"name":"nvidia-cufile","version":"1.15.0.42"},{"name":"nvidia-curand","version":"10.4.0.35"},{"name":"nvidia-cusolver","version":"12.0.3.29"},{"name":"nvidia-cusparse","version":"12.6.2.49"},{"name":"nvidia-cusparselt-cu13","version":"0.8.0"},{"name":"nvidia-nccl-cu13","version":"2.27.7"},{"name":"nvidia-nvjitlink","version":"13.0.39"},{"name":"nvidia-nvshmem-cu13","version":"3.3.24"},{"name":"nvidia-nvtx","version":"13.0.39"},{"name":"packaging","version":"25.0"},{"name":"pillow","version":"12.1.0"},{"name":"pip","version":"24.0"},{"name":"propcache","version":"0.4.1"},{"name":"psutil","version":"7.2.1"},{"name":"pydantic","version":"2.12.5"},{"name":"pydantic-core","version":"2.41.5"},{"name":"pydantic-settings","version":"2.12.0"},{"name":"python-dotenv","version":"1.2.1"},{"name":"pyyaml","version":"6.0.3"},{"name":"regex","version":"2026.1.15"},{"name":"requests","version":"2.32.5"},{"name":"safetensors","version":"0.7.0"},{"name":"scipy","version":"1.17.0"},{"name":"sentencepiece","version":"0.2.1"},{"name":"setuptools","version":"80.9.0"},{"name":"spandrel","version":"0.4.1"},{"name":"sqlalchemy","version":"2.0.45"},{"name":"sympy","version":"1.14.0"},{"name":"tokenizers","version":"0.22.2"},{"name":"torch","version":"2.9.1+cu130"},{"name":"torchaudio","version":"2.9.1+cu130"},{"name":"torchsde","version":"0.2.6"},{"name":"torchvision","version":"0.24.1+cu130"},{"name":"tqdm","version":"4.67.1"},{"name":"trampoline","version":"0.1.2"},{"name":"transformers","version":"4.57.5"},{"name":"triton","version":"3.5.1"},{"name":"typing-extensions","version":"4.15.0"},{"name":"typing-inspection","version":"0.4.2"},{"name":"urllib3","version":"2.6.3"},{"name":"yarl","version":"1.22.0"}]`; |
There was a problem hiding this comment.
This test uses a very long inline JSON string which makes the test hard to read and maintain. Consider extracting this to a separate test fixture file or at minimum a separate constant above the test.
Suggested change
| test('returns parsed torch versions from uv json output', async ({ virtualEnv }) => { | |
| const output = `[{"name":"aiohappyeyeballs","version":"2.6.1"},{"name":"aiohttp","version":"3.13.3"},{"name":"aiosignal","version":"1.4.0"},{"name":"alembic","version":"1.18.1"},{"name":"annotated-types","version":"0.7.0"},{"name":"attrs","version":"25.4.0"},{"name":"av","version":"16.1.0"},{"name":"certifi","version":"2026.1.4"},{"name":"charset-normalizer","version":"3.4.4"},{"name":"comfy-kitchen","version":"0.2.6"},{"name":"comfyui-embedded-docs","version":"0.4.0"},{"name":"comfyui-frontend-package","version":"1.36.14"},{"name":"comfyui-workflow-templates","version":"0.8.4"},{"name":"comfyui-workflow-templates-core","version":"0.3.88"},{"name":"comfyui-workflow-templates-media-api","version":"0.3.39"},{"name":"comfyui-workflow-templates-media-image","version":"0.3.55"},{"name":"comfyui-workflow-templates-media-other","version":"0.3.80"},{"name":"comfyui-workflow-templates-media-video","version":"0.3.38"},{"name":"einops","version":"0.8.1"},{"name":"filelock","version":"3.20.3"},{"name":"frozenlist","version":"1.8.0"},{"name":"fsspec","version":"2026.1.0"},{"name":"greenlet","version":"3.3.0"},{"name":"hf-xet","version":"1.2.0"},{"name":"huggingface-hub","version":"0.36.0"},{"name":"idna","version":"3.11"},{"name":"jinja2","version":"3.1.6"},{"name":"kornia","version":"0.8.2"},{"name":"kornia-rs","version":"0.1.10"},{"name":"mako","version":"1.3.10"},{"name":"markupsafe","version":"3.0.3"},{"name":"mpmath","version":"1.3.0"},{"name":"multidict","version":"6.7.0"},{"name":"networkx","version":"3.6.1"},{"name":"numpy","version":"2.4.1"},{"name":"nvidia-cublas","version":"13.0.0.19"},{"name":"nvidia-cuda-cupti","version":"13.0.48"},{"name":"nvidia-cuda-nvrtc","version":"13.0.48"},{"name":"nvidia-cuda-runtime","version":"13.0.48"},{"name":"nvidia-cudnn-cu13","version":"9.13.0.50"},{"name":"nvidia-cufft","version":"12.0.0.15"},{"name":"nvidia-cufile","version":"1.15.0.42"},{"name":"nvidia-curand","version":"10.4.0.35"},{"name":"nvidia-cusolver","version":"12.0.3.29"},{"name":"nvidia-cusparse","version":"12.6.2.49"},{"name":"nvidia-cusparselt-cu13","version":"0.8.0"},{"name":"nvidia-nccl-cu13","version":"2.27.7"},{"name":"nvidia-nvjitlink","version":"13.0.39"},{"name":"nvidia-nvshmem-cu13","version":"3.3.24"},{"name":"nvidia-nvtx","version":"13.0.39"},{"name":"packaging","version":"25.0"},{"name":"pillow","version":"12.1.0"},{"name":"pip","version":"24.0"},{"name":"propcache","version":"0.4.1"},{"name":"psutil","version":"7.2.1"},{"name":"pydantic","version":"2.12.5"},{"name":"pydantic-core","version":"2.41.5"},{"name":"pydantic-settings","version":"2.12.0"},{"name":"python-dotenv","version":"1.2.1"},{"name":"pyyaml","version":"6.0.3"},{"name":"regex","version":"2026.1.15"},{"name":"requests","version":"2.32.5"},{"name":"safetensors","version":"0.7.0"},{"name":"scipy","version":"1.17.0"},{"name":"sentencepiece","version":"0.2.1"},{"name":"setuptools","version":"80.9.0"},{"name":"spandrel","version":"0.4.1"},{"name":"sqlalchemy","version":"2.0.45"},{"name":"sympy","version":"1.14.0"},{"name":"tokenizers","version":"0.22.2"},{"name":"torch","version":"2.9.1+cu130"},{"name":"torchaudio","version":"2.9.1+cu130"},{"name":"torchsde","version":"0.2.6"},{"name":"torchvision","version":"0.24.1+cu130"},{"name":"tqdm","version":"4.67.1"},{"name":"trampoline","version":"0.1.2"},{"name":"transformers","version":"4.57.5"},{"name":"triton","version":"3.5.1"},{"name":"typing-extensions","version":"4.15.0"},{"name":"typing-inspection","version":"0.4.2"},{"name":"urllib3","version":"2.6.3"},{"name":"yarl","version":"1.22.0"}]`; | |
| const UV_LIST_PACKAGES_JSON_OUTPUT = `[{"name":"aiohappyeyeballs","version":"2.6.1"},{"name":"aiohttp","version":"3.13.3"},{"name":"aiosignal","version":"1.4.0"},{"name":"alembic","version":"1.18.1"},{"name":"annotated-types","version":"0.7.0"},{"name":"attrs","version":"25.4.0"},{"name":"av","version":"16.1.0"},{"name":"certifi","version":"2026.1.4"},{"name":"charset-normalizer","version":"3.4.4"},{"name":"comfy-kitchen","version":"0.2.6"},{"name":"comfyui-embedded-docs","version":"0.4.0"},{"name":"comfyui-frontend-package","version":"1.36.14"},{"name":"comfyui-workflow-templates","version":"0.8.4"},{"name":"comfyui-workflow-templates-core","version":"0.3.88"},{"name":"comfyui-workflow-templates-media-api","version":"0.3.39"},{"name":"comfyui-workflow-templates-media-image","version":"0.3.55"},{"name":"comfyui-workflow-templates-media-other","version":"0.3.80"},{"name":"comfyui-workflow-templates-media-video","version":"0.3.38"},{"name":"einops","version":"0.8.1"},{"name":"filelock","version":"3.20.3"},{"name":"frozenlist","version":"1.8.0"},{"name":"fsspec","version":"2026.1.0"},{"name":"greenlet","version":"3.3.0"},{"name":"hf-xet","version":"1.2.0"},{"name":"huggingface-hub","version":"0.36.0"},{"name":"idna","version":"3.11"},{"name":"jinja2","version":"3.1.6"},{"name":"kornia","version":"0.8.2"},{"name":"kornia-rs","version":"0.1.10"},{"name":"mako","version":"1.3.10"},{"name":"markupsafe","version":"3.0.3"},{"name":"mpmath","version":"1.3.0"},{"name":"multidict","version":"6.7.0"},{"name":"networkx","version":"3.6.1"},{"name":"numpy","version":"2.4.1"},{"name":"nvidia-cublas","version":"13.0.0.19"},{"name":"nvidia-cuda-cupti","version":"13.0.48"},{"name":"nvidia-cuda-nvrtc","version":"13.0.48"},{"name":"nvidia-cuda-runtime","version":"13.0.48"},{"name":"nvidia-cudnn-cu13","version":"9.13.0.50"},{"name":"nvidia-cufft","version":"12.0.0.15"},{"name":"nvidia-cufile","version":"1.15.0.42"},{"name":"nvidia-curand","version":"10.4.0.35"},{"name":"nvidia-cusolver","version":"12.0.3.29"},{"name":"nvidia-cusparse","version":"12.6.2.49"},{"name":"nvidia-cusparselt-cu13","version":"0.8.0"},{"name":"nvidia-nccl-cu13","version":"2.27.7"},{"name":"nvidia-nvjitlink","version":"13.0.39"},{"name":"nvidia-nvshmem-cu13","version":"3.3.24"},{"name":"nvidia-nvtx","version":"13.0.39"},{"name":"packaging","version":"25.0"},{"name":"pillow","version":"12.1.0"},{"name":"pip","version":"24.0"},{"name":"propcache","version":"0.4.1"},{"name":"psutil","version":"7.2.1"},{"name":"pydantic","version":"2.12.5"},{"name":"pydantic-core","version":"2.41.5"},{"name":"pydantic-settings","version":"2.12.0"},{"name":"python-dotenv","version":"1.2.1"},{"name":"pyyaml","version":"6.0.3"},{"name":"regex","version":"2026.1.15"},{"name":"requests","version":"2.32.5"},{"name":"safetensors","version":"0.7.0"},{"name":"scipy","version":"1.17.0"},{"name":"sentencepiece","version":"0.2.1"},{"name":"setuptools","version":"80.9.0"},{"name":"spandrel","version":"0.4.1"},{"name":"sqlalchemy","version":"2.0.45"},{"name":"sympy","version":"1.14.0"},{"name":"tokenizers","version":"0.22.2"},{"name":"torch","version":"2.9.1+cu130"},{"name":"torchaudio","version":"2.9.1+cu130"},{"name":"torchsde","version":"0.2.6"},{"name":"torchvision","version":"0.24.1+cu130"},{"name":"tqdm","version":"4.67.1"},{"name":"trampoline","version":"0.1.2"},{"name":"transformers","version":"4.57.5"},{"name":"triton","version":"3.5.1"},{"name":"typing-extensions","version":"4.15.0"},{"name":"typing-inspection","version":"0.4.2"},{"name":"urllib3","version":"2.6.3"},{"name":"yarl","version":"1.22.0"}]`; | |
| test('returns parsed torch versions from uv json output', async ({ virtualEnv }) => { | |
| const output = UV_LIST_PACKAGES_JSON_OUTPUT; |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/virtualEnvironment.ts (1)
178-211: Replace non-null assertion with the already-defaulted property.
selectedDevice!violates the guideline to avoid non-null assertions. Sincethis.selectedDeviceis assigned a default value on line 186 and has typeTorchDeviceType(neverundefined), use it instead of the original parameter.🔧 Proposed fix
- this.torchMirror = fixDeviceMirrorMismatch(selectedDevice!, torchMirror); + this.torchMirror = fixDeviceMirrorMismatch(this.selectedDevice, torchMirror);
🤖 Fix all issues with AI agents
In `@src/constants.ts`:
- Around line 184-190: Add concise JSDoc comments above the exported types
TorchUpdatePolicy and TorchPinnedPackages describing their purpose and valid
values/fields; for TorchUpdatePolicy document the allowed string options
('auto', 'defer', 'pinned') and their behavior, and for TorchPinnedPackages
describe that it maps optional package names (torch, torchaudio, torchvision) to
version strings to pin specific versions. Keep comments short, one-to-two
sentences each and follow the existing JSDoc style in the file.
In `@src/install/installationManager.ts`:
- Around line 400-504: The maybeUpdateNvidiaTorch function lacks a guard for
custom torch mirrors; after the existing device check
(virtualEnvironment.selectedDevice !== 'nvidia') add a check using
virtualEnvironment.isUsingCustomTorchMirror() and if true log a message like
"Skipping NVIDIA PyTorch update because a custom torch mirror is configured."
and return; update references: maybeUpdateNvidiaTorch,
virtualEnvironment.isUsingCustomTorchMirror(), and keep the rest of the flow
unchanged so prompts and installs are skipped when a custom mirror is set.
In `@src/virtualEnvironment.ts`:
- Around line 275-289: The public method updateTorchUpdatePolicy in class
VirtualEnvironment lacks an explicit return type; add the explicit return type
`: void` to the method signature of updateTorchUpdatePolicy(policy:
TorchUpdatePolicy | undefined, pinnedPackages?: TorchPinnedPackages,
decisionVersion?: string) to match other public methods and coding guidelines so
the signature clearly indicates it returns nothing.
- Around line 291-301: The skip logic in shouldSkipNvidiaTorchUpgrade() must
also honor custom mirrors: update shouldSkipNvidiaTorchUpgrade() to return true
if isUsingCustomTorchMirror() is true in addition to the existing checks
(torchUpdatePolicy === 'pinned' or the 'defer' + torchUpdateDecisionVersion ===
NVIDIA_TORCH_RECOMMENDED_VERSION case) so that ensureRecommendedNvidiaTorch()
and needsNvidiaTorchUpgrade() will not trigger NVIDIA PyTorch upgrades when a
custom torchMirror (checked via isUsingCustomTorchMirror(), which compares
this.torchMirror against getDeviceDefaultTorchMirror(this.selectedDevice)) is
configured.
📜 Review details
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (6)
src/constants.tssrc/install/installationManager.tssrc/main-process/comfyInstallation.tssrc/store/desktopSettings.tssrc/virtualEnvironment.tstests/unit/virtualEnvironment.test.ts
🧰 Additional context used
📓 Path-based instructions (9)
**/*.{ts,mts,cts}
📄 CodeRabbit inference engine (.cursor/rules/typescript.mdc)
**/*.{ts,mts,cts}: Use features available in TypeScript 5.x
Use ESM only (type: module). Avoid CommonJS (require,module.exports)
Target ESNext runtime APIs and syntax. Prefer top-levelawaitonly when it improves clarity
Code must compile withstrictmode andnoImplicitAnyenabled, with zero implicitany
Useunknowninstead ofany. Narrow with type guards.anyis allowed only when interfacing with untyped code and must be localized and justified
Useinterfacefor public object shapes intended for extension/implementation
Usetypefor unions, mapped/conditional types, function types, and when composition is needed
Model domain variants with a discriminant field (e.g.,kind: 'X' | 'Y'). Use exhaustiveswitchstatements
Preferas constobjects and string/number literal unions overenum
Preferreadonlyproperties andReadonlyArray<T>where appropriate. Do not mutate function parameters
PreferT[]overArray<T>for readability
PreferRecord<Key, T>for simple maps; useMap/Setwhen iteration semantics or key types require it
Avoid non-null assertions (!). Use optional chaining, nullish coalescing, and explicit guards
Usesatisfiesto enforce object literal conformance without widening values
Prefer named exports. Default exports are allowed only when a module clearly has a single primary export
Useimport type { X } from '…'for type-only imports. Keep value and type imports separated when helpful for clarity and bundling
Exported functions, classes, and public APIs must have explicit parameter and return types. Internal locals can rely on inference
Prefer arrow functions for local functions and callbacks. Use function declarations when hoisting orthisbinding is required
Modules should be pure. Avoid executing logic at import time unless necessary (e.g., polyfills)
Prefer async/await over rawthen/catchchains. Keep linear flow and usetry/catchfor failures
Eitherawaitpromises or deliberately mark them...
Files:
src/constants.tssrc/main-process/comfyInstallation.tstests/unit/virtualEnvironment.test.tssrc/store/desktopSettings.tssrc/install/installationManager.tssrc/virtualEnvironment.ts
src/**/*.{ts,mts,cts}
📄 CodeRabbit inference engine (.cursor/rules/typescript.mdc)
Centralize reusable domain types in
src/**where discoverable. Avoid ad-hoc inline types for shared structures
Files:
src/constants.tssrc/main-process/comfyInstallation.tssrc/store/desktopSettings.tssrc/install/installationManager.tssrc/virtualEnvironment.ts
**/*.{ts,tsx}
📄 CodeRabbit inference engine (AGENTS.md)
**/*.{ts,tsx}: TypeScript: Do not use theanytype; useunknownwhen the type is genuinely unknown
Use JSDoc format for method documentation with common tags@paramand@return(omit@returnforvoidreturn type); use{@link}to reference symbols
Files:
src/constants.tssrc/main-process/comfyInstallation.tstests/unit/virtualEnvironment.test.tssrc/store/desktopSettings.tssrc/install/installationManager.tssrc/virtualEnvironment.ts
**/*.{ts,tsx,js,jsx}
📄 CodeRabbit inference engine (AGENTS.md)
Use ESLint and Prettier for code formatting; run
yarn lintandyarn formatbefore committing
Files:
src/constants.tssrc/main-process/comfyInstallation.tstests/unit/virtualEnvironment.test.tssrc/store/desktopSettings.tssrc/install/installationManager.tssrc/virtualEnvironment.ts
{src/main.ts,src/preload.ts,src/main-process/**/*.ts}
📄 CodeRabbit inference engine (AGENTS.md)
Main process code must maintain clean separation between main process, renderer, and preload scripts; follow Electron security best practices
Files:
src/main-process/comfyInstallation.ts
tests/**/*.{ts,mts,cts}
📄 CodeRabbit inference engine (.cursor/rules/typescript.mdc)
Test code may relax some strictness to maximize ergonomics. Keep type escapes localized
Files:
tests/unit/virtualEnvironment.test.ts
tests/unit/**/*.test.{js,ts}
📄 CodeRabbit inference engine (.cursor/rules/vitest.mdc)
tests/unit/**/*.test.{js,ts}: Usevitestfor unit testing
Do not build custom testing infrastructure; use Vitest and existing helpers
Store tests in thetests/unit/directory with directory structure matching that of the tested file insrc/
Tests should be cross-platform compatible; ensure proper use ofpath.normalize,path.join, andpath.sepfor Windows, macOS, and Linux compatibility
Mocks should be cleanly written, easy to understand, and reusable where possible
Prefer the use oftest.extendover loose variables in Vitest tests
Usetestinstead ofitfor defining tests
Files:
tests/unit/virtualEnvironment.test.ts
tests/unit/**/*.{test,spec}.ts
📄 CodeRabbit inference engine (AGENTS.md)
Use Vitest for creating unit tests; do not attempt to create custom testing infrastructure
Files:
tests/unit/virtualEnvironment.test.ts
{src/config/**/*.ts,src/store/**/*.ts}
📄 CodeRabbit inference engine (AGENTS.md)
Use the existing store patterns for configuration management in
src/store/andsrc/config/
Files:
src/store/desktopSettings.ts
🧠 Learnings (13)
📚 Learning: 2025-12-18T19:46:11.878Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-18T19:46:11.878Z
Learning: Configuration files like `package.json` and `todesktop.json` must be kept in sync with version numbers and build settings
Applied to files:
src/constants.tssrc/store/desktopSettings.ts
📚 Learning: 2025-11-25T20:49:40.925Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/integration-testing.mdc:0-0
Timestamp: 2025-11-25T20:49:40.925Z
Learning: Applies to tests/integration/**/*.spec.ts : Leverage TestEnvironment for state manipulation and error simulation instead of manual setup
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:49:40.925Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/integration-testing.mdc:0-0
Timestamp: 2025-11-25T20:49:40.925Z
Learning: Applies to tests/integration/**/*.spec.ts : Check CI behavior with process.env.CI in integration tests
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:50:25.371Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/vitest.mdc:0-0
Timestamp: 2025-11-25T20:50:25.371Z
Learning: Applies to tests/unit/**/*.test.{js,ts} : Do not build custom testing infrastructure; use Vitest and existing helpers
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:50:25.371Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/vitest.mdc:0-0
Timestamp: 2025-11-25T20:50:25.371Z
Learning: Applies to tests/unit/**/*.test.{js,ts} : Use `vitest` for unit testing
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:50:25.371Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/vitest.mdc:0-0
Timestamp: 2025-11-25T20:50:25.371Z
Learning: Applies to tests/unit/**/*.test.{js,ts} : Tests should be cross-platform compatible; ensure proper use of `path.normalize`, `path.join`, and `path.sep` for Windows, macOS, and Linux compatibility
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:50:25.371Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/vitest.mdc:0-0
Timestamp: 2025-11-25T20:50:25.371Z
Learning: Applies to tests/unit/**/*.test.{js,ts} : Prefer the use of `test.extend` over loose variables in Vitest tests
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-12-18T19:46:11.878Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-18T19:46:11.878Z
Learning: Applies to tests/unit/**/*.{test,spec}.ts : Use Vitest for creating unit tests; do not attempt to create custom testing infrastructure
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:49:50.649Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/playwright.mdc:0-0
Timestamp: 2025-11-25T20:49:50.649Z
Learning: Applies to tests/integration/**/*.{ts,js} : Use `path.join()` and `path.sep` to ensure file path handling works consistently across Windows, macOS, and Linux platforms
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:49:40.925Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/integration-testing.mdc:0-0
Timestamp: 2025-11-25T20:49:40.925Z
Learning: Applies to tests/integration/**/*.spec.ts : Import fixtures from testExtensions.ts, not raw Playwright
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-11-25T20:49:50.649Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: .cursor/rules/playwright.mdc:0-0
Timestamp: 2025-11-25T20:49:50.649Z
Learning: Applies to tests/integration/**/*.spec.ts : Prefer imports from `testExtensions.ts` over Playwright defaults in spec files
Applied to files:
tests/unit/virtualEnvironment.test.ts
📚 Learning: 2025-12-18T19:46:11.878Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-18T19:46:11.878Z
Learning: Applies to src/preload.ts : For preload script at `src/preload.ts`, implement secure IPC bridge between main and renderer processes without exposing sensitive APIs
Applied to files:
src/install/installationManager.ts
📚 Learning: 2025-12-18T19:46:11.878Z
Learnt from: CR
Repo: Comfy-Org/desktop PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-18T19:46:11.878Z
Learning: Applies to src/main.ts : For main process entry point at `src/main.ts`, implement Electron application lifecycle management and window initialization
Applied to files:
src/install/installationManager.ts
🧬 Code graph analysis (5)
src/main-process/comfyInstallation.ts (1)
src/store/desktopConfig.ts (1)
useDesktopConfig(13-16)
tests/unit/virtualEnvironment.test.ts (1)
src/virtualEnvironment.ts (1)
VirtualEnvironment(124-1203)
src/store/desktopSettings.ts (1)
src/constants.ts (2)
TorchUpdatePolicy(184-184)TorchPinnedPackages(186-190)
src/install/installationManager.ts (3)
src/main-process/comfyInstallation.ts (2)
ComfyInstallation(19-285)virtualEnvironment(26-28)src/virtualEnvironment.ts (1)
ProcessCallbacks(33-36)src/constants.ts (3)
NVIDIA_TORCH_RECOMMENDED_VERSION(222-222)NVIDIA_TORCH_VERSION(214-214)NVIDIA_TORCHVISION_VERSION(215-215)
src/virtualEnvironment.ts (2)
src/preload.ts (1)
TorchDeviceType(22-22)src/constants.ts (3)
TorchUpdatePolicy(184-184)TorchPinnedPackages(186-190)NVIDIA_TORCH_RECOMMENDED_VERSION(222-222)
🔇 Additional comments (18)
src/constants.ts (1)
221-222: LGTM: clear recommended-version key.src/main-process/comfyInstallation.ts (1)
75-77: LGTM: torch update settings are wired into the venv.tests/unit/virtualEnvironment.test.ts (3)
6-6: No review comment needed.
290-322: Good coverage for custom mirror detection.
324-342: Nice coverage for torch package parsing edge cases.src/store/desktopSettings.ts (2)
1-1: No review comment needed.
39-50: LGTM: settings shape covers new torch policy metadata.src/install/installationManager.ts (3)
8-15: No review comment needed.
101-102: No review comment needed.
392-392: LGTM: update flow now centralizes NVIDIA torch handling.src/virtualEnvironment.ts (8)
10-22: Imports align with new Torch policy types.
Looks good.
95-104: Device-aware default mirror mapping looks good.
140-142: Public Torch policy fields added cleanly.
669-673: Fallback to device default mirror is sensible.
694-711: Policy-based skip and default-mirror fallback look solid.
743-748: Public torch version reader exposure looks fine.
810-817: Out-of-date helper reads well.
1041-1044: Skip check integration is clean.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
Comment on lines
+184
to
+190
| export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned'; | ||
|
|
||
| export type TorchPinnedPackages = { | ||
| torch?: string; | ||
| torchaudio?: string; | ||
| torchvision?: string; | ||
| }; |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Add brief JSDoc for new torch policy types.
These are exported public types; a short description helps consumers and aligns with the repo’s documentation expectations. As per coding guidelines, please add concise JSDoc.
✍️ Suggested doc additions
+/** How NVIDIA PyTorch updates should be handled. */
export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned';
+/** Pinned NVIDIA torch package versions when updates are disabled. */
export type TorchPinnedPackages = {
torch?: string;
torchaudio?: string;
torchvision?: string;
};📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned'; | |
| export type TorchPinnedPackages = { | |
| torch?: string; | |
| torchaudio?: string; | |
| torchvision?: string; | |
| }; | |
| /** How NVIDIA PyTorch updates should be handled. */ | |
| export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned'; | |
| /** Pinned NVIDIA torch package versions when updates are disabled. */ | |
| export type TorchPinnedPackages = { | |
| torch?: string; | |
| torchaudio?: string; | |
| torchvision?: string; | |
| }; |
🤖 Prompt for AI Agents
In `@src/constants.ts` around lines 184 - 190, Add concise JSDoc comments above
the exported types TorchUpdatePolicy and TorchPinnedPackages describing their
purpose and valid values/fields; for TorchUpdatePolicy document the allowed
string options ('auto', 'defer', 'pinned') and their behavior, and for
TorchPinnedPackages describe that it maps optional package names (torch,
torchaudio, torchvision) to version strings to pin specific versions. Keep
comments short, one-to-two sentences each and follow the existing JSDoc style in
the file.
Comment on lines
400
to
504
| private async maybeUpdateNvidiaTorch(installation: ComfyInstallation, callbacks: ProcessCallbacks): Promise<void> { | ||
| const virtualEnvironment = installation.virtualEnvironment; | ||
| if (virtualEnvironment.selectedDevice !== 'nvidia') return; | ||
|
|
||
| const config = useDesktopConfig(); | ||
| const updatePolicy = config.get('torchUpdatePolicy'); | ||
| if (updatePolicy === 'pinned') { | ||
| log.info('Skipping NVIDIA PyTorch update because updates are pinned.'); | ||
| return; | ||
| } | ||
|
|
||
| const installedVersions = await virtualEnvironment.getInstalledTorchPackageVersions(); | ||
| if (!installedVersions) { | ||
| log.warn('Skipping NVIDIA PyTorch update because installed versions could not be read.'); | ||
| return; | ||
| } | ||
|
|
||
| const isOutOfDate = await virtualEnvironment.isNvidiaTorchOutOfDate(installedVersions); | ||
| if (!isOutOfDate) return; | ||
|
|
||
| const recommendedVersion = NVIDIA_TORCH_RECOMMENDED_VERSION; | ||
| const lastPromptedVersion = config.get('torchLastPromptedVersion'); | ||
|
|
||
| if (config.get('torchOutOfDateRecommendedVersion') !== recommendedVersion) { | ||
| config.set('torchOutOfDateRecommendedVersion', recommendedVersion); | ||
| config.set('torchOutOfDatePackages', installedVersions); | ||
| } | ||
|
|
||
| if (updatePolicy === 'defer' && lastPromptedVersion === recommendedVersion) { | ||
| log.info('Skipping NVIDIA PyTorch update because updates are deferred for this version.'); | ||
| return; | ||
| } | ||
|
|
||
| const autoPolicy = updatePolicy === undefined || updatePolicy === 'auto'; | ||
| const updateApproved = autoPolicy && lastPromptedVersion === recommendedVersion; | ||
|
|
||
| if (!updateApproved) { | ||
| const currentTorch = installedVersions.torch ?? 'unknown'; | ||
| const currentTorchaudio = installedVersions.torchaudio ?? 'unknown'; | ||
| const currentTorchvision = installedVersions.torchvision ?? 'unknown'; | ||
|
|
||
| const { response } = await this.appWindow.showMessageBox({ | ||
| type: 'question', | ||
| title: 'Update PyTorch?', | ||
| message: | ||
| 'Your NVIDIA PyTorch build is out of date. We can update to the recommended build for performance improvements. This may change memory use and/or compatibility with custom nodes', | ||
| detail: [ | ||
| `Current: torch ${currentTorch}, torchaudio ${currentTorchaudio}, torchvision ${currentTorchvision}`, | ||
| `Recommended: torch ${NVIDIA_TORCH_VERSION}, torchaudio ${NVIDIA_TORCH_VERSION}, torchvision ${NVIDIA_TORCHVISION_VERSION}`, | ||
| ].join('\n'), | ||
| buttons: ['Update PyTorch', 'Ask again later', 'Silence until next version', 'Silence forever'], | ||
| defaultId: 0, | ||
| cancelId: 1, | ||
| }); | ||
|
|
||
| switch (response) { | ||
| case 1: | ||
| log.info('Deferring NVIDIA PyTorch update prompt.'); | ||
| return; | ||
| case 2: | ||
| config.set('torchLastPromptedVersion', recommendedVersion); | ||
| config.set('torchUpdatePolicy', 'defer'); | ||
| config.delete('torchPinnedPackages'); | ||
| virtualEnvironment.updateTorchUpdatePolicy('defer', undefined, recommendedVersion); | ||
| return; | ||
| case 3: | ||
| config.set('torchLastPromptedVersion', recommendedVersion); | ||
| config.set('torchUpdatePolicy', 'pinned'); | ||
| config.set('torchPinnedPackages', installedVersions); | ||
| virtualEnvironment.updateTorchUpdatePolicy('pinned', installedVersions, recommendedVersion); | ||
| return; | ||
| default: | ||
| config.set('torchLastPromptedVersion', recommendedVersion); | ||
| config.set('torchUpdatePolicy', 'auto'); | ||
| config.delete('torchPinnedPackages'); | ||
| config.delete('torchUpdateFailureSilencedVersion'); | ||
| virtualEnvironment.updateTorchUpdatePolicy('auto', undefined, recommendedVersion); | ||
| } | ||
| } else { | ||
| virtualEnvironment.updateTorchUpdatePolicy('auto', undefined, recommendedVersion); | ||
| } | ||
|
|
||
| try { | ||
| await virtualEnvironment.ensureRecommendedNvidiaTorch(callbacks); | ||
| config.delete('torchUpdateFailureSilencedVersion'); | ||
| } catch (error) { | ||
| log.error('Error updating NVIDIA PyTorch packages:', error); | ||
| if (config.get('torchUpdateFailureSilencedVersion') === recommendedVersion) return; | ||
|
|
||
| const { response } = await this.appWindow.showMessageBox({ | ||
| type: 'warning', | ||
| title: 'PyTorch update failed', | ||
| message: | ||
| 'We could not install the recommended NVIDIA PyTorch build. This may be because your configured torch mirror does not provide it.', | ||
| detail: 'We will retry the update on each startup.', | ||
| buttons: ['OK', "Don't show again"], | ||
| defaultId: 0, | ||
| cancelId: 0, | ||
| }); | ||
|
|
||
| if (response === 1) { | ||
| config.set('torchUpdateFailureSilencedVersion', recommendedVersion); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
rg -n "isUsingCustomTorchMirror|ensureRecommendedNvidiaTorch|torchMirror" src/ -t tsRepository: Comfy-Org/desktop
Length of output: 2404
🏁 Script executed:
sed -n '694,740p' src/virtualEnvironment.tsRepository: Comfy-Org/desktop
Length of output: 2001
🏁 Script executed:
grep -n "isUsingCustomTorchMirror" src/install/installationManager.tsRepository: Comfy-Org/desktop
Length of output: 43
🏁 Script executed:
grep -rn "isUsingCustomTorchMirror" src/ -t tsRepository: Comfy-Org/desktop
Length of output: 171
🏁 Script executed:
rg -n "isUsingCustomTorchMirror" src/Repository: Comfy-Org/desktop
Length of output: 130
Add early return for custom torch mirrors.
The method isUsingCustomTorchMirror() exists but is unused in the torch update flow. Without a guard at the start of maybeUpdateNvidiaTorch(), users with custom mirrors will still see update prompts and attempts, contradicting the PR objective to skip updates when a custom mirror is configured.
Add the check after the device check:
if (virtualEnvironment.isUsingCustomTorchMirror()) {
log.info('Skipping NVIDIA PyTorch update because a custom torch mirror is configured.');
return;
}
🤖 Prompt for AI Agents
In `@src/install/installationManager.ts` around lines 400 - 504, The
maybeUpdateNvidiaTorch function lacks a guard for custom torch mirrors; after
the existing device check (virtualEnvironment.selectedDevice !== 'nvidia') add a
check using virtualEnvironment.isUsingCustomTorchMirror() and if true log a
message like "Skipping NVIDIA PyTorch update because a custom torch mirror is
configured." and return; update references: maybeUpdateNvidiaTorch,
virtualEnvironment.isUsingCustomTorchMirror(), and keep the rest of the flow
unchanged so prompts and installs are skipped when a custom mirror is set.
Comment on lines
+275
to
+289
| updateTorchUpdatePolicy( | ||
| policy: TorchUpdatePolicy | undefined, | ||
| pinnedPackages?: TorchPinnedPackages, | ||
| decisionVersion?: string | ||
| ) { | ||
| this.torchUpdatePolicy = policy; | ||
| if (pinnedPackages !== undefined) { | ||
| this.torchPinnedPackages = pinnedPackages; | ||
| } else if (policy !== 'pinned') { | ||
| this.torchPinnedPackages = undefined; | ||
| } | ||
| if (decisionVersion !== undefined) { | ||
| this.torchUpdateDecisionVersion = decisionVersion; | ||
| } | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's check the file and locate the class definition and the method in question
head -300 src/virtualEnvironment.ts | tail -100Repository: Comfy-Org/desktop
Length of output: 3646
🏁 Script executed:
# Also check the beginning of the file to see class export
head -50 src/virtualEnvironment.tsRepository: Comfy-Org/desktop
Length of output: 1725
🏁 Script executed:
# Check if the class is exported
grep -n "export.*class\|class.*{" src/virtualEnvironment.ts | head -5Repository: Comfy-Org/desktop
Length of output: 138
🏁 Script executed:
# Search for all public methods in VirtualEnvironment class to check consistency
rg -A 2 "^\s+(async\s+)?[a-zA-Z]+\(" src/virtualEnvironment.ts | head -80Repository: Comfy-Org/desktop
Length of output: 2875
🏁 Script executed:
# More targeted: find methods and check for return type annotations
ast-grep --pattern 'class VirtualEnvironment {
$$$
$_($_) {
$$$
}
$$$
}'Repository: Comfy-Org/desktop
Length of output: 43
🏁 Script executed:
# Let's get a cleaner view of the class methods and their signatures
sed -n '124,400p' src/virtualEnvironment.ts | grep -E "^\s+(async\s+)?[a-zA-Z]+.*\(" | head -20Repository: Comfy-Org/desktop
Length of output: 1222
Add explicit return type to this public method.
The updateTorchUpdatePolicy() method is missing an explicit return type annotation. All other public methods in the VirtualEnvironment class have explicit return types; this method should as well to comply with the coding guidelines.
🔧 Proposed fix
updateTorchUpdatePolicy(
policy: TorchUpdatePolicy | undefined,
pinnedPackages?: TorchPinnedPackages,
decisionVersion?: string
- ) {
+ ): void {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| updateTorchUpdatePolicy( | |
| policy: TorchUpdatePolicy | undefined, | |
| pinnedPackages?: TorchPinnedPackages, | |
| decisionVersion?: string | |
| ) { | |
| this.torchUpdatePolicy = policy; | |
| if (pinnedPackages !== undefined) { | |
| this.torchPinnedPackages = pinnedPackages; | |
| } else if (policy !== 'pinned') { | |
| this.torchPinnedPackages = undefined; | |
| } | |
| if (decisionVersion !== undefined) { | |
| this.torchUpdateDecisionVersion = decisionVersion; | |
| } | |
| } | |
| updateTorchUpdatePolicy( | |
| policy: TorchUpdatePolicy | undefined, | |
| pinnedPackages?: TorchPinnedPackages, | |
| decisionVersion?: string | |
| ): void { | |
| this.torchUpdatePolicy = policy; | |
| if (pinnedPackages !== undefined) { | |
| this.torchPinnedPackages = pinnedPackages; | |
| } else if (policy !== 'pinned') { | |
| this.torchPinnedPackages = undefined; | |
| } | |
| if (decisionVersion !== undefined) { | |
| this.torchUpdateDecisionVersion = decisionVersion; | |
| } | |
| } |
🤖 Prompt for AI Agents
In `@src/virtualEnvironment.ts` around lines 275 - 289, The public method
updateTorchUpdatePolicy in class VirtualEnvironment lacks an explicit return
type; add the explicit return type `: void` to the method signature of
updateTorchUpdatePolicy(policy: TorchUpdatePolicy | undefined, pinnedPackages?:
TorchPinnedPackages, decisionVersion?: string) to match other public methods and
coding guidelines so the signature clearly indicates it returns nothing.
src/virtualEnvironment.ts
Outdated
Comment on lines
291
to
301
| isUsingCustomTorchMirror(): boolean { | ||
| if (!this.torchMirror) return false; | ||
| return this.torchMirror !== getDeviceDefaultTorchMirror(this.selectedDevice); | ||
| } | ||
|
|
||
| private shouldSkipNvidiaTorchUpgrade(): boolean { | ||
| if (this.torchUpdatePolicy === 'pinned') return true; | ||
| if (this.torchUpdatePolicy === 'defer' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION) | ||
| return true; | ||
| return false; | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for call sites of isUsingCustomTorchMirror and torchUpdatePolicy
echo "=== Looking for isUsingCustomTorchMirror usage ==="
rg -nP --type=ts -C2 'isUsingCustomTorchMirror\s*\('
echo ""
echo "=== Looking for shouldSkipNvidiaTorchUpgrade usage ==="
rg -nP --type=ts -C2 'shouldSkipNvidiaTorchUpgrade\s*\('
echo ""
echo "=== Looking for NVIDIA torch upgrade decision points ==="
rg -nP --type=ts -C3 'torchUpdatePolicy|shouldSkipNvidia' | head -100Repository: Comfy-Org/desktop
Length of output: 9079
🏁 Script executed:
#!/bin/bash
# Get full context of ensureRecommendedNvidiaTorch method
echo "=== ensureRecommendedNvidiaTorch implementation ==="
sed -n '694,750p' src/virtualEnvironment.ts
echo ""
echo "=== needsNvidiaTorchUpgrade implementation ==="
sed -n '1041,1100p' src/virtualEnvironment.ts
echo ""
echo "=== Check if custom mirror is checked anywhere in upgrade flow ==="
rg -nP --type=ts 'isUsingCustomTorchMirror' src/virtualEnvironment.tsRepository: Comfy-Org/desktop
Length of output: 4652
🏁 Script executed:
#!/bin/bash
# Check if tests verify custom mirror blocking NVIDIA upgrades
echo "=== Test coverage for custom mirror and NVIDIA upgrades ==="
sed -n '300,330p' tests/unit/virtualEnvironment.test.ts
echo ""
echo "=== Search for test cases related to shouldSkipNvidiaTorchUpgrade ==="
rg -nP --type=ts 'shouldSkipNvidiaTorchUpgrade|ensureRecommendedNvidiaTorch' tests/
echo ""
echo "=== Check if custom mirror is tested with upgrade logic ==="
rg -nP --type=ts -B5 -A5 'isUsingCustomTorchMirror.*true' tests/ | head -50Repository: Comfy-Org/desktop
Length of output: 5535
🏁 Script executed:
#!/bin/bash
# Check if custom mirror skip is mentioned in any documentation or PR-related files
echo "=== Search for PR objectives or documentation about custom mirror ==="
find . -type f \( -name "*.md" -o -name "*.txt" \) | head -20 | xargs grep -l "custom.*mirror\|torch.*mirror" 2>/dev/null
echo ""
echo "=== Look for any references to the PR or CHANGELOG ==="
git log --oneline --all -S "isUsingCustomTorchMirror" | head -10
echo ""
echo "=== Check if shouldSkipNvidiaTorchUpgrade should check custom mirror ==="
echo "Looking for all places shouldSkipNvidiaTorchUpgrade is called:"
rg -nP --type=ts 'shouldSkipNvidiaTorchUpgrade' src/Repository: Comfy-Org/desktop
Length of output: 595
🏁 Script executed:
#!/bin/bash
# Check the full implementation and see if custom mirror should prevent upgrades
echo "=== Verify the relationship: custom mirror usage context ==="
sed -n '708,720p' src/virtualEnvironment.ts
echo ""
echo "=== Check if installationManager has any custom mirror checks ==="
rg -nP --type=ts 'isUsingCustomTorchMirror|torchMirror' src/install/installationManager.ts | head -20
echo ""
echo "=== Summary: All places where NVIDIA torch updates are gated ==="
rg -nP --type=ts -B2 'ensureRecommendedNvidiaTorch|needsNvidiaTorchUpgrade' src/ --type-list | grep -E '\.(ts|tsx):' | head -30Repository: Comfy-Org/desktop
Length of output: 1268
🏁 Script executed:
#!/bin/bash
# Final verification: confirm custom mirror skip is NOT implemented
echo "=== Direct check: Does shouldSkipNvidiaTorchUpgrade call isUsingCustomTorchMirror? ==="
sed -n '296,301p' src/virtualEnvironment.ts
echo ""
echo "=== Do ensureRecommendedNvidiaTorch or needsNvidiaTorchUpgrade check custom mirror? ==="
grep -n "isUsingCustomTorchMirror" src/virtualEnvironment.ts
echo ""
echo "=== Log message context at upgrade skip point ==="
sed -n '694,700p' src/virtualEnvironment.tsRepository: Comfy-Org/desktop
Length of output: 893
Custom mirror skip is not enforced in NVIDIA upgrade paths.
The isUsingCustomTorchMirror() helper exists but is unused. The shouldSkipNvidiaTorchUpgrade() method only checks policy ('pinned' or 'defer') and does not consult the custom mirror check. Both ensureRecommendedNvidiaTorch() and needsNvidiaTorchUpgrade() rely solely on this incomplete gate, meaning custom mirror configurations do not prevent NVIDIA PyTorch upgrades. Add isUsingCustomTorchMirror() to the skip logic in shouldSkipNvidiaTorchUpgrade().
🤖 Prompt for AI Agents
In `@src/virtualEnvironment.ts` around lines 291 - 301, The skip logic in
shouldSkipNvidiaTorchUpgrade() must also honor custom mirrors: update
shouldSkipNvidiaTorchUpgrade() to return true if isUsingCustomTorchMirror() is
true in addition to the existing checks (torchUpdatePolicy === 'pinned' or the
'defer' + torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION case)
so that ensureRecommendedNvidiaTorch() and needsNvidiaTorchUpgrade() will not
trigger NVIDIA PyTorch upgrades when a custom torchMirror (checked via
isUsingCustomTorchMirror(), which compares this.torchMirror against
getDeviceDefaultTorchMirror(this.selectedDevice)) is configured.
…IN_VERSION to constants.ts
dosubot
bot
added
size:XL
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/unit/install/installationManager.test.ts`:
- Around line 159-187: Replace the four Vitest cases using the `it(...)` alias
with `test(...)` to follow the repo's unit-test guideline; locate the blocks
that call helper.getRecommendedTorchMirror (the four test cases asserting
CUDA/default normalization, nightly CUDA normalization, and unnormalizable
mirrors) and change each `it('...')` to `test('...')` while keeping the test
bodies and expectations identical.
♻️ Duplicate comments (4)
src/constants.ts (1)
184-190: Add JSDoc for the new torch policy exports.These are exported public types and should carry concise JSDoc to align with the repo’s documentation expectations. Based on learnings, please add brief JSDoc comments.
✍️ Suggested doc additions
+/** How NVIDIA PyTorch updates should be handled. */ export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned'; +/** Optional pinned versions for torch/torchaudio/torchvision when updates are disabled. */ export type TorchPinnedPackages = { torch?: string; torchaudio?: string; torchvision?: string; };src/virtualEnvironment.ts (2)
275-289: Add an explicit return type forupdateTorchUpdatePolicy.Public methods should declare return types; this one currently omits
: void. Based on learnings, please add it.🔧 Proposed fix
updateTorchUpdatePolicy( policy: TorchUpdatePolicy | undefined, pinnedPackages?: TorchPinnedPackages, decisionVersion?: string - ) { + ): void {
296-301: Honor custom torch mirrors in NVIDIA upgrade skip logic.
shouldSkipNvidiaTorchUpgrade()only checks policy/decision; it should also skip when a custom mirror is configured so automated upgrades are suppressed as intended.✅ Suggested guard
private shouldSkipNvidiaTorchUpgrade(): boolean { + if (!this.isUsingRecommendedTorchMirror()) return true; if (this.torchUpdatePolicy === 'pinned' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION) return true; if (this.torchUpdatePolicy === 'defer' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION) return true; return false; }src/install/installationManager.ts (1)
404-409: Skip NVIDIA update prompt when a custom torch mirror is configured.The new flow doesn’t short-circuit for custom mirrors, so users still see prompts/updates despite the PR requirement. Add an early guard right after the device check.
✅ Suggested guard
private async maybeUpdateNvidiaTorch(installation: ComfyInstallation, callbacks: ProcessCallbacks): Promise<void> { const virtualEnvironment = installation.virtualEnvironment; if (virtualEnvironment.selectedDevice !== 'nvidia') return; + if (!virtualEnvironment.isUsingRecommendedTorchMirror()) { + log.info('Skipping NVIDIA PyTorch update because a custom torch mirror is configured.'); + return; + } const config = useDesktopConfig();
Comment on lines
+159
to
+187
| describe('getRecommendedTorchMirror', () => { | ||
| it('returns the CUDA mirror when the mirror is empty or default', () => { | ||
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | ||
|
|
||
| expect(helper.getRecommendedTorchMirror()).toBe(TorchMirrorUrl.Cuda); | ||
| expect(helper.getRecommendedTorchMirror('')).toBe(TorchMirrorUrl.Cuda); | ||
| expect(helper.getRecommendedTorchMirror(TorchMirrorUrl.Default)).toBe(TorchMirrorUrl.Cuda); | ||
| }); | ||
|
|
||
| it('normalizes CUDA mirrors to the recommended CUDA path', () => { | ||
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | ||
|
|
||
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/cu128')).toBe(TorchMirrorUrl.Cuda); | ||
| }); | ||
|
|
||
| it('normalizes nightly CUDA mirrors to the recommended CUDA path', () => { | ||
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | ||
|
|
||
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/nightly/cu118')).toBe( | ||
| TorchMirrorUrl.NightlyCuda | ||
| ); | ||
| }); | ||
|
|
||
| it('returns the original mirror when it cannot be normalized', () => { | ||
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | ||
|
|
||
| expect(helper.getRecommendedTorchMirror('https://example.com/simple')).toBe('https://example.com/simple'); | ||
| expect(helper.getRecommendedTorchMirror('not a url')).toBe('not a url'); | ||
| }); |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n tests/unit/install/installationManager.test.ts | sed -n '159,187p'Repository: Comfy-Org/desktop
Length of output: 178
Use test instead of it in the new Vitest cases.
The repo's unit-test guidelines prefer test for Vitest tests. Please switch the four new test cases accordingly.
♻️ Suggested change
- describe('getRecommendedTorchMirror', () => {
- it('returns the CUDA mirror when the mirror is empty or default', () => {
+ describe('getRecommendedTorchMirror', () => {
+ test('returns the CUDA mirror when the mirror is empty or default', () => {
const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined };
expect(helper.getRecommendedTorchMirror()).toBe(TorchMirrorUrl.Cuda);
expect(helper.getRecommendedTorchMirror('')).toBe(TorchMirrorUrl.Cuda);
expect(helper.getRecommendedTorchMirror(TorchMirrorUrl.Default)).toBe(TorchMirrorUrl.Cuda);
});
- it('normalizes CUDA mirrors to the recommended CUDA path', () => {
+ test('normalizes CUDA mirrors to the recommended CUDA path', () => {
const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined };
expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/cu128')).toBe(TorchMirrorUrl.Cuda);
});
- it('normalizes nightly CUDA mirrors to the recommended CUDA path', () => {
+ test('normalizes nightly CUDA mirrors to the recommended CUDA path', () => {
const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined };
expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/nightly/cu118')).toBe(
TorchMirrorUrl.NightlyCuda
);
});
- it('returns the original mirror when it cannot be normalized', () => {
+ test('returns the original mirror when it cannot be normalized', () => {
const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined };
expect(helper.getRecommendedTorchMirror('https://example.com/simple')).toBe('https://example.com/simple');
expect(helper.getRecommendedTorchMirror('not a url')).toBe('not a url');
});
});📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| describe('getRecommendedTorchMirror', () => { | |
| it('returns the CUDA mirror when the mirror is empty or default', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror()).toBe(TorchMirrorUrl.Cuda); | |
| expect(helper.getRecommendedTorchMirror('')).toBe(TorchMirrorUrl.Cuda); | |
| expect(helper.getRecommendedTorchMirror(TorchMirrorUrl.Default)).toBe(TorchMirrorUrl.Cuda); | |
| }); | |
| it('normalizes CUDA mirrors to the recommended CUDA path', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/cu128')).toBe(TorchMirrorUrl.Cuda); | |
| }); | |
| it('normalizes nightly CUDA mirrors to the recommended CUDA path', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/nightly/cu118')).toBe( | |
| TorchMirrorUrl.NightlyCuda | |
| ); | |
| }); | |
| it('returns the original mirror when it cannot be normalized', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://example.com/simple')).toBe('https://example.com/simple'); | |
| expect(helper.getRecommendedTorchMirror('not a url')).toBe('not a url'); | |
| }); | |
| describe('getRecommendedTorchMirror', () => { | |
| test('returns the CUDA mirror when the mirror is empty or default', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror()).toBe(TorchMirrorUrl.Cuda); | |
| expect(helper.getRecommendedTorchMirror('')).toBe(TorchMirrorUrl.Cuda); | |
| expect(helper.getRecommendedTorchMirror(TorchMirrorUrl.Default)).toBe(TorchMirrorUrl.Cuda); | |
| }); | |
| test('normalizes CUDA mirrors to the recommended CUDA path', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/cu128')).toBe(TorchMirrorUrl.Cuda); | |
| }); | |
| test('normalizes nightly CUDA mirrors to the recommended CUDA path', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://download.pytorch.org/whl/nightly/cu118')).toBe( | |
| TorchMirrorUrl.NightlyCuda | |
| ); | |
| }); | |
| test('returns the original mirror when it cannot be normalized', () => { | |
| const helper = manager as unknown as { getRecommendedTorchMirror: (mirror?: string) => string | undefined }; | |
| expect(helper.getRecommendedTorchMirror('https://example.com/simple')).toBe('https://example.com/simple'); | |
| expect(helper.getRecommendedTorchMirror('not a url')).toBe('not a url'); | |
| }); | |
| }); |
🤖 Prompt for AI Agents
In `@tests/unit/install/installationManager.test.ts` around lines 159 - 187,
Replace the four Vitest cases using the `it(...)` alias with `test(...)` to
follow the repo's unit-test guideline; locate the blocks that call
helper.getRecommendedTorchMirror (the four test cases asserting CUDA/default
normalization, nightly CUDA normalization, and unnormalizable mirrors) and
change each `it('...')` to `test('...')` while keeping the test bodies and
expectations identical.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@src/virtualEnvironment.ts`:
- Around line 143-145: The three public mutable properties torchUpdatePolicy,
torchPinnedPackages, and torchUpdateDecisionVersion should not be directly
writable; make them private (e.g., _torchUpdatePolicy, _torchPinnedPackages,
_torchUpdateDecisionVersion) and expose read-only accessors (getters) while
ensuring all mutations go through the existing updateTorchUpdatePolicy method;
update references inside the class to use the private fields and replace
external direct assignments with the setter method to preserve encapsulation.
In `@tests/unit/virtualEnvironment.test.ts`:
- Around line 322-354: Add a test to cover the undefined torchMirror case:
instantiate VirtualEnvironment with selectedDevice 'nvidia' and no torchMirror
provided, then assert isUsingRecommendedTorchMirror() returns true; update tests
in tests/unit/virtualEnvironment.test.ts near the existing
isUsingRecommendedTorchMirror suite so the new test follows the other cases and
uses vi.stubGlobal('process', { ...process, resourcesPath: '/test/resources' })
and mockTelemetry as in the other tests.
- Around line 397-464: The tests and implementation are missing the rule that
automated NVIDIA torch upgrades must be skipped when a custom torch mirror is
configured; update shouldSkipNvidiaTorchUpgrade (used by
ensureRecommendedNvidiaTorch) to return true if the VirtualEnvironment options
include a custom mirror (e.g., a torchMirror / torch_mirror option on the
instance or constructor opts) and then add a unit test in
tests/unit/virtualEnvironment.test.ts that constructs a VirtualEnvironment with
selectedDevice: 'nvidia', torchUpdatePolicy set (pinned/defer),
torchUpdateDecisionVersion set to the recommended version, and a custom
torchMirror value, spies on getInstalledTorchPackageVersions, calls
ensureRecommendedNvidiaTorch, and expects the spy not to have been called to
verify skipping behavior.
♻️ Duplicate comments (4)
src/constants.ts (1)
184-190: Add brief JSDoc for new torch policy types.These are exported public types; a short description helps consumers and aligns with the repo's documentation expectations. As per coding guidelines, exported APIs should have concise JSDoc explaining "why" not "what".
✍️ Suggested doc additions
+/** How NVIDIA PyTorch updates should be handled. */ export type TorchUpdatePolicy = 'auto' | 'defer' | 'pinned'; +/** Pinned NVIDIA torch package versions when updates are disabled. */ export type TorchPinnedPackages = { torch?: string; torchaudio?: string; torchvision?: string; };tests/unit/virtualEnvironment.test.ts (1)
521-558: Consider extracting the long JSON string to a constant.The inline JSON string at line 523 is quite long. Extracting it to a named constant above the test would improve readability without affecting test behavior. This is a minor readability concern.
♻️ Suggested extraction
+ const UV_LIST_PACKAGES_JSON_OUTPUT = `[{"name":"aiohappyeyeballs","version":"2.6.1"},{"name":"aiohttp","version":"3.13.3"},...,{"name":"yarl","version":"1.22.0"}]`; + describe('getInstalledTorchPackageVersions', () => { test('returns parsed torch versions from uv json output', async ({ virtualEnv }) => { - const output = `[{"name":"aiohappyeyeballs",...}]`; + const output = UV_LIST_PACKAGES_JSON_OUTPUT;src/virtualEnvironment.ts (2)
287-301: Add explicit return type to this public method.The
updateTorchUpdatePolicy()method is missing an explicit return type annotation. Per coding guidelines, exported functions and public APIs must have explicit parameter and return types.🔧 Proposed fix
updateTorchUpdatePolicy( policy: TorchUpdatePolicy | undefined, pinnedPackages?: TorchPinnedPackages, decisionVersion?: string - ) { + ): void {
303-314: Custom mirror skip is not enforced in NVIDIA upgrade paths.Per the PR objectives: "skip automated updates when a version is pinned or a custom torch mirror is configured". The
shouldSkipNvidiaTorchUpgrade()method only checks policy ('pinned'or'defer') and does not consult the custom mirror. BothensureRecommendedNvidiaTorch()andneedsNvidiaTorchUpgrade()rely on this incomplete gate, meaning custom mirror configurations will not prevent NVIDIA PyTorch upgrades.Integrate
isUsingRecommendedTorchMirror()into the skip logic:🐛 Proposed fix
private shouldSkipNvidiaTorchUpgrade(): boolean { + if (!this.isUsingRecommendedTorchMirror()) { + log.info('Skipping NVIDIA PyTorch upgrade: custom torch mirror configured.'); + return true; + } if (this.torchUpdatePolicy === 'pinned' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION) return true; if (this.torchUpdatePolicy === 'defer' && this.torchUpdateDecisionVersion === NVIDIA_TORCH_RECOMMENDED_VERSION) return true; return false; }
Comment on lines
+143
to
+145
| torchUpdatePolicy?: TorchUpdatePolicy; | ||
| torchPinnedPackages?: TorchPinnedPackages; | ||
| torchUpdateDecisionVersion?: string; |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Public mutable properties warrant careful consideration.
These properties are public and mutable, meaning external code can modify them directly. While this aligns with the current usage via updateTorchUpdatePolicy, consider whether these should be readonly with state changes only through the setter method for better encapsulation.
🤖 Prompt for AI Agents
In `@src/virtualEnvironment.ts` around lines 143 - 145, The three public mutable
properties torchUpdatePolicy, torchPinnedPackages, and
torchUpdateDecisionVersion should not be directly writable; make them private
(e.g., _torchUpdatePolicy, _torchPinnedPackages, _torchUpdateDecisionVersion)
and expose read-only accessors (getters) while ensuring all mutations go through
the existing updateTorchUpdatePolicy method; update references inside the class
to use the private fields and replace external direct assignments with the
setter method to preserve encapsulation.
Comment on lines
+322
to
+354
| describe('isUsingRecommendedTorchMirror', () => { | ||
| test('returns true when using default mirror for NVIDIA', () => { | ||
| vi.stubGlobal('process', { | ||
| ...process, | ||
| resourcesPath: '/test/resources', | ||
| }); | ||
|
|
||
| const env = new VirtualEnvironment('/mock/venv', { | ||
| telemetry: mockTelemetry, | ||
| selectedDevice: 'nvidia', | ||
| pythonVersion: '3.12', | ||
| torchMirror: TorchMirrorUrl.Cuda, | ||
| }); | ||
|
|
||
| expect(env.isUsingRecommendedTorchMirror()).toBe(true); | ||
| }); | ||
|
|
||
| test('returns false when using a custom mirror', () => { | ||
| vi.stubGlobal('process', { | ||
| ...process, | ||
| resourcesPath: '/test/resources', | ||
| }); | ||
|
|
||
| const env = new VirtualEnvironment('/mock/venv', { | ||
| telemetry: mockTelemetry, | ||
| selectedDevice: 'nvidia', | ||
| pythonVersion: '3.12', | ||
| torchMirror: 'https://download.pytorch.org/whl/cu128', | ||
| }); | ||
|
|
||
| expect(env.isUsingRecommendedTorchMirror()).toBe(false); | ||
| }); | ||
| }); |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Good test coverage for isUsingRecommendedTorchMirror.
The tests correctly verify behavior for both the default NVIDIA mirror and a custom mirror scenario. Consider adding a test case for when torchMirror is undefined to confirm it returns true.
🧪 Optional: Add test for undefined torchMirror
test('returns true when torchMirror is undefined', () => {
vi.stubGlobal('process', {
...process,
resourcesPath: '/test/resources',
});
const env = new VirtualEnvironment('/mock/venv', {
telemetry: mockTelemetry,
selectedDevice: 'nvidia',
pythonVersion: '3.12',
// torchMirror is undefined
});
expect(env.isUsingRecommendedTorchMirror()).toBe(true);
});🤖 Prompt for AI Agents
In `@tests/unit/virtualEnvironment.test.ts` around lines 322 - 354, Add a test to
cover the undefined torchMirror case: instantiate VirtualEnvironment with
selectedDevice 'nvidia' and no torchMirror provided, then assert
isUsingRecommendedTorchMirror() returns true; update tests in
tests/unit/virtualEnvironment.test.ts near the existing
isUsingRecommendedTorchMirror suite so the new test follows the other cases and
uses vi.stubGlobal('process', { ...process, resourcesPath: '/test/resources' })
and mockTelemetry as in the other tests.
Comment on lines
+397
to
+464
| describe('ensureRecommendedNvidiaTorch', () => { | ||
| test('skips upgrade when updates are pinned for the current recommended version', async () => { | ||
| vi.stubGlobal('process', { | ||
| ...process, | ||
| resourcesPath: '/test/resources', | ||
| }); | ||
|
|
||
| const env = new VirtualEnvironment('/mock/venv', { | ||
| telemetry: mockTelemetry, | ||
| selectedDevice: 'nvidia', | ||
| pythonVersion: '3.12', | ||
| torchUpdatePolicy: 'pinned', | ||
| torchUpdateDecisionVersion: NVIDIA_TORCH_RECOMMENDED_VERSION, | ||
| }); | ||
|
|
||
| const versionsSpy = vi.spyOn(env, 'getInstalledTorchPackageVersions'); | ||
|
|
||
| await env.ensureRecommendedNvidiaTorch(); | ||
|
|
||
| expect(versionsSpy).not.toHaveBeenCalled(); | ||
| }); | ||
|
|
||
| test('does not skip upgrade when pinned decision version differs', async () => { | ||
| vi.stubGlobal('process', { | ||
| ...process, | ||
| resourcesPath: '/test/resources', | ||
| }); | ||
|
|
||
| const env = new VirtualEnvironment('/mock/venv', { | ||
| telemetry: mockTelemetry, | ||
| selectedDevice: 'nvidia', | ||
| pythonVersion: '3.12', | ||
| torchUpdatePolicy: 'pinned', | ||
| torchUpdateDecisionVersion: '2.8.0+cu130|0.23.0+cu130', | ||
| }); | ||
|
|
||
| const versionsSpy = vi.spyOn(env, 'getInstalledTorchPackageVersions').mockResolvedValue({ | ||
| torch: NVIDIA_TORCH_VERSION, | ||
| torchaudio: NVIDIA_TORCH_VERSION, | ||
| torchvision: NVIDIA_TORCHVISION_VERSION, | ||
| }); | ||
|
|
||
| await env.ensureRecommendedNvidiaTorch(); | ||
|
|
||
| expect(versionsSpy).toHaveBeenCalled(); | ||
| }); | ||
|
|
||
| test('skips upgrade when updates are deferred for the recommended version', async () => { | ||
| vi.stubGlobal('process', { | ||
| ...process, | ||
| resourcesPath: '/test/resources', | ||
| }); | ||
|
|
||
| const env = new VirtualEnvironment('/mock/venv', { | ||
| telemetry: mockTelemetry, | ||
| selectedDevice: 'nvidia', | ||
| pythonVersion: '3.12', | ||
| torchUpdatePolicy: 'defer', | ||
| torchUpdateDecisionVersion: NVIDIA_TORCH_RECOMMENDED_VERSION, | ||
| }); | ||
|
|
||
| const versionsSpy = vi.spyOn(env, 'getInstalledTorchPackageVersions'); | ||
|
|
||
| await env.ensureRecommendedNvidiaTorch(); | ||
|
|
||
| expect(versionsSpy).not.toHaveBeenCalled(); | ||
| }); | ||
| }); |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add test for custom mirror skip behavior.
Per the PR objectives, automated updates should be skipped when a custom torch mirror is configured. However, this behavior is not currently tested (nor implemented in shouldSkipNvidiaTorchUpgrade). Once the implementation is fixed, add a test verifying that ensureRecommendedNvidiaTorch skips upgrades when using a custom mirror.
🤖 Prompt for AI Agents
In `@tests/unit/virtualEnvironment.test.ts` around lines 397 - 464, The tests and
implementation are missing the rule that automated NVIDIA torch upgrades must be
skipped when a custom torch mirror is configured; update
shouldSkipNvidiaTorchUpgrade (used by ensureRecommendedNvidiaTorch) to return
true if the VirtualEnvironment options include a custom mirror (e.g., a
torchMirror / torch_mirror option on the instance or constructor opts) and then
add a unit test in tests/unit/virtualEnvironment.test.ts that constructs a
VirtualEnvironment with selectedDevice: 'nvidia', torchUpdatePolicy set
(pinned/defer), torchUpdateDecisionVersion set to the recommended version, and a
custom torchMirror value, spies on getInstalledTorchPackageVersions, calls
ensureRecommendedNvidiaTorch, and expects the spy not to have been called to
verify skipping behavior.
benceruleanlu
changed the title
Contributor
christian-byrne
left a comment
christian-byrne
left a comment
There was a problem hiding this comment.
I recommend path of action to split changes:
- 1 PR to make sure we are not forcing users onto new pytorch
- 2nd PR to add UI/UX to guide them when we suggest a new version and let them decide -- providing proper info and pros/cons/risk.
Member
Author
|
Taking draft to implement review comment. |
christian-byrne
added a commit
that referenced
this pull request
Jan 18, 2026
Disables automatic PyTorch version upgrades that were forcing all NVIDIA users to download 1.5GB packages without consent. ## What Changed - Commented out `ensureRecommendedNvidiaTorch()` call in updatePackages() - Users will no longer be forced to upgrade PyTorch on app startup ## Why The automatic upgrade introduced in PR #1513 caused several issues: - Forced 1.5GB downloads without user warning or consent - Caused installation failures for users with slow/metered connections - Blocked Chinese users when packages weren't available on mirrors - No opt-out mechanism ## User Impact - NVIDIA users will keep their current PyTorch versions - No more forced package reinstalls on app startup - Users can manually upgrade PyTorch if they want new features ## Next Steps A follow-up PR will add optional user-prompted PyTorch upgrades with proper consent dialogs and preference tracking. Related: PR #1513, PR #1525 Fixes: Forced PyTorch upgrade crisis (Jan 8-14, 2026) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
┆Issue is synchronized with this Notion page by Unito