Add PyPI mirror fallbacks (#1526)

This extends pip install fallback behavior when a PyPI mirror is
configured by adding multiple `--extra-index-url` values, ending with
standard PyPI, so packages missing on TUNA (e.g. comfy-kitchen) still
resolve.

Fallback order:
- https://mirrors.aliyun.com/pypi/simple/
- https://mirrors.cloud.tencent.com/pypi/simple/
- https://pypi.org/simple/

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-1526-Add-TorchMirrorUrl-Default-as-extraIndexUrl-2e76d73d365081f9b0e4fabd583bea75)
by [Unito](https://www.unito.io)

Author: benceruleanlu

src/constants.ts

@@ -191,6 +191,12 @@ export const NIGHTLY_CPU_TORCH_URL = TorchMirrorUrl.NightlyCpu;
 /** @deprecated Use {@link TorchMirrorUrl} instead. */
 export const DEFAULT_PYPI_INDEX_URL = TorchMirrorUrl.Default;
 
+export const PYPI_FALLBACK_INDEX_URLS: string[] = [
+  'https://mirrors.aliyun.com/pypi/simple/',
+  'https://mirrors.cloud.tencent.com/pypi/simple/',
+  TorchMirrorUrl.Default,
+];
+
 export const AMD_ROCM_SDK_PACKAGES: string[] = [
   'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl',
   'https://repo.radeon.com/rocm/windows/rocm-rel-7.1.1/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl',

src/virtualEnvironment.ts

@@ -15,6 +15,7 @@ import {
   NVIDIA_TORCHVISION_VERSION,
   NVIDIA_TORCH_PACKAGES,
   NVIDIA_TORCH_VERSION,
+  PYPI_FALLBACK_INDEX_URLS,
   TorchMirrorUrl,
 } from './constants';
 import { PythonImportVerificationError } from './infrastructure/pythonImportVerificationError';
@@ -45,7 +46,7 @@ export interface PythonExecutor {
 interface PipInstallConfig {
   packages: string[];
   indexUrl?: string;
-  extraIndexUrl?: string;
+  extraIndexUrls?: string[];
   prerelease?: boolean;
   upgradePackages?: boolean;
   requirementsFile?: string;
@@ -57,7 +58,7 @@ type TorchPackageVersions = Record<TorchPackageName, string | undefined>;
 
 const TORCH_PACKAGE_NAMES: TorchPackageName[] = ['torch', 'torchaudio', 'torchvision'];
 
-function getPipInstallArgs(config: PipInstallConfig): string[] {
+export function getPipInstallArgs(config: PipInstallConfig): string[] {
   const installArgs = ['pip', 'install'];
 
   if (config.upgradePackages) {
@@ -78,8 +79,10 @@ function getPipInstallArgs(config: PipInstallConfig): string[] {
     installArgs.push('--index-url', config.indexUrl);
   }
 
-  if (config.extraIndexUrl) {
-    installArgs.push('--extra-index-url', config.extraIndexUrl);
+  if (config.extraIndexUrls) {
+    for (const extraIndexUrl of config.extraIndexUrls) {
+      installArgs.push('--extra-index-url', extraIndexUrl);
+    }
   }
 
   if (config.indexStrategy) {
@@ -153,6 +156,15 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
     };
   }
 
+  /**
+   * Returns extra index URLs to use for pip installs.
+   * @returns The fallback index URLs, or `undefined` if none are configured.
+   */
+  private getPypiFallbackIndexUrls(): string[] | undefined {
+    const fallbackUrls = PYPI_FALLBACK_INDEX_URLS.filter((url) => url !== this.pypiMirror);
+    return fallbackUrls.length > 0 ? fallbackUrls : undefined;
+  }
+
   /** @todo Refactor to `using` */
   get uvPtyInstance() {
     const env = {
@@ -401,6 +413,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       indexStrategy: 'unsafe-best-match',
       packages: [],
       indexUrl: this.pypiMirror,
+      extraIndexUrls: this.getPypiFallbackIndexUrls(),
     });
     const { exitCode } = await this.runUvCommandAsync(installCmd, callbacks);
     if (exitCode !== 0) {
@@ -815,6 +828,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       requirementsFile: this.comfyUIRequirementsPath,
       packages: [],
       indexUrl: this.pypiMirror,
+      extraIndexUrls: this.getPypiFallbackIndexUrls(),
     });
     const { exitCode } = await this.runUvCommandAsync(installCmd, callbacks);
     if (exitCode !== 0) {
@@ -846,6 +860,7 @@ export class VirtualEnvironment implements HasTelemetry, PythonExecutor {
       requirementsFile: this.comfyUIManagerRequirementsPath,
       packages: [],
       indexUrl: this.pypiMirror,
+      extraIndexUrls: this.getPypiFallbackIndexUrls(),
     });
     const { exitCode } = await this.runUvCommandAsync(installCmd, callbacks);
     if (exitCode !== 0) {

tests/unit/virtualEnvironment.test.ts

@@ -3,8 +3,9 @@ import { type ChildProcess, spawn } from 'node:child_process';
 import path from 'node:path';
 import { test as baseTest, describe, expect, vi } from 'vitest';
 
+import { TorchMirrorUrl } from '@/constants';
 import type { ITelemetry } from '@/services/telemetry';
-import { VirtualEnvironment } from '@/virtualEnvironment';
+import { VirtualEnvironment, getPipInstallArgs } from '@/virtualEnvironment';
 
 vi.mock('@sentry/electron/main', () => ({
   init: vi.fn(),
@@ -151,6 +152,33 @@ test.for(allCombinations)('hasRequirements', async ({ core, manager }, { virtual
 });
 
 describe('VirtualEnvironment', () => {
+  describe('getPipInstallArgs', () => {
+    test('includes unsafe-best-match and extra index URL args', () => {
+      const args = getPipInstallArgs({
+        requirementsFile: '/tmp/requirements.txt',
+        packages: [],
+        indexUrl: 'https://mirror.example/simple/',
+        extraIndexUrls: ['https://mirror-two.example/simple/', TorchMirrorUrl.Default],
+        indexStrategy: 'unsafe-best-match',
+      });
+
+      expect(args).toEqual([
+        'pip',
+        'install',
+        '-r',
+        '/tmp/requirements.txt',
+        '--index-url',
+        'https://mirror.example/simple/',
+        '--extra-index-url',
+        'https://mirror-two.example/simple/',
+        '--extra-index-url',
+        TorchMirrorUrl.Default,
+        '--index-strategy',
+        'unsafe-best-match',
+      ]);
+    });
+  });
+
   describe('hasRequirements', () => {
     test('returns OK when all packages are installed', async ({ virtualEnv }) => {
       mockSpawnOutputOnce('Would make no changes\n');