@modular-prompt/driver 0.14.0 → 0.16.0

package/scripts/runtime-cli.js

@@ -0,0 +1,320 @@
+#!/usr/bin/env node
+
+/**
+ * Python runtime 管理 CLI
+ *
+ * setup mlx     — ~/.modular-prompt/runtimes/mlx に venv を作成
+ * setup pytorch — ~/.modular-prompt/runtimes/pytorch に cpu-minimal venv を作成
+ * setup --status
+ * cleanup mlx [--yes]
+ * cleanup --all [--yes]
+ */
+
+import { execSync } from 'child_process';
+import { existsSync, readFileSync, rmSync, mkdirSync } from 'fs';
+import { createInterface } from 'readline';
+import { dirname, join } from 'path';
+import { fileURLToPath, pathToFileURL } from 'url';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const packageRoot = join(__dirname, '..');
+const driverVersion = readPackageVersion();
+
+function runtimeModuleUrl(name) {
+  const distPath = join(packageRoot, 'dist', 'runtime', name);
+  const srcPath = join(packageRoot, 'src', 'runtime', name);
+  return pathToFileURL(existsSync(distPath) ? distPath : srcPath).href;
+}
+
+const {
+  RUNTIME_PROFILES,
+  getModularPromptHome,
+  getRuntimeDir,
+  getVenvPath,
+  getMlxPythonDir,
+  getPytorchPythonDir,
+  isRuntimeReady,
+} = await import(runtimeModuleUrl('paths-core.mjs'));
+
+const {
+  SETUP_MLX_MONOREPO,
+  SETUP_PYTORCH_MONOREPO,
+} = await import(runtimeModuleUrl('setup-commands-core.mjs'));
+
+const {
+  collectInstalledPackages,
+  readManifest,
+  writeManifest,
+} = await import(runtimeModuleUrl('manifest-core.mjs'));
+
+function readPackageVersion() {
+  try {
+    const pkg = JSON.parse(readFileSync(join(packageRoot, 'package.json'), 'utf8'));
+    return pkg.version ?? 'unknown';
+  } catch {
+    return 'unknown';
+  }
+}
+
+function ensureUv() {
+  try {
+    execSync('uv --version', { stdio: 'ignore' });
+    return;
+  } catch {
+    console.log('⚠️  uv is not installed. Installing uv...');
+    execSync('curl -LsSf https://astral.sh/uv/install.sh | sh', { stdio: 'inherit' });
+  }
+}
+
+function setupMlx() {
+  if (process.platform !== 'darwin') {
+    console.error('❌ MLX runtime is only available on macOS (Apple Silicon).');
+    console.error('   For local inference on this platform, use setup-pytorch or vLLM.');
+    process.exit(1);
+  }
+
+  const pythonDir = getMlxPythonDir(packageRoot);
+  if (!existsSync(pythonDir)) {
+    console.error(`❌ MLX Python project not found: ${pythonDir}`);
+    process.exit(1);
+  }
+
+  const venvPath = getVenvPath('mlx');
+  const runtimeDir = getRuntimeDir('mlx');
+
+  console.log('🚀 Setting up MLX runtime...\n');
+  console.log(`📁 Python project: ${pythonDir}`);
+  console.log(`📁 Runtime venv:  ${venvPath}\n`);
+
+  ensureUv();
+  mkdirSync(runtimeDir, { recursive: true });
+
+  const env = {
+    ...process.env,
+    UV_PROJECT_ENVIRONMENT: venvPath,
+  };
+
+  try {
+    execSync('uv venv --clear --python 3.13', { cwd: pythonDir, stdio: 'inherit', env });
+    execSync('uv pip install -e .', { cwd: pythonDir, stdio: 'inherit', env });
+
+    writeManifest('mlx', {
+      profile: 'mlx',
+      driverVersion,
+      platform: process.platform,
+      pythonVersion: '3.13',
+      createdAt: new Date().toISOString(),
+      packages: collectInstalledPackages(pythonDir, venvPath),
+    });
+
+    console.log('\n✅ MLX runtime setup completed.');
+    console.log(`   Home: ${getModularPromptHome()}`);
+    console.log('   You can now use MlxDriver from @modular-prompt/driver');
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    console.error('❌ Failed to setup MLX runtime:', message);
+    process.exit(1);
+  }
+}
+
+const PYTORCH_CPU_INDEX = 'https://download.pytorch.org/whl/cpu';
+const PYTORCH_PYTHON_VERSION = '3.12';
+
+function setupPytorch() {
+  const pythonDir = getPytorchPythonDir(packageRoot);
+  if (!existsSync(pythonDir)) {
+    console.error(`❌ PyTorch Python project not found: ${pythonDir}`);
+    process.exit(1);
+  }
+
+  const venvPath = getVenvPath('pytorch');
+  const runtimeDir = getRuntimeDir('pytorch');
+
+  console.log('🚀 Setting up PyTorch runtime (cpu-minimal)...\n');
+  console.log(`📁 Python project: ${pythonDir}`);
+  console.log(`📁 Runtime venv:  ${venvPath}`);
+  console.log(`📦 torch index:    ${PYTORCH_CPU_INDEX}\n`);
+
+  ensureUv();
+  mkdirSync(runtimeDir, { recursive: true });
+
+  const env = {
+    ...process.env,
+    UV_PROJECT_ENVIRONMENT: venvPath,
+  };
+
+  try {
+    execSync(`uv venv --clear --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,
+    });
+    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,
+    });
+
+    console.log('\n✅ PyTorch runtime setup completed (cpu-minimal).');
+    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');
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    console.error('❌ Failed to setup PyTorch runtime:', message);
+    process.exit(1);
+  }
+}
+
+function formatManifestDetail(manifest) {
+  const parts = [`driver ${manifest.driverVersion}`];
+  if (manifest.variant) {
+    parts.push(`variant ${manifest.variant}`);
+  }
+  const torchVersion = manifest.torchVersion ?? manifest.packages?.torch;
+  if (torchVersion) {
+    parts.push(`torch ${torchVersion}`);
+  }
+  parts.push(manifest.createdAt);
+  return ` (${parts.join(', ')})`;
+}
+
+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 icon = ready ? '✅' : '❌';
+    const runtimePath = getRuntimeDir(profile);
+    console.log(`${icon} ${profile}: ${ready ? 'ready' : 'not installed'}${detail}`);
+    console.log(`   ${runtimePath}`);
+  }
+  const setupHints = [];
+  if (!isRuntimeReady('mlx') && process.platform === 'darwin') {
+    setupHints.push(SETUP_MLX_MONOREPO);
+  }
+  if (!isRuntimeReady('pytorch')) {
+    setupHints.push(SETUP_PYTORCH_MONOREPO);
+  }
+  if (setupHints.length > 0) {
+    console.log(`\nRun: ${setupHints.join('  or  ')}`);
+  }
+}
+
+async function confirm(message) {
+  if (process.argv.includes('--yes') || process.argv.includes('-y')) {
+    return true;
+  }
+  const rl = createInterface({ input: process.stdin, output: process.stdout });
+  const answer = await new Promise((resolve) => {
+    rl.question(`${message} [y/N] `, resolve);
+  });
+  rl.close();
+  return String(answer).toLowerCase() === 'y' || String(answer).toLowerCase() === 'yes';
+}
+
+async function cleanupProfile(profile) {
+  if (!existsSync(getRuntimeDir(profile))) {
+    console.log(`ℹ️  ${profile}: nothing to clean (${getRuntimeDir(profile)} does not exist)`);
+    return;
+  }
+  const ok = await confirm(`Delete ${profile} runtime at ${getRuntimeDir(profile)}?`);
+  if (!ok) {
+    console.log('Cancelled.');
+    return;
+  }
+  rmSync(getRuntimeDir(profile), { recursive: true, force: true });
+  console.log(`✅ Removed ${profile} runtime.`);
+}
+
+async function cleanupAll() {
+  const home = getModularPromptHome();
+  if (!existsSync(home)) {
+    console.log(`ℹ️  Nothing to clean (${home} does not exist)`);
+    return;
+  }
+  const ok = await confirm(`Delete entire modular-prompt home at ${home}?`);
+  if (!ok) {
+    console.log('Cancelled.');
+    return;
+  }
+  rmSync(home, { recursive: true, force: true });
+  console.log(`✅ Removed ${home}`);
+}
+
+function printUsage() {
+  console.log(`Usage:
+  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 --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:cleanup`);
+}
+
+async function main() {
+  const [command, target] = process.argv.slice(2);
+
+  if (!command || command === '--help' || command === '-h') {
+    printUsage();
+    process.exit(0);
+  }
+
+  if (command === 'setup') {
+    if (target === '--status' || target === 'status') {
+      printStatus();
+      return;
+    }
+    if (target === 'mlx') {
+      setupMlx();
+      return;
+    }
+    if (target === 'pytorch') {
+      setupPytorch();
+      return;
+    }
+    console.error(`Unknown setup target: ${target ?? '(none)'}`);
+    printUsage();
+    process.exit(1);
+  }
+
+  if (command === 'cleanup') {
+    if (target === '--all' || target === 'all') {
+      await cleanupAll();
+      return;
+    }
+    if (target === 'mlx' || target === 'pytorch') {
+      await cleanupProfile(target);
+      return;
+    }
+    console.error(`Unknown cleanup target: ${target ?? '(none)'}`);
+    printUsage();
+    process.exit(1);
+  }
+
+  console.error(`Unknown command: ${command}`);
+  printUsage();
+  process.exit(1);
+}
+
+main().catch((error) => {
+  console.error(error);
+  process.exit(1);
+});

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

@@ -6,7 +6,28 @@ from utils.vlm_utils import detect_model_kind
 from server import Server
 
 model_name = sys.argv[1] if len(sys.argv) > 1 else "mlx-community/gemma-3-270m-it-qat-4bit"
-text_only = "--text-only" in sys.argv
+
+VALID_BACKENDS = {"auto", "lm", "vlm", "optiq"}
+
+
+def parse_backend_mode() -> str:
+    if "--backend" in sys.argv:
+        idx = sys.argv.index("--backend")
+        if idx + 1 < len(sys.argv):
+            mode = sys.argv[idx + 1]
+            if mode not in VALID_BACKENDS:
+                sys.stderr.write(
+                    f"Invalid --backend value: {mode}. "
+                    f"Use one of: {', '.join(sorted(VALID_BACKENDS))}\n"
+                )
+                sys.exit(1)
+            return mode
+    if "--text-only" in sys.argv:
+        return "lm"
+    return "auto"
+
+
+backend_mode = parse_backend_mode()
 
 drafter_model = None
 if "--drafter" in sys.argv:
@@ -25,8 +46,26 @@ if "--draft-block-size" in sys.argv:
             sys.exit(1)
 
 
-def create_backend(model_name: str, text_only: bool = False):
-    model_kind = "lm" if text_only else detect_model_kind(model_name)
+def ensure_optiq_registered() -> None:
+    import optiq  # noqa: F401 - registers vendored architectures into mlx-lm
+
+
+def create_backend(model_name: str, mode: str = "auto"):
+    if mode == "optiq":
+        ensure_optiq_registered()
+        mode = "lm"
+
+    if mode == "lm":
+        backend = MlxLmBackend()
+        backend.load(model_name)
+        return backend, "lm"
+
+    if mode == "vlm":
+        backend = MlxVlmBackend()
+        backend.load(model_name)
+        return backend, "vlm"
+
+    model_kind = detect_model_kind(model_name)
 
     if model_kind == "vlm":
         backend = MlxVlmBackend()
@@ -42,7 +81,7 @@ def create_backend(model_name: str, text_only: bool = False):
 
 
 if __name__ == "__main__":
-    backend, model_kind = create_backend(model_name, text_only)
+    backend, model_kind = create_backend(model_name, backend_mode)
 
     if drafter_model:
         backend.load_drafter(drafter_model)

package/src/mlx-ml/python/backends/cache_archive.py

@@ -0,0 +1,41 @@
+"""Streaming zip storage for MLX prompt caches."""
+
+from __future__ import annotations
+
+import zipfile
+from collections.abc import Callable
+from typing import Any
+
+
+CACHE_ENTRY_NAME = "prompt_cache.safetensors"
+
+
+def save_prompt_cache_zip(
+    file_name: str,
+    cache: Any,
+    save_impl: Callable[[Any, Any], None],
+) -> None:
+    """Save a prompt cache into a compressed zip member as it is produced.
+
+    ``save_impl`` receives the zip member's writable stream.  It must write
+    the safetensors payload to that stream instead of creating an intermediate
+    uncompressed file.
+    """
+    with zipfile.ZipFile(
+        file_name,
+        mode="w",
+        compression=zipfile.ZIP_DEFLATED,
+        allowZip64=True,
+    ) as archive:
+        with archive.open(CACHE_ENTRY_NAME, mode="w", force_zip64=True) as cache_file:
+            save_impl(cache_file, cache)
+
+
+def load_prompt_cache_zip(
+    file_name: str,
+    load_impl: Callable[[Any], Any],
+) -> Any:
+    """Load a prompt cache from the safetensors member of a zip archive."""
+    with zipfile.ZipFile(file_name, mode="r") as archive:
+        with archive.open(CACHE_ENTRY_NAME, mode="r") as cache_file:
+            return load_impl(cache_file)

package/src/mlx-ml/python/backends/mlx_lm.py

@@ -7,13 +7,29 @@ from typing import Any, Iterator
 
 from mlx_lm import load as mlx_lm_load
 from mlx_lm import stream_generate as mlx_lm_stream_generate
-from mlx_lm.models.cache import make_prompt_cache, save_prompt_cache, load_prompt_cache, trim_prompt_cache
+from mlx_lm.models.cache import (
+    make_prompt_cache,
+    save_prompt_cache as mlx_save_prompt_cache,
+    load_prompt_cache as mlx_load_prompt_cache,
+    trim_prompt_cache,
+)
 from mlx_lm.sample_utils import make_sampler
 
 from backends.base import ModelBackend
+from backends.cache_archive import load_prompt_cache_zip, save_prompt_cache_zip
 from utils.token_utils import is_eod_token
 
 
+def save_prompt_cache(file_name: str, cache: Any) -> None:
+    """Save a prompt cache as a compressed archive without an uncompressed copy."""
+    save_prompt_cache_zip(file_name, cache, mlx_save_prompt_cache)
+
+
+def load_prompt_cache(file_name: str) -> list:
+    """Load a prompt cache from its compressed safetensors archive."""
+    return load_prompt_cache_zip(file_name, mlx_load_prompt_cache)
+
+
 class MlxLmBackend(ModelBackend):
     """`mlx_lm` backend for text-only models."""
 

package/src/mlx-ml/python/handlers/__init__.py

@@ -1,6 +1,7 @@
 from handlers.cache import handle_cache_prefill
 from handlers.capabilities import handle_capabilities
-from handlers.chat import handle_chat
 from handlers.completion import handle_completion
 from handlers.format_test import handle_format_test
+from handlers.generate import handle_generate
+from handlers.render import handle_render
 from handlers.tokenize import handle_tokenize

package/src/mlx-ml/python/handlers/completion.py

@@ -1,11 +1,7 @@
 from __future__ import annotations
 
-import os
-import re
-import sys
-
 from backends.base import ModelBackend
-from handlers.cancel import poll_cancel
+from handlers.generate import handle_generate
 
 
 def handle_completion(
@@ -15,25 +11,5 @@ def handle_completion(
     images: list | None = None,
     max_image_size: int = 768,
 ) -> None:
-    """completion API の処理"""
-    if options is None:
-        options = {}
-
-    final_options = dict(options)
-    if images:
-        final_options["max_image_size"] = max_image_size
-        if os.getenv('MLX_DEBUG'):
-            display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
-            sys.stderr.write(f"--- vlm completion (images: {len(images)}, max_size: {max_image_size})\n{display_prompt}\n")
-    elif os.getenv('MLX_DEBUG'):
-        if isinstance(prompt, list):
-            sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
-        else:
-            sys.stderr.write(f"--- prompt\n{prompt}\n")
-
-    for response in backend.stream_generate(prompt, final_options, images):
-        if poll_cancel():
-            break
-        print(response.text.replace("\0", ""), end="", flush=True)
-
-    print("\n", end="\0", flush=True)
+    """completion API（後方互換）— generate に委譲"""
+    handle_generate(backend, prompt, options, images, max_image_size)