Compare commits

..

19 Commits

Author SHA1 Message Date
Evan c6fab7fa97 fix: only shutdown runners that error unrecoverably
previously we would shutdown a runner that had previously failed,
leading to loops of create -> shutdown -> create.

now when a runner critically errors, we set its status to Shutdown, and
delete it. a failed runner is expected
2026-03-25 16:46:06 +00:00
vskiwi fc1ae90111 fix: DeepSeek V3.2 warmup crash and tool calling + add catalog cards (#1769)
## Summary

DeepSeek V3.2 (`DeepseekV32ForCausalLM`) is already supported by exo's
inference engine (architecture whitelisted in `model_cards.py`, DSML
encoding added in #1548), but **doesn't work out of the box** due to two
bugs:

### Bug 1: `warmup_inference` passes empty model ID

`warmup_inference()` in `generate.py` accepts `model_id: ModelId` as a
parameter but creates `TextGenerationTaskParams(model=ModelId(""), ...)`
instead of using it. Since `_needs_dsml_encoding()` checks
`"deepseek-v3.2" in task_params.model.lower()`, the empty string never
matches → falls back to `tokenizer.apply_chat_template()` →
**ValueError** because V3.2 has no Jinja chat template.

**Fix:** `model=ModelId("")` → `model=model_id` (one line).

### Bug 2: `_needs_dsml_encoding` limited to tool calling

`_needs_dsml_encoding()` returns `True` only when `task_params.tools` is
present or tool messages exist in `chat_template_messages`. For warmup
and regular chat requests without tools → `return False` → Jinja
fallback → **ValueError**.

Unlike V3.1 (which has a `.jinja` chat template file that transformers
picks up automatically), V3.2 **has no Jinja template at all** — it uses
Python-based DSML encoding for all message types.

**Fix:** For V3.2, always return `True` — DSML encoding handles all
message types.

### Catalog cards

Added inference model cards for:
- `mlx-community/DeepSeek-V3.2-8bit`
- `mlx-community/DeepSeek-V3.2-4bit`

Parameters taken from model `config.json` on HuggingFace, storage sizes
from HF API. Capabilities include `thinking_toggle` (related: #1456).

## Notes

- The model ID string matching approach (`"deepseek-v3.2" in
model.lower()`) is acknowledged tech debt — see #1371 for the planned
architecture-based approach.

## Test plan

- [x] Start exo with DeepSeek V3.2 model → warmup should complete
without crash
- [x] Send a regular chat message (no tools) → should get a response
- [x] Send a chat message with tools → should work as before
- [x] V3.2 cards should appear in the dashboard model catalog

---------

Co-authored-by: user <user@m1.note>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
Co-authored-by: Evan <evanev7@gmail.com>
2026-03-25 16:20:35 +00:00
rltakashige 565ed41c13 Fix occasional warmup bugs by using mlx_generate (#1794)
## Motivation

Warmup occasionally had issues; e.g. #1748 and #1793 because we were
using a standard stream_generate, all of which are issues that are
resolved in the wrapper function mlx_generate.
2026-03-25 15:53:54 +00:00
DeepZima 2da740c387 Feat/static peer discovery (#1690)
**Enabling peers to be discovered in environments where mDNS is
unavailable (SSH sessions, headless servers, Docker).**

## Motivation
Exo discovers peers exclusively via mDNS, which works great on a local
network but breaks once you move beyond a single L2 broadcast domain:

- SSH sessions on macOS — TCC blocks mDNS multicast from non-GUI
sessions (#1488)
- Headless servers/rack machines — #1682 ("DGX Spark does not find other
nodes")
- Docker Compose — mDNS is often unavailable across container networks;
e.g. #1462 (E2E test framework) needs an alternative

Related works: 
#1488 (working implementation made by @AlexCheema and closed because SSH
had a GUI workaround),
#1023 (Headscale WAN then closed due to merge conflicts), 
#1656 (discovery cleanup, open). 

This PR introduces an optional bootstrap mechanism for peer discovery
while leaving the existing mDNS behavior unchanged.

## Changes
Adds two new CLI flags:

- `--bootstrap-peers` (env: `EXO_BOOTSTRAP_PEERS`) — comma-separated
libp2p multiaddrs to dial on startup and retry periodically
- `--libp2p-port` — fixed TCP port for libp2p to listen on (default:
OS-assigned). Required when bootstrap peers, so other nodes know which
port to dial.

8 files: 
- `rust/networking/src/discovery.rs`: Store bootstrap addrs, dial in
existing retry loop
- `rust/networking/src/swarm.rs`: Thread `bootstrap_peers` parameter to
`Behaviour`
- `rust/networking/examples/chatroom.rs`: Updated call site for new
create_swarm signature
- `rust/networking/tests/bootstrap_peers.rs`: Integration tests
- `rust/exo_pyo3_bindings/src/networking.rs`: Accept optional
`bootstrap_peers` in PyO3 constructor
- `rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi` : Update type stub 
- `src/exo/routing/router.py`: Pass peers to `NetworkingHandle` 
- `src/exo/main.py` : `--bootstrap-peers` CLI arg +
`EXO_BOOTSTRAP_PEERS` env var

## Why It Works

Bootstrap peers are dialed in the existing retry loop — the same path
taken by peers when mDNS-discovered. The swarm handles connection, Noise
handshake, and gossipsub mesh joining from there.

PeerId is intentionally not required in the multiaddr, the Noise
handshake discovers it.

Docker Compose example:

```yaml
services:
  exo-1:
    environment:
      EXO_BOOTSTRAP_PEERS: "/ip4/exo-2/tcp/30000"
  exo-2:
    environment:
      EXO_BOOTSTRAP_PEERS: "/ip4/exo-1/tcp/30000"
```

## Test Plan

### Manual Testing
<details>
<summary>Docker Compose config</summary>

```
services:
  exo-node1:
    build:
      context: .
      dockerfile: Dockerfile.bootstrap-test
    container_name: exo-bootstrap-node1
    hostname: exo-node1
    command: ["-q", "--libp2p-port", "30000", "--bootstrap-peers", "/ip4/172.30.20.3/tcp/30000"]
    environment:
      - EXO_LIBP2P_NAMESPACE=bootstrap-test
    ports:
      - "52415:52415"
    networks:
      bootstrap-net:
        ipv4_address: 172.30.20.2
    deploy:
      resources:
        limits:
          memory: 4g

  exo-node2:
    build:
      context: .
      dockerfile: Dockerfile.bootstrap-test
    container_name: exo-bootstrap-node2
    hostname: exo-node2
    command: ["-q", "--libp2p-port", "30000", "--bootstrap-peers", "/ip4/172.30.20.2/tcp/30000"]
    environment:
      - EXO_LIBP2P_NAMESPACE=bootstrap-test
    ports:
      - "52416:52415"
    networks:
      bootstrap-net:
        ipv4_address: 172.30.20.3
    deploy:
      resources:
        limits:
          memory: 4g

networks:
  bootstrap-net:
    driver: bridge
    ipam:
      config:
        - subnet: 172.30.20.0/24
```
</details> 

Two containers on a bridge network (`172.30.20.0/24`), fixed IPs,
`--libp2p-port 30000`, cross-referencing `--bootstrap-peers`.

Both nodes found each other and established a connection then ran the
election protocol.

### Automated Testing

4 Rust integration tests in `rust/networking/tests/bootstrap_peers.rs`
(`cargo test -p networking`):

| Test | What it verifies | Result |
|------|-----------------|--------|
| `two_nodes_connect_via_bootstrap_peers` | Node B discovers Node A via
bootstrap addr (real TCP connection) | PASS |
| `create_swarm_with_empty_bootstrap_peers` | Backward compatibility —
no bootstrap peers works | PASS |
| `create_swarm_ignores_invalid_bootstrap_addrs` | Invalid multiaddrs
silently filtered | PASS |
| `create_swarm_with_fixed_port` | `listen_port` parameter works | PASS
|

All 4 pass. The connection test takes ~6s

---------

Signed-off-by: DeepZima <deepzima@outlook.com>
Co-authored-by: Evan <evanev7@gmail.com>
2026-03-25 10:55:12 +00:00
rltakashige 7117d748ec Update dependencies including mlx 0.31.2 (#1789)
Update mlx fork to 0.31.2 and mflux to 0.17.2
2026-03-25 06:03:19 +00:00
Alex Cheema 178c617bbb Rename Nemotron to NVIDIA in model picker with logo (#1790)
## Motivation

The Nemotron model family in the model picker sidebar was displaying as
"Nemotron" with a generic checkmark icon. Since these models are
NVIDIA's Nemotron models, the category should be branded as "NVIDIA"
with the official NVIDIA logo, consistent with how other families are
branded (e.g., "llama" → "Meta", "gpt-oss" → "OpenAI").

## Changes

- **FamilySidebar.svelte**: Added `nemotron: "NVIDIA"` to the
`familyNames` mapping so the sidebar displays "NVIDIA" instead of
"Nemotron"
- **FamilyLogos.svelte**: Added the NVIDIA "eye" logo as an inline SVG
for the `nemotron` family, matching the `viewBox="0 0 24 24"` /
`fill="currentColor"` pattern used by all other brand logos
- **ModelPickerModal.svelte**: Added `"nemotron"` to the `familyOrder`
array so NVIDIA appears in a consistent position in the sidebar

## Why It Works

The model picker derives categories from the `family` field in TOML
model cards. Nemotron models already have `family = "nemotron"`, but the
three UI components (display name, logo, sort order) lacked explicit
entries for it, causing fallback behavior (auto-capitalized name,
checkmark icon, alphabetical sorting). Adding explicit entries for all
three aligns NVIDIA with the existing brand pattern.

## Test Plan

### Manual Testing
<!-- Hardware: N/A - dashboard UI change only -->
- Built dashboard successfully (`npm run build`)
- Verified the NVIDIA logo renders in the sidebar alongside existing
brand logos

### Automated Testing
- No test changes needed — this is a purely cosmetic dashboard change

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-25 01:56:59 +00:00
rltakashige 7277c90389 Fix enable thinking (#1786)
## Motivation

Minor regression from #1746 

## Why It Works

Pass thinking properly

## Test Plan

### Manual Testing
Works, it thinks
2026-03-24 16:26:25 -07:00
Evan Quiney 7ee88c1f05 override macmon in flake (#1747)
updates macmon to an upstream fork that fixes m5 max issues.
might see if the upstream version gets merged before we release.

---------

Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
2026-03-24 17:30:19 +00:00
Evan Quiney 509533d49e send error finish reason on failing to parse a tool call (#1785)
a simplification of #1757 which is now stale
2026-03-24 17:21:39 +00:00
rltakashige b6240a97e8 Prevent Qwen3.5 looping by using mlx lm fork (#1784)
## Motivation

Move back to an MLX LM to improve the Qwen 3.5 experience. 

## Test Plan

### Manual Testing
Seems to loop less from testing, no speed regressions.
2026-03-24 17:16:41 +00:00
rltakashige 6cdfbb7e8b Add HF_ENDPOINT in the app settings (#1783)
## Motivation
Some users (primarily in China) are unable to access HuggingFace.co. HF
supports the HF_ENDPOINT env variable, and we also support it, but there
is no way to easily do that from the app currently.

## Test Plan

### Manual Testing

hf-mirror works
empty field works
google.com endpoint fails
2026-03-24 17:05:49 +00:00
Evan Quiney fac6832e5f fix warmup consistency for slow machines (#1748)
a fix from pr #1643 which is now stale - should make prefill more
consistent on very slow machines

## testing
qwen-3.5-35b-a3b loads normally
gpt-oss-120b-mxfp4-q8 loads normally
2026-03-24 16:45:55 +00:00
rltakashige 7df3774ca2 Improve batch performance and stats reporting (#1777)
## Motivation

Batch generation reports incorrect statistics, as mlx lm never clears
the original stats, meaning they get polluted over time.
The dashboard also seems considerably slower than bench statistics.
We also have a large discrepancy between B=1 batch generating and
mlx_generate.
Extracting logprobs is massively expensive, causing up to a 25% slowdown
compared to pure batching.
```
[ 12:02:01.1240AM | INFO    ] step overhead: 3.49ms (next=12.49ms total=15.99ms)
[ 12:02:02.1600AM | INFO    ] step overhead: 3.23ms (next=13.01ms total=16.24ms)
[ 12:02:03.2228AM | INFO    ] step overhead: 3.28ms (next=13.38ms total=16.66ms)
[ 12:02:04.2798AM | INFO    ] step overhead: 3.25ms (next=12.84ms total=16.10ms)
[ 12:02:05.3152AM | INFO    ] step overhead: 3.18ms (next=12.61ms total=15.79ms)
[ 12:02:06.3522AM | INFO    ] step overhead: 3.41ms (next=12.83ms total=16.25ms)
[ 12:02:07.3987AM | INFO    ] step overhead: 3.38ms (next=13.14ms total=16.52ms)
[ 12:02:08.4537AM | INFO    ] step overhead: 1.84ms (next=19.44ms total=21.28ms)
```

## Changes

1. Report stats ourselves instead of using mlx lm's stats for batch
generation (they use perf_counter anyway).
2. Adjust exo bench to match
3. Improve logprobs extraction speed by 10x, improving tps for dashboard
& any requests for logprobs
4. Use an SSE comment to align the speed to the real numbers at the end
of generation
5. Patch mlx for several optimizations given our assumptions and use
cases (e.g. use vllm style RoPE).
6. Switch MLX LM version to latest main, including support for Nemotron
Super and some Qwen3.5 fixes.

## Why It Works
1. Exo bench no longer reports polluted stats
2. Exo bench now handles the reported per-request stats rather than the
aggregate stats
3. The decode speed now jumps back to a real number at the end of the
generation
4. Large batch speedup for rotating KV cache models + 1:1 matching cache
with vllm

## Test Plan

### Manual Testing
Needs testing on OpenCode and CC
Needs eval testing

### Automated Testing
Only going to show the performance optimization difference after the
accurate reporting:

**GPT OSS 20B MXFP4 Q8 (large change)**
Before:
<img width="2466" height="1534" alt="image"
src="https://github.com/user-attachments/assets/88b50637-fca2-4db4-9413-b9eee6e2057e"
/>
<img width="2410" height="1240" alt="image"
src="https://github.com/user-attachments/assets/21e5c76a-2f5f-44d2-8953-121b3ebdbd68"
/>


After:
<img width="2476" height="1472" alt="image"
src="https://github.com/user-attachments/assets/fec5cfbd-fff8-430a-b12e-a329410107a2"
/>
<img width="2454" height="1236" alt="image"
src="https://github.com/user-attachments/assets/0400344b-a4a6-42c0-a9dd-4ee91ade714a"
/>



**Qwen 3.5 35B A3B 8bit (No change)**
Before:
<img width="2414" height="1396" alt="image"
src="https://github.com/user-attachments/assets/e75f0b38-df5d-49fd-ab90-bc1667d981b3"
/>


After:
<img width="2346" height="1234" alt="image"
src="https://github.com/user-attachments/assets/eabfb59c-851f-4d88-b927-e1e699a75cc6"
/>


**Llama 3.2 1B Instruct 4bit (small change)**
Before:
<img width="2516" height="1220" alt="image"
src="https://github.com/user-attachments/assets/c2873655-acff-4536-8263-fb8aea33db80"
/>

After:
<img width="2566" height="1370" alt="image"
src="https://github.com/user-attachments/assets/15f95c75-1c2f-4474-85a2-88c4d0a32543"
/>
2026-03-24 14:03:03 +00:00
ciaranbor 248919c2a8 Fix first start in offline mode crash (#1782)
## Motivation

Running exo in offline mode on a machine where a model has never been
downloaded causes a crash

## Changes

- `download_shard` now catches `FileNotFoundError` from
`fetch_file_list_with_cache` and returns a `not_started` progress
instead of propagating the exception

## Why It Works

A status query should never crash its caller. By returning
`not_started`, the coordinator's existing offline guard (`if
self.offline:` at line 198) is reached and emits a graceful
`DownloadFailed` event. This also eliminates the warning spam on startup
where the status iterator catches the same exception for every
predefined model.

## Test Plan

### Manual Testing
- Start exo with `--offline` on a machine with no cached models, request
a model via the API — should get a graceful failure instead of a crash
2026-03-24 13:58:23 +00:00
ciaranbor 49951e1b1a Sync custom model cards across nodes (#1768)
## Motivation

Custom model cards were only saved locally on the node that handled the
API request.

## Changes

- Added AddCustomModelCard and DeleteCustomModelCard commands
- Added CustomModelCardAdded and CustomModelCardDeleted events
- Added custom_model_cards field to cluster State
- Master handles new commands by emitting corresponding events
- Workers persist model cards to disk and update the in-memory cache on
event receipt
- Separated fetch_from_hf (pure fetch) from disk persistence (now
handled by event-sourcing layer)
- Exposed add_to_card_cache() helper for the worker to update the cache

## Why It Works

- Follows the existing event-sourcing pattern: API → Command → Master →
Event → all Workers
- Every node applies the same events, so custom model cards are
consistent across the cluster

## Test Plan

### Manual Testing

Add/delete a custom model via the API on one node, verify it
appears/disappears on all nodes

### Automated Testing

Existing tests cover apply() logic; new event types follow the same
discriminated-union pattern
2026-03-24 12:51:36 +00:00
ciaranbor e06e70a835 Prefer higher model download % for placement (#1767)
## Motivation

When placing a model instance across the cluster, the master previously
only considered available RAM. This meant it could pick a node that
hasn't downloaded the model yet, even when another node already has it
(or is further along in downloading it).

## Changes

- Added download_status parameter to place_instance() in placement.py
- Added _get_node_download_fraction() to compute 0.0–1.0 download
progress per node/model
- Added _cycle_download_score() to sum download fractions across a
cycle's nodes
- Cycle selection now uses a (download_score, available_ram) tuple key —
download progress is the primary sort, RAM is the tiebreaker
- Passed self.state.downloads into place_instance() from master/main.py

## Why It Works

Python's tuple comparison gives download progress strict priority over
RAM, so a node with the model already downloaded will always be
preferred over one with more free RAM but no download.

## Test Plan

### Automated Testing

3 new tests cover: completed download preferred, higher partial progress
preferred, failed download not preferred over no-download node
2026-03-24 12:11:56 +00:00
vskiwi e9fdd8d4af improve logging: add dates to verbose stderr, match file log level to verbosity (#1772)
## Summary

- **Add ISO date to verbose stderr format**: when running with `-v`, the
stderr timestamp changes from `HH:mm:ss.SSS` to `YYYY-MM-DD
HH:mm:ss.SSS`, matching the file log format. This makes it possible to
correlate entries across days when stderr is captured by launchd,
systemd, Docker, or file redirection.
- **File log respects verbosity**: `exo.log` now uses DEBUG level when
`-v` is passed, so the persistent log has the same detail as stderr.
Previously it was always INFO regardless of verbosity.
- **Startup banner with PID**: adds a visual separator and process ID to
the startup message, making it easy to identify session boundaries in
long-running logs (e.g. `grep "Starting EXO"`).

The non-verbose (default) stderr format is unchanged — end-user terminal
experience is not affected.

## Motivation

When exo runs as a service (launchd, systemd, Docker), stderr is
typically captured to a file. Without calendar dates in the timestamp,
it is impossible to tell which day a log entry belongs to. This caused
misidentification of log entries during a multi-day RDMA debugging
session on a 4-node cluster.

The file log (`exo.log`) already had dates and rotation, but only
captured INFO level — missing the DEBUG output needed for postmortem
analysis.

## Test plan

- [x] `basedpyright` — 0 errors, 0 warnings, 0 notes
- [x] `ruff check` — all checks passed
- [x] `pytest` — 249 passed, 1 skipped

Co-authored-by: user <user@m1.note>
2026-03-23 16:06:26 +00:00
Evan Quiney 07598a3af1 teeny refactor (#1753)
api.py keeps growing. it's not tied to the master, so should have it's
own top level folder.

## testing
ci, pytest
2026-03-19 15:57:50 +00:00
ciaranbor 63f57fc193 Ciaran/dashboard download bug (#1755)
## Motivation

Download progress in the dashboard was broken: mainly treating all
download statuses as ongoing

## Changes

- Backend (apply.py): Deduplicate download progress events by model_id
instead of full shard_metadata, preventing duplicate entries per node
- Dashboard (+page.svelte): Extract shared collectDownloadStatus()
helper that both getModelDownloadStatus and getInstanceDownloadStatus
use, eliminating ~100 lines of duplicated logic. Adds proper handling
for
DownloadCompleted/DownloadFailed events, uses a Map to deduplicate
per-node entries, and introduces a typed NodeDownloadStatus with
explicit status states (downloading/completed/partial/pending)
- ModelCard: Replace single aggregate progress bar with per-node
download bars, each color-coded by status. Instance preview now scopes
download status to participating nodes only

## Why It Works

- Deduplicating by model_id in apply.py ensures each node has exactly
one download entry per model
- The perNodeMap in the frontend keeps only the latest event per node,
preventing duplicate bars
- Handling DownloadCompleted allows the UI to show finished downloads
instead of dropping them
- Scoping instance previews to assigned nodes avoids showing irrelevant
download progress
2026-03-19 14:47:45 +00:00
490 changed files with 5136 additions and 10383 deletions
-20
View File
@@ -1,20 +0,0 @@
from enum import Enum
class HarmonyEncodingName(Enum):
HARMONY_GPT_OSS = ...
class HarmonyEncoding: ...
class HarmonyError(Exception): ...
class Role(Enum):
ASSISTANT = ...
class StreamableParser:
last_content_delta: str
current_channel: str | None
current_recipient: str | None
def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
def process(self, token_id: int) -> None: ...
def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...
-17
View File
@@ -1,17 +0,0 @@
class NvmlMemoryInfo:
used: int
total: int
free: int
class NvmlUtilizationRates:
gpu: int
memory: int
def nvmlInit() -> None: ...
def nvmlShutdown() -> None: ...
def nvmlDeviceGetCount() -> int: ...
def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...
-61
View File
@@ -1,61 +0,0 @@
from typing import Any, Sequence
from torch import backends as backends
from torch import cuda as cuda
from torch import distributed as distributed
__version__: str
class version:
cuda: str
class dtype: ...
bfloat16: dtype
float16: dtype
float32: dtype
int8: dtype
int32: dtype
int64: dtype
long: dtype
float8_e4m3fn: dtype
class Tensor:
shape: Sequence[int]
dtype: dtype
def __getitem__(self, key: Any) -> Tensor: ...
def __setitem__(self, key: Any, value: Any) -> None: ...
def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
def cpu(self) -> Tensor: ...
def detach(self) -> Tensor: ...
def clone(self) -> Tensor: ...
def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
def view(self, *shape: Any) -> Tensor: ...
def squeeze(self, dim: int = ...) -> Tensor: ...
def unsqueeze(self, dim: int) -> Tensor: ...
def permute(self, *dims: int) -> Tensor: ...
def float(self) -> Tensor: ...
def numpy(self) -> Any: ...
def numel(self) -> int: ...
def nelement(self) -> int: ...
@property
def is_cuda(self) -> bool: ...
@property
def device(self) -> device: ...
def __len__(self) -> int: ...
def data_ptr(self) -> int: ...
def tolist(self) -> Any: ...
def abs(self) -> Tensor: ...
def max(self) -> Tensor: ...
def mean(self) -> Tensor: ...
def sum(self, dim: int = ...) -> Tensor: ...
def item(self) -> float: ...
def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def from_numpy(ndarray: Any) -> Tensor: ...
def inference_mode() -> Any: ...
class device:
def __init__(self, type: str, index: int = ...) -> None: ...
@@ -1 +0,0 @@
from torch.backends import cuda as cuda
@@ -1 +0,0 @@
def is_built() -> bool: ...
-10
View File
@@ -1,10 +0,0 @@
class _DeviceProperties:
total_memory: int
def is_available() -> bool: ...
def get_device_name(device: int) -> str: ...
def get_device_properties(device: int) -> _DeviceProperties: ...
def empty_cache() -> None: ...
def mem_get_info() -> tuple[int, int]: ...
def synchronize() -> None: ...
def max_memory_allocated() -> int: ...
@@ -1,2 +0,0 @@
def is_initialized() -> bool: ...
def destroy_process_group() -> None: ...
-1
View File
@@ -1 +0,0 @@
__version__: str
-2
View File
@@ -1,2 +0,0 @@
class ModelConfig:
max_model_len: int
-18
View File
@@ -1,18 +0,0 @@
from dataclasses import dataclass
@dataclass
class EngineArgs:
model: str = ...
served_model_name: str | list[str] | None = ...
tokenizer: str | None = ...
trust_remote_code: bool = ...
dtype: str = ...
seed: int = ...
max_model_len: int | None = ...
gpu_memory_utilization: float = ...
enforce_eager: bool = ...
tensor_parallel_size: int = ...
pipeline_parallel_size: int = ...
quantization: str | None = ...
load_format: str = ...
enable_sleep_mode: bool = ...
-17
View File
@@ -1,17 +0,0 @@
class CompletionOutput:
index: int
text: str
token_ids: list[int]
cumulative_logprob: float | None
logprobs: object | None
finish_reason: str | None
stop_reason: int | str | None
def finished(self) -> bool: ...
class RequestOutput:
request_id: str
prompt: str | None
prompt_token_ids: list[int] | None
outputs: list[CompletionOutput]
finished: bool
-11
View File
@@ -1,11 +0,0 @@
class SamplingParams:
n: int
temperature: float
top_p: float
top_k: int
min_p: float
seed: int | None
stop: str | list[str] | None
max_tokens: int | None
logprobs: int | None
repetition_penalty: float
@@ -1,3 +0,0 @@
from vllm.tokenizers.protocol import TokenizerLike
__all__ = ["TokenizerLike"]
@@ -1,15 +0,0 @@
from typing import Protocol
class TokenizerLike(Protocol):
@property
def eos_token_id(self) -> int: ...
@property
def vocab_size(self) -> int: ...
def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
def apply_chat_template(
self,
messages: list[dict[str, str]],
tools: list[dict[str, object]] | None = ...,
**kwargs: object,
) -> str | list[int]: ...
-1
View File
@@ -1 +0,0 @@
-1
View File
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
from collections.abc import Sequence
from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
from vllm.v1.kv_cache_interface import KVCacheConfig
class KVCacheBlocks:
blocks: tuple[Sequence[KVCacheBlock], ...]
def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
def get_block_ids(self) -> tuple[list[int], ...]: ...
class KVCacheManager:
block_pool: BlockPool
kv_cache_config: KVCacheConfig
enable_caching: bool
num_kv_cache_groups: int
coordinator: object
def __init__(self, *args: object, **kwargs: object) -> None: ...
def allocate_slots(
self, request: object, num_new_tokens: int, *args: object, **kwargs: object
) -> KVCacheBlocks | None: ...
def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks: ...
@@ -1,16 +0,0 @@
class KVCacheBlock:
block_id: int
ref_cnt: int
def __init__(self, block_id: int) -> None: ...
class FreeKVCacheBlockQueue:
def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
class BlockPool:
blocks: list[KVCacheBlock]
free_block_queue: FreeKVCacheBlockQueue
num_gpu_blocks: int
enable_caching: bool
def get_num_free_blocks(self) -> int: ...
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...
@@ -1,22 +0,0 @@
from vllm.config import ModelConfig
from vllm.engine.arg_utils import EngineArgs
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.tokenizers import TokenizerLike
class LLMEngine:
tokenizer: TokenizerLike | None
model_config: ModelConfig
@classmethod
def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
def add_request(
self,
request_id: str,
prompt: str,
params: SamplingParams,
arrival_time: float | None = ...,
) -> None: ...
def step(self) -> list[RequestOutput]: ...
def has_unfinished_requests(self) -> bool: ...
def get_tokenizer(self) -> TokenizerLike: ...
@@ -1,23 +0,0 @@
from dataclasses import dataclass
@dataclass
class KVCacheSpec:
block_size: int
num_kv_heads: int
head_size: int
@dataclass
class KVCacheGroupSpec:
layer_names: list[str]
kv_cache_spec: KVCacheSpec
@dataclass
class KVCacheTensorSpec:
shared_by: list[str]
size: int
@dataclass
class KVCacheConfig:
num_blocks: int
kv_cache_groups: list[KVCacheGroupSpec]
kv_cache_tensors: list[KVCacheTensorSpec]
-6
View File
@@ -1,6 +0,0 @@
class Request:
request_id: str
prompt_token_ids: list[int] | None
num_prompt_tokens: int
num_computed_tokens: int
num_tokens: int
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
import torch
class _CompilationConfig:
static_forward_context: dict[str, object]
class _ModelConfig:
hf_config: object
class GPUModelRunner:
kv_caches: list[torch.Tensor]
compilation_config: _CompilationConfig
model_config: _ModelConfig | None
def _allocate_kv_cache_tensors(
self, kv_cache_config: object
) -> dict[str, torch.Tensor]: ...
def initialize_kv_cache_tensors(
self, kv_cache_config: object, kernel_block_sizes: list[int]
) -> dict[str, torch.Tensor]: ...
def _reshape_kv_cache_tensors(
self,
kv_cache_config: object,
raw_tensors: dict[str, torch.Tensor],
kernel_block_sizes: list[int],
) -> dict[str, torch.Tensor]: ...
@@ -1,6 +0,0 @@
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
class Worker:
model_runner: GPUModelRunner
def determine_available_memory(self) -> int: ...
def initialize_from_config(self, kv_cache_config: object) -> None: ...
-1
View File
@@ -1 +0,0 @@
def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from typing import Protocol
import mlx.core as mx
import PIL.Image
import tqdm
from typing import Protocol
from mflux.models.common.config.config import Config
class BeforeLoopCallback(Protocol):
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.callbacks.callback import (
AfterLoopCallback,
BeforeLoopCallback,
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
import PIL.Image
import tqdm
from typing import TYPE_CHECKING
from mflux.callbacks.callback_registry import CallbackRegistry
from mflux.models.common.config.config import Config
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import Any
import mlx.core as mx
from mflux.models.common.config.model_config import ModelConfig
from tqdm import tqdm
from mflux.models.common.config.model_config import ModelConfig
logger = ...
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from functools import lru_cache
from typing import Literal
import mlx.core as mx
class ModelConfig:
precision: mx.Dtype = ...
def __init__(
@@ -2,10 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from pathlib import Path
from typing import TYPE_CHECKING, TypeAlias
import mlx.core as mx
from mlx import nn
from mflux.models.common.vae.tiling_config import TilingConfig
from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
@@ -13,7 +13,6 @@ from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreat
from mflux.models.z_image.latent_creator.z_image_latent_creator import (
ZImageLatentCreator,
)
from mlx import nn
if TYPE_CHECKING:
LatentCreatorType: TypeAlias = type[
@@ -2,8 +2,8 @@
This type stub file was generated by pyright.
"""
from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
from mlx import nn
from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
class FusedLoRALinear(nn.Module):
def __init__(
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
from collections.abc import Callable
from dataclasses import dataclass
import mlx.core as mx
import mlx.nn as nn
from collections.abc import Callable
from dataclasses import dataclass
from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
@dataclass
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from collections.abc import Callable
from dataclasses import dataclass
from typing import List, Protocol
import mlx.core as mx
@dataclass
class LoRATarget:
model_path: str
@@ -36,3 +36,4 @@ class Rule(NamedTuple):
name: str
check: str
action: QuantizationAction | PathAction | LoraAction | ConfigAction
...
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.config.model_config import ModelConfig
if TYPE_CHECKING: ...
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from abc import ABC, abstractmethod
import mlx.core as mx
from abc import ABC, abstractmethod
class BaseScheduler(ABC):
@property
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.core as mx
from typing import TYPE_CHECKING
from mflux.models.common.config.config import Config
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
@@ -4,10 +4,9 @@ This type stub file was generated by pyright.
from abc import ABC, abstractmethod
from typing import Protocol, runtime_checkable
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
from PIL import Image
from transformers import PreTrainedTokenizer
from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput
"""
This type stub file was generated by pyright.
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.common.weights.loading.weight_definition import TokenizerDefinition
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from dataclasses import dataclass
import mlx.core as mx
from dataclasses import dataclass
"""
This type stub file was generated by pyright.
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import Callable
import mlx.core as mx
from typing import Callable
class VAETiler:
@staticmethod
@@ -3,8 +3,8 @@ This type stub file was generated by pyright.
"""
import mlx.core as mx
from mflux.models.common.vae.tiling_config import TilingConfig
from mlx import nn
from mflux.models.common.vae.tiling_config import TilingConfig
class VAEUtil:
@staticmethod
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
import mlx.nn as nn
from typing import TYPE_CHECKING
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,12 +2,11 @@
This type stub file was generated by pyright.
"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, List, TypeAlias
import mlx.core as mx
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from dataclasses import dataclass
from typing import Callable, List, TYPE_CHECKING, TypeAlias
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
from mflux.models.common.tokenizer.tokenizer import BaseTokenizer
from mflux.models.depth_pro.weights.depth_pro_weight_definition import (
DepthProWeightDefinition,
)
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
@@ -2,9 +2,8 @@
This type stub file was generated by pyright.
"""
from typing import Dict, List, Optional
import mlx.core as mx
from typing import Dict, List, Optional
from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
class WeightMapper:
@@ -2,11 +2,10 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, Optional, Protocol
import mlx.core as mx
"""
This type stub file was generated by pyright.
"""
@@ -2,8 +2,7 @@
This type stub file was generated by pyright.
"""
from typing import TYPE_CHECKING, Any
from typing import Any, TYPE_CHECKING
from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType
if TYPE_CHECKING: ...
@@ -2,10 +2,9 @@
This type stub file was generated by pyright.
"""
import mlx.core as mx
from dataclasses import dataclass
from pathlib import Path
import mlx.core as mx
from PIL import Image
@dataclass
@@ -14,6 +13,7 @@ class DepthResult:
depth_array: mx.array
min_depth: float
max_depth: float
...
class DepthPro:
def __init__(self, quantize: int | None = ...) -> None: ...
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.loading.weight_definition import (
ComponentDefinition,
TokenizerDefinition,
@@ -3,7 +3,6 @@ This type stub file was generated by pyright.
"""
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import (
WeightMapping,
WeightTarget,

Some files were not shown because too many files have changed in this diff Show More