diff --git a/docs/ACCELERATION.md b/docs/ACCELERATION.md index ebfb3f24..d6c12192 100644 --- a/docs/ACCELERATION.md +++ b/docs/ACCELERATION.md @@ -198,8 +198,9 @@ 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. 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)). +- **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)). SDNQ's `triton.jit`/autotune kernels are cached separately by Triton itself (`TRITON_CACHE_DIR`, else `$TRITON_HOME/.triton/cache`, else `~/.triton/cache`); dw sets no `TRITON_*` variable and leaves that cache where it is, shared with any other Triton user on the machine. Both caches' sizes show in `get_server_info`'s `caches` and `GET /api/caches`, and `DELETE /api/caches/{name}` empties one ([SERVER.md](SERVER.md)); the next first step of each shape pays the compile or re-autotune again (100-300 s per shape on MiniMax-H3 for Triton). - **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. +- **Quantize-on-load time is a separate cost.** An SDNQ on-the-fly component re-reads and re-quantizes its bf16 weights on every cold load; `DW_QUANTIZED_CACHE=1` saves the quantized result and reads it back instead - see [QUANTIZATION.md](QUANTIZATION.md#cached-on-the-fly-dw_quantized_cache). - **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. - **TorchAO quantization needs compile to be fast** - see [QUANTIZATION.md](QUANTIZATION.md#torchao). @@ -295,7 +296,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. | +| `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. Listed as `inductor` by `GET /api/caches` and cleared with `DELETE /api/caches/inductor`. Triton's cache is not set by dw and stays at its own default (`~/.triton/cache`) - see the compile notes above. | 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/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 06e3bae0..e292d907 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -28,6 +28,7 @@ one to open. | Device translation | `dw/pipeline_processors/placement.py`: `place_component` | A device is translated by `resolve_device()` before anything reads its backend, so the MPS accommodations also fire for a translated device. | `tests/test_device_portability.py::TestPlacement::test_a_translated_device_gets_its_backends_offload_downgrade` | | On-demand wrappers | `dw/pipeline_processors/placement.py`: `apply_on_demand_placement` | The on-demand wrappers keep the wrapped signature, because H3's denoiser reads `signature(transformer.forward)`. | `tests/test_modular_pipeline.py::TestOnDemandResidency::test_the_wrapped_signature_survives` | | On-demand and offload | `dw/pipeline_processors/placement.py`: `place_component`, `has_component_group_offload` | On-demand residency and `group_offload` cannot both be set on one component, and either one stops the whole pipeline from being moved to the device. | `tests/test_modular_pipeline.py::TestOnDemandResidency::test_group_offload_and_on_demand_together_are_rejected`, `tests/test_modular_pipeline.py::TestOffloadDevice::test_component_group_offload_is_not_moved_to_device` | +| Quantized component cache | `dw/quantized_cache.py`: `load_components`, `from_pretrained`, `cache_key` | A key that can't pin the weights isn't cached, the cache never fails a job, and the save comes before pruning, so an entry is the whole component. | `tests/test_quantized_cache.py` | | LoRAs | `dw/pipeline_processors/adapters.py`: `active_loras`, `load_loras` | A `loras` entry whose `model_name` is null is switched off. | `tests/test_lora_disable.py::TestLoadLoras::test_an_all_null_entry_loads_nothing` | | Step progress | `dw/pipeline_processors/progress.py`: `reported_progress_bars` | A pipeline with no step callback reports its denoise steps through its progress bars. | `tests/test_modular_progress.py::test_every_advance_of_the_bar_is_reported` | | Community pipelines | `dw/community_pipelines/` (e.g. `pipeline_ltx2_refine.py`: `LTX2RefinePipeline`), `dw/introspection.py`: `community_pipelines`, `load_allowed_class` | A pipeline dw ships is named by its dotted path (`dw.community_pipelines..`), loads untrusted because it is under `dw`, and is listed by `list_pipelines` and resolved by `get_pipeline_signature` - only the modules in that directory, never a caller-chosen import. `LTX2RefinePipeline` VAE-encodes `video` at `width` x `height` and passes it unnormalized as `latents`, because `LTX2Pipeline.prepare_latents` normalizes 5-D latents itself. | `tests/test_ltx2_refine_pipeline.py::TestIntrospection::test_only_the_shipped_modules_resolve`, `tests/test_ltx2_refine_pipeline.py::TestEncode::test_the_latents_are_left_unnormalized_for_the_parent` | @@ -125,7 +126,7 @@ one to open. | Opening media | `dw/media.py` | This is the only module that calls `av.open`. It imports no torch and no step types. | `tests/test_media_layering.py::test_av_open_appears_only_in_media`, `tests/test_media_layering.py::test_media_imports_no_step_types` (no torch: —) | | Signal processing | `dw/dsp.py` | It measures and transforms waveforms, decides nothing, and imports nothing from `dw`. | `tests/test_media_layering.py::test_dsp_imports_nothing_from_dw` | | Assessment probes | `dw/tasks/assess.py`, `dw/assessment_rules.py` | Exactly three registered commands are assessment probes. Each measures a finished file, and every rule names a real probe field. | `tests/test_assessment_rules.py::TestRuleProbesAreRealCommands::test_assessment_is_exactly_the_three_probes`, `tests/test_assessment_rules.py::TestRuleProbesAreRealCommands::test_every_rule_names_a_registered_json_command` | -| Script check | `dw/tasks/script_check.py`: `check`; `dw/script_lines.py`: `parse_lines`, `parse_shots`, `shot_names_error` | `check_script` aligns a take's heard words to its expected lines and guards out words heard over silence and repetition loops (Whisper's hallucination over music). It builds its findings with `assessment_rules.finding()` and reads the guard's floor and window (`DEAD_AIR_FLOOR_DBFS`, `DEAD_AIR_WINDOW`) rather than copying them. It is not an assessment probe and adds nothing to `RULES`. Its literal `lines` is refused at validation through `task_problems` (`script_lines_errors`), which calls the same `parse_lines` the task runs, and a literal `shots` through the same `parse_shots` and `shot_names_error`; the parsers live in `dw/script_lines.py`, not the task, so that `task_problems` reaches them without importing `script_check` (the `import_cycles` ratchet in `scripts/arch_metrics.py`). Its shot map comes from `assess.resolve_shots`, the one the probes use, never a copy; it places each shot on the soundtrack with `assess.sample_span`, decides repeated shot names with `shots.duplicate_shot_names`, and asks `locations.is_http_url` whether the take is a URL. | `tests/test_script_check.py::TestRegistration::test_registered_as_json_command_without_assessment_flag`, `tests/test_script_check.py::TestEmptyLines::test_words_only_over_silence_are_discarded_not_found`, `tests/test_script_check.py::TestLinesErrors::test_literal_with_a_colon_is_not_a_reference` | +| Script check | `dw/tasks/script_check.py`: `check`; `dw/script_lines.py`: `parse_lines`, `parse_shots`, `shot_names_error` | `check_script` aligns a take's heard words to its expected lines and guards out words heard over silence and repetition loops (Whisper's hallucination over music). It builds its findings with `assessment_rules.finding()` and reads the guard's floor and window (`DEAD_AIR_FLOOR_DBFS`, `DEAD_AIR_WINDOW`) rather than copying them; the loudest-window measure is `dsp.loudest_window_dbfs` (`script_check.word_level` wraps it), which `transcribe_audio`'s `suspect: "silent"` mark reads with the same floor and window; the repetition-loop detector (`script_lines.looped_texts`, `REPEAT_RUN_MIN`, `REPEAT_MAX_PERIOD`) lives in `dw/script_lines.py` and backs both `check_script`'s `repetition` discard and `transcribe_audio`'s `suspect: "repetition"`, and `script_lines.suspect_runs`/`SUSPECT_LEFT_OUT` are what `plan_cuts` groups and leaves out suspect chunks by - so neither task imports the other. It is not an assessment probe and adds nothing to `RULES`. Its literal `lines` is refused at validation through `task_problems` (`script_lines_errors`), which calls the same `parse_lines` the task runs, and a literal `shots` through the same `parse_shots` and `shot_names_error`; the parsers live in `dw/script_lines.py`, not the task, so that `task_problems` reaches them without importing `script_check` (the `import_cycles` ratchet in `scripts/arch_metrics.py`). Its shot map comes from `assess.resolve_shots`, the one the probes use, never a copy; it places each shot on the soundtrack with `assess.sample_span`, decides repeated shot names with `shots.duplicate_shot_names`, and asks `locations.is_http_url` whether the take is a URL. | `tests/test_script_check.py::TestRegistration::test_registered_as_json_command_without_assessment_flag`, `tests/test_script_check.py::TestEmptyLines::test_words_only_over_silence_are_discarded_not_found`, `tests/test_script_check.py::TestLinesErrors::test_literal_with_a_colon_is_not_a_reference` | | Audio sync against a source slice | `dw/tasks/measure_sync.py`: `measure_sync`; `dw/dsp.py`: `envelope_lag` | `measure_sync` cross-correlates the onset envelopes of a shot's audio and its source slice and answers the lag and a 0-1 confidence. A positive lag means the audio is LATE (its events fall after the reference's), negative early; `lag_seconds` is None for a silent or flat track. It builds its findings (`sync_unmeasurable`, `sync_low_confidence`, `audio_out_of_sync`) with `assessment_rules.finding()` over module-local rule dicts and adds nothing to `RULES`; it reads the silence floor from `beats.SILENT_DBFS`. `envelope_lag` imports nothing from `dw`. | `tests/test_measure_sync.py::TestLag::test_late_audio_is_a_positive_lag`, `tests/test_measure_sync.py::TestLag::test_early_audio_is_a_negative_lag`, `tests/test_measure_sync.py::TestUnmeasurable::test_a_silent_reference_has_no_lag` | ## Security and architecture guardrails @@ -152,6 +153,7 @@ one to open. | Worker pool dispatch | `dw/server/jobs.py`: `JobManager._next_dispatch`; `dw/server/pool.py`: `WorkerSlot`, `check_fits`, `largest_ceiling_gb`; `dw/vram_estimate.py`: `required_vram_gb`; `dw/server/admission.py`: `pool_capacity` | `--devices` runs one worker per card; a queued job goes to a free card that fits what it needs, a job that fits no free card keeps its place while a smaller one behind it takes the card, a job bigger than every card is refused at submit, and cancel or a crash touches only the worker running that job. On a backend no `cost` entry describes, a declared `vram_estimate` is held at admission to the pool's largest card, so a job the first card cannot hold is not refused when a larger card can hold it. | `tests/test_worker_pool.py::TestConcurrency::test_two_workers_run_two_jobs_at_once`, `tests/test_worker_pool.py::TestBackfill::test_a_job_too_big_for_the_free_card_waits_while_a_smaller_one_takes_it`, `tests/test_worker_pool.py::TestFitAtSubmit::test_a_job_too_big_for_every_card_is_refused_naming_the_largest`, `tests/test_worker_pool.py::TestCancel::test_cancel_reaches_the_worker_running_the_job`, `tests/test_worker_pool.py::TestCrashIsolation::test_a_crashed_worker_fails_only_its_own_job`, `tests/test_worker_pool.py::TestOomRanking::test_the_later_started_worker_is_ranked_first`, `tests/test_worker_pool.py::TestFitAtSubmit::test_a_job_only_the_larger_card_holds_is_admitted` | | Card affinity | `dw/server/pool.py`: `choose_slot`; `dw/server/jobs.py`: `JobManager.route`; `dw/worker_protocol.py`: `workflow_identity` | Among the free cards a job fits it goes to the rerun's original card, else the card whose worker last ran its workflow identity (`workflow_identity`, the one rule the worker caches by too), else one with no worker running, else the first in `--devices` order; a cache probe and `plan.estimate` (`priced_for`) ask about the card `route` names, and memory is read and cleared per card. | `tests/test_worker_pool.py::TestIdentityAffinity::test_a_second_run_lands_on_the_card_that_ran_the_workflow_first`, `tests/test_worker_pool.py::TestRerunAffinity::test_a_rerun_lands_on_the_original_card_when_it_is_free`, `tests/test_worker_pool.py::TestProbeRouting::test_the_probe_goes_to_the_worker_the_job_would_be_routed_to`, `tests/test_worker_pool.py::TestMemoryPerCard::test_a_named_card_running_a_job_is_refused_though_the_other_is_idle` | | Per-card memory and cache | `dw/server/worker_memory.py`: `memory_status`, `clear_memory`, `probe_cache` | Each card keeps its own last memory reading; an idle card's worker is asked live, a busy one answers with the run's last report and says why it is not live, and a clear leaves a card running a job alone. Every request takes the card's lock with a bounded wait, and a card's lock is never taken while the JobManager's lock is held. | `tests/test_worker_pool.py::TestMemoryPerCard::test_a_named_card_running_a_job_is_refused_though_the_other_is_idle`, `tests/test_worker_pool.py::TestProbeRouting::test_the_probe_goes_to_the_worker_the_job_would_be_routed_to`, `tests/test_server.py::test_memory_says_why_a_reading_is_not_the_worker_s` | +| dw disk caches | `dw/disk_caches.py` | Members `inductor`, `triton`, `quantized`. Clear only a root carrying the `.dw-cache` marker, never while a job is running or queued (`quantized` also not during a model download); an entry delete takes only a 64-hex id, checked before any path is built; clear renames into `.clearing-*` then deletes. | `tests/test_disk_caches.py` | | Job-to-run link | `dw/server/job_record.py`, `dw/server/jobs.py`: `JobManager.realized` | A job records `run_id`, `run_dir` and `run_version`, and `JobManager.realized` reads that run's `workflow.json`, confined to the output root. | `tests/test_server_jobs.py::test_realized_reads_the_file_the_run_wrote`, `tests/test_server_jobs.py::test_realized_refuses_a_run_dir_that_escapes_the_output_root` | | Job history | `dw/server/job_history.py`: `JobHistory` | Finished jobs persist in `jobs.sqlite`, in WAL mode, so the Jobs view survives a restart. | `tests/test_server.py::test_job_history_survives_restart_and_reruns`, `tests/test_job_history_wal.py::test_the_jobs_database_uses_wal_mode` | | HTTP security | `dw/server/http_security.py`: `install_middleware`, `query_token_ok` | Every API route needs the token, and only a route marked `query_token_ok` takes it as `?token=`. | `tests/test_security_auth.py::TestTheTokenGate::test_every_api_spelling_needs_the_token`, `tests/test_security_auth.py::TestTheTokenGate::test_the_query_token_is_refused_where_it_is_not_allowed` | diff --git a/docs/MCP.md b/docs/MCP.md index 6126927e..2b6f884b 100644 --- a/docs/MCP.md +++ b/docs/MCP.md @@ -229,10 +229,10 @@ when no single workflow covers it. | `list_models()` | — | List what the Hugging Face model cache holds, largest first | | `get_memory()` | `device` | Get VRAM and RAM statistics: `workers` has one reading per card (the top level repeats the first card's), and `device` (`cuda:1`) asks for that card alone. `gpu_*` is the card; `host_memory_rss_mb` / `host_memory_peak_rss_mb` are the worker process's resident and high-water host memory, beside the machine's `host_memory_total_mb` / `host_memory_available_mb`; on Linux `host_memory_rss_anon_mb` / `host_memory_rss_file_mb` split the resident figure into memory the worker owns and mapped-file pages (memory-mapped checkpoint weights) the kernel can drop, so a loader that maps rather than copies raises the resident and peak figures without raising committed memory (#709) - read both, since an offloading workflow keeps its weights in host RAM and the card says little about what it holds (a host field is absent, not null, where the platform cannot measure it). `host_pinned_reserved_mb` / `host_pinned_allocated_mb`, when present, are torch's pinned-host cache - the staging buffers group offloading moves weights through, part of `host_memory_rss_mb` and invisible in every `gpu_*` figure, which is what a worker holding GB after releasing every model is usually holding (#98). `live: true` means `info` was measured now and is the worker's own memory; only live readings are comparable with each other. `live: false` with `info: null` (and `stale: false`) means nothing has been measured because nothing is resident; `live: false` with a populated `info` is a cached earlier reading, with `reason` (`job_running`, `worker_stopped`, `worker_busy`, `worker_unreachable`) and `age_seconds` - one cached while a job loads a model understates what is resident, so ask again when the card is idle rather than comparing it with a live figure | | `get_health()` | — | Check that the server is alive, and which machine answered: `version`, `device`, whether any model process is resident (`worker_alive`), and the queue depth. The engine runs one job per GPU: `workers` has one entry per card - `device`, `name`, `vram_gb`, `alive`, its own `current_job`, and the worker's `host_memory_rss_mb` while it runs - and the top-level `current_job` is the longest-running job. `worker_alive: false` is the normal idle state on a server that has not run a job since startup or the last memory clear - not a fault - the on-demand worker starts with the next job (#206) | -| `get_server_info()` | — | What this installation can do and where it keeps things: `device` (the accelerator a run will use), `version`, the `workspace` this session is working in and the workflow/asset/output/prompt `directories` of *that* workspace, the bind address and port, whether a token is required, and whether MCP is mounted. Check the device before authoring - a CUDA-only choice (bitsandbytes, `torch.compile`, flash attention) is not available on an `mps` or `cpu` server. `runtime` (#222) reports the Python version, torch version and the CUDA version torch was built against, the NVIDIA driver version (when `nvidia-smi` is reachable), and the installed versions of diffusers, transformers, accelerate, bitsandbytes, peft, safetensors and sentencepiece (`null` for one not installed) - for diagnosing an environment mismatch between boxes without shelling in | +| `get_server_info()` | — | What this installation can do and where it keeps things: `device` (the accelerator a run will use), `version`, the `workspace` this session is working in and the workflow/asset/output/prompt `directories` of *that* workspace, the bind address and port, whether a token is required, and whether MCP is mounted. Check the device before authoring - a CUDA-only choice (bitsandbytes, `torch.compile`, flash attention) is not available on an `mps` or `cpu` server. `runtime` (#222) reports the Python version, torch version and the CUDA version torch was built against, the NVIDIA driver version (when `nvidia-smi` is reachable), and the installed versions of diffusers, transformers, accelerate, bitsandbytes, peft, safetensors and sentencepiece (`null` for one not installed) - for diagnosing an environment mismatch between boxes without shelling in. `caches` lists dw's disk caches (`inductor`, `triton`, `quantized`), each with `name`, `bytes`, `files`, `enabled`, `clearable` and `disk_free`, plus `entry_count` for a cache with entries (no entry listing over MCP); it is omitted when the server has no `/api/caches`. There is no MCP tool to clear one - that is `DELETE /api/caches/{name}` or `.../entries/{entry}` ([SERVER.md](SERVER.md)) | | `list_jobs(limit=20, status=None, workspace=None)` | optional `limit` (newest N), `status` (one state or a comma-separated set of `queued`, `running`, `succeeded`, `failed`, `cancelled`), `workspace` | List queued, running and recent jobs, **newest first**. Bounded by default: the unbounded listing was over a client's tool-result limit on a server with a few months of history, which made it a tool that could not be called at all. `total` says how many matched and `truncated`/`next` say so when the answer was cut - raise `limit` or narrow with `status`. Without `workspace`, a named workspace lists its own jobs and the default one lists every job the server holds | | `list_gallery(limit=50, subfolder=None, only_orphans=False, workspace=None, folder=None, version=None)` | `limit`, `subfolder`, `only_orphans`, `workspace`, `folder`, `version` | List generated output files, newest first. A name is `//`, where `` may sit in the subfolder the step chose (`final/episode.mp4`); each entry carries `folder` (the workflow) and `subfolder` (by convention `final` or `intermediate`, `''` when the step chose none, any path the workflow wrote otherwise), and `subfolder="final"` lists only deliverables. Each entry carries `run_id` and `version` - that run's ordinal among the workflow's runs, which is how one of several runs that wrote the same basename is named to a person: the web UI labels the same file `v5`. The number is assigned when the run opens and never renumbered, so deleting a run leaves a gap rather than sliding the rest down (a failed run, or a rerun that reused every step, leaves one too - it took a number and may have nothing to list), and it is `null` under the flat output layout, which has no runs. `folder` with `version` lists that one run's files, and `output:/v5/` names one in a workflow; every other tool takes `name`. Each entry also carries a ready-made `url`, already scoped to the workspace that made it - a hand-built `/outputs/` URL 404s for anything but the default workspace. `only_orphans=True` inverts the call: instead of files, it returns run directories with no media anywhere under them (`runs`, each `{name, mtime}`) - a run whose output was deleted before `delete_output` could remove it by name, or one that failed before writing anything; `subfolder` does not apply in this mode, and `name` is exactly what `delete_output` accepts (#170). `workspace` names the workspace for this one call without switching the session to it - the same pin `run_workflow` takes, so a job run into another workspace stays reachable from the session that queued it | -| `get_gallery_metadata(name, envelope=False, workspace=None)` | `name`, `workspace` | Get the metadata embedded in a generated file — or, when `name` is an `asset:` reference, what an *input* asset holds (`source` says which; `job` is null for an asset). Reading an input's duration, frame count, fps and sample rate before a run is how a caller learns the `total_frames`, `fps` and `sample_rate` a workflow expects it to supply: the exact workflow and arguments that produced it, and, for audio/video, a `media` block (duration, rate, channels, fps, size, peak/mean dBFS) and `findings` - each level problem the server measured against `dw/audio_qc.py`'s thresholds (`full_scale`, `near_silent`), with the fix. Only an image (PNG/JPEG/WebP) embeds `metadata` this way — it is always null for audio and video, and `next` then points at `get_job_workflow(job_id)` when `job` is known, or says a kept asset carries no provenance at all when it isn't. `envelope=true` adds `media.envelope` — `rms_dbfs` and `peak_dbfs` one entry per second — which is what locates something in a track rather than measuring the whole of it. `media.shots` is set on an output joined from shots: one `{name, start_frame, num_frames, start_sample, num_samples}` per shot, as the join measured them (see the run manifest in docs/WORKFLOW_GUIDE.md), else null. `workspace` names the workspace for this one call without switching the session to it - the same pin `run_workflow` takes, so a job run into another workspace stays reachable from the session that queued it | +| `get_gallery_metadata(name, envelope=False, workspace=None, interval_seconds=None)` | `name`, `workspace` | Get the metadata embedded in a generated file — or, when `name` is an `asset:` reference, what an *input* asset holds (`source` says which; `job` is null for an asset). Reading an input's duration, frame count, fps and sample rate before a run is how a caller learns the `total_frames`, `fps` and `sample_rate` a workflow expects it to supply: the exact workflow and arguments that produced it, and, for audio/video, a `media` block (duration, rate, channels, fps, size, peak/mean dBFS) and `findings` - each level problem the server measured against `dw/audio_qc.py`'s thresholds (`full_scale`, `near_silent`), with the fix. Only an image (PNG/JPEG/WebP) embeds `metadata` this way — it is always null for audio and video, and `next` then points at `get_job_workflow(job_id)` when `job` is known, or says a kept asset carries no provenance at all when it isn't. `envelope=true` adds `media.envelope` — `rms_dbfs` and `peak_dbfs` one entry per second, or per `interval_seconds` (default 1.0, 0.05 to 3600; 0.1 places a vocal onset; a finer-than-default width that would need over 2400 bins is refused with the interval that fits) — which is what locates something in a track rather than measuring the whole of it. `media.shots` is set on an output joined from shots: one `{name, start_frame, num_frames, start_sample, num_samples}` per shot, as the join measured them (see the run manifest in docs/WORKFLOW_GUIDE.md), else null. `workspace` names the workspace for this one call without switching the session to it - the same pin `run_workflow` takes, so a job run into another workspace stays reachable from the session that queued it | ### Media diff --git a/docs/QUANTIZATION.md b/docs/QUANTIZATION.md index 0a50812d..aa89e68a 100644 --- a/docs/QUANTIZATION.md +++ b/docs/QUANTIZATION.md @@ -151,6 +151,33 @@ examples use for their large transformers and text encoders: **Examples:** [text-to-video.json](../workflows/templates/ltx2/text-to-video.json), [video-with-audio.json](../workflows/templates/minimax/video-with-audio.json) +### Cached on-the-fly (`DW_QUANTIZED_CACHE`) + +An on-the-fly load reads the bf16 weights and quantizes them on every cold load. With `DW_QUANTIZED_CACHE=1` the server keeps the result: the first cold load saves each quantized component to disk with sdnq's `save_sdnq_model` (safetensors plus JSON, no pickle), and a later load of the same key reads it back with `load_sdnq_model`, skipping the bf16 read and the quantize. Any other value leaves it off, with no events. It covers both ways a component is quantized with a `sdnq.SDNQConfig` `quantization_config`: a modular pipeline's `load_components`, one config per component (MiniMax H3), and a component loaded on its own through its class's `from_pretrained` (LTX-2.5's transformer and Gemma text encoder). A hit is placed on the config's `return_device` (LTX's transformer goes straight to `cuda`) and loads with no download. + +The variables are read by the server and inherited by its workers, so set them in the server's environment (e.g. `.env`) and restart. + +- **Where**: `DW_QUANTIZED_CACHE_DIR`, default `/dw-quantized`. `dw.serve` resolves it to an absolute path at start and pins it for the workers. Each entry is `//` with an `entry.json` (the key's inputs, size, created and last-hit times), written to a `.tmp-*` sibling and renamed into place +- **The key**: repo, subfolder, snapshot commit, quantizer class and config (minus the devices, the quantized-matmul switches and `non_blocking`, which don't change the stored weights), `torch_dtype`, and the sdnq, diffusers, transformers and torch versions. A key that can't pin the weights isn't cached +- **Not cached** (a `skipped` event with the reason): + + | Reason | When | + | ------ | ---- | + | `local_path` | the model is a local path | + | `unsupported` | a `variant`, no type hint, a single-file checkpoint, or a per-component load with a `from_pretrained` argument the key doesn't cover | + | `unpinned_snapshot` | the snapshot isn't complete on disk, or the recorded commit disagrees with the one the hub refs name - a save is only ever made under the commit an identical later lookup reads | + | `low_disk` | the save would leave free space below `min(50 GiB, 10% of the volume)` | + | `save_failed: ` | the save failed | + +- **Why a `miss`** (its `reason`): `no_entry` (nothing saved under the key yet), `incomplete_snapshot` (the snapshot isn't all on disk, so no lookup was made), or `load_failed: : ` (the entry didn't load and was deleted; the save that follows replaces it) +- **Concurrent workers**: two workers saving the same key leave one entry; the other reports `saved` ("already saved by another worker"). A corrupt entry is deleted and quantized fresh +- **What a hit holds**: the whole component as quantized. The save comes before pruning (`truncate_layers`, `remove_modules`), `sdnq_optimize`, LoRAs and offload, which all still apply after a hit +- **At server start** (only when on): a sweep deletes leftover `*.tmp-*` directories and entries built against other package versions or with an unreadable `entry.json`, and logs the count and bytes freed. It then logs the root, bytes in use and free space, warning when free space is below the floor plus the largest entry already cached +- **Clearing**: there is no other eviction. Use the server page's Caches panel (the whole cache, or one entry from its expandable table), or `DELETE /api/caches/quantized` / `DELETE /api/caches/quantized/entries/{id}`. `GET /api/caches` lists the entries (model, subfolder, size, created, last hit) and `low_disk` (free space under the floor). Clearing is refused while a job is running or queued, and while a model download is active; an off cache still lists and clears what is on disk +- **Event**: a `quantized_cache` job event, `{component, result: hit|miss|saved|skipped, reason?, seconds, bytes}`. A `miss`'s `seconds` is the time of the whole `load_components` batch, shared by the components it loaded; on the per-component path it is that one component's load + +The cache never fails a job; one that can't help steps aside. + ### Pre-quantized models Pre-quantized SDNQ repos load as complete pipelines. The `sdnq` module must be imported before loading so it can register with diffusers: diff --git a/docs/RELEASING.md b/docs/RELEASING.md index f1fa844a..80877c3c 100644 --- a/docs/RELEASING.md +++ b/docs/RELEASING.md @@ -7,6 +7,68 @@ notes from commits at tag time (see below). This section is a scratch pad for items a branch's author wants the next release note to name; clear it when a release ships. +### 0.12.0 + + + +A cold-start release. Opt in and a quantized model's second cold load +skips the bf16 read and the quantize: MiniMax H3 loads in 19 s instead of +113 s on an RTX 3090. The caches that build up on a server's disk now have +a page, and the music-video cut planner stops treating Whisper's +hallucinations as lyrics. + +**Quantized-component cache** (#815, #825, #826, #827) + +- `DW_QUANTIZED_CACHE=1` keeps each SDNQ on-the-fly-quantized component on + disk (safetensors plus JSON, under `/dw-quantized` or + `DW_QUANTIZED_CACHE_DIR`) and reads it back on the next load of the same + key. Measured on an RTX 3090: H3 cold 112.7 s -> 18.7 s, warm 67.5 s -> + 4.0 s. An H3 configuration holds about 43 GB. +- It covers a modular pipeline's `load_components` (H3) and a component + loaded on its own through `from_pretrained` (LTX-2.5's transformer and + its Gemma text encoder, saved under their own tensor names so they load + back). A hit is placed on the config's `return_device`. +- The key pins the snapshot commit, the quantizer config and the sdnq, + diffusers, transformers and torch versions; a save is refused when it + would leave less than `min(50 GiB, 10%)` free, and a `quantized_cache` + job event names every hit, miss, save and skip with its reason. The cache + never fails a job. Off by default; docs/QUANTIZATION.md has the details. + +**Disk caches** (#828, #832, #833, #834, #837) + +- `GET /api/caches` lists dw's caches of derived data - `inductor` + (`torch.compile` kernels), `triton` (SDNQ's jit and autotune kernels) + and `quantized` - with size, files, last use and the volume's free + space; `DELETE /api/caches/{name}` clears one, refused while a job runs + or a download is active. The quantized cache lists its entries and + `DELETE /api/caches/quantized/entries/{id}` removes one. +- The web UI's Server page has a Caches panel for the same, asking before + a clear and quoting what the clear costs. `get_server_info` carries the + list for an agent. + +**Transcription and cuts** (#822, #823) + +- `music-video-cuts` transcribes the vocal stem rather than the full mix, + so an instrumental intro no longer folds into the first word. Chunks + that are silent, overlong, wordless or a repetition loop are marked + suspect; `plan_cuts` leaves them out and warns once per run. +- `get_gallery_metadata`'s envelope takes `interval_seconds` (0.05 to + 3600, default 1.0) so an agent can place a cut to a frame. More than + 2400 bins is refused unless the default interval would give more. + +**Fixes** + +- The inherited cost estimate is withheld when an inline workflow runs + more steps or frames than the template it came from, instead of + quoting the template's number (#821: a 7x under-quote). +- `assess_output` shot findings carry `index` and `start_frame`, so a + repeated shot name is unambiguous (#835). +- `dialogue-short` and `music-video` expose `match_levels_dbfs` (#836). +- A default cache root that appears after the server starts is marked, + so `triton` stays clearable after the first run (#837). +- `deploy.sh` restarts a `dw.serve` run as a loaded launchd daemon on + macOS. + ### 0.11.1 diff --git a/docs/REMOTE.md b/docs/REMOTE.md index 7200039c..05ff4aee 100644 --- a/docs/REMOTE.md +++ b/docs/REMOTE.md @@ -134,6 +134,8 @@ and Host checks work unchanged behind the proxy because Caddy forwards the `Host` header as-is. nginx works the same way with `proxy_pass` and `proxy_set_header Host $host;` plus your own certificate. +A token holder can also empty dw's disk caches (`DELETE /api/caches/{name}`, or one entry via `/entries/{entry}`), as they can delete models. + The token is still the only authentication in either setup; a proxy or a VPN protects the transport, not the door. diff --git a/docs/SECURITY.md b/docs/SECURITY.md index 16db319b..fd63c0fe 100644 --- a/docs/SECURITY.md +++ b/docs/SECURITY.md @@ -271,6 +271,10 @@ SecurityError InvalidInputError — input validation failure ``` +### Disk caches + +`DELETE /api/caches/{name}` (`dw/disk_caches.py`) empties a directory whose path comes from an environment variable, so it never trusts the root: only a root carrying a `.dw-cache` marker is cleared. `dw.serve` writes the marker at start into a root that is missing or empty (a missing `triton` root is never created), or into a populated root equal to the cache's default exactly (`~/.cache/dw/torchinductor` or its `$XDG_CACHE_HOME` form; `~/.triton/cache` only - a populated `$TRITON_HOME/.triton/cache` is not grandfathered). Any other populated root is listed with `clearable: false` ("root not created by dw (no .dw-cache marker)") and its clear is a 400, until someone creates the marker by hand; deleting the marker stops clearing. `quantized` (`DW_QUANTIZED_CACHE_DIR`) follows the same marker rule, and `DELETE /api/caches/{name}/entries/{entry}` deletes one entry only if the id matches `^[0-9a-f]{64}$`, checked before any path is built. The clear refuses while a job is running or queued (for `quantized`, also while a model download is active), and renames the root's children into `.clearing-*` before deleting them. Like model deletion, it needs only the bearer token. + ## Rules - Always validate paths before file operations diff --git a/docs/SERVER.md b/docs/SERVER.md index d4a95dbb..a64d7db7 100644 --- a/docs/SERVER.md +++ b/docs/SERVER.md @@ -28,6 +28,15 @@ The server keeps a persistent GPU worker underneath: models stay loaded between runs, so re-running a workflow with a new prompt skips the load entirely. +### Environment + +The server reads these from its environment (e.g. `.env`) and its workers inherit them; changing one needs a restart. + +| Variable | Default | Effect | +| -------- | ------- | ------ | +| `DW_QUANTIZED_CACHE` | off | `1` caches SDNQ on-the-fly-quantized components on disk, so a later cold load skips the quantize. Anything else is off. See [QUANTIZATION.md](QUANTIZATION.md#cached-on-the-fly-dw_quantized_cache) | +| `DW_QUANTIZED_CACHE_DIR` | `/dw-quantized` | Where the cache lives; resolved to an absolute path at start and pinned for the workers | + ### Choosing the card `--devices cuda:1` runs the worker on that card. The worker is pinned to it @@ -178,7 +187,12 @@ all-or-nothing. See [Workspaces](WORKSPACES.md#several-workspaces-on-one-server) The sidebar lists every workspace, and `+ new` there creates one; a workspace's own Overview page is where it is deleted. Server → Status shows the resolved directories and the `claude mcp add` line for connecting an -agent from another machine, plus the queue across every workspace: +agent from another machine, plus the queue across every workspace. Its +Caches panel lists dw's disk caches (`GET /api/caches`: size, files, last +used, the volume's free space; the quantized cache also expands into a table +of its entries with a delete button each) with a Clear button per cache that asks first, +quoting what clearing costs, and shows the server's refusal when a job is +running: ![The Server page before the sidebar: address picker, generated claude mcp add line, resolved directories — the workspace list it shows now lives in the sidebar](img/ui-server-dark.png) @@ -620,6 +634,37 @@ The editor's forms come from these; they are just as usable from scripts: so no multipart parser is needed for a single-file upload - `GET /api/models`, `DELETE /api/models?repo={repo_id}` — hub cache inventory and deletion +- `GET /api/caches`, `DELETE /api/caches/{name}`, `DELETE /api/caches/{name}/entries/{entry}` - dw's disk caches of derived + data (not the hub cache above): `inductor` (`torch.compile`'s kernels), + `triton` (Triton's own cache, which holds SDNQ's `triton.jit`/autotune + kernels and is shared with any other Triton user on the machine) and + `quantized` (the SDNQ quantized-component cache, `DW_QUANTIZED_CACHE_DIR`; + `enabled` follows `DW_QUANTIZED_CACHE=1`, and an off cache still lists and clears). + A cache with entries (`quantized`) lists them inline as `entries` (`id`, + `bytes`, `files`, `unreadable`, `model`, `subfolder`, `component`, `created`, + `last_hit`; an entry with a missing or unreadable `entry.json` is + `unreadable: true` and deletable) plus `low_disk` (free space under + `min(50 GiB, 10% of the volume)`). The entry DELETE answers + `{"name", "entry", "deleted": true, "freed"}`; the id must match + `^[0-9a-f]{64}$` (else 400, before any path is built), 400 on a cache without + entries or an unmarked root, 404 for an unknown entry. For `quantized`, both + it and the whole-cache DELETE also 409 while a model download is active. The GET + answers `{"caches": [...]}`, each with `name`, `root`, `enabled` (`triton` + is false where it is not installed, as on a Mac), `bytes`, `files`, + `last_used` (newest mtime, epoch seconds), `clearable`, + `not_clearable_reason` (only when not clearable), `cost_of_clearing`, + `disk_free` and `disk_total`. The DELETE answers + `{"name", "cleared": true, "freed", "files"}`; 404 for an unknown name (the + detail names the valid ones), 409 while a job is running or queued, 400 + when the root is not clearable. It empties the root but keeps the root and + its `.dw-cache` marker: children are renamed into + `/.clearing--/`, then deleted. Clearing costs a recompile + or re-autotune on the next first step of each shape (100-300 s per shape on + MiniMax-H3 for `triton`). A populated root that an environment variable + points away from its default has no marker and is not clearable until you + create the `.dw-cache` file in it by hand (read per request, no restart); + see [SECURITY.md](SECURITY.md#disk-caches). The Server page's Caches panel + drives all three. Behind the bearer token like every `/api/*` route - `POST /api/models/download` (`{"repo_id": ...}`), `GET /api/models/downloads`, `POST /api/models/downloads/{id}/cancel` — background snapshot downloads with byte-level progress diff --git a/docs/TASKS.md b/docs/TASKS.md index 026f9298..6bcf24af 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -1970,7 +1970,7 @@ Plan a music video's cuts from a song's lyrics and beats. `transcribe_audio` (with timestamps) says when each line is sung and `analyze_beats` says where the beats fall; this turns the two into shots - each a start frame and a length - that tile the song exactly, so one prompt per shot can be written and -the list rendered. The result is JSON and nothing is built: +the list rendered. A transcript chunk `transcribe_audio` marked `suspect` `no_words`, `repetition` or `silent` is left out of the lines (no word, a decoding loop, or no voice under it - not a sung line); one marked `long` is still used. Each run of consecutive chunks with the same mark adds one warning naming its span, its chunk count and its text (clipped), so a hallucinated intro or fade is one line, not one per word. Lines squeezed to no time back to back join the next scene with one warning per run, the first three named. The result is JSON and nothing is built: ```json { @@ -2069,7 +2069,7 @@ without `beats`, `min_scene_s` above `max_scene_s`, and `fps`, `duration_s`, can't be known (no `duration_s`, no `beats` result, and no transcript to end on) and `lyrics` with no sung lines. [music-video-cuts.json](../workflows/templates/minimax/music-video-cuts.json) -chains `transcribe_audio`, `analyze_beats` and `plan_cuts`, passing H3's +chains `separate_stems`, `transcribe_audio` (of the vocal stem, so an instrumental intro is not folded into the first word), `analyze_beats` (of the full mix) and `plan_cuts`, passing H3's grid (`17`, `5`, `124`, `345`) and a `lead_s` variable (default `0.5`, 12 frames at 24 fps). Its shots are what [music-video.json](../workflows/templates/minimax/music-video.json) renders: @@ -2977,6 +2977,8 @@ Multi-channel audio is downmixed to mono and resampled to 16 kHz before transcri Set `timestamps` to `"segment"` or `"word"` to get chunk timings instead — a music video cut to the lyric, or a dialogue shot checked against its line, needs the times Whisper already produces past 30 s rather than the collapsed string. The result becomes `{"text": ..., "chunks": [{"start": ..., "end": ..., "text": ...}, ...]}`, so the step's `result.content_type` must be `"application/json"` rather than `"text/plain"`, and it's read with MCP's `get_output_text` (JSON results are text). A clip under 30 s asks Whisper for timestamps explicitly when `timestamps` is set — the 30 s long-form threshold is a separate, unrelated reason to ask. Every `start` and `end` is a number of seconds: a chunk the clip ends inside - a song cut mid-line, which Whisper leaves open-ended - ends at the clip's duration, so the transcript drops straight into [`attribute_voices`](#voice-attribution)' `lines`. `"word"` can give a zero-length chunk (`start` = `end`) on a short or clipped word; `attribute_voices` reads it as a line with no voice. +A chunk that is probably not speech carries one more field, `suspect` (#822), and is never removed. The first reason that holds wins: `"no_words"` when its text has no word at all (`—`, `♪`, `...`); `"repetition"` when it holds a word of a repetition loop (`I'm sorry, I'm sorry, ...`), found by the same detector and `REPEAT_RUN_MIN`/`REPEAT_MAX_PERIOD` [`check_script`](#script-check) discards loops with, however loud; `"silent"` when its loudest window sits at or below the dead-air floor `check_script` discards words at (`GUARD_FLOOR_DBFS`) - Whisper invents text over silence, a fade-out's end credits included; and `"long"` when a `"word"` chunk lasts more than `LONG_WORD_SECONDS` (2 s), an instrumental folded into the word. A `"segment"` chunk is never `long`. A clean input does not prevent the marks: transcribe a song's vocal stem ([`separate_stems`](#stem-separation)) so an instrumental intro is not counted as speech, and the marks catch what Whisper still invents over the stem - its leakage can keep an intro above the floor, which is what `no_words` and `repetition` are for. [`plan_cuts`](#plan_cuts) leaves `no_words`, `repetition` and `silent` chunks out of the lines and warns once per run of marked chunks. + To compare a take against the lines it was meant to speak, use [`check_script`](#script-check) rather than reading the transcript by eye. **Example:** [transcribe-audio.json](../workflows/templates/transcribe-audio.json) — Transcribe an audio file to text. diff --git a/docs/WORKER_GUIDE.md b/docs/WORKER_GUIDE.md index a90c0594..48f9b846 100644 --- a/docs/WORKER_GUIDE.md +++ b/docs/WORKER_GUIDE.md @@ -35,6 +35,13 @@ the task model cache are still emptied, and if the incoming definition cannot be prepared everything is released. `clear_memory` (`POST /api/memory/clear`) remains the way to empty the card by hand. +## Cold Loads + +With `DW_QUANTIZED_CACHE=1`, the first cold load of an SDNQ on-the-fly-quantized +component pays the quantize and a save to disk. Later cold loads - a new worker, +after a restart or after eviction - read the cached entry instead. See +[QUANTIZATION.md](QUANTIZATION.md#cached-on-the-fly-dw_quantized_cache). + ## Troubleshooting **Worker crashes**: `JobManager` detects it and starts a fresh worker on the diff --git a/docs/proposals/complete/disk-caches-complete.md b/docs/proposals/complete/disk-caches-complete.md new file mode 100644 index 00000000..aabc62f2 --- /dev/null +++ b/docs/proposals/complete/disk-caches-complete.md @@ -0,0 +1,137 @@ +# List and clear dw's own disk caches (#828) + +Written by model `claude-opus-5-5` via provider `anthropic` (close-out session, +2026-10-11). Filed after lem's HF disk filled to 100% on 2026-10-11, with no +way to see how much of the disk was dw's own data. Plan v1 and v2 were designed +on #828, and Don approved v2 with his Q4 and Q5 answers folded into the stages. +It was built as stages #832 (A), #833 (B) and #834 (C). The plan comment on #828 +stays the full design. This file records what shipped. + +## The problem + +dw's runs leave rebuildable data on disk that outlives them: +- Inductor's compiled kernels: `TORCHINDUCTOR_CACHE_DIR`, ~300 MB on lem. +- Triton's compiled and autotuned kernels: SDNQ's `triton.jit`/autotune, in + `~/.triton/cache`, 7.1 GB on lem. +- #815's quantized weights: `DW_QUANTIZED_CACHE_DIR`, ~288 GB on lem by the + end of stage C. + +Nothing reported their size, and nothing cleared them short of a shell on the +server. + +## Verdict: build, in three stages, with no new MCP tool + +- **A registry**, `dw/disk_caches.py`, with three members: + - `inductor`; + - `triton`, registered where Triton already keeps it (dw sets no `TRITON_*` + variable); + - `quantized`. +- **REST:** + - `GET /api/caches`; + - `DELETE /api/caches/{name}`; + - `DELETE /api/caches/{name}/entries/{entry}`. +- **A Caches panel** on the server page. +- **MCP reads only:** `get_server_info` gains `caches`. There is no + `clear_cache` tool, and the tool count stays 62 (Don's review of v1, Q1). + +## What shipped + +- **Registry** (`dw/disk_caches.py`). Each member has a root, a default root, + `enabled`, a `cost_of_clearing` sentence, and (for `quantized`) entries. + `scan()` reports bytes, files, `last_used`, `clearable` with a reason, and + the volume's `disk_free`/`disk_total`. The walk never follows symlinks, and a + missing root reads as 0. + - **Triton's root** is resolved without importing triton: + `TRITON_CACHE_DIR`, else `$TRITON_HOME/.triton/cache`, else + `~/.triton/cache`. `enabled` comes from `find_spec`. dw never creates the + directory. + - **Ownership marker:** a `.dw-cache` file. dw writes it into a root that is + empty or new. It also writes it into a populated root that exactly equals + that member's default; this is the grandfathering rule. Only + `~/.triton/cache` is grandfathered for Triton, not `$TRITON_HOME` (Q5). A + root without the marker is `clearable: false`, and a clear of it is a 400. + - **Clear:** each child is renamed into `.clearing--`, and that + directory is then `rmtree`d. The root and its marker are kept. Every path + goes through `validate_path(..., base_dir=root)`. +- **Routes** (`dw/server/routes/system.py`): + - 404 for an unknown cache name, and the response names the valid ones; + - 400 for a root that isn't clearable, or a malformed entry id (checked + against the key pattern before any path is built); + - 409 while a job is running or queued, and for `quantized` also while a + model download is active. + + The response models are in `api_models.py`, with the regenerated + `openapi.json` and `api-schema.ts`. +- **MCP:** `get_server_info`'s `caches` gives each member's name, bytes, files, + enabled, clearable and disk_free, plus `entry_count` for `quantized`. It is + omitted when `/api/caches` fails. +- **UI:** the Caches panel lists each cache. Its Clear button opens a confirm + that quotes the cost, and is disabled with the reason when the cache can't be + cleared. `quantized` has an expandable entry table with a delete button per + entry, and a low-disk warning row. +- **Quantized member** (stage C): + - Entries come from `/entry.json`. An unreadable entry is listed with + `unreadable: true` and can still be deleted. + - `low_disk` uses #826's `quantized_cache.floor_bytes`. + - The key check and the default root live only in `dw/quantized_cache.py`, + and `disk_caches` calls them. + - A cache that is turned off still lists and clears what is on disk. +- **Docs:** `SERVER.md`, `MCP.md`, `ACCELERATION.md`, `QUANTIZATION.md`, + `REMOTE.md`, `SECURITY.md`, and the `ARCHITECTURE.md` row *dw disk caches*. + +## Deviations from the plan + +- **No `SURFACE_BUDGET` raise.** The docstring clause measured 16 tokens, which + still fits under the existing ceiling. The measurement is recorded in + `tests/test_mcp_server.py`. +- **A missing root lists `clearable: false`** ("root does not exist (nothing to + clear)"). On mini-ai, `triton` reads `enabled: false, bytes: 0`. +- **Roots are re-read per request**, not fixed at server start. The marker is + also written when a default root first appears after start (#837 below). +- **An entry's `component` is its `subfolder`,** because `entry.json` has no + component field. +- **`DELETE /api/caches/..%2Fx` answers 404 or 405, not only 404.** The decoded + `/` matches no clear route, so nothing is touched either way. + +## Question 2 answered: clearing under a warm worker needs no restart + +This was checked on lem, on one warm worker (`cuda:1`) running the MiniMax H3 +`video-with-audio` template, with no `clear_memory` between runs. After +`DELETE triton` freed 7.08 GB (149,915 files), a new shape denoised and +succeeded. After `DELETE inductor` freed 60 MB, the same shape and then another +new shape succeeded too. The designed restart fallback wasn't needed and wasn't +built. + +`triton` did not refill under the warm worker. The worker keeps its kernels in +memory, so the cache only fills again on a cold start. This is expected, not a +defect. + +## Deferred + +- **Eviction:** clearing is manual only. If the registry shows `triton` growing + again, the first follow-up is a size cap or LRU for it. A second full-disk + incident with the quantized cache on would also reopen this. +- **A `clear_cache` MCP tool:** adding one later is additive and needs its own + approval. +- **Caching the scan:** no cache was added. If the walk over ~20k Triton + entries makes `get_server_info` slow, cache the scan for 30 s. + +## Stages and bounces + +| Stage | What | Bounces | +|---|---|---| +| #832 A | registry, `inductor`/`triton`, list and clear routes, MCP `caches` | 0. Verified on mps first time (C-F389, C-F390). | +| #833 B | the UI Caches panel | 0 tester bounces. Don's browser check found that an empty cache's confirm said "1 KB", fixed in `d83a3793`. The REST steps on lem were run by the lead, because the tester is MCP-only. | +| #834 C | `quantized` member, entry delete, `low_disk` | 1 architecture-review bounce: the key pattern and the default root were each defined twice. Fixed in `096bc1fd`. Then verified (C-F391-C-F393). | + +**After the stages:** #837. On lem, `triton` turned `clearable: false` after +its first run. The marker was only written at server start, and Triton's root +didn't exist yet then. Fixed in `16959b02`: a default root that appears after +start is now marked. When this was written, #837 was still waiting for the +tester's verification. + +**Process note:** the plan was decomposed twice, which filed a duplicate stage +set (#829-#831, closed `not planned`). Stage A was built under #829's number +and handed off under #832. + +The stage comments give no cost figures. The plan estimated ~$8-14. diff --git a/docs/proposals/complete/quantized-cache-complete.md b/docs/proposals/complete/quantized-cache-complete.md new file mode 100644 index 00000000..76bafb37 --- /dev/null +++ b/docs/proposals/complete/quantized-cache-complete.md @@ -0,0 +1,167 @@ +# On-disk cache of on-the-fly-quantized components (#815) + +Written by model `claude-opus-5-5` via provider `anthropic` (close-out session, +2026-10-11). The idea came from a load-time profile during the H3 speed A/B on +2026-10-10. Plan v1-v3 were designed on #815; stage 0's measurements and Don's +six decisions made v4, which Don approved. It was built as stages #825 (0), +#826 (A) and #827 (B). The plan comment on #815 stays the full design; this is +the record of what shipped. + +## The problem + +Every cold load of an SDNQ on-the-fly component (a fresh worker, a +`clear_memory`, a config change) read the bf16 weights, copied them to the GPU, +quantized them and copied them back, and an identical config gave identical +weights each time. The issue's profile put the H3 transformer at 46.7 s, ~34 s +of it data movement rather than quantization math. + +## Verdict: build smaller + +Built for **SDNQ only**, **opt-in**, in the two families that quantize on the +fly: MiniMax H3 (modular `load_components`) and LTX-2.5 (per-component +`from_pretrained`). + +Cut, with what would bring each back: +- **bitsandbytes, torchao, quanto:** partial serialization, some via pickled + tensor subclasses, so each needs its own security review. Trigger: a field + report of slow cold loads from a template using one. +- **GGUF:** already a file format; nothing to gain. +- **Cache-aware `plan.downloads_required`:** trigger: a run that downloads bf16 + weights a hit never reads. +- **Listing and clearing over MCP, REST or the UI:** moved to #828 (Don, + decision 5). Shipped there: the server page's Caches panel and + `DELETE /api/caches/quantized[/entries/{id}]` (`disk-caches-complete.md`). +- **Reordering pruning ahead of the cache (Q5):** stays no. The entry holds the + untruncated component; a hit takes ~1-9 s, so time doesn't trigger a revisit. + +## Plan history + +- **v1-v3:** design, a measurement gate (stage 0 must show a ~30 s H3 saving), + cache-before-pruning, and the key. +- **v4:** folded in stage 0's numbers and Don's decisions: + 1. off by default, `DW_QUANTIZED_CACHE=1` opts in; + 2. `DW_QUANTIZED_CACHE_DIR`, default `/dw-quantized`; + 3. a free-space floor of `min(50 GiB, 10%)` per save (`skipped: low_disk`); + 4. the key's commit comes from a snapshot verified complete on disk + (`skipped: unpinned_snapshot` otherwise); + 5. listing/clearing moved to #828; + 6. #825's three script fixes folded into stage A. + + Don's review fixed one thing in the decompose: the start-up low-disk warning's + margin is **the largest entry already in the cache** (zero when empty), not a + hard-coded H3 size, because model sizes don't belong in engine code. + +## What was built + +### Stage 0 (#825): measurement script + +`scripts/measure_quantized_cache.py` (03b2424d), run by Don's local session on +lem (RTX 3090) with `dw.serve` stopped. sdnq 0.2.7, diffusers 0.41.0.dev0, +transformers 5.19.0, torch 2.14.1+cu130. Seconds: + +| component | weights | fresh warm | fresh cold | cached warm | cached cold | GB on disk | +|---|---|---:|---:|---:|---:|---:| +| H3 transformer | int4 | 33.1 | 52.4 | 1.6 | 8.1 | 18.88 | +| H3 text encoder | int4 | 31.1 | 50.2 | 1.9 | 9.0 | 20.88 | +| H3 VAE | int8 | 3.3 | 10.1 | 0.5 | 1.6 | 3.15 | +| **H3 total** | | **67.5** | **112.7** | **4.0** | **18.7** | **42.91** | + +**63.5 s saved per cold load with the page cache warm, 93.9 s cold**, against +the ~30 s gate. + +### Stage A (#826): the cache module and the H3 modular path + +- `dw/quantized_cache.py`: the env switch and root (`configure_at_startup`, + called from `dw/serve.py` before workers spawn), the key, snapshot + resolution, the free-space floor and start-up disk line, atomic save with the + two-worker race, verified load (a corrupt entry is deleted and the load falls + back to a fresh quantize), the version sweep, and `quantized_cache` events. +- The hook in `dw/pipeline_processors/components.py`: look up each SDNQ + component, register hits through `update_components` (keeping `load_id`), + `load_components` for the misses only, then save each miss before + `sdnq_optimize`, LoRA and pruning. +- Snapshot resolution goes through `huggingface_hub.try_to_load_from_cache` + (offline), not a hand-written `models--` layout (the architecture review's + build-vs-buy finding). +- **Two key bugs found in acceptance, both in the quantizer config:** + - `SDNQConfig.to_dict()` returns its own lists, and sdnq appends the modules + it skips to `modules_to_not_convert` during the load, so the save key never + matched a later lookup key. Fixed with a deep-copied `config_snapshot`. + - `SDNQConfig.post_init` dedupes lists with `list(set(...))`, so their order + follows each process's hash seed: the transformer's key differed between + the two cards' workers. Fixed by sorting the set-valued arguments + (`SET_ARGUMENTS`), pinned by a test across four `PYTHONHASHSEED`s. +- #825's three script fixes (argv match excluding its own ancestors, narrowed + `allow_patterns`, fsync before eviction and page-touching inside the timed + region). +- Docs: `QUANTIZATION.md` ("Cached on-the-fly"), `ACCELERATION.md`, + `WORKER_GUIDE.md`, `SERVER.md` (both env vars), and the *Quantized component + cache* row in `ARCHITECTURE.md`. +- lem opted in through its `.env` (Don). + +**Accepted on lem** (M-F128-M-F134, M-F139, M-F140): an H3 run with all three +components hitting loaded in 72.3 s against 173.8 s for the miss-and-save run, +with per-component hits of 0.4-1.1 s; frames identical across hit and miss +runs; a text-encoder dtype change misses only that component; a LoRA applies +on a hit; the two-worker race gives one `saved` and one "already saved by +another worker". + +### Stage B (#827): the per-component path (LTX-2.5) + +- `quantized_cache.from_pretrained`, called from `load_component`'s + `model_name` branch (the path `load_optional_component` takes). It shares + `_lookup` with the modular path, so the two can't drift. A hit is placed on + the config's `return_device` (`cuda` for the LTX transformer) and skips + `download_watch`; a miss loads under it as before. +- **One bug found in acceptance:** transformers 5's `save_pretrained` writes a + model in its checkpoint's original layout by reversing the class's weight + conversion mapping, and `load_sdnq_model` reads names as they are. Gemma 4's + vision-tower mapping made every LTX text-encoder entry fail to load, get + deleted and re-save ~13 GB per run. Fixed in `_save_model` by passing + `save_original_format=False` when `save_pretrained` accepts it. H3's + Qwen3-VL has no such mapping and was never affected. +- The LTX template's "First load" wording now says the cache removes the + quantize cost after the first load. + +**Accepted on lem** (M-F135-M-F137, M-F141): loading 33.9 s and 38.0 s on a +hit against 61.6 s on the run that re-saved the text encoder; identical frames +across runs; mean step on a hit within 1.2× of the miss run; H3 still hits. +Entry sizes reported: transformer ~11.2 GB, Gemma text encoder ~13.0 GB on +disk. + +## Bounces + +| stage | architecture review | tester | +|---|---|---| +| #825 (0) | — | 0 | +| #826 (A) | 1 (build vs. buy: hand-written hub-cache layout) | 4: one environmental (lem not yet opted in), three on M-F128 (the two key bugs above; the first fix exposed the second) | +| #827 (B) | 0 | 1 (Gemma 4 entries saved in the original layout) | + +Cost per stage isn't recorded here: the stage comments name no `usage:` +figures. The plan estimated A at $6-9 and B at $3-5. + +## Deferred + +- **The LTX rows of stage 0's table** (fresh/cached, warm/cold, measured with + the fixed script and `dw.serve` stopped). lem's LTX snapshot is complete + (`da0b6909`), so it can run whenever lem is idle. Until then the template's + "about 10 minutes" fresh-quantize figure is the old, unmeasured one. +- **Listing, clearing and eviction:** #828. +- **Reporting the cache's on/off and free-space state** in `get_server_info` + or `get_health`, so a consumer can check opt-in directly rather than infer it + from events (the tester's note on #815; a natural fit for #828). +- **mps/cpu coverage:** mini-ai isn't opted in, so only the off arm (M-F139) + ran there. +- A possible simplification from B's architecture review: if + `save_sdnq_model` forwards keyword arguments to `save_pretrained`, the + temporary `save_pretrained` override in `_save_model` can go. + +## Risks that remain + +- A stale hit gives silently wrong weights. The key carries a commit from a + complete snapshot plus the sdnq, diffusers, transformers and torch versions, + and the sweep deletes old-version entries at start. The residual is an sdnq + default that changes without a version bump. +- Disk: ~43 GB per H3 config and ~24 GB per LTX config, with no eviction between + version bumps beyond the floor. +- Opt-in means only lem exercises it; unit tests carry the rest. diff --git a/docs/proposals/todo.md b/docs/proposals/todo.md index 1220b447..31886077 100644 --- a/docs/proposals/todo.md +++ b/docs/proposals/todo.md @@ -41,6 +41,22 @@ remaining deferred fix is recorded in ## Shipped since the ranking +- **List and clear dw's own disk caches** (#828, stages #832-#834), shipped + 2026-10-11. A registry (`dw/disk_caches.py`) of `inductor`, `triton` (Triton's + own `~/.triton/cache`, left in place) and the `quantized` cache, with + `GET /api/caches`, `DELETE /api/caches/{name}` and a per-entry delete for the + quantized cache, a Caches panel on the server page, and sizes in + `get_server_info`'s `caches` (no MCP clear tool). Clearing works on an idle + warm worker without a restart, and only for a root that carries dw's + `.dw-cache` marker. Eviction stays manual. Record: + `complete/disk-caches-complete.md`. +- **On-disk cache of SDNQ on-the-fly-quantized components** (#815, stages + #825-#827), shipped 2026-10-11. Opt-in with `DW_QUANTIZED_CACHE=1`, rooted at + `DW_QUANTIZED_CACHE_DIR` (default `/dw-quantized`), on both the H3 + modular path and the LTX-2.5 per-component path. Measured on lem: an H3 cold + load saves 63.5 s (page cache warm) to 93.9 s (cold). Record, including what + was cut or deferred (other backends, the LTX measurement rows; listing and + clearing shipped in #828): `complete/quantized-cache-complete.md`. - **`dw/server/jobs.py` split into `pool` / `job_results` / `worker_memory`; device ordinal and card as stored fields** (#693, stages #776-#777), shipped 2026-10-08. The `slots[0]` shims are gone diff --git a/docs/stabilization/baseline.json b/docs/stabilization/baseline.json index cd017ff0..303dabdb 100644 --- a/docs/stabilization/baseline.json +++ b/docs/stabilization/baseline.json @@ -1,10 +1,10 @@ { - "modules": 206, + "modules": 208, "modules_over_size_ceiling": 0, "functions_over_150_lines": 0, "prefix_literals": 0, "prefix_handling": 0, - "test_dw_patch_targets": 325, + "test_dw_patch_targets": 326, "claude_md_lines": 129, "duplicate_blocks": 0, "complex_functions": 2, diff --git a/dw/disk_caches.py b/dw/disk_caches.py new file mode 100644 index 00000000..a70b12c9 --- /dev/null +++ b/dw/disk_caches.py @@ -0,0 +1,471 @@ +"""dw's own disk caches: what derived data its runs leave on disk, how big it +is, and clearing it (#828). + +Each cache is a `DiskCache` in `CACHES`, named by the REST routes +(`GET /api/caches`, `DELETE /api/caches/{name}`) and summarized in the MCP +`get_server_info`'s `caches`: + +- `inductor`: torch.compile's kernels, under `TORCHINDUCTOR_CACHE_DIR`, which + `dw/__init__.py` always sets. +- `triton`: Triton's compiled and autotuned kernels (SDNQ's `triton.jit` / + `triton.autotune` matmul and attention, most of the compile disk use). + Registered where Triton itself puts it - `TRITON_CACHE_DIR`, else + `$TRITON_HOME/.triton/cache`, else `~/.triton/cache` - resolved here + without importing triton. dw sets no `TRITON_*` variable and never creates + that directory. +- `quantized`: SDNQ components quantized once and saved (#815, + `dw/quantized_cache.py`), under `DW_QUANTIZED_CACHE_DIR`, enabled by + `DW_QUANTIZED_CACHE=1`. Its entries - one `<64-hex key>/` directory each, + described by its `entry.json` - are listed and can be deleted one at a time + (`DELETE /api/caches/quantized/entries/{entry}`). A cache that is off still + lists, and still clears, what is on disk. + +**Clearing never trusts the root.** A root comes from an env var, and one +pointed at `~` by mistake must not be emptied. Only a root carrying the +`.dw-cache` marker is clearable. `mark_roots()`, run once at `dw.serve` +start, writes it into a root that is empty or missing (a missing `triton` +root is left missing), or into a populated root equal to the member's +default exactly - the cache dw (or Triton, by default) has been filling all +along. A populated root anywhere else stays unmarked and its clear is +refused, until someone marks it by hand. + +**Clear empties the root and keeps it and its marker:** every top-level child +is renamed into one `/.clearing--/` directory, which is then +deleted. A writer racing the clear finds an empty cache and rebuilds; it never +reads a half-deleted tree. An entry delete moves the one entry the same way. The +routes refuse while a job is running or queued. + +**An entry id must have a cache key's shape** (`quantized_cache.is_key`, 64 +lowercase hex) before any path is built from it. +""" + +import importlib.util +import logging +import os +import secrets +import shutil +import threading +from dataclasses import dataclass +from typing import Callable + +from . import _inductor_cache_dir_default, quantized_cache +from .security import SecurityError, validate_path + +logger = logging.getLogger("dw") + +MARKER = ".dw-cache" +_CLEARING_PREFIX = ".clearing-" +_MARKER_TEXT = ( + "This directory is a cache dw may empty from its server page or " + "DELETE /api/caches. Delete this file to stop that.\n" +) + +# One clear at a time: a second one would find the first's staging directory +# among the root's children and move it out from under its rmtree +_clear_lock = threading.Lock() + + +class UnknownCache(ValueError): + pass + + +class NotClearable(Exception): + pass + + +class NoEntries(ValueError): + """An entry operation on a cache that has no entries, or an entry id + that isn't one (a 400).""" + + +class UnknownEntry(ValueError): + pass + + +@dataclass(frozen=True) +class DiskCache: + name: str + # The absolute root, from the environment as it stands + root: Callable[[], str] + # The root with the member's env var unset - only for grandfathering + default_root: Callable[[], str] + enabled: Callable[[], bool] + cost_of_clearing: str + # Whether the cache lists entries of its own: `/<64-hex id>/` + # directories, each described by an `entry.json` + has_entries: bool = False + # Whether mark_roots may create a missing root + create_root: bool = True + # The free space under which the cache's volume reads `low_disk`, from + # the volume's total - for a cache that refuses to write below a floor + low_disk_floor: Callable[[int], int] | None = None + + +def _inductor_root(): + return os.path.abspath( + os.environ.get("TORCHINDUCTOR_CACHE_DIR") + or _inductor_cache_dir_default(os.environ) + ) + + +def _inductor_default(): + return os.path.abspath(_inductor_cache_dir_default(os.environ)) + + +def _triton_default(): + return os.path.abspath(os.path.join(os.path.expanduser("~"), ".triton", "cache")) + + +def _triton_root(): + """Where Triton caches, as triton.knobs resolves it - without importing + triton, which the server process never loads and macOS doesn't have.""" + explicit = os.environ.get("TRITON_CACHE_DIR", "").strip() + if explicit: + return os.path.abspath(os.path.expanduser(explicit)) + home = os.environ.get("TRITON_HOME", "").strip() + if home: + return os.path.abspath( + os.path.join(os.path.expanduser(home), ".triton", "cache") + ) + return _triton_default() + + +def _triton_installed(): + # find_spec on a top-level name looks for it without importing it. The + # workers run in this interpreter's environment, so this is theirs too + return importlib.util.find_spec("triton") is not None + + +CACHES = { + cache.name: cache + for cache in ( + DiskCache( + name="inductor", + root=_inductor_root, + default_root=_inductor_default, + enabled=lambda: True, + cost_of_clearing=( + "The next first step of each torch.compile'd shape compiles " + "its kernels again, up to a few minutes per shape." + ), + ), + DiskCache( + name="triton", + root=_triton_root, + default_root=_triton_default, + enabled=_triton_installed, + cost_of_clearing=( + "The next first step of each SDNQ or compiled shape " + "re-autotunes its Triton kernels (100-300 s on MiniMax-H3). " + "This directory is shared with any other Triton user on " + "this machine, whose kernels are cleared too." + ), + # A box without Triton gets no stray directory + create_root=False, + ), + DiskCache( + name="quantized", + root=quantized_cache.cache_root, + default_root=quantized_cache.default_root, + enabled=quantized_cache.enabled, + cost_of_clearing=( + "The next load of each SDNQ-quantized component reads its " + "bf16 weights and quantizes them again (minutes per component " + "on MiniMax-H3), and saves a new entry while the cache is on." + ), + has_entries=True, + # configure_at_startup creates it when the cache is on; a server + # with it off gets no stray directory + create_root=False, + low_disk_floor=quantized_cache.floor_bytes, + ), + ) +} + + +def get(name): + try: + return CACHES[name] + except (KeyError, TypeError): + raise UnknownCache( + f"Unknown cache {name!r} - the caches are: {', '.join(CACHES)}" + ) from None + + +def _children(root): + """The root's entries other than the marker.""" + with os.scandir(root) as entries: + return [entry for entry in entries if entry.name != MARKER] + + +def _has_marker(root): + return os.path.isfile(os.path.join(root, MARKER)) and not os.path.islink( + os.path.join(root, MARKER) + ) + + +def _write_marker(root): + marker = validate_path(os.path.join(root, MARKER), base_dir=root) + with open(marker, "w") as f: + f.write(_MARKER_TEXT) + + +def mark_root(cache): + """Mark one cache's root as dw's to clear, if the rule allows it. + Returns whether the root carries the marker afterwards.""" + root = cache.root() + if os.path.islink(root) or (os.path.exists(root) and not os.path.isdir(root)): + return False + if not os.path.exists(root): + if not cache.create_root: + return False + os.makedirs(root, exist_ok=True) + if _has_marker(root): + return True + if _children(root) and root != cache.default_root(): + logger.info( + f"The {cache.name} cache at {root} was not created by dw and " + f"has no {MARKER} marker, so it can't be cleared from the server" + ) + return False + _write_marker(root) + return True + + +def mark_roots(): + """Write the ownership marker into each cache root the rule allows. Run + once at server start; a root it can't write is left unclearable.""" + for cache in CACHES.values(): + try: + mark_root(cache) + except (OSError, SecurityError) as e: + logger.warning(f"Could not mark the {cache.name} cache: {e}") + + +def _walk(top): + """(bytes, files, newest mtime) of the regular files under `top`, + symlinks neither followed nor counted, the marker not counted.""" + total, files, newest = 0, 0, None + stack = [top] + while stack: + directory = stack.pop() + try: + entries = os.scandir(directory) + except OSError: + continue + with entries: + for entry in entries: + try: + if entry.is_dir(follow_symlinks=False): + stack.append(entry.path) + continue + if not entry.is_file(follow_symlinks=False): + continue + if entry.name == MARKER and directory == top: + continue + stat = entry.stat(follow_symlinks=False) + except OSError: + # A file deleted mid-walk (a racing compile's temp file) + continue + total += stat.st_size + files += 1 + if newest is None or stat.st_mtime > newest: + newest = stat.st_mtime + return total, files, newest + + +def _volume(path): + """shutil.disk_usage of the volume `path` is (or would be) on: a missing + root is measured at its nearest existing ancestor.""" + while path and not os.path.exists(path): + parent = os.path.dirname(path) + if parent == path: + return None + path = parent + try: + return shutil.disk_usage(path) + except OSError: + return None + + +def _reason(cache): + """Why `cache` can't be cleared now, or None. A root that came into being + after server start (Triton's, made by the first compile) is marked here + by the same rule mark_roots applies at startup.""" + root = cache.root() + if os.path.isdir(root) and not _has_marker(root): + try: + mark_root(cache) + except (OSError, SecurityError) as e: + logger.warning(f"Could not mark the {cache.name} cache: {e}") + return _not_clearable_reason(root) + + +def _not_clearable_reason(root): + if not os.path.isdir(root) or os.path.islink(root): + return "root does not exist (nothing to clear)" + if not _has_marker(root): + return f"root not created by dw (no {MARKER} marker)" + return None + + +def scan(cache): + """One cache as plain data.""" + root = cache.root() + present = os.path.isdir(root) and not os.path.islink(root) + total, files, newest = _walk(root) if present else (0, 0, None) + usage = _volume(root) + reason = _reason(cache) + info = { + "name": cache.name, + "root": root, + "enabled": cache.enabled(), + "bytes": total, + "files": files, + "last_used": newest, + "clearable": reason is None, + "cost_of_clearing": cache.cost_of_clearing, + "disk_free": usage.free if usage else None, + "disk_total": usage.total if usage else None, + } + if reason: + info["not_clearable_reason"] = reason + if cache.low_disk_floor and usage: + info["low_disk"] = usage.free < cache.low_disk_floor(usage.total) + if cache.has_entries: + info["entries"] = scan_entries(root) if present else [] + return info + + +def _entry_fields(path): + """What an entry's `entry.json` says of it, or None when it can't be + read. `component` falls back to the subfolder - the component's name in + a diffusers repo - since entry.json doesn't record it.""" + try: + entry = quantized_cache.read_entry(path) + except (OSError, ValueError): + return None + return { + "model": entry.get("model_name"), + "subfolder": entry.get("subfolder"), + "component": entry.get("component") or entry.get("subfolder") or None, + "created": entry.get("created"), + "last_hit": entry.get("last_hit"), + } + + +def scan_entries(root): + """Each `<64-hex id>/` directory under a cache root, newest first. An + entry whose `entry.json` is missing or unreadable is listed with + `unreadable: true`, so it can still be deleted.""" + entries = [] + with os.scandir(root) as children: + for child in children: + if not quantized_cache.is_key(child.name) or not child.is_dir( + follow_symlinks=False + ): + continue + total, files, newest = _walk(child.path) + fields = _entry_fields(child.path) + entry = {"id": child.name, "bytes": total, "files": files} + if fields is None: + entry.update( + unreadable=True, + model=None, + subfolder=None, + component=None, + created=newest, + last_hit=None, + ) + else: + entry.update(unreadable=False, **fields) + entries.append(entry) + entries.sort(key=lambda e: e["created"] or 0, reverse=True) + return entries + + +def scan_all(): + return [scan(cache) for cache in CACHES.values()] + + +def entry_cache(name, entry): + """The cache an entry route names, after checking the id's shape - so + nothing builds a path from a malformed one. Raises UnknownCache, or + NoEntries for a cache without entries or an id that isn't 64 hex.""" + cache = get(name) + if not cache.has_entries: + raise NoEntries(f"The {name} cache has no entries - clear it whole") + if not quantized_cache.is_key(entry): + raise NoEntries(f"Not an entry id: {entry!r} (64 lowercase hex characters)") + return cache + + +def _stage_out(root, sources): + """Move `sources` (paths directly under `root`) into one fresh + `.clearing-*` directory and delete it. Returns (bytes, files) removed.""" + staging = validate_path( + os.path.join(root, f"{_CLEARING_PREFIX}{os.getpid()}-{secrets.token_hex(4)}"), + base_dir=root, + ) + os.mkdir(staging) + try: + for path in sources: + try: + source = validate_path(path, base_dir=root) + target = validate_path( + os.path.join(staging, os.path.basename(path)), base_dir=root + ) + except SecurityError: + # A name validate_path refuses (one holding "..") stays + continue + try: + os.rename(source, target) + except FileNotFoundError: + continue + freed, files, _ = _walk(staging) + finally: + shutil.rmtree(staging, ignore_errors=True) + return freed, files + + +def delete_entry(name, entry): + """Delete one entry of a marked cache root. Returns `{freed, files}`. + + Raises UnknownCache, NoEntries (no entries here, or a malformed id), + NotClearable for a root without the marker, or UnknownEntry - each + before anything is moved.""" + cache = entry_cache(name, entry) + with _clear_lock: + root = cache.root() + reason = _reason(cache) + if reason: + raise NotClearable(f"The {name} cache at {root} can't be changed: {reason}") + root = validate_path(root) + path = quantized_cache.entry_path(root, entry) + if os.path.islink(path) or not os.path.isdir(path): + raise UnknownEntry(f"The {name} cache has no entry {entry}") + freed, files = _stage_out(root, [path]) + logger.info(f"Deleted {name} cache entry {entry} ({freed} bytes, {files} files)") + return {"freed": freed, "files": files} + + +def clear(name): + """Empty a marked cache root, keeping the root and its marker. Returns + `{freed, files}` for what was removed. + + Raises UnknownCache for a name not in the registry and NotClearable for a + root without the marker - before anything is moved.""" + cache = get(name) + with _clear_lock: + root = cache.root() + reason = _reason(cache) + if reason: + raise NotClearable(f"The {name} cache at {root} can't be cleared: {reason}") + root = validate_path(root) + # A link is left alone: validate_path resolves it, so the rename + # would move its target, not the link + freed, files = _stage_out( + root, [entry.path for entry in _children(root) if not entry.is_symlink()] + ) + logger.info(f"Cleared the {name} cache at {root} ({freed} bytes, {files} files)") + return {"freed": freed, "files": files} diff --git a/dw/dsp.py b/dw/dsp.py index 2db5a6c8..b7e77f7d 100644 --- a/dw/dsp.py +++ b/dw/dsp.py @@ -110,6 +110,25 @@ def rms(window): return math.sqrt(float(numpy.mean(numpy.square(window, dtype=numpy.float64)))) +def loudest_window_dbfs(mono, sample_rate, start, end, window_seconds): + """The loudest `window_seconds` window's rms over [start, end] seconds, in + dBFS, or None for silence. A span shorter than one window (a zero-length + word chunk) is measured as one window centred on it.""" + window = max(1, int(round(window_seconds * sample_rate))) + first = int(round(start * sample_rate)) + last = int(round(end * sample_rate)) + if last - first < window: + middle = (first + last) // 2 + first, last = middle - window // 2, middle - window // 2 + window + first, last = max(0, first), min(len(mono), last) + loudest = None + for offset in range(first, last, window): + level = rms(mono[offset : min(offset + window, last)]) + if level is not None and (loudest is None or level > loudest): + loudest = level + return dbfs(loudest) + + def peak(window): """The largest absolute sample of a window, or None when it has none.""" if window is None or window.size == 0: diff --git a/dw/media.py b/dw/media.py index 2dbe87e8..38780873 100644 --- a/dw/media.py +++ b/dw/media.py @@ -653,6 +653,19 @@ def _stream_info(container, audio_stream_seconds=False): return info, video, audio +# The finest envelope bin, and the most bins one answer carries, so a long +# track at a fine interval stays inside the MCP response budget (#823) +DEFAULT_ENVELOPE_INTERVAL = 1.0 +MIN_ENVELOPE_INTERVAL = 0.05 +# One bin an hour: past this the bin-sample product overflows (so does an +# infinite or NaN interval, which pydantic reads from a query string) +MAX_ENVELOPE_INTERVAL = 3600.0 +# The bin count a finer-than-default envelope may reach. The default interval +# is never refused by it: a track long enough to exceed it at 1.0 s answered +# before the cap existed and still does (#823) +MAX_ENVELOPE_BINS = 2400 + + class _AudioAccumulator: """A soundtrack's levels, gathered frame by frame as the file decodes. @@ -665,13 +678,21 @@ class _AudioAccumulator: against, so the trailing-fragment merge is the fallback. """ - def __init__(self, rate, channels, duration_seconds=None, envelope=False): + def __init__( + self, + rate, + channels, + duration_seconds=None, + envelope=False, + interval_seconds=1.0, + ): self.rate = rate + self.bin_samples = _bin_samples(rate, interval_seconds) self.channels = int(channels) self.peak = 0.0 self.total = 0.0 self.count = 0 - # One bin per second of the soundtrack: [sum of squares, sample + # One bin per interval of the soundtrack: [sum of squares, sample # count, peak] - the same numbers the whole-track level is made of, # kept per second instead of once self.bins = [] if envelope else None @@ -698,7 +719,7 @@ def add(self, decoded): else max(0, min(full_length, self.max_samples - self.elapsed)) ) if self.bins is not None: - _fill_envelope(self.bins, frame[:length], self.elapsed, self.rate) + _fill_envelope(self.bins, frame[:length], self.elapsed, self.bin_samples) if length > 0: self.chunks.append(frame[:length]) # the frame's *full* length, so a later frame's position is measured @@ -707,7 +728,7 @@ def add(self, decoded): def finish(self): if self.bins is not None and self.max_samples is None: - _merge_trailing_fragment(self.bins, self.rate) + _merge_trailing_fragment(self.bins, self.bin_samples) rms = math.sqrt(self.total / self.count) if self.count else 0.0 full = numpy.concatenate(self.chunks, axis=0) if self.chunks else None levels = { @@ -717,11 +738,11 @@ def finish(self): "true_peak_dbfs": true_peak_dbfs(full), } if self.bins is not None: - levels["envelope"] = _as_envelope(self.bins) + levels["envelope"] = _as_envelope(self.bins, self.bin_samples / self.rate) return levels -def probe_media(path, envelope=False): +def probe_media(path, envelope=False, interval_seconds=1.0): """Duration, format and level of an audio or video file, or None. Video answers fps, frame_count, width and height, plus the soundtrack's @@ -757,6 +778,12 @@ def probe_media(path, envelope=False): Args: path: The file to probe + interval_seconds: The envelope's bin width, `MIN_ENVELOPE_INTERVAL` + (0.05 s) to `MAX_ENVELOPE_INTERVAL` (3600 s); 0.1 places a cut on + a vocal onset where 1.0 cannot. A finer-than-default width that + would take more than `MAX_ENVELOPE_BINS` bins raises ValueError + naming the width that fits; the default width is never refused + for length (#823) envelope: Also report the per-second level of the soundtrack. A lossy codec's decode can run a fraction of a second past the file's reported duration - its own priming and padding - so samples past @@ -784,9 +811,15 @@ def probe_media(path, envelope=False): if not (need_frame_count or audio is not None): return info + if envelope and audio is not None: + _check_envelope_size(info.get("duration_seconds"), interval_seconds) accumulator = ( _AudioAccumulator( - audio.rate, audio.channels, info.get("duration_seconds"), envelope + audio.rate, + audio.channels, + info.get("duration_seconds"), + envelope, + interval_seconds, ) if audio is not None else None @@ -903,12 +936,12 @@ def probe_metadata(path): return info -def _fill_envelope(bins, frame, elapsed, rate): +def _fill_envelope(bins, frame, elapsed, bin_samples): """Add a decoded audio frame's samples, as (samples, channels), to the - per-second bins. + envelope bins, each `bin_samples` long. - A bin covers one second of the track regardless of how the decoder - happened to chop it, so a frame straddling a second boundary is split + A bin covers one interval of the track regardless of how the decoder + happened to chop it, so a frame straddling a bin boundary is split across the two bins rather than counted in whichever one it started in. `elapsed` is how many samples of the track came before this frame; the caller has already cut away anything past the file's reported duration (a @@ -917,11 +950,11 @@ def _fill_envelope(bins, frame, elapsed, rate): length = frame.shape[0] start = 0 while start < length: - second = (elapsed + start) // rate + second = (elapsed + start) // bin_samples while len(bins) <= second: bins.append([0.0, 0, 0.0]) # How much of this frame still belongs to the second it is in - room = int((second + 1) * rate - (elapsed + start)) + room = int((second + 1) * bin_samples - (elapsed + start)) stop = min(length, start + max(room, 1)) piece = frame[start:stop] entry = bins[second] @@ -931,7 +964,7 @@ def _fill_envelope(bins, frame, elapsed, rate): start = stop -def _merge_trailing_fragment(bins, rate): +def _merge_trailing_fragment(bins, bin_samples): """Fold a short trailing bin into the one before it, in place. Fallback for when the file's duration is unknown, so `_fill_envelope` @@ -951,17 +984,55 @@ def _merge_trailing_fragment(bins, rate): if len(bins) < 2: return total, count, peak = bins[-1] - if count >= rate: + if count >= bin_samples: return prev_total, prev_count, prev_peak = bins[-2] bins[-2] = [prev_total + total, prev_count + count, max(prev_peak, peak)] bins.pop() -def _as_envelope(bins): - """The per-second bins as the levels an agent reads.""" +def _bin_samples(rate, interval_seconds): + """Samples in one envelope bin: the interval, rounded to a whole sample.""" + return max(1, int(round(interval_seconds * rate))) + + +def _check_envelope_size(duration_seconds, interval_seconds): + """Refuse an envelope that is too fine or too long to come back whole. + + Raises ValueError, which the route turns into a 422, naming the bin width + that fits and the window-free alternative of a coarser interval (#823). + """ + if not interval_seconds >= MIN_ENVELOPE_INTERVAL: + raise ValueError( + f"interval_seconds must be at least {MIN_ENVELOPE_INTERVAL}, " + f"got {interval_seconds}" + ) + if not interval_seconds <= MAX_ENVELOPE_INTERVAL: + # Also catches inf and NaN, which no <= holds for + raise ValueError( + f"interval_seconds must be at most {MAX_ENVELOPE_INTERVAL:.0f}, " + f"got {interval_seconds}" + ) + if duration_seconds is None: + return + bins = math.ceil(duration_seconds / interval_seconds) + # Never fewer bins than the default interval gives + limit = max( + MAX_ENVELOPE_BINS, math.ceil(duration_seconds / DEFAULT_ENVELOPE_INTERVAL) + ) + if bins > limit: + fits = math.ceil(duration_seconds / limit * 100) / 100 + raise ValueError( + f"a {duration_seconds:.1f}s track at interval_seconds=" + f"{interval_seconds} is {bins} bins, over the {limit} " + f"limit; use interval_seconds={fits} or more" + ) + + +def _as_envelope(bins, interval_seconds=1.0): + """The bins as the levels an agent reads.""" return { - "interval_seconds": 1.0, + "interval_seconds": interval_seconds, "rms_dbfs": [ dbfs(math.sqrt(total / count) if count else 0.0, floor=SILENCE_DBFS) for total, count, _peak in bins diff --git a/dw/pipeline_processors/components.py b/dw/pipeline_processors/components.py index e5757543..a0aeb238 100644 --- a/dw/pipeline_processors/components.py +++ b/dw/pipeline_processors/components.py @@ -13,6 +13,7 @@ from .. import get_device_type, resolve_device from ..events import WorkflowCancelled, emit_phase from .. import download_watch +from .. import quantized_cache from huggingface_hub.errors import HfHubHTTPError from .h3_guides import insert_guides from .h3_hold import insert_audio_hold @@ -590,12 +591,25 @@ def load_component( model_name = from_pretrained_arguments.pop("model_name") logger.info(f"Loading {component_name} from model: {model_name}") emit_phase("loading", detail=f"{component_name}: {model_name}") - with download_watch.watch( - model_name, cache_dir=from_pretrained_arguments.get("cache_dir") - ): - component = component_type.from_pretrained( - model_name, **from_pretrained_arguments - ) + + def from_pretrained(): + with download_watch.watch( + model_name, + cache_dir=from_pretrained_arguments.get("cache_dir"), + ): + return component_type.from_pretrained( + model_name, **from_pretrained_arguments + ) + + # A component an SDNQ config quantizes is read from the cache on + # a hit, and neither loads nor downloads (dw/quantized_cache.py) + component = quantized_cache.from_pretrained( + component_name, + component_type, + model_name, + from_pretrained_arguments, + from_pretrained, + ) # Load from single file elif "from_single_file" in from_pretrained_arguments: @@ -636,7 +650,9 @@ def load_component( f"{component_type.__name__} does not have it" ) logger.info(f"Loading components for {component_name}") - component.load_components(**load_components_arguments) + # Before pruning, offload and sdnq_optimize, so a cached entry + # is the whole component as quantized (dw/quantized_cache.py) + quantized_cache.load_components(component, load_components_arguments) # MiniMax-H3 takes `hold_audio` on every core-denoise workflow; a # no-op on any other pipeline (dw/pipeline_processors/h3_hold.py) diff --git a/dw/quantized_cache.py b/dw/quantized_cache.py new file mode 100644 index 00000000..16224913 --- /dev/null +++ b/dw/quantized_cache.py @@ -0,0 +1,917 @@ +"""On-disk cache of SDNQ on-the-fly-quantized components (#815). + +A component an SDNQ config quantizes as it loads goes through the same steps on +every cold load: read the bf16 weights, copy them to the accelerator, quantize, +copy them back. The same config yields the same int4/int8 weights each time, so +this keeps them: the first load saves the quantized component with sdnq's own +`save_sdnq_model` (safetensors plus a JSON config - no pickle), and a later load +of the same config reads it back with `load_sdnq_model` instead. + +Off unless `DW_QUANTIZED_CACHE=1`. `DW_QUANTIZED_CACHE_DIR` moves the root from +`/dw-quantized`; `dw.serve` resolves it to an absolute path once and +pins it in the environment its workers inherit (`configure_at_startup`). + +**The key** is a sha256 over the component's repo, subfolder, the commit of the +snapshot the load reads, the quantizer config minus the arguments that do not +change the stored weights (the devices and the quantized-matmul switches), +`torch_dtype`, and the installed sdnq/diffusers/transformers/torch versions. A +stale hit silently gives the wrong weights, so the key errs complete: anything +it cannot pin (a local path, a snapshot whose files are not all on disk, a +commit the loaded object and the hub cache disagree on) is not cached. + +**The entry** is `/<64-hex key>/`: sdnq's files plus `entry.json`, the +key's inputs in readable form with the size and the created and last-hit +times. It is written to a `.tmp-*` sibling and renamed into place, so a +reader never sees half an entry and two workers saving the same key at once +leave one. Pruning, `sdnq_optimize`, LoRAs and offload hooks all run after the +save point, so an entry holds the untruncated component as quantized. + +**Two load paths** reach it, both from `load_component` +(`dw/pipeline_processors/components.py`): a modular pipeline's +`load_components()` (H3, `load_components` here), and one component's own +`from_pretrained` (LTX-2.5, `from_pretrained` here). A hit is placed on the +config's `return_device`, which a pre-quantized load would otherwise ignore. + +**Limits:** a save that would leave less free than `min(50 GiB, 10% of the +volume)` is skipped, and the server-start sweep deletes every entry built +against other package versions. Nothing else evicts; clearing is +from the server page or `DELETE /api/caches/quantized` (see dw/disk_caches.py). + +Every lookup and save emits a `quantized_cache` job event, +`{component, result: hit|miss|saved|skipped, reason?, seconds, bytes}`, and +none of it can fail a job: a cache that cannot help steps aside. +""" + +import errno +import hashlib +import importlib.metadata +import inspect +import json +import logging +import os +import re +import secrets +import shutil +import time +from functools import lru_cache, partial +from types import SimpleNamespace + +from .events import get_context +from .security import validate_path, SecurityError + +logger = logging.getLogger("dw") + +ENABLE_ENV_VAR = "DW_QUANTIZED_CACHE" +DIR_ENV_VAR = "DW_QUANTIZED_CACHE_DIR" +DEFAULT_DIR_NAME = "dw-quantized" +ENTRY_FILE = "entry.json" +EVENT = "quantized_cache" + +# A save never takes the volume below the smaller of these: on a large disk +# 50 GiB is room for a model download, on a small one 10% is +LOW_DISK_FLOOR_BYTES = 50 * 2**30 +LOW_DISK_FLOOR_FRACTION = 0.10 + +VERSIONED_PACKAGES = ("sdnq", "diffusers", "transformers", "torch") + +# Quantizer arguments that place or run the weights rather than shape them - +# the same int4 tensors come out whichever card quantized them, and a hit +# re-applies them at load +UNKEYED_ARGUMENTS = ( + "quantization_device", + "return_device", + "use_quantized_matmul", + "use_quantized_matmul_conv", + "non_blocking", +) + +# Quantizer arguments that are sets of module names held as lists (or a dict +# of them). SDNQConfig dedupes each through `list(set(...))`, so their order +# follows the process's string-hash seed - one order per worker - and the key +# sorts them, or each card would key the same config differently +SET_ARGUMENTS = ( + "modules_to_not_convert", + "modules_to_not_use_matmul", + "modules_dtype_dict", +) + +# Reasons a `skipped` event carries +LOCAL_PATH = "local_path" +UNPINNED_SNAPSHOT = "unpinned_snapshot" +LOW_DISK = "low_disk" +UNSUPPORTED = "unsupported" +ALREADY_SAVED = "already saved by another worker" + +# Reasons a `miss` event carries: why the load found nothing to read +INCOMPLETE_SNAPSHOT = "incomplete_snapshot" +NO_ENTRY = "no_entry" +LOAD_FAILED = "load_failed" + +_KEY_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_COMMIT_PATTERN = re.compile(r"^[0-9a-f]{40}$") +_TMP_MARKER = ".tmp-" +# How much of a failed load's message a `load_failed` reason carries +_REASON_DETAIL_CHARS = 200 + + +def enabled(): + """Whether the cache reads, writes and sweeps - only `DW_QUANTIZED_CACHE=1`.""" + return os.environ.get(ENABLE_ENV_VAR) == "1" + + +def default_root(): + """The root with `DW_QUANTIZED_CACHE_DIR` unset: `HF_HOME/dw-quantized`.""" + from huggingface_hub import constants + + return os.path.abspath(os.path.join(constants.HF_HOME, DEFAULT_DIR_NAME)) + + +def cache_root(): + """The cache's root directory, absolute.""" + configured = os.environ.get(DIR_ENV_VAR) + if configured: + return os.path.abspath(os.path.expanduser(configured)) + return default_root() + + +def is_key(name): + """Whether `name` has an entry key's shape (64 lowercase hex) - the check + to make before building any path from it.""" + return isinstance(name, str) and bool(_KEY_PATTERN.match(name)) + + +@lru_cache(maxsize=1) +def installed_versions(): + return { + package: importlib.metadata.version(package) for package in VERSIONED_PACKAGES + } + + +def is_sdnq_config(config): + return type(config).__name__ == "SDNQConfig" and type(config).__module__.startswith( + "sdnq" + ) + + +def emit(component, result, reason=None, seconds=None, nbytes=None): + data = { + "component": component, + "result": result, + "seconds": None if seconds is None else round(seconds, 2), + "bytes": nbytes, + } + if reason is not None: + data["reason"] = reason + logger.info( + f"Quantized cache: {component} {result}" + (f" ({reason})" if reason else "") + ) + get_context().emit(EVENT, **data) + + +# --- the key --------------------------------------------------------------- + + +def quantizer_arguments(config): + """The quantizer config as keyed: its class and every argument that shapes + the stored weights. Read before the load - sdnq clears the device fields + on the config it quantized with.""" + arguments = config_snapshot(config) + for name in UNKEYED_ARGUMENTS: + arguments.pop(name, None) + for name in SET_ARGUMENTS: + value = arguments.get(name) + if isinstance(value, list): + arguments[name] = sorted(value, key=str) + elif isinstance(value, dict): + arguments[name] = { + k: sorted(v, key=str) if isinstance(v, list) else v + for k, v in value.items() + } + return { + "class": f"{type(config).__module__}.{type(config).__qualname__}", + "arguments": arguments, + } + + +def config_snapshot(config): + """A deep copy of `config.to_dict()` in plain JSON types. `to_dict()` + hands back the config's own lists and dicts, and sdnq's load writes into + them (it appends the modules it left unquantized to + `modules_to_not_convert`), so a shallow read taken before the load turns + into the config after it - the save's key then differs from the next + lookup's, and from one card's load to another's.""" + return json.loads(json.dumps(config.to_dict(), sort_keys=True, default=str)) + + +def key_inputs(model_name, subfolder, commit, quantizer, torch_dtype): + return { + "model_name": model_name, + "subfolder": subfolder or "", + "commit": commit, + "quantizer": quantizer, + "torch_dtype": None if torch_dtype is None else str(torch_dtype), + "versions": installed_versions(), + } + + +def cache_key(inputs): + canonical = json.dumps(inputs, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +# --- the snapshot a load reads ---------------------------------------------- + + +def _hub_cache_dir(cache_dir=None): + if cache_dir: + return os.path.abspath(str(cache_dir)) + from huggingface_hub import constants + + return os.path.abspath(constants.HF_HUB_CACHE) + + +def is_local_path(model_name): + return os.path.isdir(str(model_name)) + + +def _cached_file(model_name, filename, revision, cache_dir): + """The path of `filename` in the hub cache's snapshot for `revision`, + resolved offline by huggingface_hub (a branch or tag through `refs/`, + the blob behind the symlink checked present), or None.""" + from huggingface_hub import try_to_load_from_cache + + try: + path = try_to_load_from_cache( + str(model_name), + filename, + cache_dir=_hub_cache_dir(cache_dir), + revision=revision, + ) + except Exception: + return None + return path if isinstance(path, str) else None + + +def _in_folder(subfolder, name): + return f"{subfolder}/{name}" if subfolder else name + + +def resolve_commit(model_name, revision=None, cache_dir=None, subfolder=None): + """The commit `revision` names in the hub cache, read offline: the + snapshot in which huggingface_hub finds `/config.json` for + it. None when the cache holds no such snapshot.""" + filename = _in_folder(subfolder, "config.json") + path = _cached_file(model_name, filename, revision or "main", cache_dir) + if path is None: + return None + commit = os.path.basename(os.path.normpath(path[: -len(filename)])) + return commit if _COMMIT_PATTERN.match(commit) else None + + +def _is_unvarianted_weights(name, suffix): + # 'model.safetensors' but not 'model.fp16.safetensors' - a variant is + # a file of its own, and a load without one never reads it + return name.endswith(suffix) and "." not in name[: -len(suffix)] + + +def snapshot_complete(model_name, commit, subfolder=None, cache_dir=None): + """Whether the snapshot at `commit` holds the component whole: its + `config.json` and every weight file it names, each present on disk. + huggingface_hub finds each file; which files a component needs (each + shard in the index's `weight_map`, or the unvarianted single file) is + dw's own rule.""" + from huggingface_hub.file_download import repo_folder_name + + hub = _hub_cache_dir(cache_dir) + + def present(name): + return _cached_file(model_name, _in_folder(subfolder, name), commit, hub) + + try: + # contained in the repo's own folder, symlinks resolved: a repo + # folder linked out to another disk keeps its blobs beside it + repo = os.path.join( + hub, repo_folder_name(repo_id=str(model_name), repo_type="model") + ) + config = present("config.json") + if config is None: + return False + names = os.listdir(validate_path(os.path.dirname(config), repo)) + indexes = [ + n for n in names if _is_unvarianted_weights(n, ".safetensors.index.json") + ] + if indexes: + for index in indexes: + index_path = present(index) + if index_path is None: + return False + with open(validate_path(index_path, repo)) as f: + weight_map = json.load(f).get("weight_map", {}) + if any(present(shard) is None for shard in set(weight_map.values())): + return False + return True + singles = [n for n in names if _is_unvarianted_weights(n, ".safetensors")] + return bool(singles) and all(present(n) is not None for n in singles) + except Exception: + return False + + +def lookup_commit(model_name, subfolder=None, revision=None, cache_dir=None): + """The commit to look a load up under, or None when its snapshot is not + complete on disk - the load is about to download or fail, and a key on + that commit would name files nobody read.""" + commit = resolve_commit(model_name, revision, cache_dir, subfolder) + if commit is None or not snapshot_complete( + model_name, commit, subfolder, cache_dir + ): + return None + return commit + + +def recorded_commit(component): + """The commit the library recorded on a loaded object, where it records + one (transformers' `config._commit_hash`).""" + config = getattr(component, "config", None) + commit = getattr(config, "_commit_hash", None) + if commit is None and isinstance(config, dict): + commit = config.get("_commit_hash") + return commit if isinstance(commit, str) and _COMMIT_PATTERN.match(commit) else None + + +def save_commit(component, model_name, subfolder=None, revision=None, cache_dir=None): + """The commit to save a just-loaded component under, or None when no one + snapshot can be named. It is the commit the same load's lookup would + name, so nothing is written that no later lookup reads; a recorded + commit that disagrees with it means the hub cache moved under the load.""" + commit = lookup_commit(model_name, subfolder, revision, cache_dir) + if commit is None: + return None + recorded = recorded_commit(component) + if recorded is not None and recorded != commit: + return None + return commit + + +# --- entries ---------------------------------------------------------------- + + +def entry_path(root, key): + if not is_key(key): + raise ValueError(f"not a cache key: {key!r}") + return validate_path(os.path.join(root, key), root) + + +def _tree_bytes(path): + total = 0 + for directory, _, files in os.walk(path): + for name in files: + try: + total += os.path.getsize(os.path.join(directory, name)) + except OSError: + pass + return total + + +def read_entry(path): + with open(os.path.join(path, ENTRY_FILE)) as f: + entry = json.load(f) + if not isinstance(entry, dict): + raise ValueError("entry.json is not an object") + return entry + + +def _write_json(path, data): + tmp = f"{path}{_TMP_MARKER}{os.getpid()}" + with open(tmp, "w") as f: + json.dump(data, f, indent=2, default=str) + os.replace(tmp, path) + + +def component_nbytes(component): + """The in-memory size of a component's parameters and buffers - what its + entry will take on disk, give or take headers and config.""" + total = 0 + for tensors in (component.parameters(), component.buffers()): + for tensor in tensors: + total += tensor.numel() * tensor.element_size() + return total + + +def floor_bytes(total): + return min(LOW_DISK_FLOOR_BYTES, int(LOW_DISK_FLOOR_FRACTION * total)) + + +def load_entry(root, key, model_cls, torch_dtype, device, config_arguments): + """Load a hit, placed on `device`. None on a miss. An entry that fails to + load is deleted, so the fresh quantize that follows can replace it, and + the error is raised on for the miss to name.""" + path = entry_path(root, key) + if not os.path.isdir(path): + return None + try: + from sdnq import load_sdnq_model + + component = load_sdnq_model( + path, + model_cls=model_cls, + dtype=torch_dtype, + device=device, + dequantize_fp32=config_arguments.get("dequantize_fp32"), + use_quantized_matmul=config_arguments.get("use_quantized_matmul"), + ) + component = component.to(device) + except Exception as e: + logger.warning( + f"Quantized cache entry {key} failed to load ({type(e).__name__}: {e}) " + "- deleting it and quantizing fresh" + ) + shutil.rmtree(path, ignore_errors=True) + raise + try: + entry = read_entry(path) + entry["last_hit"] = time.time() + _write_json(os.path.join(path, ENTRY_FILE), entry) + except Exception as e: + logger.debug(f"Could not record the hit on {key}: {e}") + return component + + +def save_entry(root, key, component, inputs): + """Save a just-quantized component as `key`. Returns (result, reason, + bytes) for the event: `saved`, or `skipped` with why. Never raises.""" + estimate = component_nbytes(component) + try: + os.makedirs(root, exist_ok=True) + usage = shutil.disk_usage(root) + except OSError as e: + logger.warning(f"Quantized cache: cannot reach {root}: {e}") + return "skipped", f"save_failed: {_error_name(e)}", None + floor = floor_bytes(usage.total) + if usage.free - estimate < floor: + logger.warning( + f"Quantized cache: not saving a {estimate / 2**30:.1f} GiB entry - " + f"{usage.free / 2**30:.1f} GiB free, and the floor is " + f"{floor / 2**30:.1f} GiB" + ) + return "skipped", LOW_DISK, None + + final = entry_path(root, key) + tmp = validate_path( + os.path.join(root, f"{key}{_TMP_MARKER}{os.getpid()}-{secrets.token_hex(4)}"), + root, + ) + try: + _save_model(component, tmp) + nbytes = _tree_bytes(tmp) + now = time.time() + _write_json( + os.path.join(tmp, ENTRY_FILE), + {**inputs, "key": key, "bytes": nbytes, "created": now, "last_hit": None}, + ) + try: + os.rename(tmp, final) + except OSError as e: + if e.errno in (errno.ENOTEMPTY, errno.EEXIST) and os.path.isdir(final): + shutil.rmtree(tmp, ignore_errors=True) + return "saved", ALREADY_SAVED, _tree_bytes(final) + raise + return "saved", None, nbytes + except Exception as e: + logger.warning( + f"Quantized cache: saving {key} failed ({type(e).__name__}: {e}) " + "- the job goes on without it" + ) + shutil.rmtree(tmp, ignore_errors=True) + return "skipped", f"save_failed: {_error_name(e)}", None + + +def _save_model(component, path): + """sdnq's `save_sdnq_model`, with a transformers model's tensors kept under + the names its own state_dict uses. transformers 5 saves in the checkpoint's + original layout by default (its weight conversion mapping run backwards), + while `load_sdnq_model` loads the names as they are, so a model with a + mapping - Gemma 4's vision tower, LTX-2.5's text encoder - saved that way + never loads back (#827). A diffusers model has no such switch.""" + from sdnq import save_sdnq_model + + save = getattr(component, "save_pretrained", None) + try: + own_names = "save_original_format" in inspect.signature(save).parameters + except (TypeError, ValueError): + own_names = False + if not own_names: + save_sdnq_model(component, path) + return + component.save_pretrained = partial(save, save_original_format=False) + try: + save_sdnq_model(component, path) + finally: + del component.save_pretrained + + +def _error_detail(error): + """`: `, cut short, for a reason that has + to say why a load failed, not just that it did.""" + lines = str(error).strip().splitlines() + if not lines: + return type(error).__name__ + detail = lines[0] + if len(detail) > _REASON_DETAIL_CHARS: + detail = detail[: _REASON_DETAIL_CHARS - 3] + "..." + return f"{type(error).__name__}: {detail}" + + +def _error_name(error): + code = getattr(error, "errno", None) + if code is not None and code in errno.errorcode: + return errno.errorcode[code] + return type(error).__name__ + + +# --- server start ----------------------------------------------------------- + + +def _child_path(root, name): + """The validated path of one child of the root, or None - logged once - + when it resolves outside the root (a symlink to another disk): that + child is left alone rather than failing the sweep, and with it the + server's start.""" + try: + return validate_path(os.path.join(root, name), root) + except SecurityError as e: + logger.warning( + f"Quantized cache entry {name} is skipped - it resolves outside " + f"the cache root: {e}" + ) + return None + + +def sweep(root): + """Delete every entry built against other package versions or with no + readable `entry.json`, and every leftover `*.tmp-*`. Returns (count, + bytes freed).""" + if not os.path.isdir(root): + return 0, 0 + versions = installed_versions() + count = freed = 0 + for name in os.listdir(root): + path = _child_path(root, name) + if path is None: + continue + if _TMP_MARKER in name: + stale = True + elif _KEY_PATTERN.match(name) and os.path.isdir(path): + try: + stale = read_entry(path).get("versions") != versions + except Exception: + stale = True + else: + continue + if not stale: + continue + size = _tree_bytes(path) if os.path.isdir(path) else os.path.getsize(path) + if os.path.isdir(path): + shutil.rmtree(path, ignore_errors=True) + else: + os.remove(path) + count += 1 + freed += size + return count, freed + + +def _entry_sizes(root): + sizes = [] + if not os.path.isdir(root): + return sizes + for name in os.listdir(root): + if not _KEY_PATTERN.match(name): + continue + path = _child_path(root, name) + if path is None: + continue + try: + size = int(read_entry(path)["bytes"]) + except Exception: + size = _tree_bytes(path) + sizes.append(size) + return sizes + + +def log_disk(root): + """One line for the cache's footprint and the volume's free space - a + warning when free space is already under the floor plus the largest entry + cached, so a full disk shows before the first save is skipped. Returns + whether it warned.""" + sizes = _entry_sizes(root) + usage = shutil.disk_usage(root) + floor = floor_bytes(usage.total) + margin = max(sizes, default=0) + message = ( + f"Quantized cache at {root}: {len(sizes)} entries, " + f"{sum(sizes) / 2**30:.1f} GiB in use, {usage.free / 2**30:.1f} GiB free" + ) + if usage.free < floor + margin: + logger.warning( + f"{message} - under the {floor / 2**30:.1f} GiB floor plus the largest " + f"entry ({margin / 2**30:.1f} GiB), so new entries may be skipped" + ) + return True + logger.info(message) + return False + + +def configure_at_startup(): + """Pin the root for the workers, sweep it and report the disk - at server + start, before any worker spawns. Nothing at all while the cache is off.""" + if not enabled(): + return + root = cache_root() + os.environ[DIR_ENV_VAR] = root + try: + os.makedirs(root, exist_ok=True) + count, freed = sweep(root) + logger.info( + f"Quantized cache sweep: deleted {count} stale entries, " + f"freed {freed / 2**30:.1f} GiB" + ) + log_disk(root) + except (OSError, SecurityError) as e: + logger.warning(f"Quantized cache at {root} is unusable: {e}") + + +# --- the load --------------------------------------------------------------- + +# ComponentSpec fields a load_components() keyword overrides, as diffusers +# reads them (ComponentSpec.load) +_SPEC_FIELDS = ("pretrained_model_name_or_path", "subfolder", "revision", "variant") + + +def _per_component(arguments, name): + """The keywords diffusers' load_components() hands one component: a dict + value is per name with a `default`, any other value is for all of them.""" + resolved = {} + for key, value in arguments.items(): + if not isinstance(value, dict): + resolved[key] = value + elif name in value: + resolved[key] = value[name] + elif "default" in value: + resolved[key] = value["default"] + return resolved + + +def _names_to_load(pipeline, arguments): + """The components load_components() is about to load, as it chooses them.""" + specs = pipeline._component_specs + names = arguments.get("names") + if names is None: + names = [ + name + for name, spec in specs.items() + if spec.default_creation_method == "from_pretrained" + and spec.pretrained_model_name_or_path is not None + and getattr(pipeline, name, None) is None + ] + elif isinstance(names, str): + names = [names] + workflow = arguments.get("workflow") + if workflow is not None: + expected = { + spec.name + for spec in pipeline.blocks.get_workflow(workflow).expected_components + } + names = [name for name in names if name in expected] + return [name for name in names if name in specs] + + +class _Candidate: + """One component an SDNQ config quantizes on this load, keyed.""" + + def __init__(self, name, spec, kwargs): + self.name = name + fields = {field: getattr(spec, field, None) for field in _SPEC_FIELDS} + fields.update({k: kwargs[k] for k in _SPEC_FIELDS if k in kwargs}) + self.model_name = fields["pretrained_model_name_or_path"] + self.subfolder = fields["subfolder"] + self.revision = fields["revision"] + self.variant = fields["variant"] + self.cache_dir = kwargs.get("cache_dir") + self.model_cls = spec.type_hint + self.load_id = getattr(spec, "load_id", None) + config = kwargs["quantization_config"] + self.torch_dtype = kwargs.get("torch_dtype", kwargs.get("dtype")) + # Read now: sdnq clears the devices on the config it quantized with + self.return_device = getattr(config, "return_device", None) or "cpu" + self.config_arguments = config_snapshot(config) + self.quantizer = quantizer_arguments(config) + + def unsupported(self): + if self.model_name is None or self.model_cls is None or self.variant: + return UNSUPPORTED + if is_local_path(self.model_name): + return LOCAL_PATH + return None + + def inputs(self, commit): + return key_inputs( + self.model_name, self.subfolder, commit, self.quantizer, self.torch_dtype + ) + + +def load_components(pipeline, arguments): + """`pipeline.load_components(**arguments)`, with every SDNQ-quantized + component read from the cache where an entry exists and saved to it where + one does not. With the cache off this is that call and nothing more.""" + if not enabled(): + pipeline.load_components(**arguments) + return + + try: + candidates = _candidates(pipeline, arguments) + except Exception as e: + logger.warning(f"Quantized cache: could not read the component specs: {e}") + candidates = [] + + root = cache_root() + hits = {} + misses = [] + for candidate in candidates: + reason = candidate.unsupported() + if reason is not None: + emit(candidate.name, "skipped", reason=reason) + continue + component, reason = _lookup(root, candidate) + if component is None: + misses.append((candidate, reason)) + continue + if candidate.load_id is not None: + # Keeps the spec update_components() would otherwise reset + component._diffusers_load_id = candidate.load_id + hits[candidate.name] = component + + if hits: + pipeline.update_components(**hits) + if arguments.get("names") is not None: + names = arguments["names"] + names = [names] if isinstance(names, str) else names + arguments = { + **arguments, + "names": [name for name in names if name not in hits], + } + + started = time.monotonic() + pipeline.load_components(**arguments) + seconds = time.monotonic() - started + + for candidate, reason in misses: + component = getattr(pipeline, candidate.name, None) + nbytes = _safe_nbytes(component) if component is not None else None + emit(candidate.name, "miss", reason=reason, seconds=seconds, nbytes=nbytes) + if component is None: + continue + _save(root, candidate, component) + + +# from_pretrained keywords a per-component load may carry and still be cached: +# the ones the key covers, and the ones that only say where the files are read +# from. Anything else could change the weights without changing the key +_PER_COMPONENT_KEYWORDS = frozenset( + { + "subfolder", + "revision", + "variant", + "cache_dir", + "torch_dtype", + "dtype", + "quantization_config", + "token", + "local_files_only", + "low_cpu_mem_usage", + } +) + + +def from_pretrained(name, model_cls, model_name, arguments, load): + """`load()` - one component's `model_cls.from_pretrained(model_name, + **arguments)` - with an SDNQ-quantized component read from the cache + where an entry exists, so `load` (and the download it may start) never + runs, and saved to it where one does not. With the cache off, or a load + no SDNQ config quantizes, this is `load()` and nothing more.""" + if not enabled() or not is_sdnq_config(arguments.get("quantization_config")): + return load() + + try: + spec = SimpleNamespace( + pretrained_model_name_or_path=model_name, + type_hint=model_cls, + load_id=None, + ) + candidate = _Candidate(name, spec, arguments) + reason = candidate.unsupported() + if reason is None and set(arguments) - _PER_COMPONENT_KEYWORDS: + reason = UNSUPPORTED + except Exception as e: + logger.warning(f"Quantized cache: could not key {name}: {e}") + return load() + if reason is not None: + emit(name, "skipped", reason=reason) + return load() + + root = cache_root() + component, reason = _lookup(root, candidate) + if component is not None: + return component + + started = time.monotonic() + component = load() + emit( + name, + "miss", + reason=reason, + seconds=time.monotonic() - started, + nbytes=_safe_nbytes(component), + ) + _save(root, candidate, component) + return component + + +def _lookup(root, candidate): + """(the cached component placed on its return device, None) on a hit, + emitting it; (None, why) on a miss, for the miss event to carry.""" + commit = lookup_commit( + candidate.model_name, + candidate.subfolder, + candidate.revision, + candidate.cache_dir, + ) + if commit is None: + # Not all on disk - the load downloads it, and the save after + # resolves the snapshot it ended up reading + return None, INCOMPLETE_SNAPSHOT + key = cache_key(candidate.inputs(commit)) + started = time.monotonic() + try: + component = load_entry( + root, + key, + candidate.model_cls, + candidate.torch_dtype, + candidate.return_device, + candidate.config_arguments, + ) + except Exception as e: + logger.warning(f"Quantized cache: lookup of {candidate.name} failed: {e}") + return None, f"{LOAD_FAILED}: {_error_detail(e)}" + if component is None: + return None, NO_ENTRY + emit( + candidate.name, + "hit", + seconds=time.monotonic() - started, + nbytes=_safe_nbytes(component), + ) + return component, None + + +def _candidates(pipeline, arguments): + quantization = arguments.get("quantization_config") + if quantization is None: + return [] + forwarded = {k: v for k, v in arguments.items() if k not in ("names", "workflow")} + candidates = [] + for name in _names_to_load(pipeline, arguments): + kwargs = _per_component(forwarded, name) + if not is_sdnq_config(kwargs.get("quantization_config")): + continue + candidates.append(_Candidate(name, pipeline._component_specs[name], kwargs)) + return candidates + + +def _safe_nbytes(component): + try: + return component_nbytes(component) + except Exception: + return None + + +def _save(root, candidate, component): + commit = save_commit( + component, + candidate.model_name, + candidate.subfolder, + candidate.revision, + candidate.cache_dir, + ) + if commit is None: + emit(candidate.name, "skipped", reason=UNPINNED_SNAPSHOT) + return + inputs = candidate.inputs(commit) + started = time.monotonic() + result, reason, nbytes = save_entry(root, cache_key(inputs), component, inputs) + emit( + candidate.name, + result, + reason=reason, + seconds=time.monotonic() - started, + nbytes=nbytes, + ) diff --git a/dw/script_lines.py b/dw/script_lines.py index e74fe268..5542f842 100644 --- a/dw/script_lines.py +++ b/dw/script_lines.py @@ -6,6 +6,11 @@ a cycle through `tasks.assess`, `tasks.audio_utils` and `tasks.joins` back into `task_domains` (where the check lived until #790). The task imports these names back, so a caller that reads them off `script_check` still finds them. + +It also holds what the tasks that read a transcript's words share without +importing each other (#822): the repetition-loop detector check_script +discards with and transcribe_audio marks `suspect: "repetition"` with, and +the run grouping plan_cuts warns about suspect chunks by. """ import re @@ -39,6 +44,67 @@ def normalize_words(text): return words +# A run of heard words that repeats with a period of at most +# REPEAT_MAX_PERIOD words, REPEAT_RUN_MIN or more times over, is a decoding +# loop - Whisper's hallucination over music or room tone ('Pre-pre-pre-...', +# 'thank you thank you ...'), loud enough to pass the energy guard. Speech +# repeats a word a few times; it does not loop. check_script discards such +# words and transcribe_audio marks them `suspect: "repetition"` (#822) +REPEAT_RUN_MIN = 6 +REPEAT_MAX_PERIOD = 4 + + +def repetition_loops(keys): + """The indices of `keys` (heard words) that sit in a repetition loop: a stretch that + repeats with a period of 1 to REPEAT_MAX_PERIOD and spans at least + REPEAT_RUN_MIN periods.""" + looped = set() + for period in range(1, REPEAT_MAX_PERIOD + 1): + k = period + while k < len(keys): + if keys[k] != keys[k - period]: + k += 1 + continue + first = k - period + while k < len(keys) and keys[k] == keys[k - period]: + k += 1 + if k - first >= REPEAT_RUN_MIN * period: + looped.update(range(first, k)) + return looped + + +def looped_texts(texts): + """The indices of `texts` (transcript chunks' text, in order) holding a word + of a repetition loop. Over words, not chunks: one chunk can hold a whole + loop ('Pre-pre-pre'), and a chunk with no word is never in one.""" + sequence, owner = [], [] + for index, text in enumerate(texts): + words = normalize_words(text) + owner.extend([index] * len(words)) + sequence.extend(words) + return {owner[k] for k in repetition_loops(sequence)} + + +# transcribe_audio's `suspect` reasons whose chunk plan_cuts leaves out of the +# sung lines: no word, a repetition loop, no voice under it. A "long" word is +# real lyric, only mistimed, so it stays +SUSPECT_LEFT_OUT = ("no_words", "repetition", "silent") + + +def suspect_runs(chunks): + """[(reason, [chunk, ...])]: each run of consecutive chunks carrying the + same `suspect` reason, so a long hallucination is one run, not hundreds.""" + runs, previous = [], None + for chunk in chunks: + reason = chunk.get("suspect") + if reason and reason == previous: + runs[-1][1].append(chunk) + elif reason: + runs.append((reason, [chunk])) + previous = reason + return runs + + def strip_markup(line): """A line with its H3 markup removed, and the words that markup held. diff --git a/dw/serve.py b/dw/serve.py index 7a2666c8..af52e9d1 100644 --- a/dw/serve.py +++ b/dw/serve.py @@ -327,6 +327,19 @@ def run(config): config.host, ) + # The quantized-component cache's root is pinned for the workers and swept + # here - after logging is up, so its lines reach the log, and before + # create_app spawns a worker that could be reading an entry it deletes + from .quantized_cache import configure_at_startup + + configure_at_startup() + + # Mark the cache roots dw may clear (dw/disk_caches.py) before a worker + # starts writing into a missing one + from .disk_caches import mark_roots + + mark_roots() + from .server.app import create_app from .server.app import default_ui_dir diff --git a/dw/server/api_models.py b/dw/server/api_models.py index 16b88c07..b661c37e 100644 --- a/dw/server/api_models.py +++ b/dw/server/api_models.py @@ -255,6 +255,65 @@ class ModelDeleted(ApiModel): freed: int = Field(description="Bytes freed.") +class DiskCacheEntry(ApiModel): + id: str = Field(description="The entry's 64-hex key: its directory name.") + bytes: int + files: int + unreadable: bool = Field( + description="True when its entry.json is missing or unreadable; the " + "fields below are then null, and it can still be deleted." + ) + model: str | None + subfolder: str | None + component: str | None + created: float | None = Field(description="Epoch seconds.") + last_hit: float | None = Field( + description="Epoch seconds of the last load that read it; null if none has." + ) + + +class DiskCacheInfo(ApiModel): + name: str + root: str + enabled: bool + bytes: int + files: int + last_used: float | None = Field( + description="The newest file's mtime (epoch seconds); null when empty." + ) + clearable: bool + not_clearable_reason: str | None = None + cost_of_clearing: str + disk_free: int | None + disk_total: int | None + low_disk: bool | None = Field( + default=None, + description="For a cache with a free-space floor (quantized): the " + "volume is under it, so new entries are not saved.", + ) + entries: list[DiskCacheEntry] | None = Field( + default=None, description="For a cache with entries (quantized)." + ) + + +class DiskCacheList(ApiModel): + caches: list[DiskCacheInfo] + + +class DiskCacheCleared(ApiModel): + name: str + cleared: bool + freed: int = Field(description="Bytes freed.") + files: int + + +class DiskCacheEntryDeleted(ApiModel): + name: str + entry: str + deleted: bool + freed: int = Field(description="Bytes freed.") + + class DiffusersInstall(ApiModel): version: str | None commit: str | None = Field(description="The git commit, for a git install.") diff --git a/dw/server/deps.py b/dw/server/deps.py index b183327f..2bdb8845 100644 --- a/dw/server/deps.py +++ b/dw/server/deps.py @@ -171,6 +171,24 @@ def ceiling_index(state, ws): return index +def _pipeline_steps(definition): + return sum( + 1 + for step in definition.get("steps") or [] + if isinstance(step, dict) and isinstance(step.get("pipeline"), dict) + ) + + +def _scales_past_template(definition, template, drivers, arguments): + """Whether `definition` runs the template's pipeline more times, or for + more frames, than the template run the bucket was measured on (#821).""" + if _pipeline_steps(definition) != _pipeline_steps(template): + return True + # The template as the bucket's drivers set it + bucket = {**template, "variables": {**(template.get("variables") or {}), **drivers}} + return "frame count" in inherited_differences(definition, bucket, arguments) + + def inherited_observed(state, ws, definition, arguments=None, *, card=None): """An inline workflow's `observed` block, priced by pipeline identity (#797): this box's runs of the catalog template that loads the same @@ -211,6 +229,13 @@ def inherited_observed(state, ws, definition, arguments=None, *, card=None): ) if not block or not block.get("cold_runs"): continue + # A figure is quoted only for a run of the template's own size: the + # bucket above prices the drivers the inline workflow shares by + # name, and anything else that moves the frame count or repeats the + # pipeline in more steps (#821) is a multiple the template's + # minutes do not carry - withheld rather than quoted unchanged + if _scales_past_template(definition, template, drivers, arguments): + continue key = (-block["cold_runs"], name) if best is None or key < best[0]: best = (key, name, template, block) diff --git a/dw/server/routes/media.py b/dw/server/routes/media.py index 1985b4d7..bf2c68aa 100644 --- a/dw/server/routes/media.py +++ b/dw/server/routes/media.py @@ -72,6 +72,7 @@ def gallery_metadata( request: Request, name: str, envelope: bool = False, + interval_seconds: float = 1.0, ws: Workspace = Depends(selected_workspace), ): """Generation metadata embedded in a saved image ('workflow' inside @@ -88,7 +89,11 @@ def gallery_metadata( is what says *where* in a track something is - whether a shot is still voiced at its last frame, how deep the hole at a seam goes. Opt-in: a ten-minute track is 600 numbers, and the default call has - to stay small. + to stay small. `interval_seconds` (default 1.0, 0.05 to 3600) is + the bin width: 0.1 places a cut on a vocal onset. A track that + would need more than 2400 bins at a finer-than-default width is + refused (422) with the interval that fits; the default width is + never refused for length (#823). `name` may also be an 'asset:' reference, and then it is the input asset of that name that is described rather than an output (#127). @@ -127,11 +132,14 @@ def gallery_metadata( job, run_id, version = job_provenance(request.app.state, name, ws) metadata = read_embedded_metadata(path) extension = os.path.splitext(path)[1].lower() - media = ( - probe_media(path, envelope=envelope) - if MEDIA_KINDS.get(extension) in ("audio", "video") - else None - ) + try: + media = ( + probe_media(path, envelope=envelope, interval_seconds=interval_seconds) + if MEDIA_KINDS.get(extension) in ("audio", "video") + else None + ) + except ValueError as error: + raise HTTPException(status_code=422, detail=str(error)) from error if media is not None and source == "output": # Where each shot of a joined video sits, as the run that wrote # it recorded (dw/shots.py) - null for a file not joined from shots diff --git a/dw/server/routes/system.py b/dw/server/routes/system.py index 5f053114..10711f72 100644 --- a/dw/server/routes/system.py +++ b/dw/server/routes/system.py @@ -20,6 +20,9 @@ PipelineNames, TaskList, DiffusersStatus, + DiskCacheCleared, + DiskCacheEntryDeleted, + DiskCacheList, HealthInfo, MemoryCleared, MemoryStatus, @@ -29,6 +32,7 @@ ModelDownloads, ServerInfo, ) +from ... import disk_caches from ...hub_cache import delete_model, scan_models from ...introspection import ( describe_class, @@ -255,6 +259,87 @@ def delete_cached_model(request: Request, repo: str): return {"repo_id": repo, "deleted": True, "freed": freed} +# ------------------------------------------------------ dw's disk caches + + +@router.get( + "/api/caches", response_model=DiskCacheList, response_model_exclude_unset=True +) +def list_caches(): + """dw's own disk caches (dw/disk_caches.py): each one's root, size, and + whether it can be cleared from here.""" + return {"caches": disk_caches.scan_all()} + + +@router.delete( + "/api/caches/{name}", + response_model=DiskCacheCleared, + response_model_exclude_unset=True, +) +def clear_cache(request: Request, name: str): + """Empty one cache's root, keeping the root and its marker. Refused + while anything may be writing into it (`_refuse_while_writing`).""" + try: + disk_caches.get(name) + except disk_caches.UnknownCache as e: + raise HTTPException(status_code=404, detail=str(e)) + _refuse_while_writing(request, name) + try: + result = disk_caches.clear(name) + except disk_caches.NotClearable as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"name": name, "cleared": True, **result} + + +@router.delete( + "/api/caches/{name}/entries/{entry}", + response_model=DiskCacheEntryDeleted, + response_model_exclude_unset=True, +) +def delete_cache_entry(request: Request, name: str, entry: str): + """Delete one entry of a cache that has entries (quantized). The id is + checked to be 64 hex before any path is built from it.""" + try: + disk_caches.entry_cache(name, entry) + except disk_caches.UnknownCache as e: + raise HTTPException(status_code=404, detail=str(e)) + except disk_caches.NoEntries as e: + raise HTTPException(status_code=400, detail=str(e)) + _refuse_while_writing(request, name) + try: + result = disk_caches.delete_entry(name, entry) + except disk_caches.UnknownEntry as e: + raise HTTPException(status_code=404, detail=str(e)) + except disk_caches.NotClearable as e: + raise HTTPException(status_code=400, detail=str(e)) + return {"name": name, "entry": entry, "deleted": True, "freed": result["freed"]} + + +# Caches that live beside the hub cache (under HF_HOME): changed only when no +# download is writing there either, the rule DELETE /api/models follows +_DOWNLOAD_SENSITIVE_CACHES = frozenset({"quantized"}) + + +def _refuse_while_writing(request, name): + """409 while a job is running or queued - a compile, autotune or + quantized save in progress writes into the root, and a delete mid-write + could leave a truncated artifact a later lookup trusts - and, for the + quantized cache, while a model download is active.""" + state = request.app.state + if state.job_manager.is_busy(): + raise HTTPException( + status_code=409, + detail="A job is running or queued - clearing a cache it may be " + "writing would leave it a truncated artifact", + ) + if name in _DOWNLOAD_SENSITIVE_CACHES and state.downloads.is_active(): + raise HTTPException( + status_code=409, + detail="A model download is in progress - the quantized cache " + "can't be changed until it finishes", + ) + + # ------------------------------------------------------ diffusers update diff --git a/dw/tasks/assess.py b/dw/tasks/assess.py index a820d62a..0f3a4bf3 100644 --- a/dw/tasks/assess.py +++ b/dw/tasks/assess.py @@ -386,7 +386,7 @@ def _shot_span_findings(probe, records, media): span reaches past the file - clipped silently otherwise (#425).""" findings = [] overrun_names = [] - for shot in records or []: + for index, shot in enumerate(records or []): overrun = _span_overrun(media, shot) if overrun is None: continue @@ -399,7 +399,11 @@ def _shot_span_findings(probe, records, media): { "rule": "shot_span_overrun", "severity": "warning", - "at": {"shot": shot.get("name")}, + "at": { + "shot": shot.get("name"), + "index": index, + "start_frame": shot.get("start_frame"), + }, "value": overrun, "threshold": 0, "says": f"shot {shot.get('name')!r} reaches {detail} past the file's end", @@ -519,12 +523,18 @@ def shots_answer(media, records, source): loudest = max(voiced, key=lambda shot: shot["rms_dbfs"]) quietest = min(voiced, key=lambda shot: shot["rms_dbfs"]) rms_range = _round(loudest["rms_dbfs"] - quietest["rms_dbfs"]) - at = {"between": [loudest["name"], quietest["name"]]} + # Names repeat once a cut re-joins earlier episodes, so the shots' + # positions in `shots` and their start frames say which is meant + at = { + "between": [loudest["name"], quietest["name"]], + "indices": [measured.index(loudest), measured.index(quietest)], + "start_frames": [loudest["start_frame"], quietest["start_frame"]], + } elif voiced: rms_range = 0.0 answer = {"shots": measured, "rms_range_db": rms_range, "has_audio": True} dead_air_findings = [] - for shot in measured: + for index, shot in enumerate(measured): # A shot that is quiet throughout, on purpose, holds no gap to find - # the same guard seam_hole applies to a seam's two sides (#465) quiet = shot["rms_dbfs"] is not None and shot["rms_dbfs"] <= HOLE_VOICED_DBFS @@ -532,7 +542,12 @@ def shots_answer(media, records, source): _findings( "analyze_shots", shot, - {"shot": shot["name"], "seconds": shot["dead_air_at"]}, + { + "shot": shot["name"], + "index": index, + "start_frame": shot["start_frame"], + "seconds": shot["dead_air_at"], + }, skip={"shot_dead_air"} if quiet else (), ) ) diff --git a/dw/tasks/audio_transcription.py b/dw/tasks/audio_transcription.py index 897b5e99..d72410e2 100644 --- a/dw/tasks/audio_transcription.py +++ b/dw/tasks/audio_transcription.py @@ -21,7 +21,10 @@ from .. import preferred_task_dtype from ..events import emit_warning from .model_cache import cached_model, hf_pipeline_placement -from ..dsp import resample_waveform +from ..assessment_rules import DEAD_AIR_FLOOR_DBFS +from ..dsp import loudest_window_dbfs, resample_waveform +from ..script_lines import looped_texts, normalize_words +from .assess import DEAD_AIR_WINDOW from .audio_utils import waveform_and_rate logger = logging.getLogger("dw") @@ -69,6 +72,9 @@ def transcribe_audio(audio, device="cpu", sample_rate=None, **kwargs): word with no audible energy keeps Whisper's bounds - the step's result content_type must then be "application/json" rather than "text/plain", since the return shape follows the argument. + A chunk that is probably invented carries a `suspect` reason + ("no_words", "repetition", "silent" or "long"), marked and + never dropped (#822). Returns: Transcribed text string, or (with timestamps set) a @@ -134,9 +140,55 @@ def load_pipe(): chunks = _numeric_chunks(result.get("chunks", []), duration) if timestamps == "word": chunks = _trim_to_speech(chunks, mono) + chunks = _mark_suspect(chunks, mono, long_word=timestamps == "word") return {"text": text, "chunks": chunks} +# A word (timestamps="word") lasting more than this is marked `suspect: "long"`: +# sung or spoken words do not hold for seconds; a span that long is an +# instrumental folded into the word, or an invented one (#822) +LONG_WORD_SECONDS = 2.0 + + +def _mark_suspect(chunks, mono, long_word): + """Chunks with `suspect` set on those that are probably not speech (#822). + + Marked, never removed. The first reason that holds wins: + "no_words": the text has no word at all ('—', '♪', '...'). + "repetition": a word of a repetition loop ('I'm sorry, I'm sorry, ...'), + found by `script_lines.looped_texts`, the detector check_script discards + loops with - Whisper's hallucination over music or a fade, however loud. + "silent": the chunk's loudest window sits at or below the dead-air floor + script_check guards with (`DEAD_AIR_FLOOR_DBFS`, measured by + `dsp.loudest_window_dbfs`) - Whisper invents text over silence, a + fade-out's end credits included. + "long": a word lasting more than LONG_WORD_SECONDS (`long_word` only, + since a segment is a line and runs long by nature). + """ + looped = looped_texts([chunk["text"] for chunk in chunks]) + marked = [] + for index, chunk in enumerate(chunks): + if not chunk["text"]: + marked.append(chunk) + continue + if not normalize_words(chunk["text"]): + reason = "no_words" + elif index in looped: + reason = "repetition" + else: + level = loudest_window_dbfs( + mono, _ASR_SAMPLE_RATE, chunk["start"], chunk["end"], DEAD_AIR_WINDOW + ) + if level is None or level <= DEAD_AIR_FLOOR_DBFS: + reason = "silent" + elif long_word and chunk["end"] - chunk["start"] > LONG_WORD_SECONDS: + reason = "long" + else: + reason = None + marked.append({**chunk, "suspect": reason} if reason else chunk) + return marked + + def _numeric_chunks(raw_chunks, duration): """The {start, end, text} chunks, with every time a number of seconds. diff --git a/dw/tasks/cuts.py b/dw/tasks/cuts.py index ce30c63c..c5613dca 100644 --- a/dw/tasks/cuts.py +++ b/dw/tasks/cuts.py @@ -34,18 +34,11 @@ import math import re +from ..script_lines import SUSPECT_LEFT_OUT, suspect_runs from ..variable_constraints import aligned, aligned_down from .registry import register_command -from ..task_domains import ( - NON_NEGATIVE, - POSITIVE, - real_number, - whole_number, -) -from ..task_problems import ( - cuts_errors, - transcript_problem, -) +from ..task_domains import NON_NEGATIVE, POSITIVE, real_number, whole_number +from ..task_problems import cuts_errors, transcript_problem logger = logging.getLogger("dw") @@ -107,9 +100,10 @@ def _chunks(transcript, duration): end = chunk.get("end") start = _number(start, f"transcript chunks[{index}].start", required=True) end = _number(end, f"transcript chunks[{index}].end") - parsed.append( - {"start": start, "end": end, "text": str(chunk.get("text") or "").strip()} - ) + row = {"start": start, "end": end, "text": str(chunk.get("text") or "").strip()} + if chunk.get("suspect"): + row["suspect"] = str(chunk["suspect"]) + parsed.append(row) parsed.sort(key=lambda chunk: chunk["start"]) for index, chunk in enumerate(parsed): if chunk["end"] is None: @@ -120,6 +114,22 @@ def _chunks(transcript, duration): return parsed +def _suspect_warnings(chunks): + """One warning per run of chunks transcribe_audio marked with the same + `suspect` reason (#822), naming its span, size and text.""" + return [ + f"{COMMAND}: {len(run)} transcript chunk(s) {run[0]['start']:.2f}-" + f"{max(c['end'] or c['start'] for c in run):.2f} s " + f"{_clipped(' '.join(c['text'] for c in run))!r} marked suspect ({reason}); " + + ("left out of the lines" if reason in SUSPECT_LEFT_OUT else "still used") + for reason, run in suspect_runs(chunks) + ] + + +def _clipped(text, limit=80): + return text if len(text) <= limit else text[: limit - 3] + "..." + + def _beats(beats): """(beat seconds ascending, bpm or None, song duration or None) from an analyze_beats result or a bare list of seconds.""" @@ -626,18 +636,23 @@ def _joined(first, second): def _merge_empty(scenes, warnings, why): - index = 0 + """Each scene of no length merged into the next (the last into the one + before), with one warning per run of lines squeezed out in one place.""" + index, squeezed = 0, {} while index < len(scenes) and len(scenes) > 1: - scene = scenes[index] - if scene["end"] - scene["start"] > 1e-9: + if scenes[index]["end"] - scenes[index]["start"] > 1e-9: index += 1 continue - if scene["lyrics"]: - warnings.append( - f"{COMMAND}: {why} left no time for " - f"{' / '.join(scene['lyrics'])!r}; it joins the next scene" - ) + # Merged back to back at one index, the scene holds every line before it + if scenes[index]["lyrics"]: + squeezed[index] = scenes[index]["lyrics"] _merge_into(scenes, index, index + 1 if index + 1 < len(scenes) else index - 1) + for lines in squeezed.values(): + more = f" / ... ({len(lines) - 3} more)" if len(lines) > 3 else "" + warnings.append( + f"{COMMAND}: {why} left no time for " + f"{_clipped(' / '.join(lines[:3])) + more!r}; it joins the next scene" + ) return scenes @@ -932,7 +947,9 @@ def plan_cuts( Args: transcript: transcribe_audio's result with 'timestamps' set - a {text, chunks} dict of {start, end, text} chunks in seconds. A - null 'end' runs to the next chunk's start, or the song's end + null 'end' runs to the next chunk's start, or the song's end. A + chunk marked `suspect` no_words, repetition or silent is left + out of the lines; every suspect run is warned about lyrics: The song's own lyrics - one line per line, a blank line or a lone section tag ("[Chorus]") between stanzas - or a list of lines. The shots carry these lines verbatim and in order; the @@ -1008,6 +1025,8 @@ def plan_cuts( warnings = [] duration, chunks = _song_length(args, transcript, beats, beats_duration, warnings) beat_times = [beat for beat in beat_times if 0.0 < beat < duration] + warnings.extend(_suspect_warnings(chunks)) + chunks = [chunk for chunk in chunks if chunk.get("suspect") not in SUSPECT_LEFT_OUT] if lyrics is not None and not (isinstance(lyrics, str) and not lyrics.strip()): lines = _align_lyrics( diff --git a/dw/tasks/script_check.py b/dw/tasks/script_check.py index 8ab72481..1832b78f 100644 --- a/dw/tasks/script_check.py +++ b/dw/tasks/script_check.py @@ -50,6 +50,9 @@ from ..locations import is_http_url from ..script_lines import ( COMMAND, + REPEAT_MAX_PERIOD, + REPEAT_RUN_MIN, + looped_texts, normalize_words, parse_lines, parse_shots, @@ -78,13 +81,8 @@ # GUARD_FLOOR_DBFS - still voiced when the file stops, not a word that # finished into a quiet tail CLIP_TAIL_SECONDS = 0.25 -# A run of heard words that repeats with a period of at most -# REPEAT_MAX_PERIOD words, REPEAT_RUN_MIN or more times over, is discarded -# as a decoding loop - Whisper's hallucination over music or room tone -# ('Pre-pre-pre-...', 'thank you thank you ...'), loud enough to pass the -# energy guard. Speech repeats a word a few times; it does not loop -REPEAT_RUN_MIN = 6 -REPEAT_MAX_PERIOD = 4 +# REPEAT_RUN_MIN and REPEAT_MAX_PERIOD, the repetition-loop guard, live in +# dw/script_lines.py beside the detector, which transcribe_audio shares THRESHOLDS = { "similarity": DEFAULT_SIMILARITY, @@ -189,46 +187,14 @@ def shot_at(spans, seconds): def word_level(mono, sample_rate, start, end): """The loudest GUARD_WINDOW_SECONDS window's rms over [start, end], in - dBFS, or None for silence. A span shorter than one window (a zero-length - word chunk) is measured as one window centred on it.""" - window = max(1, int(round(GUARD_WINDOW_SECONDS * sample_rate))) - first = int(round(start * sample_rate)) - last = int(round(end * sample_rate)) - if last - first < window: - middle = (first + last) // 2 - first, last = middle - window // 2, middle - window // 2 + window - first, last = max(0, first), min(len(mono), last) - loudest = None - for offset in range(first, last, window): - level = dsp.rms(mono[offset : min(offset + window, last)]) - if level is not None and (loudest is None or level > loudest): - loudest = level - return dsp.dbfs(loudest) + dBFS, or None for silence (`dsp.loudest_window_dbfs`).""" + return dsp.loudest_window_dbfs(mono, sample_rate, start, end, GUARD_WINDOW_SECONDS) def _round(value, places=3): return None if value is None else round(float(value), places) -def repetition_loops(keys): - """The indices of `keys` (heard words) that sit in a repetition loop: a stretch that - repeats with a period of 1 to REPEAT_MAX_PERIOD and spans at least - REPEAT_RUN_MIN periods.""" - looped = set() - for period in range(1, REPEAT_MAX_PERIOD + 1): - k = period - while k < len(keys): - if keys[k] != keys[k - period]: - k += 1 - continue - first = k - period - while k < len(keys) and keys[k] == keys[k - period]: - k += 1 - if k - first >= REPEAT_RUN_MIN * period: - looped.update(range(first, k)) - return looped - - def guard_words(chunks, mono, sample_rate): """Split transcribe_audio's word chunks into (heard, discarded). @@ -237,15 +203,14 @@ def guard_words(chunks, mono, sample_rate): [{text, start, end, level_dbfs, reason}] for a chunk in a repetition loop (`reason` "repetition") or at or below the floor ("below_floor"). """ - spoken, sequence, owner = [], [], [] + spoken = [] for chunk in chunks: words = normalize_words(chunk.get("text", "")) if words: - owner.extend([len(spoken)] * len(words)) - sequence.extend(words) spoken.append((chunk, words)) - # Over words, not chunks: one chunk can hold a whole loop ('Pre-pre-pre') - looped = {owner[k] for k in repetition_loops(sequence)} + # The loop detector is script_lines' (shared with transcribe_audio's + # `suspect: "repetition"`); a chunk with no word is in no loop + looped = looped_texts([chunk.get("text", "") for chunk, _ in spoken]) heard, discarded = [], [] for index, (chunk, words) in enumerate(spoken): diff --git a/dw_mcp/catalog.py b/dw_mcp/catalog.py index 869a4fd1..c6a2d184 100644 --- a/dw_mcp/catalog.py +++ b/dw_mcp/catalog.py @@ -282,7 +282,9 @@ def list_gallery( return client.get_json("/api/gallery", params=params, workspace=workspace) -def get_gallery_metadata(client, name, envelope=False, workspace=None): +def get_gallery_metadata( + client, name, envelope=False, workspace=None, interval_seconds=None +): """Metadata embedded in a saved file: the full workflow that made it - the exact workflow, arguments and seed, so a result can be reproduced or a failed run's definition edited and re-run - plus the job that @@ -308,10 +310,15 @@ def get_gallery_metadata(client, name, envelope=False, workspace=None): With envelope=True the soundtrack's level is reported second by second as well, which is what locates something in a track rather than only measuring the whole of it. Opt-in: it is one number per second per - measure, and the default answer has to stay small.""" + measure, and the default answer has to stay small. `interval_seconds` + (0.05 to 3600, default 1.0) narrows the bins; a finer-than-default width + that would need over 2400 of them is refused with the interval that fits.""" + params = {"envelope": "true"} if envelope else None + if envelope and interval_seconds is not None: + params["interval_seconds"] = interval_seconds body = client.get_json( api_path("api", "gallery", name, "metadata"), - params={"envelope": "true"} if envelope else None, + params=params, workspace=workspace, ) media = body.get("media") diff --git a/dw_mcp/tools_catalog.py b/dw_mcp/tools_catalog.py index 77eaa2f7..cb5e9d93 100644 --- a/dw_mcp/tools_catalog.py +++ b/dw_mcp/tools_catalog.py @@ -226,7 +226,8 @@ def get_server_info( workspace. `trust_workflows` reports the posture a submitted workflow is read under: false - the default - means the file is untrusted input, so an out-of-ecosystem import, remote code, and a - media location outside the workspace's roots are all refused.""" + media location outside the workspace's roots are all refused. + `caches` sizes dw's disk caches (cleared from the server page).""" return workspaces.server_info(self.client) def list_jobs( @@ -311,37 +312,43 @@ def list_gallery( ) def get_gallery_metadata( - self, name: str, envelope: bool = False, workspace: str | None = None + self, + name: str, + envelope: bool = False, + workspace: str | None = None, + interval_seconds: float | None = None, ) -> dict: """Get the metadata embedded in a generated file: the exact - workflow, arguments and seed that produced it - the definition, - not a summary, so a result can be reproduced or a failed run's - definition edited and re-run. Only an image embeds it; for audio + workflow, arguments and seed that produced it, so a result can be + reproduced or edited and re-run. Only an image embeds it; for audio and video `metadata` is null and `next` names - `get_job_workflow(job_id)` when known, else a kept asset has no - provenance. `media` itself carries duration, sample rate, - channels, fps, size and level - the checks an agent that cannot - listen makes on a deliverable; `media.shots` places a joined + `get_job_workflow(job_id)` when known. `media` carries duration, + sample rate, channels, fps, size and level - how an agent that + cannot listen checks a deliverable; `media.shots` places a joined video's shots. `envelope=true` adds that level second by second (`media.envelope.rms_dbfs` / `peak_dbfs`), which says *where* in a track something is: a shot's last frame, a seam's hole, where a - score goes quiet. Leave it off unless it's about position - a long - track is a long list. + score goes quiet. Leave it off unless it's about position. + `interval_seconds` (default 1.0, 0.05-3600) is the bin width; 0.1 + places an onset. Over 2400 bins is refused. - `findings` lists each level problem the server measured - full - scale, near silence - with its threshold and the fix. + `findings` lists each level problem measured (full scale, near + silence) with its fix. - `name` may be an `asset:` reference instead of a gallery name, and - then it describes that input asset - how many frames a shot is, - whether two shots share an fps, whether a score reaches the length - of the cut it will lie under. Check before running: frame counts - and rates are arguments the caller supplies, and a wrong one is a - failed job or, worse, silence padded onto the end of a track. + `name` may be an `asset:` reference, and then it describes that + input asset - a shot's frame count, whether two shots share an fps, + whether a score reaches the cut's length. Check before running: a + wrong frame count or rate is a failed job or silence padded onto a + track. `workspace` pins this call to another workspace.""" return catalog.get_gallery_metadata( - self.client, name, envelope=envelope, workspace=workspace + self.client, + name, + envelope=envelope, + workspace=workspace, + interval_seconds=interval_seconds, ) def list_guides( diff --git a/dw_mcp/workspaces.py b/dw_mcp/workspaces.py index 14baaf40..237d5509 100644 --- a/dw_mcp/workspaces.py +++ b/dw_mcp/workspaces.py @@ -195,4 +195,35 @@ def server_info(client): # every request), so its directories are already this workspace's info = catalog.get_server_info(client) info["workspace"] = client.workspace + caches = _cache_sizes(client) + if caches is not None: + info["caches"] = caches return info + + +# What get_server_info carries of each cache: enough to see what dw's disk use +# is. The root, cost and last-used time stay in GET /api/caches, for the page +# that clears them +_CACHE_FIELDS = ("name", "bytes", "files", "enabled", "clearable", "disk_free") + + +def _cache_sizes(client): + """dw's disk caches (#828), or None from a server without /api/caches - + the rest of get_server_info doesn't depend on them.""" + try: + listing = client.get_json("/api/caches") + except DwApiError: + return None + caches = listing.get("caches") if isinstance(listing, dict) else None + if not isinstance(caches, list): + return None + return [_cache_summary(cache) for cache in caches if isinstance(cache, dict)] + + +def _cache_summary(cache): + summary = {field: cache.get(field) for field in _CACHE_FIELDS} + # A cache with entries (quantized) is counted, not listed: the listing is + # GET /api/caches, for the page that deletes them + if isinstance(cache.get("entries"), list): + summary["entry_count"] = len(cache["entries"]) + return summary diff --git a/plugins/dw/.claude-plugin/plugin.json b/plugins/dw/.claude-plugin/plugin.json index b07eebf1..7a6c095c 100644 --- a/plugins/dw/.claude-plugin/plugin.json +++ b/plugins/dw/.claude-plugin/plugin.json @@ -1,7 +1,7 @@ { "name": "dw", "description": "Compose MiniMax H3 video, MiniMax Music 3, LTX-2.5 and Kandinsky 6.0 workflows over a dw MCP server, and cut a multi-episode series from them: which template fits which shape, the hard rules, cost, and how to judge the output. Prompt format comes from the vendors' own guides.", - "version": "0.11.1", + "version": "0.12.0-alpha.1", "author": { "name": "Don Kackman" }, diff --git a/pyproject.toml b/pyproject.toml index 977b1e8b..66011ddf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ name = "diffusers-workflow" # runtime (static in TOML because importing dw at build time would drag # torch into the build environment). Release tags must match it - see # docs/RELEASING.md -version = "0.11.1" +version = "0.12.0-alpha.1" description = "Declarative workflow engine and web UI for Hugging Face Diffusers" readme = "README.md" requires-python = ">=3.10" diff --git a/scripts/deploy.sh b/scripts/deploy.sh index 7387af89..1b98aedd 100755 --- a/scripts/deploy.sh +++ b/scripts/deploy.sh @@ -20,7 +20,11 @@ # running job is the caller's own. # 4. restart. With the systemd user unit installed (scripts/dw-serve.service) # that is `systemctl --user restart dw-serve` and the log is the journal. -# Without it: SIGTERM the running dw.serve (never SIGKILL; a stuck one is +# With a launchd daemon loaded (macOS, DW_LAUNCHD_LABEL): SIGTERM the +# running dw.serve and let launchd's KeepAlive start the new one - no +# second copy, which would race launchd's for the port and leave the +# loser crash-looping every few seconds. The daemon's plist owns the +# arguments and the log then, not this script. Otherwise: SIGTERM the running dw.serve (never SIGKILL; a stuck one is # a bug to report, so after 30 s this fails loudly instead of # escalating), then start the new server as its own window of the # `dw-serve` screen session (created if absent), stdout+stderr appended @@ -33,7 +37,9 @@ # DW_PORT (8765), DW_WORKSPACE (~/diffusers-workspace), DW_HOST (0.0.0.0), # DW_MCP_MAX_WAIT_SECONDS (1800, the per-call wait_for_job cap), # DW_NODE_BIN (the directory holding npm, when it is not on a -# non-interactive PATH and not in one of the places find_npm looks). +# non-interactive PATH and not in one of the places find_npm looks), +# DW_LAUNCHD_LABEL (com.local.dw, the system launchd daemon that may +# supervise dw.serve on macOS). set -euo pipefail DW_DIR="${DW_DIR:-$HOME/diffusers-workflow}" @@ -45,6 +51,7 @@ DW_WORKSPACE="${DW_WORKSPACE:-$HOME/diffusers-workspace}" # default is 55). #377: one call covers a long render. DW_MCP_MAX_WAIT_SECONDS="${DW_MCP_MAX_WAIT_SECONDS:-1800}" LOG="$HOME/dw-serve.log" +DW_LAUNCHD_LABEL="${DW_LAUNCHD_LABEL:-com.local.dw}" SCREEN_SESSION="dw-serve" HEALTH="http://localhost:$DW_PORT/api/health" @@ -78,6 +85,37 @@ listener_pids() { } server_running() { pgrep -if 'python[0-9.]* -m dw\.serve' >/dev/null; } any_alive() { local p; for p in "$@"; do kill -0 "$p" 2>/dev/null && return 0; done; return 1; } +# launchctl print reads a system daemon's state without root +launchd_loaded() { + [ "$(uname -s)" = Darwin ] && launchctl print "system/$DW_LAUNCHD_LABEL" >/dev/null 2>&1 +} +# The daemon's own stdout path, so a failed start points at the right log +launchd_log() { + launchctl print "system/$DW_LAUNCHD_LABEL" 2>/dev/null \ + | sed -n 's/^[[:space:]]*stdout path = //p' | head -n 1 +} + +# SIGTERM whatever listens on DW_PORT and wait for it to exit +stop_server() { + local pids + pids="$(listener_pids)" + if [ -z "$pids" ]; then + say "no dw.serve running" + return 0 + fi + say "stopping dw.serve (pid $pids)" + # shellcheck disable=SC2086 # one word per pid + kill -TERM $pids + for _ in $(seq 1 30); do + # shellcheck disable=SC2086 + any_alive $pids || break + sleep 1 + done + # shellcheck disable=SC2086 + if any_alive $pids; then + say "dw.serve did not exit within 30 s of SIGTERM (pid $pids); not escalating - report this"; exit 1 + fi +} # `ssh lem deploy.sh` runs a non-interactive shell, and a per-user node # install is usually put on PATH by ~/.bashrc *after* its "not interactive, @@ -154,26 +192,22 @@ if systemctl --user cat dw-serve >/dev/null 2>&1; then systemctl --user restart dw-serve LOG_HINT="journalctl --user -u dw-serve -n 200" tail_log() { journalctl --user -u dw-serve -n 20 --no-pager; } -else -LOG_HINT="$LOG" -tail_log() { tail -n 20 "$LOG"; } -pids="$(listener_pids)" -if [ -n "$pids" ]; then - say "stopping dw.serve (pid $pids)" - # shellcheck disable=SC2086 # one word per pid - kill -TERM $pids +elif launchd_loaded; then + LOG_HINT="$(launchd_log)"; LOG_HINT="${LOG_HINT:-$LOG}" + tail_log() { tail -n 20 "$LOG_HINT"; } + say "restarting via launchd daemon system/$DW_LAUNCHD_LABEL (KeepAlive starts the new one)" + stop_server + # launchd respawns at once, or after its throttle interval (10 s default) + # for a job that ran briefly; the health loop below expects a process for _ in $(seq 1 30); do - # shellcheck disable=SC2086 - any_alive $pids || break + server_running && break sleep 1 done - # shellcheck disable=SC2086 - if any_alive $pids; then - say "dw.serve did not exit within 30 s of SIGTERM (pid $pids); not escalating - report this"; exit 1 - fi + server_running || { say "launchd did not start dw.serve within 30 s of the stop; check: launchctl print system/$DW_LAUNCHD_LABEL"; exit 1; } else - say "no dw.serve running" -fi +LOG_HINT="$LOG" +tail_log() { tail -n 20 "$LOG"; } +stop_server # start the new one in its own screen window cmd="cd $DW_DIR && source venv/bin/activate && DW_MCP_MAX_WAIT_SECONDS=$DW_MCP_MAX_WAIT_SECONDS exec python -m dw.serve --host $DW_HOST --port $DW_PORT --mcp --token $DW_TOKEN --workspace $DW_WORKSPACE --examples-dir $DW_DIR/workflows >> $LOG 2>&1" diff --git a/scripts/measure_quantized_cache.py b/scripts/measure_quantized_cache.py new file mode 100644 index 00000000..8df5e16f --- /dev/null +++ b/scripts/measure_quantized_cache.py @@ -0,0 +1,541 @@ +#!/usr/bin/env python +"""Measure what an on-disk cache of SDNQ on-the-fly-quantized components would save. + +Stage 0 of #815 (#825): the gate for building the quantized cache. For the +MiniMax H3 default config (transformer, text encoder, VAE, as in +workflows/templates/minimax/video-with-audio.json) and the LTX-2.5 default +(transformer, text encoder, as in workflows/templates/ltx2/text-to-video.json) +it times, per component: + +- fresh: `from_pretrained` with an on-the-fly `SDNQConfig`, the way dw loads + it today (read bf16, quantize on `quantization_device`, land on + `return_device`); +- cached: `load_sdnq_model` of the same component saved with + `save_sdnq_model`, plus placement on `return_device`. + +Each is timed with the page cache warm and cold. Cold means evicted per file, +without root: `os.posix_fadvise(fd, 0, 0, POSIX_FADV_DONTNEED)` on every +source safetensors (fresh) or every saved file (cached) right before the +load. No `drop_caches`, no sudo. Eviction is advisory: a page another process +holds may stay resident, so cold numbers may read slightly warm. + +It imports only torch, diffusers, transformers, huggingface_hub and sdnq - no +dw engine code - and reads the weights from the local Hugging Face cache only +(the models must already be downloaded; nothing is fetched). + +**Stop `dw.serve` first.** A running worker's resident models distort the +page cache and the timings, and both cards share host RAM; the script refuses +to run while a process running `-m dw.serve` exists (this process and its +ancestors excepted). + +Run it on lem from the repo root, with the venv active: + + cd ~/diffusers-workflow && source ./activate + python scripts/measure_quantized_cache.py + +Options: `--family h3|ltx|all` (default all), `--device cuda:0` (default +`cuda`; replaces "cuda" in the configs' quantization/return devices), +`--scratch DIR` (default `$HF_HOME/dw-quantized-measure`), `--keep` to leave +the saved entries in the scratch directory (by default each is deleted once +measured). It prints one table at the end: per component the fresh warm/cold +and cached warm/cold seconds and the bytes on disk, the eviction method, +per-family totals, and the H3 saving against the ~30 s gate. +""" + +import argparse +import gc +import importlib.metadata +import json +import os +import shutil +import subprocess +import sys +import time +from dataclasses import dataclass, field +from pathlib import Path + +GATE_SECONDS = 30.0 +EVICTION_METHOD = "posix_fadvise DONTNEED (per file, no root)" + + +@dataclass +class ComponentSpec: + family: str + name: str + repo: str + subfolder: str + cls_path: str # "diffusers.X" or "transformers.X" + sdnq_args: dict + timings: dict = field(default_factory=dict) + bytes_on_disk: int | None = None + error: str | None = None + + +# The defaults as the templates ship them (variables resolved). +H3_REPO = "MiniMaxAI/MiniMax-H3" +LTX_REPO = "Lightricks/LTX-2.5-Diffusers" + + +def default_specs(): + return [ + ComponentSpec( + "h3", + "transformer", + H3_REPO, + "transformer", + "diffusers.MiniMaxH3Transformer3DModel", + { + "weights_dtype": "int4", + "quantization_device": "cuda", + "return_device": "cpu", + "use_quantized_matmul": True, + "dequantize_fp32": False, + "modules_to_not_convert": [ + "proj_in", + "audio_proj_in", + "context_embedder", + "time_embedder", + "time_proj", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ], + }, + ), + ComponentSpec( + "h3", + "text_encoder", + H3_REPO, + "text_encoder", + "transformers.Qwen3VLForConditionalGeneration", + { + "weights_dtype": "int4", + "quantization_device": "cuda", + "return_device": "cpu", + "dequantize_fp32": False, + "modules_to_not_convert": [".model.visual", "lm_head"], + }, + ), + ComponentSpec( + "h3", + "vae", + H3_REPO, + "vae", + "diffusers.AutoencoderKLMiniMaxH3", + { + "weights_dtype": "int8", + "quant_conv": True, + "use_quantized_matmul_conv": True, + "quantization_device": "cuda", + "return_device": "cpu", + "dequantize_fp32": False, + }, + ), + ComponentSpec( + "ltx", + "transformer", + LTX_REPO, + "transformer", + "diffusers.LTX2VideoTransformer3DModel", + { + "weights_dtype": "uint4", + "quantization_device": "cuda", + "return_device": "cuda", + "use_quantized_matmul": True, + "dequantize_fp32": False, + }, + ), + ComponentSpec( + "ltx", + "text_encoder", + LTX_REPO, + "text_encoder", + "transformers.Gemma4UnifiedForConditionalGeneration", + { + "weights_dtype": "int8", + "quantization_device": "cuda", + "return_device": "cpu", + "dequantize_fp32": False, + "modules_to_not_convert": ["embed_vision", "embed_audio", "lm_head"], + }, + ), + ] + + +def ancestor_pids(): + """This process's pid and every ancestor's, so the launching shell never matches.""" + pids = set() + pid = os.getpid() + while pid > 0 and pid not in pids: + pids.add(pid) + result = subprocess.run( + ["ps", "-o", "ppid=", "-p", str(pid)], capture_output=True, text=True + ) + try: + pid = int(result.stdout.strip()) + except ValueError: + break + return pids + + +def refuse_if_server_running(): + # the server's own invocation is `python -m dw.serve`; a bare "dw.serve" would + # also match any unrelated command line that merely mentions it + result = subprocess.run( + ["pgrep", "-f", "--", r"-m dw\.serve"], capture_output=True, text=True + ) + skip = ancestor_pids() + pids = [p for p in result.stdout.split() if p.isdigit() and int(p) not in skip] + if pids: + sys.exit( + f"dw.serve is running (pid {', '.join(pids)}). Stop it first: its resident models distort the page cache and timings." + ) + + +def resolve_class(cls_path): + module_name, _, cls_name = cls_path.partition(".") + module = __import__(module_name) + cls = getattr(module, cls_name, None) + if cls is None: + raise RuntimeError(f"{cls_path} is not in the installed {module_name}") + return cls + + +def source_dir(spec): + from huggingface_hub import snapshot_download + + prefix = f"{spec.subfolder}/" if spec.subfolder else "" + # only what a load reads: the config, the shard index, and the shards the + # index names - a folder can hold more than one shard set (LTX-2.5's + # transformer/ kept its 8-shard files after going back to 4), and asking + # for every *.safetensors under local_files_only demands the unused set + snapshot = snapshot_download( + spec.repo, + allow_patterns=[f"{prefix}config.json", f"{prefix}*.index.json"], + local_files_only=True, + ) + path = Path(snapshot) / spec.subfolder + indexes = sorted(path.glob("*.index.json")) + if indexes: + with open(indexes[0]) as f: + shards = sorted(set(json.load(f)["weight_map"].values())) + weights = [f"{prefix}{name}" for name in shards] + else: + weights = [f"{prefix}*.safetensors"] + snapshot_download(spec.repo, allow_patterns=weights, local_files_only=True) + if not any(path.glob("*.safetensors")): + raise RuntimeError(f"no safetensors under {path}: download {spec.repo} first") + return path + + +def files_under(path): + # resolve() follows the HF cache's snapshot symlinks to the blobs; a shard + # set the index does not name was never downloaded and is not a file here + return [p.resolve() for p in sorted(Path(path).rglob("*")) if p.is_file()] + + +def evict(paths): + for path in paths: + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) # DONTNEED does not drop dirty pages + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + + +def warm(paths): + # read every byte once, so the warm run starts with the files resident + for path in paths: + with open(path, "rb") as f: + while f.read(64 * 1024 * 1024): + pass + + +def sync(torch): + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def release(torch): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def device_args(sdnq_args, device): + args = dict(sdnq_args) + for key in ("quantization_device", "return_device"): + if args.get(key) == "cuda": + args[key] = device + return args + + +def touch_pages(model, torch): + """Read one byte per 4 KiB page of every CPU tensor, so lazily mmapped + safetensors are actually read inside the timed region.""" + total = 0 + for t in list(model.parameters()) + list(model.buffers()): + try: + if t.device.type != "cpu" or t.numel() == 0: + continue + flat = t.detach().reshape(-1).view(torch.uint8) + total += int(flat[::4096].sum()) + except Exception: + # a quantized tensor subclass may not expose a plain view: read its + # inner tensors instead, else skip it + flatten = getattr(t, "__tensor_flatten__", None) + for name in flatten()[0] if flatten else (): + inner = getattr(t, name, None) + if isinstance(inner, torch.Tensor) and inner.device.type == "cpu": + try: + flat = inner.detach().reshape(-1).view(torch.uint8) + total += int(flat[::4096].sum()) + except Exception: + pass + return total + + +def load_fresh(spec, cls, sdnq_args, torch): + from sdnq import SDNQConfig + + dtype_kwarg = ( + "dtype" if spec.cls_path.startswith("transformers.") else "torch_dtype" + ) + kwargs = { + dtype_kwarg: torch.bfloat16, + "quantization_config": SDNQConfig(**sdnq_args), + } + model = cls.from_pretrained( + spec.repo, subfolder=spec.subfolder, local_files_only=True, **kwargs + ) + touch_pages(model, torch) + return model + + +def load_cached(entry, cls, sdnq_args, torch): + from sdnq import load_sdnq_model + + return_device = sdnq_args["return_device"] + model = load_sdnq_model( + str(entry), + model_cls=cls, + dtype=torch.bfloat16, + device=return_device, + dequantize_fp32=sdnq_args.get("dequantize_fp32"), + use_quantized_matmul=sdnq_args.get("use_quantized_matmul"), + ) + model = model.to(return_device) + touch_pages(model, torch) + return model + + +def timed(fn, torch): + sync(torch) + start = time.perf_counter() + model = fn() + sync(torch) + return model, time.perf_counter() - start + + +def measure(spec, scratch, device, keep, torch): + from sdnq import save_sdnq_model + + cls = resolve_class(spec.cls_path) + sdnq_args = device_args(spec.sdnq_args, device) + src_files = files_under(source_dir(spec)) + entry = scratch / f"{spec.family}-{spec.name}" + if entry.exists(): + shutil.rmtree(entry) + + def step(label, fn, files, cold): + (evict if cold else warm)(files) + model, seconds = timed(fn, torch) + spec.timings[label] = seconds + print(f" {spec.family}/{spec.name} {label}: {seconds:.1f} s", flush=True) + return model + + print( + f"{spec.family}/{spec.name}: {spec.repo} {spec.subfolder} ({spec.cls_path})", + flush=True, + ) + model = step( + "fresh_cold", + lambda: load_fresh(spec, cls, sdnq_args, torch), + src_files, + cold=True, + ) + del model + release(torch) + model = step( + "fresh_warm", + lambda: load_fresh(spec, cls, sdnq_args, torch), + src_files, + cold=False, + ) + save_sdnq_model(model, str(entry)) + del model + release(torch) + entry_files = files_under(entry) + spec.bytes_on_disk = sum(p.stat().st_size for p in entry_files) + print( + f" {spec.family}/{spec.name} saved: {spec.bytes_on_disk / 1e9:.2f} GB", + flush=True, + ) + # the fresh loads leave the source pages warm; evict them so they don't + # crowd the cached runs out of RAM + evict(src_files) + + model = step( + "cached_cold", + lambda: load_cached(entry, cls, sdnq_args, torch), + entry_files, + cold=True, + ) + del model + release(torch) + model = step( + "cached_warm", + lambda: load_cached(entry, cls, sdnq_args, torch), + entry_files, + cold=False, + ) + del model + release(torch) + evict(entry_files) + if not keep: + shutil.rmtree(entry) + + +def fmt_seconds(value): + return "-" if value is None else f"{value:.1f}" + + +def print_table(specs, device): + columns = ("fresh_warm", "fresh_cold", "cached_warm", "cached_cold") + print() + print(f"eviction: {EVICTION_METHOD}; device: {device}") + versions = { + pkg: importlib.metadata.version(pkg) + for pkg in ("sdnq", "diffusers", "transformers", "torch") + } + print("versions: " + ", ".join(f"{k} {v}" for k, v in versions.items())) + print() + header = ( + f"| {'component':<20} | {'weights':<7} | " + + " | ".join(f"{c:>11}" for c in columns) + + f" | {'GB on disk':>10} |" + ) + print(header) + print( + "|" + "|".join("-" * (len(cell) + 0) for cell in header.split("|")[1:-1]) + "|" + ) + totals = {} + for spec in specs: + row_name = f"{spec.family}/{spec.name}" + if spec.error: + print( + f"| {row_name:<20} | {spec.sdnq_args['weights_dtype']:<7} | error: {spec.error}" + ) + continue + cells = " | ".join(f"{fmt_seconds(spec.timings.get(c)):>11}" for c in columns) + gb = ( + f"{spec.bytes_on_disk / 1e9:.2f}" if spec.bytes_on_disk is not None else "-" + ) + print( + f"| {row_name:<20} | {spec.sdnq_args['weights_dtype']:<7} | {cells} | {gb:>10} |" + ) + family = totals.setdefault( + spec.family, {c: 0.0 for c in columns} | {"bytes": 0} + ) + for c in columns: + family[c] += spec.timings.get(c, 0.0) + family["bytes"] += spec.bytes_on_disk or 0 + for name, family in totals.items(): + complete = all(not s.error for s in specs if s.family == name) + label = f"{name} total" + ("" if complete else " (partial)") + cells = " | ".join(f"{family[c]:>11.1f}" for c in columns) + print(f"| {label:<20} | {'':<7} | {cells} | {family['bytes'] / 1e9:>10.2f} |") + print() + h3 = totals.get("h3") + if h3 is None: + return + if any(s.error for s in specs if s.family == "h3"): + print("H3 gate: not decided - a component failed (see the table)") + return + saving_warm = h3["fresh_warm"] - h3["cached_warm"] + saving_cold = h3["fresh_cold"] - h3["cached_cold"] + for label, saving in ( + ("page cache warm", saving_warm), + ("page cache cold", saving_cold), + ): + verdict = "clears" if saving >= GATE_SECONDS else "misses" + print( + f"H3 saving per cold load, {label}: {saving:.1f} s - {verdict} the ~{GATE_SECONDS:.0f} s gate" + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument("--family", choices=("h3", "ltx", "all"), default="all") + parser.add_argument( + "--device", + default="cuda", + help="replaces 'cuda' in the configs' devices, e.g. cuda:1", + ) + parser.add_argument( + "--scratch", + type=Path, + default=None, + help="where saved entries go (default $HF_HOME/dw-quantized-measure)", + ) + parser.add_argument("--keep", action="store_true", help="keep the saved entries") + args = parser.parse_args() + + if not hasattr(os, "posix_fadvise"): + sys.exit( + "os.posix_fadvise is unavailable on this platform; run this on lem (Linux)." + ) + refuse_if_server_running() + + import sdnq # noqa: F401 - registers the SDNQ quantizer with diffusers and transformers + import torch + from huggingface_hub import constants + + scratch = args.scratch or Path(constants.HF_HOME) / "dw-quantized-measure" + scratch.mkdir(parents=True, exist_ok=True) + specs = [s for s in default_specs() if args.family in ("all", s.family)] + + if torch.cuda.is_available(): + # pay CUDA context and allocator setup before the first timed load + torch.ones(1, device=args.device).add_(1) + sync(torch) + + for spec in specs: + try: + measure(spec, scratch, args.device, args.keep, torch) + except Exception as e: # report and go on to the next component + spec.error = f"{type(e).__name__}: {e}" + print(f" {spec.family}/{spec.name} failed: {spec.error}", flush=True) + release(torch) + + print_table(specs, args.device) + print() + print( + json.dumps( + [ + { + "component": f"{s.family}/{s.name}", + **s.timings, + "bytes": s.bytes_on_disk, + "error": s.error, + } + for s in specs + ] + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index a6fcbb18..c920660b 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -153,6 +153,9 @@ def test_the_committed_document_is_current(): ("get", "/api/models/downloads"), ("post", "/api/models/downloads/{download_id}/cancel"), ("delete", "/api/models"), + ("get", "/api/caches"), + ("delete", "/api/caches/{name}"), + ("delete", "/api/caches/{name}/entries/{entry}"), ("get", "/api/system/diffusers"), ("post", "/api/system/diffusers/update"), ("get", "/api/jobs"), diff --git a/tests/test_assess.py b/tests/test_assess.py index e77bd74f..a6e5f2c6 100644 --- a/tests/test_assess.py +++ b/tests/test_assess.py @@ -156,6 +156,11 @@ def test_step_flags_one_seam_and_shot_spread(self): ] assert len(spread_findings) == 1 assert shots_answer["rms_range_db"] >= 6.0 + # The finding says which shots, by position and start frame, not name alone + at = spread_findings[0]["at"] + table = shots_answer["shots"] + assert [table[i]["name"] for i in at["indices"]] == at["between"] + assert at["start_frames"] == [table[i]["start_frame"] for i in at["indices"]] # JSON-serializable, per requirement 10 json.dumps(seams) @@ -978,7 +983,7 @@ def test_analyze_seams_warns_and_reports_a_finding(self, monkeypatch): overrun = [f for f in answer["findings"] if f["rule"] == "shot_span_overrun"] assert len(overrun) == 1 - assert overrun[0]["at"] == {"shot": "b"} + assert overrun[0]["at"]["shot"] == "b" assert overrun[0]["value"]["frames"] == 176 assert any(w[1].get("kind") == "shot_span_overrun" for w in warnings) @@ -991,7 +996,7 @@ def test_analyze_shots_warns_and_reports_a_finding(self, monkeypatch): overrun = [f for f in answer["findings"] if f["rule"] == "shot_span_overrun"] assert len(overrun) == 1 - assert overrun[0]["at"] == {"shot": "b"} + assert overrun[0]["at"]["shot"] == "b" assert any(w[1].get("kind") == "shot_span_overrun" for w in warnings) json.dumps(answer) @@ -1001,7 +1006,7 @@ def test_analyze_sync_drift_warns_and_reports_a_finding(self, monkeypatch): overrun = [f for f in answer["findings"] if f["rule"] == "shot_span_overrun"] assert len(overrun) == 1 - assert overrun[0]["at"] == {"shot": "b"} + assert overrun[0]["at"]["shot"] == "b" assert any(w[1].get("kind") == "shot_span_overrun" for w in warnings) json.dumps(answer) diff --git a/tests/test_disk_caches.py b/tests/test_disk_caches.py new file mode 100644 index 00000000..a7028917 --- /dev/null +++ b/tests/test_disk_caches.py @@ -0,0 +1,760 @@ +"""dw's disk-cache registry (dw/disk_caches.py) and its routes (#828).""" + +import builtins +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +from dw import _inductor_cache_dir_default, disk_caches +from dw.disk_caches import ( + MARKER, + NoEntries, + NotClearable, + UnknownCache, + UnknownEntry, +) +from dw.server.app import create_app +from dw.server.jobs import JobManager + +from tests.test_server import ScriptedWorkerManager, success_script + + +@pytest.fixture +def env(tmp_path, monkeypatch): + """HOME and every cache variable pointed into tmp_path, triton gone from + sys.modules so the never-imported assertions mean something.""" + home = tmp_path / "home" + home.mkdir() + monkeypatch.setenv("HOME", str(home)) + monkeypatch.delenv("TRITON_CACHE_DIR", raising=False) + monkeypatch.delenv("TRITON_HOME", raising=False) + monkeypatch.delenv("XDG_CACHE_HOME", raising=False) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", str(tmp_path / "inductor")) + monkeypatch.setenv("DW_QUANTIZED_CACHE_DIR", str(tmp_path / "quantized")) + monkeypatch.delenv("DW_QUANTIZED_CACHE", raising=False) + monkeypatch.delitem(sys.modules, "triton", raising=False) + return home + + +def populate(root, sizes=(10, 20, 30)): + """Files in the root and a nested directory; returns (bytes, files).""" + root.mkdir(parents=True, exist_ok=True) + (root / "sub" / "deep").mkdir(parents=True, exist_ok=True) + paths = [root / "a.bin", root / "sub" / "b.bin", root / "sub" / "deep" / "c.bin"] + for path, size in zip(paths, sizes): + path.write_bytes(b"x" * size) + return sum(sizes), len(paths) + + +# ------------------------------------------------------------------ scan + + +def test_scan_counts_bytes_files_and_newest_mtime(env, tmp_path): + root = tmp_path / "inductor" + total, count = populate(root) + (root / MARKER).write_text("marked") + os.utime(root / "a.bin", (1000, 1000)) + os.utime(root / "sub" / "b.bin", (3000, 3000)) + os.utime(root / "sub" / "deep" / "c.bin", (2000, 2000)) + os.utime(root / MARKER, (9000, 9000)) + + info = disk_caches.scan(disk_caches.get("inductor")) + + assert info["bytes"] == total + assert info["files"] == count + assert info["last_used"] == 3000 + assert info["clearable"] is True + assert "not_clearable_reason" not in info + + +def test_scan_missing_root(env, tmp_path): + info = disk_caches.scan(disk_caches.get("inductor")) + + assert not (tmp_path / "inductor").exists() + assert info["bytes"] == 0 + assert info["files"] == 0 + assert info["last_used"] is None + assert info["clearable"] is False + assert isinstance(info["disk_free"], int) + assert isinstance(info["disk_total"], int) + + +def test_scan_does_not_follow_symlinks(env, tmp_path): + outside = tmp_path / "outside" + outside.mkdir() + (outside / "big.bin").write_bytes(b"x" * 100_000) + target_file = tmp_path / "target.bin" + target_file.write_bytes(b"x" * 50_000) + root = tmp_path / "inductor" + total, count = populate(root) + os.symlink(outside, root / "dirlink") + os.symlink(target_file, root / "filelink") + + info = disk_caches.scan(disk_caches.get("inductor")) + + assert info["bytes"] == total + assert info["files"] == count + + +# ---------------------------------------------------------------- triton + + +def test_triton_root_cache_dir_wins(env, tmp_path, monkeypatch): + monkeypatch.setenv("TRITON_CACHE_DIR", str(tmp_path / "tc")) + monkeypatch.setenv("TRITON_HOME", str(tmp_path / "th")) + assert disk_caches.get("triton").root() == str(tmp_path / "tc") + + +def test_triton_root_from_triton_home(env, tmp_path, monkeypatch): + monkeypatch.setenv("TRITON_HOME", str(tmp_path / "th")) + assert disk_caches.get("triton").root() == str( + tmp_path / "th" / ".triton" / "cache" + ) + + +def test_triton_root_defaults_under_home(env): + assert disk_caches.get("triton").root() == str(env / ".triton" / "cache") + + +def test_triton_is_never_imported(env, monkeypatch): + real_import = builtins.__import__ + + def refuse_triton(name, *args, **kwargs): + if name == "triton" or name.startswith("triton."): + raise AssertionError("triton was imported") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", refuse_triton) + triton = disk_caches.get("triton") + triton.root() + disk_caches.scan(triton) + disk_caches.mark_roots() + monkeypatch.undo() + + assert "triton" not in sys.modules + + +@pytest.mark.parametrize("installed", [True, False]) +def test_triton_enabled_follows_find_spec(env, monkeypatch, installed): + monkeypatch.setattr( + "dw.disk_caches.importlib.util.find_spec", + lambda name: object() if installed and name == "triton" else None, + ) + assert disk_caches.get("triton").enabled() is installed + + +def test_missing_triton_root_is_not_created(env): + triton = disk_caches.get("triton") + + assert disk_caches.mark_root(triton) is False + disk_caches.mark_roots() + + assert not (env / ".triton").exists() + info = disk_caches.scan(triton) + assert info["bytes"] == 0 + assert info["files"] == 0 + assert info["clearable"] is False + + +# -------------------------------------------------------------- inductor + + +def test_inductor_root_from_env(env, tmp_path): + assert disk_caches.get("inductor").root() == str(tmp_path / "inductor") + + +def test_inductor_default_is_xdg_dw_torchinductor(env, tmp_path, monkeypatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + monkeypatch.delenv("TORCHINDUCTOR_CACHE_DIR") + cache = disk_caches.get("inductor") + + expected = str(tmp_path / "xdg" / "dw" / "torchinductor") + assert _inductor_cache_dir_default(os.environ) == expected + assert cache.default_root() == expected + assert cache.root() == expected + + +def test_missing_inductor_root_is_created_and_marked(env, tmp_path): + assert disk_caches.mark_root(disk_caches.get("inductor")) is True + + assert (tmp_path / "inductor" / MARKER).is_file() + + +# ---------------------------------------------------------- marker rule + + +def test_empty_root_gets_marker(env, tmp_path): + (tmp_path / "inductor").mkdir() + assert disk_caches.mark_root(disk_caches.get("inductor")) is True + assert (tmp_path / "inductor" / MARKER).is_file() + + +def test_populated_inductor_default_is_grandfathered(env, tmp_path, monkeypatch): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + default = _inductor_cache_dir_default(os.environ) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", default) + populate(tmp_path / "xdg" / "dw" / "torchinductor") + + assert disk_caches.mark_root(disk_caches.get("inductor")) is True + assert os.path.isfile(os.path.join(default, MARKER)) + + +def test_populated_triton_default_is_grandfathered(env): + root = env / ".triton" / "cache" + populate(root) + + assert disk_caches.mark_root(disk_caches.get("triton")) is True + assert (root / MARKER).is_file() + + +def test_populated_non_default_root_is_not_marked_or_cleared(env, tmp_path): + root = tmp_path / "inductor" + populate(root) + + assert disk_caches.mark_root(disk_caches.get("inductor")) is False + assert not (root / MARKER).exists() + info = disk_caches.scan(disk_caches.get("inductor")) + assert info["clearable"] is False + assert MARKER in info["not_clearable_reason"] + with pytest.raises(NotClearable): + disk_caches.clear("inductor") + assert (root / "a.bin").exists() + assert (root / "sub" / "deep" / "c.bin").exists() + + +def test_triton_home_populated_is_not_grandfathered(env, tmp_path, monkeypatch): + monkeypatch.setenv("TRITON_HOME", str(tmp_path / "th")) + root = tmp_path / "th" / ".triton" / "cache" + populate(root) + + assert disk_caches.mark_root(disk_caches.get("triton")) is False + assert not (root / MARKER).exists() + assert disk_caches.scan(disk_caches.get("triton"))["clearable"] is False + with pytest.raises(NotClearable): + disk_caches.clear("triton") + assert (root / "a.bin").exists() + + +def test_mark_roots_swallows_oserror(env, monkeypatch, caplog): + def refuse(root): + raise PermissionError("read-only") + + monkeypatch.setattr(disk_caches, "_write_marker", refuse) + with caplog.at_level("WARNING", logger="dw"): + disk_caches.mark_roots() + + assert "Could not mark the inductor cache" in caplog.text + + +# ----------------------------------------------------------------- clear + + +def marked_root(tmp_path): + root = tmp_path / "inductor" + result = populate(root) + (root / MARKER).write_text("marked") + return root, result + + +def test_clear_empties_root_and_keeps_marker(env, tmp_path): + root, (total, count) = marked_root(tmp_path) + + result = disk_caches.clear("inductor") + + assert result == {"freed": total, "files": count} + assert root.is_dir() + assert [p.name for p in root.iterdir()] == [MARKER] + + +def test_clear_keeps_a_file_written_during_the_clear(env, tmp_path, monkeypatch): + root, _ = marked_root(tmp_path) + real_rmtree = disk_caches.shutil.rmtree + + def rmtree_after_a_write(path, *args, **kwargs): + (root / "new.bin").write_bytes(b"new") + return real_rmtree(path, *args, **kwargs) + + monkeypatch.setattr("dw.disk_caches.shutil.rmtree", rmtree_after_a_write) + disk_caches.clear("inductor") + + assert (root / "new.bin").exists() + assert not (root / "a.bin").exists() + assert not (root / "sub").exists() + + +def test_clear_removes_a_stale_clearing_dir(env, tmp_path): + root, _ = marked_root(tmp_path) + stale = root / ".clearing-123-abcd" + stale.mkdir() + (stale / "old.bin").write_bytes(b"x" * 5) + + disk_caches.clear("inductor") + + assert sorted(p.name for p in root.iterdir()) == [MARKER] + + +@pytest.mark.parametrize("name", ["nope", "../x"]) +def test_unknown_cache_names_the_known_ones(name): + with pytest.raises(UnknownCache) as raised: + disk_caches.clear(name) + assert "inductor" in str(raised.value) + assert "triton" in str(raised.value) + + +# ----------------------------------------------------------------- routes + + +@pytest.fixture(params=[None]) +def api(env, tmp_path, request, monkeypatch): + # No UI unless a test asks for one: with ui/dist built, its StaticFiles + # mount at / answers a path no route matches, so a route test would + # depend on whether this checkout has built the UI + ui_dir = request.param + if ui_dir is None: + monkeypatch.setattr("dw.server.app.default_ui_dir", lambda: None) + else: + ui_dir = tmp_path / ui_dir + ui_dir.mkdir() + (ui_dir / "index.html").write_text("") + workflow_dir = tmp_path / "workflows" + workflow_dir.mkdir() + prompt_dir = tmp_path / "prompts" + prompt_dir.mkdir() + manager = JobManager( + str(tmp_path / "outputs"), + worker_manager=ScriptedWorkerManager(success_script), + history_path=str(tmp_path / "jobs.sqlite"), + ) + app = create_app( + workflow_dir=str(workflow_dir), + output_dir=str(tmp_path / "outputs"), + job_manager=manager, + prompt_dir=str(prompt_dir), + ui_dir=str(ui_dir) if ui_dir else None, + ) + with TestClient(app, base_url="http://localhost") as client: + client.manager = manager + yield client + + +def tree(root): + return sorted(str(p.relative_to(root)) for p in root.rglob("*")) + + +class TestCacheRoutes: + def test_list_reports_both_caches(self, api, tmp_path): + root, (total, count) = marked_root(tmp_path) + + response = api.get("/api/caches") + + assert response.status_code == 200 + caches = {c["name"]: c for c in response.json()["caches"]} + assert set(caches) == {"inductor", "triton", "quantized"} + keys = { + "name", + "root", + "enabled", + "bytes", + "files", + "last_used", + "clearable", + "cost_of_clearing", + "disk_free", + "disk_total", + } + for cache in caches.values(): + assert keys <= set(cache) + assert cache["cost_of_clearing"] + inductor, triton = caches["inductor"], caches["triton"] + assert inductor["root"] == str(root) + assert inductor["bytes"] == total + assert inductor["files"] == count + assert inductor["clearable"] is True + assert "not_clearable_reason" not in inductor + assert "shared" in triton["cost_of_clearing"] + assert triton["clearable"] is False + assert triton["not_clearable_reason"] + + def test_clear_empties_a_marked_root(self, api, tmp_path): + root, (total, count) = marked_root(tmp_path) + + response = api.delete("/api/caches/inductor") + + assert response.status_code == 200 + assert response.json() == { + "name": "inductor", + "cleared": True, + "freed": total, + "files": count, + } + assert tree(root) == [MARKER] + + def test_clear_refused_while_busy(self, api, tmp_path): + root, _ = marked_root(tmp_path) + before = tree(root) + api.manager.is_busy = lambda: True + + response = api.delete("/api/caches/inductor") + + assert response.status_code == 409 + detail = response.json()["detail"] + assert "running" in detail or "queued" in detail + assert tree(root) == before + + def test_clear_refused_for_unmarked_root(self, api, tmp_path): + root = tmp_path / "inductor" + populate(root) + before = tree(root) + + response = api.delete("/api/caches/inductor") + + assert response.status_code == 400 + assert MARKER in response.json()["detail"] + assert tree(root) == before + + def test_unknown_name_is_404(self, api, tmp_path): + marked_root(tmp_path) + + response = api.delete("/api/caches/nope") + + assert response.status_code == 404 + assert "inductor" in response.json()["detail"] + assert "triton" in response.json()["detail"] + + # The decoded name has a "/", so no cache route matches: without a UI + # nothing does (404); with one, the UI's StaticFiles mount at / takes the + # path and refuses a DELETE (405). Either way the handler never runs. + @pytest.mark.parametrize( + "api, status", [(None, 404), ("ui", 405)], indirect=["api"] + ) + def test_traversal_name_touches_nothing(self, api, tmp_path, status): + root, _ = marked_root(tmp_path) + victim = tmp_path / "x" / "victim.bin" + victim.parent.mkdir() + victim.write_bytes(b"x" * 8) + before_root = tree(root) + before_parent = tree(tmp_path) + + response = api.delete("/api/caches/..%2Fx") + + assert response.status_code == status + assert tree(root) == before_root + assert victim.read_bytes() == b"x" * 8 + # The server's own workspace lives under tmp_path too (outputs, the + # job history's -shm/-wal), and it adds files there on its own + # schedule, so the parent is checked for losses, not for equality + assert set(before_parent) <= set(tree(tmp_path)) + + +# ------------------------------------------------------- quantized entries + +KEY_A = "a" * 64 +KEY_B = "b" * 64 + + +def quantized_entry(root, key, sizes=(100, 200), entry=True, **fields): + """An entry directory as quantized_cache.save_entry leaves it.""" + path = root / key + path.mkdir(parents=True) + for i, size in enumerate(sizes): + (path / f"model-{i}.safetensors").write_bytes(b"q" * size) + if entry is True: + record = { + "model_name": "MiniMaxAI/MiniMax-H3", + "subfolder": "vae", + "commit": "c" * 40, + "key": key, + "bytes": sum(sizes), + "created": 1700000000.0, + "last_hit": None, + **fields, + } + (path / "entry.json").write_text(json.dumps(record)) + elif entry is not None: + (path / "entry.json").write_text(entry) + return path + + +def quantized_root(tmp_path): + root = tmp_path / "quantized" + root.mkdir() + (root / MARKER).write_text("marked") + return root + + +def test_quantized_root_and_enabled_follow_the_env(env, tmp_path, monkeypatch): + cache = disk_caches.get("quantized") + assert cache.root() == str(tmp_path / "quantized") + assert cache.enabled() is False + monkeypatch.setenv("DW_QUANTIZED_CACHE", "1") + assert cache.enabled() is True + + +def test_entries_parsed_from_entry_json(env, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A, sizes=(100, 200), last_hit=1700000500.0) + quantized_entry( + root, KEY_B, sizes=(50,), subfolder="transformer", created=1700000900.0 + ) + (root / "not-a-key").mkdir() + (root / f"{KEY_A}.tmp-1-abcd").mkdir() + + info = disk_caches.scan(disk_caches.get("quantized")) + + by_id = {e["id"]: e for e in info["entries"]} + assert set(by_id) == {KEY_A, KEY_B} + # newest first + assert [e["id"] for e in info["entries"]] == [KEY_B, KEY_A] + a = by_id[KEY_A] + assert a["model"] == "MiniMaxAI/MiniMax-H3" + assert a["subfolder"] == "vae" + assert a["component"] == "vae" + assert a["created"] == 1700000000.0 + assert a["last_hit"] == 1700000500.0 + assert a["unreadable"] is False + # the files on disk, entry.json included + entry_json = (root / KEY_A / "entry.json").stat().st_size + assert a["bytes"] == 300 + entry_json + assert a["files"] == 3 + assert by_id[KEY_B]["component"] == "transformer" + + +@pytest.mark.parametrize("entry", [None, "{not json", "[1, 2]"]) +def test_unreadable_entry_is_listed_and_deletable(env, tmp_path, entry): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A, sizes=(40,), entry=entry) + + [listed] = disk_caches.scan(disk_caches.get("quantized"))["entries"] + + assert listed["id"] == KEY_A + assert listed["unreadable"] is True + assert listed["model"] is None + assert listed["component"] is None + + result = disk_caches.delete_entry("quantized", KEY_A) + + assert result["freed"] >= 40 + assert not (root / KEY_A).exists() + + +def test_off_cache_still_lists_and_clears(env, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + + info = disk_caches.scan(disk_caches.get("quantized")) + assert info["enabled"] is False + assert [e["id"] for e in info["entries"]] == [KEY_A] + assert info["clearable"] is True + + disk_caches.clear("quantized") + + assert [p.name for p in root.iterdir()] == [MARKER] + + +def test_missing_quantized_root_lists_no_entries_and_is_not_created(env, tmp_path): + disk_caches.mark_roots() + + info = disk_caches.scan(disk_caches.get("quantized")) + + assert info["entries"] == [] + assert info["bytes"] == 0 + assert not (tmp_path / "quantized").exists() + + +def test_populated_quantized_default_is_grandfathered(env, tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HOME", str(tmp_path / "hf")) + monkeypatch.delenv("DW_QUANTIZED_CACHE_DIR") + root = tmp_path / "hf" / "dw-quantized" + quantized_entry(root, KEY_A) + + assert disk_caches.mark_root(disk_caches.get("quantized")) is True + assert (root / MARKER).is_file() + + assert disk_caches.get("quantized").default_root() == str(root) + + +def test_populated_non_default_quantized_root_is_not_marked(env, tmp_path): + root = tmp_path / "quantized" + quantized_entry(root, KEY_A) + + assert disk_caches.mark_root(disk_caches.get("quantized")) is False + with pytest.raises(NotClearable): + disk_caches.delete_entry("quantized", KEY_A) + assert (root / KEY_A).is_dir() + + +def test_delete_entry_removes_only_that_entry(env, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A, sizes=(100,)) + quantized_entry(root, KEY_B, sizes=(7,)) + + result = disk_caches.delete_entry("quantized", KEY_A) + + assert result["freed"] >= 100 + assert sorted(p.name for p in root.iterdir()) == sorted([MARKER, KEY_B]) + + +@pytest.mark.parametrize( + "entry", + ["abc", "A" * 64, "a" * 63, "a" * 65, "../" + "a" * 61, "g" * 64, "", None], +) +def test_malformed_id_refused_before_any_path_is_built( + env, tmp_path, monkeypatch, entry +): + quantized_root(tmp_path) + built = [] + monkeypatch.setattr( + "dw.disk_caches.validate_path", + lambda *a, **k: built.append(a) or a[0], + ) + monkeypatch.setattr( + "dw.quantized_cache.validate_path", + lambda *a, **k: built.append(a) or a[0], + ) + monkeypatch.setattr( + "dw.disk_caches.os.path.join", lambda *a: built.append(a) or "/nowhere" + ) + + with pytest.raises(NoEntries): + disk_caches.delete_entry("quantized", entry) + assert built == [] + + +def test_entry_ops_on_a_cache_without_entries(env, tmp_path): + marked_root(tmp_path) + with pytest.raises(NoEntries): + disk_caches.delete_entry("inductor", KEY_A) + + +def test_unknown_entry(env, tmp_path): + quantized_root(tmp_path) + with pytest.raises(UnknownEntry): + disk_caches.delete_entry("quantized", KEY_A) + + +@pytest.mark.parametrize("free, low", [(10 * 2**30, True), (60 * 2**30, False)]) +def test_low_disk_on_both_arms_of_the_floor(env, tmp_path, monkeypatch, free, low): + quantized_root(tmp_path) + total = 1000 * 2**30 # floor = min(50 GiB, 100 GiB) = 50 GiB + usage = disk_caches.shutil._ntuple_diskusage(total, total - free, free) + monkeypatch.setattr("dw.disk_caches.shutil.disk_usage", lambda path: usage) + + info = disk_caches.scan(disk_caches.get("quantized")) + + assert info["low_disk"] is low + assert "low_disk" not in disk_caches.scan(disk_caches.get("inductor")) + + +def test_low_disk_reads_the_quantized_floor(env): + from dw import quantized_cache + + assert disk_caches.get("quantized").low_disk_floor is quantized_cache.floor_bytes + + +class TestEntryRoutes: + def test_list_carries_entries_and_low_disk(self, api, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + + caches = {c["name"]: c for c in api.get("/api/caches").json()["caches"]} + + quantized = caches["quantized"] + assert isinstance(quantized["low_disk"], bool) + [entry] = quantized["entries"] + assert entry["id"] == KEY_A + assert entry["component"] == "vae" + assert "entries" not in caches["inductor"] + + def test_delete_entry(self, api, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A, sizes=(100,)) + + response = api.delete(f"/api/caches/quantized/entries/{KEY_A}") + + assert response.status_code == 200 + body = response.json() + assert body["name"] == "quantized" + assert body["entry"] == KEY_A + assert body["deleted"] is True + assert body["freed"] >= 100 + assert not (root / KEY_A).exists() + + @pytest.mark.parametrize("entry", ["abc", "A" * 64, "z" * 64]) + def test_malformed_id_is_400(self, api, tmp_path, entry): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + before = tree(root) + + response = api.delete(f"/api/caches/quantized/entries/{entry}") + + assert response.status_code == 400 + assert tree(root) == before + + def test_cache_without_entries_is_400(self, api, tmp_path): + marked_root(tmp_path) + response = api.delete(f"/api/caches/inductor/entries/{KEY_A}") + assert response.status_code == 400 + + def test_unknown_entry_is_404(self, api, tmp_path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + before = tree(root) + + response = api.delete(f"/api/caches/quantized/entries/{KEY_B}") + + assert response.status_code == 404 + assert tree(root) == before + + def test_unknown_cache_is_404(self, api, tmp_path): + response = api.delete(f"/api/caches/nope/entries/{KEY_A}") + assert response.status_code == 404 + + @pytest.mark.parametrize("path", [f"/entries/{KEY_A}", ""]) + def test_refused_while_busy(self, api, tmp_path, path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + before = tree(root) + api.manager.is_busy = lambda: True + + response = api.delete(f"/api/caches/quantized{path}") + + assert response.status_code == 409 + assert tree(root) == before + + @pytest.mark.parametrize("path", [f"/entries/{KEY_A}", ""]) + def test_refused_while_a_download_is_active(self, api, tmp_path, path): + root = quantized_root(tmp_path) + quantized_entry(root, KEY_A) + before = tree(root) + api.app.state.downloads.is_active = lambda: True + + response = api.delete(f"/api/caches/quantized{path}") + + assert response.status_code == 409 + assert "download" in response.json()["detail"] + assert tree(root) == before + + def test_a_download_does_not_block_other_caches(self, api, tmp_path): + marked_root(tmp_path) + api.app.state.downloads.is_active = lambda: True + + assert api.delete("/api/caches/inductor").status_code == 200 + + +def test_triton_default_created_after_start_is_clearable(env): + triton = disk_caches.get("triton") + disk_caches.mark_roots() + assert disk_caches.scan(triton)["clearable"] is False + + # a compile creates the root after the server started + populate(env / ".triton" / "cache") + + info = disk_caches.scan(triton) + assert info["clearable"] is True + assert info["files"] > 0 + assert disk_caches.clear("triton")["files"] > 0 diff --git a/tests/test_inherited_cost.py b/tests/test_inherited_cost.py index dcd4c17c..249f37af 100644 --- a/tests/test_inherited_cost.py +++ b/tests/test_inherited_cost.py @@ -193,3 +193,33 @@ def test_the_estimate_is_inherited_and_the_validate_answer_warns(self, tmp_path) def test_an_observed_estimate_does_not_warn(self): assert _inherited_cost_warnings({"basis": "observed"}) == [] + + +class TestNotQuotedUnchangedWhenItScales: + """#821: the template's minutes are for one template-sized run.""" + + HISTORY = {"templates/t": [run(8 * MINUTE, {"num_frames": 124})]} + + def test_several_steps_of_one_pipeline_are_not_priced_as_one_run(self): + inline = definition("inline") + inline["steps"] = [{**inline["steps"][0], "id": f"shot{n}"} for n in range(3)] + assert ( + inherited(inline, [("templates/t", definition("t"))], self.HISTORY) is None + ) + + def test_a_literal_longer_frame_count_is_not_priced_as_the_templates(self): + inline = definition("inline") + inline["variables"] = {} + inline["steps"][0]["run"] = {"num_frames": 899} + assert ( + inherited(inline, [("templates/t", definition("t"))], self.HISTORY) is None + ) + + def test_a_shared_driver_name_still_prices_its_bucket(self): + block = inherited( + definition("inline"), + [("templates/t", definition("t"))], + self.HISTORY, + {"num_frames": 124}, + ) + assert block["cold_minutes"] == 8.0 diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index f5520325..9f16d234 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -771,6 +771,9 @@ def handler(request): # The full get_workflow also reads the variables route for `observed` (#786) if name == "get_workflow": seen = seen[-1:] + # get_server_info also reads /api/caches for `caches` (#828) + if name == "get_server_info": + seen = seen[:1] assert seen == [(method, path)] assert result.content @@ -1562,6 +1565,10 @@ def test_the_stated_tool_count_is_the_registered_one(): # ceiling moves to 14_265 with 10 of headroom. # 2026-10-07 (#682): get_job and list_jobs name a job's `device`, paid for by # trimming their own text - measured at 14_264.75; the ceiling stays. +# 2026-10-11 (#828 stage A): get_server_info's one clause on `caches` (16 +# tokens) - measured 14_246.25 before and 14_262.25 after (9_281.75 / +# 4_470.25 / 510.25). Text trimmed since 10-07 left room, so the ceiling +# stays, 2.75 of headroom. SURFACE_BUDGET = 14_265 diff --git a/tests/test_mcp_workspaces.py b/tests/test_mcp_workspaces.py index c110e52e..e6f7a74c 100644 --- a/tests/test_mcp_workspaces.py +++ b/tests/test_mcp_workspaces.py @@ -245,18 +245,96 @@ def handler(request): }, }, ) + if request.url.path == "/api/caches": + return httpx.Response(200, json={"caches": []}) raise AssertionError(f"unexpected request {request.url.path}") client = DwClient(transport=httpx.MockTransport(handler)) client.workspace = "shots" result = server_info(client) - assert seen == [("/api/server", "shots")] + assert seen == [("/api/server", "shots"), ("/api/caches", "shots")] assert result["workspace"] == "shots" assert result["directories"]["workflows"] == "/studio/shots/workflows" assert result["directories"]["prompts"] == "/home/user/prompts" assert result["device"] == "cuda" + @staticmethod + def info_client(caches_response): + def handler(request): + if request.url.path == "/api/server": + return httpx.Response(200, json={"device": "cuda", "version": "0.1.0"}) + if request.url.path == "/api/caches": + return caches_response + raise AssertionError(f"unexpected request {request.url.path}") + + return DwClient(transport=httpx.MockTransport(handler)) + + def test_caches_carry_only_the_size_fields(self): + listing = { + "caches": [ + { + "name": "inductor", + "root": "/home/user/.cache/dw/torchinductor", + "enabled": True, + "bytes": 1234, + "files": 5, + "last_used": 1700000000.0, + "clearable": True, + "cost_of_clearing": "recompiles", + "disk_free": 99, + "disk_total": 1000, + } + ] + } + client = self.info_client(httpx.Response(200, json=listing)) + + result = server_info(client) + + assert result["caches"] == [ + { + "name": "inductor", + "bytes": 1234, + "files": 5, + "enabled": True, + "clearable": True, + "disk_free": 99, + } + ] + + def test_a_cache_with_entries_carries_only_their_count(self): + listing = { + "caches": [ + { + "name": "quantized", + "bytes": 300, + "files": 3, + "enabled": True, + "clearable": True, + "disk_free": 99, + "low_disk": False, + "entries": [{"id": "a" * 64}, {"id": "b" * 64}], + } + ] + } + client = self.info_client(httpx.Response(200, json=listing)) + + [quantized] = server_info(client)["caches"] + + assert quantized["entry_count"] == 2 + assert "entries" not in quantized + + @pytest.mark.parametrize("status", [404, 500]) + def test_caches_omitted_when_the_listing_fails(self, status): + client = self.info_client(httpx.Response(status, json={"detail": "no"})) + + result = server_info(client) + + assert "caches" not in result + assert result["device"] == "cuda" + assert result["version"] == "0.1.0" + assert result["workspace"] == DEFAULT_WORKSPACE + class TestPerCallPin: """`workspace=` on an output-side tool: for this one call, without diff --git a/tests/test_media_info.py b/tests/test_media_info.py index f91ec4a2..ff5c17d2 100644 --- a/tests/test_media_info.py +++ b/tests/test_media_info.py @@ -267,6 +267,65 @@ def test_one_entry_per_second_of_the_track(self, tmp_path): assert len(envelope["rms_dbfs"]) == 4 assert len(envelope["peak_dbfs"]) == 4 + def test_a_finer_interval_gives_finer_bins(self, tmp_path): + self.write_gapped_wav(tmp_path / "score.wav", seconds=4.0) + + envelope = probe_media( + str(tmp_path / "score.wav"), envelope=True, interval_seconds=0.1 + )["envelope"] + + assert envelope["interval_seconds"] == pytest.approx(0.1) + assert len(envelope["rms_dbfs"]) == 40 + # the silent second (index 2) is ten silent 0.1 s bins + assert all(db == pytest.approx(-120.0) for db in envelope["rms_dbfs"][20:30]) + assert envelope["rms_dbfs"][19] > -20.0 + + def test_an_interval_below_the_floor_is_refused(self, tmp_path): + write_wav(tmp_path / "score.wav", seconds=2.0) + + with pytest.raises(ValueError, match="at least 0.05"): + probe_media( + str(tmp_path / "score.wav"), envelope=True, interval_seconds=0.01 + ) + + def test_a_track_over_the_bin_cap_is_refused_with_the_interval_that_fits( + self, tmp_path, monkeypatch + ): + write_wav(tmp_path / "score.wav", seconds=4.0) + monkeypatch.setattr("dw.media.MAX_ENVELOPE_BINS", 20) + + with pytest.raises(ValueError, match="interval_seconds=0.2 or more"): + probe_media( + str(tmp_path / "score.wav"), envelope=True, interval_seconds=0.1 + ) + + def test_the_bin_cap_never_refuses_the_default_interval( + self, tmp_path, monkeypatch + ): + """What master answered at interval_seconds=1.0 stays answered: the + cap bounds a finer envelope, not a long track at the default.""" + write_wav(tmp_path / "score.wav", seconds=4.0) + monkeypatch.setattr("dw.media.MAX_ENVELOPE_BINS", 2) + + envelope = probe_media(str(tmp_path / "score.wav"), envelope=True)["envelope"] + assert len(envelope["rms_dbfs"]) == 4 + + with pytest.raises(ValueError, match="interval_seconds=1.0 or more"): + probe_media( + str(tmp_path / "score.wav"), envelope=True, interval_seconds=0.5 + ) + + @pytest.mark.parametrize("interval", [math.inf, math.nan, 1e305, 3601.0]) + def test_an_interval_past_the_ceiling_is_refused(self, tmp_path, interval): + """inf and 1e305 would overflow the bin-sample product; NaN fails + the floor check first, which is as good a refusal.""" + write_wav(tmp_path / "score.wav", seconds=2.0) + + with pytest.raises(ValueError, match="at most 3600|at least 0.05"): + probe_media( + str(tmp_path / "score.wav"), envelope=True, interval_seconds=interval + ) + def test_it_finds_the_second_the_sound_stops_in(self, tmp_path): self.write_gapped_wav(tmp_path / "score.wav", seconds=4.0, silent_second=2) diff --git a/tests/test_plan_cuts.py b/tests/test_plan_cuts.py index da807a6d..0e9d512c 100644 --- a/tests/test_plan_cuts.py +++ b/tests/test_plan_cuts.py @@ -55,6 +55,95 @@ def plan(**kwargs): return plan_cuts(**kwargs) +def test_suspect_chunks_become_warnings_naming_span_and_text(): + # #822: transcribe_audio marks a chunk it doubts; plan_cuts says so + chunks = [dict(c) for c in CHUNKS] + chunks.append( + {"start": 26.0, "end": 29.0, "text": "Редактор субтитров", "suspect": "silent"} + ) + chunks[0]["suspect"] = "long" + result = plan(transcript={"text": "x", "chunks": chunks}) + suspect = [w for w in result["warnings"] if "suspect" in w] + assert len(suspect) == 2 + assert any( + "26.00-29.00" in w and "Редактор субтитров" in w and "(silent)" in w + for w in suspect + ) + assert any("4.00-7.00" in w and "(long)" in w for w in suspect) + + +def test_a_hallucinated_intro_and_fade_are_left_out_of_the_lines(): + # #822 bounce (s23): hundreds of '—' chunks over the intro and an + # "I'm sorry" loop over the fade were laid out as sung shots from frame 0, + # and the warnings named every chunk + intro = [ + {"start": 0.0, "end": 0.36, "text": "—", "suspect": "no_words"}, + {"start": 0.36, "end": 4.8, "text": "—", "suspect": "no_words"}, + ] + [ + { + "start": 4.8 + i * 0.1, + "end": 4.8 + i * 0.1, + "text": "—", + "suspect": "no_words", + } + for i in range(200) + ] + sung_lines = [ + {"start": 30.0, "end": 33.0, "text": "Soyuz nerushimy"}, + {"start": 34.0, "end": 37.0, "text": "respublik svobodnykh"}, + ] + fade = [ + { + "start": 40.0 + i * 0.2, + "end": 40.0 + i * 0.2, + "text": t, + "suspect": "repetition", + } + for i, t in enumerate(["I'm", "sorry,"] * 120) + ] + result = plan( + transcript={"text": "x", "chunks": intro + sung_lines + fade}, duration_s=90.0 + ) + lyrics = [s["lyric"] for s in result["shots"] if s["lyric"]] + assert lyrics == ["Soyuz nerushimy", "respublik svobodnykh"] + first_vocal = next(s for s in result["shots"] if s["kind"] == "vocal") + assert first_vocal["start_frame"] == 30 * FPS + suspect = [w for w in result["warnings"] if "suspect" in w] + assert len(suspect) == 2 + assert any( + "202 transcript chunk(s) 0.00-" in w and "(no_words)" in w for w in suspect + ) + assert any( + "240 transcript chunk(s) 40.00-" in w and "(repetition)" in w for w in suspect + ) + assert all("left out of the lines" in w for w in suspect) + assert sum(len(w) for w in result["warnings"]) < 2000 + + +def test_a_long_word_is_still_sung(): + chunks = [dict(c) for c in CHUNKS] + chunks[0]["suspect"] = "long" + result = plan(transcript={"text": "x", "chunks": chunks}) + assert "Walking in the moonlight tonight" in sung(result) + assert any("(long); still used" in w for w in result["warnings"]) + + +def test_lines_squeezed_out_back_to_back_make_one_bounded_warning(): + # #822 bounce: a run of zero-length lines warned once per line, each + # warning repeating every line before it - quadratic in the run + squeezed = [{"start": 10.0, "end": 10.0, "text": f"word{i}"} for i in range(150)] + chunks = CHUNKS[:1] + squeezed + CHUNKS[2:] + result = plan(transcript={"text": "x", "chunks": chunks}) + no_time = [w for w in result["warnings"] if "left no time" in w] + assert len(no_time) == 1 + assert "word0 / word1 / word2 / ... (147 more)" in no_time[0] + assert len(no_time[0]) < 300 + + +def test_unmarked_chunks_add_no_suspect_warning(): + assert not any("suspect" in w for w in plan()["warnings"]) + + def sung(result): return "\n".join(s["lyric"] for s in result["shots"] if s["lyric"]).split("\n") diff --git a/tests/test_quantized_cache.py b/tests/test_quantized_cache.py new file mode 100644 index 00000000..878c4429 --- /dev/null +++ b/tests/test_quantized_cache.py @@ -0,0 +1,1597 @@ +""" +Unit tests for the on-disk cache of SDNQ-quantized components (#815, #826). + +Covers the key, the hub-snapshot rule, the disk floor, save/load failure +handling, the sweep, the switch and directory, the load_components() +orchestration against a fake modular pipeline, one real sdnq round trip on +CPU, and the two call sites (load_component and the server start). +""" + +import errno +import json +import logging +import os +import subprocess +import sys +from collections import namedtuple +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from huggingface_hub.file_download import repo_folder_name +from torch import nn + +import dw.quantized_cache as qc +from dw.events import RunContext, activate_context, deactivate_context + +sdnq = pytest.importorskip("sdnq") +from sdnq import SDNQConfig # noqa: E402 + +GIB = 2**30 +COMMIT = "a" * 40 +OTHER_COMMIT = "b" * 40 +KEY = "c" * 64 +VERSIONS = {"sdnq": "1", "diffusers": "2", "transformers": "3", "torch": "4"} +Usage = namedtuple("Usage", "total used free") + + +@pytest.fixture(autouse=True) +def _clean_environment(monkeypatch): + # set then delete, so monkeypatch restores whatever the run began with + for name in (qc.ENABLE_ENV_VAR, qc.DIR_ENV_VAR): + monkeypatch.setenv(name, "x") + monkeypatch.delenv(name) + monkeypatch.setattr(qc, "installed_versions", lambda: dict(VERSIONS)) + + +@pytest.fixture +def events(): + collected = [] + token = activate_context(RunContext(on_event=collected.append)) + yield collected + deactivate_context(token) + + +@pytest.fixture +def cache_on(monkeypatch, tmp_path): + root = tmp_path / "cache" + monkeypatch.setenv(qc.ENABLE_ENV_VAR, "1") + monkeypatch.setenv(qc.DIR_ENV_VAR, str(root)) + return root + + +def make_snapshot( + hub, + model="org/name", + commit=COMMIT, + subfolder=None, + files=("model.safetensors",), + config=True, + ref="main", +): + repo = hub / repo_folder_name(repo_id=model, repo_type="model") + (repo / "refs").mkdir(parents=True, exist_ok=True) + if ref: + (repo / "refs" / ref).write_text(commit) + folder = repo / "snapshots" / commit + if subfolder: + folder = folder / subfolder + folder.mkdir(parents=True, exist_ok=True) + if config: + (folder / "config.json").write_text("{}") + for name in files: + (folder / name).write_bytes(b"weights") + return folder + + +def disk(total_gib, free_gib): + return lambda _path: Usage(total_gib * GIB, 0, int(free_gib * GIB)) + + +def fake_save(component, path): + os.makedirs(path) + with open(os.path.join(path, "model.safetensors"), "wb") as f: + f.write(b"x" * 100) + + +def entry_json(directory, **fields): + entry = {"versions": dict(VERSIONS), "bytes": 10, **fields} + directory.mkdir(parents=True, exist_ok=True) + (directory / qc.ENTRY_FILE).write_text(json.dumps(entry)) + + +def base_inputs(**overrides): + values = { + "model_name": "org/name", + "subfolder": "vae", + "commit": COMMIT, + "quantizer": qc.quantizer_arguments(SDNQConfig(weights_dtype="int8")), + "torch_dtype": torch.bfloat16, + } + values.update(overrides) + return qc.key_inputs(**values) + + +# 1. The key + + +class TestKey: + def test_identical_inputs_give_identical_key(self): + assert qc.cache_key(base_inputs()) == qc.cache_key(base_inputs()) + assert qc._KEY_PATTERN.match(qc.cache_key(base_inputs())) + + @pytest.mark.parametrize( + "overrides", + [ + {"model_name": "org/other"}, + {"subfolder": "unet"}, + {"commit": OTHER_COMMIT}, + {"quantizer": qc.quantizer_arguments(SDNQConfig(weights_dtype="int4"))}, + {"torch_dtype": torch.float16}, + ], + ) + def test_each_input_changes_the_key(self, overrides): + assert qc.cache_key(base_inputs(**overrides)) != qc.cache_key(base_inputs()) + + def test_package_version_changes_the_key(self, monkeypatch): + before = qc.cache_key(base_inputs()) + monkeypatch.setattr( + qc, "installed_versions", lambda: {**VERSIONS, "sdnq": "other"} + ) + assert qc.cache_key(base_inputs()) != before + + def test_installed_versions_reads_the_four_packages(self, monkeypatch): + monkeypatch.undo() + qc.installed_versions.cache_clear() + versions = qc.installed_versions() + assert set(versions) == set(qc.VERSIONED_PACKAGES) + assert all(isinstance(v, str) and v for v in versions.values()) + + def test_placement_arguments_do_not_change_the_key(self): + plain = SDNQConfig(weights_dtype="int8") + placed = SDNQConfig( + weights_dtype="int8", + quantization_device="cpu", + return_device="cpu", + use_quantized_matmul=True, + ) + a = base_inputs(quantizer=qc.quantizer_arguments(plain)) + b = base_inputs(quantizer=qc.quantizer_arguments(placed)) + assert qc.cache_key(a) == qc.cache_key(b) + + def test_the_key_does_not_follow_the_process_hash_seed(self): + # SDNQConfig dedupes modules_to_not_convert through a set, so each + # worker process holds it in its own order - H3's transformer keyed + # differently on each card until the key sorted it + script = ( + "import json\n" + "from sdnq import SDNQConfig\n" + "import dw.quantized_cache as qc\n" + "config = SDNQConfig(weights_dtype='int4', use_quantized_matmul=True," + " modules_to_not_convert=[f'block_{i}' for i in range(9)]," + " modules_dtype_dict={'int8': ['a', 'b', 'c', 'd']})\n" + "print(json.dumps([config.modules_to_not_convert," + " qc.quantizer_arguments(config)]))\n" + ) + orders, quantizers = set(), set() + for seed in ("1", "2", "3", "4"): + out = subprocess.run( + [sys.executable, "-c", script], + env={**os.environ, "PYTHONHASHSEED": seed}, + capture_output=True, + text=True, + check=True, + ).stdout.splitlines()[-1] + order, quantizer = json.loads(out) + orders.add(tuple(order)) + quantizers.add(json.dumps(quantizer, sort_keys=True)) + assert len(orders) > 1, "the seeds did not reorder the list" + assert len(quantizers) == 1 + + def test_quantizer_arguments_drop_only_the_unkeyed_ones(self): + described = qc.quantizer_arguments(SDNQConfig(weights_dtype="int8")) + assert described["class"].endswith("SDNQConfig") + assert described["arguments"]["weights_dtype"] == "int8" + for name in qc.UNKEYED_ARGUMENTS: + assert name not in described["arguments"] + + def test_empty_subfolder_and_none_are_the_same(self): + assert qc.cache_key(base_inputs(subfolder=None)) == qc.cache_key( + base_inputs(subfolder="") + ) + + def test_entry_path_refuses_a_non_key(self, tmp_path): + with pytest.raises(ValueError): + qc.entry_path(str(tmp_path), "../escape") + + +# 2. The snapshot rule + + +class TestSnapshot: + def test_complete_single_safetensors(self, tmp_path): + make_snapshot(tmp_path, subfolder="vae") + assert qc.lookup_commit("org/name", "vae", None, tmp_path) == COMMIT + + def test_complete_without_subfolder(self, tmp_path): + make_snapshot(tmp_path) + assert qc.lookup_commit("org/name", None, "main", tmp_path) == COMMIT + + def test_sharded_index_with_every_shard(self, tmp_path): + folder = make_snapshot(tmp_path, files=("a-1.safetensors", "a-2.safetensors")) + index = {"weight_map": {"x": "a-1.safetensors", "y": "a-2.safetensors"}} + (folder / "model.safetensors.index.json").write_text(json.dumps(index)) + assert qc.lookup_commit("org/name", None, None, tmp_path) == COMMIT + + def test_sharded_index_with_a_missing_shard(self, tmp_path): + folder = make_snapshot(tmp_path, files=("a-1.safetensors",)) + index = {"weight_map": {"x": "a-1.safetensors", "y": "a-2.safetensors"}} + (folder / "model.safetensors.index.json").write_text(json.dumps(index)) + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_sharded_index_with_a_dangling_symlink(self, tmp_path): + folder = make_snapshot(tmp_path, files=("a-1.safetensors",)) + os.symlink(tmp_path / "no-such-blob", folder / "a-2.safetensors") + index = {"weight_map": {"x": "a-1.safetensors", "y": "a-2.safetensors"}} + (folder / "model.safetensors.index.json").write_text(json.dumps(index)) + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_symlink_to_a_real_blob_is_complete(self, tmp_path): + folder = make_snapshot(tmp_path, files=()) + blob = tmp_path / "blob" + blob.write_bytes(b"w") + os.symlink(blob, folder / "model.safetensors") + assert qc.lookup_commit("org/name", None, None, tmp_path) == COMMIT + + def test_dangling_single_file(self, tmp_path): + folder = make_snapshot(tmp_path, files=()) + os.symlink(tmp_path / "gone", folder / "model.safetensors") + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_missing_config(self, tmp_path): + make_snapshot(tmp_path, config=False) + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_variant_file_only(self, tmp_path): + make_snapshot(tmp_path, files=("model.fp16.safetensors",)) + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_no_weights_at_all(self, tmp_path): + make_snapshot(tmp_path, files=()) + assert qc.lookup_commit("org/name", None, None, tmp_path) is None + + def test_commit_revision_names_its_snapshot(self, tmp_path): + make_snapshot(tmp_path, commit=OTHER_COMMIT, ref=None) + assert qc.resolve_commit("org/name", OTHER_COMMIT, tmp_path) == OTHER_COMMIT + assert qc.resolve_commit("org/name", COMMIT, tmp_path) is None + + def test_resolution_goes_through_huggingface_hub(self, tmp_path, monkeypatch): + import huggingface_hub + + calls = [] + real = huggingface_hub.try_to_load_from_cache + + def spy(repo_id, filename, **kwargs): + calls.append((repo_id, filename, kwargs["revision"])) + return real(repo_id, filename, **kwargs) + + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", spy) + make_snapshot(tmp_path, subfolder="vae") + assert qc.lookup_commit("org/name", "vae", None, tmp_path) == COMMIT + assert ("org/name", "vae/config.json", "main") in calls + assert ("org/name", "vae/model.safetensors", COMMIT) in calls + + def test_unknown_ref_or_repo(self, tmp_path): + assert qc.resolve_commit("org/name", "main", tmp_path) is None + make_snapshot(tmp_path) + assert qc.resolve_commit("org/name", "dev", tmp_path) is None + + def test_ref_that_is_not_a_commit(self, tmp_path): + make_snapshot(tmp_path) + (tmp_path / "models--org--name" / "refs" / "main").write_text("garbage") + assert qc.resolve_commit("org/name", None, tmp_path) is None + + def test_repo_folder_linked_to_another_disk(self, tmp_path): + # a big repo moved off the hub's volume and symlinked back in + other = tmp_path / "other" + folder = make_snapshot(other, subfolder="transformer", files=()) + index = {"weight_map": {"x": "a-1.safetensors", "y": "a-2.safetensors"}} + (folder / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps(index) + ) + for shard in ("a-1.safetensors", "a-2.safetensors"): + (folder / shard).write_bytes(b"w") + hub = tmp_path / "hub" + hub.mkdir() + name = repo_folder_name(repo_id="org/name", repo_type="model") + os.symlink(other / name, hub / name) + assert qc.lookup_commit("org/name", "transformer", None, hub) == COMMIT + assert ( + qc.save_commit(nn.Linear(1, 1), "org/name", "transformer", cache_dir=hub) + == COMMIT + ) + + def test_hub_root_linked_to_another_disk(self, tmp_path): + make_snapshot(tmp_path / "real", subfolder="vae") + os.symlink(tmp_path / "real", tmp_path / "hub") + assert qc.lookup_commit("org/name", "vae", None, tmp_path / "hub") == COMMIT + + def test_blob_outside_the_repo_folder_is_not_complete(self, tmp_path): + folder = make_snapshot(tmp_path / "hub", files=()) + blob = tmp_path / "elsewhere.json" + blob.write_text(json.dumps({"weight_map": {}})) + os.symlink(blob, folder / "model.safetensors.index.json") + assert qc.lookup_commit("org/name", None, None, tmp_path / "hub") is None + + def test_hub_cache_defaults_to_the_hub_constant(self, tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path)) + make_snapshot(tmp_path) + assert qc.lookup_commit("org/name") == COMMIT + + +class TestSaveCommit: + @staticmethod + def loaded(commit): + return SimpleNamespace(config=SimpleNamespace(_commit_hash=commit)) + + def test_recorded_that_disagrees_with_refs_is_unpinned(self, tmp_path): + make_snapshot(tmp_path, commit=OTHER_COMMIT) + assert ( + qc.save_commit(self.loaded(COMMIT), "org/name", cache_dir=tmp_path) is None + ) + + def test_recorded_that_agrees_is_used(self, tmp_path): + make_snapshot(tmp_path) + assert ( + qc.save_commit(self.loaded(COMMIT), "org/name", cache_dir=tmp_path) + == COMMIT + ) + + def test_recorded_with_no_ref_is_unpinned(self, tmp_path): + # no lookup could name it, so a save under it is never read + assert ( + qc.save_commit(self.loaded(COMMIT), "org/name", cache_dir=tmp_path) is None + ) + + def test_recorded_over_an_incomplete_snapshot_is_unpinned(self, tmp_path): + make_snapshot(tmp_path, files=()) + assert ( + qc.save_commit(self.loaded(COMMIT), "org/name", cache_dir=tmp_path) is None + ) + + @pytest.mark.parametrize("recorded", [True, False]) + @pytest.mark.parametrize("complete", [True, False]) + def test_a_save_is_always_one_a_lookup_reads(self, tmp_path, recorded, complete): + make_snapshot( + tmp_path, subfolder="vae", files=("model.safetensors",) if complete else () + ) + component = self.loaded(COMMIT) if recorded else nn.Linear(1, 1) + saved = qc.save_commit(component, "org/name", "vae", cache_dir=tmp_path) + looked_up = qc.lookup_commit("org/name", "vae", None, tmp_path) + assert saved is None or saved == looked_up + + def test_recorded_in_a_dict_config(self, tmp_path): + make_snapshot(tmp_path) + component = SimpleNamespace(config={"_commit_hash": COMMIT}) + assert qc.save_commit(component, "org/name", cache_dir=tmp_path) == COMMIT + + def test_nothing_recorded_and_incomplete_is_unpinned(self, tmp_path): + make_snapshot(tmp_path, files=()) + assert qc.save_commit(nn.Linear(1, 1), "org/name", cache_dir=tmp_path) is None + + def test_nothing_recorded_and_complete_uses_the_ref(self, tmp_path): + make_snapshot(tmp_path, subfolder="vae") + assert ( + qc.save_commit(nn.Linear(1, 1), "org/name", "vae", cache_dir=tmp_path) + == COMMIT + ) + + def test_malformed_recorded_commit_is_ignored(self, tmp_path): + assert ( + qc.save_commit(self.loaded("short"), "org/name", cache_dir=tmp_path) is None + ) + + +# 3. The floor + + +class TestFloor: + def test_floor_is_ten_percent_of_a_small_volume(self): + assert qc.floor_bytes(100 * GIB) == 10 * GIB + + def test_floor_caps_at_fifty_gib_on_a_large_volume(self): + assert qc.floor_bytes(4000 * GIB) == 50 * GIB + assert qc.floor_bytes(500 * GIB) == 50 * GIB + + def test_low_disk_skips_and_writes_nothing(self, tmp_path, monkeypatch): + component = nn.Linear(100, 100) + estimate = qc.component_nbytes(component) + monkeypatch.setattr(qc.shutil, "disk_usage", disk(100, 10 + 0)) + # free - estimate < floor (10 GiB) for any estimate > 0 + outcome = qc.save_entry(str(tmp_path), KEY, component, {}) + assert estimate > 0 + assert outcome == ("skipped", qc.LOW_DISK, None) + assert os.listdir(tmp_path) == [] + + def test_one_byte_short_of_the_floor_skips(self, tmp_path, monkeypatch): + component = nn.Linear(100, 100) + estimate = qc.component_nbytes(component) + floor = 10 * GIB + monkeypatch.setattr( + qc.shutil, + "disk_usage", + lambda _p: Usage(100 * GIB, 0, floor + estimate - 1), + ) + assert qc.save_entry(str(tmp_path), KEY, component, {})[0:2] == ( + "skipped", + qc.LOW_DISK, + ) + + def test_ample_disk_saves(self, tmp_path, monkeypatch): + component = nn.Linear(100, 100) + estimate = qc.component_nbytes(component) + monkeypatch.setattr( + qc.shutil, + "disk_usage", + lambda _p: Usage(100 * GIB, 0, 10 * GIB + estimate), + ) + monkeypatch.setattr(sdnq, "save_sdnq_model", fake_save) + result, reason, nbytes = qc.save_entry( + str(tmp_path), KEY, component, {"model_name": "org/name"} + ) + assert (result, reason) == ("saved", None) + assert nbytes == 100 + entry = qc.read_entry(str(tmp_path / KEY)) + assert entry["key"] == KEY and entry["bytes"] == 100 + assert entry["model_name"] == "org/name" and entry["last_hit"] is None + assert sorted(os.listdir(tmp_path)) == [KEY] + + def test_unreachable_root_is_skipped(self, tmp_path): + blocker = tmp_path / "file" + blocker.write_text("x") + outcome = qc.save_entry(str(blocker / "sub"), KEY, nn.Linear(1, 1), {}) + assert outcome[0] == "skipped" and outcome[1].startswith("save_failed") + + +# 4. The start-up warning + + +class TestLogDisk: + @pytest.fixture(autouse=True) + def _logging(self, caplog): + caplog.set_level(logging.INFO, logger="dw") + self.caplog = caplog + + def warned(self): + return [r for r in self.caplog.records if r.levelno == logging.WARNING] + + def test_empty_cache_below_the_floor_warns(self, tmp_path, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(100, 9)) + assert qc.log_disk(str(tmp_path)) is True + assert len(self.warned()) == 1 + + def test_empty_cache_just_above_the_floor_is_quiet(self, tmp_path, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(100, 10.5)) + assert qc.log_disk(str(tmp_path)) is False + assert self.warned() == [] + + def test_margin_is_the_largest_entry(self, tmp_path, monkeypatch): + entry_json(tmp_path / ("1" * 64), bytes=2 * GIB) + entry_json(tmp_path / ("2" * 64), bytes=5 * GIB) + # floor is 10 GiB, the largest entry 5 GiB + monkeypatch.setattr(qc.shutil, "disk_usage", disk(100, 14)) + assert qc.log_disk(str(tmp_path)) is True + assert len(self.warned()) == 1 + monkeypatch.setattr(qc.shutil, "disk_usage", disk(100, 16)) + assert qc.log_disk(str(tmp_path)) is False + assert len(self.warned()) == 1 + + def test_entry_without_readable_json_is_sized_from_disk( + self, tmp_path, monkeypatch + ): + broken = tmp_path / ("3" * 64) + broken.mkdir() + (broken / "w").write_bytes(b"x" * 10) + assert qc._entry_sizes(str(tmp_path)) == [10] + + +# 5. Failures + + +class TestSaveFailures: + @pytest.fixture(autouse=True) + def _roomy(self, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + + def test_enospc(self, tmp_path, monkeypatch): + def full(component, path): + os.makedirs(path) + (tmp_path / "scratch").write_text("") + with open(os.path.join(path, "partial"), "w") as f: + f.write("x") + raise OSError(errno.ENOSPC, "No space left on device") + + monkeypatch.setattr(sdnq, "save_sdnq_model", full) + root = tmp_path / "root" + outcome = qc.save_entry(str(root), KEY, nn.Linear(2, 2), {}) + assert outcome == ("skipped", "save_failed: ENOSPC", None) + assert os.listdir(root) == [] + + def test_generic_exception_names_its_type(self, tmp_path, monkeypatch): + def boom(component, path): + os.makedirs(path) + raise RuntimeError("nope") + + monkeypatch.setattr(sdnq, "save_sdnq_model", boom) + outcome = qc.save_entry(str(tmp_path), KEY, nn.Linear(2, 2), {}) + assert outcome == ("skipped", "save_failed: RuntimeError", None) + assert os.listdir(tmp_path) == [] + + +# 6. The race + + +class TestRace: + def test_another_worker_wins(self, tmp_path, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + root = str(tmp_path) + + def racing(component, path): + fake_save(component, path) + # another worker's finished entry lands first + winner = os.path.join(root, KEY) + os.makedirs(winner) + with open(os.path.join(winner, "model.safetensors"), "wb") as f: + f.write(b"y" * 7) + + monkeypatch.setattr(sdnq, "save_sdnq_model", racing) + outcome = qc.save_entry(root, KEY, nn.Linear(2, 2), {}) + assert outcome == ("saved", qc.ALREADY_SAVED, 7) + assert os.listdir(root) == [KEY] + + +# 7. A corrupt entry + + +class TestLoadEntry: + def test_missing_entry_is_a_miss(self, tmp_path): + assert qc.load_entry(str(tmp_path), KEY, nn.Linear, None, "cpu", {}) is None + + def test_failing_load_deletes_the_entry(self, tmp_path, monkeypatch): + entry_json(tmp_path / KEY) + monkeypatch.setattr( + sdnq, "load_sdnq_model", MagicMock(side_effect=ValueError("corrupt")) + ) + with pytest.raises(ValueError): + qc.load_entry(str(tmp_path), KEY, nn.Linear, None, "cpu", {}) + assert not (tmp_path / KEY).exists() + + def test_hit_passes_arguments_and_records_last_hit(self, tmp_path, monkeypatch): + entry_json(tmp_path / KEY) + module = nn.Linear(2, 2) + loader = MagicMock(return_value=module) + monkeypatch.setattr(sdnq, "load_sdnq_model", loader) + loaded = qc.load_entry( + str(tmp_path), + KEY, + nn.Linear, + torch.float32, + "cpu", + {"dequantize_fp32": True, "use_quantized_matmul": False}, + ) + assert loaded is module + kwargs = loader.call_args.kwargs + assert kwargs["model_cls"] is nn.Linear and kwargs["dtype"] is torch.float32 + assert kwargs["dequantize_fp32"] is True + assert kwargs["use_quantized_matmul"] is False + assert qc.read_entry(str(tmp_path / KEY))["last_hit"] is not None + + +# 8. The sweep + + +class TestSweep: + def test_nothing_without_a_root(self, tmp_path): + assert qc.sweep(str(tmp_path / "absent")) == (0, 0) + + def test_deletes_only_what_is_stale(self, tmp_path): + tmp_dir = tmp_path / f"{'1' * 64}.tmp-9-abcd" + tmp_dir.mkdir() + (tmp_dir / "w").write_bytes(b"x" * 10) + (tmp_path / f"{'2' * 64}.tmp-9-file").write_bytes(b"x" * 5) + stale = tmp_path / ("3" * 64) + entry_json(stale, versions={**VERSIONS, "torch": "old"}) + (stale / "w").write_bytes(b"x" * 20) + garbage = tmp_path / ("4" * 64) + garbage.mkdir() + (garbage / qc.ENTRY_FILE).write_text("{not json") + (garbage / "w").write_bytes(b"x" * 1) + no_json = tmp_path / ("5" * 64) + no_json.mkdir() + (no_json / "w").write_bytes(b"x" * 2) + not_object = tmp_path / ("7" * 64) + not_object.mkdir() + (not_object / qc.ENTRY_FILE).write_text("[1]") + current = tmp_path / ("6" * 64) + entry_json(current) + (current / "w").write_bytes(b"x" * 99) + (tmp_path / "notes.txt").write_text("keep") + (tmp_path / "other").mkdir() + + count, freed = qc.sweep(str(tmp_path)) + + assert count == 6 + assert freed >= 10 + 5 + 20 + 1 + 2 + assert sorted(os.listdir(tmp_path)) == sorted(["6" * 64, "notes.txt", "other"]) + assert (current / "w").exists() + + def test_a_symlink_out_of_the_root_is_skipped_not_fatal(self, tmp_path): + """An entry linked to another disk resolves outside the root, which + validate_path refuses with a SecurityError - that one entry is left + alone and the sweep carries on, rather than the server never + starting.""" + root = tmp_path / "root" + root.mkdir() + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + (root / ("1" * 64)).symlink_to(elsewhere) + stale = root / ("3" * 64) + entry_json(stale, versions={**VERSIONS, "torch": "old"}) + (stale / "w").write_bytes(b"x" * 20) + + count, freed = qc.sweep(str(root)) + + assert count == 1 + assert freed >= 20 + assert (root / ("1" * 64)).is_symlink() + assert elsewhere.exists() + assert qc._entry_sizes(str(root)) == [] + + +# 9. The switch and the directory + + +class TestSwitchAndDirectory: + @pytest.mark.parametrize( + "value,expected", + [("1", True), ("0", False), ("true", False), ("", False), (None, False)], + ) + def test_enabled_only_for_one(self, monkeypatch, value, expected): + if value is not None: + monkeypatch.setenv(qc.ENABLE_ENV_VAR, value) + assert qc.enabled() is expected + + def test_default_root_is_under_hf_home(self, tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HOME", str(tmp_path / "hf")) + assert qc.cache_root() == str(tmp_path / "hf" / "dw-quantized") + + def test_configured_root_is_made_absolute(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv(qc.DIR_ENV_VAR, "rel/cache") + assert qc.cache_root() == str(tmp_path / "rel" / "cache") + + def test_startup_with_the_cache_off_does_nothing(self, tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HOME", str(tmp_path / "hf")) + qc.configure_at_startup() + assert not (tmp_path / "hf").exists() + assert qc.DIR_ENV_VAR not in os.environ + + def test_startup_with_the_cache_off_leaves_a_configured_dir( + self, tmp_path, monkeypatch + ): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv(qc.DIR_ENV_VAR, "rel") + monkeypatch.setenv(qc.ENABLE_ENV_VAR, "0") + qc.configure_at_startup() + assert os.environ[qc.DIR_ENV_VAR] == "rel" + assert not (tmp_path / "rel").exists() + + def test_startup_with_the_cache_on_pins_and_creates(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv(qc.ENABLE_ENV_VAR, "1") + monkeypatch.setenv(qc.DIR_ENV_VAR, "rel") + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + stale = tmp_path / "rel" / ("8" * 64) + entry_json(stale, versions={"torch": "old"}) + qc.configure_at_startup() + assert os.environ[qc.DIR_ENV_VAR] == str(tmp_path / "rel") + assert (tmp_path / "rel").is_dir() + assert not stale.exists() + + def test_startup_with_the_default_root_pins_it(self, tmp_path, monkeypatch): + from huggingface_hub import constants + + monkeypatch.setattr(constants, "HF_HOME", str(tmp_path / "hf")) + monkeypatch.setenv(qc.ENABLE_ENV_VAR, "1") + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + qc.configure_at_startup() + assert os.environ[qc.DIR_ENV_VAR] == str(tmp_path / "hf" / "dw-quantized") + assert (tmp_path / "hf" / "dw-quantized").is_dir() + + def test_unusable_root_does_not_raise(self, tmp_path, monkeypatch): + blocker = tmp_path / "file" + blocker.write_text("x") + monkeypatch.setenv(qc.ENABLE_ENV_VAR, "1") + monkeypatch.setenv(qc.DIR_ENV_VAR, str(blocker / "sub")) + qc.configure_at_startup() + + +# 10. Orchestration + + +class FakePipeline: + """A modular pipeline as the cache sees one.""" + + def __init__(self, specs): + self._component_specs = {spec.name: spec for spec in specs} + for name in self._component_specs: + setattr(self, name, None) + self.load_calls = [] + self.loaded = [] + self.updates = [] + + def load_components(self, **kwargs): + self.load_calls.append(kwargs) + names = kwargs.get("names") + for name in self._component_specs: + if names is not None and name not in names: + continue + if getattr(self, name) is None: + setattr(self, name, nn.Linear(4, 4)) + self.loaded.append(name) + + def update_components(self, **kwargs): + self.updates.append(kwargs) + for name, component in kwargs.items(): + setattr(self, name, component) + + +def make_spec(name, subfolder=None, model="org/name", **fields): + values = { + "name": name, + "default_creation_method": "from_pretrained", + "pretrained_model_name_or_path": model, + "subfolder": subfolder or name, + "revision": None, + "variant": None, + "type_hint": nn.Linear, + "load_id": f"{model}|{subfolder or name}", + } + values.update(fields) + return SimpleNamespace(**values) + + +def results(events): + return [(e["component"], e["result"]) for e in events if e["event"] == qc.EVENT] + + +class TestLoadComponents: + @pytest.fixture + def hub(self, tmp_path): + hub = tmp_path / "hub" + for sub in ("vae", "text_encoder"): + make_snapshot(hub, subfolder=sub) + return hub + + @pytest.fixture(autouse=True) + def _fakes(self, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + monkeypatch.setattr(sdnq, "save_sdnq_model", fake_save) + + def arguments(self, hub, **extra): + return { + "quantization_config": SDNQConfig(weights_dtype="int8"), + "torch_dtype": torch.bfloat16, + "cache_dir": str(hub), + **extra, + } + + def pipeline(self): + return FakePipeline([make_spec("vae"), make_spec("text_encoder")]) + + def test_cache_off_is_a_plain_passthrough(self, hub, events, tmp_path): + pipeline = self.pipeline() + arguments = self.arguments(hub) + qc.load_components(pipeline, arguments) + assert pipeline.load_calls == [arguments] + assert pipeline.loaded == ["vae", "text_encoder"] + assert events == [] + assert not (tmp_path / "cache").exists() + + def test_first_load_misses_then_saves(self, hub, events, cache_on): + pipeline = self.pipeline() + qc.load_components(pipeline, self.arguments(hub)) + assert results(events) == [ + ("vae", "miss"), + ("vae", "saved"), + ("text_encoder", "miss"), + ("text_encoder", "saved"), + ] + for event in events: + assert {"component", "result", "seconds", "bytes"} <= set(event) + saved = [e for e in events if e["result"] == "saved"] + assert all(e["bytes"] == 100 for e in saved) + entries = [n for n in os.listdir(cache_on) if qc._KEY_PATTERN.match(n)] + assert len(entries) == 2 + assert pipeline.updates == [] + + def test_second_load_hits(self, hub, events, cache_on, monkeypatch): + qc.load_components(self.pipeline(), self.arguments(hub)) + events.clear() + + cached = nn.Linear(4, 4) + loader = MagicMock(side_effect=lambda *a, **k: nn.Linear(4, 4)) + monkeypatch.setattr(sdnq, "load_sdnq_model", loader) + pipeline = self.pipeline() + qc.load_components(pipeline, self.arguments(hub)) + + assert results(events) == [("vae", "hit"), ("text_encoder", "hit")] + assert all(e["bytes"] == qc.component_nbytes(cached) for e in events) + assert len(pipeline.updates) == 1 + assert set(pipeline.updates[0]) == {"vae", "text_encoder"} + for name, component in pipeline.updates[0].items(): + spec = pipeline._component_specs[name] + assert component._diffusers_load_id == spec.load_id + assert pipeline.loaded == [] + assert loader.call_count == 2 + + def test_named_components_have_the_hit_removed( + self, hub, events, cache_on, monkeypatch + ): + qc.load_components(self.pipeline(), self.arguments(hub)) + events.clear() + monkeypatch.setattr(sdnq, "load_sdnq_model", lambda *a, **k: nn.Linear(4, 4)) + + pipeline = self.pipeline() + qc.load_components(pipeline, self.arguments(hub, names=["vae"])) + + assert results(events) == [("vae", "hit")] + assert pipeline.load_calls[0]["names"] == [] + assert pipeline.loaded == [] + + def test_a_string_name_is_accepted(self, hub, events, cache_on, monkeypatch): + qc.load_components(self.pipeline(), self.arguments(hub)) + events.clear() + monkeypatch.setattr(sdnq, "load_sdnq_model", lambda *a, **k: nn.Linear(4, 4)) + pipeline = self.pipeline() + qc.load_components(pipeline, self.arguments(hub, names="vae")) + assert pipeline.load_calls[0]["names"] == [] + + def test_a_failing_entry_falls_back_to_a_fresh_load( + self, hub, events, cache_on, monkeypatch + ): + qc.load_components(self.pipeline(), self.arguments(hub)) + events.clear() + monkeypatch.setattr( + sdnq, "load_sdnq_model", MagicMock(side_effect=OSError("bad")) + ) + pipeline = self.pipeline() + qc.load_components(pipeline, self.arguments(hub)) + assert pipeline.loaded == ["vae", "text_encoder"] + assert {r for r in results(events) if r[1] == "miss"} == { + ("vae", "miss"), + ("text_encoder", "miss"), + } + misses = [e for e in events if e["result"] == "miss"] + assert {e["reason"] for e in misses} == {"load_failed: OSError: bad"} + # the bad entry went, so the save is a fresh one, not a lost race + saved = [e for e in events if e["result"] == "saved"] + assert saved and all("reason" not in e for e in saved) + + def test_a_miss_names_why_nothing_was_read(self, tmp_path, events, cache_on): + hub = tmp_path / "hub" + make_snapshot(hub, subfolder="vae") + make_snapshot(hub, subfolder="text_encoder", files=()) + qc.load_components(self.pipeline(), self.arguments(hub)) + reasons = { + e["component"]: e.get("reason") for e in events if e["result"] == "miss" + } + assert reasons == {"vae": "no_entry", "text_encoder": "incomplete_snapshot"} + assert ("text_encoder", "skipped") in results(events) + + def test_a_repo_linked_off_the_hub_hits_on_the_second_load( + self, tmp_path, events, cache_on, monkeypatch + ): + # the bounce: run 1 saved under the recorded commit, and run 2's + # lookup, refused by the containment check, never read it + other = tmp_path / "other" + for sub in ("vae", "text_encoder"): + make_snapshot(other, subfolder=sub) + hub = tmp_path / "hub" + hub.mkdir() + name = repo_folder_name(repo_id="org/name", repo_type="model") + os.symlink(other / name, hub / name) + + class Recording(FakePipeline): + def load_components(self, names=None, **kwargs): + super().load_components(names=names, **kwargs) + for name in self.loaded: + getattr(self, name).config = SimpleNamespace(_commit_hash=COMMIT) + + specs = [make_spec("vae"), make_spec("text_encoder")] + qc.load_components(Recording(specs), self.arguments(hub)) + assert [r for r in results(events) if r[1] == "saved"] == [ + ("vae", "saved"), + ("text_encoder", "saved"), + ] + events.clear() + monkeypatch.setattr(sdnq, "load_sdnq_model", lambda *a, **k: nn.Linear(4, 4)) + qc.load_components(Recording(specs), self.arguments(hub)) + assert results(events) == [("vae", "hit"), ("text_encoder", "hit")] + + def test_a_different_config_is_a_different_entry( + self, hub, events, cache_on, monkeypatch + ): + qc.load_components(self.pipeline(), self.arguments(hub)) + events.clear() + arguments = self.arguments(hub, quantization_config=SDNQConfig("int4")) + qc.load_components(self.pipeline(), arguments) + assert ("vae", "hit") not in results(events) + assert ("vae", "saved") in results(events) + + def test_per_component_configs_are_read(self, hub, events, cache_on): + arguments = self.arguments( + hub, + quantization_config={ + "default": None, + "vae": SDNQConfig(weights_dtype="int8"), + }, + ) + qc.load_components(self.pipeline(), arguments) + assert {name for name, _ in results(events)} == {"vae"} + + def test_a_non_sdnq_config_is_left_alone(self, hub, events, cache_on): + pipeline = self.pipeline() + arguments = self.arguments(hub, quantization_config=object()) + qc.load_components(pipeline, arguments) + assert events == [] + assert pipeline.loaded == ["vae", "text_encoder"] + assert not cache_on.exists() + + def test_no_quantization_config_is_left_alone(self, hub, events, cache_on): + pipeline = self.pipeline() + qc.load_components(pipeline, {"torch_dtype": torch.bfloat16}) + assert events == [] and pipeline.loaded == ["vae", "text_encoder"] + + def test_variant_is_unsupported(self, hub, events, cache_on): + pipeline = FakePipeline([make_spec("vae", variant="fp16")]) + qc.load_components(pipeline, self.arguments(hub)) + assert results(events) == [("vae", "skipped")] + assert events[0]["reason"] == qc.UNSUPPORTED + assert pipeline.loaded == ["vae"] + + def test_a_variant_keyword_is_unsupported(self, hub, events, cache_on): + pipeline = FakePipeline([make_spec("vae")]) + qc.load_components(pipeline, self.arguments(hub, variant="fp16")) + assert events[0]["reason"] == qc.UNSUPPORTED + + def test_local_path_is_skipped(self, hub, events, cache_on, tmp_path): + local = tmp_path / "local-model" + local.mkdir() + pipeline = FakePipeline([make_spec("vae", model=str(local), subfolder="vae")]) + qc.load_components(pipeline, self.arguments(hub)) + assert results(events) == [("vae", "skipped")] + assert events[0]["reason"] == qc.LOCAL_PATH + assert pipeline.loaded == ["vae"] + + def test_incomplete_snapshot_loads_but_is_not_saved( + self, events, cache_on, tmp_path + ): + pipeline = FakePipeline([make_spec("vae")]) + qc.load_components(pipeline, self.arguments(tmp_path / "empty-hub")) + assert results(events) == [("vae", "miss"), ("vae", "skipped")] + assert events[1]["reason"] == qc.UNPINNED_SNAPSHOT + assert pipeline.loaded == ["vae"] + + def test_components_already_loaded_are_not_candidates(self, hub, events, cache_on): + pipeline = self.pipeline() + pipeline.vae = nn.Linear(1, 1) + qc.load_components(pipeline, self.arguments(hub)) + assert {name for name, _ in results(events)} == {"text_encoder"} + + def test_unreadable_specs_fall_through_to_the_plain_load( + self, hub, events, cache_on + ): + pipeline = self.pipeline() + pipeline._component_specs = None + pipeline.load_components = MagicMock() + qc.load_components(pipeline, self.arguments(hub)) + pipeline.load_components.assert_called_once() + assert events == [] + + def test_a_load_that_writes_into_its_config_still_hits_next_time( + self, hub, events, cache_on, monkeypatch + ): + # sdnq appends the modules it left unquantized to the config's own + # modules_to_not_convert as it loads, and which ones can vary by + # card: the key must be the config as the workflow gave it + class WritingPipeline(FakePipeline): + def __init__(self, specs, card): + super().__init__(specs) + self.card = card + + def load_components(self, **kwargs): + config = kwargs["quantization_config"] + config.modules_to_not_convert.append(f"skipped_on_{self.card}") + super().load_components(**kwargs) + + def load(card): + pipeline = WritingPipeline([make_spec("vae")], card) + qc.load_components(pipeline, self.arguments(hub)) + return pipeline + + monkeypatch.setattr(sdnq, "load_sdnq_model", lambda *a, **k: nn.Linear(4, 4)) + load("cuda:0") + assert results(events) == [("vae", "miss"), ("vae", "saved")] + events.clear() + + pipeline = load("cuda:1") + assert results(events) == [("vae", "hit")] + assert pipeline.loaded == [] + entries = [n for n in os.listdir(cache_on) if qc._KEY_PATTERN.match(n)] + assert len(entries) == 1 + + def test_a_failing_save_never_fails_the_load( + self, hub, events, cache_on, monkeypatch + ): + def boom(component, path): + raise OSError(errno.ENOSPC, "full") + + monkeypatch.setattr(sdnq, "save_sdnq_model", boom) + pipeline = FakePipeline([make_spec("vae")]) + qc.load_components(pipeline, self.arguments(hub)) + assert pipeline.loaded == ["vae"] + assert events[-1]["reason"] == "save_failed: ENOSPC" + + +# 11. A real round trip + + +class TestRealRoundTrip: + def test_save_entry_then_load_entry_reproduces_the_forward_pass( + self, tmp_path, monkeypatch + ): + from diffusers.models import UNet2DModel + + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + source = tmp_path / "source" + UNet2DModel( + sample_size=8, + in_channels=3, + out_channels=3, + block_out_channels=(32, 64), + layers_per_block=1, + down_block_types=("DownBlock2D", "AttnDownBlock2D"), + up_block_types=("AttnUpBlock2D", "UpBlock2D"), + norm_num_groups=8, + attention_head_dim=8, + ).save_pretrained(source) + config = SDNQConfig( + weights_dtype="int8", + quant_conv=True, + minimum_allowed_numel=256, + minimum_allowed_channel_size=8, + ) + inputs = qc.key_inputs( + "org/tiny", + None, + COMMIT, + qc.quantizer_arguments(config), + torch.float32, + ) + arguments = qc.config_snapshot(config) + quantized = UNet2DModel.from_pretrained( + source, quantization_config=config, torch_dtype=torch.float32 + ) + # the quantizer really ran: something is stored as int8 + assert any(p.dtype == torch.int8 for p in quantized.state_dict().values()) + # and it wrote into the config it was given, which the key read + # before the load does not see + assert config.modules_to_not_convert + assert inputs["quantizer"]["arguments"]["modules_to_not_convert"] == [] + assert qc.quantizer_arguments(config) != inputs["quantizer"] + + root = str(tmp_path / "cache") + key = qc.cache_key(inputs) + result, reason, nbytes = qc.save_entry(root, key, quantized, inputs) + assert (result, reason) == ("saved", None) and nbytes > 0 + assert qc.read_entry(os.path.join(root, key))["model_name"] == "org/tiny" + + loaded = qc.load_entry(root, key, UNet2DModel, torch.float32, "cpu", arguments) + assert loaded is not None + + sample = torch.randn(1, 3, 8, 8, generator=torch.Generator().manual_seed(0)) + step = torch.tensor([1]) + with torch.no_grad(): + expected = quantized(sample, step).sample + actual = loaded(sample, step).sample + assert torch.allclose(expected, actual, atol=1e-4) + assert qc.read_entry(os.path.join(root, key))["last_hit"] is not None + + +# 12. The save point and the server start + + +class TestCallSites: + def configuration(self, component): + component_type = MagicMock() + component_type.__name__ = "MockModularPipeline" + component_type.from_pretrained.return_value = component + return { + "component_type": component_type, + "load_components": {"torch_dtype": torch.bfloat16}, + "components": {"vae": {"truncate_layers": {"layers": 1}}}, + } + + def test_load_component_goes_through_the_cache_and_before_pruning(self): + from dw.pipeline_processors import components + + component = MagicMock() + component.vae = MagicMock() + order = [] + configuration = self.configuration(component) + + with ( + patch.object( + components.quantized_cache, + "load_components", + side_effect=lambda *a: order.append("cache"), + ) as cache, + patch.object( + components, + "truncate_module_lists", + side_effect=lambda *a: order.append("truncate"), + ), + patch.object( + components, + "replace_modules_with_identity", + side_effect=lambda *a: order.append("identity"), + ), + ): + loaded = components.load_component( + "pipeline", configuration, {"model_name": "m"}, "cpu" + ) + components.configure_components(loaded, configuration, "cpu") + + cache.assert_called_once_with(component, {"torch_dtype": torch.bfloat16}) + component.load_components.assert_not_called() + assert order == ["cache", "truncate", "identity"] + + def test_the_cache_off_still_reaches_load_components(self): + from dw.pipeline_processors import components + + component = MagicMock() + components.load_component( + "pipeline", self.configuration(component), {"model_name": "m"}, "cpu" + ) + component.load_components.assert_called_once_with(torch_dtype=torch.bfloat16) + + def test_serve_run_configures_the_cache_before_the_app_exists(self): + import dw.server.app + import dw.serve as serve + + order = [] + config = SimpleNamespace( + host="127.0.0.1", + port=1, + token=None, + log_level="INFO", + workflow_dir=".", + output_dir=None, + prompt_dir=None, + asset_dir=None, + examples_dirs=[], + workspace_root=None, + mcp=False, + devices=None, + ) + with ( + patch("dw.startup"), + patch.object( + qc, "configure_at_startup", side_effect=lambda: order.append("cache") + ), + patch.object( + dw.server.app, + "create_app", + side_effect=lambda **kw: order.append("app"), + ), + patch.object(dw.server.app, "default_ui_dir", return_value=None), + patch("uvicorn.run", side_effect=lambda *a, **k: order.append("run")), + ): + serve.run(config) + + assert order == ["cache", "app", "run"] + + +# 13. The per-component path (#827) + + +class PlacedModule(nn.Linear): + """A loaded component that records where it was put.""" + + def __init__(self): + super().__init__(4, 4) + self.placed_on = [] + + def to(self, *args, **kwargs): + self.placed_on.append(args[0] if args else kwargs.get("device")) + return self + + +class TestFromPretrained: + @pytest.fixture + def hub(self, tmp_path): + hub = tmp_path / "hub" + make_snapshot(hub, subfolder="transformer") + return hub + + @pytest.fixture(autouse=True) + def _fakes(self, monkeypatch): + monkeypatch.setattr(qc.shutil, "disk_usage", disk(1000, 900)) + monkeypatch.setattr(sdnq, "save_sdnq_model", fake_save) + + def arguments(self, hub, **extra): + return { + "subfolder": "transformer", + "torch_dtype": torch.bfloat16, + "cache_dir": str(hub), + "quantization_config": SDNQConfig( + weights_dtype="uint4", + quantization_device="cuda", + return_device="cuda", + use_quantized_matmul=True, + ), + **extra, + } + + def load(self, hub, loader, **extra): + return qc.from_pretrained( + "transformer", nn.Linear, "org/name", self.arguments(hub, **extra), loader + ) + + def test_cache_off_is_a_plain_passthrough(self, hub, events, tmp_path): + fresh = nn.Linear(4, 4) + loader = MagicMock(return_value=fresh) + assert self.load(hub, loader) is fresh + loader.assert_called_once_with() + assert events == [] + assert not (tmp_path / "cache").exists() + + def test_a_load_no_sdnq_config_quantizes_is_untouched(self, hub, events, cache_on): + fresh = nn.Linear(4, 4) + loader = MagicMock(return_value=fresh) + arguments = self.arguments(hub) + del arguments["quantization_config"] + assert ( + qc.from_pretrained("vae", nn.Linear, "org/name", arguments, loader) is fresh + ) + assert events == [] + + def test_a_miss_loads_and_saves(self, hub, events, cache_on): + fresh = nn.Linear(4, 4) + loader = MagicMock(return_value=fresh) + assert self.load(hub, loader) is fresh + loader.assert_called_once_with() + assert results(events) == [("transformer", "miss"), ("transformer", "saved")] + assert events[0]["reason"] == qc.NO_ENTRY + assert events[0]["bytes"] == qc.component_nbytes(fresh) + assert events[1]["bytes"] == 100 + (entry,) = [n for n in os.listdir(cache_on) if qc._KEY_PATTERN.match(n)] + stored = qc.read_entry(str(cache_on / entry)) + assert stored["model_name"] == "org/name" + assert stored["subfolder"] == "transformer" + assert stored["commit"] == COMMIT + assert "return_device" not in stored["quantizer"]["arguments"] + + def test_a_hit_skips_the_load_and_is_placed_on_return_device( + self, hub, events, cache_on, monkeypatch + ): + self.load(hub, MagicMock(return_value=nn.Linear(4, 4))) + events.clear() + + cached = PlacedModule() + sdnq_load = MagicMock(return_value=cached) + monkeypatch.setattr(sdnq, "load_sdnq_model", sdnq_load) + loader = MagicMock() + assert self.load(hub, loader) is cached + + loader.assert_not_called() + assert results(events) == [("transformer", "hit")] + assert events[0]["bytes"] == qc.component_nbytes(cached) + kwargs = sdnq_load.call_args.kwargs + assert kwargs["model_cls"] is nn.Linear + assert kwargs["dtype"] is torch.bfloat16 + assert kwargs["device"] == "cuda" + assert kwargs["use_quantized_matmul"] is True + assert cached.placed_on == ["cuda"] + + def test_a_different_weights_dtype_misses(self, hub, events, cache_on): + self.load(hub, MagicMock(return_value=nn.Linear(4, 4))) + events.clear() + arguments = self.arguments(hub) + arguments["quantization_config"] = SDNQConfig(weights_dtype="int8") + qc.from_pretrained( + "transformer", nn.Linear, "org/name", arguments, lambda: nn.Linear(4, 4) + ) + assert results(events) == [("transformer", "miss"), ("transformer", "saved")] + + def test_a_corrupt_entry_falls_back_to_a_fresh_load( + self, hub, events, cache_on, monkeypatch + ): + self.load(hub, MagicMock(return_value=nn.Linear(4, 4))) + events.clear() + monkeypatch.setattr( + sdnq, "load_sdnq_model", MagicMock(side_effect=ValueError("corrupt")) + ) + fresh = nn.Linear(4, 4) + assert self.load(hub, MagicMock(return_value=fresh)) is fresh + assert results(events) == [("transformer", "miss"), ("transformer", "saved")] + assert events[0]["reason"] == f"{qc.LOAD_FAILED}: ValueError: corrupt" + + def test_a_resharded_refs_main_with_missing_shards_gives_no_lookup( + self, tmp_path, events, cache_on, monkeypatch + ): + # The LTX-2.5 re-shard: refs/main names a commit whose index lists + # shards that are not on disk + hub = tmp_path / "hub" + folder = make_snapshot( + hub, subfolder="transformer", files=("m-1-of-2.safetensors",) + ) + index = { + "weight_map": {"a": "m-1-of-2.safetensors", "b": "m-2-of-2.safetensors"} + } + (folder / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps(index) + ) + sdnq_load = MagicMock() + monkeypatch.setattr(sdnq, "load_sdnq_model", sdnq_load) + lookup = MagicMock(wraps=qc.load_entry) + monkeypatch.setattr(qc, "load_entry", lookup) + fresh = nn.Linear(4, 4) + + assert self.load(hub, MagicMock(return_value=fresh)) is fresh + + lookup.assert_not_called() + sdnq_load.assert_not_called() + assert results(events) == [ + ("transformer", "miss"), + ("transformer", "skipped"), + ] + assert events[0]["reason"] == qc.INCOMPLETE_SNAPSHOT + assert events[1]["reason"] == qc.UNPINNED_SNAPSHOT + assert not cache_on.exists() or os.listdir(cache_on) == [] + + def test_a_local_path_is_skipped(self, tmp_path, events, cache_on): + fresh = nn.Linear(4, 4) + result = qc.from_pretrained( + "transformer", + nn.Linear, + str(tmp_path), + self.arguments(tmp_path), + lambda: fresh, + ) + assert result is fresh + assert results(events) == [("transformer", "skipped")] + assert events[0]["reason"] == qc.LOCAL_PATH + + def test_a_keyword_the_key_does_not_cover_is_skipped(self, hub, events, cache_on): + fresh = nn.Linear(4, 4) + loader = MagicMock(return_value=fresh) + assert self.load(hub, loader, ignore_mismatched_sizes=True) is fresh + loader.assert_called_once_with() + assert results(events) == [("transformer", "skipped")] + assert events[0]["reason"] == qc.UNSUPPORTED + + def test_a_variant_is_skipped(self, hub, events, cache_on): + self.load(hub, lambda: nn.Linear(4, 4), variant="fp16") + assert results(events) == [("transformer", "skipped")] + assert events[0]["reason"] == qc.UNSUPPORTED + + +class TestPerComponentCallSite: + """load_component, the path load_optional_component takes, end to end.""" + + def configuration(self, component_type): + return {"component_type": component_type, "preserve_device_placement": True} + + def test_download_watch_is_used_on_a_miss_and_skipped_on_a_hit(self, cache_on): + from dw.pipeline_processors import components + + component_type = MagicMock() + component_type.__name__ = "MockTransformer" + component_type.from_pretrained.return_value = nn.Linear(4, 4) + cached = nn.Linear(4, 4) + lookups = iter([(None, qc.NO_ENTRY), (cached, None)]) + + with ( + patch.object(components.download_watch, "watch") as watch, + patch.object(qc, "_lookup", side_effect=lambda *a: next(lookups)), + patch.object(qc, "_save"), + patch.object(qc, "is_local_path", return_value=False), + ): + for _ in range(2): + components.load_component( + "transformer", + self.configuration(component_type), + { + "model_name": "org/name", + "subfolder": "transformer", + "quantization_config": SDNQConfig(weights_dtype="uint4"), + }, + "cpu", + ) + + # the miss loaded under the watch, the hit reached neither + watch.assert_called_once() + assert watch.call_args.args == ("org/name",) + component_type.from_pretrained.assert_called_once() + + def test_the_cache_off_loads_under_the_watch(self): + from dw.pipeline_processors import components + + component_type = MagicMock() + component_type.__name__ = "MockTransformer" + with patch.object(components.download_watch, "watch") as watch: + components.load_component( + "transformer", + self.configuration(component_type), + { + "model_name": "org/name", + "quantization_config": SDNQConfig(weights_dtype="uint4"), + }, + "cpu", + ) + watch.assert_called_once() + component_type.from_pretrained.assert_called_once_with( + "org/name", quantization_config=ANY_CONFIG + ) + + def test_a_real_component_misses_saves_then_hits_with_the_same_output( + self, tmp_path, cache_on, events + ): + # A real diffusers model in a real hub-cache layout, quantized by sdnq + # on the first load and read back from the cache on the second + from diffusers.models import UNet2DModel + + from dw.pipeline_processors import components + + hub = tmp_path / "hub" + folder = make_snapshot(hub, subfolder="unet", files=(), config=False) + UNet2DModel( + sample_size=8, + in_channels=3, + out_channels=3, + block_out_channels=(32, 64), + layers_per_block=1, + down_block_types=("DownBlock2D", "AttnDownBlock2D"), + up_block_types=("AttnUpBlock2D", "UpBlock2D"), + norm_num_groups=8, + attention_head_dim=8, + ).save_pretrained(folder) + + def load(): + return components.load_component( + "unet", + {"component_type": UNet2DModel}, + { + "model_name": "org/name", + "subfolder": "unet", + "cache_dir": str(hub), + "local_files_only": True, + "torch_dtype": torch.float32, + "quantization_config": SDNQConfig( + weights_dtype="int8", + quant_conv=True, + minimum_allowed_numel=256, + minimum_allowed_channel_size=8, + return_device="cpu", + ), + }, + "cpu", + ) + + first = load() + assert results(events) == [("unet", "miss"), ("unet", "saved")] + events.clear() + with patch.object(UNet2DModel, "from_pretrained") as fresh: + second = load() + fresh.assert_not_called() + assert results(events) == [("unet", "hit")] + assert any(p.dtype == torch.int8 for p in second.state_dict().values()) + + sample = torch.randn(1, 3, 8, 8, generator=torch.Generator().manual_seed(0)) + step = torch.tensor([1]) + with torch.no_grad(): + expected = first(sample, step).sample + actual = second(sample, step).sample + assert torch.allclose(expected, actual, atol=1e-4) + + +class _AnyConfig: + def __eq__(self, other): + return qc.is_sdnq_config(other) + + +ANY_CONFIG = _AnyConfig() + + +class TestTransformersRoundTrip: + """An entry is saved under the model's own tensor names, so a transformers + model with a weight conversion mapping loads back (#827: LTX-2.5's Gemma 4 + text encoder saved in its checkpoint layout and failed every lookup).""" + + @staticmethod + def gemma4(): + transformers = pytest.importorskip("transformers") + if not hasattr(transformers, "Gemma4UnifiedForConditionalGeneration"): + pytest.skip("this transformers has no Gemma 4") + config = transformers.Gemma4UnifiedConfig( + text_config=dict( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + global_head_dim=32, + vocab_size=512, + hidden_size_per_layer_input=16, + vocab_size_per_layer_input=512, + num_kv_shared_layers=0, + ), + vision_config=dict( + mm_embed_dim=64, mm_posemb_size=32, output_proj_dims=64, patch_size=4 + ), + audio_config=dict(audio_embed_dim=32), + ) + return transformers.Gemma4UnifiedForConditionalGeneration, config + + def test_a_model_with_a_conversion_mapping_saves_and_loads_back(self, tmp_path): + model_cls, config = self.gemma4() + torch.manual_seed(0) + source = tmp_path / "source" + model_cls(config).to(torch.bfloat16).save_pretrained(source) + arguments = { + "weights_dtype": "int8", + "dequantize_fp32": False, + "modules_to_not_convert": ["embed_vision", "embed_audio", "lm_head"], + "minimum_allowed_numel": 256, + "minimum_allowed_channel_size": 8, + } + quantized = model_cls.from_pretrained( + source, + quantization_config=SDNQConfig(**arguments), + torch_dtype=torch.bfloat16, + ) + root = tmp_path / "cache" + result, reason, _ = qc.save_entry(str(root), KEY, quantized, {}) + assert (result, reason) == ("saved", None) + # the save left the model's own save_pretrained in place + assert "save_pretrained" not in vars(quantized) + + loaded = qc.load_entry( + str(root), KEY, model_cls, torch.bfloat16, "cpu", arguments + ) + expected = quantized.state_dict() + actual = loaded.state_dict() + assert set(actual) == set(expected) + assert any(t.dtype == torch.int8 for t in actual.values()) + for name, tensor in expected.items(): + assert torch.equal(actual[name], tensor), name + + def test_a_load_failure_reason_carries_the_first_line_of_the_message(self): + error = RuntimeError("Error(s) in loading state_dict:\n\tMissing key(s)") + assert qc._error_detail(error) == ( + "RuntimeError: Error(s) in loading state_dict:" + ) + long = qc._error_detail(ValueError("x" * 500)) + assert long.startswith("ValueError: x") and long.endswith("...") + assert len(long) == len("ValueError: ") + qc._REASON_DETAIL_CHARS + assert qc._error_detail(OSError()) == "OSError" diff --git a/tests/test_server.py b/tests/test_server.py index a494c22b..fb4c1080 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -2397,6 +2397,36 @@ def test_gallery_metadata_reports_a_level_envelope_on_request(server, tmp_path): ) +def test_gallery_metadata_envelope_takes_an_interval_and_refuses_a_bad_one( + server, tmp_path +): + from tests.test_media_info import write_wav + + with server(success_script) as client: + write_wav(tmp_path / "outputs" / "score-gen.0-0.0.wav", seconds=3.0) + url = "/api/gallery/score-gen.0-0.0.wav/metadata" + + fine = client.get(url, params={"envelope": "true", "interval_seconds": 0.1}) + assert fine.json()["media"]["envelope"]["interval_seconds"] == pytest.approx( + 0.1 + ) + assert len(fine.json()["media"]["envelope"]["rms_dbfs"]) == 30 + + too_fine = client.get( + url, params={"envelope": "true", "interval_seconds": 0.01} + ) + assert too_fine.status_code == 422 + assert "at least 0.05" in too_fine.json()["detail"] + + # pydantic reads "inf" as a float; it must be a 422, not an + # OverflowError 500 in the binning + infinite = client.get( + url, params={"envelope": "true", "interval_seconds": "inf"} + ) + assert infinite.status_code == 422 + assert "at most 3600" in infinite.json()["detail"] + + def test_gallery_metadata_survives_a_damaged_track(server, tmp_path): """A track that opens fine but fails partway through decode (damage past the header) must not 500 the metadata route - it should fall back diff --git a/tests/test_server_assess.py b/tests/test_server_assess.py index 9e1762f5..ffef5c3d 100644 --- a/tests/test_server_assess.py +++ b/tests/test_server_assess.py @@ -273,7 +273,7 @@ def test_an_overrunning_shot_is_reported_once_not_per_probe(server, tmp_path): overrun = [f for f in body["findings"] if f["rule"] == "shot_span_overrun"] assert len(overrun) == 1 - assert overrun[0]["at"] == {"shot": "s2"} + assert overrun[0]["at"]["shot"] == "s2" def test_a_path_outside_the_outputs_is_refused(server): diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index 9cdd5593..e0ef2220 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -1127,3 +1127,20 @@ def test_pipeline_definition_key_moves_with_runtime_settings(): c = _definition(arguments={"prompt": "other"}, seed=7) assert pipeline_definition_key(a) != pipeline_definition_key(b) assert pipeline_definition_key(a) == pipeline_definition_key(c) + + +@pytest.mark.parametrize("name", ["dialogue-short", "music-video"]) +def test_minimax_join_passes_match_levels_dbfs_through(name): + """#836: the join's RMS target is an argument, not the fixed -20 that + clip-holds every H3 shot.""" + path = pathlib.Path("workflows/templates/minimax") / f"{name}.json" + definition = json.loads(path.read_text()) + assert definition["variables"]["match_levels_dbfs"] is None + joins = [ + s["task"]["arguments"] + for s in definition["steps"] + if "match_levels" in s.get("task", {}).get("arguments", {}) + ] + assert joins + for args in joins: + assert args["match_levels_dbfs"] == "variable:match_levels_dbfs" diff --git a/tests/test_transcribe_audio.py b/tests/test_transcribe_audio.py index dd3ae0c0..21dc0959 100644 --- a/tests/test_transcribe_audio.py +++ b/tests/test_transcribe_audio.py @@ -5,6 +5,8 @@ import numpy +from dw.tasks import audio_transcription + from dw.tasks.audio_transcription import ( transcribe_audio, WHISPER_DEFAULT_MODEL, @@ -24,6 +26,14 @@ def _waveform(self, channels=1, seconds=1.0, sample_rate=16000): tone = numpy.zeros((channels, samples), dtype=numpy.float32) return tone, sample_rate + def _voiced(self, seconds=1.0, sample_rate=16000): + """A waveform loud throughout, so no chunk of it is marked silent.""" + tone, rate = self._waveform(seconds=seconds, sample_rate=sample_rate) + tone[0] = 0.3 * numpy.sin( + 2 * numpy.pi * 220 * numpy.arange(tone.shape[1]) / rate + ) + return tone, rate + def _mock_pipe( self, mock_pipeline, text="hello world", kind="seq2seq_whisper", chunks=None ): @@ -182,7 +192,7 @@ def test_timestamps_word_requests_word_level_and_returns_chunks( {"text": " world", "timestamp": (0.5, 1.0)}, ], ) - waveform, rate = self._waveform() + waveform, rate = self._voiced() result = transcribe_audio( waveform, device="cpu", sample_rate=rate, timestamps="word" @@ -280,7 +290,7 @@ def test_timestamps_segment_requests_true_and_returns_chunks(self, mock_pipeline "hello world", chunks=[{"text": "hello world", "timestamp": (0.0, 1.0)}], ) - waveform, rate = self._waveform() + waveform, rate = self._voiced() result = transcribe_audio( waveform, device="cpu", sample_rate=rate, timestamps="segment" @@ -304,7 +314,7 @@ def test_an_open_ended_last_chunk_ends_at_the_clips_duration(self, mock_pipeline {"text": " cut off", "timestamp": (2.5, None)}, ], ) - waveform, rate = self._waveform(seconds=4.0) + waveform, rate = self._voiced(seconds=4.0) result = transcribe_audio( waveform, device="cpu", sample_rate=rate, timestamps="segment" @@ -318,6 +328,108 @@ def test_an_open_ended_last_chunk_ends_at_the_clips_duration(self, mock_pipeline ], ) + @patch.object(audio_transcription, "hf_pipeline") + def test_a_word_over_silence_is_marked_silent_not_dropped(self, mock_pipeline): + # #822: end credits Whisper invents over a fade. 3 s voiced, 3 s silent + waveform, rate = self._voiced(seconds=6.0) + waveform[0, 48000:] = 0.0 + self._mock_pipe( + mock_pipeline, + "la credits", + chunks=[ + {"text": " la", "timestamp": (0.5, 1.0)}, + {"text": " credits", "timestamp": (4.0, 5.0)}, + ], + ) + + result = transcribe_audio( + waveform, device="cpu", sample_rate=rate, timestamps="word" + ) + + la, credits = result["chunks"] + self.assertNotIn("suspect", la) + self.assertEqual(credits["suspect"], "silent") + self.assertEqual(credits["text"], "credits") + + @patch.object(audio_transcription, "hf_pipeline") + def test_a_voiced_word_over_two_seconds_is_marked_long(self, mock_pipeline): + waveform, rate = self._voiced(seconds=6.0) + self._mock_pipe( + mock_pipeline, + "Union fine", + chunks=[ + {"text": " Union", "timestamp": (0.0, 4.0)}, + {"text": " fine", "timestamp": (4.0, 5.0)}, + ], + ) + + result = transcribe_audio( + waveform, device="cpu", sample_rate=rate, timestamps="word" + ) + + self.assertEqual([c.get("suspect") for c in result["chunks"]], ["long", None]) + + @patch.object(audio_transcription, "hf_pipeline") + def test_a_repetition_loop_is_marked_however_loud(self, mock_pipeline): + # #822 bounce: Whisper's "I'm sorry, I'm sorry, ..." over a fade on + # the vocal stem, loud enough to pass the floor + waveform, rate = self._voiced(seconds=8.0) + waveform[0, int(5.2 * rate) :] = 0.0 # no tail for a resumed pass + loop = [ + {"text": f" {word}", "timestamp": (2.0 + i * 0.2, 2.2 + i * 0.2)} + for i, word in enumerate(["I'm", "sorry,"] * 8) + ] + self._mock_pipe( + mock_pipeline, + "Union I'm sorry", + chunks=[{"text": " Union", "timestamp": (0.5, 1.0)}] + loop, + ) + + result = transcribe_audio( + waveform, device="cpu", sample_rate=rate, timestamps="word" + ) + + marks = [c.get("suspect") for c in result["chunks"]] + self.assertEqual(marks, [None] + ["repetition"] * 16) + + @patch.object(audio_transcription, "hf_pipeline") + def test_a_chunk_with_no_word_is_marked_no_words(self, mock_pipeline): + # #822 bounce: '—' chunks over the intro, where the vocal stem's + # leakage keeps the span above the floor + waveform, rate = self._voiced(seconds=6.0) + waveform[0, int(3.5 * rate) :] = 0.0 # no tail for a resumed pass + self._mock_pipe( + mock_pipeline, + "— — Union", + chunks=[ + {"text": " —", "timestamp": (0.0, 0.36)}, + {"text": " —", "timestamp": (1.0, 1.0)}, + {"text": " Union", "timestamp": (3.0, 3.5)}, + ], + ) + + result = transcribe_audio( + waveform, device="cpu", sample_rate=rate, timestamps="segment" + ) + + marks = [c.get("suspect") for c in result["chunks"]] + self.assertEqual(marks, ["no_words", "no_words", None]) + + @patch.object(audio_transcription, "hf_pipeline") + def test_a_long_segment_is_not_marked_long(self, mock_pipeline): + waveform, rate = self._voiced(seconds=6.0) + self._mock_pipe( + mock_pipeline, + "a line", + chunks=[{"text": "a line", "timestamp": (0.0, 5.0)}], + ) + + result = transcribe_audio( + waveform, device="cpu", sample_rate=rate, timestamps="segment" + ) + + self.assertNotIn("suspect", result["chunks"][0]) + @patch("dw.tasks.audio_transcription.hf_pipeline") def test_the_duration_is_the_clips_whatever_its_sample_rate(self, mock_pipeline): # The duration comes from the 16 kHz mono array the model heard, so a diff --git a/ui/src/lib/CachesPanel.svelte b/ui/src/lib/CachesPanel.svelte new file mode 100644 index 00000000..aff6cae7 --- /dev/null +++ b/ui/src/lib/CachesPanel.svelte @@ -0,0 +1,364 @@ + + + +
+

Caches

+ {#if cachesError} +

{cachesError}

+ {:else if !caches} +

loading caches…

+ {:else if caches.length === 0} +

No disk caches registered.

+ {:else} + {#if clearError} + + {/if} +
+ + + + + + + + + + + + + + + {#each caches as cache (cache.name)} + {@const pct = usedPct(cache)} + + + + + + + + + + + {#if cache.low_disk} + + + + + {/if} + {#if Array.isArray(cache.entries)} + {@const open = expanded.includes(cache.name)} + + + + + {#if open} + + + + + {/if} + {/if} + {#if !cache.clearable && cache.not_clearable_reason} + + + + + {/if} + {/each} + +
cacherootsizefileslast usedenableddisk
{cache.name}{cache.root}{cache.bytes ? formatBytes(cache.bytes) : '0'}{cache.files} + {cache.last_used ? formatMtime(cache.last_used) : '—'} + {cache.enabled ? 'yes' : 'no'} + {#if pct != null} + + + {gbFromBytes(cache.disk_free ?? 0)} / {gbFromBytes( + cache.disk_total ?? 0, + )} GB free + + {:else} + — + {/if} + + +
+ volume under the {cache.name} + cache's free-space floor - new entries are not saved +
+ +
+
+ + + + + + + + + + + + + + + {#each cache.entries as entry (entry.id)} + + + + + + + + + + + {/each} + +
componentmodelsubfoldersizecreatedlast hitid
+ {entry.unreadable + ? 'unreadable' + : entry.component} + {entry.model || '—'}{entry.subfolder || '—'}{cacheSize(entry.bytes)} + {entry.created + ? formatMtime(entry.created) + : '—'} + + {entry.last_hit != null + ? formatMtime(entry.last_hit) + : '—'} + + {entry.id.slice(0, 12)} + + +
+
+
{cache.not_clearable_reason}
+
+ {/if} +
+ + diff --git a/ui/src/lib/api.ts b/ui/src/lib/api.ts index 9aa74cda..d9fee80d 100644 --- a/ui/src/lib/api.ts +++ b/ui/src/lib/api.ts @@ -4,6 +4,9 @@ import type { ClassList, Deleted, DiffusersStatus, + DiskCacheCleared, + DiskCacheEntryDeleted, + DiskCacheList, EnhancerPresets, GalleryList, GalleryMetadata, @@ -251,6 +254,22 @@ export const api = { }), health: () => request('/api/health'), server: () => request('/api/server'), + /** dw's own disk caches (dw/disk_caches.py) - Inductor's and Triton's + * compiled kernels - with each one's size and whether it can be cleared. */ + listCaches: () => request('/api/caches'), + /** Empty one cache's root. 409 (an `ApiError`) while a job is running or + * queued; 400 when the root is not dw's to clear. */ + clearCache: (name: string) => + request(`/api/caches/${encodeURIComponent(name)}`, { + method: 'DELETE', + }), + /** Delete one entry of a cache that has them (today `quantized`). 404 for + * an unknown entry; 409 (an `ApiError`) while a job or download is active. */ + deleteCacheEntry: (name: string, entry: string) => + request( + `/api/caches/${encodeURIComponent(name)}/entries/${encodeURIComponent(entry)}`, + { method: 'DELETE' }, + ), // Loads the whole gallery in one request, like listWorkflows/listPrompts - // the limit just needs to exceed any real output directory's file count gallery: () => request('/api/gallery?limit=100000'), diff --git a/ui/src/lib/generated/api-schema.ts b/ui/src/lib/generated/api-schema.ts index c24a2d7f..dd681ac6 100644 --- a/ui/src/lib/generated/api-schema.ts +++ b/ui/src/lib/generated/api-schema.ts @@ -124,6 +124,69 @@ export interface paths { patch?: never; trace?: never; }; + "/api/caches": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Caches + * @description dw's own disk caches (dw/disk_caches.py): each one's root, size, and + * whether it can be cleared from here. + */ + get: operations["list_caches_api_caches_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/caches/{name}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Clear Cache + * @description Empty one cache's root, keeping the root and its marker. Refused + * while anything may be writing into it (`_refuse_while_writing`). + */ + delete: operations["clear_cache_api_caches__name__delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/caches/{name}/entries/{entry}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Delete Cache Entry + * @description Delete one entry of a cache that has entries (quantized). The id is + * checked to be 64 hex before any path is built from it. + */ + delete: operations["delete_cache_entry_api_caches__name__entries__entry__delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/api/classes": { parameters: { query?: never; @@ -485,7 +548,11 @@ export interface paths { * is what says *where* in a track something is - whether a shot is * still voiced at its last frame, how deep the hole at a seam goes. * Opt-in: a ten-minute track is 600 numbers, and the default call has - * to stay small. + * to stay small. `interval_seconds` (default 1.0, 0.05 to 3600) is + * the bin width: 0.1 places a cut on a vocal onset. A track that + * would need more than 2400 bins at a finer-than-default width is + * refused (422) with the interval that fits; the default width is + * never refused for length (#823). * * `name` may also be an 'asset:' reference, and then it is the input * asset of that name that is described rather than an output (#127). @@ -1869,6 +1936,110 @@ export interface components { /** Version */ version: string | null; }; + /** DiskCacheCleared */ + DiskCacheCleared: { + /** Cleared */ + cleared: boolean; + /** Files */ + files: number; + /** + * Freed + * @description Bytes freed. + */ + freed: number; + /** Name */ + name: string; + }; + /** DiskCacheEntry */ + DiskCacheEntry: { + /** Bytes */ + bytes: number; + /** Component */ + component: string | null; + /** + * Created + * @description Epoch seconds. + */ + created: number | null; + /** Files */ + files: number; + /** + * Id + * @description The entry's 64-hex key: its directory name. + */ + id: string; + /** + * Last Hit + * @description Epoch seconds of the last load that read it; null if none has. + */ + last_hit: number | null; + /** Model */ + model: string | null; + /** Subfolder */ + subfolder: string | null; + /** + * Unreadable + * @description True when its entry.json is missing or unreadable; the fields below are then null, and it can still be deleted. + */ + unreadable: boolean; + }; + /** DiskCacheEntryDeleted */ + DiskCacheEntryDeleted: { + /** Deleted */ + deleted: boolean; + /** Entry */ + entry: string; + /** + * Freed + * @description Bytes freed. + */ + freed: number; + /** Name */ + name: string; + }; + /** DiskCacheInfo */ + DiskCacheInfo: { + /** Bytes */ + bytes: number; + /** Clearable */ + clearable: boolean; + /** Cost Of Clearing */ + cost_of_clearing: string; + /** Disk Free */ + disk_free: number | null; + /** Disk Total */ + disk_total: number | null; + /** Enabled */ + enabled: boolean; + /** + * Entries + * @description For a cache with entries (quantized). + */ + entries?: components["schemas"]["DiskCacheEntry"][] | null; + /** Files */ + files: number; + /** + * Last Used + * @description The newest file's mtime (epoch seconds); null when empty. + */ + last_used: number | null; + /** + * Low Disk + * @description For a cache with a free-space floor (quantized): the volume is under it, so new entries are not saved. + */ + low_disk?: boolean | null; + /** Name */ + name: string; + /** Not Clearable Reason */ + not_clearable_reason?: string | null; + /** Root */ + root: string; + }; + /** DiskCacheList */ + DiskCacheList: { + /** Caches */ + caches: components["schemas"]["DiskCacheInfo"][]; + }; /** DiskUsage */ DiskUsage: { /** Bytes */ @@ -3492,6 +3663,89 @@ export interface operations { }; }; }; + list_caches_api_caches_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DiskCacheList"]; + }; + }; + }; + }; + clear_cache_api_caches__name__delete: { + parameters: { + query?: never; + header?: never; + path: { + name: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DiskCacheCleared"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_cache_entry_api_caches__name__entries__entry__delete: { + parameters: { + query?: never; + header?: never; + path: { + name: string; + entry: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DiskCacheEntryDeleted"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; classes_api_classes_get: { parameters: { query: { @@ -3902,6 +4156,7 @@ export interface operations { parameters: { query?: { envelope?: boolean; + interval_seconds?: number; workspace?: string | null; }; header?: never; diff --git a/ui/src/lib/generated/openapi.json b/ui/src/lib/generated/openapi.json index f3613bae..7c5b9851 100644 --- a/ui/src/lib/generated/openapi.json +++ b/ui/src/lib/generated/openapi.json @@ -571,6 +571,295 @@ "title": "DiffusersStatus", "type": "object" }, + "DiskCacheCleared": { + "additionalProperties": false, + "properties": { + "cleared": { + "title": "Cleared", + "type": "boolean" + }, + "files": { + "title": "Files", + "type": "integer" + }, + "freed": { + "description": "Bytes freed.", + "title": "Freed", + "type": "integer" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "cleared", + "freed", + "files" + ], + "title": "DiskCacheCleared", + "type": "object" + }, + "DiskCacheEntry": { + "additionalProperties": false, + "properties": { + "bytes": { + "title": "Bytes", + "type": "integer" + }, + "component": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Component" + }, + "created": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Epoch seconds.", + "title": "Created" + }, + "files": { + "title": "Files", + "type": "integer" + }, + "id": { + "description": "The entry's 64-hex key: its directory name.", + "title": "Id", + "type": "string" + }, + "last_hit": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "Epoch seconds of the last load that read it; null if none has.", + "title": "Last Hit" + }, + "model": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model" + }, + "subfolder": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Subfolder" + }, + "unreadable": { + "description": "True when its entry.json is missing or unreadable; the fields below are then null, and it can still be deleted.", + "title": "Unreadable", + "type": "boolean" + } + }, + "required": [ + "id", + "bytes", + "files", + "unreadable", + "model", + "subfolder", + "component", + "created", + "last_hit" + ], + "title": "DiskCacheEntry", + "type": "object" + }, + "DiskCacheEntryDeleted": { + "additionalProperties": false, + "properties": { + "deleted": { + "title": "Deleted", + "type": "boolean" + }, + "entry": { + "title": "Entry", + "type": "string" + }, + "freed": { + "description": "Bytes freed.", + "title": "Freed", + "type": "integer" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name", + "entry", + "deleted", + "freed" + ], + "title": "DiskCacheEntryDeleted", + "type": "object" + }, + "DiskCacheInfo": { + "additionalProperties": false, + "properties": { + "bytes": { + "title": "Bytes", + "type": "integer" + }, + "clearable": { + "title": "Clearable", + "type": "boolean" + }, + "cost_of_clearing": { + "title": "Cost Of Clearing", + "type": "string" + }, + "disk_free": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Disk Free" + }, + "disk_total": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Disk Total" + }, + "enabled": { + "title": "Enabled", + "type": "boolean" + }, + "entries": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/DiskCacheEntry" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "description": "For a cache with entries (quantized).", + "title": "Entries" + }, + "files": { + "title": "Files", + "type": "integer" + }, + "last_used": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "description": "The newest file's mtime (epoch seconds); null when empty.", + "title": "Last Used" + }, + "low_disk": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "For a cache with a free-space floor (quantized): the volume is under it, so new entries are not saved.", + "title": "Low Disk" + }, + "name": { + "title": "Name", + "type": "string" + }, + "not_clearable_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Not Clearable Reason" + }, + "root": { + "title": "Root", + "type": "string" + } + }, + "required": [ + "name", + "root", + "enabled", + "bytes", + "files", + "last_used", + "clearable", + "cost_of_clearing", + "disk_free", + "disk_total" + ], + "title": "DiskCacheInfo", + "type": "object" + }, + "DiskCacheList": { + "additionalProperties": false, + "properties": { + "caches": { + "items": { + "$ref": "#/components/schemas/DiskCacheInfo" + }, + "title": "Caches", + "type": "array" + } + }, + "required": [ + "caches" + ], + "title": "DiskCacheList", + "type": "object" + }, "DiskUsage": { "additionalProperties": false, "properties": { @@ -4657,6 +4946,114 @@ "summary": "Delete Asset" } }, + "/api/caches": { + "get": { + "description": "dw's own disk caches (dw/disk_caches.py): each one's root, size, and\nwhether it can be cleared from here.", + "operationId": "list_caches_api_caches_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DiskCacheList" + } + } + }, + "description": "Successful Response" + } + }, + "summary": "List Caches" + } + }, + "/api/caches/{name}": { + "delete": { + "description": "Empty one cache's root, keeping the root and its marker. Refused\nwhile anything may be writing into it (`_refuse_while_writing`).", + "operationId": "clear_cache_api_caches__name__delete", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DiskCacheCleared" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Clear Cache" + } + }, + "/api/caches/{name}/entries/{entry}": { + "delete": { + "description": "Delete one entry of a cache that has entries (quantized). The id is\nchecked to be 64 hex before any path is built from it.", + "operationId": "delete_cache_entry_api_caches__name__entries__entry__delete", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "path", + "name": "entry", + "required": true, + "schema": { + "title": "Entry", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DiskCacheEntryDeleted" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Delete Cache Entry" + } + }, "/api/classes": { "get": { "description": "Class names of one kind (pipelines, models, schedulers,\nquantization) - the pickers' data source.", @@ -5577,7 +5974,7 @@ }, "/api/gallery/{name}/metadata": { "get": { - "description": "Generation metadata embedded in a saved image ('workflow' inside\nit is the full definition the editor can reopen), plus the job that\nproduced the file when history remembers one, plus - for audio and\nvideo - what the file itself holds: duration, format and level,\nwhich is how an agent that cannot listen checks a track. Only an\nimage embeds 'metadata' this way - it is always null for audio and\nvideo, since neither format has a slot this writer uses; recover\nthe recipe from 'job' (GET /api/jobs/{id}/workflow) when one is\nknown, or from nothing when it isn't (a kept asset has no job).\n\n`envelope=true` adds the soundtrack's level second by second, which\nis what says *where* in a track something is - whether a shot is\nstill voiced at its last frame, how deep the hole at a seam goes.\nOpt-in: a ten-minute track is 600 numbers, and the default call has\nto stay small.\n\n`name` may also be an 'asset:' reference, and then it is the input\nasset of that name that is described rather than an output (#127).\nThe numbers here - duration, frame count, fps, sample rate - are\nwhat decide whether a call will work at all, and for a file the\ncaller is about to *consume* they were previously unobtainable:\nthe only way to read a wav's length was to run a job that copied it\ninto the output directory. `job` is null for an asset with no\nrecorded provenance, and `source` says which of the two roots\nanswered.\n\nA kept asset (`keep_output`) is not always provenance-blind: when\nthe file it was kept from still had a job in history, `keep_output`\nrecorded `{job, run_id, version}` beside it, and this route reads it\nback the same way it reads a kept file's shots (#556). `run_id`/\n`version` are that source run's own - `job.id` is still the one to\npass `get_job_workflow`, since a pruned run leaves no `workflow.json`\nof its own to reopen.", + "description": "Generation metadata embedded in a saved image ('workflow' inside\nit is the full definition the editor can reopen), plus the job that\nproduced the file when history remembers one, plus - for audio and\nvideo - what the file itself holds: duration, format and level,\nwhich is how an agent that cannot listen checks a track. Only an\nimage embeds 'metadata' this way - it is always null for audio and\nvideo, since neither format has a slot this writer uses; recover\nthe recipe from 'job' (GET /api/jobs/{id}/workflow) when one is\nknown, or from nothing when it isn't (a kept asset has no job).\n\n`envelope=true` adds the soundtrack's level second by second, which\nis what says *where* in a track something is - whether a shot is\nstill voiced at its last frame, how deep the hole at a seam goes.\nOpt-in: a ten-minute track is 600 numbers, and the default call has\nto stay small. `interval_seconds` (default 1.0, 0.05 to 3600) is\nthe bin width: 0.1 places a cut on a vocal onset. A track that\nwould need more than 2400 bins at a finer-than-default width is\nrefused (422) with the interval that fits; the default width is\nnever refused for length (#823).\n\n`name` may also be an 'asset:' reference, and then it is the input\nasset of that name that is described rather than an output (#127).\nThe numbers here - duration, frame count, fps, sample rate - are\nwhat decide whether a call will work at all, and for a file the\ncaller is about to *consume* they were previously unobtainable:\nthe only way to read a wav's length was to run a job that copied it\ninto the output directory. `job` is null for an asset with no\nrecorded provenance, and `source` says which of the two roots\nanswered.\n\nA kept asset (`keep_output`) is not always provenance-blind: when\nthe file it was kept from still had a job in history, `keep_output`\nrecorded `{job, run_id, version}` beside it, and this route reads it\nback the same way it reads a kept file's shots (#556). `run_id`/\n`version` are that source run's own - `job.id` is still the one to\npass `get_job_workflow`, since a pruned run leaves no `workflow.json`\nof its own to reopen.", "operationId": "gallery_metadata_api_gallery__name__metadata_get", "parameters": [ { @@ -5599,6 +5996,16 @@ "type": "boolean" } }, + { + "in": "query", + "name": "interval_seconds", + "required": false, + "schema": { + "default": 1.0, + "title": "Interval Seconds", + "type": "number" + } + }, { "in": "query", "name": "workspace", diff --git a/ui/src/lib/pages/ServerPage.svelte b/ui/src/lib/pages/ServerPage.svelte index 33649dd8..dd21db02 100644 --- a/ui/src/lib/pages/ServerPage.svelte +++ b/ui/src/lib/pages/ServerPage.svelte @@ -5,6 +5,7 @@ import CopyButton from '../CopyButton.svelte' import WorkerList from '../WorkerList.svelte' import CardMemoryList from '../CardMemoryList.svelte' + import CachesPanel from '../CachesPanel.svelte' import { cardReadings } from '../cardMemory' import { addressLabel, @@ -305,6 +306,8 @@ + + {:else if !error}

loading server details…

{/if} diff --git a/ui/src/lib/pages/ServerPage.test.ts b/ui/src/lib/pages/ServerPage.test.ts index 0dab3462..b212ad75 100644 --- a/ui/src/lib/pages/ServerPage.test.ts +++ b/ui/src/lib/pages/ServerPage.test.ts @@ -60,13 +60,70 @@ const health = vi.hoisted(() => const memory = vi.hoisted(() => vi.fn((): Promise => Promise.reject(new Error('503'))), ) +const TRITON_COST = + 'The next first step of each SDNQ or compiled shape re-autotunes. ' + + 'This directory is shared with any other Triton user on this machine.' +const cacheList = () => ({ + caches: [ + { + name: 'inductor', + root: '/home/don/.cache/dw/torchinductor', + enabled: true, + bytes: 273 * 1024 ** 2, + files: 1200, + last_used: 1760000000, + clearable: true, + cost_of_clearing: 'The next compiled run recompiles its kernels.', + disk_free: 400 * 1024 ** 3, + disk_total: 1000 * 1024 ** 3, + }, + { + name: 'triton', + root: '/home/don/.triton/cache', + enabled: true, + bytes: Math.round(7.1 * 1024 ** 3), + files: 19683, + last_used: 1760000000, + clearable: true, + cost_of_clearing: TRITON_COST, + disk_free: 400 * 1024 ** 3, + disk_total: 1000 * 1024 ** 3, + }, + ] as Record[], +}) +const listCaches = vi.hoisted(() => + vi.fn((): Promise => Promise.resolve({ caches: [] })), +) +const clearCache = vi.hoisted(() => + vi.fn((name: string): Promise => + Promise.resolve({ name, cleared: true, freed: 0, files: 0 }), + ), +) +const confirmDialog = vi.hoisted(() => + vi.fn((...args: [message: string, options?: unknown]): Promise => + Promise.resolve(args.length > 0), + ), +) +const deleteCacheEntry = vi.hoisted(() => + vi.fn((name: string, entry: string): Promise => + Promise.resolve({ name, entry, deleted: true, freed: 0 }), + ), +) vi.mock('../api', () => ({ api: { server: () => server(), health: () => health(), memory: () => memory(), + listCaches: () => listCaches(), + clearCache: (name: string) => clearCache(name), + deleteCacheEntry: (name: string, entry: string) => + deleteCacheEntry(name, entry), }, })) +vi.mock('../confirm.svelte', () => ({ + confirmDialog: (message: string, options?: unknown) => + confirmDialog(message, options), +})) /** Render with a given payload and let the promises settle, returning the * page's text with whitespace collapsed - the prose wraps in the markup, so @@ -88,6 +145,11 @@ beforeEach(() => { afterEach(() => { server.mockClear() health.mockClear() + listCaches.mockReset() + listCaches.mockResolvedValue({ caches: [] }) + clearCache.mockClear() + deleteCacheEntry.mockClear() + confirmDialog.mockClear() localStorage.clear() }) @@ -207,3 +269,219 @@ it('shows memory for every card, not just the first', async () => { expect(text).toContain('cuda:0 1.0 / 24.0 GB') expect(text).toContain('cuda:1 18.2 / 24.0 GB 12s ago') }) + +/** Render with the cache list given, returning the container. */ +async function showCaches(list = cacheList()) { + listCaches.mockResolvedValue(list) + state.info = base + state.fail = '' + const { container } = render(ServerPage) + for (let i = 0; i < 5; i++) await new Promise((r) => setTimeout(r, 0)) + return container +} +const settle = async () => { + for (let i = 0; i < 5; i++) await new Promise((r) => setTimeout(r, 0)) +} +const flat = (el: Element) => (el.textContent ?? '').replace(/\s+/g, ' ') +const clearButton = (container: Element, name: string) => + [...container.querySelectorAll('.caches tbody tr')] + .find((tr) => tr.querySelector('td')?.textContent === name)! + .querySelector('button') as HTMLButtonElement + +it('lists each cache with its root, size, files and the disk it is on', async () => { + const container = await showCaches() + const text = flat(container.querySelector('.caches')!) + expect(text).toContain('Caches') + expect(text).toContain('inductor') + expect(text).toContain('/home/don/.cache/dw/torchinductor') + expect(text).toContain('273.0 MB') + expect(text).toContain('1200') + expect(text).toContain('triton') + expect(text).toContain('/home/don/.triton/cache') + expect(text).toContain('7.10 GB') + expect(text).toContain('19683') + expect(text).toContain('400.0 / 1000.0 GB free') + // the bar is the used share of the volume + const bar = container.querySelector('.caches .diskbar > span') as HTMLElement + expect(bar.getAttribute('style')).toContain('width: 60%') +}) + +it('asks before clearing, quoting the cost and the shared-cache warning', async () => { + const container = await showCaches() + clearButton(container, 'triton').click() + await settle() + expect(confirmDialog).toHaveBeenCalledOnce() + const message = confirmDialog.mock.calls[0][0] + expect(message).toContain('triton') + expect(message).toContain('7.10 GB') + expect(message).toContain('re-autotunes') + expect(message).toContain('shared with any other Triton user') + expect(clearCache).toHaveBeenCalledWith('triton') + // the list is read again after the clear, so the size shown drops + expect(listCaches).toHaveBeenCalledTimes(2) +}) + +it('names an empty cache as 0 bytes in the confirm, not 1 KB', async () => { + const list = cacheList() + list.caches[0] = { ...list.caches[0], bytes: 0, files: 0 } + const container = await showCaches(list) + clearButton(container, 'inductor').click() + await settle() + const message = confirmDialog.mock.calls[0][0] + expect(message).toContain('(0 bytes, 0 files)') + expect(message).not.toContain('1 KB') +}) + +it('does not clear when the confirm is declined', async () => { + confirmDialog.mockResolvedValueOnce(false) + const container = await showCaches() + clearButton(container, 'inductor').click() + await settle() + expect(confirmDialog).toHaveBeenCalledOnce() + expect(clearCache).not.toHaveBeenCalled() +}) + +it('disables Clear with the reason when a cache is not clearable', async () => { + const list = cacheList() + list.caches[0] = { + ...list.caches[0], + clearable: false, + not_clearable_reason: 'root not created by dw (no .dw-cache marker)', + } + const container = await showCaches(list) + const button = clearButton(container, 'inductor') + expect(button.disabled).toBe(true) + expect(button.title).toBe('root not created by dw (no .dw-cache marker)') + expect(flat(container)).toContain( + 'root not created by dw (no .dw-cache marker)', + ) + expect(clearButton(container, 'triton').disabled).toBe(false) +}) + +it("shows the server's text when a clear is refused as busy", async () => { + const busy = + 'A job is running or queued - clearing a cache it may be writing ' + + 'would leave it a truncated artifact' + clearCache.mockRejectedValueOnce( + Object.assign(new Error(busy), { status: 409 }), + ) + const container = await showCaches() + clearButton(container, 'inductor').click() + await settle() + expect(flat(container)).toContain(busy) + // and the sizes are re-read, unchanged + expect(flat(container)).toContain('273.0 MB') +}) + +const ID_A = 'a'.repeat(64) +const ID_B = 'b'.repeat(64) +const quantizedList = (extra: Record = {}) => { + const list = cacheList() + list.caches.push({ + name: 'quantized', + root: '/home/don/.cache/dw/quantized', + enabled: true, + bytes: 5 * 1024 ** 3, + files: 12, + last_used: 1760000000, + clearable: true, + cost_of_clearing: 'The next load re-quantizes the component.', + disk_free: 400 * 1024 ** 3, + disk_total: 1000 * 1024 ** 3, + low_disk: false, + entries: [ + { + id: ID_A, + bytes: 3 * 1024 ** 3, + files: 8, + unreadable: false, + model: 'org/flux', + subfolder: 'transformer', + component: 'transformer', + created: 1760000000, + last_hit: null, + }, + { + id: ID_B, + bytes: 2 * 1024 ** 3, + files: 4, + unreadable: true, + model: '', + subfolder: '', + component: '', + created: 1760000000, + last_hit: 1760000500, + }, + ], + ...extra, + }) + return list +} +const toggleButton = (c: Element) => + c.querySelector('.entrytoggle') as HTMLButtonElement +const entryDelete = (c: Element, i: number) => + c + .querySelectorAll('.entrytable tbody tr') + [i].querySelector('button') as HTMLButtonElement + +it('shows an entry count that expands to the entry table', async () => { + const container = await showCaches(quantizedList()) + expect(flat(toggleButton(container))).toContain('2 entries') + expect(container.querySelector('.entrytable')).toBeNull() + toggleButton(container).click() + await settle() + const text = flat(container.querySelector('.entrytable')!) + expect(text).toContain('org/flux') + expect(text).toContain('transformer') + expect(text).toContain('3.00 GB') + expect(text).toContain('unreadable') + expect(text).toContain('a'.repeat(12)) + expect(text).not.toContain('a'.repeat(13)) + expect(container.querySelector(`code[title="${ID_A}"]`)).not.toBeNull() + expect(text).toContain('—') +}) + +it('deletes an entry after the confirm, then re-reads the caches', async () => { + const container = await showCaches(quantizedList()) + toggleButton(container).click() + await settle() + entryDelete(container, 0).click() + await settle() + const message = confirmDialog.mock.calls[0][0] + expect(message).toContain('transformer') + expect(message).toContain('org/flux') + expect(message).toContain('3.00 GB') + expect(message).toContain('re-quantizes') + expect(deleteCacheEntry).toHaveBeenCalledWith('quantized', ID_A) + expect(listCaches).toHaveBeenCalledTimes(2) +}) + +it("shows the server's text when an entry delete is refused", async () => { + const busy = 'A model download is active - try again when it finishes' + deleteCacheEntry.mockRejectedValueOnce( + Object.assign(new Error(busy), { status: 409 }), + ) + const container = await showCaches(quantizedList()) + toggleButton(container).click() + await settle() + entryDelete(container, 1).click() + await settle() + expect(flat(container)).toContain(busy) +}) + +it('disables entry Delete when the cache is not clearable', async () => { + const container = await showCaches( + quantizedList({ clearable: false, not_clearable_reason: 'no marker' }), + ) + toggleButton(container).click() + await settle() + expect(entryDelete(container, 0).disabled).toBe(true) + expect(entryDelete(container, 1).disabled).toBe(true) +}) + +it('warns when the volume is under the free-space floor', async () => { + const container = await showCaches(quantizedList({ low_disk: true })) + expect(flat(container)).toContain('free-space floor') + const ok = await showCaches(quantizedList()) + expect(flat(ok.querySelector('.caches')!)).not.toContain('free-space floor') +}) diff --git a/ui/src/lib/types.ts b/ui/src/lib/types.ts index 0637d1cc..ed18f539 100644 --- a/ui/src/lib/types.ts +++ b/ui/src/lib/types.ts @@ -17,6 +17,11 @@ export type ModelDownload = Schemas['ModelDownload'] export type ModelDownloads = Schemas['ModelDownloads'] export type ModelDeleted = Schemas['ModelDeleted'] export type DiffusersStatus = Schemas['DiffusersStatus'] +export type DiskCacheInfo = Schemas['DiskCacheInfo'] +export type DiskCacheList = Schemas['DiskCacheList'] +export type DiskCacheCleared = Schemas['DiskCacheCleared'] +export type DiskCacheEntry = Schemas['DiskCacheEntry'] +export type DiskCacheEntryDeleted = Schemas['DiskCacheEntryDeleted'] export type JobSummary = Schemas['JobSummary'] export type ManifestEntry = Schemas['ManifestEntry'] export type JobDetail = Schemas['JobDetail'] diff --git a/workflows/templates/ltx2/text-to-video.json b/workflows/templates/ltx2/text-to-video.json index c3e721d9..d3f88fd3 100644 --- a/workflows/templates/ltx2/text-to-video.json +++ b/workflows/templates/ltx2/text-to-video.json @@ -1,6 +1,6 @@ { "id": "LTX2", - "description": "The baseline for the LTX-2.5 family: text to video with a soundtrack generated alongside it, on a single 24GB card. The distilled transformer runs a fixed eight-step schedule with guidance off, quantized per component so it stays resident while the Gemma text encoder and the connectors are group offloaded around it. First load adds about 10 minutes.", + "description": "The baseline for the LTX-2.5 family: text to video with a soundtrack generated alongside it, on a single 24GB card. The distilled transformer runs a fixed eight-step schedule with guidance off, quantized per component so it stays resident while the Gemma text encoder and the connectors are group offloaded around it. First load adds about 10 minutes to quantize; with the server's quantized cache on (DW_QUANTIZED_CACHE=1) only the first load pays it, and later loads read the quantized weights from disk.", "cost": [ {"device": "cuda", "name": "RTX 3090", "vram_gb": 24, "minutes": 1.8}, {"device": "mps", "name": "M5 Pro 64GB", "vram_gb": 64, "minutes": 6} diff --git a/workflows/templates/minimax/dialogue-short.json b/workflows/templates/minimax/dialogue-short.json index 6fc3fd58..996b7581 100644 --- a/workflows/templates/minimax/dialogue-short.json +++ b/workflows/templates/minimax/dialogue-short.json @@ -1,6 +1,6 @@ { "id": "MiniMaxH3SitcomShort", - "description": "A digital short built the way television is built: from cuts, not from one long take. Chained generation degrades with length - every segment conditions on the previous segment's output, so artifacts compound and identity drifts. A scene cut resets that completely: each shot here is generated fresh from the same two character portraits, so shot five is exactly as clean as shot one and the scene can run as long as the script does. Two Z-Image steps draw the cast (the second reuses the first's loaded pipeline - identical configurations share one model - and 'release_pipeline' frees it before the video model loads). The shots are one 'for_each' step over the 'shots' list: one entry per shot, carrying its 'name', its 'prompt', its 'references' and its 'num_frames' - the tag entry runs 141 frames where the others run 124, since length is per-shot. The list is an argument, so a six-shot scene is one more entry, not another file, and the members are named for their entries ('shot@react'); the loaded MiniMax-H3 is reused across all of them. Character consistency across cuts comes from referencing the same portraits in every shot; voice consistency comes from repeating each character's voice description verbatim in every prompt, and, when 'character_a_voice' / 'character_b_voice' name a clip ('asset:cast/priya.wav'), from the audio reference each entry lists for whoever speaks in it - an entry's reference says 'variable:character_a_voice', so one variable sets the voice in every shot that character has. Both default to null, and a reference whose file is null is left out of the list, so a run that names no voice generates exactly what it generated before the variables existed. A shot where both speak lists both; if that reads worse than one, name one voice and leave the other null. The variables are named for roles rather than for the cast of this example - 'character_a', the 'react' entry - because the beats are the reusable part and the sketch is not. The soundscape writes the laugh track. A final 'concat_videos' task is the editor, gathering the shots in list order into one episode - hard cuts, no trims, no seams to hide, because nothing was carried between them. Only the picture cuts hard: 'audio_bleed_ms' rings each shot's laugh track on over the silent opening of the next, the way a live audience carries across a cut. It and 'seam_fade_ms' are variables, so a seam is re-tuned with an argument rather than a copy of the workflow; 1800 ms is where a five-shot cut measured best, since a generated shot opens on more silence than it looks; 'audio_bleed_gain_db' ducks it, left at 0 for no change. The two are mutually exclusive, not additive: at a hard cut with 'audio_bleed_ms' non-zero the join takes the bleed path and 'seam_fade_ms' has no effect at all, so a plain declick fade instead of a bleed means setting 'audio_bleed_ms' to 0, not just passing 'seam_fade_ms'. Shots generated independently also drift in loudness - 10 dB between two shots of one scene is ordinary - and no seam control can hide a level jump, because it is either side of the cut rather than at it; 'match_levels' ('rms' for perceived level, 'peak' for the loudest sample) evens the shots out before they are joined, and left null, as it is by default, a wide spread is warned about in the log rather than passing in silence. An episode can be cast from portraits that already exist rather than drawn here: a shot entry's subject reference takes 'from_file' the way its voice references do ('asset:cast/priya.jpg'), so a recurring cast carries across episodes by name instead of being redrawn each time - which is the whole point of a cast. An episode cast entirely from files does not pay for them: the two Z-Image steps save nothing ('save': false - the shots are the only thing that reads them), so once no shot references them the engine drops them from the run and warns that it did, and the minute and two model loads they cost are not spent (#122). Cast some shots from files and not others and they still run, because something still reads them.", + "description": "A digital short built the way television is built: from cuts, not from one long take. Chained generation degrades with length - every segment conditions on the previous segment's output, so artifacts compound and identity drifts. A scene cut resets that completely: each shot here is generated fresh from the same two character portraits, so shot five is exactly as clean as shot one and the scene can run as long as the script does. Two Z-Image steps draw the cast (the second reuses the first's loaded pipeline - identical configurations share one model - and 'release_pipeline' frees it before the video model loads). The shots are one 'for_each' step over the 'shots' list: one entry per shot, carrying its 'name', its 'prompt', its 'references' and its 'num_frames' - the tag entry runs 141 frames where the others run 124, since length is per-shot. The list is an argument, so a six-shot scene is one more entry, not another file, and the members are named for their entries ('shot@react'); the loaded MiniMax-H3 is reused across all of them. Character consistency across cuts comes from referencing the same portraits in every shot; voice consistency comes from repeating each character's voice description verbatim in every prompt, and, when 'character_a_voice' / 'character_b_voice' name a clip ('asset:cast/priya.wav'), from the audio reference each entry lists for whoever speaks in it - an entry's reference says 'variable:character_a_voice', so one variable sets the voice in every shot that character has. Both default to null, and a reference whose file is null is left out of the list, so a run that names no voice generates exactly what it generated before the variables existed. A shot where both speak lists both; if that reads worse than one, name one voice and leave the other null. The variables are named for roles rather than for the cast of this example - 'character_a', the 'react' entry - because the beats are the reusable part and the sketch is not. The soundscape writes the laugh track. A final 'concat_videos' task is the editor, gathering the shots in list order into one episode - hard cuts, no trims, no seams to hide, because nothing was carried between them. Only the picture cuts hard: 'audio_bleed_ms' rings each shot's laugh track on over the silent opening of the next, the way a live audience carries across a cut. It and 'seam_fade_ms' are variables, so a seam is re-tuned with an argument rather than a copy of the workflow; 1800 ms is where a five-shot cut measured best, since a generated shot opens on more silence than it looks; 'audio_bleed_gain_db' ducks it, left at 0 for no change. The two are mutually exclusive, not additive: at a hard cut with 'audio_bleed_ms' non-zero the join takes the bleed path and 'seam_fade_ms' has no effect at all, so a plain declick fade instead of a bleed means setting 'audio_bleed_ms' to 0, not just passing 'seam_fade_ms'. Shots generated independently also drift in loudness - 10 dB between two shots of one scene is ordinary - and no seam control can hide a level jump, because it is either side of the cut rather than at it; 'match_levels' ('rms' for perceived level, 'peak' for the loudest sample) evens the shots out before they are joined, and left null, as it is by default, a wide spread is warned about in the log rather than passing in silence. 'match_levels_dbfs' sets the level match_levels moves every shot to, below its own default of -20; H3 shots are hot enough that a shot whose gain to -20 would clip is held short of it instead (the run warns), leaving the shots unevenly apart, so lower it (-24 clears the template's own shots) when that warning fires. An episode can be cast from portraits that already exist rather than drawn here: a shot entry's subject reference takes 'from_file' the way its voice references do ('asset:cast/priya.jpg'), so a recurring cast carries across episodes by name instead of being redrawn each time - which is the whole point of a cast. An episode cast entirely from files does not pay for them: the two Z-Image steps save nothing ('save': false - the shots are the only thing that reads them), so once no shot references them the engine drops them from the run and warns that it did, and the minute and two model loads they cost are not spent (#122). Cast some shots from files and not others and they still run, because something still reads them.", "summary": "A multi-shot dialogue short: Z-Image draws the cast, each shot is generated fresh from the same portraits, then cut.", "cost": [ { @@ -131,6 +131,7 @@ "audio_bleed_gain_db": 0.0, "seam_fade_ms": null, "match_levels": null, + "match_levels_dbfs": null, "width": 960, "height": 544, "num_inference_steps": 9, @@ -370,7 +371,8 @@ "audio_bleed_gain_db": "variable:audio_bleed_gain_db", "fps": 24, "seam_fade_ms": "variable:seam_fade_ms", - "match_levels": "variable:match_levels" + "match_levels": "variable:match_levels", + "match_levels_dbfs": "variable:match_levels_dbfs" } }, "result": { diff --git a/workflows/templates/minimax/music-video-cuts.json b/workflows/templates/minimax/music-video-cuts.json index be8dfff2..85606d47 100644 --- a/workflows/templates/minimax/music-video-cuts.json +++ b/workflows/templates/minimax/music-video-cuts.json @@ -1,6 +1,6 @@ { "id": "MiniMaxMusicVideoCuts", - "description": "Plans a music video's cuts from a song and its lyrics. It transcribes the song with timestamps ('transcribe', a Whisper-class model) and finds its beats ('beats'), then 'plan_cuts' turns the two into a shot list ('plan') covering the whole song, to the length 'beats' measured. Nothing is generated: the result is one JSON document of shots, each {name, start_frame, num_frames, lead_frames, cut_frames, lyric, kind}, tiling the song exactly at 'fps', so the agent writes one prompt per shot and renders the list with music-video.json, whose entries are {name, prompt, start_frame, num_frames, lead_frames, cut_frames}: add a 'prompt' to each entry and drop 'lyric' and 'kind' to render the plan as it stands. The shots are sized for H3 - 'num_frames' on its 17n+5 grid between 124 and 345 frames, with 'lead_s' seconds of run-up before each cut that music-video trims off - and the result's 'render_frames' is the total frames that will be generated, lead included, for quoting the cost. 'lyrics' is the song's own words, one line per line with a blank line or a lone section tag such as [Chorus] between stanzas: the shots carry those lines verbatim and the transcript only times them, since Whisper mishears sung words. Left null, the shots carry the transcript's text. 'segment_by' is 'line' (a shot per sung line), 'stanza' or 'beat' (a cut every few beats). 'min_scene_s' merges a shot that is too short into its neighbour and 'max_scene_s' splits one that is too long, on a beat. 'vocal_tail_s' holds a sung shot past its last word, 'include_instrumental_gaps' gives a long silence its own instrumental shot, and 'snap_to_beats' moves every cut to its nearest beat. 'timestamps' must stay 'segment' (a Whisper segment per line) or be 'word'; plan_cuts refuses a transcript without timings.", + "description": "Plans a music video's cuts from a song and its lyrics. It separates the vocal stem ('stems', htdemucs) so the transcript is of the voice alone - transcribing the full mix folds an instrumental intro into the first word and lets Whisper invent text over a fade-out - then transcribes that stem with timestamps ('transcribe', a Whisper-class model) and finds the beats of the full mix ('beats'), then 'plan_cuts' turns the two into a shot list ('plan') covering the whole song, to the length 'beats' measured. Nothing is generated: the result is one JSON document of shots, each {name, start_frame, num_frames, lead_frames, cut_frames, lyric, kind}, tiling the song exactly at 'fps', so the agent writes one prompt per shot and renders the list with music-video.json, whose entries are {name, prompt, start_frame, num_frames, lead_frames, cut_frames}: add a 'prompt' to each entry and drop 'lyric' and 'kind' to render the plan as it stands. The shots are sized for H3 - 'num_frames' on its 17n+5 grid between 124 and 345 frames, with 'lead_s' seconds of run-up before each cut that music-video trims off - and the result's 'render_frames' is the total frames that will be generated, lead included, for quoting the cost. 'lyrics' is the song's own words, one line per line with a blank line or a lone section tag such as [Chorus] between stanzas: the shots carry those lines verbatim and the transcript only times them, since Whisper mishears sung words. Left null, the shots carry the transcript's text. 'segment_by' is 'line' (a shot per sung line), 'stanza' or 'beat' (a cut every few beats). 'min_scene_s' merges a shot that is too short into its neighbour and 'max_scene_s' splits one that is too long, on a beat. 'vocal_tail_s' holds a sung shot past its last word, 'include_instrumental_gaps' gives a long silence its own instrumental shot, and 'snap_to_beats' moves every cut to its nearest beat. 'timestamps' must stay 'segment' (a Whisper segment per line) or be 'word'; plan_cuts refuses a transcript without timings.", "cost_drivers": [ "fps" ], @@ -19,12 +19,25 @@ "lead_s": 0.5 }, "steps": [ + { + "name": "stems", + "task": { + "command": "separate_stems", + "arguments": { + "audio": "variable:song" + } + }, + "result": { + "content_type": "audio/wav", + "subfolder": "intermediate" + } + }, { "name": "transcribe", "task": { "command": "transcribe_audio", "arguments": { - "audio": "variable:song", + "audio": "previous_result:stems.vocals", "timestamps": "variable:timestamps", "model_name": "variable:model_name" } diff --git a/workflows/templates/minimax/music-video.json b/workflows/templates/minimax/music-video.json index 008add00..1d966cb6 100644 --- a/workflows/templates/minimax/music-video.json +++ b/workflows/templates/minimax/music-video.json @@ -80,6 +80,7 @@ "audio_bleed_gain_db": 0.0, "seam_fade_ms": null, "match_levels": null, + "match_levels_dbfs": null, "width": 960, "height": 544, "num_inference_steps": 9, @@ -364,7 +365,8 @@ "audio_bleed_ms": "variable:audio_bleed_ms", "audio_bleed_gain_db": "variable:audio_bleed_gain_db", "seam_fade_ms": "variable:seam_fade_ms", - "match_levels": "variable:match_levels" + "match_levels": "variable:match_levels", + "match_levels_dbfs": "variable:match_levels_dbfs" } } },