@modular-prompt/driver 0.15.0 → 0.17.0

package/scripts/runtime-cli.js

@@ -4,13 +4,14 @@
  * Python runtime 管理 CLI
  *
  * setup mlx     — ~/.modular-prompt/runtimes/mlx に venv を作成
- * setup pytorch — ~/.modular-prompt/runtimes/pytorch に cpu-minimal venv を作成
+ * setup pytorch — ~/.modular-prompt/runtimes/pytorch に選択した variant の venv を作成
+ * sync pytorch  — PyTorch runtime のコードと依存を更新
  * setup --status
  * cleanup mlx [--yes]
  * cleanup --all [--yes]
  */
 
-import { execSync } from 'child_process';
+import { execFileSync, execSync } from 'child_process';
 import { existsSync, readFileSync, rmSync, mkdirSync } from 'fs';
 import { createInterface } from 'readline';
 import { dirname, join } from 'path';
@@ -32,16 +33,29 @@ const {
   getRuntimeDir,
   getVenvPath,
   getMlxPythonDir,
-  getPytorchPythonDir,
+  getPytorchRuntimePythonDir,
+  getPytorchTemplateDir,
+  PYTORCH_DEFAULT_VARIANT,
   isRuntimeReady,
 } = await import(runtimeModuleUrl('paths-core.mjs'));
 
+const {
+  SETUP_MLX_MONOREPO,
+  SETUP_PYTORCH_MONOREPO,
+  SYNC_PYTORCH_CLI,
+} = await import(runtimeModuleUrl('setup-commands-core.mjs'));
+
 const {
   collectInstalledPackages,
   readManifest,
   writeManifest,
 } = await import(runtimeModuleUrl('manifest-core.mjs'));
 
