Add Apple Silicon local LLM inference
This commit is contained in:
@@ -17,6 +17,7 @@
|
||||
!/llm_compatibility.py
|
||||
!/llm_resolver.py
|
||||
!/llm_staging.py
|
||||
!/local_llm_runtime.py
|
||||
!/modal_executor_node.py
|
||||
!/modal_endpoint_node.py
|
||||
!/modal_llm_node.py
|
||||
|
||||
@@ -10,7 +10,7 @@ ComfyUI Modal-Sync is a ComfyUI custom node extension for running selected parts
|
||||
Modal-Sync provides:
|
||||
|
||||
- a ComfyUI frontend extension with a `Run on Modal` toggle, workflow-level GPU selection, and remote execution overlays
|
||||
- a resident `Modal LLM` node for text, image, video, and bounded file understanding on the same GPU worker as ComfyUI
|
||||
- a resident `Modal LLM` node for text, image, video, and bounded file understanding on either Apple Silicon or the same Modal GPU worker as ComfyUI
|
||||
- a `Modal Endpoint Chat` node for prompt, image, and file inference through Modal hosted-model endpoints
|
||||
- a queue route at `/modal/queue_prompt` that intercepts normal prompt submission
|
||||
- queue-time graph partitioning and proxy-node rewrite for selected remote regions
|
||||
@@ -112,7 +112,20 @@ For credential safety, the node accepts only HTTPS `modal.direct` origins, refus
|
||||
|
||||
### Running A Resident LLM Beside ComfyUI
|
||||
|
||||
`Modal LLM` is a V3 node modelled on ComfyUI's built-in OpenAI chat node, but inference happens inside the same Modal GPU worker that executes the surrounding remote ComfyUI component. Enable `Run on Modal` for the node. It deliberately refuses local execution so a missed toggle cannot download or load a multi-gigabyte model into the desktop ComfyUI process.
|
||||
`Modal LLM` is a V3 multimodal language-model node with two execution targets controlled by the existing node toggle:
|
||||
|
||||
- Leave `Run on Modal` disabled to run the node locally on Apple Silicon through MLX-VLM. This is intended for MacBooks whose unified memory can hold an LLM while heavier image or video nodes run on Modal.
|
||||
- Enable `Run on Modal` to run the node inside the same Modal GPU worker as its surrounding remote component.
|
||||
|
||||
The queue planner treats an unmarked Modal LLM as an ordinary local node, including when it sits between two remote regions. Inputs and outputs cross those region boundaries through the existing transport layer; the model itself stays on the Mac.
|
||||
|
||||
Apple-local inference requires macOS on Apple Silicon and the optional pinned runtime in the Python environment that launches ComfyUI:
|
||||
|
||||
```bash
|
||||
uv pip install --python <comfyui-venv>/bin/python "mlx-vlm==0.6.15" "psutil>=7,<8"
|
||||
```
|
||||
|
||||
Repository developers can instead run `uv sync --extra local-apple --extra remote --group test`. Other local platforms receive an actionable error and can enable `Run on Modal`. Local snapshots are stored under `<ComfyUI models>/modal_llm`; set `COMFY_MODAL_LOCAL_LLM_STORAGE_ROOT` to use another writable location.
|
||||
|
||||
The node accepts:
|
||||
|
||||
@@ -122,15 +135,15 @@ The node accepts:
|
||||
- UTF-8 text or text-based PDF values from `OpenAI ChatGPT Input Files`
|
||||
- bounded generation, sampling, seed, video-frame, VRAM-reserve, and residency controls
|
||||
|
||||
It returns three strings: the clean final `response`, compact `metadata_json`, and a separate `reasoning` channel. The reasoning output is appended after the two original outputs so existing response and metadata links retain their indices. For thinking models, `max_new_tokens` covers reasoning and final-response tokens together; a generation that reaches the limit before its reasoning terminator produces an empty response and keeps all partial work in `reasoning`. Telemetry includes total output tokens, time to first token, live/final tokens per second, detected reasoning tokens and parser, cold/warm model status, media counts, GPU memory, the resident LLM profiles, and ComfyUI-managed models still resident after generation. The node's ordinary progress panel carries compact stages for profile/input preparation, CPU snapshot download, processor and weight loading, vLLM engine warmup, prefill, and generation. Known file, shard, and token totals use numeric bars; compilation and warmup use an indeterminate pulse. Prompt cancellation remains wired through ComfyUI's interruption hook without a separate inference server.
|
||||
It returns three strings: the clean final `response`, compact `metadata_json`, and a separate `reasoning` channel. The reasoning output is appended after the two original outputs so existing response and metadata links retain their indices. For thinking models, `max_new_tokens` covers reasoning and final-response tokens together. The MLX backend caps thinking at half that budget so it retains room for a final answer; the remote backend retains its existing model-native thinking behavior. Telemetry includes execution target, device, total output tokens, time to first token, live/final tokens per second, detected reasoning tokens and parser, cold/warm model status, media counts, available memory, resident LLM profiles, and ComfyUI-managed models still resident after generation. The node's ordinary progress panel carries compact stages for profile/input preparation, snapshot download, processor and weight loading, engine warmup where applicable, prefill, and generation. Prompt cancellation remains wired through ComfyUI's interruption hook without a separate inference server.
|
||||
|
||||
Modal LLM uses immutable schema-v2 model profiles. Checked-in entries in [`llm_profiles.json`](llm_profiles.json) remain available, while generated profiles content-address the exact Hugging Face revision, compatibility-policy version, backend, quantization, reasoning parser, context/media limits, expected VRAM, and remote-code policy. The compatibility registry currently selects Transformers for Muse-Glimmer and vLLM for Qwen3.5, including block-FP8 and NVIDIA ModelOpt NVFP4 checkpoints. Qwen3.5 profiles select the `qwen3` reasoning parser. Unknown architectures fail compatibility inspection before weight download.
|
||||
Modal LLM uses immutable schema-v2 model profiles. Checked-in entries in [`llm_profiles.json`](llm_profiles.json) remain available, while generated profiles content-address the execution target, exact Hugging Face revision, compatibility-policy version, backend, quantization, reasoning parser, context/media limits, expected memory, and remote-code policy. The local registry selects MLX-VLM for reviewed SmolVLM, Muse-Glimmer, and Qwen3.5 architectures; the Modal registry selects Transformers or vLLM for its reviewed CUDA formats. Qwen3.5 profiles select the shared `qwen3` reasoning parser. Unknown or target-incompatible architectures and quantizations fail compatibility inspection before weight download.
|
||||
|
||||
Enter either a curated profile name or a Hugging Face `owner/model` ID directly in the Modal LLM node. On first execution, the deployed CPU `ModelStager` resolves and downloads it, commits the generated manifest and snapshot to the shared Volume, and the dispatcher rewrites only that execution payload to the immutable generated profile ID before requesting the GPU worker. Use `owner/model@revision` to resolve an explicit branch, tag, or commit; an unqualified ID pins the repository's current exact commit and does not silently update on later warm executions.
|
||||
Enter either a curated profile name or a Hugging Face `owner/model` ID directly in the node. Use `owner/model@revision` to resolve an explicit branch, tag, or commit. Each target performs metadata and compatibility inspection before downloading safetensors, pins the exact commit in a target-specific generated profile, and reuses its completed snapshot on later executions. Public models need no token. For gated local models, set `HF_TOKEN` in ComfyUI's environment; for gated Modal models, include it in the selected Modal secret collection.
|
||||
|
||||
Before any GPU call, the local dispatcher finds fixed LLM profiles in the remote prompt and calls the deployed CPU-only `ModelStager`. The stager downloads a missing immutable snapshot into `<volume>/llm_models/<profile>/<revision>`, writes a completion marker, commits the shared Modal Volume, and only then allows the selected GPU worker to start or reload that volume revision. A linked/dynamic `model_profile` is rejected because it cannot be staged before GPU allocation.
|
||||
For remote execution, the local dispatcher finds fixed LLM profiles in the prompt and calls the deployed CPU-only `ModelStager`. It writes the immutable snapshot and completion marker to the shared Modal Volume before the selected GPU worker starts. A linked/dynamic remote `model_profile` is rejected because it cannot be staged before GPU allocation. For local execution, resolution and staging happen directly in the ComfyUI process before MLX allocation.
|
||||
|
||||
The warm worker holds a process-global, serialized LRU of Transformers and vLLM models. A cache hit reuses both processor and weights across node and workflow executions. Before a cold load, it asks ComfyUI to release idle managed models, evicts older resident LLMs if necessary, and enforces the requested free-VRAM reserve. This makes co-residency intentional while allowing ComfyUI's image/video loaders to use the rest of a large GPU. Both backends pass generated token IDs through one model-aware reasoning abstraction. It prefers an engine-native separated field when available and otherwise splits exact model boundary tokens before decoding each clean channel. Transformers currently returns generated IDs without a reasoning field. The resident vLLM path owns one persistent event-loop thread and one `AsyncLLM` per resident model; cumulative `generate()` outputs provide live token counts while preserving the same in-process co-residency and LRU behavior. vLLM workers use the native top-k/top-p sampler and explicit Triton full attention because FlashInfer's optional sampler and Blackwell TRT-LLM attention kernels otherwise attempt first-run `nvcc` JIT compilation, while FlashAttention 4 rejects Qwen3.5's used-sequence prefill at head dimension 256. Qwen's model-specific GDN kernel remains available.
|
||||
Each target holds a process-global serialized LRU. A cache hit reuses the processor and weights across node and workflow executions. Before a cold load, the manager asks ComfyUI to release idle managed models, evicts older resident LLMs if necessary, and enforces the target's reserve: `local_reserve_free_memory_gb` defaults to 4 GiB for macOS unified memory, while `reserve_free_vram_gb` remains the Modal GPU control. Local MLX defaults to one resident model and can be changed with `COMFY_MODAL_LOCAL_LLM_MAX_RESIDENT_MODELS`; `COMFY_MODAL_LOCAL_LLM_RESERVE_FREE_GB` supplies the programmatic default when the runner is called without a node value. All backends pass generated token IDs through the same model-aware reasoning abstraction.
|
||||
|
||||
For a single-user B300 session, start conservatively with one container and one active component:
|
||||
|
||||
|
||||
@@ -404,3 +404,7 @@
|
||||
- For a freshly auto-deployed LLM app, stage pinned Hugging Face snapshots through the CPU-only class immediately after deployment and before the first `RemoteEngine` lookup. A runtime-version method belongs to the GPU class and can allocate the B300, defeating CPU-first staging if it is called first.
|
||||
- Transformers can load a SmolVLM2 snapshot successfully and still fail while constructing its processor because `num2words` is a model-specific optional dependency rather than a core Transformers dependency. Pin it in the Modal runtime image and cover model construction with a live canary, not just registry and staging unit tests.
|
||||
- Reasoning-enabled Qwen models can consume a small `max_new_tokens` allowance entirely before emitting `</think>`, correctly leaving the separate response output empty. Live inference canaries need enough combined reasoning-and-response budget, and vLLM canaries should assert positive streamed TTFT and token-rate telemetry rather than treating any generated token as proof that the async progress contract worked.
|
||||
- The Modal marker already provides the correct local/remote placement contract for an LLM node. Once the node's explicit local rejection is removed, an unmarked LLM remains a normal local scheduler node even between two remote regions; this preserves the existing transport and warning behavior without a second backend selector widget.
|
||||
- Apple-local generated profiles must include their execution target in the content digest so CUDA and MLX compatibility decisions can never share an identity. Preserve the historical Modal digest shape by omitting the default `modal` target, or existing immutable profile ids and saved payloads change unnecessarily.
|
||||
- Sampled ComfyUI video is already reduced to timestamped PIL frames by the shared input layer. Passing that ordered frame list to MLX-VLM as multiple image inputs avoids backend-specific video-container handling while preserving the bounded visual evidence described in the prompt.
|
||||
- MLX-VLM exposes a thinking-budget criterion that can force Qwen's closing reasoning token. Capping Apple-local thinking at half of `max_new_tokens` prevents smaller reasoning models from repeatedly consuming the entire allowance without a final response while preserving the separate reasoning channel.
|
||||
|
||||
+50
-54
@@ -1,79 +1,75 @@
|
||||
# Resident Modal LLM Architecture
|
||||
# Local And Modal Resident LLM Architecture
|
||||
|
||||
Modal LLM runs multimodal language-model inference inside the persistent RemoteEngine process. It is not an HTTP client and does not create an independently scaling inference service.
|
||||
The `Modal LLM` V3 node runs multimodal language-model inference in the current ComfyUI execution process. An unmarked node runs locally on Apple Silicon through MLX-VLM. A node marked `Run on Modal` runs in the persistent RemoteEngine process on the workflow's selected Modal GPU. Neither target creates a separate HTTP inference service.
|
||||
|
||||
## Request Path
|
||||
## Target Selection And Request Path
|
||||
|
||||
1. The V3 node exposes text, IMAGE, native VIDEO, and OPENAI_INPUT_FILES inputs.
|
||||
2. Queue rewriting places the marked node in a normal remote component.
|
||||
3. The local dispatcher scans the complete payload, including split phases, for curated profiles or Hugging Face model IDs.
|
||||
4. For a model ID, the CPU resolver inspects Hub metadata and `config.json`, validates architecture, quantization, safetensors inventory, security status, and download budget, then pins the exact commit in a content-addressed schema-v2 manifest.
|
||||
5. The CPU ModelStager downloads only the approved revision-pinned safetensors, processor, and tokenizer assets to the shared Modal Volume.
|
||||
6. The stager writes a completion marker and commits the Volume.
|
||||
7. The dispatcher rewrites the model ID to the immutable generated profile ID, and the GPU worker reloads that exact Volume revision before executing the component.
|
||||
8. The node converts ComfyUI tensors directly to processor inputs, samples video frames uniformly, and extracts bounded text from supported files.
|
||||
9. ResidentLLMManager reuses or loads the profile under a process lock. Structured progress covers processor setup, known checkpoint shard totals, vLLM engine warmup, prefill, and generation.
|
||||
10. A model-aware output parser separates generated reasoning from final content using exact token boundaries, with a native engine field taking precedence when available.
|
||||
11. The clean response, JSON telemetry, and separate reasoning string return through the normal Modal-Sync node-output transport.
|
||||
1. The node exposes text, IMAGE, native VIDEO, and `OPENAI_INPUT_FILES` inputs plus the existing `Run on Modal` toggle.
|
||||
2. Queue rewriting leaves an unmarked node in the local graph. Marking it places the node in a normal remote component.
|
||||
3. Both targets inspect a Hugging Face repository and `config.json` before weight download, validate the target-specific architecture and quantization policy, and pin the exact commit in a content-addressed schema-v2 profile.
|
||||
4. Local execution downloads the approved snapshot beneath `<ComfyUI models>/modal_llm` and loads it through pinned MLX-VLM. Modal execution uses the CPU-only ModelStager and shared Modal Volume before allocating a GPU worker.
|
||||
5. The node converts ComfyUI tensors directly to processor inputs, samples video uniformly into timestamped frames, and extracts bounded text from supported files. The local adapter supplies sampled video frames as an ordered image sequence because the shared input layer has already done the sampling and timestamp annotation.
|
||||
6. A target-specific process-global `ResidentLLMManager` reuses or loads the immutable profile under a process lock.
|
||||
7. A shared model-aware parser separates generated reasoning from final content using exact token boundaries, with a native engine field taking precedence when available.
|
||||
8. The clean response, JSON telemetry, and separate reasoning string return through normal ComfyUI node outputs.
|
||||
|
||||
An unmarked LLM may sit between remote regions. The planner keeps it local and creates the same transport boundaries it would for any other local node; only the values cross those boundaries.
|
||||
|
||||
## Apple-Local Runtime
|
||||
|
||||
Local execution currently requires Apple Silicon macOS and exactly `mlx-vlm==0.6.15`. Install it into the interpreter that launches ComfyUI:
|
||||
|
||||
```bash
|
||||
uv pip install --python <comfyui-venv>/bin/python "mlx-vlm==0.6.15" "psutil>=7,<8"
|
||||
```
|
||||
|
||||
The default snapshot root follows ComfyUI's model directory and registers the `modal_llm` model folder. `COMFY_MODAL_LOCAL_LLM_STORAGE_ROOT` overrides it. The compatibility registry currently permits unquantized or native MLX checkpoints for SmolVLM, Muse-Glimmer, and Qwen3.5 architectures. CUDA-oriented FP8, ModelOpt FP4, and unknown formats are rejected before download; users should choose an unquantized repository or an `mlx-community` conversion.
|
||||
|
||||
For Qwen3.5, MLX-VLM receives an explicit thinking budget equal to half of `max_new_tokens`. This preserves the existing separate reasoning output while ensuring that a verbose small model cannot consume the entire allowance before emitting a final answer.
|
||||
|
||||
## Residency And Memory
|
||||
|
||||
The manager is module-global in the warm RemoteEngine process. Its LRU is keyed by immutable profile ID and retains up to COMFY_MODAL_LLM_MAX_RESIDENT_MODELS.
|
||||
The shared manager implements a serialized, immutable-profile LRU for both targets. Before a cold load it asks ComfyUI to release idle managed models, evicts older resident LLMs as needed, and rejects the load when the model estimate plus configured reserve does not fit.
|
||||
|
||||
Before loading a new LLM, the manager:
|
||||
Apple-local admission uses `psutil.virtual_memory().available` and total system unified memory. The node's `local_reserve_free_memory_gb` defaults to 4 GiB, the local LRU defaults to one model, and evictions clear MLX caches. `COMFY_MODAL_LOCAL_LLM_MAX_RESIDENT_MODELS` changes the LRU limit. `COMFY_MODAL_LOCAL_LLM_RESERVE_FREE_GB` supplies the default to programmatic calls that omit a reserve.
|
||||
|
||||
1. asks ComfyUI's memory manager to release idle managed models if needed;
|
||||
2. checks torch.cuda.mem_get_info();
|
||||
3. evicts least-recently used LLMs until the profile estimate plus reserve fits;
|
||||
4. fails with measured free/total VRAM rather than attempting an unsafe load.
|
||||
Modal admission keeps the existing `torch.cuda.mem_get_info()` path, `reserve_free_vram_gb` control, and `COMFY_MODAL_LLM_MAX_RESIDENT_MODELS` setting. Executions are serialized within each process; retaining more than one model does not promise simultaneous kernels.
|
||||
|
||||
ComfyUI sees the LLM's real CUDA allocation when calculating free memory, even though it does not own the Transformers object. The configurable reserve gives subsequent image/video nodes headroom. Executions are serialized within a worker; co-residency does not promise simultaneous CUDA kernels.
|
||||
Telemetry exposes `execution_target`, `device`, generic available/total-memory fields, cache status, resident profiles, and ComfyUI-managed model names. Existing Modal GPU memory keys remain present for backward compatibility.
|
||||
|
||||
## Security And Reproducibility
|
||||
|
||||
- Repository IDs are metadata-inspected on CPU before any weight download or GPU allocation.
|
||||
- Every profile pins a 40-character repository commit.
|
||||
- Generated profile IDs include the SHA-256 digest of all runtime-defining fields.
|
||||
- `trust_remote_code` must remain false.
|
||||
- The remote model loader uses `local_files_only=true`.
|
||||
- The CPU stager is the only network download path.
|
||||
- Repository IDs are metadata-inspected before weight download or accelerator allocation.
|
||||
- Every generated profile pins a 40-character repository commit and includes the execution target in its content identity.
|
||||
- Existing Modal generated-profile identities remain unchanged for compatibility; `execution_target=modal` is intentionally omitted from their historical digest shape.
|
||||
- `trust_remote_code` remains false.
|
||||
- Snapshot allow-patterns omit Python source and pickle weight formats.
|
||||
- Text and PDF input sizes are bounded before prompt construction.
|
||||
- The compatibility policy participates in the remote runtime fingerprint.
|
||||
- Both loaders consume only the completed local snapshot path.
|
||||
- Text and PDF sizes are bounded before prompt construction.
|
||||
- Public models need no token. Local gated access reads `HF_TOKEN` from ComfyUI's environment; remote gated access reads it from the selected Modal secret collection.
|
||||
|
||||
Compatibility-policy changes remain deployment changes. New repositories using a reviewed architecture and quantization can generate immutable profiles without changing the image.
|
||||
Changing target requires resolving the original Hugging Face reference again. A generated profile created for Modal cannot be silently loaded through MLX, and vice versa.
|
||||
|
||||
## Current Backend Boundary
|
||||
## Backend Boundary
|
||||
|
||||
The compatibility registry selects Transformers for Muse-Glimmer and vLLM for Qwen3.5. Qwen block-FP8 and NVIDIA ModelOpt NVFP4 checkpoints share the vLLM adapter. Generated vLLM profiles set an explicit KV-cache byte budget and a conservative 32K default context instead of allowing the serving engine to reserve nearly all available GPU memory.
|
||||
The Modal compatibility registry selects Transformers for Muse-Glimmer and vLLM for Qwen3.5, including reviewed block-FP8 and NVIDIA ModelOpt NVFP4 checkpoints. The Apple registry selects MLX-VLM for reviewed unquantized or MLX-format SmolVLM, Muse-Glimmer, and Qwen3.5 checkpoints.
|
||||
|
||||
Reasoning extraction is deliberately model-aware and backend-neutral:
|
||||
Reasoning extraction remains backend-neutral:
|
||||
|
||||
- A profile selects a top-level immutable `reasoning_parser`; known architectures provide a compatibility fallback for previously generated profiles.
|
||||
- A profile selects an immutable `reasoning_parser`; known architectures provide a compatibility fallback for older profiles.
|
||||
- Each backend supplies raw generated token IDs and any native separated reasoning field to the same parser contract.
|
||||
- A native field wins when an engine supplies one. This supports future engine APIs without changing the node contract.
|
||||
- Transformers generation has no standard separated reasoning result, so reasoning models retain boundary tokens until the parser splits the token sequence and decodes each channel with special tokens removed.
|
||||
- Resident vLLM uses one persistent `AsyncLLM` event loop per loaded backend. Its cumulative async generation stream exposes token IDs and text but not the OpenAI server's parsed reasoning response, so the same token-aware fallback applies.
|
||||
- Qwen3.5 chat templating explicitly enables thinking, and its parser treats tokens before `</think>` as reasoning. If generation stops before that terminator, the entire partial output is reasoning and `response` remains empty.
|
||||
- Transformers and MLX retain reasoning boundary tokens until parsing; resident vLLM uses its cumulative async token stream.
|
||||
- The node appends `reasoning` after the pre-existing `response` and `metadata_json` outputs, preserving saved workflow link indices.
|
||||
- `output_tokens` counts all generated tokens and `reasoning_tokens` counts reasoning content without boundary markers.
|
||||
|
||||
The node appends `reasoning` after the pre-existing `response` and `metadata_json` outputs to preserve saved workflow link indices. `output_tokens` remains the total generated count; `reasoning_tokens` counts reasoning content tokens excluding boundary markers.
|
||||
## Progress And Cancellation
|
||||
|
||||
## Progress And Streaming
|
||||
Both staging paths pass a custom progress adapter into Hugging Face snapshot download. Local and remote execution emit the same typed stages for profile inspection, input preparation, memory admission, processor/model loading, and token generation. The standard ComfyUI progress bar therefore works locally, while the existing streamed remote event path continues to feed Modal overlays. Work without a real completion fraction remains indeterminate.
|
||||
|
||||
The CPU ModelStager passes a custom `tqdm_class` into Hugging Face `snapshot_download` and streams aggregate file progress before GPU allocation. GPU execution emits the same typed progress contract for profile resolution, multimodal input preparation, GPU-memory admission, processor load, known safetensor shard totals, engine initialization, prefill, and decoding. Work without an upstream completion fraction uses an indeterminate pulse instead of a fabricated percentage.
|
||||
Each generated token checks ComfyUI's interruption hook through the node progress callback. Remote vLLM also aborts its active async request before propagating cancellation. Final time-to-first-token and throughput are retained in `metadata_json`.
|
||||
|
||||
The vLLM backend keeps `AsyncLLM` and its output task attached to a persistent event-loop thread for the lifetime of the resident model. Each cumulative request output updates generated token count, time to first token, elapsed time, and live tokens per second. Cancellation raised by ComfyUI aborts the active vLLM request before propagating. Final TTFT and throughput are retained in `metadata_json`.
|
||||
## Live Validation
|
||||
|
||||
Every selectable NVIDIA GPU runtime uses Python 3.13 and pins Torch 2.13.0, torchvision 0.28.0, Transformers 5.15.0, and the CUDA 13.0 vLLM 0.27.1 wheel. This covers Modal's T4-through-Blackwell choices and ensures a workflow's GPU price/capacity selection never removes its LLM backend. Python 3.13 is necessary because a FlashInfer communications module imported by vLLM's kernel warmup evaluates `array.array[int]` at runtime. The initial vLLM policy uses a 12 GiB BF16 KV cache, one request at a time, and eager execution to keep graph-capture allocations predictable beside ComfyUI. The worker sets `VLLM_USE_FLASHINFER_SAMPLER=0` and selects Triton full attention so optional FlashInfer sampling and Blackwell TRT-LLM attention kernels do not require an `nvcc` JIT compiler, while avoiding FlashAttention 4's unsupported Qwen3.5 used-sequence prefill at head dimension 256. Qwen's separate GDN kernel remains available. Torch and vLLM imports are validated while building every image; model-specific capacity remains enforced separately from backend availability.
|
||||
On 2026-08-19, Apple-local validation on a 64 GiB Mac loaded the exact `mlx-community/Qwen3.5-2B-4bit` revision `674aaa7240b91e8012fcad5d791b7dfe5ba90207`, produced the requested text response, accepted a real `[1,32,32,3]` ComfyUI IMAGE tensor and identified its dominant red colour, then reported a resident cache hit on the second request. The text pass generated at approximately 193 tokens per second. This validates one representative 4-bit MLX model, not every model that fits in unified memory.
|
||||
|
||||
## Live GPU Validation
|
||||
|
||||
On 2026-08-19, the generated-profile canary resolved exact Hub revisions, staged immutable snapshots, loaded each model, accepted a real ComfyUI IMAGE input, and generated non-empty output for:
|
||||
|
||||
- `orcarouter/Qwen3.8-27B-Uncensored-FP8` at `9228df5c6c9c509e1019f83b4e085cf643118bac` through vLLM FP8;
|
||||
- `meta-models/Muse-Glimmer-30B` at `a4e59da52a7bc87ae7251dd5545c0dd437c44b68` through Transformers;
|
||||
- `Blackfrost-AI/Qwen3.8-27B-ABLITERATED-NVFP4` at `faf7945020c138c8ef864ab1644273f3158f85fa` through vLLM ModelOpt FP4.
|
||||
|
||||
The live co-residency canary can select a curated profile or Hub ID through `COMFY_MODAL_LLM_CANARY_PROFILE`. It retains that LLM, executes a real ComfyUI VAE encode on the same single B300 worker, and then requires the second LLM call to report a resident cache hit. The 2026-08-19 run passed with the generated Orcarouter FP8 profile retained across a real Flux VAE encode on one B300 worker.
|
||||
|
||||
The same date's isolated RTX PRO 6000 canary rebuilt the universal image from its custom-node requirements, validated the final pinned dependency versions, and passed both the remote runtime handshake and real generated-profile inference for `Blackfrost-AI/Qwen3.8-27B-ABLITERATED-NVFP4`. This specifically covers the non-B300 path that previously omitted vLLM. It does not imply that every model fits or supports every selectable GPU; profile compatibility and measured-VRAM admission remain independent checks.
|
||||
The same date's Modal canaries validated generated vLLM FP8, Transformers, and vLLM ModelOpt FP4 profiles on B300 and RTX PRO 6000 workers. Model-specific capacity and compatibility remain separate from backend availability on both targets.
|
||||
|
||||
+46
-4
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import logging
|
||||
from typing import Any, Mapping
|
||||
from typing import Any, Literal, Mapping
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -13,6 +13,8 @@ LLM_COMPATIBILITY_POLICY_VERSION = 1
|
||||
TRANSFORMERS_RUNTIME_VERSION = "5.15.0"
|
||||
VLLM_RUNTIME_VERSION = "0.27.1"
|
||||
VLLM_TORCH_VERSION = "2.13.0"
|
||||
LOCAL_MLX_VLM_VERSION = "0.6.15"
|
||||
LLMExecutionTarget = Literal["modal", "local_apple"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -84,6 +86,11 @@ def _quantization_method(config: Mapping[str, Any]) -> str:
|
||||
return "modelopt_fp4"
|
||||
if method == "fp8":
|
||||
return "fp8"
|
||||
mlx_quantization = _configuration_mapping(config, "quantization")
|
||||
mlx_bits = mlx_quantization.get("bits")
|
||||
mlx_mode = str(mlx_quantization.get("mode") or "").strip().lower()
|
||||
if mlx_bits:
|
||||
return f"mlx_{mlx_mode or 'quantized'}_{int(mlx_bits)}bit"
|
||||
return method or "none"
|
||||
|
||||
|
||||
@@ -92,7 +99,10 @@ def _estimated_vram_gb(
|
||||
) -> float:
|
||||
"""Estimate model plus runtime allocations before the first measured load."""
|
||||
artifact_gib = artifact_bytes / 1024**3
|
||||
runtime_overhead_gib = 12.0 if backend == "vllm" else 8.0
|
||||
if backend == "mlx_vlm":
|
||||
runtime_overhead_gib = 2.0
|
||||
else:
|
||||
runtime_overhead_gib = 12.0 if backend == "vllm" else 8.0
|
||||
multiplier = 1.12 if quantization_method != "none" else 1.08
|
||||
return round(artifact_gib * multiplier + runtime_overhead_gib, 1)
|
||||
|
||||
@@ -101,6 +111,7 @@ def resolve_compatibility(
|
||||
config: Mapping[str, Any],
|
||||
*,
|
||||
artifact_bytes: int,
|
||||
execution_target: LLMExecutionTarget = "modal",
|
||||
) -> LLMCompatibilityDecision:
|
||||
"""Select a reviewed backend or reject an unknown model before weight download."""
|
||||
architecture = _architecture(config)
|
||||
@@ -109,7 +120,35 @@ def resolve_compatibility(
|
||||
default_context_tokens = min(advertised_context_tokens, 32768)
|
||||
quantization_method = _quantization_method(config)
|
||||
|
||||
if architecture == "MuseGlimmerForConditionalGeneration":
|
||||
if execution_target == "local_apple":
|
||||
if quantization_method != "none" and not quantization_method.startswith(
|
||||
"mlx_"
|
||||
):
|
||||
raise ValueError(
|
||||
f"Quantization {quantization_method!r} is not an MLX checkpoint "
|
||||
"format supported by the Apple-local LLM policy. Choose an "
|
||||
"unquantized or mlx-community conversion of this model."
|
||||
)
|
||||
if architecture == "SmolVLMForConditionalGeneration":
|
||||
modalities = frozenset({"text", "image", "video", "file"})
|
||||
reasoning_parser = "none"
|
||||
elif architecture == "MuseGlimmerForConditionalGeneration":
|
||||
modalities = frozenset({"text", "image", "file"})
|
||||
reasoning_parser = "none"
|
||||
elif architecture == "Qwen3_5ForConditionalGeneration":
|
||||
modalities = frozenset({"text", "image", "video", "file"})
|
||||
reasoning_parser = "qwen3"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Architecture {architecture!r} is not supported by the Apple-local "
|
||||
"LLM compatibility policy "
|
||||
f"v{LLM_COMPATIBILITY_POLICY_VERSION}. Add and validate an MLX-VLM "
|
||||
"adapter before downloading its weights."
|
||||
)
|
||||
backend = "mlx_vlm"
|
||||
backend_options = ()
|
||||
requirements = (f"mlx-vlm=={LOCAL_MLX_VLM_VERSION}",)
|
||||
elif architecture == "MuseGlimmerForConditionalGeneration":
|
||||
if quantization_method != "none":
|
||||
raise ValueError(
|
||||
"Muse-Glimmer is currently validated only for its unquantized "
|
||||
@@ -171,8 +210,9 @@ def resolve_compatibility(
|
||||
runtime_requirements=requirements,
|
||||
)
|
||||
logger.info(
|
||||
"Resolved Modal LLM compatibility architecture=%s backend=%s "
|
||||
"Resolved LLM compatibility target=%s architecture=%s backend=%s "
|
||||
"quantization=%s policy=%d.",
|
||||
execution_target,
|
||||
decision.architecture,
|
||||
decision.backend,
|
||||
decision.quantization_method,
|
||||
@@ -185,6 +225,8 @@ __all__ = [
|
||||
"LLM_COMPATIBILITY_POLICY_VERSION",
|
||||
"LLM_PROFILE_SCHEMA_VERSION",
|
||||
"LLMCompatibilityDecision",
|
||||
"LLMExecutionTarget",
|
||||
"LOCAL_MLX_VLM_VERSION",
|
||||
"TRANSFORMERS_RUNTIME_VERSION",
|
||||
"VLLM_RUNTIME_VERSION",
|
||||
"VLLM_TORCH_VERSION",
|
||||
|
||||
+46
-3
@@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from functools import lru_cache
|
||||
import json
|
||||
import logging
|
||||
@@ -14,11 +14,15 @@ if __package__:
|
||||
from .llm_compatibility import (
|
||||
LLM_COMPATIBILITY_POLICY_VERSION,
|
||||
LLM_PROFILE_SCHEMA_VERSION,
|
||||
LOCAL_MLX_VLM_VERSION,
|
||||
LLMExecutionTarget,
|
||||
)
|
||||
else: # pragma: no cover - stable cloud entrypoint imports top-level modules.
|
||||
from llm_compatibility import (
|
||||
LLM_COMPATIBILITY_POLICY_VERSION,
|
||||
LLM_PROFILE_SCHEMA_VERSION,
|
||||
LOCAL_MLX_VLM_VERSION,
|
||||
LLMExecutionTarget,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -31,7 +35,8 @@ _PROFILE_ID_PATTERN = re.compile(r"^[a-z0-9][a-z0-9._-]*$")
|
||||
_REVISION_PATTERN = re.compile(r"^[0-9a-f]{40}$")
|
||||
_SUPPORTED_DTYPES = frozenset({"bfloat16", "float16", "float32"})
|
||||
_SUPPORTED_MODALITIES = frozenset({"text", "image", "video", "file"})
|
||||
_SUPPORTED_BACKENDS = frozenset({"transformers", "vllm"})
|
||||
_SUPPORTED_BACKENDS = frozenset({"transformers", "vllm", "mlx_vlm"})
|
||||
_SUPPORTED_EXECUTION_TARGETS = frozenset({"modal", "local_apple"})
|
||||
_SUPPORTED_REASONING_PARSERS = frozenset({"", "none", "qwen3"})
|
||||
_PROFILE_DIGEST_PATTERN = re.compile(r"^[0-9a-f]{64}$")
|
||||
|
||||
@@ -66,6 +71,7 @@ class LLMModelProfile:
|
||||
reasoning_parser: str = ""
|
||||
backend_options: tuple[tuple[str, str | int | float | bool], ...] = ()
|
||||
runtime_requirements: tuple[str, ...] = ()
|
||||
execution_target: str = "modal"
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, value: Mapping[str, Any]) -> "LLMModelProfile":
|
||||
@@ -139,6 +145,9 @@ class LLMModelProfile:
|
||||
str(requirement)
|
||||
for requirement in value.get("runtime_requirements", [])
|
||||
),
|
||||
execution_target=str(value.get("execution_target", "modal"))
|
||||
.strip()
|
||||
.lower(),
|
||||
)
|
||||
profile.validate()
|
||||
return profile
|
||||
@@ -160,6 +169,11 @@ class LLMModelProfile:
|
||||
f"Modal LLM profile {self.profile_id!r} has unsupported backend "
|
||||
f"{self.backend!r}."
|
||||
)
|
||||
if self.execution_target not in _SUPPORTED_EXECUTION_TARGETS:
|
||||
raise ValueError(
|
||||
f"Modal LLM profile {self.profile_id!r} has unsupported execution "
|
||||
f"target {self.execution_target!r}."
|
||||
)
|
||||
if self.reasoning_parser not in _SUPPORTED_REASONING_PARSERS:
|
||||
raise ValueError(
|
||||
f"Modal LLM profile {self.profile_id!r} has unsupported reasoning "
|
||||
@@ -218,7 +232,7 @@ class LLMModelProfile:
|
||||
|
||||
def to_mapping(self) -> dict[str, Any]:
|
||||
"""Return the stable JSON representation used by generated manifests."""
|
||||
return {
|
||||
mapping = {
|
||||
"id": self.profile_id,
|
||||
"display_name": self.display_name,
|
||||
"repository": self.repository,
|
||||
@@ -246,6 +260,34 @@ class LLMModelProfile:
|
||||
"backend_options": dict(self.backend_options),
|
||||
"runtime_requirements": list(self.runtime_requirements),
|
||||
}
|
||||
if self.execution_target != "modal":
|
||||
mapping["execution_target"] = self.execution_target
|
||||
return mapping
|
||||
|
||||
|
||||
def profile_for_execution_target(
|
||||
profile: LLMModelProfile,
|
||||
execution_target: LLMExecutionTarget,
|
||||
) -> LLMModelProfile:
|
||||
"""Return a curated profile adapted to one execution target."""
|
||||
if profile.source == "generated":
|
||||
if profile.execution_target != execution_target:
|
||||
raise ValueError(
|
||||
f"Generated LLM profile {profile.profile_id!r} targets "
|
||||
f"{profile.execution_target!r}, not {execution_target!r}. Resolve the "
|
||||
"original Hugging Face model reference for this execution target."
|
||||
)
|
||||
return profile
|
||||
if execution_target == "modal":
|
||||
return profile
|
||||
adapted = replace(
|
||||
profile,
|
||||
backend="mlx_vlm",
|
||||
execution_target="local_apple",
|
||||
runtime_requirements=(f"mlx-vlm=={LOCAL_MLX_VLM_VERSION}",),
|
||||
)
|
||||
adapted.validate()
|
||||
return adapted
|
||||
|
||||
|
||||
def _mapping_value(value: Mapping[str, Any], name: str) -> Mapping[str, Any]:
|
||||
@@ -449,5 +491,6 @@ __all__ = [
|
||||
"llm_profile_ids_from_payload",
|
||||
"llm_profile_options",
|
||||
"load_llm_profiles",
|
||||
"profile_for_execution_target",
|
||||
"rewrite_llm_model_references",
|
||||
]
|
||||
|
||||
+26
-6
@@ -18,6 +18,7 @@ if __package__:
|
||||
from .llm_compatibility import (
|
||||
LLM_COMPATIBILITY_POLICY_VERSION,
|
||||
LLM_PROFILE_SCHEMA_VERSION,
|
||||
LLMExecutionTarget,
|
||||
resolve_compatibility,
|
||||
)
|
||||
from .llm_profiles import (
|
||||
@@ -29,6 +30,7 @@ else: # pragma: no cover - stable cloud entrypoint imports top-level modules.
|
||||
from llm_compatibility import (
|
||||
LLM_COMPATIBILITY_POLICY_VERSION,
|
||||
LLM_PROFILE_SCHEMA_VERSION,
|
||||
LLMExecutionTarget,
|
||||
resolve_compatibility,
|
||||
)
|
||||
from llm_profiles import (
|
||||
@@ -106,6 +108,13 @@ def _token() -> str | None:
|
||||
return os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
|
||||
|
||||
|
||||
def _hugging_face_token_guidance(execution_target: LLMExecutionTarget) -> str:
|
||||
"""Describe where to configure a Hugging Face token for one target."""
|
||||
if execution_target == "local_apple":
|
||||
return "set HF_TOKEN in the local ComfyUI environment"
|
||||
return "add HF_TOKEN to the Modal Secret"
|
||||
|
||||
|
||||
def _repository_files(model_info: Any) -> tuple[tuple[str, int], ...]:
|
||||
"""Return stable repository filenames and known sizes from model metadata."""
|
||||
siblings = _field(model_info, "siblings", ()) or ()
|
||||
@@ -206,9 +215,14 @@ def _generated_profile(
|
||||
revision: str,
|
||||
config: Mapping[str, Any],
|
||||
artifact_bytes: int,
|
||||
execution_target: LLMExecutionTarget,
|
||||
) -> LLMModelProfile:
|
||||
"""Build a content-addressed profile from reviewed compatibility metadata."""
|
||||
decision = resolve_compatibility(config, artifact_bytes=artifact_bytes)
|
||||
decision = resolve_compatibility(
|
||||
config,
|
||||
artifact_bytes=artifact_bytes,
|
||||
execution_target=execution_target,
|
||||
)
|
||||
identity: dict[str, Any] = {
|
||||
"schema_version": LLM_PROFILE_SCHEMA_VERSION,
|
||||
"compatibility_policy_version": LLM_COMPATIBILITY_POLICY_VERSION,
|
||||
@@ -227,6 +241,8 @@ def _generated_profile(
|
||||
"backend_options": dict(decision.backend_options),
|
||||
"runtime_requirements": list(decision.runtime_requirements),
|
||||
}
|
||||
if execution_target != "modal":
|
||||
identity["execution_target"] = execution_target
|
||||
digest = hashlib.sha256(_canonical_profile_identity(identity)).hexdigest()
|
||||
profile_id = generated_profile_id(digest)
|
||||
return LLMModelProfile.from_mapping(
|
||||
@@ -291,6 +307,7 @@ def resolve_model_profile(
|
||||
api: HuggingFaceApi | None = None,
|
||||
hf_hub_download: Callable[..., str] | None = None,
|
||||
max_download_bytes: int | None = None,
|
||||
execution_target: LLMExecutionTarget = "modal",
|
||||
) -> ResolvedModelProfile:
|
||||
"""Inspect one Hugging Face ID and persist its immutable generated profile."""
|
||||
reference = HuggingFaceModelReference.parse(model_reference)
|
||||
@@ -309,10 +326,10 @@ def resolve_model_profile(
|
||||
token=token,
|
||||
)
|
||||
except (HfHubHTTPError, OSError, ValueError) as exc:
|
||||
token_guidance = _hugging_face_token_guidance(execution_target)
|
||||
raise ValueError(
|
||||
f"Unable to access Hugging Face model {reference.display()!r}. For "
|
||||
"gated or "
|
||||
"private models, add HF_TOKEN to the Modal Secret: "
|
||||
f"gated or private models, {token_guidance}: "
|
||||
f"{exc}"
|
||||
) from exc
|
||||
revision = str(_field(model_info, "sha", "")).lower()
|
||||
@@ -337,16 +354,18 @@ def resolve_model_profile(
|
||||
hf_hub_download=hf_hub_download,
|
||||
)
|
||||
except (HfHubHTTPError, OSError, ValueError) as exc:
|
||||
token_guidance = _hugging_face_token_guidance(execution_target)
|
||||
raise ValueError(
|
||||
f"Unable to download config.json for Hugging Face model "
|
||||
f"{reference.display()!r}. For gated or private models, add "
|
||||
f"HF_TOKEN to the Modal Secret: {exc}"
|
||||
f"{reference.display()!r}. For gated or private models, "
|
||||
f"{token_guidance}: {exc}"
|
||||
) from exc
|
||||
profile = _generated_profile(
|
||||
reference=reference,
|
||||
revision=revision,
|
||||
config=config,
|
||||
artifact_bytes=artifact_bytes,
|
||||
execution_target=execution_target,
|
||||
)
|
||||
manifest_path, manifest_created = _write_generated_manifest(
|
||||
storage_root,
|
||||
@@ -355,11 +374,12 @@ def resolve_model_profile(
|
||||
security_scan_complete=scan_complete,
|
||||
)
|
||||
logger.info(
|
||||
"Resolved Modal LLM model=%s revision=%s profile=%s backend=%s "
|
||||
"Resolved LLM model=%s revision=%s profile=%s target=%s backend=%s "
|
||||
"weights_gib=%.2f manifest_created=%s.",
|
||||
reference.display(),
|
||||
revision,
|
||||
profile.profile_id,
|
||||
execution_target,
|
||||
profile.backend,
|
||||
artifact_bytes / 1024**3,
|
||||
manifest_created,
|
||||
|
||||
@@ -0,0 +1,524 @@
|
||||
"""Apple Silicon local inference for the Modal LLM ComfyUI node."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
from importlib import metadata
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import platform
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable, Sequence
|
||||
|
||||
if __package__:
|
||||
from .llm_compatibility import LOCAL_MLX_VLM_VERSION
|
||||
from .llm_profiles import (
|
||||
LLMModelProfile,
|
||||
get_llm_profile,
|
||||
profile_for_execution_target,
|
||||
)
|
||||
from .llm_reasoning import (
|
||||
ReasoningOutputParser,
|
||||
create_reasoning_parser,
|
||||
reasoning_chat_template_kwargs,
|
||||
)
|
||||
from .llm_resolver import HuggingFaceModelReference, resolve_model_profile
|
||||
from .llm_staging import LLMStagingProgress, stage_model_profile
|
||||
from .modal_llm_runtime import (
|
||||
BackendGenerationResult,
|
||||
LLMGenerationSettings,
|
||||
LLMInferenceResult,
|
||||
LLMProgressCallback,
|
||||
LLMProgressEvent,
|
||||
PreparedLLMInputs,
|
||||
ResidentLLMManager,
|
||||
_coerce_positive_int,
|
||||
prepare_llm_inputs,
|
||||
)
|
||||
else: # pragma: no cover - ComfyUI loads this module as part of the package.
|
||||
from llm_compatibility import LOCAL_MLX_VLM_VERSION
|
||||
from llm_profiles import (
|
||||
LLMModelProfile,
|
||||
get_llm_profile,
|
||||
profile_for_execution_target,
|
||||
)
|
||||
from llm_reasoning import (
|
||||
ReasoningOutputParser,
|
||||
create_reasoning_parser,
|
||||
reasoning_chat_template_kwargs,
|
||||
)
|
||||
from llm_resolver import HuggingFaceModelReference, resolve_model_profile
|
||||
from llm_staging import LLMStagingProgress, stage_model_profile
|
||||
from modal_llm_runtime import (
|
||||
BackendGenerationResult,
|
||||
LLMGenerationSettings,
|
||||
LLMInferenceResult,
|
||||
LLMProgressCallback,
|
||||
LLMProgressEvent,
|
||||
PreparedLLMInputs,
|
||||
ResidentLLMManager,
|
||||
_coerce_positive_int,
|
||||
prepare_llm_inputs,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_BYTES_PER_GIB = 1024**3
|
||||
_LOCAL_STORAGE_ENV = "COMFY_MODAL_LOCAL_LLM_STORAGE_ROOT"
|
||||
_LOCAL_RESERVE_ENV = "COMFY_MODAL_LOCAL_LLM_RESERVE_FREE_GB"
|
||||
_LOCAL_RESIDENT_MODELS_ENV = "COMFY_MODAL_LOCAL_LLM_MAX_RESIDENT_MODELS"
|
||||
_DEFAULT_LOCAL_RESERVE_GB = 4.0
|
||||
_DEFAULT_LOCAL_RESIDENT_MODELS = 1
|
||||
LOCAL_MLX_VLM_SPEC = f"mlx-vlm=={LOCAL_MLX_VLM_VERSION}"
|
||||
|
||||
|
||||
def _local_install_command() -> str:
|
||||
"""Return the exact command that installs the validated local backend."""
|
||||
return (
|
||||
f"uv pip install --python {sys.executable!r} "
|
||||
f"{LOCAL_MLX_VLM_SPEC!r} 'psutil>=7,<8'"
|
||||
)
|
||||
|
||||
|
||||
def ensure_local_apple_runtime_available() -> None:
|
||||
"""Fail clearly unless the validated MLX-VLM runtime is usable."""
|
||||
system = platform.system()
|
||||
machine = platform.machine().lower()
|
||||
if system != "Darwin" or machine not in {"arm64", "aarch64"}:
|
||||
raise RuntimeError(
|
||||
"Local Modal LLM inference currently requires Apple Silicon macOS. "
|
||||
f"Detected {system} {machine}. Enable 'Run on Modal' on this node for "
|
||||
"the supported remote backend."
|
||||
)
|
||||
try:
|
||||
installed_version = metadata.version("mlx-vlm")
|
||||
except metadata.PackageNotFoundError as exc:
|
||||
raise RuntimeError(
|
||||
"Apple-local LLM inference requires the optional MLX-VLM backend. "
|
||||
f"Install it into ComfyUI's Python environment and restart ComfyUI: "
|
||||
f"{_local_install_command()}"
|
||||
) from exc
|
||||
if installed_version != LOCAL_MLX_VLM_VERSION:
|
||||
raise RuntimeError(
|
||||
f"Apple-local LLM inference requires {LOCAL_MLX_VLM_SPEC}, but "
|
||||
f"mlx-vlm=={installed_version} is installed. Run: "
|
||||
f"{_local_install_command()}"
|
||||
)
|
||||
try:
|
||||
__import__("mlx.core")
|
||||
__import__("mlx_vlm")
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name not in {"mlx", "mlx.core", "mlx_vlm"}:
|
||||
raise
|
||||
raise RuntimeError(
|
||||
"The Apple-local MLX-VLM backend is incomplete. Reinstall it and "
|
||||
f"restart ComfyUI: {_local_install_command()}"
|
||||
) from exc
|
||||
|
||||
|
||||
def local_llm_storage_root() -> Path:
|
||||
"""Return the writable ComfyUI model root for immutable local LLM snapshots."""
|
||||
configured = os.getenv(_LOCAL_STORAGE_ENV)
|
||||
if configured is not None:
|
||||
normalized = configured.strip()
|
||||
if not normalized:
|
||||
raise ValueError(f"{_LOCAL_STORAGE_ENV} must not be blank.")
|
||||
return Path(normalized).expanduser().resolve()
|
||||
try:
|
||||
import folder_paths
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name != "folder_paths":
|
||||
raise
|
||||
raise RuntimeError(
|
||||
f"Unable to locate ComfyUI's model directory. Set {_LOCAL_STORAGE_ENV} "
|
||||
"to a writable directory."
|
||||
) from exc
|
||||
default_path = Path(folder_paths.models_dir).resolve() / "modal_llm"
|
||||
folder_paths.add_model_folder_path(
|
||||
"modal_llm",
|
||||
str(default_path),
|
||||
is_default=True,
|
||||
)
|
||||
return Path(folder_paths.get_folder_paths("modal_llm")[0]).resolve()
|
||||
|
||||
|
||||
def _read_non_negative_float(name: str, default: float) -> float:
|
||||
"""Read one non-negative floating-point environment value."""
|
||||
raw_value = os.getenv(name)
|
||||
value = default if raw_value is None else float(raw_value)
|
||||
if value < 0:
|
||||
raise ValueError(f"{name} cannot be negative.")
|
||||
return value
|
||||
|
||||
|
||||
def _read_positive_int(name: str, default: int) -> int:
|
||||
"""Read one positive integer environment value."""
|
||||
raw_value = os.getenv(name)
|
||||
value = default if raw_value is None else int(raw_value)
|
||||
if value <= 0:
|
||||
raise ValueError(f"{name} must be positive.")
|
||||
return value
|
||||
|
||||
|
||||
def _apple_memory_info() -> tuple[int, int]:
|
||||
"""Return system-available and total unified memory in bytes."""
|
||||
import psutil
|
||||
|
||||
memory = psutil.virtual_memory()
|
||||
return int(memory.available), int(memory.total)
|
||||
|
||||
|
||||
def _clear_apple_caches() -> None:
|
||||
"""Release MLX cache buffers after a resident-model eviction."""
|
||||
import mlx.core as mx
|
||||
|
||||
mx.clear_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def _stage_progress_callback(
|
||||
progress_callback: LLMProgressCallback,
|
||||
) -> Callable[[LLMStagingProgress], None]:
|
||||
"""Translate storage progress into the shared node progress contract."""
|
||||
|
||||
def report(progress: LLMStagingProgress) -> None:
|
||||
"""Forward one local snapshot event."""
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage=progress.stage,
|
||||
message=progress.message,
|
||||
value=progress.value,
|
||||
maximum=progress.maximum,
|
||||
unit=progress.unit,
|
||||
indeterminate=progress.indeterminate,
|
||||
)
|
||||
)
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def resolve_and_stage_local_profile(
|
||||
model_reference: str,
|
||||
storage_root: str | Path,
|
||||
*,
|
||||
progress_callback: LLMProgressCallback,
|
||||
) -> LLMModelProfile:
|
||||
"""Resolve one local model reference and stage its pinned snapshot."""
|
||||
normalized_reference = model_reference.strip()
|
||||
if not normalized_reference:
|
||||
raise ValueError("Modal LLM model_profile must not be blank.")
|
||||
try:
|
||||
profile = get_llm_profile(normalized_reference, storage_root=storage_root)
|
||||
except ValueError:
|
||||
if normalized_reference.startswith("hf-"):
|
||||
raise
|
||||
HuggingFaceModelReference.parse(normalized_reference)
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="profile",
|
||||
message="Inspecting Apple-local model compatibility",
|
||||
indeterminate=True,
|
||||
)
|
||||
)
|
||||
profile = resolve_model_profile(
|
||||
normalized_reference,
|
||||
storage_root,
|
||||
execution_target="local_apple",
|
||||
).profile
|
||||
else:
|
||||
profile = profile_for_execution_target(profile, "local_apple")
|
||||
stage_model_profile(
|
||||
profile.profile_id,
|
||||
storage_root,
|
||||
profile=profile,
|
||||
progress_callback=_stage_progress_callback(progress_callback),
|
||||
)
|
||||
return profile
|
||||
|
||||
|
||||
class MLXVLMBackend:
|
||||
"""Run a pinned multimodal model locally through MLX-VLM."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
profile: LLMModelProfile,
|
||||
snapshot_path: Path,
|
||||
progress_callback: LLMProgressCallback,
|
||||
) -> None:
|
||||
"""Load processor and weights from one immutable local snapshot."""
|
||||
ensure_local_apple_runtime_available()
|
||||
import mlx.core as mx
|
||||
from mlx_vlm import load
|
||||
|
||||
if profile.backend != "mlx_vlm":
|
||||
raise ValueError(
|
||||
f"Apple-local profile {profile.profile_id!r} selects backend "
|
||||
f"{profile.backend!r}, not 'mlx_vlm'."
|
||||
)
|
||||
self.profile = profile
|
||||
self.snapshot_path = snapshot_path
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="processor",
|
||||
message="Loading MLX processor",
|
||||
indeterminate=True,
|
||||
)
|
||||
)
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="weights",
|
||||
message="Loading model into unified memory",
|
||||
indeterminate=True,
|
||||
)
|
||||
)
|
||||
mx.reset_peak_memory()
|
||||
self.model, self.processor = load(
|
||||
str(snapshot_path),
|
||||
lazy=False,
|
||||
strict=True,
|
||||
)
|
||||
tokenizer = getattr(self.processor, "tokenizer", self.processor)
|
||||
self.reasoning_parser: ReasoningOutputParser = create_reasoning_parser(
|
||||
profile,
|
||||
tokenizer,
|
||||
)
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="weights",
|
||||
message="MLX model ready",
|
||||
value=1,
|
||||
maximum=1,
|
||||
unit="model",
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _media(prepared_inputs: PreparedLLMInputs) -> list[Any]:
|
||||
"""Return PIL frames that MLX-VLM should treat as image inputs."""
|
||||
media = list(prepared_inputs.images)
|
||||
if prepared_inputs.video is not None:
|
||||
media.extend(prepared_inputs.video.frames)
|
||||
return media
|
||||
|
||||
@staticmethod
|
||||
def _messages(prepared_inputs: PreparedLLMInputs) -> list[dict[str, str]]:
|
||||
"""Build system and user messages without embedding binary media."""
|
||||
messages: list[dict[str, str]] = []
|
||||
if prepared_inputs.system_prompt:
|
||||
messages.append(
|
||||
{"role": "system", "content": prepared_inputs.system_prompt}
|
||||
)
|
||||
messages.append({"role": "user", "content": prepared_inputs.prompt})
|
||||
return messages
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prepared_inputs: PreparedLLMInputs,
|
||||
settings: LLMGenerationSettings,
|
||||
progress_callback: LLMProgressCallback,
|
||||
) -> BackendGenerationResult:
|
||||
"""Stream one MLX generation while preserving reasoning token IDs."""
|
||||
from mlx_vlm.generate import stream_generate
|
||||
from mlx_vlm.prompt_utils import apply_chat_template
|
||||
|
||||
media = self._media(prepared_inputs)
|
||||
thinking_options = reasoning_chat_template_kwargs(self.profile)
|
||||
generation_options: dict[str, bool | int] = dict(thinking_options)
|
||||
if thinking_options.get("enable_thinking"):
|
||||
generation_options["thinking_budget"] = settings.max_new_tokens // 2
|
||||
prompt = apply_chat_template(
|
||||
self.processor,
|
||||
self.model.config,
|
||||
self._messages(prepared_inputs),
|
||||
num_images=len(media),
|
||||
**thinking_options,
|
||||
)
|
||||
started_at = time.perf_counter()
|
||||
first_token_at: float | None = None
|
||||
raw_segments: list[str] = []
|
||||
token_ids: list[int] = []
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
tokens_per_second: float | None = None
|
||||
for generation in stream_generate(
|
||||
self.model,
|
||||
self.processor,
|
||||
prompt,
|
||||
image=media or None,
|
||||
max_tokens=settings.max_new_tokens,
|
||||
temperature=settings.temperature,
|
||||
top_p=settings.top_p,
|
||||
seed=settings.seed,
|
||||
skip_special_tokens=False,
|
||||
**generation_options,
|
||||
):
|
||||
now = time.perf_counter()
|
||||
input_tokens = int(generation.prompt_tokens)
|
||||
output_tokens = int(generation.generation_tokens)
|
||||
if input_tokens + settings.max_new_tokens > self.profile.max_context_tokens:
|
||||
raise ValueError(
|
||||
f"Modal LLM request requires up to "
|
||||
f"{input_tokens + settings.max_new_tokens} tokens; profile "
|
||||
f"{self.profile.profile_id!r} is capped at "
|
||||
f"{self.profile.max_context_tokens}."
|
||||
)
|
||||
if generation.token is not None:
|
||||
if first_token_at is None:
|
||||
first_token_at = now
|
||||
token = generation.token
|
||||
token_ids.append(int(token.item() if hasattr(token, "item") else token))
|
||||
raw_segments.append(str(generation.text or ""))
|
||||
tokens_per_second = float(generation.generation_tps or 0.0) or None
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="generating",
|
||||
message="Generating locally with MLX",
|
||||
value=output_tokens,
|
||||
maximum=settings.max_new_tokens,
|
||||
unit="tokens",
|
||||
elapsed_seconds=now - started_at,
|
||||
time_to_first_token_seconds=(
|
||||
first_token_at - started_at
|
||||
if first_token_at is not None
|
||||
else None
|
||||
),
|
||||
tokens_per_second=tokens_per_second,
|
||||
)
|
||||
)
|
||||
raw_text = "".join(raw_segments).strip()
|
||||
reasoning_output = self.reasoning_parser.extract(raw_text, token_ids)
|
||||
return BackendGenerationResult(
|
||||
text=reasoning_output.response,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
reasoning=reasoning_output.reasoning,
|
||||
reasoning_tokens=reasoning_output.reasoning_tokens,
|
||||
reasoning_parser=reasoning_output.parser,
|
||||
time_to_first_token_seconds=(
|
||||
first_token_at - started_at if first_token_at is not None else None
|
||||
),
|
||||
tokens_per_second=tokens_per_second,
|
||||
)
|
||||
|
||||
def unload(self) -> None:
|
||||
"""Release model references and reclaim MLX cache buffers."""
|
||||
model = self.model
|
||||
processor = self.processor
|
||||
self.model = None
|
||||
self.processor = None
|
||||
del model, processor
|
||||
_clear_apple_caches()
|
||||
|
||||
|
||||
_LOCAL_MANAGER: ResidentLLMManager | None = None
|
||||
_LOCAL_MANAGER_ROOT: Path | None = None
|
||||
_LOCAL_MANAGER_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def get_local_resident_llm_manager(storage_root: str | Path) -> ResidentLLMManager:
|
||||
"""Return the process-global Apple-local resident model manager."""
|
||||
global _LOCAL_MANAGER, _LOCAL_MANAGER_ROOT
|
||||
resolved_root = Path(storage_root).resolve()
|
||||
with _LOCAL_MANAGER_LOCK:
|
||||
if _LOCAL_MANAGER is None or _LOCAL_MANAGER_ROOT != resolved_root:
|
||||
if _LOCAL_MANAGER is not None:
|
||||
_LOCAL_MANAGER.unload_all()
|
||||
_LOCAL_MANAGER = ResidentLLMManager(
|
||||
storage_root=resolved_root,
|
||||
backend_factory=MLXVLMBackend,
|
||||
max_resident_models=_read_positive_int(
|
||||
_LOCAL_RESIDENT_MODELS_ENV,
|
||||
_DEFAULT_LOCAL_RESIDENT_MODELS,
|
||||
),
|
||||
memory_info=_apple_memory_info,
|
||||
empty_cache=_clear_apple_caches,
|
||||
execution_target="local_apple",
|
||||
device_name="metal",
|
||||
memory_label="unified memory",
|
||||
)
|
||||
_LOCAL_MANAGER_ROOT = resolved_root
|
||||
return _LOCAL_MANAGER
|
||||
|
||||
|
||||
def run_local_llm_inference(
|
||||
*,
|
||||
prompt: str,
|
||||
model_profile: str,
|
||||
images: Any | None,
|
||||
video: Any | None,
|
||||
files: Sequence[Any] | None,
|
||||
system_prompt: str,
|
||||
max_new_tokens: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
seed: int,
|
||||
video_frames: int,
|
||||
reserve_free_vram_gb: float | None,
|
||||
keep_model_loaded: bool,
|
||||
progress_callback: LLMProgressCallback,
|
||||
) -> LLMInferenceResult:
|
||||
"""Resolve, stage, and execute one Apple-local LLM request."""
|
||||
ensure_local_apple_runtime_available()
|
||||
storage_root = local_llm_storage_root()
|
||||
profile = resolve_and_stage_local_profile(
|
||||
model_profile,
|
||||
storage_root,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
settings = LLMGenerationSettings(
|
||||
max_new_tokens=_coerce_positive_int(
|
||||
max_new_tokens,
|
||||
"max_new_tokens",
|
||||
32768,
|
||||
),
|
||||
temperature=float(temperature),
|
||||
top_p=float(top_p),
|
||||
seed=int(seed),
|
||||
)
|
||||
if not 0.0 <= settings.temperature <= 2.0:
|
||||
raise ValueError("temperature must be between 0.0 and 2.0.")
|
||||
if not 0.0 < settings.top_p <= 1.0:
|
||||
raise ValueError("top_p must be greater than 0.0 and at most 1.0.")
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="inputs",
|
||||
message="Preparing local multimodal inputs",
|
||||
indeterminate=True,
|
||||
)
|
||||
)
|
||||
prepared_inputs = prepare_llm_inputs(
|
||||
prompt=prompt,
|
||||
system_prompt=system_prompt,
|
||||
images=images,
|
||||
video=video,
|
||||
files=files,
|
||||
video_frames=video_frames,
|
||||
profile=profile,
|
||||
)
|
||||
reserve_gb = (
|
||||
_read_non_negative_float(_LOCAL_RESERVE_ENV, _DEFAULT_LOCAL_RESERVE_GB)
|
||||
if reserve_free_vram_gb is None
|
||||
else float(reserve_free_vram_gb)
|
||||
)
|
||||
return get_local_resident_llm_manager(storage_root).infer(
|
||||
profile=profile,
|
||||
prepared_inputs=prepared_inputs,
|
||||
generation_settings=settings,
|
||||
reserve_free_vram_gb=reserve_gb,
|
||||
keep_model_loaded=keep_model_loaded,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LOCAL_MLX_VLM_SPEC",
|
||||
"MLXVLMBackend",
|
||||
"ensure_local_apple_runtime_available",
|
||||
"get_local_resident_llm_manager",
|
||||
"local_llm_storage_root",
|
||||
"resolve_and_stage_local_profile",
|
||||
"run_local_llm_inference",
|
||||
]
|
||||
+38
-14
@@ -12,9 +12,11 @@ import comfy.utils
|
||||
from comfy_api.latest import io
|
||||
|
||||
if __package__:
|
||||
from .local_llm_runtime import run_local_llm_inference
|
||||
from .llm_profiles import MODAL_LLM_NODE_ID, llm_profile_options
|
||||
from .modal_llm_runtime import LLMProgressEvent, run_modal_llm_inference
|
||||
else: # pragma: no cover - the stable remote runtime imports this module top-level.
|
||||
from local_llm_runtime import run_local_llm_inference
|
||||
from llm_profiles import MODAL_LLM_NODE_ID, llm_profile_options
|
||||
from modal_llm_runtime import LLMProgressEvent, run_modal_llm_inference
|
||||
|
||||
@@ -35,7 +37,7 @@ def _modal_llm_primary_inputs() -> list[io.Input]:
|
||||
default=llm_profile_options()[0],
|
||||
tooltip=(
|
||||
"A curated profile or Hugging Face ID such as owner/model. The first run "
|
||||
"inspects, pins, and stages a compatible model on the Modal Volume."
|
||||
"inspects and pins a compatible model for the selected local or Modal target."
|
||||
),
|
||||
),
|
||||
io.Image.Input(
|
||||
@@ -118,19 +120,34 @@ def _modal_llm_advanced_inputs() -> list[io.Input]:
|
||||
max=256.0,
|
||||
step=1.0,
|
||||
advanced=True,
|
||||
tooltip="VRAM kept free for ComfyUI image/video models before loading this LLM.",
|
||||
tooltip=(
|
||||
"Accelerator or unified memory kept free for ComfyUI and the operating "
|
||||
"system before loading this LLM."
|
||||
),
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"keep_model_loaded",
|
||||
default=True,
|
||||
advanced=True,
|
||||
tooltip="Keep the model resident for subsequent nodes in this warm Modal container.",
|
||||
tooltip="Keep the model resident for subsequent requests in this process.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"local_reserve_free_memory_gb",
|
||||
default=4.0,
|
||||
min=0.0,
|
||||
max=256.0,
|
||||
step=1.0,
|
||||
advanced=True,
|
||||
tooltip=(
|
||||
"Apple unified memory kept free for ComfyUI and macOS before loading "
|
||||
"a local LLM. Ignored when Run on Modal is enabled."
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class ModalLLM(io.ComfyNode):
|
||||
"""Generate text with a resident multimodal Transformers model on Modal."""
|
||||
"""Generate text with a resident local or Modal multimodal model."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
@@ -142,8 +159,8 @@ class ModalLLM(io.ComfyNode):
|
||||
essentials_category="Text Generation",
|
||||
description=(
|
||||
"Run text, image, video, and bounded file understanding through a "
|
||||
"revision-pinned language model resident beside ComfyUI on the Modal GPU. "
|
||||
"Enable Run on Modal for this node."
|
||||
"revision-pinned language model. Leave Run on Modal disabled for "
|
||||
"Apple-local MLX inference, or enable it for a Modal GPU worker."
|
||||
),
|
||||
inputs=[*_modal_llm_primary_inputs(), *_modal_llm_advanced_inputs()],
|
||||
outputs=[
|
||||
@@ -172,14 +189,15 @@ class ModalLLM(io.ComfyNode):
|
||||
video_frames: int = 12,
|
||||
reserve_free_vram_gb: float = 24.0,
|
||||
keep_model_loaded: bool = True,
|
||||
local_reserve_free_memory_gb: float = 4.0,
|
||||
unique_id: str | None = None,
|
||||
) -> io.NodeOutput:
|
||||
"""Run one cancellation-aware resident inference request."""
|
||||
if os.getenv("COMFY_MODAL_REMOTE_WORKER") != "1":
|
||||
raise RuntimeError(
|
||||
"Modal LLM only runs inside the Modal worker. Enable 'Run on Modal' for this "
|
||||
"node and choose a GPU workflow target."
|
||||
)
|
||||
is_remote = os.getenv("COMFY_MODAL_REMOTE_WORKER") == "1"
|
||||
execution_target = "modal" if is_remote else "local_apple"
|
||||
inference_runner = (
|
||||
run_modal_llm_inference if is_remote else run_local_llm_inference
|
||||
)
|
||||
progress_bar = comfy.utils.ProgressBar(max_new_tokens, node_id=unique_id)
|
||||
|
||||
def report_progress(progress: LLMProgressEvent) -> None:
|
||||
@@ -219,14 +237,16 @@ class ModalLLM(io.ComfyNode):
|
||||
PromptServer.instance.send_sync("modal_llm_progress", payload, None)
|
||||
|
||||
logger.info(
|
||||
"Starting Modal LLM profile=%s max_new_tokens=%d images=%s video=%s files=%d.",
|
||||
"Starting LLM target=%s profile=%s max_new_tokens=%d images=%s "
|
||||
"video=%s files=%d.",
|
||||
execution_target,
|
||||
model_profile,
|
||||
max_new_tokens,
|
||||
getattr(images, "shape", None),
|
||||
video is not None,
|
||||
len(files or ()),
|
||||
)
|
||||
result = run_modal_llm_inference(
|
||||
result = inference_runner(
|
||||
prompt=prompt,
|
||||
model_profile=model_profile,
|
||||
images=images,
|
||||
@@ -238,7 +258,11 @@ class ModalLLM(io.ComfyNode):
|
||||
top_p=top_p,
|
||||
seed=seed,
|
||||
video_frames=video_frames,
|
||||
reserve_free_vram_gb=reserve_free_vram_gb,
|
||||
reserve_free_vram_gb=(
|
||||
reserve_free_vram_gb
|
||||
if is_remote
|
||||
else local_reserve_free_memory_gb
|
||||
),
|
||||
keep_model_loaded=keep_model_loaded,
|
||||
progress_callback=report_progress,
|
||||
)
|
||||
|
||||
+32
-12
@@ -1127,6 +1127,9 @@ class ResidentLLMManager:
|
||||
empty_cache: Callable[[], None] | None = None,
|
||||
snapshot_ready: Callable[[str | Path, LLMModelProfile], bool] = is_model_snapshot_staged,
|
||||
comfy_memory_release: Callable[[int], None] | None = None,
|
||||
execution_target: str = "modal",
|
||||
device_name: str = "cuda",
|
||||
memory_label: str = "GPU memory",
|
||||
) -> None:
|
||||
"""Configure the shared model cache and injectable hardware operations."""
|
||||
if max_resident_models <= 0:
|
||||
@@ -1138,6 +1141,9 @@ class ResidentLLMManager:
|
||||
self._empty_cache = empty_cache or self._cuda_empty_cache
|
||||
self._snapshot_ready = snapshot_ready
|
||||
self._comfy_memory_release = comfy_memory_release or self._release_comfy_memory
|
||||
self.execution_target = execution_target
|
||||
self.device_name = device_name
|
||||
self.memory_label = memory_label
|
||||
self._models: OrderedDict[str, ResidentModel] = OrderedDict()
|
||||
self._lock = threading.RLock()
|
||||
|
||||
@@ -1176,7 +1182,8 @@ class ResidentLLMManager:
|
||||
"""Unload one resident backend and release cached CUDA allocations."""
|
||||
resident = self._models.pop(profile_id)
|
||||
logger.info(
|
||||
"Evicting resident Modal LLM profile=%s allocated_gib=%.3f.",
|
||||
"Evicting resident LLM target=%s profile=%s allocated_gib=%.3f.",
|
||||
self.execution_target,
|
||||
profile_id,
|
||||
resident.allocated_bytes / _BYTES_PER_GIB,
|
||||
)
|
||||
@@ -1194,10 +1201,11 @@ class ResidentLLMManager:
|
||||
free_bytes, total_bytes = self._memory_info()
|
||||
if free_bytes < required_bytes:
|
||||
raise RuntimeError(
|
||||
f"Modal LLM profile {profile.profile_id!r} needs approximately "
|
||||
f"LLM profile {profile.profile_id!r} needs approximately "
|
||||
f"{profile.estimated_vram_gb:.1f} GiB plus {reserve_free_vram_gb:.1f} GiB reserve, "
|
||||
f"but only {free_bytes / _BYTES_PER_GIB:.1f} of "
|
||||
f"{total_bytes / _BYTES_PER_GIB:.1f} GiB is free."
|
||||
f"{total_bytes / _BYTES_PER_GIB:.1f} GiB of {self.memory_label} is "
|
||||
"available."
|
||||
)
|
||||
|
||||
def _load(
|
||||
@@ -1223,16 +1231,16 @@ class ResidentLLMManager:
|
||||
return cached, True
|
||||
if not self._snapshot_ready(self.storage_root, profile):
|
||||
raise RuntimeError(
|
||||
f"Modal LLM profile {profile.profile_id!r} is not staged at "
|
||||
f"{model_snapshot_path(self.storage_root, profile)}. The CPU ModelStager must "
|
||||
"complete before GPU inference starts."
|
||||
f"LLM profile {profile.profile_id!r} is not staged at "
|
||||
f"{model_snapshot_path(self.storage_root, profile)}. Model staging must "
|
||||
f"complete before {self.execution_target} inference starts."
|
||||
)
|
||||
while len(self._models) >= self.max_resident_models:
|
||||
self._evict(next(iter(self._models)))
|
||||
progress_callback(
|
||||
LLMProgressEvent(
|
||||
stage="memory",
|
||||
message="Preparing GPU memory",
|
||||
message=f"Preparing {self.memory_label}",
|
||||
indeterminate=True,
|
||||
)
|
||||
)
|
||||
@@ -1254,7 +1262,9 @@ class ResidentLLMManager:
|
||||
)
|
||||
self._models[profile.profile_id] = resident
|
||||
logger.info(
|
||||
"Loaded resident Modal LLM profile=%s measured_allocation_gib=%.3f residents=%s.",
|
||||
"Loaded resident LLM target=%s profile=%s measured_allocation_gib=%.3f "
|
||||
"residents=%s.",
|
||||
self.execution_target,
|
||||
profile.profile_id,
|
||||
resident.allocated_bytes / _BYTES_PER_GIB,
|
||||
list(self._models),
|
||||
@@ -1300,6 +1310,8 @@ class ResidentLLMManager:
|
||||
comfy_loaded_model_names = _comfy_loaded_model_names()
|
||||
metadata = {
|
||||
"backend": profile.backend,
|
||||
"execution_target": self.execution_target,
|
||||
"device": self.device_name,
|
||||
"profile": profile.profile_id,
|
||||
"repository": profile.repository,
|
||||
"revision": profile.revision,
|
||||
@@ -1329,14 +1341,22 @@ class ResidentLLMManager:
|
||||
"video_frame_count": (
|
||||
len(prepared_inputs.video.frames) if prepared_inputs.video is not None else 0
|
||||
),
|
||||
"gpu_total_gib": total_bytes / _BYTES_PER_GIB,
|
||||
"gpu_free_before_gib": before_free / _BYTES_PER_GIB,
|
||||
"gpu_free_after_gib": after_free / _BYTES_PER_GIB,
|
||||
"memory_total_gib": total_bytes / _BYTES_PER_GIB,
|
||||
"memory_available_before_gib": before_free / _BYTES_PER_GIB,
|
||||
"memory_available_after_gib": after_free / _BYTES_PER_GIB,
|
||||
"resident_profiles": resident_ids,
|
||||
"comfy_loaded_model_count": len(comfy_loaded_model_names),
|
||||
"comfy_loaded_model_names": comfy_loaded_model_names,
|
||||
}
|
||||
logger.info("Completed Modal LLM inference: %s", metadata)
|
||||
if self.execution_target == "modal":
|
||||
metadata.update(
|
||||
{
|
||||
"gpu_total_gib": total_bytes / _BYTES_PER_GIB,
|
||||
"gpu_free_before_gib": before_free / _BYTES_PER_GIB,
|
||||
"gpu_free_after_gib": after_free / _BYTES_PER_GIB,
|
||||
}
|
||||
)
|
||||
logger.info("Completed LLM inference: %s", metadata)
|
||||
return LLMInferenceResult(
|
||||
text=generation_result.text,
|
||||
metadata=metadata,
|
||||
|
||||
@@ -13,6 +13,10 @@ classifiers = [
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-apple = [
|
||||
"mlx-vlm==0.6.15; sys_platform == 'darwin' and platform_machine == 'arm64'",
|
||||
"psutil>=7,<8; sys_platform == 'darwin' and platform_machine == 'arm64'",
|
||||
]
|
||||
remote = [
|
||||
"modal==1.4.2",
|
||||
]
|
||||
|
||||
@@ -162,6 +162,12 @@ def modal_llm_runtime_module(extension_package: object) -> object:
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.modal_llm_runtime")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def local_llm_runtime_module(extension_package: object) -> object:
|
||||
"""Return the Apple-local resident LLM runtime module."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.local_llm_runtime")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def llm_reasoning_module(extension_package: object) -> object:
|
||||
"""Return the backend-neutral LLM reasoning-output module."""
|
||||
|
||||
@@ -124,6 +124,14 @@ class _FakeRemoteAudioNode:
|
||||
OUTPUT_IS_LIST = (False,)
|
||||
|
||||
|
||||
class _FakeTextNode:
|
||||
"""Fake text node used for local LLM boundary planning."""
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
OUTPUT_IS_LIST = (False,)
|
||||
|
||||
|
||||
class _FakeRemoteImageConsumerNode:
|
||||
"""Fake remote node that consumes IMAGE and produces IMAGE."""
|
||||
|
||||
@@ -150,6 +158,50 @@ class _FakeRemoteArtifactWriterNode:
|
||||
OUTPUT_NODE = False
|
||||
|
||||
|
||||
def test_planner_keeps_unmarked_llm_local_between_remote_text_nodes(
|
||||
api_intercept_module: Any,
|
||||
) -> None:
|
||||
"""Transportable text boundaries must not absorb a local LLM into Modal."""
|
||||
prompt = {
|
||||
"1": {"class_type": "RemoteTextSource", "inputs": {}},
|
||||
"2": {
|
||||
"class_type": "ModalLLM",
|
||||
"inputs": {"prompt": ["1", 0], "model_profile": "owner/model"},
|
||||
},
|
||||
"3": {
|
||||
"class_type": "RemoteTextConsumer",
|
||||
"inputs": {"prompt": ["2", 0]},
|
||||
},
|
||||
}
|
||||
workflow = {
|
||||
"nodes": [
|
||||
{"id": 1, "properties": {"is_modal_remote": True}},
|
||||
{"id": 2, "properties": {"is_modal_remote": False}},
|
||||
{"id": 3, "properties": {"is_modal_remote": True}},
|
||||
]
|
||||
}
|
||||
fake_nodes_module = SimpleNamespace(
|
||||
NODE_CLASS_MAPPINGS={
|
||||
"RemoteTextSource": _FakeTextNode,
|
||||
"ModalLLM": _FakeTextNode,
|
||||
"RemoteTextConsumer": _FakeTextNode,
|
||||
},
|
||||
NODE_DISPLAY_NAME_MAPPINGS={},
|
||||
)
|
||||
|
||||
analysis = api_intercept_module.analyze_remote_node_selection(
|
||||
prompt,
|
||||
workflow,
|
||||
[],
|
||||
settings=SimpleNamespace(marker_property="is_modal_remote"),
|
||||
nodes_module=fake_nodes_module,
|
||||
)
|
||||
|
||||
assert analysis.resolved_remote_node_ids == ["1", "3"]
|
||||
assert analysis.sandwiched_local_node_ids == ["2"]
|
||||
assert analysis.added_node_ids == []
|
||||
|
||||
|
||||
def _artifact_finalizer_node_id(summary: Any) -> str:
|
||||
"""Return the finalizer id after asserting that prompt rewrite attached it."""
|
||||
finalizer_node_id = summary.artifact_finalizer_node_id
|
||||
|
||||
+467
-16
@@ -125,6 +125,79 @@ def test_compatibility_policy_selects_requested_model_backends(
|
||||
assert decision.reasoning_parser == ("qwen3" if backend == "vllm" else "none")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("architecture", "quantization", "reasoning_parser"),
|
||||
[
|
||||
("SmolVLMForConditionalGeneration", {}, "none"),
|
||||
("MuseGlimmerForConditionalGeneration", {}, "none"),
|
||||
(
|
||||
"Qwen3_5ForConditionalGeneration",
|
||||
{"bits": 4, "group_size": 64, "mode": "affine"},
|
||||
"qwen3",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_apple_local_compatibility_selects_mlx_vlm(
|
||||
llm_compatibility_module: Any,
|
||||
architecture: str,
|
||||
quantization: dict[str, Any],
|
||||
reasoning_parser: str,
|
||||
) -> None:
|
||||
"""Reviewed Apple-local architectures should resolve to the MLX adapter."""
|
||||
decision = llm_compatibility_module.resolve_compatibility(
|
||||
{
|
||||
"architectures": [architecture],
|
||||
"dtype": "bfloat16",
|
||||
"text_config": {"max_position_embeddings": 65536},
|
||||
"quantization": quantization,
|
||||
},
|
||||
artifact_bytes=2 * 1024**3,
|
||||
execution_target="local_apple",
|
||||
)
|
||||
|
||||
assert decision.backend == "mlx_vlm"
|
||||
assert decision.reasoning_parser == reasoning_parser
|
||||
assert decision.runtime_requirements == ("mlx-vlm==0.6.15",)
|
||||
assert decision.estimated_vram_gb < 5
|
||||
|
||||
|
||||
def test_apple_local_compatibility_rejects_cuda_quantization(
|
||||
llm_compatibility_module: Any,
|
||||
) -> None:
|
||||
"""Apple resolution should fail before downloading CUDA-specific weights."""
|
||||
with pytest.raises(ValueError, match="not an MLX checkpoint format"):
|
||||
llm_compatibility_module.resolve_compatibility(
|
||||
{
|
||||
"architectures": ["Qwen3_5ForConditionalGeneration"],
|
||||
"text_config": {"max_position_embeddings": 32768},
|
||||
"quantization_config": {"quant_method": "fp8"},
|
||||
},
|
||||
artifact_bytes=1024,
|
||||
execution_target="local_apple",
|
||||
)
|
||||
|
||||
|
||||
def test_curated_profile_adapts_to_apple_without_mutating_modal_profile(
|
||||
llm_profiles_module: Any,
|
||||
) -> None:
|
||||
"""One saved curated id should resolve independently for both targets."""
|
||||
modal_profile = llm_profiles_module.get_llm_profile(
|
||||
"smolvlm2-2.2b-instruct"
|
||||
)
|
||||
|
||||
local_profile = llm_profiles_module.profile_for_execution_target(
|
||||
modal_profile,
|
||||
"local_apple",
|
||||
)
|
||||
|
||||
assert modal_profile.backend == "transformers"
|
||||
assert modal_profile.execution_target == "modal"
|
||||
assert local_profile.profile_id == modal_profile.profile_id
|
||||
assert local_profile.repository == modal_profile.repository
|
||||
assert local_profile.backend == "mlx_vlm"
|
||||
assert local_profile.execution_target == "local_apple"
|
||||
|
||||
|
||||
class _FakeReasoningTokenizer:
|
||||
"""Decode deterministic IDs while modelling Qwen's special think tokens."""
|
||||
|
||||
@@ -311,6 +384,57 @@ def test_cpu_resolver_pins_and_persists_generated_profile(
|
||||
assert loaded == first.profile
|
||||
|
||||
|
||||
def test_local_resolver_creates_target_specific_mlx_profile(
|
||||
llm_resolver_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""The same user-facing model reference should get an Apple-local manifest."""
|
||||
config_path = tmp_path / "mlx-config.json"
|
||||
config_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"architectures": ["Qwen3_5ForConditionalGeneration"],
|
||||
"dtype": "bfloat16",
|
||||
"text_config": {"max_position_embeddings": 65536},
|
||||
"quantization": {
|
||||
"bits": 4,
|
||||
"group_size": 64,
|
||||
"mode": "affine",
|
||||
},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
class FakeApi:
|
||||
"""Return one immutable MLX-format repository."""
|
||||
|
||||
def model_info(self, repo_id: str, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Return compatible metadata without downloading model weights."""
|
||||
del repo_id, kwargs
|
||||
return {
|
||||
"sha": "7" * 40,
|
||||
"siblings": [
|
||||
{"rfilename": "config.json", "size": 1000},
|
||||
{"rfilename": "model.safetensors", "size": 2 * 1024**3},
|
||||
],
|
||||
"securityStatus": {"scansDone": True, "filesWithIssues": []},
|
||||
}
|
||||
|
||||
result = llm_resolver_module.resolve_model_profile(
|
||||
"mlx-community/Qwen3.5-2B-4bit",
|
||||
tmp_path,
|
||||
api=FakeApi(),
|
||||
hf_hub_download=lambda **kwargs: str(config_path),
|
||||
execution_target="local_apple",
|
||||
)
|
||||
|
||||
assert result.profile.backend == "mlx_vlm"
|
||||
assert result.profile.execution_target == "local_apple"
|
||||
assert result.profile.quantization_method == "mlx_affine_4bit"
|
||||
assert result.profile.runtime_requirements == ("mlx-vlm==0.6.15",)
|
||||
|
||||
|
||||
def test_cpu_resolver_rejects_unknown_architecture_before_weights(
|
||||
llm_resolver_module: Any,
|
||||
tmp_path: Path,
|
||||
@@ -390,6 +514,30 @@ def test_cpu_resolver_wraps_gated_config_download_error(
|
||||
)
|
||||
|
||||
|
||||
def test_local_resolver_reports_local_hugging_face_token_location(
|
||||
llm_resolver_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Gated local models should not direct users to a Modal secret."""
|
||||
|
||||
class DeniedApi:
|
||||
"""Reject model metadata as a gated repository would."""
|
||||
|
||||
def model_info(self, repo_id: str, **kwargs: Any) -> Any:
|
||||
"""Raise a deterministic authorization error."""
|
||||
del repo_id, kwargs
|
||||
raise OSError("401 Unauthorized")
|
||||
|
||||
with pytest.raises(ValueError, match="set HF_TOKEN in the local ComfyUI"):
|
||||
llm_resolver_module.resolve_model_profile(
|
||||
"owner/gated",
|
||||
tmp_path,
|
||||
api=DeniedApi(),
|
||||
hf_hub_download=lambda **kwargs: "unused",
|
||||
execution_target="local_apple",
|
||||
)
|
||||
|
||||
|
||||
def test_cpu_stager_writes_completion_marker_and_reuses_snapshot(
|
||||
llm_staging_module: Any,
|
||||
tmp_path: Path,
|
||||
@@ -911,6 +1059,234 @@ def test_vllm_backend_aborts_when_progress_callback_cancels(
|
||||
assert aborted_requests[0].startswith("modal-llm-")
|
||||
|
||||
|
||||
def test_local_storage_uses_comfyui_model_directory(
|
||||
local_llm_runtime_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Local snapshots should follow ComfyUI's configured model root."""
|
||||
paths: list[str] = []
|
||||
fake_folder_paths = SimpleNamespace(
|
||||
models_dir=str(tmp_path / "models"),
|
||||
add_model_folder_path=lambda name, path, is_default=False: paths.append(
|
||||
f"{name}:{path}:{is_default}"
|
||||
),
|
||||
get_folder_paths=lambda name: [str(tmp_path / "models" / name)],
|
||||
)
|
||||
monkeypatch.delenv("COMFY_MODAL_LOCAL_LLM_STORAGE_ROOT", raising=False)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths)
|
||||
|
||||
storage_root = local_llm_runtime_module.local_llm_storage_root()
|
||||
|
||||
assert storage_root == (tmp_path / "models" / "modal_llm").resolve()
|
||||
assert paths == [
|
||||
f"modal_llm:{(tmp_path / 'models' / 'modal_llm').resolve()}:True"
|
||||
]
|
||||
|
||||
|
||||
def test_local_runtime_rejects_non_apple_hardware_actionably(
|
||||
local_llm_runtime_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Unsupported local hardware should direct the workflow to Modal."""
|
||||
monkeypatch.setattr(local_llm_runtime_module.platform, "system", lambda: "Linux")
|
||||
monkeypatch.setattr(local_llm_runtime_module.platform, "machine", lambda: "x86_64")
|
||||
|
||||
with pytest.raises(RuntimeError, match="Enable 'Run on Modal'"):
|
||||
local_llm_runtime_module.ensure_local_apple_runtime_available()
|
||||
|
||||
|
||||
def test_local_runtime_requires_the_pinned_mlx_version(
|
||||
local_llm_runtime_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A drifted MLX adapter should fail with an exact repair command."""
|
||||
monkeypatch.setattr(local_llm_runtime_module.platform, "system", lambda: "Darwin")
|
||||
monkeypatch.setattr(local_llm_runtime_module.platform, "machine", lambda: "arm64")
|
||||
monkeypatch.setattr(
|
||||
local_llm_runtime_module.metadata,
|
||||
"version",
|
||||
lambda distribution: "0.1.0",
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="mlx-vlm==0.6.15"):
|
||||
local_llm_runtime_module.ensure_local_apple_runtime_available()
|
||||
|
||||
|
||||
def test_local_curated_profile_is_staged_with_mlx_backend(
|
||||
local_llm_runtime_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A saved curated id should stage locally without changing the workflow."""
|
||||
staged: list[tuple[str, Any, Any]] = []
|
||||
monkeypatch.setattr(
|
||||
local_llm_runtime_module,
|
||||
"stage_model_profile",
|
||||
lambda profile_id, storage_root, **kwargs: staged.append(
|
||||
(profile_id, Path(storage_root), kwargs["profile"])
|
||||
),
|
||||
)
|
||||
|
||||
profile = local_llm_runtime_module.resolve_and_stage_local_profile(
|
||||
"smolvlm2-2.2b-instruct",
|
||||
tmp_path,
|
||||
progress_callback=lambda progress: None,
|
||||
)
|
||||
|
||||
assert profile.backend == "mlx_vlm"
|
||||
assert profile.execution_target == "local_apple"
|
||||
assert staged == [(profile.profile_id, tmp_path, profile)]
|
||||
|
||||
|
||||
def test_mlx_backend_streams_multimodal_reasoning_and_progress(
|
||||
local_llm_runtime_module: Any,
|
||||
modal_llm_runtime_module: Any,
|
||||
llm_profiles_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""MLX streaming should retain the shared outputs and reasoning contract."""
|
||||
import types
|
||||
|
||||
observed: dict[str, Any] = {}
|
||||
|
||||
class FakeMX:
|
||||
"""Expose only the MLX lifecycle methods used by the adapter."""
|
||||
|
||||
@staticmethod
|
||||
def reset_peak_memory() -> None:
|
||||
"""Record peak-memory reset."""
|
||||
observed["reset_peak_memory"] = True
|
||||
|
||||
@staticmethod
|
||||
def clear_cache() -> None:
|
||||
"""Record cache release."""
|
||||
observed["clear_cache"] = True
|
||||
|
||||
processor = SimpleNamespace(tokenizer=_FakeReasoningTokenizer())
|
||||
model = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_5"))
|
||||
|
||||
def fake_load(path: str, **kwargs: Any) -> tuple[Any, Any]:
|
||||
"""Return a deterministic resident MLX model and processor."""
|
||||
observed["load"] = (path, kwargs)
|
||||
return model, processor
|
||||
|
||||
def fake_apply_chat_template(
|
||||
selected_processor: Any,
|
||||
config: Any,
|
||||
messages: Any,
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Capture the structured prompt passed to MLX-VLM."""
|
||||
observed["chat"] = (selected_processor, config, messages, kwargs)
|
||||
return "rendered prompt"
|
||||
|
||||
def fake_stream_generate(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Yield exact reasoning boundary and response tokens."""
|
||||
observed["generate"] = (args, kwargs)
|
||||
segments = [
|
||||
("<think>", 10),
|
||||
("consider", 1),
|
||||
("</think>", 11),
|
||||
("final answer", 3),
|
||||
]
|
||||
for index, (text, token) in enumerate(segments, start=1):
|
||||
yield SimpleNamespace(
|
||||
text=text,
|
||||
token=token,
|
||||
prompt_tokens=5,
|
||||
generation_tokens=index,
|
||||
generation_tps=10.0 + index,
|
||||
)
|
||||
|
||||
mlx_module = types.ModuleType("mlx")
|
||||
mlx_module.__path__ = []
|
||||
mlx_core_module = types.ModuleType("mlx.core")
|
||||
mlx_core_module.reset_peak_memory = FakeMX.reset_peak_memory
|
||||
mlx_core_module.clear_cache = FakeMX.clear_cache
|
||||
mlx_vlm_module = types.ModuleType("mlx_vlm")
|
||||
mlx_vlm_module.__path__ = []
|
||||
mlx_vlm_module.load = fake_load
|
||||
mlx_generate_module = types.ModuleType("mlx_vlm.generate")
|
||||
mlx_generate_module.stream_generate = fake_stream_generate
|
||||
mlx_prompt_module = types.ModuleType("mlx_vlm.prompt_utils")
|
||||
mlx_prompt_module.apply_chat_template = fake_apply_chat_template
|
||||
monkeypatch.setitem(sys.modules, "mlx", mlx_module)
|
||||
monkeypatch.setitem(sys.modules, "mlx.core", mlx_core_module)
|
||||
monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm_module)
|
||||
monkeypatch.setitem(sys.modules, "mlx_vlm.generate", mlx_generate_module)
|
||||
monkeypatch.setitem(sys.modules, "mlx_vlm.prompt_utils", mlx_prompt_module)
|
||||
monkeypatch.setattr(
|
||||
local_llm_runtime_module,
|
||||
"ensure_local_apple_runtime_available",
|
||||
lambda: None,
|
||||
)
|
||||
profile = replace(
|
||||
llm_profiles_module.get_llm_profile("smolvlm2-2.2b-instruct"),
|
||||
backend="mlx_vlm",
|
||||
architecture="Qwen3_5ForConditionalGeneration",
|
||||
reasoning_parser="qwen3",
|
||||
execution_target="local_apple",
|
||||
max_context_tokens=128,
|
||||
)
|
||||
progress: list[Any] = []
|
||||
backend = local_llm_runtime_module.MLXVLMBackend(
|
||||
profile,
|
||||
tmp_path,
|
||||
progress.append,
|
||||
)
|
||||
prepared = modal_llm_runtime_module.PreparedLLMInputs(
|
||||
prompt="describe",
|
||||
system_prompt="be concise",
|
||||
images=("image-a",),
|
||||
video=modal_llm_runtime_module.PreparedVideo(
|
||||
frames=("frame-a", "frame-b"),
|
||||
timestamps_seconds=(0.0, 1.0),
|
||||
),
|
||||
file_characters=0,
|
||||
file_count=0,
|
||||
)
|
||||
|
||||
result = backend.generate(
|
||||
prepared,
|
||||
modal_llm_runtime_module.LLMGenerationSettings(
|
||||
max_new_tokens=8,
|
||||
temperature=0.0,
|
||||
top_p=0.95,
|
||||
seed=7,
|
||||
),
|
||||
progress.append,
|
||||
)
|
||||
backend.unload()
|
||||
|
||||
assert result.text == "final answer"
|
||||
assert result.reasoning == "consider"
|
||||
assert result.reasoning_tokens == 1
|
||||
assert result.input_tokens == 5
|
||||
assert result.output_tokens == 4
|
||||
assert observed["chat"][2] == [
|
||||
{"role": "system", "content": "be concise"},
|
||||
{"role": "user", "content": "describe"},
|
||||
]
|
||||
assert observed["chat"][3]["num_images"] == 3
|
||||
assert observed["generate"][1]["image"] == [
|
||||
"image-a",
|
||||
"frame-a",
|
||||
"frame-b",
|
||||
]
|
||||
assert observed["generate"][1]["seed"] == 7
|
||||
assert observed["generate"][1]["enable_thinking"] is True
|
||||
assert observed["generate"][1]["thinking_budget"] == 4
|
||||
assert [event.value for event in progress if event.stage == "generating"] == [
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
]
|
||||
assert observed["clear_cache"] is True
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeBackend:
|
||||
"""Record resident backend inference and unload behavior."""
|
||||
@@ -1020,18 +1396,73 @@ def test_resident_manager_reuses_and_lru_evicts_models(
|
||||
)
|
||||
|
||||
assert first.metadata["cache_hit"] is False
|
||||
assert first.metadata["execution_target"] == "modal"
|
||||
assert first.metadata["device"] == "cuda"
|
||||
assert first.metadata["memory_total_gib"] == 256
|
||||
assert first.metadata["gpu_total_gib"] == 256
|
||||
assert second.metadata["cache_hit"] is True
|
||||
assert backends[base_profile.profile_id].generate_calls == 2
|
||||
assert backends[base_profile.profile_id].unloaded is True
|
||||
assert manager.resident_profiles() == (second_profile.profile_id,)
|
||||
|
||||
|
||||
def test_modal_llm_node_is_v3_remote_only_and_returns_metadata(
|
||||
def test_resident_manager_reports_apple_unified_memory_metadata(
|
||||
modal_llm_runtime_module: Any,
|
||||
llm_profiles_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Local inference telemetry should not mislabel unified memory as VRAM."""
|
||||
profile = replace(
|
||||
llm_profiles_module.get_llm_profile("smolvlm2-2.2b-instruct"),
|
||||
backend="mlx_vlm",
|
||||
execution_target="local_apple",
|
||||
)
|
||||
manager = modal_llm_runtime_module.ResidentLLMManager(
|
||||
storage_root=tmp_path,
|
||||
backend_factory=lambda *args: _FakeBackend(profile.profile_id),
|
||||
memory_info=lambda: (48 * 1024**3, 64 * 1024**3),
|
||||
empty_cache=lambda: None,
|
||||
snapshot_ready=lambda storage_root, selected_profile: True,
|
||||
comfy_memory_release=lambda required_bytes: None,
|
||||
execution_target="local_apple",
|
||||
device_name="metal",
|
||||
memory_label="unified memory",
|
||||
)
|
||||
|
||||
result = manager.infer(
|
||||
profile=profile,
|
||||
prepared_inputs=modal_llm_runtime_module.PreparedLLMInputs(
|
||||
prompt="hello",
|
||||
system_prompt="",
|
||||
images=(),
|
||||
video=None,
|
||||
file_characters=0,
|
||||
file_count=0,
|
||||
),
|
||||
generation_settings=modal_llm_runtime_module.LLMGenerationSettings(
|
||||
max_new_tokens=8,
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
seed=0,
|
||||
),
|
||||
reserve_free_vram_gb=1,
|
||||
keep_model_loaded=False,
|
||||
progress_callback=lambda progress: None,
|
||||
)
|
||||
|
||||
assert result.metadata["execution_target"] == "local_apple"
|
||||
assert result.metadata["device"] == "metal"
|
||||
assert result.metadata["memory_total_gib"] == 64
|
||||
assert result.metadata["memory_available_before_gib"] == 48
|
||||
assert "gpu_total_gib" not in result.metadata
|
||||
|
||||
|
||||
def test_modal_llm_node_routes_between_local_and_remote_runtimes(
|
||||
modal_llm_node_module: Any,
|
||||
modal_llm_runtime_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The V3 node should refuse accidental local model loading."""
|
||||
"""The V3 node toggle state should select local or remote execution."""
|
||||
schema = modal_llm_node_module.ModalLLM.define_schema()
|
||||
assert schema.node_id == "ModalLLM"
|
||||
assert [output.display_name for output in schema.outputs] == [
|
||||
@@ -1040,32 +1471,52 @@ def test_modal_llm_node_is_v3_remote_only_and_returns_metadata(
|
||||
"reasoning",
|
||||
]
|
||||
|
||||
monkeypatch.delenv("COMFY_MODAL_REMOTE_WORKER", raising=False)
|
||||
with pytest.raises(RuntimeError, match="Enable 'Run on Modal'"):
|
||||
modal_llm_node_module.ModalLLM.execute(
|
||||
prompt="hello",
|
||||
model_profile="smolvlm2-2.2b-instruct",
|
||||
calls: list[tuple[str, float]] = []
|
||||
|
||||
def fake_result(target: str, **kwargs: Any) -> Any:
|
||||
"""Return deterministic output while recording the selected runtime."""
|
||||
calls.append((target, kwargs["reserve_free_vram_gb"]))
|
||||
return modal_llm_runtime_module.LLMInferenceResult(
|
||||
text=f"{target}-done",
|
||||
metadata={
|
||||
"output_tokens": 1,
|
||||
"cache_hit": True,
|
||||
"execution_target": target,
|
||||
},
|
||||
reasoning="working",
|
||||
)
|
||||
|
||||
monkeypatch.delenv("COMFY_MODAL_REMOTE_WORKER", raising=False)
|
||||
monkeypatch.setattr(
|
||||
modal_llm_node_module,
|
||||
"run_local_llm_inference",
|
||||
lambda **kwargs: fake_result("local_apple", **kwargs),
|
||||
)
|
||||
local_output = modal_llm_node_module.ModalLLM.execute(
|
||||
prompt="hello",
|
||||
model_profile="smolvlm2-2.2b-instruct",
|
||||
max_new_tokens=4,
|
||||
local_reserve_free_memory_gb=3.0,
|
||||
)
|
||||
monkeypatch.setenv("COMFY_MODAL_REMOTE_WORKER", "1")
|
||||
monkeypatch.setattr(
|
||||
modal_llm_node_module,
|
||||
"run_modal_llm_inference",
|
||||
lambda **kwargs: modal_llm_runtime_module.LLMInferenceResult(
|
||||
text="done",
|
||||
metadata={"output_tokens": 1, "cache_hit": True},
|
||||
reasoning="working",
|
||||
),
|
||||
lambda **kwargs: fake_result("modal", **kwargs),
|
||||
)
|
||||
node_output = modal_llm_node_module.ModalLLM.execute(
|
||||
remote_output = modal_llm_node_module.ModalLLM.execute(
|
||||
prompt="hello",
|
||||
model_profile="smolvlm2-2.2b-instruct",
|
||||
max_new_tokens=4,
|
||||
reserve_free_vram_gb=17.0,
|
||||
)
|
||||
|
||||
assert node_output.result[0] == "done"
|
||||
assert json.loads(node_output.result[1])["cache_hit"] is True
|
||||
assert node_output.result[2] == "working"
|
||||
assert local_output.result[0] == "local_apple-done"
|
||||
assert json.loads(local_output.result[1])["execution_target"] == "local_apple"
|
||||
assert remote_output.result[0] == "modal-done"
|
||||
assert json.loads(remote_output.result[1])["execution_target"] == "modal"
|
||||
assert remote_output.result[2] == "working"
|
||||
assert calls == [("local_apple", 3.0), ("modal", 17.0)]
|
||||
|
||||
|
||||
def test_remote_dispatch_stages_llm_once_and_forces_volume_reload(
|
||||
|
||||
@@ -24,6 +24,21 @@ def test_bootstrap_pin_matches_declared_remote_extra(modal_sdk_module: Any) -> N
|
||||
]
|
||||
|
||||
|
||||
def test_local_apple_extra_matches_validated_mlx_runtime(
|
||||
local_llm_runtime_module: Any,
|
||||
) -> None:
|
||||
"""The opt-in Apple backend should pin the adapter validated by the runtime."""
|
||||
project_metadata = tomllib.loads(
|
||||
(REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8")
|
||||
)
|
||||
|
||||
marker = "sys_platform == 'darwin' and platform_machine == 'arm64'"
|
||||
assert project_metadata["project"]["optional-dependencies"]["local-apple"] == [
|
||||
f"{local_llm_runtime_module.LOCAL_MLX_VLM_SPEC}; {marker}",
|
||||
f"psutil>=7,<8; {marker}",
|
||||
]
|
||||
|
||||
|
||||
def test_remote_mode_installs_and_reimports_missing_modal_sdk(
|
||||
modal_sdk_module: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
Reference in New Issue
Block a user