diff --git a/docs/ACCELERATION.md b/docs/ACCELERATION.md index 82a3b220..ed1450f2 100644 --- a/docs/ACCELERATION.md +++ b/docs/ACCELERATION.md @@ -186,7 +186,7 @@ Compile a component once it is fully configured - the graph captures final dtype Typical gains are 1.3-1.5x on diffusion transformers, and compilation stacks with the caches above. Notes: -- **First run pays the compile cost.** The [server](SERVER.md)'s persistent worker keeps compiled pipelines loaded between runs, so the cost is paid once per session rather than once per generation. +- **First run pays the compile cost.** The [server](SERVER.md)'s persistent worker keeps compiled pipelines loaded between runs, so the cost is paid once per session rather than once per generation. A new session still retraces, but the kernels it built come back from the on-disk cache, which dw keeps across reboots (`TORCHINDUCTOR_CACHE_DIR`, see [Environment Defaults](#environment-defaults)). - **Pin the attention backend** on a compiled component (`"attention_backend"` in the same `components` entry) rather than using the pipeline-level per-call context manager, which forces recompiles. - **Composes with offloading**: apply `group_offload` and `compile` on the same component and the offload hooks are installed first, as required. Skipped with a warning on MPS. - **Don't combine `fullgraph` with a `cache`**: the cache hooks decide skip-or-compute per step, a data-dependent branch diffusers wraps in `torch.compiler.disable` - it needs the graph break that `fullgraph: true` forbids. Compile with the default (partial) graph mode when a cache is active. @@ -279,6 +279,7 @@ Set automatically at import unless already present in the environment (export yo | `PYTORCH_CUDA_ALLOC_CONF` | `expandable_segments:True` | Lets the CUDA allocator grow segments instead of fragmenting fixed-size ones. Multi-step workflows churn differently-shaped allocations (generate, upscale, interpolate); fragmentation is what OOMs a card that nominally has room. | | `HF_ENABLE_PARALLEL_LOADING` | `true` | Loads sharded checkpoints in parallel - faster cold starts. | | `PYTORCH_MPS_HIGH_WATERMARK_RATIO` | `0.0` | MPS only - use all available unified memory. | +| `TORCHINDUCTOR_CACHE_DIR` | `$XDG_CACHE_HOME/dw/torchinductor` (`~/.cache/dw/torchinductor`) | Keeps `torch.compile`'s on-disk cache across reboots - Inductor's own default is under `/tmp`. A new worker still retraces, but a cache hit skips codegen and kernel builds, the bulk of a cold compile. Entries are per shape, so a new size or frame count compiles once more. | For faster model downloads, optionally `pip install hf_transfer` and set `HF_HUB_ENABLE_HF_TRANSFER=1`. Not enabled automatically - it bypasses the Python HTTP stack and breaks some proxy setups. diff --git a/dw/__init__.py b/dw/__init__.py index 3a76faf5..3337b5ac 100644 --- a/dw/__init__.py +++ b/dw/__init__.py @@ -48,6 +48,24 @@ def _parallel_loading_default(platform): if "HF_ENABLE_PARALLEL_LOADING" not in os.environ: os.environ["HF_ENABLE_PARALLEL_LOADING"] = _parallel_loading_default(sys.platform) + +def _inductor_cache_dir_default(environ): + """Where torch.compile keeps its on-disk cache. + + Inductor's own default is /tmp/torchinductor_, which a reboot wipes. + A fresh worker always retraces, but a cache hit skips codegen and kernel + builds - most of a cold compile, about 90s of MiniMax-H3's first step on a + 3090 - so the user's cache directory keeps that paid once per shape rather + than once per boot. An empty XDG_CACHE_HOME counts as unset, per the spec.""" + base = environ.get("XDG_CACHE_HOME") or os.path.join( + os.path.expanduser("~"), ".cache" + ) + return os.path.join(base, "dw", "torchinductor") + + +if "TORCHINDUCTOR_CACHE_DIR" not in os.environ: + os.environ["TORCHINDUCTOR_CACHE_DIR"] = _inductor_cache_dir_default(os.environ) + # Suppress all common library warnings before any imports warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) diff --git a/tests/test_inductor_cache_default.py b/tests/test_inductor_cache_default.py new file mode 100644 index 00000000..79a84071 --- /dev/null +++ b/tests/test_inductor_cache_default.py @@ -0,0 +1,58 @@ +"""torch.compile's on-disk cache survives a reboot by default. + +Inductor caches compiled graphs and generated kernels under +/tmp/torchinductor_ unless TORCHINDUCTOR_CACHE_DIR says otherwise, and +/tmp is wiped on reboot. A fresh worker still has to trace, but a cache hit +skips codegen and kernel builds - the bulk of a cold compile, about 90s of +MiniMax-H3's first step on a 3090. dw points the cache at the user's cache +directory instead, so that cost is paid once per shape rather than once per +boot. +""" + +import os +import subprocess +import sys + +import dw + + +def test_follows_xdg_cache_home(): + environ = {"XDG_CACHE_HOME": os.path.join("/srv", "cache")} + assert dw._inductor_cache_dir_default(environ) == os.path.join( + "/srv", "cache", "dw", "torchinductor" + ) + + +def test_falls_back_to_home_cache(): + expected = os.path.join(os.path.expanduser("~"), ".cache", "dw", "torchinductor") + assert dw._inductor_cache_dir_default({}) == expected + + +def test_an_empty_xdg_cache_home_is_ignored(): + # The XDG spec treats an empty value as unset + expected = os.path.join(os.path.expanduser("~"), ".cache", "dw", "torchinductor") + assert dw._inductor_cache_dir_default({"XDG_CACHE_HOME": ""}) == expected + + +def _env_after_import(preset=None): + code = "import os, dw; print(os.environ['TORCHINDUCTOR_CACHE_DIR'])" + env = {k: v for k, v in os.environ.items() if k != "TORCHINDUCTOR_CACHE_DIR"} + if preset is not None: + env["TORCHINDUCTOR_CACHE_DIR"] = preset + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env=env, + timeout=120, + ) + assert result.returncode == 0, result.stderr + return result.stdout.strip().splitlines()[-1] + + +def test_import_applies_the_default(): + assert _env_after_import() == dw._inductor_cache_dir_default(os.environ) + + +def test_an_explicit_setting_wins(): + assert _env_after_import(preset="/var/tmp/inductor") == "/var/tmp/inductor" diff --git a/ui/CLAUDE.md b/ui/CLAUDE.md index bc8c6ff8..0090ac5c 100644 --- a/ui/CLAUDE.md +++ b/ui/CLAUDE.md @@ -8,4 +8,4 @@ Checks, from `ui/`: `npm run check`, `npm run lint`, `npm test`, and `npx playwr - Show the proof: a list of things that produce images shows the images. - Every token pair should pass WCAG AA in both themes. -- `--live` marks machine state and nothing else; where it may appear is held by `scripts/design-rules.test.ts`. +- `--live` marks machine state and `--select` the user's place (focus, the open workspace, the active section), each nothing else; where they may appear is held by `scripts/design-rules.test.ts`. diff --git a/ui/scripts/design-rules.test.ts b/ui/scripts/design-rules.test.ts index 91b483b4..21791a44 100644 --- a/ui/scripts/design-rules.test.ts +++ b/ui/scripts/design-rules.test.ts @@ -1,6 +1,6 @@ // @vitest-environment node -// Where the --live state colour may appear, held mechanically (ui/CLAUDE.md; -// app.css's header says why --live means machine state and nothing else) +// Where the --live state colour and the --select place colour may appear, +// held mechanically (ui/CLAUDE.md; app.css's header says what each means) import { readFileSync } from 'node:fs' import { dirname, join } from 'node:path' import { fileURLToPath } from 'node:url' @@ -11,8 +11,8 @@ const UI = join(dirname(fileURLToPath(import.meta.url)), '..') // The places --live may appear: the header's running link and VRAM // pressure, a running job's row and status chip, the job page's progress -// bar and current step, the flow view's active step, a model download in -// flight, and the focus ring +// bar and current step, the flow view's active step, and a model download +// in flight const LIVE_FILES = [ 'src/App.svelte', 'src/app.css', @@ -29,3 +29,15 @@ it('uses --live only where machine state is shown', () => { .map((f) => f.slice(UI.length + 1).replaceAll('\\', '/')) expect(using).toEqual(LIVE_FILES) }) + +// The places --select may appear: the focus ring and the sidebar's open +// workspace and active section. It is chrome only - never a surface behind +// an image +const SELECT_FILES = ['src/app.css', 'src/lib/Sidebar.svelte'] + +it('uses --select only for the focus ring and the nav', () => { + const using = sourceFiles() + .filter((f) => readFileSync(f, 'utf8').includes('var(--select)')) + .map((f) => f.slice(UI.length + 1).replaceAll('\\', '/')) + expect(using).toEqual(SELECT_FILES) +}) diff --git a/ui/src/App.svelte b/ui/src/App.svelte index 656b2adb..04709092 100644 --- a/ui/src/App.svelte +++ b/ui/src/App.svelte @@ -1,6 +1,6 @@