+const {
+  seedPytorchTemplate,
+  syncPytorchTemplate,
+} = await import(runtimeModuleUrl('pytorch-template-core.mjs'));
+
 function readPackageVersion() {
   try {
     const pkg = JSON.parse(readFileSync(join(packageRoot, 'package.json'), 'utf8'));
@@ -90,7 +104,7 @@ function setupMlx() {
   };
 
   try {
-    execSync('uv venv --python 3.13', { cwd: pythonDir, stdio: 'inherit', env });
+    execSync('uv venv --clear --python 3.13', { cwd: pythonDir, stdio: 'inherit', env });
     execSync('uv pip install -e .', { cwd: pythonDir, stdio: 'inherit', env });
 
     writeManifest('mlx', {
@@ -113,22 +127,171 @@ function setupMlx() {
 }
 
 const PYTORCH_CPU_INDEX = 'https://download.pytorch.org/whl/cpu';
+const PYTORCH_CUDA_INDEX_BASE = 'https://download.pytorch.org/whl';
+const PYTORCH_DEFAULT_CUDA_VERSION = '12.4';
+const PYTORCH_TORCH_VERSION = '2.9.1';
 const PYTORCH_PYTHON_VERSION = '3.12';
 
-function setupPytorch() {
-  const pythonDir = getPytorchPythonDir(packageRoot);
-  if (!existsSync(pythonDir)) {
-    console.error(`❌ PyTorch Python project not found: ${pythonDir}`);
+function resolveVenvPython(venvPath) {
+  return process.platform === 'win32'
+    ? join(venvPath, 'Scripts', 'python.exe')
+    : join(venvPath, 'bin', 'python');
+}
+
+function resolveCudaVersion(value = PYTORCH_DEFAULT_CUDA_VERSION) {
+  const raw = String(value).trim().toLowerCase();
+  const match =
+    raw.match(/^(?:cu)?(\d{1,2})\.(\d{1,2})$/) ??
+    raw.match(/^(?:cu)?(\d{2})(\d{1,2})$/);
+  if (!match) {
+    throw new Error(
+      `Invalid CUDA version "${value}". Use a version such as 12.4 or cu124.`,
+    );
+  }
+
+  const major = Number(match[1]);
+  const minor = Number(match[2]);
+  if (major < 1 || minor > 99) {
+    throw new Error(
+      `Invalid CUDA version "${value}". Use a version such as 12.4 or cu124.`,
+    );
+  }
+
+  return {
+    version: `${major}.${minor}`,
+    index: `${PYTORCH_CUDA_INDEX_BASE}/cu${major}${minor}`,
+  };
+}
+
+function resolvePytorchIndex(variant, cudaVersion) {
+  if (variant !== 'cuda') {
+    return { index: PYTORCH_CPU_INDEX };
+  }
+  return resolveCudaVersion(cudaVersion);
+}
+
+function hasNvidiaGpu() {
+  try {
+    const output = execFileSync(
+      'nvidia-smi',
+      ['--query-gpu=name', '--format=csv,noheader,nounits'],
+      { encoding: 'utf8', stdio: ['ignore', 'pipe', 'ignore'] },
+    );
+    return output.trim().length > 0;
+  } catch {
+    return false;
+  }
+}
+
+function checkCudaAvailability(venvPath) {
+  const venvPython = resolveVenvPython(venvPath);
+  if (!existsSync(venvPython)) {
+    return null;
+  }
+
+  try {
+    const output = execFileSync(
+      venvPython,
+      ['-c', 'import torch; print(torch.cuda.is_available())'],
+      { encoding: 'utf8', stdio: ['ignore', 'pipe', 'ignore'] },
+    );
+    const value = output.trim().split(/\s+/).at(-1)?.toLowerCase();
+    if (value === 'true') {
+      return true;
+    }
+    if (value === 'false') {
+      return false;
+    }
+  } catch {
+    // A partially installed runtime or an unavailable torch import is reported as unknown.
+  }
+  return null;
+}
+
+function installPytorchProject(
+  pythonDir,
+  venvPath,
+  env,
+  { installTorch = false, torchIndex = PYTORCH_CPU_INDEX } = {},
+) {
+  const venvPython = resolveVenvPython(venvPath);
+  if (installTorch) {
+    execFileSync(
+      'uv',
+      [
+        'pip',
+        'install',
+        '--python',
+        venvPython,
+        `torch==${PYTORCH_TORCH_VERSION}`,
+        '--index-url',
+        torchIndex,
+      ],
+      { cwd: pythonDir, stdio: 'inherit', env },
+    );
+  }
+  execFileSync(
+    'uv',
+    ['pip', 'install', '--python', venvPython, '.'],
+    { cwd: pythonDir, stdio: 'inherit', env },
+  );
+}
+
+function writePytorchManifest(previousManifest, variant, packages, cudaVersion) {
+  const manifest = {
+    ...(previousManifest ?? {}),
+    profile: 'pytorch',
+    variant,
+    driverVersion,
+    platform: previousManifest?.platform ?? process.platform,
+    pythonVersion: previousManifest?.pythonVersion ?? PYTORCH_PYTHON_VERSION,
+    createdAt: new Date().toISOString(),
+  };
+
+  if (variant === 'cuda') {
+    manifest.cudaVersion = resolveCudaVersion(
+      cudaVersion ?? previousManifest?.cudaVersion ?? PYTORCH_DEFAULT_CUDA_VERSION,
+    ).version;
+  } else {
+    delete manifest.cudaVersion;
+  }
+
+  if (packages) {
+    if (packages.torch) {
+      manifest.torchVersion = packages.torch;
+    }
+    manifest.packages = packages;
+  }
+  if (!manifest.torchVersion) {
+    manifest.torchVersion = PYTORCH_TORCH_VERSION;
+  }
+  writeManifest('pytorch', manifest);
+}
+
+function setupPytorch(variant = PYTORCH_DEFAULT_VARIANT, cudaVersion) {
+  const templateDir = getPytorchTemplateDir(packageRoot, variant);
+  if (!existsSync(templateDir)) {
+    console.error(`❌ PyTorch template not found for variant ${variant}: ${templateDir}`);
     process.exit(1);
   }
 
+  const pytorchIndex = resolvePytorchIndex(variant, cudaVersion);
+  if (variant === 'cuda' && !hasNvidiaGpu()) {
+    console.warn(
+      '⚠️  NVIDIA GPU/driver was not detected. Continuing with the CUDA runtime; ' +
+        'verify torch.cuda.is_available() before running inference.',
+    );
+  }
+
+  const pythonDir = getPytorchRuntimePythonDir();
   const venvPath = getVenvPath('pytorch');
   const runtimeDir = getRuntimeDir('pytorch');
 
-  console.log('🚀 Setting up PyTorch runtime (cpu-minimal)...\n');
+  console.log(`🚀 Setting up PyTorch runtime (${variant})...\n`);
+  console.log(`📁 Template:       ${templateDir}`);
   console.log(`📁 Python project: ${pythonDir}`);
   console.log(`📁 Runtime venv:  ${venvPath}`);
-  console.log(`📦 torch index:    ${PYTORCH_CPU_INDEX}\n`);
+  console.log(`📦 torch index:    ${pytorchIndex.index}\n`);
 
   ensureUv();
   mkdirSync(runtimeDir, { recursive: true });
@@ -139,34 +302,24 @@ function setupPytorch() {
   };
 
   try {
-    execSync(`uv venv --python ${PYTORCH_PYTHON_VERSION}`, { cwd: pythonDir, stdio: 'inherit', env });
-    const venvPython =
-      process.platform === 'win32'
-        ? join(venvPath, 'Scripts', 'python.exe')
-        : join(venvPath, 'bin', 'python');
-    execSync(`uv pip install --python "${venvPython}" "torch==2.9.1" --index-url ${PYTORCH_CPU_INDEX}`, {
-      cwd: pythonDir,
-      stdio: 'inherit',
-      env,
+    seedPytorchTemplate(templateDir, pythonDir);
+    execSync(`uv venv --clear --python ${PYTORCH_PYTHON_VERSION}`, { cwd: pythonDir, stdio: 'inherit', env });
+    installPytorchProject(pythonDir, venvPath, env, {
+      installTorch: true,
+      torchIndex: pytorchIndex.index,
     });
-    execSync(`uv pip install --python "${venvPython}" -e .`, { cwd: pythonDir, stdio: 'inherit', env });
 
     const packages = collectInstalledPackages(pythonDir, venvPath);
-    writeManifest('pytorch', {
-      profile: 'pytorch',
-      variant: 'cpu-minimal',
-      driverVersion,
-      platform: process.platform,
-      pythonVersion: PYTORCH_PYTHON_VERSION,
-      torchVersion: packages?.torch,
-      createdAt: new Date().toISOString(),
-      packages,
-    });
+    writePytorchManifest(null, variant, packages, pytorchIndex.version);
 
-    console.log('\n✅ PyTorch runtime setup completed (cpu-minimal).');
+    console.log(`\n✅ PyTorch runtime setup completed (${variant}).`);
     console.log(`   Home: ${getModularPromptHome()}`);
     console.log('   You can now use PyTorchDriver from @modular-prompt/driver');
-    console.log('   For CUDA / custom environments, see docs/LOCAL_MODEL_SETUP.md');
+    const localModelSetupDoc = join(packageRoot, 'docs', 'LOCAL_MODEL_SETUP.md');
+    const docHint = existsSync(localModelSetupDoc)
+      ? localModelSetupDoc
+      : './docs/LOCAL_MODEL_SETUP.md in @modular-prompt/driver';
+    console.log(`   For CUDA / custom environments, see ${docHint}`);
   } catch (error) {
     const message = error instanceof Error ? error.message : String(error);
     console.error('❌ Failed to setup PyTorch runtime:', message);
@@ -174,11 +327,71 @@ function setupPytorch() {
   }
 }
 
+function syncPytorch(requestedVariant) {
+  const pythonDir = getPytorchRuntimePythonDir();
+  const venvPath = getVenvPath('pytorch');
+  if (!isRuntimeReady('pytorch')) {
+    console.error(
+      `❌ PyTorch runtime is not ready at ${getRuntimeDir('pytorch')}. ` +
+        'Run: modular-prompt-runtime setup pytorch',
+    );
+    process.exit(1);
+  }
+
+  const manifest = readManifest('pytorch');
+  if (
+    requestedVariant &&
+    manifest?.variant &&
+    requestedVariant !== manifest.variant
+  ) {
+    console.error(
+      `❌ PyTorch runtime variant mismatch: runtime is ${manifest.variant}, ` +
+        `but ${requestedVariant} was requested. Sync without --variant or rerun setup pytorch.`,
+    );
+    process.exit(1);
+  }
+  const variant = manifest?.variant ?? requestedVariant ?? PYTORCH_DEFAULT_VARIANT;
+  const templateDir = getPytorchTemplateDir(packageRoot, variant);
+  if (!existsSync(templateDir)) {
+    console.error(`❌ PyTorch template not found for variant ${variant}: ${templateDir}`);
+    process.exit(1);
+  }
+
+  console.log(`🔄 Syncing PyTorch runtime (${variant})...\n`);
+  console.log(`📁 Template:       ${templateDir}`);
+  console.log(`📁 Python project: ${pythonDir}`);
+  console.log(`📁 Runtime venv:  ${venvPath}\n`);
+
+  ensureUv();
+  const env = {
+    ...process.env,
+    UV_PROJECT_ENVIRONMENT: venvPath,
+  };
+
+  try {
+    syncPytorchTemplate(templateDir, pythonDir);
+    installPytorchProject(pythonDir, venvPath, env);
+
+    const packages = collectInstalledPackages(pythonDir, venvPath);
+    writePytorchManifest(manifest, variant, packages);
+
+    console.log('\n✅ PyTorch runtime sync completed.');
+    console.log(`   Runtime driver version: ${driverVersion}`);
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    console.error('❌ Failed to sync PyTorch runtime:', message);
+    process.exit(1);
+  }
+}
+
 function formatManifestDetail(manifest) {
   const parts = [`driver ${manifest.driverVersion}`];
   if (manifest.variant) {
     parts.push(`variant ${manifest.variant}`);
   }
+  if (manifest.cudaVersion) {
+    parts.push(`CUDA ${manifest.cudaVersion}`);
+  }
   const torchVersion = manifest.torchVersion ?? manifest.packages?.torch;
   if (torchVersion) {
     parts.push(`torch ${torchVersion}`);
@@ -191,19 +404,35 @@ function printStatus() {
   console.log(`modular-prompt home: ${getModularPromptHome()}\n`);
   for (const profile of RUNTIME_PROFILES) {
     const ready = isRuntimeReady(profile);
-    const manifest = ready ? readManifest(profile) : null;
-    const detail = manifest ? formatManifestDetail(manifest) : '';
+    const manifest = readManifest(profile);
+    const detail = ready && manifest ? formatManifestDetail(manifest) : '';
     const icon = ready ? '✅' : '❌';
     const runtimePath = getRuntimeDir(profile);
     console.log(`${icon} ${profile}: ${ready ? 'ready' : 'not installed'}${detail}`);
     console.log(`   ${runtimePath}`);
+    if (profile === 'pytorch' && manifest?.variant === 'cuda') {
+      const cudaAvailable = ready ? checkCudaAvailability(getVenvPath('pytorch')) : null;
+      const cudaStatus =
+        cudaAvailable === true
+          ? 'available'
+          : cudaAvailable === false
+            ? 'unavailable'
+            : 'unknown';
+      console.log(`   CUDA: ${cudaStatus}`);
+    }
+    if (profile === 'pytorch' && manifest && manifest.driverVersion !== driverVersion) {
+      console.log(
+        `   ⚠️ driver version differs (installed ${manifest.driverVersion}, current ${driverVersion}). ` +
+          `Run: ${SYNC_PYTORCH_CLI}`,
+      );
+    }
   }
   const setupHints = [];
   if (!isRuntimeReady('mlx') && process.platform === 'darwin') {
-    setupHints.push('pnpm run setup-mlx -w @modular-prompt/driver');
+    setupHints.push(SETUP_MLX_MONOREPO);
   }
   if (!isRuntimeReady('pytorch')) {
-    setupHints.push('pnpm run setup-pytorch -w @modular-prompt/driver');
+    setupHints.push(SETUP_PYTORCH_MONOREPO);
   }
   if (setupHints.length > 0) {
     console.log(`\nRun: ${setupHints.join('  or  ')}`);
@@ -253,19 +482,49 @@ async function cleanupAll() {
 
 function printUsage() {
   console.log(`Usage:
-  node scripts/runtime-cli.js setup mlx         Set up MLX Python runtime (macOS only)
-  node scripts/runtime-cli.js setup pytorch     Set up PyTorch runtime (cpu-minimal)
-  node scripts/runtime-cli.js setup --status    Show runtime status
-  node scripts/runtime-cli.js cleanup mlx       Remove MLX runtime
-  node scripts/runtime-cli.js cleanup pytorch   Remove PyTorch runtime
-  node scripts/runtime-cli.js cleanup --all     Remove entire ~/.modular-prompt
-  node scripts/runtime-cli.js cleanup ... --yes   Skip confirmation
-
-  npm scripts: setup-mlx, setup-pytorch, runtime:status, runtime:cleanup`);
+  modular-prompt-runtime setup mlx         Set up MLX Python runtime (macOS only)
+  modular-prompt-runtime setup pytorch     Set up PyTorch runtime (cpu-minimal)
+  modular-prompt-runtime setup pytorch --variant <variant> [--cuda <version>]
+  modular-prompt-runtime sync pytorch      Sync PyTorch code and dependencies
+  modular-prompt-runtime sync pytorch --variant <variant>
+  modular-prompt-runtime setup --status    Show runtime status
+  modular-prompt-runtime cleanup mlx       Remove MLX runtime
+  modular-prompt-runtime cleanup pytorch   Remove PyTorch runtime
+  modular-prompt-runtime cleanup --all     Remove entire ~/.modular-prompt
+  modular-prompt-runtime cleanup ... --yes   Skip confirmation
+
+  npm scripts: setup-mlx, setup-pytorch, runtime:status, runtime:sync-pytorch, runtime:cleanup`);
+}
+
+function parseVariant(args) {
+  return parseOptionValue(args, '--variant');
+}
+
+function parseCudaVersion(args) {
+  return parseOptionValue(args, '--cuda');
+}
+
+function parseOptionValue(args, option) {
+  const inlinePrefix = `${option}=`;
+  const inline = args.find((arg) => arg.startsWith(inlinePrefix));
+  const index = args.indexOf(option);
+  if (!inline && index === -1) {
+    return undefined;
+  }
+  if (inline && index !== -1) {
+    throw new Error(`Specify ${option} only once`);
+  }
+
+  const value = inline ? inline.slice(inlinePrefix.length) : args[index + 1];
+  if (!value || value.startsWith('-')) {
+    throw new Error(`Missing value for ${option}`);
+  }
+  return value;
 }
 
 async function main() {
-  const [command, target] = process.argv.slice(2);
+  const args = process.argv.slice(2);
+  const [command, target] = args;
 
   if (!command || command === '--help' || command === '-h') {
     printUsage();
@@ -282,7 +541,13 @@ async function main() {
       return;
     }
     if (target === 'pytorch') {
-      setupPytorch();
+      const setupArgs = args.slice(2);
+      const variant = parseVariant(setupArgs) ?? PYTORCH_DEFAULT_VARIANT;
+      const cudaVersion = parseCudaVersion(setupArgs);
+      if (cudaVersion && variant !== 'cuda') {
+        throw new Error('--cuda can only be used with --variant cuda');
+      }
+      setupPytorch(variant, cudaVersion);
       return;
     }
     console.error(`Unknown setup target: ${target ?? '(none)'}`);
@@ -290,6 +555,16 @@ async function main() {
     process.exit(1);
   }
 
+  if (command === 'sync') {
+    if (target === 'pytorch') {
+      syncPytorch(parseVariant(args.slice(2)));
+      return;
+    }
+    console.error(`Unknown sync target: ${target ?? '(none)'}`);
+    printUsage();
+    process.exit(1);
+  }
+
   if (command === 'cleanup') {
     if (target === '--all' || target === 'all') {
       await cleanupAll();

package/scripts/runtime-cli.test.ts

@@ -0,0 +1,163 @@
+import { execFileSync } from 'node:child_process';
+import {
+  chmodSync,
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from 'node:fs';
+import { dirname, join, resolve } from 'node:path';
+import { fileURLToPath } from 'node:url';
+import { tmpdir } from 'node:os';
+import { describe, expect, it } from 'vitest';
+
+const packageRoot = resolve(dirname(fileURLToPath(import.meta.url)), '..');
+const runtimeCli = join(packageRoot, 'scripts', 'runtime-cli.js');
+const driverVersion = (
+  JSON.parse(readFileSync(join(packageRoot, 'package.json'), 'utf8')) as { version: string }
+).version;
+
+describe.skipIf(process.platform === 'win32')('runtime CLI PyTorch setup and sync', () => {
+  it('seeds, preserves customization, and syncs the runtime project', () => {
+    const temporaryDirectory = mkdtempSync(join(tmpdir(), 'modular-prompt-pytorch-cli-'));
+    const fakeBinDirectory = join(temporaryDirectory, 'bin');
+    const fakeUvPath = join(fakeBinDirectory, 'uv');
+    const uvLogPath = join(temporaryDirectory, 'uv.log');
+    const runtimePythonDir = join(
+      temporaryDirectory,
+      'runtimes',
+      'pytorch',
+      'python',
+    );
+
+    try {
+      mkdirSync(fakeBinDirectory, { recursive: true });
+      writeFileSync(
+        fakeUvPath,
+        `#!/usr/bin/env node
+import { appendFileSync, mkdirSync, writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+
+const args = process.argv.slice(2);
+appendFileSync(process.env.FAKE_UV_LOG, args.join(' ') + '\\n');
+
+if (args[0] === '--version') {
+  process.exit(0);
+}
+if (args[0] === 'venv') {
+  const environment = process.env.UV_PROJECT_ENVIRONMENT;
+  mkdirSync(join(environment, 'bin'), { recursive: true });
+  writeFileSync(join(environment, 'bin', 'python'), '');
+  process.exit(0);
+}
+if (args[0] === 'pip' && args[1] === 'list') {
+  process.stdout.write('[]');
+  process.exit(0);
+}
+if (args[0] === 'pip' && args[1] === 'install') {
+  process.exit(0);
+}
+process.exit(1);
+`,
+      );
+      chmodSync(fakeUvPath, 0o755);
+
+      const env = {
+        ...process.env,
+        MODULAR_PROMPT_HOME: temporaryDirectory,
+        FAKE_UV_LOG: uvLogPath,
+        PATH: `${fakeBinDirectory}:${process.env.PATH ?? ''}`,
+      };
+      const runCli = (...args: string[]) =>
+        execFileSync(process.execPath, [runtimeCli, ...args], {
+          encoding: 'utf8',
+          env,
+        });
+
+      expect(() => runCli('setup', 'pytorch', '--cuda', '12.4')).toThrow(
+        /--cuda can only be used with --variant cuda/,
+      );
+      expect(
+        () => runCli('setup', 'pytorch', '--variant', 'cuda', '--cuda', '12'),
+      ).toThrow(/Invalid CUDA version/);
+      runCli('setup', 'pytorch');
+      expect(existsSync(join(runtimePythonDir, 'pyproject.toml'))).toBe(true);
+      expect(existsSync(join(runtimePythonDir, '__main__.py'))).toBe(true);
+      expect(existsSync(join(runtimePythonDir, 'backends', 'base.py'))).toBe(true);
+      const manifestPath = join(
+        temporaryDirectory,
+        'runtimes',
+        'pytorch',
+        'manifest.json',
+      );
+      expect(JSON.parse(readFileSync(manifestPath, 'utf8'))).toMatchObject({
+        profile: 'pytorch',
+        variant: 'cpu-minimal',
+        driverVersion,
+      });
+      expect(() => runCli('sync', 'pytorch', '--variant', 'other')).toThrow(
+        /variant mismatch/,
+      );
+
+      writeFileSync(join(runtimePythonDir, 'pyproject.toml'), 'user dependencies\n');
+      runCli('setup', 'pytorch');
+      expect(readFileSync(join(runtimePythonDir, 'pyproject.toml'), 'utf8')).toBe(
+        'user dependencies\n',
+      );
+
+      writeFileSync(join(runtimePythonDir, 'backends', 'base.py'), 'user code\n');
+      runCli('sync', 'pytorch');
+      expect(readFileSync(join(runtimePythonDir, 'pyproject.toml'), 'utf8')).toBe(
+        'user dependencies\n',
+      );
+      expect(readFileSync(join(runtimePythonDir, 'backends', 'base.py'), 'utf8')).toBe(
+        readFileSync(
+          join(packageRoot, 'src', 'pytorch', 'templates', 'cpu-minimal', 'backends', 'base.py'),
+          'utf8',
+        ),
+      );
+      expect(JSON.parse(readFileSync(manifestPath, 'utf8'))).toMatchObject({
+        profile: 'pytorch',
+        variant: 'cpu-minimal',
+        driverVersion,
+      });
+
+      runCli('setup', 'pytorch', '--variant=cuda', '--cuda=12.1');
+      expect(JSON.parse(readFileSync(manifestPath, 'utf8'))).toMatchObject({
+        profile: 'pytorch',
+        variant: 'cuda',
+        cudaVersion: '12.1',
+        torchVersion: '2.9.1',
+        driverVersion,
+      });
+      expect(readFileSync(join(runtimePythonDir, '__main__.py'), 'utf8')).toContain(
+        'os.environ.get("PYTORCH_DEVICE", "cuda")',
+      );
+      expect(
+        readFileSync(
+          join(packageRoot, 'src', 'pytorch', 'templates', 'cuda', 'pyproject.toml'),
+          'utf8',
+        ),
+      ).toContain('https://download.pytorch.org/whl/cu124');
+
+      runCli('setup', 'pytorch', '--variant', 'cuda');
+      expect(JSON.parse(readFileSync(manifestPath, 'utf8'))).toMatchObject({
+        variant: 'cuda',
+        cudaVersion: '12.4',
+      });
+
+      const uvLog = readFileSync(uvLogPath, 'utf8');
+      expect(uvLog).toContain('venv --clear --python 3.12');
+      expect(uvLog).toContain('pip install');
+      expect(uvLog).toContain('torch==2.9.1');
+      expect(uvLog).toContain('https://download.pytorch.org/whl/cu121');
+      expect(uvLog).toContain('https://download.pytorch.org/whl/cu124');
+      expect(uvLog).toContain(' .');
+      expect(uvLog).not.toContain(' -e ');
+    } finally {
+      rmSync(temporaryDirectory, { recursive: true, force: true });
+    }
+  });
+});

package/src/mlx-ml/python/__main__.py

@@ -90,7 +90,7 @@ if __name__ == "__main__":
 
     capabilities = get_capabilities(backend.get_tokenizer())
     capabilities["model_kind"] = model_kind
-    if model_kind == "lm":
+    if model_kind in {"lm", "vlm"} and "cache_prefill" not in capabilities["methods"]:
         capabilities["methods"].append("cache_prefill")
 
     server = Server(backend, capabilities)