From f9a828fbc7745cb3caa663bde481615aa3e818ae Mon Sep 17 00:00:00 2001 From: Lion Rayonnant <106342136+lionrayonnant@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:25:00 +0200 Subject: [PATCH] feat: add AMD ROCm support (Linux + Windows) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds AMD GPU detection and a PyTorch/ROCm redirect for extension setup, on top of the existing NVIDIA/MPS/CPU paths. - electron/main/gpu-detect.ts: detects AMD GPUs (KFD topology on Linux, Win32_VideoController on Windows), resolves the ROCm pip index and requirements. NVIDIA keeps detection priority; explicit overrides (MODLY_TORCH_FLAVOR, MODLY_ROCM_GFX, MODLY_ROCM_INDEX, MODLY_ROCM_TORCH_SPEC) are available for machines the auto-detection gets wrong. - electron/main/setup-launcher.ts: extracted from ipc-handlers.ts, adds a compatibility shim that redirects an extension's pip torch install to ROCm wheels — needed because most third-party extension setup.py scripts predate AMD support and hardcode a CUDA index. - api/routers/extensions.py: the FastAPI-side GPU detection no longer mistakes a ROCm build's device capability for CUDA compute capability (both answer torch.cuda.get_device_capability the same way), and reads the AMD compute target from the same KFD topology as the Electron side. - electron/main/copy-runtime.ts: fixes an unrelated but blocking AppImage bug found while verifying this end-to-end — fs.cp rewrote the bundled Python runtime's relative symlinks into absolute paths pointing at the ephemeral AppImage mount, so every extension venv died on the next launch. verbatimSymlinks keeps them relative. - docs/running-on-amd-rocm.md, arch/decisions/AMD-ROCM-SUPPORT.md: usage, verified configuration, and known limitations. Verified end-to-end on a Radeon RX 9060 XT (gfx1200): detection, ROCm wheel install (torch 2.13.0+rocm7.2), and a full image-to-3D generation through hunyuan3d-mini all complete successfully on the GPU. --- README.md | 3 + api/routers/extensions.py | 81 +++++- arch/decisions/AMD-ROCM-SUPPORT.md | 99 +++++++ arch/decisions/README.md | 1 + docs/running-on-amd-rocm.md | 193 +++++++++++++ electron/main/copy-runtime.test.mjs | 93 ++++++ electron/main/copy-runtime.ts | 35 +++ electron/main/gpu-detect.test.mjs | 230 +++++++++++++++ electron/main/gpu-detect.ts | 390 ++++++++++++++++++++++++++ electron/main/ipc-handlers.ts | 191 +++---------- electron/main/python-setup.ts | 5 +- electron/main/setup-launcher.test.mjs | 224 +++++++++++++++ electron/main/setup-launcher.ts | 231 +++++++++++++++ 13 files changed, 1611 insertions(+), 165 deletions(-) create mode 100644 arch/decisions/AMD-ROCM-SUPPORT.md create mode 100644 docs/running-on-amd-rocm.md create mode 100644 electron/main/copy-runtime.test.mjs create mode 100644 electron/main/copy-runtime.ts create mode 100644 electron/main/gpu-detect.test.mjs create mode 100644 electron/main/gpu-detect.ts create mode 100644 electron/main/setup-launcher.test.mjs create mode 100644 electron/main/setup-launcher.ts diff --git a/README.md b/README.md index 11184ddc..293368e3 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,9 @@ npm run build ## Platform notes +- AMD GPUs are supported through ROCm on Linux and Windows: a Radeon card is detected + automatically and extensions are steered to ROCm PyTorch wheels, with no ROCm install + required. See [docs/running-on-amd-rocm.md](docs/running-on-amd-rocm.md). - macOS support targets Apple Silicon only. - macOS uses native window controls. Windows and Linux keep the existing custom controls. - The top bar includes a live RAM indicator sourced from the main process. diff --git a/api/routers/extensions.py b/api/routers/extensions.py index 2313d3b3..ce549035 100644 --- a/api/routers/extensions.py +++ b/api/routers/extensions.py @@ -1,6 +1,9 @@ import asyncio +import json +import re import subprocess import sys +from pathlib import Path from fastapi import APIRouter, HTTPException router = APIRouter(tags=["extensions"]) @@ -43,14 +46,29 @@ async def setup_extension(ext_id: str): return {"status": "skipped", "reason": "no setup.py"} # Detect GPU compute capability - gpu_sm = _detect_gpu_sm() + gpu_sm = _detect_gpu_sm() + gfx_target = _detect_gfx_target() + + # Pass arguments as JSON so setup.py sees torch_flavor. Note this endpoint is + # a fallback: Electron normally runs setup.py itself, and only that path gets + # the ROCm index rewriting for extensions that ignore torch_flavor. + args = json.dumps({ + "python_exe": sys.executable, + "ext_dir": str(ext_dir), + "gpu_sm": gpu_sm, + "cuda_version": 0, + "accelerator": "rocm" if gfx_target else ("cuda" if gpu_sm else "cpu"), + "torch_flavor": "rocm" if gfx_target else ("cuda" if gpu_sm else "cpu"), + "gfx_target": gfx_target, + "platform": sys.platform, + }) # Run setup.py using Modly's embedded Python (sys.executable) loop = asyncio.get_running_loop() result = await loop.run_in_executor( None, lambda: subprocess.run( - [sys.executable, str(setup_py), sys.executable, str(ext_dir), str(gpu_sm)], + [sys.executable, str(setup_py), args], capture_output=True, text=True, ) @@ -60,9 +78,10 @@ async def setup_extension(ext_id: str): raise HTTPException(500, f"setup.py failed:\n{result.stderr}") return { - "status": "ok", - "gpu_sm": gpu_sm, - "output": result.stdout, + "status": "ok", + "gpu_sm": gpu_sm, + "gfx_target": gfx_target, + "output": result.stdout, } @@ -74,12 +93,62 @@ async def extension_errors(): def _detect_gpu_sm() -> int: - """Returns GPU compute capability as integer (e.g. 86 for SM 8.6), or 0 if no GPU.""" + """ + Returns GPU compute capability as integer (e.g. 86 for SM 8.6), or 0 if no GPU. + + Returns 0 on ROCm too. PyTorch's HIP build answers the whole torch.cuda API, + so get_device_capability() happily reports (12, 0) for a gfx1200 Radeon — + indistinguishable from an sm_120 Blackwell, which would send setup.py to the + CUDA 12.8 index. AMD cards are identified by _detect_gfx_target() instead. + """ try: import torch + if torch.version.hip: + return 0 if torch.cuda.is_available(): major, minor = torch.cuda.get_device_capability(0) return major * 10 + minor except Exception: pass return 0 + + +def _detect_gfx_target() -> str: + """ + Returns the ROCm compute target (e.g. "gfx1200"), or "" when there is no AMD GPU. + + Reads the kernel's KFD topology rather than asking torch: this process runs + in Modly's main venv, which has no torch at all (see api/requirements.txt). + The amdgpu driver publishes the target on its own, so no ROCm install is + needed either. Mirrors electron/main/gpu-detect.ts. + """ + kfd_nodes = Path("/sys/class/kfd/kfd/topology/nodes") + if not Path("/dev/kfd").exists() or not kfd_nodes.is_dir(): + return "" + + def _prop(text: str, key: str) -> int: + match = re.search(rf"^{key}\s+(\d+)\s*$", text, re.M) + return int(match.group(1)) if match else 0 + + try: + nodes = sorted(kfd_nodes.iterdir(), key=lambda p: int(p.name) if p.name.isdigit() else 0) + except OSError: + return "" + + for node in nodes: + try: + text = (node / "properties").read_text() + except OSError: + continue + # Node 0 is the CPU node (simd_count 0) and carries no compute target. + if _prop(text, "simd_count") <= 0: + continue + # major*10000 + minor*100 + step, minor and step read as hex digits. + version = _prop(text, "gfx_target_version") + if version <= 0: + continue + major, minor, step = version // 10000, (version % 10000) // 100, version % 100 + if major <= 0 or minor > 15 or step > 15: + continue + return f"gfx{major}{minor:x}{step:x}" + return "" diff --git a/arch/decisions/AMD-ROCM-SUPPORT.md b/arch/decisions/AMD-ROCM-SUPPORT.md new file mode 100644 index 00000000..7cd4ee13 --- /dev/null +++ b/arch/decisions/AMD-ROCM-SUPPORT.md @@ -0,0 +1,99 @@ +# AMD-ROCM-SUPPORT + +- Status: proposed +- Date: 2026-08-17 + +## Decision + +Modly supports AMD Radeon GPUs through ROCm on Linux and Windows. Detection is +automatic and requires no ROCm installation on the user's machine. + +Scope and operating rules: + +- Detection is centralised in `electron/main/gpu-detect.ts` and produces, for AMD + machines, a compute target plus the pip index and requirements an extension's + torch install must end up using. +- NVIDIA keeps detection priority. On a machine with both vendors the existing + CUDA behaviour is unchanged. +- AMD machines report `gpu_sm = 0` and `cuda_version = 0`, never a synthesised + compute capability. +- Extensions are told the flavour via a `torch_flavor` setup argument. Extensions + that ignore it are corrected by a rewrite shim in `electron/main/setup-launcher.ts`. +- The ROCm wheel source differs by platform: `download.pytorch.org/whl/rocm7.2` + on Linux, `repo.amd.com/rocm/whl-multi-arch/` on Windows. +- Every automatic choice has an environment-variable override. See + `docs/running-on-amd-rocm.md`. + +## Context + +Modly never installs PyTorch itself. Each extension ships a `setup.py` that +creates its own venv and installs torch from an index it hardcodes — and those +scripts are third-party code in separate GitHub repositories that Modly cannot +edit. Before this work `detectGpuInfo()` only probed `nvidia-smi`, so an AMD +machine was reported as `accelerator: 'cpu'` with `gpu_sm: 0`, which sent every +extension down its legacy CUDA 11.8 branch and installed a torch that cannot see +the GPU at all. + +That leaves two distinct problems, and both have to be solved: + +- Extensions that *do* understand AMD were never told. The official + `modly-hunyuan3d-mini-extension` has accepted a `torch_flavor: "rocm"` argument + for some time; Modly simply never sent it. +- Extensions that don't understand AMD — `triposg`, `trellis2`, and the rest — + hardcode `--index-url .../whl/cu124` and have no branch to select. Passing an + argument achieves nothing for them. + +The wheel sources are also not symmetric across platforms. `download.pytorch.org` +publishes no ROCm wheels for Windows at all; AMD's own multi-arch index does, but +there the compute target is selected by a pip extra (`torch[device-gfx1200]`) +rather than by the index URL, which means Windows needs the compute target +*before* the install, not after. + +## Consequences + +- **A rewrite shim is unavoidable.** Correcting extensions we cannot edit means + intercepting their pip calls. The launcher already patched `subprocess` for two + other compatibility fixes, so ROCm redirection joins those rather than + introducing a new mechanism. +- **The decision is made in TypeScript, applied in Python.** The launcher is an + inline Python string that cannot be unit-tested in isolation, so index and + requirement resolution lives in `gpu-detect.ts` and reaches the launcher as + environment variables. The launcher is separately exercised end-to-end by + `setup-launcher.test.mjs`, which runs it against the command shapes the + official extensions actually use. +- **`gpu_sm = 0` is load-bearing, not a placeholder.** Extensions written before + `torch_flavor` branch on that number, and 0 selects their most conservative + path. It also keeps them off `rembg[gpu]`, whose `onnxruntime-gpu` is + CUDA-only. Reporting a synthesised capability instead would break both. + PyTorch's HIP build answers the whole `torch.cuda` API, so + `get_device_capability()` reports `(12, 0)` for a gfx1200 Radeon — + indistinguishable from an sm_120 Blackwell. `api/routers/extensions.py` guards + on `torch.version.hip` for this reason. +- **Compute-target discovery is platform-specific.** Linux reads + `gfx_target_version` from the kernel's KFD topology, which needs no ROCm + install and no external binary. Windows has no equivalent, so it maps PCI + device ids from `Win32_VideoController` through a table keyed by silicon. That + table is a maintenance surface: new AMD silicon needs an entry, and an unmapped + AMD card falls back to CPU with an actionable message rather than guessing a + wheel. +- **The Linux and Windows torch versions diverge.** Linux gets unpinned wheels + from the pytorch.org ROCm index (currently torch 2.11+); Windows gets a pinned + pair from AMD's index. Extension code written against torch 2.6/2.7 may not + survive that jump, which is why `MODLY_ROCM_INDEX` and `MODLY_ROCM_TORCH_SPEC` + exist as first-class escape hatches rather than debug affordances. +- **Linux is verified, Windows is not.** On a Radeon RX 9060 XT (gfx1200), + `torch 2.13.0+rocm7.2` loads, rocBLAS and MIOpen kernels execute, 14 GB of the + card's 16 GB allocates and reads back cleanly ([ROCm #6295](https://github.com/ROCm/ROCm/issues/6295), + which reports this card capped near 8 GB, did not reproduce), and a full + image-to-3D generation completes through the normal `ExtensionProcess` path in + 221 s. For Windows the wheel URLs, `cp311` availability and index layout were + checked, but no end-to-end run has been performed. +- **This work also required fixing an unrelated AppImage bug** to be verifiable + at all. `ensureStableEmbeddedPython()` copied the bundled runtime with + `fs.cp`, which rewrites relative symlinks into absolute paths pointing back at + the ephemeral `/tmp/.mount_Modly-XXXXXX/` mount — so the "stable" copy was not + stable, and every extension venv built from it died on the next launch with a + misleading `No module named 'PIL'`. See `electron/main/copy-runtime.ts`. +- **Texture generation is out of scope.** `api/texture_baker` already carries a + HIP build path, but those native extensions are not built as part of standard + extension setup, so texture generation is not covered by this ADR. diff --git a/arch/decisions/README.md b/arch/decisions/README.md index 063df577..925441ea 100644 --- a/arch/decisions/README.md +++ b/arch/decisions/README.md @@ -8,3 +8,4 @@ single reviewable document. Current ADRs: - [APPLE-SILICON-SUPPORT](./APPLE-SILICON-SUPPORT.md) +- [AMD-ROCM-SUPPORT](./AMD-ROCM-SUPPORT.md) diff --git a/docs/running-on-amd-rocm.md b/docs/running-on-amd-rocm.md new file mode 100644 index 00000000..5fe60f8a --- /dev/null +++ b/docs/running-on-amd-rocm.md @@ -0,0 +1,193 @@ +# Running Modly on an AMD GPU (ROCm) + +Modly's default GPU path is NVIDIA/CUDA (plus Metal/MPS on Apple Silicon). This page +covers the AMD path: how Modly detects a Radeon card, which PyTorch wheels it steers +extensions to, and what to do when the automatic choice is wrong. + +Nothing here needs a manual setup: install the app, install an extension, and the AMD +path is taken automatically when an AMD GPU is present. + +--- + +## 1. Requirements + +| | Linux | Windows | +|---|---|---| +| GPU | RDNA 2 or newer discrete Radeon (see the table below) | same | +| Driver | in-tree `amdgpu` kernel driver — any recent distro kernel | Adrenalin with the ROCm runtime (26.2.2 or newer) | +| ROCm install | **not required** — the PyTorch ROCm wheels bundle their own runtime | not required | +| Device access | `/dev/kfd` and `/dev/dri/renderD*` must be readable by your user | n/a | + +On most distributions `/dev/kfd` is world-accessible. If it is not, add yourself to the +`render` group and log back in: + +```bash +ls -l /dev/kfd # crw-rw-rw- → nothing to do +sudo usermod -aG render "$USER" +``` + +Integrated (APU) graphics are not mapped by the Windows detection table. They work on +Linux whenever the kernel publishes a compute target, but they are not a target Modly +tests against. + +--- + +## 2. How detection works + +Detection lives in `electron/main/gpu-detect.ts` and runs before an extension's +`setup.py`, in this order: + +1. `MODLY_TORCH_FLAVOR` (`cuda` / `rocm` / `cpu`) — an explicit override wins over everything. +2. Apple Silicon → MPS. +3. `nvidia-smi` → CUDA. **NVIDIA keeps priority**: on a machine with both vendors, + nothing about the existing CUDA behaviour changes. +4. AMD: + - **Linux** — reads `gfx_target_version` from the kernel's KFD topology + (`/sys/class/kfd/kfd/topology/nodes/*/properties`). No ROCm install and no external + binary involved; the `amdgpu` driver publishes this on its own. + - **Windows** — reads the PCI device id from `Win32_VideoController` via PowerShell + and maps it to a compute target. +5. Otherwise → CPU. + +You can check what was detected in the logs (Settings → Logs), on the line beginning +`[ext-setup] accelerator=`: + +``` +[ext-setup] accelerator=rocm gfx=gfx1200 torch_index=https://download.pytorch.org/whl/rocm7.2 +``` + +### Supported compute targets + +| Silicon | Compute target | Cards | +|---|---|---| +| Navi 44 | `gfx1200` | RX 9060 XT | +| Navi 48 | `gfx1201` | RX 9070, RX 9070 XT, RX 9070 GRE, AI PRO R9700 | +| Navi 31 | `gfx1100` | RX 7900 XT/XTX/GRE, PRO W7800/W7900 | +| Navi 32 | `gfx1101` | RX 7700 XT, RX 7800 XT, PRO W7700 | +| Navi 33 | `gfx1102` | RX 7600 series, PRO W7500/W7600 | +| Navi 21 | `gfx1030` | RX 6800/6800 XT/6900 XT/6950 XT, PRO W6800 | +| Navi 22/23/24 | `gfx1031` / `gfx1032` / `gfx1034` | RX 6700 / 6600 / 6400 series | + +On Linux the target is read from the kernel, so any AMD GPU the kernel knows about is +picked up — the table above only bounds the *Windows* mapping. AMD officially supports +`gfx1030`, `gfx110x` and `gfx120x`; the RDNA 2 mid-range entries work in practice but +are not part of AMD's supported matrix. + +--- + +## 3. Which wheels get installed + +Modly does not install PyTorch itself — each extension's `setup.py` does, from an index +it hardcodes. On an AMD machine Modly corrects that choice two ways: + +- It passes `torch_flavor: "rocm"` in the setup arguments. Extensions that know about + AMD (the official `hunyuan3d-mini` one does) branch on it themselves. +- For every extension that doesn't, a compatibility shim in the setup launcher rewrites + the pip command: the CUDA (or CPU) index is swapped for the ROCm one, and the pinned + `torch==…` requirements are replaced. Non-torch installs are left untouched. + +| Platform | Index | Requirements | +|---|---|---| +| Linux | `https://download.pytorch.org/whl/rocm7.2` | `torch`, `torchvision`, unpinned | +| Windows | `https://repo.amd.com/rocm/whl-multi-arch/` | `torch[device-gfxNNNN]==2.11.0+rocm7.14.0` and matching `torchvision` | + +The two differ because `download.pytorch.org` publishes no ROCm wheels for Windows at +all. AMD's multi-arch index does, and there the compute target is selected by a pip +extra rather than by the index URL. + +`sm` and `cuda_version` are deliberately reported as `0` for AMD. Extensions written +before `torch_flavor` existed branch on those numbers, and `0` sends them down their +most conservative path — which also keeps them off `rembg[gpu]`, whose `onnxruntime-gpu` +is CUDA-only. + +--- + +## 4. Verifying an install + +After installing an extension (or running **Repair** on it from the Models page), check +what actually landed in its venv: + +```bash +# Linux; adjust to your extensions directory +EXT=~/Documents/Modly/extensions/hunyuan3d-mini +"$EXT/venv/bin/python" -c "import torch; print(torch.__version__, '| hip', torch.version.hip, '| avail', torch.cuda.is_available(), '|', torch.cuda.get_device_name(0))" +``` + +Expected: a `+rocm…` version, a non-null `hip`, `True`, and your card's name. A version +ending in `+cu118` or `+cu124` means the AMD path was not taken — check the detection +line in the logs. + +Then confirm real VRAM is usable, not just that the device opens: + +```bash +"$EXT/venv/bin/python" -c "import torch; x = torch.empty(5_000_000_000, dtype=torch.float16, device='cuda'); torch.cuda.synchronize(); print('10 GB allocated OK')" +``` + +--- + +## 5. Escape hatches + +All of these are environment variables read at detection time — set them before +launching Modly. + +| Variable | Effect | +|---|---| +| `MODLY_TORCH_FLAVOR` | `cuda` / `rocm` / `cpu`. Forces the path, skipping detection entirely. | +| `MODLY_ROCM_GFX` | Forces the compute target (e.g. `gfx1201`). Needed for an AMD card the Windows table doesn't map. | +| `MODLY_ROCM_INDEX` | Overrides the pip index, e.g. `https://download.pytorch.org/whl/rocm6.4` to fall back to torch 2.8/2.9. | +| `MODLY_ROCM_TORCH_SPEC` | Overrides the requirements entirely, space-separated: `"torch==2.9.1 torchvision==0.24.1"`. | +| `HSA_OVERRIDE_GFX_VERSION` | ROCm's own override, for cards without native kernels (e.g. `10.3.0` on an unsupported RDNA 2 part). Not needed on RDNA 3/4. | + +After changing any of these, run **Repair** on the extension so its venv is rebuilt. + +--- + +## 6. Troubleshooting + +| Symptom | Cause | Fix | +|---|---|---| +| Logs say `accelerator=cpu` on an AMD machine | `/dev/kfd` missing or unreadable (Linux), or an unmapped PCI id (Windows) | Check `ls -l /dev/kfd`; on Windows set `MODLY_ROCM_GFX` | +| `torch.__version__` ends in `+cu118` | Extension venv predates AMD support | Run **Repair** on the extension | +| `torch.cuda.is_available()` is `False` with a ROCm build | Device nodes not accessible from the process | Add your user to the `render` group, log back in | +| `HIP error: invalid device function` | Wheel has no kernels for your card | Set `HSA_OVERRIDE_GFX_VERSION` to a supported nearby target | +| Extension imports fail after install (`diffusers`/`transformers`) | The ROCm index only carries recent torch (2.11+), newer than some extensions expect | `MODLY_ROCM_INDEX=https://download.pytorch.org/whl/rocm6.4`, then **Repair** | +| Allocations above ~8 GB segfault on a 16 GB RX 9060 XT | [ROCm issue #6295](https://github.com/ROCm/ROCm/issues/6295) — did not reproduce on torch 2.13.0+rocm7.2 (see below) | If you hit it, pin an older stack via `MODLY_ROCM_INDEX` | + +--- + +## 7. Verified configuration and limitations + +The Linux path was measured end to end on a Radeon RX 9060 XT (Navi 44, gfx1200, +16 GB), CachyOS kernel 7.1.8, with the stack this page installs by default: + +- `torch 2.13.0+rocm7.2` / `torchvision 0.28.0+rocm7.2`, HIP runtime 7.2.53211 +- Detection reported `gfx1200`; `torch.cuda.is_available()` is `True` and + `get_device_properties(0).gcnArchName` is `gfx1200` +- rocBLAS (fp16 matmul) and MIOpen (conv2d) kernels both execute — no + `no kernel image is available` failures +- Allocations of 4/6/8/10/12/14 GB all succeeded, filled and read back. + [ROCm #6295](https://github.com/ROCm/ROCm/issues/6295), which reports this exact + card capped near 8 GB with a segfault, **did not reproduce** on this stack. + It remains open upstream, so it may still affect older ROCm builds. +- A full image-to-3D generation with `hunyuan3d-mini` completed through the normal + `ExtensionProcess` subprocess path: model load 18 s, generation 221 s, 15.8 MB GLB. + +### Attention kernels + +During generation PyTorch emits: + +> Mem Efficient attention on Current AMD GPU is still experimental. Enable it with +> `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` + +Generation works without it — `scaled_dot_product_attention` falls back to a +slower path. Setting `TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1` before launching +Modly enables the memory-efficient kernels, at the cost of running code AMD still +labels experimental on RDNA 3/4. Modly does not set it for you. + +Limitations: + +- **Windows is untested** by the Modly maintainers. The wheel URLs and the `cp311` + availability were verified, but no end-to-end run has been done on that path. +- **Texture generation** relies on optional native extensions (`texture_baker`, + `uv_unwrapper`) whose CUDA kernels have a HIP path but are not built as part of the + standard extension setup. diff --git a/electron/main/copy-runtime.test.mjs b/electron/main/copy-runtime.test.mjs new file mode 100644 index 00000000..fd8b9ba4 --- /dev/null +++ b/electron/main/copy-runtime.test.mjs @@ -0,0 +1,93 @@ +/** + * Guards the symlink property that makes the AppImage's "stable" Python copy + * actually stable. Getting this wrong is silent at copy time and only breaks on + * the *next* launch, once the source mount is gone — so it needs a test that + * checks the links rather than the copy succeeding. + */ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { + mkdtempSync, mkdirSync, writeFileSync, symlinkSync, readlinkSync, rmSync, existsSync, +} from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve, isAbsolute } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-copy-test-')), 'copy-runtime.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/main/copy-runtime.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +const { copyRuntimeTree } = loadModule() + +/** Builds a miniature of the bundled runtime: a real binary plus relative links. */ +function makeRuntime(root) { + mkdirSync(join(root, 'bin'), { recursive: true }) + mkdirSync(join(root, 'lib'), { recursive: true }) + writeFileSync(join(root, 'bin', 'python3.11'), '#!/bin/sh\n', 'utf8') + symlinkSync('python3.11', join(root, 'bin', 'python3')) + symlinkSync('python3.11', join(root, 'bin', 'python')) + writeFileSync(join(root, 'lib', 'libpython3.11.so.1.0'), '', 'utf8') + symlinkSync('libpython3.11.so.1.0', join(root, 'lib', 'libpython3.11.so')) +} + +test('copyRuntimeTree keeps relative symlinks relative', async () => { + const dir = mkdtempSync(join(tmpdir(), 'modly-runtime-')) + const source = join(dir, 'mount', 'python-embed') + const dest = join(dir, 'stable', 'python-embed') + makeRuntime(source) + + await copyRuntimeTree(source, dest) + + for (const link of ['bin/python3', 'bin/python', 'lib/libpython3.11.so']) { + const target = readlinkSync(join(dest, link)) + assert.ok( + !isAbsolute(target), + `${link} was rewritten to an absolute path (${target}); it would point back at the source mount`, + ) + assert.ok(!target.includes(source), `${link} still references the source tree`) + } + + rmSync(dir, { recursive: true, force: true }) +}) + +test('the copy survives the source being deleted', async () => { + // This is the actual failure mode: the AppImage mount disappears between + // launches, and every venv built from the copy dies with it. + const dir = mkdtempSync(join(tmpdir(), 'modly-runtime-')) + const source = join(dir, 'mount', 'python-embed') + const dest = join(dir, 'stable', 'python-embed') + makeRuntime(source) + + await copyRuntimeTree(source, dest) + rmSync(join(dir, 'mount'), { recursive: true, force: true }) + + // existsSync follows symlinks, so this is false for a dangling link — exactly + // the check generator_registry.py's _venv_python(...).exists() performs. + assert.ok(existsSync(join(dest, 'bin', 'python3')), 'bin/python3 is dangling after the source went away') + assert.ok(existsSync(join(dest, 'lib', 'libpython3.11.so')), 'libpython3.11.so is dangling') +}) + +test('copyRuntimeTree copies regular files and directory structure', async () => { + const dir = mkdtempSync(join(tmpdir(), 'modly-runtime-')) + const source = join(dir, 'mount', 'python-embed') + const dest = join(dir, 'stable', 'python-embed') + makeRuntime(source) + + await copyRuntimeTree(source, dest) + + assert.ok(existsSync(join(dest, 'bin', 'python3.11'))) + assert.ok(existsSync(join(dest, 'lib', 'libpython3.11.so.1.0'))) + + rmSync(dir, { recursive: true, force: true }) +}) diff --git a/electron/main/copy-runtime.ts b/electron/main/copy-runtime.ts new file mode 100644 index 00000000..eb062f5c --- /dev/null +++ b/electron/main/copy-runtime.ts @@ -0,0 +1,35 @@ +/** + * Copying the bundled Python runtime out of an ephemeral AppImage mount. + * + * Kept apart from python-setup.ts (which imports electron) so the symlink + * behaviour this depends on can be tested — it is subtle, silent when wrong, + * and only shows up on the *next* launch. + */ + +import { cp } from 'fs/promises' + +/** + * Copies a self-contained runtime tree, keeping relative symlinks relative. + * + * `verbatimSymlinks` is the whole point. Without it `fs.cp` resolves a relative + * link (`bin/python3 -> python3.11`) into an absolute path pointing back at the + * *source* tree. When the source is an AppImage mount at + * /tmp/.mount_Modly-XXXXXX/ — which is a different path on every launch — the + * copy silently keeps a hard dependency on a directory that is about to vanish: + * + * - `bin/python3` in the "stable" copy points into the old mount + * - `sys._base_executable` of any venv made from it inherits that path + * - every extension venv records it in pyvenv.cfg and as a bin/python symlink + * - on the next launch the mount is gone, so every extension venv is dead + * + * The user-visible symptom is remote from the cause: the extension registry + * finds no usable venv, falls back to importing generator.py in the main API + * process, and reports a missing third-party module such as `No module named 'PIL'`. + */ +export async function copyRuntimeTree(source: string, destination: string): Promise { + await cp(source, destination, { + recursive: true, + preserveTimestamps: true, + verbatimSymlinks: true, + }) +} diff --git a/electron/main/gpu-detect.test.mjs b/electron/main/gpu-detect.test.mjs new file mode 100644 index 00000000..c4a2edcb --- /dev/null +++ b/electron/main/gpu-detect.test.mjs @@ -0,0 +1,230 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-gpu-test-')), 'gpu-detect.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/main/gpu-detect.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +const mod = loadModule() + +// ─── KFD topology parsing ───────────────────────────────────────────────────── + +test('formatGfxTarget decodes gfx_target_version across GPU generations', () => { + // Encoding is major*10000 + minor*100 + step, minor and step read as hex + // digits. 120000 is the value this machine's RX 9060 XT actually reports. + assert.equal(mod.formatGfxTarget(120000), 'gfx1200') // Navi 44 / RX 9060 XT + assert.equal(mod.formatGfxTarget(120001), 'gfx1201') // Navi 48 + assert.equal(mod.formatGfxTarget(110000), 'gfx1100') // Navi 31 + assert.equal(mod.formatGfxTarget(110002), 'gfx1102') // Navi 33 + assert.equal(mod.formatGfxTarget(100300), 'gfx1030') // Navi 21 + assert.equal(mod.formatGfxTarget(90402), 'gfx942') // MI300, hex step + assert.equal(mod.formatGfxTarget(90010), 'gfx90a') // MI200, hex step +}) + +test('formatGfxTarget rejects the "no GPU" and malformed encodings', () => { + assert.equal(mod.formatGfxTarget(0), null) + assert.equal(mod.formatGfxTarget(-1), null) + assert.equal(mod.formatGfxTarget(1.5), null) +}) + +test('parseKfdGfxTarget skips the CPU node and reads the first GPU', () => { + // Verbatim shape of /sys/class/kfd/kfd/topology/nodes/*/properties + const cpuNode = 'cpu_cores_count 32\nsimd_count 0\ngfx_target_version 0\n' + const gpuNode = 'cpu_cores_count 0\nsimd_count 64\ngfx_target_version 120000\n' + + assert.equal(mod.parseKfdGfxTarget([cpuNode, gpuNode]), 'gfx1200') +}) + +test('parseKfdGfxTarget returns null without a usable GPU node', () => { + assert.equal(mod.parseKfdGfxTarget([]), null) + assert.equal(mod.parseKfdGfxTarget(['simd_count 0\ngfx_target_version 0\n']), null) + // A node advertising SIMDs but no compute target is not something we can target + assert.equal(mod.parseKfdGfxTarget(['simd_count 64\ngfx_target_version 0\n']), null) +}) + +test('parseKfdGfxTarget does not confuse a key with its prefix', () => { + // simd_count must not be satisfied by e.g. "max_simd_count" + const node = 'max_simd_count 999\nsimd_count 64\ngfx_target_version 110002\n' + assert.equal(mod.parseKfdGfxTarget([node]), 'gfx1102') +}) + +// ─── nvidia-smi parsing (non-regression) ────────────────────────────────────── + +test('parseNvidiaSmi maps compute cap and driver version to CUDA version', () => { + assert.deepEqual(mod.parseNvidiaSmi('8.6, 551.61\n'), { sm: 86, cudaVersion: 124 }) + assert.deepEqual(mod.parseNvidiaSmi('12.0, 572.16\n'), { sm: 120, cudaVersion: 128 }) + assert.deepEqual(mod.parseNvidiaSmi('6.1, 470.82\n'), { sm: 61, cudaVersion: 118 }) +}) + +test('parseNvidiaSmi falls back to sm 86 on an unparseable compute cap', () => { + assert.deepEqual(mod.parseNvidiaSmi('N/A, 551.61\n'), { sm: 86, cudaVersion: 124 }) +}) + +test('parseNvidiaSmi returns null on empty output', () => { + assert.equal(mod.parseNvidiaSmi(''), null) + assert.equal(mod.parseNvidiaSmi(' \n'), null) +}) + +// ─── Windows adapter mapping ────────────────────────────────────────────────── + +test('parseWindowsVideoControllers accepts both the object and array JSON shapes', () => { + const single = mod.parseWindowsVideoControllers( + '{"Name":"AMD Radeon RX 9060 XT","PNPDeviceID":"PCI\\\\VEN_1002&DEV_7590&SUBSYS_06391043&REV_C0\\\\4&1"}', + ) + assert.deepEqual(single, [ + { name: 'AMD Radeon RX 9060 XT', pnpDeviceId: 'PCI\\VEN_1002&DEV_7590&SUBSYS_06391043&REV_C0\\4&1' }, + ]) + + const many = mod.parseWindowsVideoControllers( + '[{"Name":"A","PNPDeviceID":"PCI\\\\VEN_1002&DEV_7550"},{"Name":"B","PNPDeviceID":"PCI\\\\VEN_8086&DEV_1234"}]', + ) + assert.equal(many.length, 2) +}) + +test('parseWindowsVideoControllers survives malformed PowerShell output', () => { + assert.deepEqual(mod.parseWindowsVideoControllers('not json'), []) + assert.deepEqual(mod.parseWindowsVideoControllers('{"Name":"No id"}'), []) +}) + +test('parseAmdPciDeviceId only matches AMD vendor ids', () => { + assert.equal(mod.parseAmdPciDeviceId('PCI\\VEN_1002&DEV_7590&SUBSYS_0'), '7590') + assert.equal(mod.parseAmdPciDeviceId('PCI\\VEN_10DE&DEV_2684'), null) +}) + +test('resolveWindowsGfxTarget maps device ids by silicon, not by marketing range', () => { + const target = (deviceId) => + mod.resolveWindowsGfxTarget([{ name: 'x', pnpDeviceId: `PCI\\VEN_1002&DEV_${deviceId}` }]).gfxTarget + + assert.equal(target('7590'), 'gfx1200') // Navi 44 + assert.equal(target('7550'), 'gfx1201') // Navi 48 + assert.equal(target('744C'), 'gfx1100') // Navi 31, uppercase from WMI + assert.equal(target('747e'), 'gfx1101') // Navi 32 + // 0x73f0 sells as "RX 7600M XT" but is Navi 33 — it must not land with its + // 0x73xx RDNA2 neighbours. + assert.equal(target('73f0'), 'gfx1102') + assert.equal(target('73bf'), 'gfx1030') // Navi 21 +}) + +test('resolveWindowsGfxTarget reports unmapped AMD cards instead of guessing', () => { + const result = mod.resolveWindowsGfxTarget([ + { name: 'AMD Radeon RX 9999', pnpDeviceId: 'PCI\\VEN_1002&DEV_FFFF' }, + ]) + assert.equal(result.gfxTarget, null) + assert.deepEqual(result.amdAdapters, ['AMD Radeon RX 9999']) +}) + +test('resolveWindowsGfxTarget ignores non-AMD adapters', () => { + const result = mod.resolveWindowsGfxTarget([ + { name: 'NVIDIA RTX 4090', pnpDeviceId: 'PCI\\VEN_10DE&DEV_2684' }, + ]) + assert.equal(result.gfxTarget, null) + assert.deepEqual(result.amdAdapters, []) +}) + +// ─── ROCm wheel source resolution ───────────────────────────────────────────── + +test('resolveRocmTorchSpec uses the unpinned pytorch.org index on Linux', () => { + const { indexUrl, specs } = mod.resolveRocmTorchSpec('linux', 'gfx1200', {}) + assert.equal(indexUrl, 'https://download.pytorch.org/whl/rocm7.2') + assert.deepEqual(specs, ['torch', 'torchvision']) +}) + +test('resolveRocmTorchSpec uses AMD\'s index with a device extra on Windows', () => { + // download.pytorch.org publishes no ROCm wheels for Windows at all. + const { indexUrl, specs } = mod.resolveRocmTorchSpec('win32', 'gfx1200', {}) + assert.equal(indexUrl, 'https://repo.amd.com/rocm/whl-multi-arch/') + assert.deepEqual(specs, [ + 'torch[device-gfx1200]==2.11.0+rocm7.14.0', + 'torchvision[device-gfx1200]==0.26.0+rocm7.14.0', + ]) +}) + +test('resolveRocmTorchSpec honours MODLY_ROCM_INDEX and MODLY_ROCM_TORCH_SPEC', () => { + const rolledBack = mod.resolveRocmTorchSpec('linux', 'gfx1200', { + MODLY_ROCM_INDEX: 'https://download.pytorch.org/whl/rocm6.4', + MODLY_ROCM_TORCH_SPEC: 'torch==2.8.0 torchvision==0.23.0', + }) + assert.equal(rolledBack.indexUrl, 'https://download.pytorch.org/whl/rocm6.4') + assert.deepEqual(rolledBack.specs, ['torch==2.8.0', 'torchvision==0.23.0']) +}) + +test('resolveRocmTorchSpec stays unpinned on Windows without a compute target', () => { + // No target means no device extra to ask for; better an install that fails + // loudly than one silently pinned to the wrong architecture. + const { specs } = mod.resolveRocmTorchSpec('win32', null, {}) + assert.deepEqual(specs, ['torch', 'torchvision']) +}) + +test('torchFlavorFor maps accelerators to the setup.py argument', () => { + assert.equal(mod.torchFlavorFor('rocm'), 'rocm') + assert.equal(mod.torchFlavorFor('cuda'), 'cuda') + assert.equal(mod.torchFlavorFor('cpu'), 'cpu') + assert.equal(mod.torchFlavorFor('mps'), 'cpu') +}) + +// ─── Detection precedence ───────────────────────────────────────────────────── + +test('detectGpuInfo keeps Apple Silicon on MPS', async () => { + const info = await mod.detectGpuInfo({ env: {}, platform: 'darwin', arch: 'arm64' }) + assert.equal(info.accelerator, 'mps') +}) + +test('detectGpuInfo forced to rocm resolves wheels without probing hardware', async () => { + const info = await mod.detectGpuInfo({ + env: { MODLY_TORCH_FLAVOR: 'rocm', MODLY_ROCM_GFX: 'gfx1201' }, + platform: 'linux', + arch: 'x64', + }) + assert.equal(info.accelerator, 'rocm') + assert.equal(info.gfxTarget, 'gfx1201') + assert.equal(info.torchIndexUrl, 'https://download.pytorch.org/whl/rocm7.2') + // sm/cudaVersion stay at 0 so CUDA-era extensions take their most + // conservative branch (and keep off rembg[gpu], which is CUDA-only). + assert.equal(info.sm, 0) + assert.equal(info.cudaVersion, 0) +}) + +test('detectGpuInfo forced to rocm overrides MPS on Apple Silicon', async () => { + const info = await mod.detectGpuInfo({ + env: { MODLY_TORCH_FLAVOR: 'rocm', MODLY_ROCM_GFX: 'gfx1200' }, + platform: 'darwin', + arch: 'arm64', + }) + assert.equal(info.accelerator, 'rocm') +}) + +test('detectGpuInfo forced to cpu short-circuits everything', async () => { + const info = await mod.detectGpuInfo({ + env: { MODLY_TORCH_FLAVOR: 'cpu', MODLY_ROCM_GFX: 'gfx1200' }, + platform: 'linux', + arch: 'x64', + }) + assert.deepEqual(info, { sm: 0, cudaVersion: 0, accelerator: 'cpu' }) +}) + +test('detectGpuInfo falls back to CPU when forced to rocm on Windows with no target', async () => { + const logs = [] + const info = await mod.detectGpuInfo({ + env: { MODLY_TORCH_FLAVOR: 'rocm' }, + platform: 'win32', + arch: 'x64', + onLog: (line) => logs.push(line), + }) + assert.equal(info.accelerator, 'cpu') + assert.ok(logs.some((l) => l.includes('MODLY_ROCM_GFX'))) +}) diff --git a/electron/main/gpu-detect.ts b/electron/main/gpu-detect.ts new file mode 100644 index 00000000..62201151 --- /dev/null +++ b/electron/main/gpu-detect.ts @@ -0,0 +1,390 @@ +/** + * GPU detection and PyTorch flavour resolution. + * + * Deliberately free of electron imports: the parsing and resolution helpers are + * pure and get unit-tested by bundling this file directly (gpu-detect.test.mjs). + * + * Modly never installs PyTorch itself — each extension's setup.py does, from an + * index it picks on its own. What we produce here is the information that lets + * that choice land on the right wheels: the accelerator, and for AMD the ROCm + * compute target plus the pip index/requirements the setup must end up using. + */ + +import { spawn } from 'child_process' +import { existsSync, readdirSync, readFileSync } from 'fs' +import { join } from 'path' + +export type Accelerator = 'cuda' | 'rocm' | 'mps' | 'cpu' +export type TorchFlavor = 'cuda' | 'rocm' | 'cpu' + +export interface GpuInfo { + sm: number + cudaVersion: number + accelerator: Accelerator + /** ROCm compute target ("gfx1200"). Only set when accelerator is 'rocm'. */ + gfxTarget?: string + /** pip --index-url the torch install has to come from (ROCm only). */ + torchIndexUrl?: string + /** pip requirements replacing whatever torch pins an extension hardcodes. */ + torchSpecs?: string[] +} + +// ─── ROCm wheel sources ─────────────────────────────────────────────────────── +// +// Linux and Windows need different indexes. download.pytorch.org publishes no +// ROCm wheels for Windows at all, so Windows has to go through AMD's multi-arch +// index, where the compute target is selected by a pip extra +// (torch[device-gfx1200]) instead of by the index URL. + +const ROCM_LINUX_INDEX = 'https://download.pytorch.org/whl/rocm7.2' +const ROCM_WINDOWS_INDEX = 'https://repo.amd.com/rocm/whl-multi-arch/' + +// Pinned because AMD's index carries several ROCm builds side by side; this is +// the newest pair published for cp311 (Modly's embedded Python) on Windows. +const ROCM_WINDOWS_TORCH = '2.11.0+rocm7.14.0' +const ROCM_WINDOWS_TORCHVISION = '0.26.0+rocm7.14.0' + +const KFD_TOPOLOGY_DIR = '/sys/class/kfd/kfd/topology/nodes' + +/** + * AMD PCI device id → ROCm compute target, for Windows where there is no KFD + * topology to read. Keyed by silicon rather than by marketing name: 0x73f0 is + * sold as an "RX 7600M XT" but is Navi 33, so it belongs with gfx1102, not with + * its 0x73xx neighbours. Device ids come from the pci.ids database. + */ +const WINDOWS_PCI_GFX_TARGETS: Record = { + // Navi 21 (RDNA 2) + '73a1': 'gfx1030', '73a2': 'gfx1030', '73a3': 'gfx1030', '73a5': 'gfx1030', + '73ab': 'gfx1030', '73ae': 'gfx1030', '73af': 'gfx1030', '73bf': 'gfx1030', + // Navi 22 (RDNA 2) + '73c3': 'gfx1031', '73ce': 'gfx1031', '73df': 'gfx1031', + // Navi 23 (RDNA 2) + '73e0': 'gfx1032', '73e1': 'gfx1032', '73e3': 'gfx1032', '73ef': 'gfx1032', + '73ff': 'gfx1032', + // Navi 24 (RDNA 2) + '7421': 'gfx1034', '7422': 'gfx1034', '7423': 'gfx1034', '7424': 'gfx1034', + '743f': 'gfx1034', + // Navi 31 (RDNA 3) + '7448': 'gfx1100', '7449': 'gfx1100', '744a': 'gfx1100', '744b': 'gfx1100', + '744c': 'gfx1100', '745e': 'gfx1100', + // Navi 32 (RDNA 3) + '7460': 'gfx1101', '7461': 'gfx1101', '7470': 'gfx1101', '747e': 'gfx1101', + // Navi 33 (RDNA 3) + '73f0': 'gfx1102', '7480': 'gfx1102', '7481': 'gfx1102', '7483': 'gfx1102', + '7487': 'gfx1102', '7489': 'gfx1102', '748b': 'gfx1102', '7499': 'gfx1102', + '749f': 'gfx1102', + // Navi 44 / Navi 48 (RDNA 4) + '7590': 'gfx1200', + '7550': 'gfx1201', '7551': 'gfx1201', +} + +// ─── Pure parsing helpers ───────────────────────────────────────────────────── + +/** + * Parses `nvidia-smi --query-gpu=compute_cap,driver_version --format=csv,noheader`. + * Returns null when the output carries no usable line. + */ +export function parseNvidiaSmi(stdout: string): { sm: number; cudaVersion: number } | null { + const line = stdout.trim().split('\n')[0]?.trim() // e.g. "8.6, 551.61" + if (!line) return null + + const parts = line.split(',').map((s) => s.trim()) + const sm = Math.round(parseFloat(parts[0] ?? '') * 10) // → 86 + + // Derive max supported CUDA version from driver version + // Driver ≥ 520 → CUDA 11.8, ≥ 525 → 12.0, ≥ 530 → 12.1, ≥ 535 → 12.2, + // ≥ 545 → 12.3, ≥ 550 → 12.4, ≥ 555 → 12.5, ≥ 560 → 12.6 + const driverMajor = parseInt((parts[1] ?? '').split('.')[0] ?? '0', 10) + let cudaVersion = 118 // safe minimum + if (driverMajor >= 570) cudaVersion = 128 // Blackwell (RTX 50xx, sm_120) + else if (driverMajor >= 560) cudaVersion = 126 + else if (driverMajor >= 555) cudaVersion = 125 + else if (driverMajor >= 550) cudaVersion = 124 + else if (driverMajor >= 545) cudaVersion = 123 + else if (driverMajor >= 535) cudaVersion = 122 + else if (driverMajor >= 530) cudaVersion = 121 + else if (driverMajor >= 525) cudaVersion = 120 + else if (driverMajor >= 520) cudaVersion = 118 + + return { sm: isNaN(sm) ? 86 : sm, cudaVersion } +} + +/** + * Decodes the KFD `gfx_target_version` integer (major*10000 + minor*100 + step, + * with minor and step read as hex digits) into a compute target name: + * 120000 → gfx1200, 90402 → gfx942, 90010 → gfx90a. + */ +export function formatGfxTarget(version: number): string | null { + if (!Number.isInteger(version) || version <= 0) return null + const major = Math.floor(version / 10000) + const minor = Math.floor((version % 10000) / 100) + const step = version % 100 + if (major <= 0 || minor > 15 || step > 15) return null + return `gfx${major}${minor.toString(16)}${step.toString(16)}` +} + +/** + * Picks the compute target out of the KFD topology node `properties` files. + * Node 0 is the CPU node (simd_count 0) and is skipped; the first real GPU wins. + */ +export function parseKfdGfxTarget(nodeProperties: string[]): string | null { + for (const text of nodeProperties) { + if (readKfdProperty(text, 'simd_count') <= 0) continue + const target = formatGfxTarget(readKfdProperty(text, 'gfx_target_version')) + if (target) return target + } + return null +} + +function readKfdProperty(text: string, key: string): number { + const match = new RegExp(`^${key}\\s+(\\d+)\\s*$`, 'm').exec(text) + return match ? parseInt(match[1], 10) : 0 +} + +export interface VideoController { + name: string + pnpDeviceId: string +} + +/** + * Parses the JSON emitted by `Get-CimInstance Win32_VideoController | ConvertTo-Json`. + * PowerShell emits a bare object rather than an array when there is one adapter. + */ +export function parseWindowsVideoControllers(stdout: string): VideoController[] { + let parsed: unknown + try { + parsed = JSON.parse(stdout) + } catch { + return [] + } + const list = Array.isArray(parsed) ? parsed : [parsed] + return list + .filter((entry): entry is Record => !!entry && typeof entry === 'object') + .map((entry) => ({ + name: typeof entry['Name'] === 'string' ? entry['Name'] : '', + pnpDeviceId: typeof entry['PNPDeviceID'] === 'string' ? entry['PNPDeviceID'] : '', + })) + .filter((c) => c.pnpDeviceId !== '') +} + +/** Extracts the PCI device id from a PNPDeviceID, e.g. `PCI\VEN_1002&DEV_7590&…` → "7590". */ +export function parseAmdPciDeviceId(pnpDeviceId: string): string | null { + const match = /VEN_1002&DEV_([0-9A-F]{4})/i.exec(pnpDeviceId) + return match ? match[1].toLowerCase() : null +} + +/** + * Resolves a compute target from the installed display adapters. Returns the + * AMD adapters it saw as well, so an unmapped card can be reported by name + * instead of silently falling back to CPU. + */ +export function resolveWindowsGfxTarget( + controllers: VideoController[], +): { gfxTarget: string | null; amdAdapters: string[] } { + const amdAdapters: string[] = [] + let gfxTarget: string | null = null + + for (const controller of controllers) { + const deviceId = parseAmdPciDeviceId(controller.pnpDeviceId) + if (!deviceId) continue + amdAdapters.push(controller.name || `PCI 1002:${deviceId}`) + gfxTarget ??= WINDOWS_PCI_GFX_TARGETS[deviceId] ?? null + } + + return { gfxTarget, amdAdapters } +} + +/** + * The pip index and requirements an extension's torch install has to end up + * using on this platform. `MODLY_ROCM_INDEX` and `MODLY_ROCM_TORCH_SPEC` + * (space-separated requirements) override either half — the escape hatch when a + * newer torch breaks an extension and you need to drop back to, say, rocm6.4. + */ +export function resolveRocmTorchSpec( + platform: string, + gfxTarget: string | null, + env: NodeJS.ProcessEnv = process.env, +): { indexUrl: string; specs: string[] } { + const indexOverride = env['MODLY_ROCM_INDEX']?.trim() + const specOverride = env['MODLY_ROCM_TORCH_SPEC']?.trim() + + const isWindows = platform === 'win32' + const indexUrl = indexOverride || (isWindows ? ROCM_WINDOWS_INDEX : ROCM_LINUX_INDEX) + + if (specOverride) return { indexUrl, specs: specOverride.split(/\s+/).filter(Boolean) } + + // AMD's multi-arch index ships one torch per compute target, selected by a + // pip extra. The pytorch.org ROCm index bakes the targets into a single + // wheel, so there is nothing to select and nothing to pin. + if (isWindows && gfxTarget) { + return { + indexUrl, + specs: [ + `torch[device-${gfxTarget}]==${ROCM_WINDOWS_TORCH}`, + `torchvision[device-${gfxTarget}]==${ROCM_WINDOWS_TORCHVISION}`, + ], + } + } + return { indexUrl, specs: ['torch', 'torchvision'] } +} + +/** The `torch_flavor` value extension setup.py scripts branch on. */ +export function torchFlavorFor(accelerator: Accelerator): TorchFlavor { + if (accelerator === 'rocm') return 'rocm' + if (accelerator === 'cuda') return 'cuda' + return 'cpu' +} + +export function describeGpuInfo(info: GpuInfo): string { + const bits = [`accelerator=${info.accelerator}`] + if (info.accelerator === 'cuda') bits.push(`sm=${info.sm}`, `cuda=${info.cudaVersion}`) + if (info.gfxTarget) bits.push(`gfx=${info.gfxTarget}`) + if (info.torchIndexUrl) bits.push(`torch_index=${info.torchIndexUrl}`) + return bits.join(' ') +} + +// ─── Detection ──────────────────────────────────────────────────────────────── + +function cpuInfo(): GpuInfo { + return { sm: 0, cudaVersion: 0, accelerator: 'cpu' } +} + +/** + * AMD keeps sm/cudaVersion at 0 on purpose. Extensions that predate `torch_flavor` + * branch on those two numbers, and 0 sends them down their most conservative + * path — which also keeps them off `rembg[gpu]`, whose onnxruntime-gpu is + * CUDA-only. Their torch install is then corrected by the ROCm setup shim. + */ +function rocmInfo( + gfxTarget: string | null, + platform: string, + env: NodeJS.ProcessEnv, +): GpuInfo { + const { indexUrl, specs } = resolveRocmTorchSpec(platform, gfxTarget, env) + return { + sm: 0, + cudaVersion: 0, + accelerator: 'rocm', + ...(gfxTarget ? { gfxTarget } : {}), + torchIndexUrl: indexUrl, + torchSpecs: specs, + } +} + +function readFlavorOverride(env: NodeJS.ProcessEnv): TorchFlavor | null { + const raw = env['MODLY_TORCH_FLAVOR']?.trim().toLowerCase() + return raw === 'cuda' || raw === 'rocm' || raw === 'cpu' ? raw : null +} + +function queryNvidiaSmi(): Promise<{ sm: number; cudaVersion: number } | null> { + return new Promise((resolve) => { + // Query compute cap + driver version in one call + const proc = spawn('nvidia-smi', ['--query-gpu=compute_cap,driver_version', '--format=csv,noheader'], { + stdio: ['ignore', 'pipe', 'ignore'], + }) + let out = '' + proc.stdout?.on('data', (d: Buffer) => { out += d.toString() }) + proc.on('close', (code) => resolve(code === 0 ? parseNvidiaSmi(out) : null)) + proc.on('error', () => resolve(null)) + }) +} + +/** + * Reads the compute target straight out of the kernel's KFD topology. This + * needs no ROCm installation and no external binary — the amdgpu driver alone + * publishes it, which is exactly the state of a machine that has only ever run + * PyTorch ROCm wheels (they bundle their own runtime). + */ +function readKfdGfxTarget(): string | null { + if (!existsSync('/dev/kfd')) return null + try { + const nodes = readdirSync(KFD_TOPOLOGY_DIR).sort((a, b) => Number(a) - Number(b)) + const properties = nodes.map((node) => { + try { + return readFileSync(join(KFD_TOPOLOGY_DIR, node, 'properties'), 'utf-8') + } catch { + return '' + } + }) + return parseKfdGfxTarget(properties) + } catch { + return null + } +} + +function queryWindowsVideoControllers(): Promise { + return new Promise((resolve) => { + const proc = spawn('powershell', [ + '-NoProfile', '-NonInteractive', '-Command', + 'Get-CimInstance Win32_VideoController | Select-Object Name,PNPDeviceID | ConvertTo-Json -Compress', + ], { stdio: ['ignore', 'pipe', 'ignore'] }) + let out = '' + proc.stdout?.on('data', (d: Buffer) => { out += d.toString() }) + proc.on('close', (code) => resolve(code === 0 ? parseWindowsVideoControllers(out) : [])) + proc.on('error', () => resolve([])) + }) +} + +export interface DetectOptions { + env?: NodeJS.ProcessEnv + platform?: string + arch?: string + onLog?: (line: string) => void +} + +export async function detectGpuInfo(options: DetectOptions = {}): Promise { + const env = options.env ?? process.env + const platform = options.platform ?? process.platform + const arch = options.arch ?? process.arch + const log = options.onLog ?? (() => {}) + + const forced = readFlavorOverride(env) + if (forced) log(`[gpu-detect] MODLY_TORCH_FLAVOR=${forced} — skipping auto-detection`) + + if (forced === 'cpu') return cpuInfo() + + if (forced === 'rocm') { + const gfxTarget = await resolveGfxTarget(env, platform) + if (!gfxTarget && platform === 'win32') { + log('[gpu-detect] ROCm forced on Windows but no compute target found — set MODLY_ROCM_GFX (e.g. gfx1200)') + return cpuInfo() + } + return rocmInfo(gfxTarget, platform, env) + } + + if (platform === 'darwin' && arch === 'arm64') { + return { sm: 0, cudaVersion: 0, accelerator: 'mps' } + } + + // NVIDIA keeps priority: on a machine with both, nothing about the existing + // CUDA behaviour changes. + const nvidia = await queryNvidiaSmi() + if (nvidia) return { ...nvidia, accelerator: 'cuda' } + if (forced === 'cuda') return { sm: 0, cudaVersion: 0, accelerator: 'cuda' } + + const gfxTarget = await resolveGfxTarget(env, platform) + if (gfxTarget) { + log(`[gpu-detect] AMD GPU detected — compute target ${gfxTarget}`) + return rocmInfo(gfxTarget, platform, env) + } + + if (platform === 'win32') { + const { amdAdapters } = resolveWindowsGfxTarget(await queryWindowsVideoControllers()) + if (amdAdapters.length > 0) { + log( + `[gpu-detect] AMD GPU found (${amdAdapters.join(', ')}) but its ROCm compute target is unknown. ` + + 'Falling back to CPU — set MODLY_ROCM_GFX (e.g. gfx1201) to force one.', + ) + } + } + + return cpuInfo() +} + +async function resolveGfxTarget(env: NodeJS.ProcessEnv, platform: string): Promise { + const override = env['MODLY_ROCM_GFX']?.trim() + if (override) return override + if (platform === 'linux') return readKfdGfxTarget() + if (platform === 'win32') return resolveWindowsGfxTarget(await queryWindowsVideoControllers()).gfxTarget + return null +} diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index 14fc984a..f4ad43e3 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -33,67 +33,20 @@ import { resolvePathWithinRoot, } from './extension-path-guard' import { validateInstallManifest } from './extension-install-utils' +import { detectGpuInfo, describeGpuInfo, torchFlavorFor, type GpuInfo } from './gpu-detect' +import { SETUP_LAUNCHER_SOURCE } from './setup-launcher' import { registerWorkspaceAssetLibraryIpcHandlers } from './artifact-registry-service' import { updatesSupported } from './updater' type WindowGetter = () => BrowserWindow | null const pExecFile = promisify(execFile) -// ─── GPU detect (best-effort, no Python required) ───────────────────────────── - -interface GpuInfo { - sm: number - cudaVersion: number - accelerator: 'cuda' | 'mps' | 'cpu' -} - -function detectGpuInfo(): Promise { - if (process.platform === 'darwin' && process.arch === 'arm64') { - return Promise.resolve({ sm: 0, cudaVersion: 0, accelerator: 'mps' }) - } - - return new Promise((resolve) => { - // Query compute cap + driver version in one call - const proc = spawn('nvidia-smi', ['--query-gpu=compute_cap,driver_version', '--format=csv,noheader'], { - stdio: ['ignore', 'pipe', 'ignore'], - }) - let out = '' - proc.stdout?.on('data', (d: Buffer) => { out += d.toString() }) - proc.on('close', (code) => { - if (code === 0) { - const line = out.trim().split('\n')[0].trim() // e.g. "8.6, 551.61" - const parts = line.split(',').map(s => s.trim()) - const sm = Math.round(parseFloat(parts[0] ?? '') * 10) // → 86 - // Derive max supported CUDA version from driver version - // Driver ≥ 520 → CUDA 11.8, ≥ 525 → 12.0, ≥ 530 → 12.1, ≥ 535 → 12.2, - // ≥ 545 → 12.3, ≥ 550 → 12.4, ≥ 555 → 12.5, ≥ 560 → 12.6 - const driverMajor = parseInt((parts[1] ?? '').split('.')[0] ?? '0', 10) - let cudaVersion = 118 // safe minimum - if (driverMajor >= 570) cudaVersion = 128 // Blackwell (RTX 50xx, sm_120) - else if (driverMajor >= 560) cudaVersion = 126 - else if (driverMajor >= 555) cudaVersion = 125 - else if (driverMajor >= 550) cudaVersion = 124 - else if (driverMajor >= 545) cudaVersion = 123 - else if (driverMajor >= 535) cudaVersion = 122 - else if (driverMajor >= 530) cudaVersion = 121 - else if (driverMajor >= 525) cudaVersion = 120 - else if (driverMajor >= 520) cudaVersion = 118 - resolve({ sm: isNaN(sm) ? 86 : sm, cudaVersion, accelerator: 'cuda' }) - } else { - resolve({ sm: 0, cudaVersion: 0, accelerator: 'cpu' }) - } - }) - proc.on('error', () => resolve({ sm: 0, cudaVersion: 0, accelerator: 'cpu' })) - }) -} - // ─── Run an extension's setup.py directly (no FastAPI needed) ───────────────── function runExtensionSetup( - extDir: string, - gpuSm: number, - cudaVersion: number, - onLog?: (line: string) => void, + extDir: string, + gpu: GpuInfo, + onLog?: (line: string) => void, ): Promise { return new Promise((resolve, reject) => { const userData = app.getPath('userData') @@ -107,113 +60,34 @@ function runExtensionSetup( const pipCacheDir = join(getSettings(userData).dependenciesDir, 'pip-cache') try { mkdirSync(pipCacheDir, { recursive: true }) } catch { /* pip creates it too */ } - const accelerator = process.platform === 'darwin' && process.arch === 'arm64' ? 'mps' : gpuSm > 0 ? 'cuda' : 'cpu' + const torchFlavor = torchFlavorFor(gpu.accelerator) const args = JSON.stringify({ python_exe: pythonExe, ext_dir: extDir, - gpu_sm: gpuSm, - cuda_version: cudaVersion, - accelerator, + gpu_sm: gpu.sm, + cuda_version: gpu.cudaVersion, + accelerator: gpu.accelerator, + // Extensions that know about AMD branch on torch_flavor (the official + // hunyuan3d-mini one does). Those that don't get corrected by the ROCm + // shim in setup-launcher.ts instead. + torch_flavor: torchFlavor, + gfx_target: gpu.gfxTarget ?? '', + torch_index_url: gpu.torchIndexUrl ?? '', platform: process.platform, arch: process.arch, }) - const launcher = ` -import runpy -import subprocess -import sys - -setup_py = sys.argv[1] -setup_args = sys.argv[2:] - -_original_run = subprocess.run -_original_check_call = subprocess.check_call -_original_check_output = subprocess.check_output - -def _is_cuda_torch_index(value): - return isinstance(value, str) and value.startswith("https://download.pytorch.org/whl/cu") - -def _mentions_torch(command): - if not isinstance(command, (list, tuple)): - return False - return any(str(part).startswith(("torch==", "torchvision==", "torchaudio==")) for part in command) - -def _rewrite_command(command): - if sys.platform != "darwin" or not _mentions_torch(command): - return command - if not isinstance(command, (list, tuple)): - return command - - rewritten = [] - changed = False - i = 0 - while i < len(command): - part = command[i] - text = str(part) - if text in ("--index-url", "-i", "--extra-index-url") and i + 1 < len(command) and _is_cuda_torch_index(str(command[i + 1])): - changed = True - i += 2 - continue - if text.startswith("--index-url=") or text.startswith("--extra-index-url="): - value = text.split("=", 1)[1] - if _is_cuda_torch_index(value): - changed = True - i += 1 - continue - rewritten.append(part) - i += 1 - - if changed: - print("[Modly setup compat] Removed CUDA-only PyTorch index on macOS; pip will use macOS wheels.", file=sys.stderr) - return rewritten - return command - -def _is_pip_command(command): - if not isinstance(command, (list, tuple)): - return False - return any("pip" in str(part).lower() for part in command[:3]) - -def _strip_no_cache(command): - # Extension setup scripts often hardcode --no-cache-dir, which forces pip to - # re-download multi-GB wheels on every retry. Modly provides a shared cache - # via PIP_CACHE_DIR, so drop the flag and let pip use it. - if not _is_pip_command(command): - return command - if not any(str(part) == "--no-cache-dir" for part in command): - return command - print("[Modly setup compat] Removed --no-cache-dir so pip reuses the shared wheel cache.", file=sys.stderr) - return [part for part in command if str(part) != "--no-cache-dir"] - -def _transform_command(command): - return _strip_no_cache(_rewrite_command(command)) - -def _patched_run(*args, **kwargs): - args = list(args) - if args: - args[0] = _transform_command(args[0]) - return _original_run(*args, **kwargs) - -def _patched_check_call(*args, **kwargs): - args = list(args) - if args: - args[0] = _transform_command(args[0]) - return _original_check_call(*args, **kwargs) - -def _patched_check_output(*args, **kwargs): - args = list(args) - if args: - args[0] = _transform_command(args[0]) - return _original_check_output(*args, **kwargs) - -subprocess.run = _patched_run -subprocess.check_call = _patched_check_call -subprocess.check_output = _patched_check_output - -sys.argv = [setup_py] + setup_args -runpy.run_path(setup_py, run_name="__main__") -` + const launcher = SETUP_LAUNCHER_SOURCE + // The rewrite decision itself is made in gpu-detect.ts (and unit-tested + // there); the launcher above only applies what these carry. const proc = spawn(pythonExe, ['-c', launcher, setupPy, args], { stdio: ['ignore', 'pipe', 'pipe'], - env: { ...process.env, PIP_CACHE_DIR: pipCacheDir }, + env: { + ...process.env, + PIP_CACHE_DIR: pipCacheDir, + MODLY_TORCH_FLAVOR: torchFlavor, + MODLY_TORCH_INDEX_URL: gpu.torchIndexUrl ?? '', + MODLY_TORCH_SPECS: JSON.stringify(gpu.torchSpecs ?? []), + }, }) const handleLine = (line: string) => { if (line) onLog?.(line) } @@ -1149,8 +1023,9 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // 7a. Python process extension: run setup.py if present (same as model extensions) if (existsSync(join(destDir, 'setup.py'))) { emit({ step: 'setting_up', message: 'Setting up Python environment…' }) - const { sm: gpuSm, cudaVersion } = await detectGpuInfo() - await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { + const gpu = await detectGpuInfo({ onLog: (line) => logger.info(line) }) + logger.info(`[ext-setup] ${describeGpuInfo(gpu)}`) + await runExtensionSetup(destDir, gpu, (line) => { logger.info(`[ext-setup] ${line}`) emit({ step: 'setting_up', message: line }) }) @@ -1185,8 +1060,9 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // 7c. Model extension: run setup.py directly (no FastAPI required) if (existsSync(join(destDir, 'setup.py'))) { emit({ step: 'setting_up', message: 'Setting up Python environment…' }) - const { sm: gpuSm, cudaVersion } = await detectGpuInfo() - await runExtensionSetup(destDir, gpuSm, cudaVersion, (line) => { + const gpu = await detectGpuInfo({ onLog: (line) => logger.info(line) }) + logger.info(`[ext-setup] ${describeGpuInfo(gpu)}`) + await runExtensionSetup(destDir, gpu, (line) => { logger.info(`[ext-setup] ${line}`) emit({ step: 'setting_up', message: line }) }) @@ -1279,8 +1155,9 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (!existsSync(join(extDir, 'setup.py'))) { return { success: false, error: 'setup.py is missing from the extension folder — the install looks incomplete. Uninstall the extension and install it again.' } } - const { sm: gpuSm, cudaVersion } = await detectGpuInfo() - await runExtensionSetup(extDir, gpuSm, cudaVersion, (line) => logger.info(`[ext-repair] ${line}`)) + const gpu = await detectGpuInfo({ onLog: (line) => logger.info(line) }) + logger.info(`[ext-repair] ${describeGpuInfo(gpu)}`) + await runExtensionSetup(extDir, gpu, (line) => logger.info(`[ext-repair] ${line}`)) try { await axios.post(`${API_BASE_URL}/extensions/reload`, {}, { timeout: 10_000 }) } catch { /* ignore if Python is not running yet */ } diff --git a/electron/main/python-setup.ts b/electron/main/python-setup.ts index 3b85db1b..13e86a25 100644 --- a/electron/main/python-setup.ts +++ b/electron/main/python-setup.ts @@ -1,10 +1,11 @@ import { BrowserWindow, app } from 'electron' import { existsSync, readFileSync, writeFileSync } from 'fs' -import { cp, rm, mkdir } from 'fs/promises' +import { rm, mkdir } from 'fs/promises' import { join } from 'path' import { spawn, execSync } from 'child_process' import { createHash } from 'crypto' import { getSettings } from './settings-store' +import { copyRuntimeTree } from './copy-runtime' const SETUP_VERSION = 3 @@ -176,7 +177,7 @@ async function ensureStableEmbeddedPython(userData: string, win: BrowserWindow): await rm(stableDir, { recursive: true, force: true }) } await mkdir(stableDir, { recursive: true }) - await cp(getEmbeddedPythonDir(), stableDir, { recursive: true, preserveTimestamps: true }) + await copyRuntimeTree(getEmbeddedPythonDir(), stableDir) writeFileSync(versionFile, currentVersion, 'utf-8') console.log('[PythonSetup] Python runtime ready at:', stableDir) } diff --git a/electron/main/setup-launcher.test.mjs b/electron/main/setup-launcher.test.mjs new file mode 100644 index 00000000..c4a40524 --- /dev/null +++ b/electron/main/setup-launcher.test.mjs @@ -0,0 +1,224 @@ +/** + * Runs the setup launcher for real, with a stub standing in for pip, and checks + * what the extension's pip invocation was rewritten into. + * + * The commands exercised below are copied from the official extensions' + * setup.py, so a change that breaks AMD installs fails here rather than after a + * multi-gigabyte download on a user's machine. + */ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, writeFileSync, readFileSync, existsSync } from 'node:fs' +import { spawnSync } from 'node:child_process' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-launcher-test-')), 'setup-launcher.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/main/setup-launcher.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +const { SETUP_LAUNCHER_SOURCE } = loadModule() + +function findPython() { + for (const candidate of ['python3', 'python']) { + const probe = spawnSync(candidate, ['--version'], { stdio: 'ignore' }) + if (probe.status === 0) return candidate + } + return null +} + +const PYTHON = findPython() + +// The launcher's _is_pip_command looks for "pip" in the first three tokens, so +// the stub is named pip.py — matching how extensions invoke "/bin/pip". +const FAKE_PIP = ` +import json, sys +with open(sys.argv[1], "w") as handle: + json.dump(sys.argv[2:], handle) +` + +/** + * Runs a pip command through the launcher and returns the argv the stub + * actually received, i.e. the command after every rewrite. + */ +function runThroughLauncher(pipArgs, env = {}) { + const dir = mkdtempSync(join(tmpdir(), 'modly-launcher-run-')) + const fakePip = join(dir, 'pip.py') + const capture = join(dir, 'captured.json') + writeFileSync(fakePip, FAKE_PIP, 'utf8') + + // Stands in for the extension's setup.py: issues one pip call, exactly as the + // real ones do, and lets the launcher's patched subprocess handle it. + const setupPy = join(dir, 'setup.py') + writeFileSync(setupPy, [ + 'import subprocess, sys', + `PIP = ${JSON.stringify([fakePip, capture])}`, + `subprocess.run([sys.executable] + PIP + ${JSON.stringify(pipArgs)}, check=True)`, + ].join('\n'), 'utf8') + + const result = spawnSync(PYTHON, ['-c', SETUP_LAUNCHER_SOURCE, setupPy, '{}'], { + encoding: 'utf8', + env: { ...process.env, ...env }, + }) + assert.equal(result.status, 0, `launcher failed:\n${result.stderr}`) + assert.ok(existsSync(capture), `pip stub was never invoked:\n${result.stderr}`) + + // Drop the stub's own two arguments; keep the pip command itself. + return { argv: JSON.parse(readFileSync(capture, 'utf8')), stderr: result.stderr } +} + +const ROCM_ENV = { + MODLY_TORCH_FLAVOR: 'rocm', + MODLY_TORCH_INDEX_URL: 'https://download.pytorch.org/whl/rocm7.2', + MODLY_TORCH_SPECS: JSON.stringify(['torch', 'torchvision']), +} + +test('ROCm shim redirects a CUDA-pinned install (triposg / trellis2 shape)', { skip: !PYTHON }, () => { + const { argv, stderr } = runThroughLauncher( + ['install', 'torch==2.6.0', 'torchvision==0.21.0', '--index-url', 'https://download.pytorch.org/whl/cu124'], + ROCM_ENV, + ) + + assert.deepEqual(argv, [ + 'install', + '--index-url', 'https://download.pytorch.org/whl/rocm7.2', + 'torch', 'torchvision', + ]) + assert.match(stderr, /Redirected PyTorch to ROCm wheels/) +}) + +test('ROCm shim rescues the CPU fallback hunyuan3d-mini forces on Windows', { skip: !PYTHON }, () => { + // That setup.py hardcodes the CPU index when torch_flavor is rocm on Windows; + // without this rewrite an AMD Windows user would silently get CPU-only torch. + const { argv } = runThroughLauncher( + ['install', 'torch==2.6.0', 'torchvision==0.21.0', '--index-url', 'https://download.pytorch.org/whl/cpu'], + { + ...ROCM_ENV, + MODLY_TORCH_INDEX_URL: 'https://repo.amd.com/rocm/whl-multi-arch/', + MODLY_TORCH_SPECS: JSON.stringify([ + 'torch[device-gfx1200]==2.11.0+rocm7.14.0', + 'torchvision[device-gfx1200]==0.26.0+rocm7.14.0', + ]), + }, + ) + + assert.deepEqual(argv, [ + 'install', + '--index-url', 'https://repo.amd.com/rocm/whl-multi-arch/', + 'torch[device-gfx1200]==2.11.0+rocm7.14.0', + 'torchvision[device-gfx1200]==0.26.0+rocm7.14.0', + ]) +}) + +test('ROCm shim leaves an extension that already chose ROCm alone', { skip: !PYTHON }, () => { + // hunyuan3d-mini's own rocm branch on Linux. Its index is kept; only the + // requirements are normalised to what Modly resolved. + const { argv } = runThroughLauncher( + ['install', 'torch', 'torchvision', '--index-url', 'https://download.pytorch.org/whl/rocm7.2'], + ROCM_ENV, + ) + + assert.deepEqual(argv, [ + 'install', + 'torch', 'torchvision', + '--index-url', 'https://download.pytorch.org/whl/rocm7.2', + ]) +}) + +test('ROCm shim replaces the pinned direct wheel URLs of the ARM64 path', { skip: !PYTHON }, () => { + const { argv } = runThroughLauncher( + [ + 'install', '--retries', '10', + '--extra-index-url', 'https://download.pytorch.org/whl/cu128', + 'https://download-r2.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp311-cp311-manylinux_2_28_aarch64.whl', + 'https://download-r2.pytorch.org/whl/cu128/torchvision-0.22.0-cp311-cp311-manylinux_2_28_aarch64.whl', + ], + ROCM_ENV, + ) + + assert.deepEqual(argv, [ + 'install', '--retries', '10', + '--index-url', 'https://download.pytorch.org/whl/rocm7.2', + 'torch', 'torchvision', + ]) +}) + +test('ROCm shim handles the --index-url=VALUE spelling', { skip: !PYTHON }, () => { + const { argv } = runThroughLauncher( + ['install', '--index-url=https://download.pytorch.org/whl/cu124', 'torch==2.6.0'], + ROCM_ENV, + ) + + assert.deepEqual(argv, [ + 'install', + '--index-url', 'https://download.pytorch.org/whl/rocm7.2', + 'torch', 'torchvision', + ]) +}) + +test('ROCm shim leaves non-torch installs untouched', { skip: !PYTHON }, () => { + // The bulk of every setup.py: core deps, rembg, etc. Nothing here is ours to + // rewrite, and an extra --index-url would send them to the ROCm index. + const original = ['install', 'Pillow', 'numpy', 'trimesh', 'rembg', 'onnxruntime'] + const { argv } = runThroughLauncher(original, ROCM_ENV) + assert.deepEqual(argv, original) +}) + +test('ROCm shim keeps PyPI reachable when torch is mixed with other packages', { skip: !PYTHON }, () => { + // The ROCm index mirrors torch's dependency closure only: asking it for + // trimesh returns 403, so a mixed install would hard-fail without this. + const { argv } = runThroughLauncher( + ['install', 'torch==2.6.0', 'trimesh', 'diffusers', '--index-url', 'https://download.pytorch.org/whl/cu124'], + ROCM_ENV, + ) + + assert.deepEqual(argv, [ + 'install', + '--index-url', 'https://download.pytorch.org/whl/rocm7.2', + '--extra-index-url', 'https://pypi.org/simple', + 'torch', 'torchvision', + 'trimesh', 'diffusers', + ]) +}) + +test('ROCm shim does not add PyPI for a torch-only install', { skip: !PYTHON }, () => { + // Keeping PyPI out of a pure torch install avoids pip ever preferring a + // plain CUDA wheel over the ROCm one. + const { argv } = runThroughLauncher( + ['install', '--retries', '10', 'torch==2.6.0', '--index-url', 'https://download.pytorch.org/whl/cu124'], + ROCM_ENV, + ) + assert.ok(!argv.includes('--extra-index-url')) +}) + +test('shim is inert on a CUDA machine', { skip: !PYTHON }, () => { + const original = ['install', 'torch==2.6.0', '--index-url', 'https://download.pytorch.org/whl/cu124'] + const { argv } = runThroughLauncher(original, { + MODLY_TORCH_FLAVOR: 'cuda', + MODLY_TORCH_INDEX_URL: '', + MODLY_TORCH_SPECS: '[]', + }) + assert.deepEqual(argv, original) +}) + +test('--no-cache-dir is still stripped alongside the ROCm rewrite', { skip: !PYTHON }, () => { + const { argv } = runThroughLauncher( + ['install', '--no-cache-dir', 'torch==2.6.0', '--index-url', 'https://download.pytorch.org/whl/cu124'], + ROCM_ENV, + ) + + assert.ok(!argv.includes('--no-cache-dir')) + assert.ok(argv.includes('https://download.pytorch.org/whl/rocm7.2')) +}) diff --git a/electron/main/setup-launcher.ts b/electron/main/setup-launcher.ts new file mode 100644 index 00000000..4207e4a6 --- /dev/null +++ b/electron/main/setup-launcher.ts @@ -0,0 +1,231 @@ +/** + * Python launcher used to run an extension's setup.py. + * + * Extension setup scripts are third-party code we cannot edit, and they install + * PyTorch themselves from an index they hardcode. The launcher wraps them: it + * patches subprocess so every pip invocation passes through a few corrections + * before it runs — dropping CUDA-only indexes on macOS, keeping the shared wheel + * cache alive, and redirecting torch to ROCm wheels on AMD machines. + * + * Kept in its own module so setup-launcher.test.mjs can execute it for real + * against the command shapes the official extensions actually use. + */ + +export const SETUP_LAUNCHER_SOURCE = ` +import json +import os +import re +import runpy +import subprocess +import sys + +setup_py = sys.argv[1] +setup_args = sys.argv[2:] + +_original_run = subprocess.run +_original_check_call = subprocess.check_call +_original_check_output = subprocess.check_output + +_TORCH_REQ_RE = re.compile(r"^(torch|torchvision|torchaudio)(\\[[^\\]]*\\])?\\s*([<>=!~].*)?$", re.I) + +def _is_cuda_torch_index(value): + return isinstance(value, str) and value.startswith("https://download.pytorch.org/whl/cu") + +def _is_torch_requirement(text): + # Matches "torch", "torch==2.6.0", "torch[device-gfx1200]==2.11.0+rocm7.14.0", + # and the pinned direct wheel URLs the ARM64 install path uses. + if _TORCH_REQ_RE.match(text): + return True + if text.startswith(("http://", "https://")): + return any(seg in text for seg in ("/torch-", "/torchvision-", "/torchaudio-")) + return False + +def _mentions_torch(command): + if not isinstance(command, (list, tuple)): + return False + return any(_is_torch_requirement(str(part)) for part in command) + +def _rewrite_command(command): + if sys.platform != "darwin" or not _mentions_torch(command): + return command + if not isinstance(command, (list, tuple)): + return command + + rewritten = [] + changed = False + i = 0 + while i < len(command): + part = command[i] + text = str(part) + if text in ("--index-url", "-i", "--extra-index-url") and i + 1 < len(command) and _is_cuda_torch_index(str(command[i + 1])): + changed = True + i += 2 + continue + if text.startswith("--index-url=") or text.startswith("--extra-index-url="): + value = text.split("=", 1)[1] + if _is_cuda_torch_index(value): + changed = True + i += 1 + continue + rewritten.append(part) + i += 1 + + if changed: + print("[Modly setup compat] Removed CUDA-only PyTorch index on macOS; pip will use macOS wheels.", file=sys.stderr) + return rewritten + return command + +def _is_pip_command(command): + if not isinstance(command, (list, tuple)): + return False + return any("pip" in str(part).lower() for part in command[:3]) + +def _strip_no_cache(command): + # Extension setup scripts often hardcode --no-cache-dir, which forces pip to + # re-download multi-GB wheels on every retry. Modly provides a shared cache + # via PIP_CACHE_DIR, so drop the flag and let pip use it. + if not _is_pip_command(command): + return command + if not any(str(part) == "--no-cache-dir" for part in command): + return command + print("[Modly setup compat] Removed --no-cache-dir so pip reuses the shared wheel cache.", file=sys.stderr) + return [part for part in command if str(part) != "--no-cache-dir"] + +# ROCm redirect. Most extension setup.py scripts predate AMD support and +# hardcode a CUDA index (hunyuan3d-mini even forces the CPU index on Windows), +# so on an AMD machine we swap the whole torch install for the ROCm one Modly +# resolved. An index the extension already pointed at ROCm is left alone. +_ROCM_INDEX = os.environ.get("MODLY_TORCH_INDEX_URL", "") +try: + _ROCM_SPECS = json.loads(os.environ.get("MODLY_TORCH_SPECS", "[]")) +except ValueError: + _ROCM_SPECS = [] + +def _is_pytorch_index(value): + return isinstance(value, str) and "download.pytorch.org/whl/" in value + +def _is_rocm_index(value): + return isinstance(value, str) and "/whl/rocm" in value + +_PIP_VALUE_FLAGS = ( + "--index-url", "-i", "--extra-index-url", "--find-links", "-f", + "--retries", "--timeout", "--cache-dir", "--target", "-t", + "--requirement", "-r", "--constraint", "-c", "--progress-bar", + "--proxy", "--cert", "--client-cert", "--trusted-host", "--log", + "--no-binary", "--only-binary", "--prefix", "--root", "--upgrade-strategy", + "--python-version", "--platform", "--abi", "--implementation", +) + +def _has_non_torch_requirement(command): + # Only look past the subcommand, so the interpreter and pip executable + # paths ahead of it are never mistaken for requirements. + texts = [str(part) for part in command] + start = None + for index, text in enumerate(texts): + if text in ("install", "download", "wheel"): + start = index + 1 + break + if start is None: + return False + + skip_next = False + for text in texts[start:]: + if skip_next: + skip_next = False + continue + if text in _PIP_VALUE_FLAGS: + skip_next = True + continue + if text.startswith("-") or _is_torch_requirement(text): + continue + return True + return False + +def _rewrite_rocm(command): + if os.environ.get("MODLY_TORCH_FLAVOR") != "rocm" or not _ROCM_INDEX or not _ROCM_SPECS: + return command + if not isinstance(command, (list, tuple)): + return command + if not _is_pip_command(command) or not _mentions_torch(command): + return command + + rewritten = [] + insert_at = None + keeps_rocm_index = False + changed = False + i = 0 + while i < len(command): + text = str(command[i]) + if text in ("--index-url", "-i", "--extra-index-url") and i + 1 < len(command): + value = str(command[i + 1]) + if _is_rocm_index(value): + keeps_rocm_index = True + elif _is_pytorch_index(value): + changed = True + i += 2 + continue + rewritten.extend(command[i:i + 2]) + i += 2 + continue + if text.split("=", 1)[0] in ("--index-url", "--extra-index-url") and "=" in text: + value = text.split("=", 1)[1] + if _is_rocm_index(value): + keeps_rocm_index = True + elif _is_pytorch_index(value): + changed = True + i += 1 + continue + elif _is_torch_requirement(text): + if insert_at is None: + insert_at = len(rewritten) + changed = True + i += 1 + continue + rewritten.append(command[i]) + i += 1 + + if not changed: + return command + if insert_at is None: + insert_at = len(rewritten) + injected = list(_ROCM_SPECS) + if not keeps_rocm_index: + index_args = ["--index-url", _ROCM_INDEX] + if _has_non_torch_requirement(command): + # The ROCm index only mirrors torch's own dependency closure + # (numpy, pillow…), not application packages like trimesh or + # diffusers. A pip call that mixes both still needs PyPI reachable. + index_args += ["--extra-index-url", "https://pypi.org/simple"] + injected = index_args + injected + rewritten[insert_at:insert_at] = injected + print("[Modly setup compat] Redirected PyTorch to ROCm wheels: " + " ".join(injected), file=sys.stderr) + return rewritten + +def _transform_command(command): + return _strip_no_cache(_rewrite_rocm(_rewrite_command(command))) + +def _patched_run(*args, **kwargs): + args = list(args) + if args: + args[0] = _transform_command(args[0]) + return _original_run(*args, **kwargs) + +def _patched_check_call(*args, **kwargs): + args = list(args) + if args: + args[0] = _transform_command(args[0]) + return _original_check_call(*args, **kwargs) + +def _patched_check_output(*args, **kwargs): + args = list(args) + if args: + args[0] = _transform_command(args[0]) + return _original_check_output(*args, **kwargs) + +subprocess.run = _patched_run +subprocess.check_call = _patched_check_call +subprocess.check_output = _patched_check_output + +sys.argv = [setup_py] + setup_args +runpy.run_path(setup_py, run_name="__main__") +`