Compare commits
75 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9f184de322 | |||
| 8485805042 | |||
| 4de8f801c7 | |||
| 5777bf3c39 | |||
| 886192f1e6 | |||
| d914acd64e | |||
| 37296c8249 | |||
| 28817d3ee3 | |||
| 0e1b9501c3 | |||
| f0d4ccbeb3 | |||
| 858dc808df | |||
| 635118ef24 | |||
| dc0bb5e13b | |||
| 152a27ea5d | |||
| db36bd5ac6 | |||
| 639243aa09 | |||
| db73c4fd5d | |||
| eaed92952c | |||
| ba611f9cd0 | |||
| eab3e0b456 | |||
| c4e874e97d | |||
| e23c3a3026 | |||
| 190e63e56d | |||
| 7660893538 | |||
| 9a2d2a4a7c | |||
| bea64fe85e | |||
| 14526d281a | |||
| 73e50df827 | |||
| 2b417f28ff | |||
| b65982ddd7 | |||
| 2fe689315b | |||
| 644c5573ce | |||
| 12c3015f52 | |||
| 365dd68d9a | |||
| d3d129581e | |||
| c90a0cec78 | |||
| e8c1337168 | |||
| 7024ddcf3e | |||
| dc89ba662a | |||
| 5cd96b5060 | |||
| 05986f77aa | |||
| dab7ed4821 | |||
| 2261014715 | |||
| 61d2a2b6cf | |||
| 0ff99a2c40 | |||
| fbb80e1cc9 | |||
| 8d94eab6c6 | |||
| f370452d7e | |||
| a4c2aa2b87 | |||
| 7312c535b4 | |||
| 18717023ad | |||
| 1780e4ade4 | |||
| ab9273e723 | |||
| 71e48c0f62 | |||
| 42da58c297 | |||
| 6b5a705959 | |||
| 6b54a27019 | |||
| e01f50a5cd | |||
| 1093080214 | |||
| 1a2b8b044a | |||
| dc8d42b4dc | |||
| d484b062e8 | |||
| e32b649d2f | |||
| bddad7e79c | |||
| addf73a144 | |||
| a16ff2c047 | |||
| 3006c8ea4e | |||
| f662c129dd | |||
| c45ff9ad43 | |||
| 7031901ae5 | |||
| cf648a53b8 | |||
| 94b2ce6922 | |||
| 423ed0f07f | |||
| ed001f2409 | |||
| 4c4c6ce99f |
@@ -215,6 +215,22 @@ class StreamContext:
|
||||
traceback: object | None = ...,
|
||||
) -> None: ...
|
||||
|
||||
def device_info() -> dict[str, str | int]:
|
||||
"""
|
||||
Get information about the GPU device and system settings.
|
||||
|
||||
Currently returns:
|
||||
|
||||
* ``architecture``
|
||||
* ``max_buffer_size``
|
||||
* ``max_recommended_working_set_size``
|
||||
* ``memory_size``
|
||||
* ``resource_limit``
|
||||
|
||||
Returns:
|
||||
dict: A dictionary with string keys and string or integer values.
|
||||
"""
|
||||
|
||||
def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
|
||||
"""
|
||||
Element-wise absolute value.
|
||||
|
||||
@@ -73,9 +73,11 @@ class GenerationResponse:
|
||||
finish_reason: Optional[str] = ...
|
||||
|
||||
def maybe_quantize_kv_cache(
|
||||
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
|
||||
): # -> None:
|
||||
...
|
||||
prompt_cache: Any,
|
||||
quantized_kv_start: int | None,
|
||||
kv_group_size: int | None,
|
||||
kv_bits: int | None,
|
||||
) -> None: ...
|
||||
def generate_step(
|
||||
prompt: mx.array,
|
||||
model: nn.Module,
|
||||
|
||||
@@ -16,7 +16,7 @@ class Cache(Protocol):
|
||||
self, keys: mx.array, values: mx.array
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
|
||||
@@ -92,13 +92,14 @@ class _BaseCache(Cache):
|
||||
values: mx.array
|
||||
offset: int
|
||||
@property
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
@property
|
||||
def meta_state(self) -> Literal[""]: ...
|
||||
@meta_state.setter
|
||||
def meta_state(self, v) -> None: ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def is_trimmable(self) -> Literal[False]: ...
|
||||
@classmethod
|
||||
def from_state(cls, state, meta_state) -> Self: ...
|
||||
@@ -114,15 +115,13 @@ class ConcatenateKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]:
|
||||
...
|
||||
@property
|
||||
def state(self): # -> tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -132,10 +131,7 @@ class QuantizedKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> Any:
|
||||
...
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
): # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -147,8 +143,7 @@ class QuantizedKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -160,18 +155,18 @@ class KVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
) -> tuple[array, array]: ...
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
def make_mask(
|
||||
self, *args: Any, **kwargs: Any
|
||||
) -> mx.array | Literal["causal"] | None: ...
|
||||
|
||||
class RotatingKVCache(_BaseCache):
|
||||
step = ...
|
||||
@@ -183,8 +178,7 @@ class RotatingKVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
): # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -196,8 +190,7 @@ class RotatingKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
@@ -212,8 +205,7 @@ class ArraysCache(_BaseCache):
|
||||
...
|
||||
def __getitem__(self, idx): ...
|
||||
@property
|
||||
def state(self): # -> list[Any | array] | list[array]:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -227,8 +219,7 @@ class ArraysCache(_BaseCache):
|
||||
In-place extend this cache with the other cache.
|
||||
"""
|
||||
|
||||
def make_mask(self, N: int): # -> array | None:
|
||||
...
|
||||
def make_mask(self, N: int) -> mx.array | None: ...
|
||||
|
||||
class MambaCache(ArraysCache):
|
||||
def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
|
||||
@@ -239,8 +230,7 @@ class ChunkedKVCache(KVCache):
|
||||
...
|
||||
def update_and_fetch(self, keys, values): # -> tuple[array, array]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
@property
|
||||
def meta_state(self): # -> tuple[str, ...]:
|
||||
...
|
||||
@@ -253,10 +243,9 @@ class CacheList(_BaseCache):
|
||||
def __getitem__(self, idx): ...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n): ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
@property
|
||||
def state(self): # -> list[Any]:
|
||||
...
|
||||
def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_next import (
|
||||
Qwen3NextAttention as Attention,
|
||||
Qwen3NextMLP as MLP,
|
||||
Qwen3NextRMSNormGated as RMSNormGated,
|
||||
Qwen3NextSparseMoeBlock,
|
||||
)
|
||||
|
||||
SparseMoeBlock = Qwen3NextSparseMoeBlock
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class TextModelArgs:
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int
|
||||
linear_num_value_heads: int
|
||||
linear_num_key_heads: int
|
||||
linear_key_head_dim: int
|
||||
linear_value_head_dim: int
|
||||
linear_conv_kernel_dim: int
|
||||
tie_word_embeddings: bool
|
||||
attention_bias: bool
|
||||
head_dim: Optional[int]
|
||||
full_attention_interval: int
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
decoder_sparse_step: int
|
||||
shared_expert_intermediate_size: int
|
||||
moe_intermediate_size: int
|
||||
norm_topk_prob: bool
|
||||
rope_parameters: Optional[dict[str, Any]]
|
||||
partial_rotary_factor: float
|
||||
rope_theta: float
|
||||
rope_scaling: Optional[dict[str, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class GatedDeltaNet(nn.Module):
|
||||
hidden_size: int
|
||||
num_v_heads: int
|
||||
num_k_heads: int
|
||||
head_k_dim: int
|
||||
head_v_dim: int
|
||||
key_dim: int
|
||||
value_dim: int
|
||||
conv_kernel_size: int
|
||||
conv_dim: int
|
||||
conv1d: nn.Conv1d
|
||||
in_proj_qkv: nn.Linear
|
||||
in_proj_z: nn.Linear
|
||||
in_proj_b: nn.Linear
|
||||
in_proj_a: nn.Linear
|
||||
dt_bias: mx.array
|
||||
A_log: mx.array
|
||||
norm: RMSNormGated
|
||||
out_proj: nn.Linear
|
||||
|
||||
def __init__(self, config: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
is_linear: bool
|
||||
linear_attn: GatedDeltaNet
|
||||
self_attn: Attention
|
||||
input_layernorm: nn.RMSNorm
|
||||
post_attention_layernorm: nn.RMSNorm
|
||||
mlp: MLP | SparseMoeBlock
|
||||
|
||||
def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3_5TextModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[DecoderLayer]
|
||||
norm: nn.RMSNorm
|
||||
ssm_idx: int
|
||||
fa_idx: int
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class TextModel(nn.Module):
|
||||
args: TextModelArgs
|
||||
model_type: str
|
||||
model: Qwen3_5TextModel
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
language_model: TextModel
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
@@ -0,0 +1,19 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(Qwen3_5Model):
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@@ -7,6 +7,15 @@ import mlx.nn as nn
|
||||
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
class Qwen3NextRMSNormGated(nn.Module):
|
||||
eps: float
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
|
||||
def __call__(
|
||||
self, hidden_states: mx.array, gate: mx.array | None = None
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3NextMLP(nn.Module):
|
||||
gate_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
|
||||
Generated
+2258
-1831
File diff suppressed because it is too large
Load Diff
+1
-2
@@ -3,7 +3,6 @@ resolver = "3"
|
||||
members = [
|
||||
"rust/networking",
|
||||
"rust/exo_pyo3_bindings",
|
||||
"rust/util",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
@@ -24,7 +23,6 @@ opt-level = 3
|
||||
[workspace.dependencies]
|
||||
## Crate members as common dependencies
|
||||
networking = { path = "rust/networking" }
|
||||
util = { path = "rust/util" }
|
||||
|
||||
# Macro dependecies
|
||||
extend = "1.2"
|
||||
@@ -34,6 +32,7 @@ delegate = "0.13"
|
||||
keccak-const = "0.2"
|
||||
|
||||
# Async dependencies
|
||||
async-stream = "0.3"
|
||||
tokio = "1.46"
|
||||
futures-lite = "2.6.1"
|
||||
futures-timer = "3.0"
|
||||
|
||||
+215
-148
@@ -15,14 +15,23 @@ struct ContentView: View {
|
||||
@EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var settingsWindowController: SettingsWindowController
|
||||
@State private var focusedNode: NodeViewModel?
|
||||
@State private var deletingInstanceIDs: Set<String> = []
|
||||
@State private var showAllNodes = false
|
||||
@State private var showAllInstances = false
|
||||
@State private var baseURLCopied = false
|
||||
@State private var showAdvanced = false
|
||||
@State private var showDebugInfo = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
private enum BugReportPhase: Equatable {
|
||||
case idle
|
||||
case prompting
|
||||
case sending(String)
|
||||
case success(String)
|
||||
case failure(String)
|
||||
}
|
||||
@State private var bugReportPhase: BugReportPhase = .idle
|
||||
@State private var bugReportUserDescription: String = ""
|
||||
@State private var uninstallInProgress = false
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@@ -258,139 +267,79 @@ struct ContentView: View {
|
||||
VStack(alignment: .leading, spacing: 0) {
|
||||
if controller.status != .stopped {
|
||||
dashboardButton
|
||||
baseURLRow
|
||||
Divider()
|
||||
.padding(.vertical, 8)
|
||||
} else {
|
||||
Divider()
|
||||
.padding(.vertical, 4)
|
||||
}
|
||||
advancedSection
|
||||
.padding(.bottom, 8)
|
||||
controlButton(title: "Quit", tint: .secondary) {
|
||||
HoverButton(
|
||||
title: "Settings",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "gear"
|
||||
) {
|
||||
settingsWindowController.open(
|
||||
controller: controller,
|
||||
updater: updater,
|
||||
networkStatusService: networkStatusService,
|
||||
thunderboltBridgeService: thunderboltBridgeService,
|
||||
stateService: stateService
|
||||
)
|
||||
}
|
||||
HoverButton(
|
||||
title: "Check for Updates",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.triangle.2.circlepath"
|
||||
) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
.padding(.bottom, 8)
|
||||
HoverButton(title: "Quit", tint: .secondary) {
|
||||
controller.stop()
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var advancedSection: some View {
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
HStack {
|
||||
Text("Advanced")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
Spacer()
|
||||
collapseButton(isExpanded: $showAdvanced)
|
||||
}
|
||||
.animation(nil, value: showAdvanced)
|
||||
if showAdvanced {
|
||||
VStack(alignment: .leading, spacing: 8) {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("Cluster Namespace")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
TextField("optional", text: $pendingNamespace)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.customNamespace = pendingNamespace
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingNamespace == controller.customNamespace)
|
||||
}
|
||||
}
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("HuggingFace Token")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
SecureField("optional", text: $pendingHFToken)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingHFToken = controller.hfToken
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.hfToken = pendingHFToken
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingHFToken == controller.hfToken)
|
||||
}
|
||||
}
|
||||
Divider()
|
||||
HStack {
|
||||
Toggle(
|
||||
"Enable Image Models (experimental)", isOn: $pendingEnableImageModels
|
||||
)
|
||||
.toggleStyle(.switch)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
Button("Save & Restart") {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingEnableImageModels == controller.enableImageModels)
|
||||
}
|
||||
HoverButton(title: "Check for Updates", small: true) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
debugSection
|
||||
HoverButton(title: "Uninstall", tint: .red, small: true) {
|
||||
showUninstallConfirmationAlert()
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
.transition(.opacity)
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.25), value: showAdvanced)
|
||||
}
|
||||
|
||||
private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void)
|
||||
-> some View
|
||||
{
|
||||
HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action)
|
||||
}
|
||||
|
||||
private var dashboardButton: some View {
|
||||
Button {
|
||||
HoverButton(
|
||||
title: "Web Dashboard",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.up.right"
|
||||
) {
|
||||
guard let url = URL(string: "http://localhost:52415/") else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
} label: {
|
||||
HStack {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
.imageScale(.small)
|
||||
Text("Dashboard")
|
||||
.fontWeight(.medium)
|
||||
Spacer()
|
||||
}
|
||||
.padding(.vertical, 8)
|
||||
.padding(.horizontal, 10)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 8, style: .continuous)
|
||||
.fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.padding(.bottom, 4)
|
||||
}
|
||||
|
||||
private var baseURLRow: some View {
|
||||
HStack(spacing: 6) {
|
||||
Image(systemName: "link")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(.secondary)
|
||||
Text("localhost:52415/v1")
|
||||
.font(.system(.caption, design: .monospaced))
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
NSPasteboard.general.clearContents()
|
||||
NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string)
|
||||
baseURLCopied = true
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
|
||||
baseURLCopied = false
|
||||
}
|
||||
} label: {
|
||||
Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(baseURLCopied ? .green : .secondary)
|
||||
.contentTransition(.symbolEffect(.replace))
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.help("Copy API base URL")
|
||||
}
|
||||
.padding(.vertical, 4)
|
||||
.padding(.horizontal, 8)
|
||||
}
|
||||
|
||||
private func collapseButton(isExpanded: Binding<Bool>) -> some View {
|
||||
@@ -611,39 +560,115 @@ struct ContentView: View {
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Button {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
if bugReportInFlight {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
switch bugReportPhase {
|
||||
case .idle:
|
||||
Button {
|
||||
bugReportPhase = .prompting
|
||||
bugReportUserDescription = ""
|
||||
} label: {
|
||||
HStack {
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
}
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
.padding(.vertical, 6)
|
||||
.padding(.horizontal, 8)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color.accentColor.opacity(0.12))
|
||||
)
|
||||
}
|
||||
.padding(.vertical, 6)
|
||||
.padding(.horizontal, 8)
|
||||
.buttonStyle(.plain)
|
||||
|
||||
case .prompting:
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
Text("What's the issue? (optional)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
TextEditor(text: $bugReportUserDescription)
|
||||
.font(.caption2)
|
||||
.frame(height: 60)
|
||||
.overlay(
|
||||
RoundedRectangle(cornerRadius: 4)
|
||||
.stroke(Color.secondary.opacity(0.3), lineWidth: 1)
|
||||
)
|
||||
HStack(spacing: 8) {
|
||||
Button("Send") {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.borderedProminent)
|
||||
.controlSize(.small)
|
||||
Button("Cancel") {
|
||||
bugReportPhase = .idle
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.bordered)
|
||||
.controlSize(.small)
|
||||
}
|
||||
}
|
||||
.padding(8)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color.accentColor.opacity(0.12))
|
||||
.fill(Color.accentColor.opacity(0.06))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.disabled(bugReportInFlight)
|
||||
|
||||
if let message = bugReportMessage {
|
||||
Text(message)
|
||||
case .sending(let message):
|
||||
HStack(spacing: 6) {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
case .success(let message):
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
Button {
|
||||
openGitHubIssue()
|
||||
} label: {
|
||||
HStack(spacing: 4) {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
.imageScale(.small)
|
||||
Text("Create GitHub Issue")
|
||||
.font(.caption2)
|
||||
}
|
||||
}
|
||||
.buttonStyle(.bordered)
|
||||
.controlSize(.small)
|
||||
Button("Done") {
|
||||
bugReportPhase = .idle
|
||||
bugReportUserDescription = ""
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.plain)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
|
||||
case .failure(let message):
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.red)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
Button("Dismiss") {
|
||||
bugReportPhase = .idle
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.plain)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.2), value: bugReportPhase)
|
||||
}
|
||||
|
||||
private var processToggleBinding: Binding<Bool> {
|
||||
@@ -687,16 +712,58 @@ struct ContentView: View {
|
||||
}
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportInFlight = true
|
||||
bugReportMessage = "Collecting logs..."
|
||||
bugReportPhase = .sending("Collecting logs...")
|
||||
let service = BugReportService()
|
||||
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
do {
|
||||
let outcome = try await service.sendReport(isManual: true)
|
||||
bugReportMessage = outcome.message
|
||||
let outcome = try await service.sendReport(
|
||||
isManual: true,
|
||||
userDescription: description.isEmpty ? nil : description
|
||||
)
|
||||
if outcome.success {
|
||||
bugReportPhase = .success(outcome.message)
|
||||
} else {
|
||||
bugReportPhase = .failure(outcome.message)
|
||||
}
|
||||
} catch {
|
||||
bugReportMessage = error.localizedDescription
|
||||
bugReportPhase = .failure(error.localizedDescription)
|
||||
}
|
||||
}
|
||||
|
||||
private func openGitHubIssue() {
|
||||
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
|
||||
var bodyParts: [String] = []
|
||||
bodyParts.append("## Describe the bug")
|
||||
bodyParts.append("")
|
||||
if !description.isEmpty {
|
||||
bodyParts.append(description)
|
||||
} else {
|
||||
bodyParts.append("A clear and concise description of what the bug is.")
|
||||
}
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Environment")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
|
||||
bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Additional context")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
|
||||
|
||||
let body = bodyParts.joined(separator: "\n")
|
||||
|
||||
var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
|
||||
components.queryItems = [
|
||||
URLQueryItem(name: "template", value: "bug_report.md"),
|
||||
URLQueryItem(name: "title", value: "[BUG] "),
|
||||
URLQueryItem(name: "body", value: body),
|
||||
URLQueryItem(name: "labels", value: "bug"),
|
||||
]
|
||||
|
||||
if let url = components.url {
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
bugReportInFlight = false
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
|
||||
@@ -21,7 +21,9 @@ struct EXOApp: App {
|
||||
@StateObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@StateObject private var updater: SparkleUpdater
|
||||
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@StateObject private var settingsWindowController: SettingsWindowController
|
||||
private let terminationObserver: TerminationObserver
|
||||
private let firstLaunchPopout = FirstLaunchPopout()
|
||||
private let ciContext = CIContext(options: nil)
|
||||
|
||||
init() {
|
||||
@@ -43,12 +45,13 @@ struct EXOApp: App {
|
||||
_updater = StateObject(wrappedValue: updater)
|
||||
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
|
||||
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
|
||||
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
|
||||
enableLaunchAtLoginIfNeeded()
|
||||
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
|
||||
NetworkSetupHelper.promptAndInstallIfNeeded()
|
||||
// Check local network access periodically (warning disappears when user grants permission)
|
||||
localNetwork.startPeriodicChecking(interval: 10)
|
||||
controller.scheduleLaunch(after: 15)
|
||||
controller.scheduleLaunch(after: 5)
|
||||
service.startPolling()
|
||||
networkStatus.startPolling()
|
||||
}
|
||||
@@ -62,8 +65,19 @@ struct EXOApp: App {
|
||||
.environmentObject(localNetworkChecker)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(settingsWindowController)
|
||||
} label: {
|
||||
menuBarIcon
|
||||
.onReceive(controller.$isFirstLaunchReady) { ready in
|
||||
if ready {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 3.0) {
|
||||
self.firstLaunchPopout.onComplete = { [weak controller] in
|
||||
controller?.markOnboardingCompleted()
|
||||
}
|
||||
self.firstLaunchPopout.show()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.menuBarExtraStyle(.window)
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ import Foundation
|
||||
private let customNamespaceKey = "EXOCustomNamespace"
|
||||
private let hfTokenKey = "EXOHFToken"
|
||||
private let enableImageModelsKey = "EXOEnableImageModels"
|
||||
private let offlineModeKey = "EXOOfflineMode"
|
||||
private let onboardingCompletedKey = "EXOOnboardingCompleted"
|
||||
|
||||
@MainActor
|
||||
final class ExoProcessController: ObservableObject {
|
||||
@@ -59,6 +61,17 @@ final class ExoProcessController: ObservableObject {
|
||||
UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
|
||||
}
|
||||
}
|
||||
@Published var offlineMode: Bool = {
|
||||
return UserDefaults.standard.bool(forKey: offlineModeKey)
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
|
||||
}
|
||||
}
|
||||
|
||||
/// Fires once when EXO transitions to `.running` for the very first time (fresh install).
|
||||
@Published private(set) var isFirstLaunchReady = false
|
||||
|
||||
private var process: Process?
|
||||
private var runtimeDirectoryURL: URL?
|
||||
@@ -113,6 +126,9 @@ final class ExoProcessController: ObservableObject {
|
||||
try child.run()
|
||||
process = child
|
||||
status = .running
|
||||
|
||||
// Show welcome popout on every launch
|
||||
isFirstLaunchReady = true
|
||||
} catch {
|
||||
process = nil
|
||||
status = .failed(message: "Launch error")
|
||||
@@ -164,6 +180,17 @@ final class ExoProcessController: ObservableObject {
|
||||
launch()
|
||||
}
|
||||
|
||||
/// Mark onboarding as completed (user interacted with the welcome popout).
|
||||
func markOnboardingCompleted() {
|
||||
UserDefaults.standard.set(true, forKey: onboardingCompletedKey)
|
||||
}
|
||||
|
||||
/// Reset onboarding so the welcome popout appears on next launch.
|
||||
func resetOnboarding() {
|
||||
UserDefaults.standard.removeObject(forKey: onboardingCompletedKey)
|
||||
isFirstLaunchReady = false
|
||||
}
|
||||
|
||||
func scheduleLaunch(after seconds: TimeInterval) {
|
||||
cancelPendingLaunch()
|
||||
let start = max(1, Int(ceil(seconds)))
|
||||
@@ -249,6 +276,9 @@ final class ExoProcessController: ObservableObject {
|
||||
if enableImageModels {
|
||||
environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
|
||||
}
|
||||
if offlineMode {
|
||||
environment["EXO_OFFLINE"] = "true"
|
||||
}
|
||||
|
||||
var paths: [String] = []
|
||||
if let existing = environment["PATH"], !existing.isEmpty {
|
||||
|
||||
@@ -38,7 +38,8 @@ struct BugReportService {
|
||||
func sendReport(
|
||||
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
|
||||
now: Date = Date(),
|
||||
isManual: Bool = false
|
||||
isManual: Bool = false,
|
||||
userDescription: String? = nil
|
||||
) async throws -> BugReportOutcome {
|
||||
let timestamp = Self.runTimestampString(now)
|
||||
let dayPrefix = Self.dayPrefixString(now)
|
||||
@@ -60,7 +61,8 @@ struct BugReportService {
|
||||
ifconfig: ifconfigText,
|
||||
debugInfo: debugInfo,
|
||||
isManual: isManual,
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus,
|
||||
userDescription: userDescription
|
||||
)
|
||||
|
||||
let eventLogFiles = readAllEventLogs()
|
||||
@@ -306,7 +308,8 @@ struct BugReportService {
|
||||
ifconfig: String,
|
||||
debugInfo: DebugInfo,
|
||||
isManual: Bool,
|
||||
clusterTbBridgeStatus: [[String: Any]]?
|
||||
clusterTbBridgeStatus: [[String: Any]]?,
|
||||
userDescription: String? = nil
|
||||
) -> Data? {
|
||||
let system = readSystemMetadata()
|
||||
let exo = readExoMetadata()
|
||||
@@ -323,6 +326,9 @@ struct BugReportService {
|
||||
if let tbStatus = clusterTbBridgeStatus {
|
||||
payload["cluster_thunderbolt_bridge"] = tbStatus
|
||||
}
|
||||
if let desc = userDescription, !desc.isEmpty {
|
||||
payload["user_description"] = desc
|
||||
}
|
||||
return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// A popover callout anchored to the menu bar icon on every launch,
|
||||
/// pointing the user to the web dashboard with an arrow connecting to the icon.
|
||||
@MainActor
|
||||
final class FirstLaunchPopout {
|
||||
private var popover: NSPopover?
|
||||
private var countdownTask: Task<Void, Never>?
|
||||
private static let dashboardURL = "http://localhost:52415/"
|
||||
|
||||
/// Called when the user completes onboarding (clicks Open Dashboard or dismisses).
|
||||
var onComplete: (() -> Void)?
|
||||
|
||||
func show() {
|
||||
guard popover == nil else { return }
|
||||
|
||||
// The status bar button may not exist yet on first launch; retry generously.
|
||||
showWithRetry(attemptsRemaining: 15)
|
||||
}
|
||||
|
||||
private func showWithRetry(attemptsRemaining: Int) {
|
||||
guard attemptsRemaining > 0 else {
|
||||
// Exhausted retries — fall back to just opening the dashboard directly.
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
return
|
||||
}
|
||||
|
||||
guard let button = Self.findStatusItemButton() else {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in
|
||||
self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
let pop = NSPopover()
|
||||
pop.behavior = .applicationDefined
|
||||
pop.animates = true
|
||||
pop.contentSize = NSSize(width: 280, height: 120)
|
||||
pop.contentViewController = NSHostingController(
|
||||
rootView: WelcomeCalloutView(
|
||||
countdownDuration: 10,
|
||||
onDismiss: { [weak self] in
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
},
|
||||
onOpen: { [weak self] in
|
||||
self?.openDashboard()
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
self.popover = pop
|
||||
pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
|
||||
|
||||
// Auto-open dashboard after 10s then dismiss
|
||||
countdownTask = Task {
|
||||
try? await Task.sleep(nanoseconds: 10_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
dismiss()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func dismiss() {
|
||||
countdownTask?.cancel()
|
||||
countdownTask = nil
|
||||
guard let pop = popover else { return }
|
||||
popover = nil
|
||||
pop.performClose(nil)
|
||||
}
|
||||
|
||||
private func openDashboard() {
|
||||
guard let url = URL(string: Self.dashboardURL) else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
|
||||
/// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra.
|
||||
/// Walks the view hierarchy to find the actual button rather than the content view.
|
||||
private static func findStatusItemButton() -> NSView? {
|
||||
for window in NSApp.windows {
|
||||
let className = NSStringFromClass(type(of: window))
|
||||
// Match NSStatusBarWindow or any internal SwiftUI status bar window
|
||||
guard className.contains("StatusBar") || className.contains("MenuBarExtra") else {
|
||||
continue
|
||||
}
|
||||
if let content = window.contentView {
|
||||
if let button = findButton(in: content) {
|
||||
return button
|
||||
}
|
||||
// Fall back to the content view itself if it has a non-zero frame
|
||||
if content.frame.width > 0 {
|
||||
return content
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
/// Recursively searches the view hierarchy for an NSStatusBarButton.
|
||||
private static func findButton(in view: NSView) -> NSView? {
|
||||
let className = NSStringFromClass(type(of: view))
|
||||
if className.contains("StatusBarButton") || className.contains("StatusItem") {
|
||||
return view
|
||||
}
|
||||
for subview in view.subviews {
|
||||
if let found = findButton(in: subview) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal welcome callout — friendly pointer, not a wall of text.
|
||||
/// Rendered inside the NSPopover which provides its own chrome and arrow.
|
||||
private struct WelcomeCalloutView: View {
|
||||
let countdownDuration: Int
|
||||
let onDismiss: () -> Void
|
||||
let onOpen: () -> Void
|
||||
@State private var countdown: Int
|
||||
@State private var timerTask: Task<Void, Never>?
|
||||
|
||||
init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) {
|
||||
self.countdownDuration = countdownDuration
|
||||
self.onDismiss = onDismiss
|
||||
self.onOpen = onOpen
|
||||
self._countdown = State(initialValue: countdownDuration)
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading, spacing: 10) {
|
||||
HStack(alignment: .top) {
|
||||
Text("EXO is running")
|
||||
.font(.system(.headline, design: .rounded))
|
||||
.fontWeight(.semibold)
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
onDismiss()
|
||||
} label: {
|
||||
Image(systemName: "xmark.circle.fill")
|
||||
.font(.system(size: 14))
|
||||
.foregroundStyle(.tertiary)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
}
|
||||
|
||||
Text("Run your first model here:")
|
||||
.font(.system(.subheadline, design: .default))
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Button {
|
||||
onOpen()
|
||||
} label: {
|
||||
Label("Open Dashboard", systemImage: "arrow.up.right.square")
|
||||
.font(.system(.caption, design: .default))
|
||||
.fontWeight(.medium)
|
||||
}
|
||||
.buttonStyle(.borderedProminent)
|
||||
.tint(.accentColor)
|
||||
.controlSize(.small)
|
||||
|
||||
Spacer()
|
||||
|
||||
if countdown > 0 {
|
||||
Text("Launching in \(countdown) secs...")
|
||||
.font(.system(.caption2, design: .default))
|
||||
.foregroundColor(.secondary.opacity(0.6))
|
||||
.monospacedDigit()
|
||||
}
|
||||
}
|
||||
}
|
||||
.padding(14)
|
||||
.onAppear {
|
||||
startCountdown()
|
||||
}
|
||||
.onDisappear {
|
||||
timerTask?.cancel()
|
||||
timerTask = nil
|
||||
}
|
||||
}
|
||||
|
||||
private func startCountdown() {
|
||||
timerTask = Task {
|
||||
while countdown > 0 {
|
||||
try? await Task.sleep(nanoseconds: 1_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
countdown -= 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Native macOS Settings window following Apple HIG.
|
||||
/// Organized into General, Model, Advanced, and About sections.
|
||||
struct SettingsView: View {
|
||||
@EnvironmentObject private var controller: ExoProcessController
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var networkStatusService: NetworkStatusService
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var stateService: ClusterStateService
|
||||
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@State private var pendingEnableImageModels = false
|
||||
@State private var pendingOfflineMode = false
|
||||
@State private var needsRestart = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
@State private var uninstallInProgress = false
|
||||
|
||||
var body: some View {
|
||||
TabView {
|
||||
generalTab
|
||||
.tabItem {
|
||||
Label("General", systemImage: "gear")
|
||||
}
|
||||
modelTab
|
||||
.tabItem {
|
||||
Label("Model", systemImage: "cube")
|
||||
}
|
||||
advancedTab
|
||||
.tabItem {
|
||||
Label("Advanced", systemImage: "wrench.and.screwdriver")
|
||||
}
|
||||
aboutTab
|
||||
.tabItem {
|
||||
Label("About", systemImage: "info.circle")
|
||||
}
|
||||
}
|
||||
.frame(width: 450, height: 400)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
pendingHFToken = controller.hfToken
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
pendingOfflineMode = controller.offlineMode
|
||||
needsRestart = false
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - General Tab
|
||||
|
||||
private var generalTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Cluster Namespace") {
|
||||
TextField("default", text: $pendingNamespace)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 200)
|
||||
}
|
||||
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
LabeledContent("HuggingFace Token") {
|
||||
SecureField("optional", text: $pendingHFToken)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 200)
|
||||
}
|
||||
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
Toggle("Offline Mode", isOn: $pendingOfflineMode)
|
||||
Text("Skip internet checks and use only locally available models.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyGeneralSettings()
|
||||
}
|
||||
.disabled(!hasGeneralChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Model Tab
|
||||
|
||||
private var modelTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels)
|
||||
Text("Allow text-to-image and image-to-image models in the model picker.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyModelSettings()
|
||||
}
|
||||
.disabled(!hasModelChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Advanced Tab
|
||||
|
||||
private var advancedTab: some View {
|
||||
Form {
|
||||
Section("Onboarding") {
|
||||
HStack {
|
||||
VStack(alignment: .leading) {
|
||||
Text("Reset Onboarding")
|
||||
Text("Opens the dashboard and resets the onboarding wizard.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
Spacer()
|
||||
Button("Reset") {
|
||||
guard let url = URL(string: "http://localhost:52415/?reset-onboarding")
|
||||
else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Section("Debug Info") {
|
||||
LabeledContent("Thunderbolt Bridge") {
|
||||
Text(thunderboltStatusText)
|
||||
.foregroundColor(thunderboltStatusColor)
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
clusterThunderboltBridgeView
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
interfaceIpList
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
rdmaStatusView
|
||||
}
|
||||
|
||||
sendBugReportButton
|
||||
}
|
||||
|
||||
Section("Danger Zone") {
|
||||
Button(role: .destructive) {
|
||||
showUninstallConfirmationAlert()
|
||||
} label: {
|
||||
HStack {
|
||||
Text("Uninstall EXO")
|
||||
Spacer()
|
||||
Image(systemName: "trash")
|
||||
.imageScale(.small)
|
||||
}
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - About Tab
|
||||
|
||||
private var aboutTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Version") {
|
||||
Text(buildTag)
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
LabeledContent("Commit") {
|
||||
Text(buildCommit)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
}
|
||||
|
||||
Section {
|
||||
Button("Check for Updates") {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Debug Info Views (moved from ContentView)
|
||||
|
||||
private var thunderboltStatusText: String {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled):
|
||||
return "Disabled"
|
||||
case .some(.deleted):
|
||||
return "Deleted"
|
||||
case .some(.enabled):
|
||||
return "Enabled"
|
||||
case nil:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
private var thunderboltStatusColor: Color {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled), .some(.deleted):
|
||||
return .green
|
||||
case .some(.enabled):
|
||||
return .red
|
||||
case nil:
|
||||
return .secondary
|
||||
}
|
||||
}
|
||||
|
||||
private var clusterThunderboltBridgeView: some View {
|
||||
let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if bridgeStatuses.isEmpty {
|
||||
Text("Cluster TB Bridge: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster TB Bridge Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = bridgeStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText =
|
||||
!status.exists
|
||||
? "N/A"
|
||||
: (status.enabled ? "Enabled" : "Disabled")
|
||||
let color: Color =
|
||||
!status.exists
|
||||
? .secondary
|
||||
: (status.enabled ? .red : .green)
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var interfaceIpList: some View {
|
||||
let statuses = networkStatusService.status.interfaceStatuses
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
Text("Interfaces (en0–en7):")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
if statuses.isEmpty {
|
||||
Text(" Unknown")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
ForEach(statuses, id: \.interfaceName) { status in
|
||||
let ipText = status.ipAddress ?? "No IP"
|
||||
Text(" \(status.interfaceName): \(ipText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(status.ipAddress == nil ? .red : .green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var rdmaStatusView: some View {
|
||||
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
let localDevices = networkStatusService.status.localRdmaDevices
|
||||
let localPorts = networkStatusService.status.localRdmaActivePorts
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if rdmaStatuses.isEmpty {
|
||||
Text("Cluster RDMA: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster RDMA Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = rdmaStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText = status.enabled ? "Enabled" : "Disabled"
|
||||
let color: Color = status.enabled ? .green : .orange
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !localDevices.isEmpty {
|
||||
Text(" Local Devices: \(localDevices.joined(separator: ", "))")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
if !localPorts.isEmpty {
|
||||
Text(" Local Active Ports:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(localPorts, id: \.device) { port in
|
||||
Text(" \(port.device) port \(port.port): \(port.state)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Button {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
if bugReportInFlight {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
}
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
}
|
||||
}
|
||||
.disabled(bugReportInFlight)
|
||||
|
||||
if let message = bugReportMessage {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Actions
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportInFlight = true
|
||||
bugReportMessage = "Collecting logs..."
|
||||
let service = BugReportService()
|
||||
do {
|
||||
let outcome = try await service.sendReport(isManual: true)
|
||||
bugReportMessage = outcome.message
|
||||
} catch {
|
||||
bugReportMessage = error.localizedDescription
|
||||
}
|
||||
bugReportInFlight = false
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
let alert = NSAlert()
|
||||
alert.messageText = "Uninstall EXO"
|
||||
alert.informativeText = """
|
||||
This will remove EXO and all its system components:
|
||||
|
||||
• Network configuration daemon
|
||||
• Launch at login registration
|
||||
• EXO network location
|
||||
|
||||
The app will be moved to Trash.
|
||||
"""
|
||||
alert.alertStyle = .warning
|
||||
alert.addButton(withTitle: "Uninstall")
|
||||
alert.addButton(withTitle: "Cancel")
|
||||
|
||||
if let uninstallButton = alert.buttons.first {
|
||||
uninstallButton.hasDestructiveAction = true
|
||||
}
|
||||
|
||||
let response = alert.runModal()
|
||||
if response == .alertFirstButtonReturn {
|
||||
performUninstall()
|
||||
}
|
||||
}
|
||||
|
||||
private func performUninstall() {
|
||||
uninstallInProgress = true
|
||||
|
||||
controller.cancelPendingLaunch()
|
||||
controller.stop()
|
||||
stateService.stopPolling()
|
||||
|
||||
DispatchQueue.global(qos: .utility).async {
|
||||
do {
|
||||
try NetworkSetupHelper.uninstall()
|
||||
|
||||
DispatchQueue.main.async {
|
||||
LaunchAtLoginHelper.disable()
|
||||
self.moveAppToTrash()
|
||||
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
DispatchQueue.main.async {
|
||||
let errorAlert = NSAlert()
|
||||
errorAlert.messageText = "Uninstall Failed"
|
||||
errorAlert.informativeText = error.localizedDescription
|
||||
errorAlert.alertStyle = .critical
|
||||
errorAlert.addButton(withTitle: "OK")
|
||||
errorAlert.runModal()
|
||||
self.uninstallInProgress = false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func moveAppToTrash() {
|
||||
guard let appURL = Bundle.main.bundleURL as URL? else { return }
|
||||
do {
|
||||
try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
|
||||
} catch {
|
||||
// If we can't trash the app, that's OK - user can do it manually
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private var hasGeneralChanges: Bool {
|
||||
pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
|
||||
|| pendingOfflineMode != controller.offlineMode
|
||||
}
|
||||
|
||||
private var hasModelChanges: Bool {
|
||||
pendingEnableImageModels != controller.enableImageModels
|
||||
}
|
||||
|
||||
private func applyGeneralSettings() {
|
||||
controller.customNamespace = pendingNamespace
|
||||
controller.hfToken = pendingHFToken
|
||||
controller.offlineMode = pendingOfflineMode
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func applyModelSettings() {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func restartIfRunning() {
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Manages a standalone native macOS Settings window.
|
||||
/// Ensures only one instance exists and brings it to front on repeated opens.
|
||||
@MainActor
|
||||
final class SettingsWindowController: ObservableObject {
|
||||
private var window: NSWindow?
|
||||
|
||||
func open(
|
||||
controller: ExoProcessController,
|
||||
updater: SparkleUpdater,
|
||||
networkStatusService: NetworkStatusService,
|
||||
thunderboltBridgeService: ThunderboltBridgeService,
|
||||
stateService: ClusterStateService
|
||||
) {
|
||||
if let existing = window, existing.isVisible {
|
||||
existing.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
return
|
||||
}
|
||||
|
||||
let settingsView = SettingsView()
|
||||
.environmentObject(controller)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(networkStatusService)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(stateService)
|
||||
|
||||
let hostingView = NSHostingView(rootView: settingsView)
|
||||
|
||||
let newWindow = NSWindow(
|
||||
contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
|
||||
styleMask: [.titled, .closable],
|
||||
backing: .buffered,
|
||||
defer: false
|
||||
)
|
||||
newWindow.title = "EXO Settings"
|
||||
newWindow.contentView = hostingView
|
||||
newWindow.center()
|
||||
newWindow.isReleasedWhenClosed = false
|
||||
newWindow.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
|
||||
window = newWindow
|
||||
}
|
||||
}
|
||||
+1
-3
@@ -2,6 +2,4 @@
|
||||
#
|
||||
# Lists the suite files to include. Each file defines benchmarks
|
||||
# with shared constraints, topology, and default args.
|
||||
include = [
|
||||
"single-m3-ultra.toml",
|
||||
]
|
||||
include = ["single-m3-ultra.toml"]
|
||||
|
||||
+11
-2
@@ -75,7 +75,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
|
||||
model_path = Path(
|
||||
snapshot_download(
|
||||
model_id,
|
||||
allow_patterns=["*.json", "*.py", "*.tiktoken"],
|
||||
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
|
||||
)
|
||||
)
|
||||
|
||||
@@ -338,7 +338,7 @@ def main() -> int:
|
||||
)
|
||||
|
||||
logger.info("Planning phase: checking downloads...")
|
||||
run_planning_phase(
|
||||
download_duration_s = run_planning_phase(
|
||||
client,
|
||||
full_model_id,
|
||||
selected[0],
|
||||
@@ -346,6 +346,10 @@ def main() -> int:
|
||||
args.timeout,
|
||||
settle_deadline,
|
||||
)
|
||||
if download_duration_s is not None:
|
||||
logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
|
||||
else:
|
||||
logger.info("Download: model already cached")
|
||||
|
||||
all_rows: list[dict[str, Any]] = []
|
||||
|
||||
@@ -409,6 +413,11 @@ def main() -> int:
|
||||
"pp_tokens": actual_pp_tokens,
|
||||
"tg": tg,
|
||||
"repeat_index": r,
|
||||
**(
|
||||
{"download_duration_s": download_duration_s}
|
||||
if download_duration_s is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
runs.append(row)
|
||||
|
||||
+18
-6
@@ -289,8 +289,12 @@ def run_planning_phase(
|
||||
danger_delete: bool,
|
||||
timeout: float,
|
||||
settle_deadline: float | None,
|
||||
) -> None:
|
||||
"""Check disk space and ensure model is downloaded before benchmarking."""
|
||||
) -> float | None:
|
||||
"""Check disk space and ensure model is downloaded before benchmarking.
|
||||
|
||||
Returns the wall-clock download duration in seconds if a fresh download
|
||||
was needed, or None if the model was already cached on all nodes.
|
||||
"""
|
||||
# Get model size from /models
|
||||
models = client.request_json("GET", "/models") or {}
|
||||
model_bytes = 0
|
||||
@@ -303,7 +307,7 @@ def run_planning_phase(
|
||||
logger.warning(
|
||||
f"Could not determine size for {full_model_id}, skipping disk check"
|
||||
)
|
||||
return
|
||||
return None
|
||||
|
||||
# Get nodes from preview
|
||||
inner = unwrap_instance(preview["instance"])
|
||||
@@ -314,6 +318,8 @@ def run_planning_phase(
|
||||
downloads = state.get("downloads", {})
|
||||
node_disk = state.get("nodeDisk", {})
|
||||
|
||||
needs_download = False
|
||||
|
||||
for node_id in node_ids:
|
||||
node_downloads = downloads.get(node_id, [])
|
||||
|
||||
@@ -329,6 +335,8 @@ def run_planning_phase(
|
||||
if already_downloaded:
|
||||
continue
|
||||
|
||||
needs_download = True
|
||||
|
||||
# Wait for disk info if settle_deadline is set
|
||||
disk_info = node_disk.get(node_id, {})
|
||||
backoff = _SETTLE_INITIAL_BACKOFF_S
|
||||
@@ -357,16 +365,17 @@ def run_planning_phase(
|
||||
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
|
||||
)
|
||||
|
||||
# Delete from smallest to largest
|
||||
# Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
|
||||
completed = [
|
||||
(
|
||||
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
|
||||
"modelId"
|
||||
],
|
||||
p["DownloadCompleted"]["totalBytes"]["inBytes"],
|
||||
p["DownloadCompleted"]["total"]["inBytes"],
|
||||
)
|
||||
for p in node_downloads
|
||||
if "DownloadCompleted" in p
|
||||
and not p["DownloadCompleted"].get("readOnly", False)
|
||||
]
|
||||
for del_model, size in sorted(completed, key=lambda x: x[1]):
|
||||
logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
|
||||
@@ -379,6 +388,7 @@ def run_planning_phase(
|
||||
raise RuntimeError(f"Could not free enough space on {node_id}")
|
||||
|
||||
# Start downloads (idempotent)
|
||||
download_t0 = time.perf_counter() if needs_download else None
|
||||
for node_id in node_ids:
|
||||
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
|
||||
shard = runner_to_shard[runner_id]
|
||||
@@ -421,7 +431,9 @@ def run_planning_phase(
|
||||
if not done:
|
||||
all_done = False
|
||||
if all_done:
|
||||
return
|
||||
if download_t0 is not None:
|
||||
return time.perf_counter() - download_t0
|
||||
return None
|
||||
time.sleep(1)
|
||||
|
||||
raise TimeoutError("Downloads did not complete in time")
|
||||
|
||||
@@ -4,12 +4,13 @@ version = "0.1.0"
|
||||
description = "Benchmarking tool for exo distributed inference"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"httpx>=0.27.0",
|
||||
"loguru>=0.7.3",
|
||||
"transformers>=5.0.0",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"tiktoken>=0.12.0",
|
||||
"jinja2>=3.1.0",
|
||||
"httpx>=0.27.0",
|
||||
"loguru>=0.7.3",
|
||||
"transformers>=5.0.0",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"tiktoken>=0.12.0",
|
||||
"jinja2>=3.1.0",
|
||||
"protobuf>=5.29.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -2,10 +2,10 @@
|
||||
#
|
||||
# Shared constraints applied to ALL benchmarks in this file.
|
||||
constraints = [
|
||||
"All(MacOsBuild(=25D125))",
|
||||
"Hosts(=1)",
|
||||
"All(Chip(m3_ultra))",
|
||||
"All(GpuCores(=80))",
|
||||
"All(MacOsBuild(=25D125))",
|
||||
"Hosts(=1)",
|
||||
"All(Chip(m3_ultra))",
|
||||
"All(GpuCores(=80))",
|
||||
]
|
||||
|
||||
[topology]
|
||||
|
||||
@@ -202,6 +202,23 @@
|
||||
filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5));
|
||||
}
|
||||
|
||||
/* Onboarding step 2: connection line between devices */
|
||||
.onboarding-connection-line {
|
||||
stroke: oklch(0.85 0.18 85 / 0.5);
|
||||
stroke-width: 1.5px;
|
||||
stroke-dasharray: 6, 6;
|
||||
animation: flowAnimation 1s linear infinite;
|
||||
filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4));
|
||||
}
|
||||
|
||||
/* Onboarding step 4: red connection line for disconnect */
|
||||
.onboarding-connection-line-red {
|
||||
stroke: rgba(220, 38, 38, 0.7);
|
||||
stroke-width: 1.5px;
|
||||
stroke-dasharray: 6, 6;
|
||||
filter: drop-shadow(0 0 2px rgba(220, 38, 38, 0.3));
|
||||
}
|
||||
|
||||
.graph-link-active {
|
||||
stroke: oklch(0.85 0.18 85 / 0.8);
|
||||
stroke-width: 2px;
|
||||
@@ -320,3 +337,51 @@ input:focus, textarea:focus {
|
||||
transform: translate(400px, 400px);
|
||||
}
|
||||
}
|
||||
|
||||
/* Onboarding smooth fade-in */
|
||||
@keyframes onb-fade-in {
|
||||
from { opacity: 0; transform: translateY(8px); }
|
||||
to { opacity: 1; transform: none; }
|
||||
}
|
||||
|
||||
/* Opacity-only fade-in — no transform, safe for containers with fixed-position children */
|
||||
@keyframes onb-fade-opacity {
|
||||
from { opacity: 0; }
|
||||
to { opacity: 1; }
|
||||
}
|
||||
|
||||
/* Respect reduced motion preference */
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
.shooting-star,
|
||||
.shooting-star::before {
|
||||
animation: none !important;
|
||||
opacity: 0 !important;
|
||||
}
|
||||
.graph-link {
|
||||
animation: none;
|
||||
}
|
||||
.status-pulse {
|
||||
animation: none;
|
||||
}
|
||||
.cursor-blink {
|
||||
animation: none;
|
||||
}
|
||||
.onboarding-connection-line {
|
||||
animation: none;
|
||||
}
|
||||
.onboarding-connection-line-red {
|
||||
animation: none;
|
||||
}
|
||||
[style*="onb-fade-in"],
|
||||
[style*="onb-fade-opacity"] {
|
||||
animation: none !important;
|
||||
opacity: 1 !important;
|
||||
}
|
||||
*,
|
||||
*::before,
|
||||
*::after {
|
||||
transition-duration: 0.01ms !important;
|
||||
animation-duration: 0.01ms !important;
|
||||
animation-iteration-count: 1 !important;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,8 +7,6 @@
|
||||
editingImage,
|
||||
clearEditingImage,
|
||||
selectedChatModel,
|
||||
setSelectedChatModel,
|
||||
instances,
|
||||
ttftMs,
|
||||
tps,
|
||||
totalTokens,
|
||||
@@ -29,6 +27,19 @@
|
||||
showModelSelector?: boolean;
|
||||
modelTasks?: Record<string, string[]>;
|
||||
modelCapabilities?: Record<string, string[]>;
|
||||
onSend?: () => void;
|
||||
onAutoSend?: (
|
||||
content: string,
|
||||
files?: {
|
||||
id: string;
|
||||
name: string;
|
||||
type: string;
|
||||
textContent?: string;
|
||||
preview?: string;
|
||||
}[],
|
||||
) => void;
|
||||
onOpenModelPicker?: () => void;
|
||||
modelDisplayOverride?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
@@ -39,6 +50,10 @@
|
||||
showModelSelector = false,
|
||||
modelTasks = {},
|
||||
modelCapabilities = {},
|
||||
onSend,
|
||||
onAutoSend,
|
||||
onOpenModelPicker,
|
||||
modelDisplayOverride,
|
||||
}: Props = $props();
|
||||
|
||||
let message = $state("");
|
||||
@@ -49,27 +64,12 @@
|
||||
const thinkingEnabled = $derived(thinkingEnabledStore());
|
||||
let loading = $derived(isLoading());
|
||||
const currentModel = $derived(selectedChatModel());
|
||||
const instanceData = $derived(instances());
|
||||
const currentTtft = $derived(ttftMs());
|
||||
const currentTps = $derived(tps());
|
||||
const currentTokens = $derived(totalTokens());
|
||||
const currentEditingImage = $derived(editingImage());
|
||||
const isEditMode = $derived(currentEditingImage !== null);
|
||||
|
||||
// Custom dropdown state
|
||||
let isModelDropdownOpen = $state(false);
|
||||
let dropdownButtonRef: HTMLButtonElement | undefined = $state();
|
||||
let dropdownPosition = $derived(() => {
|
||||
if (!dropdownButtonRef || !isModelDropdownOpen)
|
||||
return { top: 0, left: 0, width: 0 };
|
||||
const rect = dropdownButtonRef.getBoundingClientRect();
|
||||
return {
|
||||
top: rect.top,
|
||||
left: rect.left,
|
||||
width: rect.width,
|
||||
};
|
||||
});
|
||||
|
||||
// Accept all supported file types
|
||||
const acceptString = getAcceptString(["image", "text", "pdf"]);
|
||||
|
||||
@@ -122,73 +122,14 @@
|
||||
uploadedFiles.length > 0),
|
||||
);
|
||||
|
||||
// Extract available models from running instances
|
||||
const availableModels = $derived(() => {
|
||||
const models: Array<{ id: string; label: string; isImageModel: boolean }> =
|
||||
[];
|
||||
for (const [, instance] of Object.entries(instanceData)) {
|
||||
const modelId = getInstanceModelId(instance);
|
||||
if (
|
||||
modelId &&
|
||||
modelId !== "Unknown" &&
|
||||
!models.some((m) => m.id === modelId)
|
||||
) {
|
||||
models.push({
|
||||
id: modelId,
|
||||
label: modelId.split("/").pop() || modelId,
|
||||
isImageModel: modelSupportsImageGeneration(modelId),
|
||||
});
|
||||
}
|
||||
}
|
||||
return models;
|
||||
});
|
||||
|
||||
// Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
|
||||
let previousModelIds: Set<string> = new Set();
|
||||
|
||||
// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
|
||||
$effect(() => {
|
||||
const models = availableModels();
|
||||
const currentModelIds = new Set(models.map((m) => m.id));
|
||||
|
||||
if (models.length > 0) {
|
||||
// Find newly added models (in current but not in previous)
|
||||
const newModels = models.filter((m) => !previousModelIds.has(m.id));
|
||||
|
||||
// If no model selected, select the first available
|
||||
if (!currentModel) {
|
||||
setSelectedChatModel(models[0].id);
|
||||
}
|
||||
// If current model is stale (no longer has a running instance), reset to first available
|
||||
else if (!models.some((m) => m.id === currentModel)) {
|
||||
setSelectedChatModel(models[0].id);
|
||||
}
|
||||
// If a new model was just added, select it
|
||||
else if (newModels.length > 0 && previousModelIds.size > 0) {
|
||||
setSelectedChatModel(newModels[0].id);
|
||||
}
|
||||
} else {
|
||||
// No instances running - clear the selected model
|
||||
if (currentModel) {
|
||||
setSelectedChatModel("");
|
||||
}
|
||||
}
|
||||
|
||||
// Update previous model IDs for next comparison
|
||||
previousModelIds = currentModelIds;
|
||||
});
|
||||
|
||||
function getInstanceModelId(instanceWrapped: unknown): string {
|
||||
if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
|
||||
const keys = Object.keys(instanceWrapped as Record<string, unknown>);
|
||||
if (keys.length === 1) {
|
||||
const instance = (instanceWrapped as Record<string, unknown>)[
|
||||
keys[0]
|
||||
] as { shardAssignments?: { modelId?: string } };
|
||||
return instance?.shardAssignments?.modelId || "";
|
||||
}
|
||||
return "";
|
||||
}
|
||||
// Short label for the currently selected model
|
||||
const currentModelLabel = $derived(
|
||||
currentModel
|
||||
? currentModel.split("/").pop() || currentModel
|
||||
: modelDisplayOverride
|
||||
? modelDisplayOverride.split("/").pop() || modelDisplayOverride
|
||||
: "",
|
||||
);
|
||||
|
||||
async function handleFiles(files: File[]) {
|
||||
if (files.length === 0) return;
|
||||
@@ -275,6 +216,15 @@
|
||||
uploadedFiles = [];
|
||||
resetTextareaHeight();
|
||||
|
||||
// When onAutoSend is provided, the parent controls all send logic
|
||||
// (including launching non-running models before sending)
|
||||
if (onAutoSend) {
|
||||
onAutoSend(content, files);
|
||||
onSend?.();
|
||||
setTimeout(() => textareaRef?.focus(), 10);
|
||||
return;
|
||||
}
|
||||
|
||||
// Use image editing if in edit mode
|
||||
if (isEditMode && currentEditingImage && content) {
|
||||
editImage(content, currentEditingImage.imageDataUrl);
|
||||
@@ -306,6 +256,8 @@
|
||||
);
|
||||
}
|
||||
|
||||
onSend?.();
|
||||
|
||||
// Refocus the textarea after sending
|
||||
setTimeout(() => textareaRef?.focus(), 10);
|
||||
}
|
||||
@@ -416,7 +368,7 @@
|
||||
{/if}
|
||||
|
||||
<!-- Model selector (when enabled) -->
|
||||
{#if showModelSelector && availableModels().length > 0}
|
||||
{#if showModelSelector}
|
||||
<div
|
||||
class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
|
||||
>
|
||||
@@ -425,33 +377,22 @@
|
||||
class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
|
||||
>MODEL:</span
|
||||
>
|
||||
<!-- Custom dropdown -->
|
||||
<!-- Model button — opens the full model picker -->
|
||||
<div class="relative flex-1 max-w-xs">
|
||||
<button
|
||||
bind:this={dropdownButtonRef}
|
||||
type="button"
|
||||
onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
|
||||
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
|
||||
? 'border-exo-yellow/70'
|
||||
: ''}"
|
||||
onclick={() => onOpenModelPicker?.()}
|
||||
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
|
||||
>
|
||||
{#if availableModels().find((m) => m.id === currentModel)}
|
||||
<span class="text-exo-yellow truncate"
|
||||
>{availableModels().find((m) => m.id === currentModel)
|
||||
?.label}</span
|
||||
>
|
||||
{:else if availableModels().length > 0}
|
||||
<span class="text-exo-yellow truncate"
|
||||
>{availableModels()[0].label}</span
|
||||
{#if currentModelLabel}
|
||||
<span class="text-exo-yellow truncate">{currentModelLabel}</span
|
||||
>
|
||||
{:else}
|
||||
<span class="text-exo-light-gray/50">— SELECT MODEL —</span>
|
||||
{/if}
|
||||
</button>
|
||||
<div
|
||||
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
|
||||
? 'rotate-180'
|
||||
: ''}"
|
||||
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
|
||||
>
|
||||
<svg
|
||||
class="w-3 h-3 text-exo-yellow/60"
|
||||
@@ -468,78 +409,6 @@
|
||||
</svg>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{#if isModelDropdownOpen}
|
||||
<!-- Backdrop to close dropdown -->
|
||||
<button
|
||||
type="button"
|
||||
class="fixed inset-0 z-[9998] cursor-default"
|
||||
onclick={() => (isModelDropdownOpen = false)}
|
||||
aria-label="Close dropdown"
|
||||
></button>
|
||||
|
||||
<!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
|
||||
<div
|
||||
class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
|
||||
style="bottom: calc(100vh - {dropdownPosition()
|
||||
.top}px + 4px); left: {dropdownPosition()
|
||||
.left}px; width: {dropdownPosition().width}px;"
|
||||
>
|
||||
<div class="py-1">
|
||||
{#each availableModels() as model}
|
||||
<button
|
||||
type="button"
|
||||
onclick={() => {
|
||||
setSelectedChatModel(model.id);
|
||||
isModelDropdownOpen = false;
|
||||
}}
|
||||
class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
|
||||
model.id
|
||||
? 'bg-transparent text-exo-yellow'
|
||||
: 'text-exo-light-gray hover:text-exo-yellow'}"
|
||||
>
|
||||
{#if currentModel === model.id}
|
||||
<svg
|
||||
class="w-3 h-3 flex-shrink-0"
|
||||
fill="currentColor"
|
||||
viewBox="0 0 20 20"
|
||||
>
|
||||
<path
|
||||
fill-rule="evenodd"
|
||||
d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
|
||||
clip-rule="evenodd"
|
||||
/>
|
||||
</svg>
|
||||
{:else}
|
||||
<span class="w-3"></span>
|
||||
{/if}
|
||||
{#if model.isImageModel}
|
||||
<svg
|
||||
class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
aria-label="Image generation model"
|
||||
>
|
||||
<rect
|
||||
x="3"
|
||||
y="3"
|
||||
width="18"
|
||||
height="18"
|
||||
rx="2"
|
||||
ry="2"
|
||||
/>
|
||||
<circle cx="8.5" cy="8.5" r="1.5" />
|
||||
<polyline points="21 15 16 10 5 21" />
|
||||
</svg>
|
||||
{/if}
|
||||
<span class="truncate flex-1">{model.label}</span>
|
||||
</button>
|
||||
{/each}
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
<!-- Thinking toggle -->
|
||||
{#if modelSupportsThinking()}
|
||||
|
||||
@@ -807,8 +807,8 @@
|
||||
>
|
||||
AWAITING INPUT
|
||||
</p>
|
||||
<p class="text-sm sm:text-xs text-exo-light-gray tracking-wider mt-1">
|
||||
ENTER A QUERY TO BEGIN
|
||||
<p class="text-xs text-white/30 tracking-wider mt-1.5 font-mono">
|
||||
Type a message below · Shift+Enter for newline
|
||||
</p>
|
||||
</div>
|
||||
{/if}
|
||||
@@ -823,6 +823,7 @@
|
||||
onclick={scrollToBottom}
|
||||
class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10"
|
||||
title="Scroll to bottom"
|
||||
aria-label="Scroll to bottom of messages"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5"
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
<script lang="ts" module>
|
||||
export interface ChatModelInfo {
|
||||
id: string;
|
||||
name: string;
|
||||
base_model: string;
|
||||
storage_size_megabytes: number;
|
||||
capabilities: string[];
|
||||
family: string;
|
||||
quantization: string;
|
||||
}
|
||||
|
||||
// Auto mode tier list (for when user just starts typing)
|
||||
export const AUTO_TIERS: string[][] = [
|
||||
// Tier 1 (frontier)
|
||||
["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
|
||||
// Tier 2 (excellent)
|
||||
[
|
||||
"Kimi K2",
|
||||
"Qwen3 235B",
|
||||
"MiniMax M2.5",
|
||||
"Step 3.5 Flash",
|
||||
"Qwen3 Next 80B",
|
||||
],
|
||||
// Tier 3 (great)
|
||||
[
|
||||
"GLM 4.7",
|
||||
"MiniMax M2.1",
|
||||
"Qwen3 Coder 480B",
|
||||
"GLM 4.5 Air",
|
||||
"Llama 3.3 70B",
|
||||
],
|
||||
// Tier 4 (good)
|
||||
["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
|
||||
// Tier 5 (small/fast)
|
||||
[
|
||||
"Llama 3.1 8B",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
"Llama 3.2 1B",
|
||||
],
|
||||
];
|
||||
|
||||
/** Return the tier index (0 = best) for a base_model name. */
|
||||
export function getAutoTierIndex(baseModel: string): number {
|
||||
for (let i = 0; i < AUTO_TIERS.length; i++) {
|
||||
if (AUTO_TIERS[i].includes(baseModel)) return i;
|
||||
}
|
||||
return AUTO_TIERS.length; // not in any tier → lowest priority
|
||||
}
|
||||
|
||||
/** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
|
||||
export function pickAutoModel(
|
||||
modelList: ChatModelInfo[],
|
||||
memoryGB: number,
|
||||
): ChatModelInfo | null {
|
||||
for (const tier of AUTO_TIERS) {
|
||||
const candidates: ChatModelInfo[] = [];
|
||||
for (const baseModel of tier) {
|
||||
const variants = modelList
|
||||
.filter(
|
||||
(m) =>
|
||||
m.base_model === baseModel &&
|
||||
(m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
|
||||
(m.storage_size_megabytes || 0) > 0,
|
||||
)
|
||||
.sort(
|
||||
(a, b) =>
|
||||
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
|
||||
);
|
||||
if (variants[0]) candidates.push(variants[0]);
|
||||
}
|
||||
if (candidates.length > 0) {
|
||||
candidates.sort(
|
||||
(a, b) =>
|
||||
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
|
||||
);
|
||||
return candidates[0];
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
</script>
|
||||
|
||||
<script lang="ts">
|
||||
interface CategoryRecommendation {
|
||||
category: "coding" | "writing" | "agentic" | "biggest";
|
||||
label: string;
|
||||
model: ChatModelInfo | null;
|
||||
tooltip: string;
|
||||
}
|
||||
|
||||
interface Props {
|
||||
models: ChatModelInfo[];
|
||||
clusterLabel: string;
|
||||
totalMemoryGB: number;
|
||||
onSelect: (modelId: string, category: string) => void;
|
||||
onAddModel: () => void;
|
||||
class?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
models,
|
||||
clusterLabel,
|
||||
totalMemoryGB,
|
||||
onSelect,
|
||||
onAddModel,
|
||||
class: className = "",
|
||||
}: Props = $props();
|
||||
|
||||
// --- Hardcoded Rankings ---
|
||||
const CODING_RANKING = [
|
||||
"Qwen3 Coder Next",
|
||||
"Qwen3 Coder 480B",
|
||||
"Qwen3 30B",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
const WRITING_RANKING = [
|
||||
"Kimi K2.5",
|
||||
"Kimi K2",
|
||||
"Qwen3 Next 80B",
|
||||
"Llama 3.3 70B",
|
||||
"MiniMax M2.5",
|
||||
"GLM 4.5 Air",
|
||||
"GLM 4.7 Flash",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
const AGENTIC_RANKING = [
|
||||
"DeepSeek V3.1",
|
||||
"GLM-5",
|
||||
"Qwen3 235B",
|
||||
"Step 3.5 Flash",
|
||||
"GLM 4.7",
|
||||
"MiniMax M2.1",
|
||||
"GPT-OSS 120B",
|
||||
"Llama 3.3 70B",
|
||||
"Llama 3.1 70B",
|
||||
"GLM 4.7 Flash",
|
||||
"GPT-OSS 20B",
|
||||
"Qwen3 30B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
function getModelSizeGB(m: ChatModelInfo): number {
|
||||
return (m.storage_size_megabytes || 0) / 1024;
|
||||
}
|
||||
|
||||
function fitsInMemory(m: ChatModelInfo): boolean {
|
||||
return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
|
||||
}
|
||||
|
||||
/** For a given base_model name, find the biggest quant variant that fits in memory. */
|
||||
function pickBestVariant(baseModel: string): ChatModelInfo | null {
|
||||
const variants = models
|
||||
.filter((m) => m.base_model === baseModel && fitsInMemory(m))
|
||||
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
|
||||
return variants[0] ?? null;
|
||||
}
|
||||
|
||||
/** Walk a ranked list of base_model names, return the first that has a fitting variant. */
|
||||
function pickFromRanking(ranking: string[]): ChatModelInfo | null {
|
||||
for (const baseModel of ranking) {
|
||||
const pick = pickBestVariant(baseModel);
|
||||
if (pick) return pick;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Pick the single biggest model that fits. */
|
||||
function pickBiggest(): ChatModelInfo | null {
|
||||
const fitting = models
|
||||
.filter((m) => fitsInMemory(m))
|
||||
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
|
||||
return fitting[0] ?? null;
|
||||
}
|
||||
|
||||
const recommendations = $derived.by((): CategoryRecommendation[] => {
|
||||
return [
|
||||
{
|
||||
category: "coding",
|
||||
label: "Best for Coding",
|
||||
model: pickFromRanking(CODING_RANKING),
|
||||
tooltip:
|
||||
"Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
|
||||
},
|
||||
{
|
||||
category: "writing",
|
||||
label: "Best for Writing",
|
||||
model: pickFromRanking(WRITING_RANKING),
|
||||
tooltip: "Ranked by creative writing quality and instruction following",
|
||||
},
|
||||
{
|
||||
category: "agentic",
|
||||
label: "Best Agentic",
|
||||
model: pickFromRanking(AGENTIC_RANKING),
|
||||
tooltip: "Ranked by reasoning, planning, and tool-use capability",
|
||||
},
|
||||
{
|
||||
category: "biggest",
|
||||
label: "Biggest",
|
||||
model: pickBiggest(),
|
||||
tooltip: "Largest model that fits in your available memory",
|
||||
},
|
||||
];
|
||||
});
|
||||
|
||||
function formatSize(mb: number): string {
|
||||
const gb = mb / 1024;
|
||||
if (gb >= 100) return `${Math.round(gb)} GB`;
|
||||
return `${gb.toFixed(1)} GB`;
|
||||
}
|
||||
|
||||
// Category icons (SVG paths)
|
||||
const categoryIcons: Record<string, string> = {
|
||||
coding:
|
||||
"M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
|
||||
writing:
|
||||
"M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
|
||||
agentic:
|
||||
"M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
|
||||
biggest:
|
||||
"M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
|
||||
};
|
||||
|
||||
let hoveredTooltip = $state<string | null>(null);
|
||||
let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
|
||||
|
||||
function showTooltip(category: string, e: MouseEvent | FocusEvent) {
|
||||
hoveredTooltip = category;
|
||||
const target = e.currentTarget as HTMLElement;
|
||||
const rect = target.getBoundingClientRect();
|
||||
tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
|
||||
}
|
||||
|
||||
function hideTooltip() {
|
||||
hoveredTooltip = null;
|
||||
tooltipAnchor = null;
|
||||
}
|
||||
</script>
|
||||
|
||||
<div class="flex flex-col items-center justify-center gap-6 {className}">
|
||||
<!-- Header -->
|
||||
<div class="text-center">
|
||||
<p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
|
||||
Recommended for your
|
||||
</p>
|
||||
<p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
|
||||
</div>
|
||||
|
||||
<!-- Category Cards Grid -->
|
||||
<div class="grid grid-cols-2 gap-3 w-full max-w-md">
|
||||
{#each recommendations as rec}
|
||||
{#if rec.model}
|
||||
<button
|
||||
type="button"
|
||||
onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
|
||||
class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
|
||||
>
|
||||
<!-- Category icon + label -->
|
||||
<div class="flex items-center gap-2 w-full">
|
||||
<svg
|
||||
class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="1.5"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={categoryIcons[rec.category]}
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
|
||||
>
|
||||
{rec.label}
|
||||
</span>
|
||||
<!-- Info tooltip -->
|
||||
<div class="ml-auto flex-shrink-0">
|
||||
<span
|
||||
role="button"
|
||||
tabindex="-1"
|
||||
class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
|
||||
onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
|
||||
onmouseleave={() => hideTooltip()}
|
||||
onclick={(e: MouseEvent) => {
|
||||
e.stopPropagation();
|
||||
if (hoveredTooltip === rec.category) {
|
||||
hideTooltip();
|
||||
} else {
|
||||
showTooltip(rec.category, e);
|
||||
}
|
||||
}}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<circle cx="12" cy="12" r="10" />
|
||||
<path d="M12 16v-4m0-4h.01" />
|
||||
</svg>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Model name + size -->
|
||||
<div class="w-full">
|
||||
<p class="text-sm text-white font-mono truncate">
|
||||
{rec.model.base_model}
|
||||
</p>
|
||||
<p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
|
||||
{formatSize(rec.model.storage_size_megabytes)}
|
||||
{#if rec.model.quantization}
|
||||
<span class="text-exo-light-gray/40"
|
||||
>· {rec.model.quantization}</span
|
||||
>
|
||||
{/if}
|
||||
</p>
|
||||
</div>
|
||||
</button>
|
||||
{:else}
|
||||
<!-- No model fits for this category -->
|
||||
<div
|
||||
class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
|
||||
>
|
||||
<div class="flex items-center gap-2">
|
||||
<svg
|
||||
class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="1.5"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={categoryIcons[rec.category]}
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
|
||||
>{rec.label}</span
|
||||
>
|
||||
</div>
|
||||
<p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
|
||||
</div>
|
||||
{/if}
|
||||
{/each}
|
||||
</div>
|
||||
|
||||
<!-- Add Model Button -->
|
||||
<button
|
||||
type="button"
|
||||
onclick={onAddModel}
|
||||
class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
|
||||
</svg>
|
||||
Add Model
|
||||
</button>
|
||||
|
||||
<!-- Auto hint -->
|
||||
<p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
|
||||
Or just start typing — we'll pick the best model automatically
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
|
||||
{#if hoveredTooltip && tooltipAnchor}
|
||||
{@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
|
||||
{#if rec}
|
||||
<div
|
||||
class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
|
||||
style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
|
||||
8}px; transform: translate(-50%, -100%);"
|
||||
>
|
||||
{rec.tooltip}
|
||||
</div>
|
||||
{/if}
|
||||
{/if}
|
||||
@@ -2,7 +2,6 @@
|
||||
import {
|
||||
conversations,
|
||||
activeConversationId,
|
||||
createConversation,
|
||||
loadConversation,
|
||||
deleteConversation,
|
||||
deleteAllConversations,
|
||||
@@ -17,9 +16,15 @@
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
onNewChat?: () => void;
|
||||
onSelectConversation?: () => void;
|
||||
}
|
||||
|
||||
let { class: className = "" }: Props = $props();
|
||||
let {
|
||||
class: className = "",
|
||||
onNewChat,
|
||||
onSelectConversation,
|
||||
}: Props = $props();
|
||||
|
||||
const conversationList = $derived(conversations());
|
||||
const activeId = $derived(activeConversationId());
|
||||
@@ -42,10 +47,11 @@
|
||||
);
|
||||
|
||||
function handleNewChat() {
|
||||
createConversation();
|
||||
onNewChat?.();
|
||||
}
|
||||
|
||||
function handleSelectConversation(id: string) {
|
||||
onSelectConversation?.();
|
||||
loadConversation(id);
|
||||
}
|
||||
|
||||
@@ -185,11 +191,7 @@
|
||||
|
||||
let instanceType: string | null = null;
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (
|
||||
instanceTag === "MlxIbvInstance" ||
|
||||
instanceTag === "MlxJacclInstance"
|
||||
)
|
||||
instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
@@ -307,7 +309,7 @@
|
||||
<div class="py-2">
|
||||
<div class="px-4 py-2">
|
||||
<span
|
||||
class="text-sm text-white/70 font-mono tracking-wider uppercase"
|
||||
class="text-xs text-exo-light-gray font-mono tracking-wider uppercase"
|
||||
>
|
||||
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"}
|
||||
</span>
|
||||
@@ -376,39 +378,37 @@
|
||||
onkeydown={(e) =>
|
||||
e.key === "Enter" &&
|
||||
handleSelectConversation(conversation.id)}
|
||||
class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer
|
||||
class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer
|
||||
{activeId === conversation.id
|
||||
? 'bg-transparent border border-exo-yellow/30'
|
||||
: 'hover:border-exo-yellow/20 border border-transparent'}"
|
||||
? 'bg-exo-yellow/5 border border-exo-yellow/30'
|
||||
: 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}"
|
||||
>
|
||||
<div class="flex-1 min-w-0 pr-2">
|
||||
<div
|
||||
class="text-sm truncate {activeId === conversation.id
|
||||
class="text-sm font-medium truncate {activeId ===
|
||||
conversation.id
|
||||
? 'text-exo-yellow'
|
||||
: 'text-white/90'}"
|
||||
: 'text-white'}"
|
||||
>
|
||||
{conversation.name}
|
||||
</div>
|
||||
<div class="text-sm text-white/50 mt-0.5">
|
||||
<div class="text-xs text-white/60 mt-0.5">
|
||||
{formatDate(conversation.updatedAt)}
|
||||
</div>
|
||||
<div class="text-sm text-white/70 truncate">
|
||||
<div class="text-xs text-exo-light-gray truncate">
|
||||
{info.modelLabel}
|
||||
</div>
|
||||
<div class="text-xs text-white/60 font-mono">
|
||||
Strategy: <span class="text-white/80"
|
||||
>{info.strategyLabel}</span
|
||||
>
|
||||
</div>
|
||||
{#if stats}
|
||||
<div class="text-xs text-white/60 font-mono mt-1">
|
||||
{#if stats.ttftMs}<span class="text-white/40">TTFT</span>
|
||||
{stats.ttftMs.toFixed(
|
||||
0,
|
||||
)}ms{/if}{#if stats.ttftMs && stats.tps}<span
|
||||
class="text-white/30 mx-1.5">•</span
|
||||
>{/if}{#if stats.tps}{stats.tps.toFixed(1)}
|
||||
<span class="text-white/40">tok/s</span>{/if}
|
||||
<div class="text-xs text-white/70 font-mono mt-1">
|
||||
{#if stats.ttftMs}<span class="text-white/50">TTFT</span>
|
||||
<span class="text-exo-yellow/80"
|
||||
>{stats.ttftMs.toFixed(0)}ms</span
|
||||
>{/if}{#if stats.ttftMs && stats.tps}<span
|
||||
class="text-white/30 mx-1.5">·</span
|
||||
>{/if}{#if stats.tps}<span class="text-exo-yellow/80"
|
||||
>{stats.tps.toFixed(1)}</span
|
||||
>
|
||||
<span class="text-white/50">tok/s</span>{/if}
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
<script lang="ts">
|
||||
import { isConnected } from "$lib/stores/app.svelte";
|
||||
import { slide } from "svelte/transition";
|
||||
|
||||
const connected = $derived(isConnected());
|
||||
</script>
|
||||
|
||||
{#if !connected}
|
||||
<div
|
||||
transition:slide={{ duration: 200 }}
|
||||
class="relative z-50 flex items-center justify-center gap-2 px-4 py-2 bg-red-950/80 border-b border-red-500/30"
|
||||
role="alert"
|
||||
aria-live="assertive"
|
||||
>
|
||||
<div class="w-2 h-2 bg-red-500 rounded-full animate-pulse"></div>
|
||||
<span class="text-xs font-mono text-red-300 tracking-wider uppercase">
|
||||
Connection lost — Reconnecting to backend…
|
||||
</span>
|
||||
</div>
|
||||
{/if}
|
||||
@@ -0,0 +1,277 @@
|
||||
<script lang="ts">
|
||||
/**
|
||||
* DeviceIcon — renders a device icon as an SVG <g> element.
|
||||
* Uses the exact same proportional math as TopologyGraph.svelte
|
||||
* so that devices look identical in both the topology view and
|
||||
* the onboarding animation.
|
||||
*
|
||||
* Must be placed inside an <svg> element.
|
||||
*/
|
||||
|
||||
interface Props {
|
||||
/** "macbook pro" | "mac studio" | "mac mini" etc. */
|
||||
deviceType: string;
|
||||
/** Center X coordinate in SVG space */
|
||||
cx: number;
|
||||
/** Center Y coordinate in SVG space */
|
||||
cy: number;
|
||||
/** Base sizing factor (equivalent to TopologyGraph's nodeRadius) */
|
||||
size?: number;
|
||||
/** RAM usage 0–100 */
|
||||
ramPercent?: number;
|
||||
/** Unique id suffix for clip-path ids */
|
||||
uid?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
deviceType,
|
||||
cx,
|
||||
cy,
|
||||
size = 60,
|
||||
ramPercent = 60,
|
||||
uid = "dev",
|
||||
}: Props = $props();
|
||||
|
||||
// Apple logo path — same constant used by TopologyGraph
|
||||
const APPLE_LOGO_PATH =
|
||||
"M788.1 340.9c-5.8 4.5-108.2 62.2-108.2 190.5 0 148.4 130.3 200.9 134.2 202.2-.6 3.2-20.7 71.9-68.7 141.9-42.8 61.6-87.5 123.1-155.5 123.1s-85.5-39.5-164-39.5c-76.5 0-103.7 40.8-165.9 40.8s-105.6-57-155.5-127C46.7 790.7 0 663 0 541.8c0-194.4 126.4-297.5 250.8-297.5 66.1 0 121.2 43.4 162.7 43.4 39.5 0 101.1-46 176.3-46 28.5 0 130.9 2.6 198.3 99.2zm-234-181.5c31.1-36.9 53.1-88.1 53.1-139.3 0-7.1-.6-14.3-1.9-20.1-50.6 1.9-110.8 33.7-147.1 75.8-28.5 32.4-55.1 83.6-55.1 135.5 0 7.8 1.3 15.6 1.9 18.1 3.2.6 8.4 1.3 13.6 1.3 45.4 0 102.5-30.4 135.5-71.3z";
|
||||
const LOGO_NATIVE_WIDTH = 814;
|
||||
const LOGO_NATIVE_HEIGHT = 1000;
|
||||
|
||||
const wireColor = "rgba(179,179,179,0.8)";
|
||||
const strokeWidth = 1.5;
|
||||
|
||||
const modelLower = $derived(deviceType.toLowerCase());
|
||||
|
||||
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
|
||||
const studioW = $derived(size * 1.25);
|
||||
const studioH = $derived(size * 0.85);
|
||||
const studioX = $derived(cx - studioW / 2);
|
||||
const studioY = $derived(cy - studioH / 2);
|
||||
const studioCorner = 4;
|
||||
const studioTopH = $derived(studioH * 0.15);
|
||||
|
||||
// Studio front panel details
|
||||
const studioSlotH = $derived(studioH * 0.14);
|
||||
const studioVSlotW = $derived(studioW * 0.05);
|
||||
const studioVSlotY = $derived(
|
||||
studioY + studioTopH + (studioH - studioTopH) * 0.6,
|
||||
);
|
||||
const studioVSlot1X = $derived(studioX + studioW * 0.18);
|
||||
const studioVSlot2X = $derived(studioX + studioW * 0.28);
|
||||
const studioHSlotW = $derived(studioW * 0.2);
|
||||
const studioHSlotX = $derived(studioX + studioW * 0.5 - studioHSlotW / 2);
|
||||
|
||||
// Studio memory fill
|
||||
const studioMemTotalH = $derived(studioH - studioTopH);
|
||||
const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
|
||||
|
||||
// ── MacBook dimensions (same ratios as TopologyGraph) ──
|
||||
const mbW = $derived((size * 1.6 * 0.85) / 1.15);
|
||||
const mbH = $derived(size * 0.85);
|
||||
const mbX = $derived(cx - mbW / 2);
|
||||
const mbY = $derived(cy - mbH / 2);
|
||||
|
||||
const mbScreenH = $derived(mbH * 0.7);
|
||||
const mbBaseH = $derived(mbH * 0.3);
|
||||
const mbScreenW = $derived(mbW * 0.85);
|
||||
const mbScreenX = $derived(cx - mbScreenW / 2);
|
||||
const mbBezel = 3;
|
||||
|
||||
// MacBook memory fill
|
||||
const mbMemTotalH = $derived(mbScreenH - mbBezel * 2);
|
||||
const mbMemH = $derived((ramPercent / 100) * mbMemTotalH);
|
||||
|
||||
// Apple logo sizing
|
||||
const mbLogoTargetH = $derived(mbScreenH * 0.22);
|
||||
const mbLogoScale = $derived(mbLogoTargetH / LOGO_NATIVE_HEIGHT);
|
||||
const mbLogoX = $derived(cx - (LOGO_NATIVE_WIDTH * mbLogoScale) / 2);
|
||||
const mbLogoY = $derived(
|
||||
mbY + mbScreenH / 2 - (LOGO_NATIVE_HEIGHT * mbLogoScale) / 2,
|
||||
);
|
||||
|
||||
// MacBook base (trapezoidal)
|
||||
const mbBaseY = $derived(mbY + mbScreenH);
|
||||
const mbBaseTopW = $derived(mbScreenW);
|
||||
const mbBaseBottomW = $derived(mbW);
|
||||
const mbBaseTopX = $derived(cx - mbBaseTopW / 2);
|
||||
const mbBaseBottomX = $derived(cx - mbBaseBottomW / 2);
|
||||
|
||||
// Keyboard
|
||||
const mbKbX = $derived(mbBaseTopX + 6);
|
||||
const mbKbY = $derived(mbBaseY + 3);
|
||||
const mbKbW = $derived(mbBaseTopW - 12);
|
||||
const mbKbH = $derived(mbBaseH * 0.55);
|
||||
|
||||
// Trackpad
|
||||
const mbTpW = $derived(mbBaseTopW * 0.4);
|
||||
const mbTpX = $derived(cx - mbTpW / 2);
|
||||
const mbTpY = $derived(mbBaseY + mbKbH + 5);
|
||||
const mbTpH = $derived(mbBaseH * 0.3);
|
||||
|
||||
// Clip IDs
|
||||
const screenClipId = $derived(`di-screen-${uid}`);
|
||||
const studioClipId = $derived(`di-studio-${uid}`);
|
||||
</script>
|
||||
|
||||
{#if modelLower === "mac studio" || modelLower === "mac mini"}
|
||||
<!-- Mac Studio / Mac Mini -->
|
||||
<defs>
|
||||
<clipPath id={studioClipId}>
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY + studioTopH}
|
||||
width={studioW}
|
||||
height={studioH - studioTopH}
|
||||
rx={studioCorner - 1}
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<!-- Main body -->
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY}
|
||||
width={studioW}
|
||||
height={studioH}
|
||||
rx={studioCorner}
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
|
||||
<!-- Memory fill -->
|
||||
{#if ramPercent > 0}
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
|
||||
width={studioW}
|
||||
height={studioMemH}
|
||||
fill="rgba(255,215,0,0.75)"
|
||||
clip-path="url(#{studioClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Top surface divider -->
|
||||
<line
|
||||
x1={studioX}
|
||||
y1={studioY + studioTopH}
|
||||
x2={studioX + studioW}
|
||||
y2={studioY + studioTopH}
|
||||
stroke="rgba(179,179,179,0.3)"
|
||||
stroke-width="0.5"
|
||||
/>
|
||||
|
||||
<!-- Front panel: vertical slots -->
|
||||
<rect
|
||||
x={studioVSlot1X - studioVSlotW / 2}
|
||||
y={studioVSlotY}
|
||||
width={studioVSlotW}
|
||||
height={studioSlotH}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1.5"
|
||||
/>
|
||||
<rect
|
||||
x={studioVSlot2X - studioVSlotW / 2}
|
||||
y={studioVSlotY}
|
||||
width={studioVSlotW}
|
||||
height={studioSlotH}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1.5"
|
||||
/>
|
||||
|
||||
<!-- Horizontal slot (SD card) -->
|
||||
<rect
|
||||
x={studioHSlotX}
|
||||
y={studioVSlotY}
|
||||
width={studioHSlotW}
|
||||
height={studioSlotH * 0.6}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1"
|
||||
/>
|
||||
{:else}
|
||||
<!-- MacBook Pro -->
|
||||
<defs>
|
||||
<clipPath id={screenClipId}>
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<!-- Screen outer frame -->
|
||||
<rect
|
||||
x={mbScreenX}
|
||||
y={mbY}
|
||||
width={mbScreenW}
|
||||
height={mbScreenH}
|
||||
rx="3"
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
|
||||
<!-- Screen inner (dark) -->
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
fill="#0a0a12"
|
||||
/>
|
||||
|
||||
<!-- Memory fill on screen -->
|
||||
{#if ramPercent > 0}
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbMemH}
|
||||
fill="rgba(255,215,0,0.85)"
|
||||
clip-path="url(#{screenClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Apple logo -->
|
||||
<path
|
||||
d={APPLE_LOGO_PATH}
|
||||
transform="translate({mbLogoX}, {mbLogoY}) scale({mbLogoScale})"
|
||||
fill="#FFFFFF"
|
||||
opacity="0.9"
|
||||
/>
|
||||
|
||||
<!-- Keyboard base (trapezoidal) -->
|
||||
<path
|
||||
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
|
||||
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
|
||||
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
|
||||
fill="#2c2c2c"
|
||||
stroke={wireColor}
|
||||
stroke-width="1"
|
||||
/>
|
||||
|
||||
<!-- Keyboard area -->
|
||||
<rect
|
||||
x={mbKbX}
|
||||
y={mbKbY}
|
||||
width={mbKbW}
|
||||
height={mbKbH}
|
||||
fill="rgba(0,0,0,0.2)"
|
||||
rx="2"
|
||||
/>
|
||||
|
||||
<!-- Trackpad -->
|
||||
<rect
|
||||
x={mbTpX}
|
||||
y={mbTpY}
|
||||
width={mbTpW}
|
||||
height={mbTpH}
|
||||
fill="rgba(255,255,255,0.08)"
|
||||
rx="2"
|
||||
/>
|
||||
{/if}
|
||||
@@ -6,6 +6,10 @@
|
||||
export let showSidebarToggle = false;
|
||||
export let sidebarVisible = true;
|
||||
export let onToggleSidebar: (() => void) | null = null;
|
||||
export let downloadProgress: {
|
||||
count: number;
|
||||
percentage: number;
|
||||
} | null = null;
|
||||
|
||||
function handleHome(): void {
|
||||
if (onHome) {
|
||||
@@ -33,13 +37,17 @@
|
||||
<div class="absolute left-6 top-1/2 -translate-y-1/2">
|
||||
<button
|
||||
onclick={handleToggleSidebar}
|
||||
class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer"
|
||||
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer"
|
||||
title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
|
||||
aria-label={sidebarVisible
|
||||
? "Hide conversation sidebar"
|
||||
: "Show conversation sidebar"}
|
||||
aria-pressed={sidebarVisible}
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 {sidebarVisible
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-medium-gray'}"
|
||||
: 'text-exo-light-gray'}"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
@@ -75,18 +83,19 @@
|
||||
<img
|
||||
src="/exo-logo.png"
|
||||
alt="EXO"
|
||||
class="h-18 drop-shadow-[0_0_20px_rgba(255,215,0,0.5)]"
|
||||
class="h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
|
||||
/>
|
||||
</button>
|
||||
|
||||
<!-- Right: Home + Downloads -->
|
||||
<div
|
||||
<nav
|
||||
class="absolute right-6 top-1/2 -translate-y-1/2 flex items-center gap-4"
|
||||
aria-label="Main navigation"
|
||||
>
|
||||
{#if showHome}
|
||||
<button
|
||||
onclick={handleHome}
|
||||
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
title="Back to topology view"
|
||||
>
|
||||
<svg
|
||||
@@ -107,23 +116,59 @@
|
||||
{/if}
|
||||
<a
|
||||
href="/#/downloads"
|
||||
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
title="View downloads overview"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M12 3v12" />
|
||||
<path d="M7 12l5 5 5-5" />
|
||||
<path d="M5 21h14" />
|
||||
</svg>
|
||||
{#if downloadProgress}
|
||||
<!-- Compact download progress indicator -->
|
||||
<div class="relative w-4 h-4 flex-shrink-0">
|
||||
<svg class="w-4 h-4 -rotate-90" viewBox="0 0 20 20">
|
||||
<circle
|
||||
cx="10"
|
||||
cy="10"
|
||||
r="8"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
opacity="0.2"
|
||||
/>
|
||||
<circle
|
||||
cx="10"
|
||||
cy="10"
|
||||
r="8"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-dasharray={2 * Math.PI * 8}
|
||||
stroke-dashoffset={2 *
|
||||
Math.PI *
|
||||
8 *
|
||||
(1 - downloadProgress.percentage / 100)}
|
||||
class="text-blue-400 transition-all duration-300"
|
||||
/>
|
||||
</svg>
|
||||
<div
|
||||
class="absolute inset-0 flex items-center justify-center text-[6px] font-mono text-blue-400"
|
||||
>
|
||||
{downloadProgress.count}
|
||||
</div>
|
||||
</div>
|
||||
{:else}
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M12 3v12" />
|
||||
<path d="M7 12l5 5 5-5" />
|
||||
<path d="M5 21h14" />
|
||||
</svg>
|
||||
{/if}
|
||||
Downloads
|
||||
</a>
|
||||
</div>
|
||||
</nav>
|
||||
</header>
|
||||
|
||||
@@ -21,7 +21,7 @@
|
||||
} | null;
|
||||
nodes?: Record<string, NodeInfo>;
|
||||
sharding?: "Pipeline" | "Tensor";
|
||||
runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl";
|
||||
runtime?: "MlxRing" | "MlxJaccl";
|
||||
onLaunch?: () => void;
|
||||
tags?: string[];
|
||||
apiPreview?: PlacementPreview | null;
|
||||
@@ -348,7 +348,7 @@
|
||||
// Debug mode state
|
||||
const isDebugMode = $derived(debugMode());
|
||||
const topology = $derived(topologyData());
|
||||
const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl");
|
||||
const isRdma = $derived(runtime === "MlxJaccl");
|
||||
|
||||
// Get interface name for an IP from node data
|
||||
function getInterfaceForIp(nodeId: string, ip?: string): string | null {
|
||||
@@ -567,20 +567,46 @@
|
||||
<div class="flex items-center gap-1.5 mb-2">
|
||||
<span
|
||||
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
|
||||
title={sharding === "Pipeline"
|
||||
? "Pipeline: splits model into sequential stages across devices. Lower network overhead."
|
||||
: "Tensor: splits each layer across devices. Best with high-bandwidth connections (Thunderbolt)."}
|
||||
>
|
||||
{sharding}
|
||||
</span>
|
||||
<span
|
||||
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
|
||||
title={runtime === "MlxRing"
|
||||
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
|
||||
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
|
||||
>
|
||||
{runtime === "MlxRing"
|
||||
? "MLX Ring"
|
||||
: runtime === "MlxIbv" || runtime === "MlxJaccl"
|
||||
: runtime === "MlxJaccl"
|
||||
? "MLX RDMA"
|
||||
: runtime}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<!-- Download Status -->
|
||||
{#if isDownloading && progress}
|
||||
<div class="mb-2 space-y-1">
|
||||
<div class="flex items-center justify-between text-xs font-mono">
|
||||
<span class="text-blue-400 tracking-wider uppercase">Downloading</span
|
||||
>
|
||||
<span class="text-white/60"
|
||||
>{percentage.toFixed(1)}% · {formatSpeed(progress.speed)}
|
||||
· {formatEta(progress.etaMs)}</span
|
||||
>
|
||||
</div>
|
||||
<div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
|
||||
<div
|
||||
class="h-full bg-blue-500/70 transition-all duration-300"
|
||||
style="width: {percentage}%"
|
||||
></div>
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
<!-- Mini Topology Preview -->
|
||||
{#if placementPreview().nodes.length > 0}
|
||||
{@const preview = placementPreview()}
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
capabilities: string[];
|
||||
sizeRange: { min: number; max: number } | null;
|
||||
downloadedOnly: boolean;
|
||||
readyOnly: boolean;
|
||||
}
|
||||
|
||||
type ModelFilterPopoverProps = {
|
||||
@@ -190,34 +191,58 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Downloaded only -->
|
||||
<!-- Availability filters -->
|
||||
<div>
|
||||
<h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
<div class="flex flex-wrap gap-1.5">
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
|
||||
>
|
||||
<path
|
||||
class="text-white/40"
|
||||
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
|
||||
/>
|
||||
<path class="text-green-400" d="m9 13 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Downloaded</span>
|
||||
</button>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path
|
||||
class="text-white/40"
|
||||
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
|
||||
/>
|
||||
<path class="text-green-400" d="m9 13 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Downloaded</span>
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, readyOnly: !filters.readyOnly })}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<circle cx="12" cy="12" r="10" />
|
||||
<path d="m9 12 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Ready</span>
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Size range -->
|
||||
|
||||
@@ -41,6 +41,7 @@
|
||||
onShowInfo: (group: ModelGroup) => void;
|
||||
downloadStatusMap?: Map<string, DownloadAvailability>;
|
||||
launchedAt?: number;
|
||||
instanceStatuses?: Record<string, { status: string; statusClass: string }>;
|
||||
};
|
||||
|
||||
let {
|
||||
@@ -56,6 +57,7 @@
|
||||
onShowInfo,
|
||||
downloadStatusMap,
|
||||
launchedAt,
|
||||
instanceStatuses = {},
|
||||
}: ModelPickerGroupProps = $props();
|
||||
|
||||
// Group-level download status: show if any variant is downloaded
|
||||
@@ -92,6 +94,12 @@
|
||||
const anyVariantFits = $derived(
|
||||
group.variants.some((v) => canModelFit(v.id)),
|
||||
);
|
||||
// Check if any variant has an active instance (ready, loading, downloading)
|
||||
const anyVariantHasInstance = $derived(
|
||||
instanceStatuses
|
||||
? group.variants.some((v) => instanceStatuses[v.id] != null)
|
||||
: false,
|
||||
);
|
||||
const groupFitStatus = $derived.by((): ModelFitStatus => {
|
||||
let hasClusterCapacityOnly = false;
|
||||
for (const variant of group.variants) {
|
||||
@@ -122,16 +130,35 @@
|
||||
!group.hasMultipleVariants &&
|
||||
group.variants.some((v) => v.id === selectedModelId),
|
||||
);
|
||||
|
||||
// Group-level instance status: show the "best" status across all variants
|
||||
const groupInstanceStatus = $derived.by(() => {
|
||||
if (!instanceStatuses) return null;
|
||||
const readyStatuses = ["READY", "LOADED", "RUNNING"];
|
||||
const loadingStatuses = ["LOADING", "WARMING UP"];
|
||||
let bestStatus: { status: string; statusClass: string } | null = null;
|
||||
for (const variant of group.variants) {
|
||||
const s = instanceStatuses[variant.id];
|
||||
if (!s) continue;
|
||||
if (readyStatuses.includes(s.status)) return s; // Ready is best
|
||||
if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
|
||||
bestStatus = s;
|
||||
}
|
||||
}
|
||||
return bestStatus;
|
||||
});
|
||||
</script>
|
||||
|
||||
<div
|
||||
class="border-b border-white/5 last:border-b-0 {!anyVariantFits
|
||||
class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
|
||||
!anyVariantHasInstance
|
||||
? 'opacity-50'
|
||||
: ''}"
|
||||
>
|
||||
<!-- Main row -->
|
||||
<div
|
||||
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
|
||||
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
|
||||
anyVariantHasInstance
|
||||
? 'hover:bg-white/5 cursor-pointer'
|
||||
: 'cursor-not-allowed'} {isMainSelected
|
||||
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
|
||||
@@ -141,7 +168,7 @@
|
||||
onToggleExpand();
|
||||
} else {
|
||||
const modelId = group.variants[0]?.id;
|
||||
if (modelId && canModelFit(modelId)) {
|
||||
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
|
||||
onSelectModel(modelId);
|
||||
}
|
||||
}
|
||||
@@ -155,7 +182,7 @@
|
||||
onToggleExpand();
|
||||
} else {
|
||||
const modelId = group.variants[0]?.id;
|
||||
if (modelId && canModelFit(modelId)) {
|
||||
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
|
||||
onSelectModel(modelId);
|
||||
}
|
||||
}
|
||||
@@ -346,6 +373,47 @@
|
||||
</span>
|
||||
{/if}
|
||||
|
||||
<!-- Instance status badge -->
|
||||
{#if groupInstanceStatus}
|
||||
{#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
|
||||
<span class="flex-shrink-0" title="Running">
|
||||
<svg
|
||||
class="w-3 h-3 text-green-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if groupInstanceStatus.status === "DOWNLOADING"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Downloading">
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-blue-400"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
|
||||
<polyline points="7 10 12 15 17 10" />
|
||||
<line x1="12" y1="15" x2="12" y2="3" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Loading">
|
||||
<svg
|
||||
class="w-3 h-3 text-yellow-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Check mark if selected (single-variant) -->
|
||||
{#if isMainSelected}
|
||||
<svg
|
||||
@@ -420,20 +488,30 @@
|
||||
{#each group.variants as variant}
|
||||
{@const fitStatus = getModelFitStatus(variant.id)}
|
||||
{@const modelCanFit = canModelFit(variant.id)}
|
||||
{@const variantHasInstance = instanceStatuses[variant.id] != null}
|
||||
{@const isSelected = selectedModelId === variant.id}
|
||||
<button
|
||||
type="button"
|
||||
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
|
||||
<div
|
||||
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
|
||||
!variantHasInstance
|
||||
? 'opacity-50 cursor-not-allowed'
|
||||
: 'cursor-pointer'} {isSelected
|
||||
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
|
||||
: 'border-l-2 border-transparent'}"
|
||||
disabled={!modelCanFit}
|
||||
role="button"
|
||||
tabindex="0"
|
||||
onclick={() => {
|
||||
if (modelCanFit) {
|
||||
if (modelCanFit || variantHasInstance) {
|
||||
onSelectModel(variant.id);
|
||||
}
|
||||
}}
|
||||
onkeydown={(e) => {
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
if (modelCanFit) {
|
||||
onSelectModel(variant.id);
|
||||
}
|
||||
}
|
||||
}}
|
||||
>
|
||||
<!-- Quantization badge -->
|
||||
<span
|
||||
@@ -478,6 +556,48 @@
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Instance status badge -->
|
||||
{#if instanceStatuses[variant.id]}
|
||||
{@const instStatus = instanceStatuses[variant.id]}
|
||||
{#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
|
||||
<span class="flex-shrink-0" title="Running">
|
||||
<svg
|
||||
class="w-3 h-3 text-green-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if instStatus.status === "DOWNLOADING"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Downloading">
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-blue-400"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
|
||||
<polyline points="7 10 12 15 17 10" />
|
||||
<line x1="12" y1="15" x2="12" y2="3" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Loading">
|
||||
<svg
|
||||
class="w-3 h-3 text-yellow-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Check mark if selected -->
|
||||
{#if isSelected}
|
||||
<svg
|
||||
@@ -490,7 +610,36 @@
|
||||
/>
|
||||
</svg>
|
||||
{/if}
|
||||
</button>
|
||||
|
||||
<!-- Info button -->
|
||||
<button
|
||||
type="button"
|
||||
class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
|
||||
onclick={(e) => {
|
||||
e.stopPropagation();
|
||||
onShowInfo({
|
||||
id: variant.id,
|
||||
name: variant.name || variant.id,
|
||||
capabilities: group.capabilities,
|
||||
family: group.family,
|
||||
variants: [variant],
|
||||
smallestVariant: variant,
|
||||
hasMultipleVariants: false,
|
||||
});
|
||||
}}
|
||||
title="View variant details"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4 text-white/30 hover:text-white/50"
|
||||
viewBox="0 0 24 24"
|
||||
fill="currentColor"
|
||||
>
|
||||
<path
|
||||
d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
|
||||
/>
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
{/each}
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
@@ -36,6 +36,7 @@
|
||||
capabilities: string[];
|
||||
sizeRange: { min: number; max: number } | null;
|
||||
downloadedOnly: boolean;
|
||||
readyOnly: boolean;
|
||||
}
|
||||
|
||||
interface HuggingFaceModel {
|
||||
@@ -49,6 +50,11 @@
|
||||
|
||||
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
|
||||
|
||||
export type InstanceStatus = {
|
||||
status: string;
|
||||
statusClass: string;
|
||||
};
|
||||
|
||||
type ModelPickerModalProps = {
|
||||
isOpen: boolean;
|
||||
models: ModelInfo[];
|
||||
@@ -75,6 +81,7 @@
|
||||
macmon_info?: { memory?: { ram_total?: number } };
|
||||
}
|
||||
>;
|
||||
instanceStatuses?: Record<string, InstanceStatus>;
|
||||
};
|
||||
|
||||
let {
|
||||
@@ -96,6 +103,7 @@
|
||||
usedMemoryGB,
|
||||
downloadsData,
|
||||
topologyNodes,
|
||||
instanceStatuses = {},
|
||||
}: ModelPickerModalProps = $props();
|
||||
|
||||
// Local state
|
||||
@@ -107,6 +115,7 @@
|
||||
capabilities: [],
|
||||
sizeRange: null,
|
||||
downloadedOnly: false,
|
||||
readyOnly: false,
|
||||
});
|
||||
let infoGroup = $state<ModelGroup | null>(null);
|
||||
|
||||
@@ -440,6 +449,16 @@
|
||||
);
|
||||
}
|
||||
|
||||
// Filter to ready/running models only
|
||||
if (filters.readyOnly) {
|
||||
result = result.filter((g) =>
|
||||
g.variants.some((v) => {
|
||||
const s = instanceStatuses[v.id];
|
||||
return s && s.statusClass === "ready";
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Sort: fits-now first, then fits-cluster-capacity, then too-large
|
||||
result.sort((a, b) => {
|
||||
const getGroupFitRank = (group: ModelGroup): number => {
|
||||
@@ -512,6 +531,18 @@
|
||||
);
|
||||
});
|
||||
|
||||
// Split filtered groups into recommended (fits_now) and others for visual separation
|
||||
const recommendedGroups = $derived(
|
||||
filteredGroups.filter((g) =>
|
||||
g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
|
||||
),
|
||||
);
|
||||
const otherGroups = $derived(
|
||||
filteredGroups.filter(
|
||||
(g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
|
||||
),
|
||||
);
|
||||
|
||||
function toggleGroupExpanded(groupId: string) {
|
||||
const next = new Set(expandedGroups);
|
||||
if (next.has(groupId)) {
|
||||
@@ -538,13 +569,19 @@
|
||||
}
|
||||
|
||||
function clearFilters() {
|
||||
filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
|
||||
filters = {
|
||||
capabilities: [],
|
||||
sizeRange: null,
|
||||
downloadedOnly: false,
|
||||
readyOnly: false,
|
||||
};
|
||||
}
|
||||
|
||||
const hasActiveFilters = $derived(
|
||||
filters.capabilities.length > 0 ||
|
||||
filters.sizeRange !== null ||
|
||||
filters.downloadedOnly,
|
||||
filters.downloadedOnly ||
|
||||
filters.readyOnly,
|
||||
);
|
||||
</script>
|
||||
|
||||
@@ -813,6 +850,7 @@
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
{/if}
|
||||
@@ -840,7 +878,34 @@
|
||||
{/if}
|
||||
</div>
|
||||
{:else}
|
||||
{#each filteredGroups as group}
|
||||
<!-- Recommended for your cluster -->
|
||||
{#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()}
|
||||
<div
|
||||
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-green-950/60 border-b border-green-500/20 backdrop-blur-sm"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-green-400 flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono text-green-400 tracking-wider uppercase"
|
||||
>Recommended for your cluster</span
|
||||
>
|
||||
<span class="text-xs font-mono text-green-400/50"
|
||||
>— fits in available memory</span
|
||||
>
|
||||
</div>
|
||||
{/if}
|
||||
{#each recommendedGroups as group}
|
||||
<ModelPickerGroup
|
||||
{group}
|
||||
isExpanded={expandedGroups.has(group.id)}
|
||||
@@ -853,6 +918,34 @@
|
||||
{onToggleFavorite}
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
<!-- Other models -->
|
||||
{#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()}
|
||||
<div
|
||||
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-exo-dark-gray/80 border-y border-exo-medium-gray/20 backdrop-blur-sm"
|
||||
>
|
||||
<span
|
||||
class="text-xs font-mono text-white/40 tracking-wider uppercase"
|
||||
>Other models</span
|
||||
>
|
||||
</div>
|
||||
{/if}
|
||||
{#each otherGroups as group}
|
||||
<ModelPickerGroup
|
||||
{group}
|
||||
isExpanded={expandedGroups.has(group.id)}
|
||||
isFavorite={favorites.has(group.id)}
|
||||
{selectedModelId}
|
||||
{canModelFit}
|
||||
{getModelFitStatus}
|
||||
onToggleExpand={() => toggleGroupExpanded(group.id)}
|
||||
onSelectModel={handleSelect}
|
||||
{onToggleFavorite}
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
{/if}
|
||||
@@ -875,6 +968,11 @@
|
||||
>Downloaded</span
|
||||
>
|
||||
{/if}
|
||||
{#if filters.readyOnly}
|
||||
<span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
|
||||
>Ready</span
|
||||
>
|
||||
{/if}
|
||||
{#if filters.sizeRange}
|
||||
<span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
|
||||
{filters.sizeRange.min}GB - {filters.sizeRange.max}GB
|
||||
|
||||
@@ -14,6 +14,21 @@
|
||||
: 0,
|
||||
);
|
||||
|
||||
const etaText = $derived.by(() => {
|
||||
if (progress.processed <= 0 || progress.total <= 0) return null;
|
||||
const elapsedMs = performance.now() - progress.startedAt;
|
||||
if (elapsedMs < 200) return null; // need a minimum sample window
|
||||
const tokensPerMs = progress.processed / elapsedMs;
|
||||
const remainingTokens = progress.total - progress.processed;
|
||||
const remainingMs = remainingTokens / tokensPerMs;
|
||||
const remainingSec = Math.ceil(remainingMs / 1000);
|
||||
if (remainingSec <= 0) return null;
|
||||
if (remainingSec < 60) return `~${remainingSec}s remaining`;
|
||||
const mins = Math.floor(remainingSec / 60);
|
||||
const secs = remainingSec % 60;
|
||||
return `~${mins}m ${secs}s remaining`;
|
||||
});
|
||||
|
||||
function formatTokenCount(count: number | undefined): string {
|
||||
if (count == null) return "0";
|
||||
if (count >= 1000) {
|
||||
@@ -40,8 +55,11 @@
|
||||
style="width: {percentage}%"
|
||||
></div>
|
||||
</div>
|
||||
<div class="text-right text-xs text-exo-light-gray/70 mt-0.5 font-mono">
|
||||
{percentage}%
|
||||
<div
|
||||
class="flex items-center justify-between text-xs text-exo-light-gray/70 mt-0.5 font-mono"
|
||||
>
|
||||
<span>{etaText ?? ""}</span>
|
||||
<span>{percentage}%</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
<script lang="ts">
|
||||
import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte";
|
||||
import { fly, fade } from "svelte/transition";
|
||||
import { flip } from "svelte/animate";
|
||||
|
||||
const items = $derived(toasts());
|
||||
|
||||
const typeStyles: Record<
|
||||
Toast["type"],
|
||||
{ border: string; icon: string; iconColor: string }
|
||||
> = {
|
||||
success: {
|
||||
border: "border-l-green-500",
|
||||
icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z",
|
||||
iconColor: "text-green-400",
|
||||
},
|
||||
error: {
|
||||
border: "border-l-red-500",
|
||||
icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z",
|
||||
iconColor: "text-red-400",
|
||||
},
|
||||
warning: {
|
||||
border: "border-l-yellow-500",
|
||||
icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z",
|
||||
iconColor: "text-yellow-400",
|
||||
},
|
||||
info: {
|
||||
border: "border-l-blue-500",
|
||||
icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z",
|
||||
iconColor: "text-blue-400",
|
||||
},
|
||||
};
|
||||
</script>
|
||||
|
||||
{#if items.length > 0}
|
||||
<div
|
||||
class="fixed bottom-6 right-6 z-[9999] flex flex-col gap-2 pointer-events-none"
|
||||
role="log"
|
||||
aria-live="polite"
|
||||
aria-label="Notifications"
|
||||
>
|
||||
{#each items as toast (toast.id)}
|
||||
{@const style = typeStyles[toast.type]}
|
||||
<div
|
||||
class="pointer-events-auto max-w-sm w-80 bg-exo-dark-gray/95 backdrop-blur-sm border border-exo-medium-gray/60 border-l-[3px] {style.border} rounded shadow-lg shadow-black/40"
|
||||
in:fly={{ x: 80, duration: 250 }}
|
||||
out:fade={{ duration: 150 }}
|
||||
animate:flip={{ duration: 200 }}
|
||||
role="alert"
|
||||
>
|
||||
<div class="flex items-start gap-3 px-4 py-3">
|
||||
<!-- Icon -->
|
||||
<svg
|
||||
class="w-5 h-5 flex-shrink-0 mt-0.5 {style.iconColor}"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke-width="1.5"
|
||||
stroke="currentColor"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={style.icon}
|
||||
/>
|
||||
</svg>
|
||||
|
||||
<!-- Message -->
|
||||
<p class="flex-1 text-sm text-white/90 font-mono leading-snug">
|
||||
{toast.message}
|
||||
</p>
|
||||
|
||||
<!-- Dismiss button -->
|
||||
<button
|
||||
onclick={() => dismissToast(toast.id)}
|
||||
class="flex-shrink-0 p-0.5 text-white/40 hover:text-white/80 transition-colors cursor-pointer"
|
||||
aria-label="Dismiss notification"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke-width="2"
|
||||
stroke="currentColor"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M6 18L18 6M6 6l12 12"
|
||||
/>
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<!-- Auto-dismiss progress bar -->
|
||||
{#if toast.duration > 0}
|
||||
<div class="h-0.5 bg-white/5 rounded-b overflow-hidden">
|
||||
<div
|
||||
class="h-full {style.border.replace('border-l-', 'bg-')}/60"
|
||||
style="animation: shrink {toast.duration}ms linear forwards"
|
||||
></div>
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
{/each}
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
<style>
|
||||
@keyframes shrink {
|
||||
from {
|
||||
width: 100%;
|
||||
}
|
||||
to {
|
||||
width: 0%;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
@@ -12,3 +12,4 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
|
||||
export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
|
||||
export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
|
||||
export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
|
||||
export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
|
||||
|
||||
@@ -168,7 +168,7 @@ export interface ModelDownloadStatus {
|
||||
export interface PlacementPreview {
|
||||
model_id: string;
|
||||
sharding: "Pipeline" | "Tensor";
|
||||
instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl";
|
||||
instance_meta: "MlxRing" | "MlxJaccl";
|
||||
instance: unknown | null;
|
||||
memory_delta_by_node: Record<string, number> | null;
|
||||
error: string | null;
|
||||
@@ -219,7 +219,6 @@ interface RawStateResponse {
|
||||
string,
|
||||
{
|
||||
MlxRingInstance?: Instance;
|
||||
MlxIbvInstance?: Instance;
|
||||
MlxJacclInstance?: Instance;
|
||||
}
|
||||
>;
|
||||
@@ -250,6 +249,11 @@ interface RawStateResponse {
|
||||
>;
|
||||
// Thunderbolt bridge cycles (nodes with bridge enabled forming loops)
|
||||
thunderboltBridgeCycles?: string[][];
|
||||
// Disk usage per node
|
||||
nodeDisk?: Record<
|
||||
string,
|
||||
{ total: { inBytes: number }; available: { inBytes: number } }
|
||||
>;
|
||||
}
|
||||
|
||||
export interface MessageAttachment {
|
||||
@@ -276,6 +280,8 @@ export interface TokenData {
|
||||
export interface PrefillProgress {
|
||||
processed: number;
|
||||
total: number;
|
||||
/** Timestamp (performance.now()) when prefill started. */
|
||||
startedAt: number;
|
||||
}
|
||||
|
||||
export interface Message {
|
||||
@@ -583,6 +589,12 @@ class AppStore {
|
||||
// Image editing state
|
||||
editingImage = $state<EditingImage | null>(null);
|
||||
|
||||
/** True when the backend is reachable. */
|
||||
isConnected = $state<boolean>(true);
|
||||
/** Number of consecutive fetch failures. */
|
||||
private consecutiveFailures = 0;
|
||||
private static readonly CONNECTION_LOST_THRESHOLD = 3;
|
||||
|
||||
private fetchInterval: ReturnType<typeof setInterval> | null = null;
|
||||
private previewsInterval: ReturnType<typeof setInterval> | null = null;
|
||||
private lastConversationPersistTs = 0;
|
||||
@@ -835,6 +847,10 @@ class AppStore {
|
||||
this.thinkingEnabled = conversation.enableThinking ?? true;
|
||||
this.refreshConversationModelFromInstances();
|
||||
|
||||
// Sync global selection to the loaded conversation's model so reactive
|
||||
// effects in +page.svelte can determine the correct chat launch state.
|
||||
this.selectedChatModel = conversation.modelId || "";
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -905,11 +921,7 @@ class AppStore {
|
||||
|
||||
let instanceType: string | null = null;
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (
|
||||
instanceTag === "MlxIbvInstance" ||
|
||||
instanceTag === "MlxJacclInstance"
|
||||
)
|
||||
instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
@@ -1288,7 +1300,19 @@ class AppStore {
|
||||
// Thunderbolt bridge status per node
|
||||
this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
|
||||
this.lastUpdate = Date.now();
|
||||
// Connection recovered
|
||||
if (!this.isConnected) {
|
||||
this.isConnected = true;
|
||||
}
|
||||
this.consecutiveFailures = 0;
|
||||
} catch (error) {
|
||||
this.consecutiveFailures++;
|
||||
if (
|
||||
this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD &&
|
||||
this.isConnected
|
||||
) {
|
||||
this.isConnected = false;
|
||||
}
|
||||
console.error("Error fetching state:", error);
|
||||
}
|
||||
}
|
||||
@@ -1652,11 +1676,12 @@ class AppStore {
|
||||
if (!reader) throw new Error("No response body");
|
||||
|
||||
let fullContent = prefixText;
|
||||
let streamedThinking = "";
|
||||
const collectedTokens: TokenData[] = [...tokensToKeep];
|
||||
|
||||
interface ChatCompletionChunk {
|
||||
choices?: Array<{
|
||||
delta?: { content?: string };
|
||||
delta?: { content?: string; reasoning_content?: string };
|
||||
logprobs?: {
|
||||
content?: Array<{
|
||||
token: string;
|
||||
@@ -1677,6 +1702,7 @@ class AppStore {
|
||||
(parsed) => {
|
||||
const choice = parsed.choices?.[0];
|
||||
const delta = choice?.delta?.content;
|
||||
const thinkingDelta = choice?.delta?.reasoning_content;
|
||||
|
||||
// Collect logprobs data
|
||||
const logprobsContent = choice?.logprobs?.content;
|
||||
@@ -1695,7 +1721,11 @@ class AppStore {
|
||||
}
|
||||
}
|
||||
|
||||
if (delta) {
|
||||
if (thinkingDelta) {
|
||||
streamedThinking += thinkingDelta;
|
||||
}
|
||||
|
||||
if (delta || thinkingDelta) {
|
||||
if (firstTokenTime === null) {
|
||||
firstTokenTime = performance.now();
|
||||
this.ttftMs = firstTokenTime - requestStartTime;
|
||||
@@ -1709,9 +1739,14 @@ class AppStore {
|
||||
this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000;
|
||||
}
|
||||
|
||||
fullContent += delta;
|
||||
const { displayContent, thinkingContent } =
|
||||
if (delta) {
|
||||
fullContent += delta;
|
||||
}
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(fullContent);
|
||||
const combinedThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
|
||||
if (this.activeConversationId === targetConversationId) {
|
||||
this.currentResponse = displayContent;
|
||||
@@ -1723,7 +1758,7 @@ class AppStore {
|
||||
messageId,
|
||||
(m) => {
|
||||
m.content = displayContent;
|
||||
m.thinking = thinkingContent || undefined;
|
||||
m.thinking = combinedThinking || undefined;
|
||||
m.tokens = [...collectedTokens];
|
||||
},
|
||||
);
|
||||
@@ -1735,11 +1770,14 @@ class AppStore {
|
||||
|
||||
// Final update
|
||||
if (this.conversationExists(targetConversationId)) {
|
||||
const { displayContent, thinkingContent } =
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(fullContent);
|
||||
const finalThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
this.updateConversationMessage(targetConversationId, messageId, (m) => {
|
||||
m.content = displayContent;
|
||||
m.thinking = thinkingContent || undefined;
|
||||
m.thinking = finalThinking || undefined;
|
||||
m.tokens = [...collectedTokens];
|
||||
if (this.ttftMs !== null) m.ttftMs = this.ttftMs;
|
||||
if (this.tps !== null) m.tps = this.tps;
|
||||
@@ -1815,7 +1853,7 @@ class AppStore {
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content =
|
||||
"Error: No model available. Please launch an instance first.";
|
||||
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.";
|
||||
},
|
||||
);
|
||||
this.syncActiveMessagesIfNeeded(targetConversationId);
|
||||
@@ -1847,11 +1885,12 @@ class AppStore {
|
||||
}
|
||||
|
||||
let streamedContent = "";
|
||||
let streamedThinking = "";
|
||||
const collectedTokens: TokenData[] = [];
|
||||
|
||||
interface ChatCompletionChunk {
|
||||
choices?: Array<{
|
||||
delta?: { content?: string };
|
||||
delta?: { content?: string; reasoning_content?: string };
|
||||
logprobs?: {
|
||||
content?: Array<{
|
||||
token: string;
|
||||
@@ -1872,6 +1911,7 @@ class AppStore {
|
||||
(parsed) => {
|
||||
const choice = parsed.choices?.[0];
|
||||
const delta = choice?.delta?.content;
|
||||
const thinkingDelta = choice?.delta?.reasoning_content;
|
||||
|
||||
// Collect logprobs data
|
||||
const logprobsContent = choice?.logprobs?.content;
|
||||
@@ -1890,10 +1930,19 @@ class AppStore {
|
||||
}
|
||||
}
|
||||
|
||||
if (delta) {
|
||||
streamedContent += delta;
|
||||
const { displayContent, thinkingContent } =
|
||||
if (thinkingDelta) {
|
||||
streamedThinking += thinkingDelta;
|
||||
}
|
||||
|
||||
if (delta || thinkingDelta) {
|
||||
if (delta) {
|
||||
streamedContent += delta;
|
||||
}
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(streamedContent);
|
||||
const combinedThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
|
||||
// Only update currentResponse if target conversation is active
|
||||
if (this.activeConversationId === targetConversationId) {
|
||||
@@ -1906,7 +1955,7 @@ class AppStore {
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content = displayContent;
|
||||
msg.thinking = thinkingContent || undefined;
|
||||
msg.thinking = combinedThinking || undefined;
|
||||
msg.tokens = [...collectedTokens];
|
||||
},
|
||||
);
|
||||
@@ -1918,14 +1967,17 @@ class AppStore {
|
||||
|
||||
// Final cleanup of the message (if conversation still exists)
|
||||
if (this.conversationExists(targetConversationId)) {
|
||||
const { displayContent, thinkingContent } =
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(streamedContent);
|
||||
const finalThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
this.updateConversationMessage(
|
||||
targetConversationId,
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content = displayContent;
|
||||
msg.thinking = thinkingContent || undefined;
|
||||
msg.thinking = finalThinking || undefined;
|
||||
msg.tokens = [...collectedTokens];
|
||||
},
|
||||
);
|
||||
@@ -2272,7 +2324,7 @@ class AppStore {
|
||||
const modelToUse = this.getModelForRequest();
|
||||
if (!modelToUse) {
|
||||
throw new Error(
|
||||
"No model selected and no running instances available. Please launch an instance first.",
|
||||
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.",
|
||||
);
|
||||
}
|
||||
|
||||
@@ -2317,10 +2369,11 @@ class AppStore {
|
||||
}
|
||||
|
||||
let streamedContent = "";
|
||||
let streamedThinking = "";
|
||||
|
||||
interface ChatCompletionChunk {
|
||||
choices?: Array<{
|
||||
delta?: { content?: string };
|
||||
delta?: { content?: string; reasoning_content?: string };
|
||||
logprobs?: {
|
||||
content?: Array<{
|
||||
token: string;
|
||||
@@ -2348,6 +2401,7 @@ class AppStore {
|
||||
|
||||
const choice = parsed.choices?.[0];
|
||||
const tokenContent = choice?.delta?.content;
|
||||
const thinkingContent = choice?.delta?.reasoning_content;
|
||||
|
||||
// Collect logprobs data
|
||||
const logprobsContent = choice?.logprobs?.content;
|
||||
@@ -2366,7 +2420,11 @@ class AppStore {
|
||||
}
|
||||
}
|
||||
|
||||
if (tokenContent) {
|
||||
if (thinkingContent) {
|
||||
streamedThinking += thinkingContent;
|
||||
}
|
||||
|
||||
if (tokenContent || thinkingContent) {
|
||||
// Track first token for TTFT
|
||||
if (firstTokenTime === null) {
|
||||
firstTokenTime = performance.now();
|
||||
@@ -2383,11 +2441,16 @@ class AppStore {
|
||||
this.tps = (tokenCount / elapsed) * 1000;
|
||||
}
|
||||
|
||||
streamedContent += tokenContent;
|
||||
if (tokenContent) {
|
||||
streamedContent += tokenContent;
|
||||
}
|
||||
|
||||
// Strip thinking tags for display and extract thinking content
|
||||
const { displayContent, thinkingContent } =
|
||||
// Use stripThinkingTags as fallback for any <think> tags still in content
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(streamedContent);
|
||||
const combinedThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
|
||||
// Only update currentResponse if target conversation is active
|
||||
if (this.activeConversationId === targetConversationId) {
|
||||
@@ -2400,7 +2463,7 @@ class AppStore {
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content = displayContent;
|
||||
msg.thinking = thinkingContent || undefined;
|
||||
msg.thinking = combinedThinking || undefined;
|
||||
msg.tokens = [...collectedTokens];
|
||||
},
|
||||
);
|
||||
@@ -2420,6 +2483,7 @@ class AppStore {
|
||||
this.prefillProgress = {
|
||||
processed: inner.processed_tokens,
|
||||
total: inner.total_tokens,
|
||||
startedAt: this.prefillProgress?.startedAt ?? performance.now(),
|
||||
};
|
||||
},
|
||||
},
|
||||
@@ -2436,14 +2500,17 @@ class AppStore {
|
||||
|
||||
// Final cleanup of the message (if conversation still exists)
|
||||
if (this.conversationExists(targetConversationId)) {
|
||||
const { displayContent, thinkingContent } =
|
||||
const { displayContent, thinkingContent: tagThinking } =
|
||||
this.stripThinkingTags(streamedContent);
|
||||
const finalThinking = [streamedThinking, tagThinking]
|
||||
.filter(Boolean)
|
||||
.join("\n\n");
|
||||
this.updateConversationMessage(
|
||||
targetConversationId,
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content = displayContent;
|
||||
msg.thinking = thinkingContent || undefined;
|
||||
msg.thinking = finalThinking || undefined;
|
||||
msg.tokens = [...collectedTokens];
|
||||
// Store performance metrics on the message
|
||||
if (this.ttftMs !== null) {
|
||||
@@ -3091,6 +3158,23 @@ class AppStore {
|
||||
return (await response.json()) as TraceStatsResponse;
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete traces by task IDs
|
||||
*/
|
||||
async deleteTraces(
|
||||
taskIds: string[],
|
||||
): Promise<{ deleted: string[]; notFound: string[] }> {
|
||||
const response = await fetch("/v1/traces/delete", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ taskIds }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to delete traces: ${response.status}`);
|
||||
}
|
||||
return await response.json();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the URL for the raw trace file (for Perfetto)
|
||||
*/
|
||||
@@ -3200,6 +3284,9 @@ export const setChatSidebarVisible = (visible: boolean) =>
|
||||
appStore.setChatSidebarVisible(visible);
|
||||
export const refreshState = () => appStore.fetchState();
|
||||
|
||||
// Connection status
|
||||
export const isConnected = () => appStore.isConnected;
|
||||
|
||||
// Node identities (for OS version mismatch detection)
|
||||
export const nodeIdentities = () => appStore.nodeIdentities;
|
||||
|
||||
@@ -3231,3 +3318,5 @@ export const fetchTraceStats = (taskId: string) =>
|
||||
appStore.fetchTraceStats(taskId);
|
||||
export const getTraceRawUrl = (taskId: string) =>
|
||||
appStore.getTraceRawUrl(taskId);
|
||||
export const deleteTraces = (taskIds: string[]) =>
|
||||
appStore.deleteTraces(taskIds);
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Toast notification store - Global notification system for the EXO dashboard.
|
||||
*
|
||||
* Usage:
|
||||
* import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte";
|
||||
* addToast({ type: "success", message: "Model launched" });
|
||||
* addToast({ type: "error", message: "Connection lost", persistent: true });
|
||||
*/
|
||||
|
||||
type ToastType = "success" | "error" | "warning" | "info";
|
||||
|
||||
export interface Toast {
|
||||
id: string;
|
||||
type: ToastType;
|
||||
message: string;
|
||||
/** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */
|
||||
duration: number;
|
||||
createdAt: number;
|
||||
}
|
||||
|
||||
interface ToastInput {
|
||||
type: ToastType;
|
||||
message: string;
|
||||
/** If true, toast stays until manually dismissed. Default: false. */
|
||||
persistent?: boolean;
|
||||
/** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */
|
||||
duration?: number;
|
||||
}
|
||||
|
||||
const DEFAULT_DURATIONS: Record<ToastType, number> = {
|
||||
success: 4000,
|
||||
info: 4000,
|
||||
warning: 6000,
|
||||
error: 6000,
|
||||
};
|
||||
|
||||
let toastList = $state<Toast[]>([]);
|
||||
const timers = new Map<string, ReturnType<typeof setTimeout>>();
|
||||
|
||||
function generateId(): string {
|
||||
return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||||
}
|
||||
|
||||
export function addToast(input: ToastInput): string {
|
||||
const id = generateId();
|
||||
const duration = input.persistent
|
||||
? 0
|
||||
: (input.duration ?? DEFAULT_DURATIONS[input.type]);
|
||||
|
||||
const toast: Toast = {
|
||||
id,
|
||||
type: input.type,
|
||||
message: input.message,
|
||||
duration,
|
||||
createdAt: Date.now(),
|
||||
};
|
||||
|
||||
toastList = [...toastList, toast];
|
||||
|
||||
if (duration > 0) {
|
||||
const timer = setTimeout(() => dismissToast(id), duration);
|
||||
timers.set(id, timer);
|
||||
}
|
||||
|
||||
return id;
|
||||
}
|
||||
|
||||
export function dismissToast(id: string): void {
|
||||
const timer = timers.get(id);
|
||||
if (timer) {
|
||||
clearTimeout(timer);
|
||||
timers.delete(id);
|
||||
}
|
||||
toastList = toastList.filter((t) => t.id !== id);
|
||||
}
|
||||
|
||||
/** Dismiss all toasts matching a message (useful for dedup). */
|
||||
export function dismissByMessage(message: string): void {
|
||||
const matching = toastList.filter((t) => t.message === message);
|
||||
for (const t of matching) {
|
||||
dismissToast(t.id);
|
||||
}
|
||||
}
|
||||
|
||||
export function toasts(): Toast[] {
|
||||
return toastList;
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
<script lang="ts">
|
||||
import "../app.css";
|
||||
import ToastContainer from "$lib/components/ToastContainer.svelte";
|
||||
import ConnectionBanner from "$lib/components/ConnectionBanner.svelte";
|
||||
|
||||
let { children } = $props();
|
||||
</script>
|
||||
@@ -10,5 +12,7 @@
|
||||
</svelte:head>
|
||||
|
||||
<div class="min-h-screen bg-background text-foreground">
|
||||
<ConnectionBanner />
|
||||
{@render children?.()}
|
||||
<ToastContainer />
|
||||
</div>
|
||||
|
||||
+3427
-381
File diff suppressed because it is too large
Load Diff
@@ -29,7 +29,12 @@
|
||||
etaMs: number;
|
||||
modelDirectory?: string;
|
||||
}
|
||||
| { kind: "pending"; modelDirectory?: string }
|
||||
| {
|
||||
kind: "pending";
|
||||
downloaded: number;
|
||||
total: number;
|
||||
modelDirectory?: string;
|
||||
}
|
||||
| { kind: "failed"; modelDirectory?: string }
|
||||
| { kind: "not_present" };
|
||||
|
||||
@@ -74,7 +79,6 @@
|
||||
if (typeof value === "number") return value;
|
||||
if (value && typeof value === "object") {
|
||||
const v = value as Record<string, unknown>;
|
||||
if (typeof v.in_bytes === "number") return v.in_bytes;
|
||||
if (typeof v.inBytes === "number") return v.inBytes;
|
||||
}
|
||||
return 0;
|
||||
@@ -231,23 +235,14 @@
|
||||
undefined;
|
||||
let cell: CellStatus;
|
||||
if (tag === "DownloadCompleted") {
|
||||
const totalBytes = getBytes(
|
||||
payload.total_bytes ?? payload.totalBytes,
|
||||
);
|
||||
const totalBytes = getBytes(payload.total);
|
||||
cell = { kind: "completed", totalBytes, modelDirectory };
|
||||
} else if (tag === "DownloadOngoing") {
|
||||
const rawProgress =
|
||||
payload.download_progress ?? payload.downloadProgress ?? {};
|
||||
const prog = rawProgress as Record<string, unknown>;
|
||||
const totalBytes = getBytes(
|
||||
prog.total_bytes ??
|
||||
prog.totalBytes ??
|
||||
payload.total_bytes ??
|
||||
payload.totalBytes,
|
||||
);
|
||||
const downloadedBytes = getBytes(
|
||||
prog.downloaded_bytes ?? prog.downloadedBytes,
|
||||
);
|
||||
const totalBytes = getBytes(prog.total ?? payload.total);
|
||||
const downloadedBytes = getBytes(prog.downloaded);
|
||||
const speed = (prog.speed as number) ?? 0;
|
||||
const etaMs =
|
||||
(prog.eta_ms as number) ?? (prog.etaMs as number) ?? 0;
|
||||
@@ -265,7 +260,20 @@
|
||||
} else if (tag === "DownloadFailed") {
|
||||
cell = { kind: "failed", modelDirectory };
|
||||
} else {
|
||||
cell = { kind: "pending", modelDirectory };
|
||||
const downloaded = getBytes(
|
||||
payload.downloaded ??
|
||||
payload.downloaded_bytes ??
|
||||
payload.downloadedBytes,
|
||||
);
|
||||
const total = getBytes(
|
||||
payload.total ?? payload.total_bytes ?? payload.totalBytes,
|
||||
);
|
||||
cell = {
|
||||
kind: "pending",
|
||||
downloaded,
|
||||
total,
|
||||
modelDirectory,
|
||||
};
|
||||
}
|
||||
|
||||
const existing = row.cells[nodeId];
|
||||
@@ -275,14 +283,51 @@
|
||||
}
|
||||
}
|
||||
|
||||
function rowSortKey(row: ModelRow): number {
|
||||
// in progress (4) -> completed (3) -> paused (2) -> not started (1) -> not present (0)
|
||||
let best = 0;
|
||||
for (const cell of Object.values(row.cells)) {
|
||||
let score = 0;
|
||||
if (cell.kind === "downloading") score = 4;
|
||||
else if (cell.kind === "completed") score = 3;
|
||||
else if (cell.kind === "pending" && cell.downloaded > 0)
|
||||
score = 2; // paused
|
||||
else if (cell.kind === "pending" || cell.kind === "failed") score = 1; // not started
|
||||
if (score > best) best = score;
|
||||
}
|
||||
return best;
|
||||
}
|
||||
|
||||
function totalCompletedBytes(row: ModelRow): number {
|
||||
let total = 0;
|
||||
for (const cell of Object.values(row.cells)) {
|
||||
if (cell.kind === "completed") total += cell.totalBytes;
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
const rows = Array.from(rowMap.values()).sort((a, b) => {
|
||||
const aCompleted = Object.values(a.cells).filter(
|
||||
(c) => c.kind === "completed",
|
||||
).length;
|
||||
const bCompleted = Object.values(b.cells).filter(
|
||||
(c) => c.kind === "completed",
|
||||
).length;
|
||||
if (aCompleted !== bCompleted) return bCompleted - aCompleted;
|
||||
const aPriority = rowSortKey(a);
|
||||
const bPriority = rowSortKey(b);
|
||||
if (aPriority !== bPriority) return bPriority - aPriority;
|
||||
// Within completed or paused, sort by biggest size first
|
||||
if (aPriority === 3 && bPriority === 3) {
|
||||
const sizeDiff = totalCompletedBytes(b) - totalCompletedBytes(a);
|
||||
if (sizeDiff !== 0) return sizeDiff;
|
||||
}
|
||||
if (aPriority === 2 && bPriority === 2) {
|
||||
const aSize = Math.max(
|
||||
...Object.values(a.cells).map((c) =>
|
||||
c.kind === "pending" ? c.total : 0,
|
||||
),
|
||||
);
|
||||
const bSize = Math.max(
|
||||
...Object.values(b.cells).map((c) =>
|
||||
c.kind === "pending" ? c.total : 0,
|
||||
),
|
||||
);
|
||||
if (aSize !== bSize) return bSize - aSize;
|
||||
}
|
||||
return a.modelId.localeCompare(b.modelId);
|
||||
});
|
||||
|
||||
@@ -367,7 +412,7 @@
|
||||
<div>{col.label}</div>
|
||||
{#if col.diskAvailable != null}
|
||||
<div
|
||||
class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
|
||||
class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
|
||||
>
|
||||
{formatBytes(col.diskAvailable)} free
|
||||
</div>
|
||||
@@ -391,7 +436,7 @@
|
||||
</div>
|
||||
{#if row.prettyName}
|
||||
<div
|
||||
class="text-[10px] text-exo-light-gray/60"
|
||||
class="text-[10px] text-white/60"
|
||||
title={row.modelId}
|
||||
>
|
||||
{row.modelId}
|
||||
@@ -405,7 +450,7 @@
|
||||
title="View model details"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4 text-white/30 hover:text-white/60"
|
||||
class="w-4 h-4 text-white/60 hover:text-white/80"
|
||||
viewBox="0 0 24 24"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -424,11 +469,11 @@
|
||||
<td class="px-4 py-3 text-center align-middle">
|
||||
{#if cell.kind === "completed"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title="Completed ({formatBytes(cell.totalBytes)})"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 text-green-400"
|
||||
class="w-7 h-7 text-green-400"
|
||||
viewBox="0 0 20 20"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -438,18 +483,18 @@
|
||||
clip-rule="evenodd"
|
||||
></path>
|
||||
</svg>
|
||||
<span class="text-[10px] text-exo-light-gray/70"
|
||||
<span class="text-xs text-white/70"
|
||||
>{formatBytes(cell.totalBytes)}</span
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
|
||||
class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
|
||||
onclick={() =>
|
||||
deleteDownload(col.nodeId, row.modelId)}
|
||||
title="Delete from this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
@@ -472,11 +517,11 @@
|
||||
cell.speed,
|
||||
)} - ETA {formatEta(cell.etaMs)}"
|
||||
>
|
||||
<span class="text-exo-yellow text-xs font-medium"
|
||||
<span class="text-exo-yellow text-sm font-medium"
|
||||
>{clampPercent(cell.percentage).toFixed(1)}%</span
|
||||
>
|
||||
<div
|
||||
class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
|
||||
class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
|
||||
>
|
||||
<div
|
||||
class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
|
||||
@@ -485,24 +530,93 @@
|
||||
).toFixed(1)}%"
|
||||
></div>
|
||||
</div>
|
||||
<span class="text-[9px] text-exo-light-gray/60"
|
||||
<span class="text-[10px] text-white/70"
|
||||
>{formatSpeed(cell.speed)}</span
|
||||
>
|
||||
</div>
|
||||
{:else if cell.kind === "pending"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
title="Download pending"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title={cell.downloaded > 0
|
||||
? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
|
||||
: "Download pending"}
|
||||
>
|
||||
<span class="text-exo-light-gray/50 text-sm">...</span>
|
||||
{#if cell.downloaded > 0 && cell.total > 0}
|
||||
<span class="text-white/70 text-xs"
|
||||
>{formatBytes(cell.downloaded)} / {formatBytes(
|
||||
cell.total,
|
||||
)}</span
|
||||
>
|
||||
<div
|
||||
class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
|
||||
>
|
||||
<div
|
||||
class="h-full bg-exo-light-gray/40 rounded-full"
|
||||
style="width: {(
|
||||
(cell.downloaded / cell.total) *
|
||||
100
|
||||
).toFixed(1)}%"
|
||||
></div>
|
||||
</div>
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Resume download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
></path>
|
||||
</svg>
|
||||
</button>
|
||||
{:else}
|
||||
<span class="text-white/50 text-[10px]">paused</span
|
||||
>
|
||||
{/if}
|
||||
{:else if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Start download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-6 h-6"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
></path>
|
||||
</svg>
|
||||
</button>
|
||||
{:else}
|
||||
<span class="text-white/40 text-sm">...</span>
|
||||
{/if}
|
||||
</div>
|
||||
{:else if cell.kind === "failed"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title="Download failed"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 text-red-400"
|
||||
class="w-7 h-7 text-red-400"
|
||||
viewBox="0 0 20 20"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -515,13 +629,13 @@
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Retry download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
@@ -547,13 +661,13 @@
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Download to this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import {
|
||||
listTraces,
|
||||
getTraceRawUrl,
|
||||
deleteTraces,
|
||||
type TraceListItem,
|
||||
} from "$lib/stores/app.svelte";
|
||||
import HeaderNav from "$lib/components/HeaderNav.svelte";
|
||||
@@ -10,6 +11,51 @@
|
||||
let traces = $state<TraceListItem[]>([]);
|
||||
let loading = $state(true);
|
||||
let error = $state<string | null>(null);
|
||||
let selectedIds = $state<Set<string>>(new Set());
|
||||
let deleting = $state(false);
|
||||
|
||||
let allSelected = $derived(
|
||||
traces.length > 0 && selectedIds.size === traces.length,
|
||||
);
|
||||
|
||||
function toggleSelect(taskId: string) {
|
||||
const next = new Set(selectedIds);
|
||||
if (next.has(taskId)) {
|
||||
next.delete(taskId);
|
||||
} else {
|
||||
next.add(taskId);
|
||||
}
|
||||
selectedIds = next;
|
||||
}
|
||||
|
||||
function toggleSelectAll() {
|
||||
if (allSelected) {
|
||||
selectedIds = new Set();
|
||||
} else {
|
||||
selectedIds = new Set(traces.map((t) => t.taskId));
|
||||
}
|
||||
}
|
||||
|
||||
async function handleDelete() {
|
||||
if (selectedIds.size === 0) return;
|
||||
const count = selectedIds.size;
|
||||
if (
|
||||
!confirm(
|
||||
`Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
|
||||
)
|
||||
)
|
||||
return;
|
||||
deleting = true;
|
||||
try {
|
||||
await deleteTraces([...selectedIds]);
|
||||
selectedIds = new Set();
|
||||
await refresh();
|
||||
} catch (e) {
|
||||
error = e instanceof Error ? e.message : "Failed to delete traces";
|
||||
} finally {
|
||||
deleting = false;
|
||||
}
|
||||
}
|
||||
|
||||
function formatBytes(bytes: number): string {
|
||||
if (!bytes || bytes <= 0) return "0B";
|
||||
@@ -109,6 +155,16 @@
|
||||
</h1>
|
||||
</div>
|
||||
<div class="flex items-center gap-3">
|
||||
{#if selectedIds.size > 0}
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
|
||||
onclick={handleDelete}
|
||||
disabled={deleting}
|
||||
>
|
||||
{deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
|
||||
</button>
|
||||
{/if}
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
|
||||
@@ -143,14 +199,41 @@
|
||||
</div>
|
||||
{:else}
|
||||
<div class="space-y-3">
|
||||
<div class="flex items-center gap-2 px-1">
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono uppercase transition-colors {allSelected
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-light-gray hover:text-exo-yellow'}"
|
||||
onclick={toggleSelectAll}
|
||||
>
|
||||
{allSelected ? "Deselect all" : "Select all"}
|
||||
</button>
|
||||
</div>
|
||||
{#each traces as trace}
|
||||
{@const isSelected = selectedIds.has(trace.taskId)}
|
||||
<!-- svelte-ignore a11y_no_static_element_interactions -->
|
||||
<div
|
||||
class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
|
||||
role="button"
|
||||
tabindex="0"
|
||||
class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
|
||||
? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
|
||||
: 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
|
||||
onclick={() => toggleSelect(trace.taskId)}
|
||||
onkeydown={(e) => {
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
toggleSelect(trace.taskId);
|
||||
}
|
||||
}}
|
||||
>
|
||||
<div class="min-w-0 flex-1">
|
||||
<a
|
||||
href="#/traces/{trace.taskId}"
|
||||
class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
|
||||
class="text-sm font-mono transition-colors truncate block {isSelected
|
||||
? 'text-exo-yellow'
|
||||
: 'text-white hover:text-exo-yellow'}"
|
||||
onclick={(e) => e.stopPropagation()}
|
||||
>
|
||||
{trace.taskId}
|
||||
</a>
|
||||
@@ -160,7 +243,11 @@
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-center gap-2 shrink-0">
|
||||
<!-- svelte-ignore a11y_click_events_have_key_events -->
|
||||
<div
|
||||
class="flex items-center gap-2 shrink-0"
|
||||
onclick={(e) => e.stopPropagation()}
|
||||
>
|
||||
<a
|
||||
href="#/traces/{trace.taskId}"
|
||||
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
|
||||
|
||||
+22
@@ -183,6 +183,27 @@ Same schema as `/v1/chat/completions`.
|
||||
**Response:**
|
||||
Chat completion plus benchmarking metrics.
|
||||
|
||||
### Cancel Command
|
||||
|
||||
**POST** `/v1/cancel/{command_id}`
|
||||
|
||||
Cancels an active generation command (text or image). Notifies workers and closes the stream.
|
||||
|
||||
**Path parameters:**
|
||||
|
||||
* `command_id`: string, ID of the command to cancel
|
||||
|
||||
**Response (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"message": "Command cancelled.",
|
||||
"command_id": "cmd-abc-123"
|
||||
}
|
||||
```
|
||||
|
||||
Returns 404 if the command is not found or already completed.
|
||||
|
||||
## 5. Image Generation & Editing
|
||||
|
||||
### Image Generation
|
||||
@@ -266,6 +287,7 @@ GET /v1/models
|
||||
|
||||
POST /v1/chat/completions
|
||||
POST /bench/chat/completions
|
||||
POST /v1/cancel/{command_id}
|
||||
|
||||
POST /v1/images/generations
|
||||
POST /bench/images/generations
|
||||
|
||||
@@ -108,6 +108,7 @@
|
||||
package = pkgsSwift.swiftPackages.swift-format;
|
||||
};
|
||||
shfmt.enable = true;
|
||||
taplo.enable = true;
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
+3
-3
@@ -41,7 +41,7 @@ let
|
||||
|
||||
mlx = stdenv.mkDerivation rec {
|
||||
pname = "mlx";
|
||||
version = let v = "0.30.7.dev20260218+14841977"; in
|
||||
version = let v = "0.30.7.dev20260225+257d5692"; in
|
||||
assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
|
||||
v;
|
||||
pyproject = true;
|
||||
@@ -49,8 +49,8 @@ let
|
||||
src = fetchFromGitHub {
|
||||
owner = "rltakashige";
|
||||
repo = "mlx-jaccl-fix-small-recv";
|
||||
rev = "1484197707f35186ad3bd614357c7c47fdf86ebc";
|
||||
hash = "sha256-FupCMoK/SF/ldfKuvMSAKECcOP8c+ANgkQlPZttDsLk=";
|
||||
rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
|
||||
hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
|
||||
};
|
||||
|
||||
patches = [
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 6.8 KiB |
Executable
+112
@@ -0,0 +1,112 @@
|
||||
#!/usr/bin/env bash
|
||||
# create-dmg.sh — Build a polished macOS DMG installer for EXO
|
||||
#
|
||||
# Usage:
|
||||
# ./packaging/dmg/create-dmg.sh <app-path> <output-dmg> [volume-name]
|
||||
#
|
||||
# Example:
|
||||
# ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO"
|
||||
#
|
||||
# Creates a DMG with:
|
||||
# - Custom background image with drag-to-Applications arrow
|
||||
# - App icon on left, Applications alias on right
|
||||
# - Proper window size and icon positioning
|
||||
set -euo pipefail
|
||||
|
||||
APP_PATH="${1:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
|
||||
OUTPUT_DMG="${2:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
|
||||
VOLUME_NAME="${3:-EXO}"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py"
|
||||
TEMP_DIR="$(mktemp -d)"
|
||||
DMG_STAGING="${TEMP_DIR}/dmg-root"
|
||||
TEMP_DMG="${TEMP_DIR}/temp.dmg"
|
||||
BACKGROUND_PNG="${TEMP_DIR}/background.png"
|
||||
|
||||
cleanup() { rm -rf "$TEMP_DIR"; }
|
||||
trap cleanup EXIT
|
||||
|
||||
echo "==> Creating DMG installer for ${VOLUME_NAME}"
|
||||
|
||||
# ── Step 1: Generate background image ────────────────────────────────────────
|
||||
if command -v python3 &>/dev/null; then
|
||||
python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG"
|
||||
echo " Background image generated"
|
||||
else
|
||||
echo " Warning: python3 not found, skipping custom background"
|
||||
BACKGROUND_PNG=""
|
||||
fi
|
||||
|
||||
# ── Step 2: Prepare staging directory ─────────────────────────────────────────
|
||||
mkdir -p "$DMG_STAGING"
|
||||
cp -R "$APP_PATH" "$DMG_STAGING/"
|
||||
ln -s /Applications "$DMG_STAGING/Applications"
|
||||
|
||||
# ── Step 3: Create writable DMG ──────────────────────────────────────────────
|
||||
# Calculate required size (app size + 20MB headroom)
|
||||
APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1)
|
||||
DMG_SIZE_KB=$((APP_SIZE_KB + 20480))
|
||||
|
||||
hdiutil create \
|
||||
-volname "$VOLUME_NAME" \
|
||||
-size "${DMG_SIZE_KB}k" \
|
||||
-fs HFS+ \
|
||||
-layout SPUD \
|
||||
"$TEMP_DMG"
|
||||
|
||||
# ── Step 4: Mount and configure ──────────────────────────────────────────────
|
||||
MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}')
|
||||
echo " Mounted at: $MOUNT_DIR"
|
||||
|
||||
# Copy contents
|
||||
cp -R "$DMG_STAGING/"* "$MOUNT_DIR/"
|
||||
|
||||
# Add background image
|
||||
if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then
|
||||
mkdir -p "$MOUNT_DIR/.background"
|
||||
cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png"
|
||||
fi
|
||||
|
||||
# ── Step 5: Configure window appearance via AppleScript ──────────────────────
|
||||
# Window: 800×400, app icon on left, Applications on right (matches Ollama layout)
|
||||
# Background image is 1600×740 (2× retina for 800×400 logical window).
|
||||
APP_NAME="$(basename "$APP_PATH")"
|
||||
|
||||
osascript <<APPLESCRIPT
|
||||
tell application "Finder"
|
||||
tell disk "$VOLUME_NAME"
|
||||
open
|
||||
set current view of container window to icon view
|
||||
set toolbar visible of container window to false
|
||||
set statusbar visible of container window to false
|
||||
set bounds of container window to {200, 120, 1000, 520}
|
||||
set opts to icon view options of container window
|
||||
set icon size of opts to 128
|
||||
set text size of opts to 12
|
||||
set arrangement of opts to not arranged
|
||||
if exists file ".background:background.png" then
|
||||
set background picture of opts to file ".background:background.png"
|
||||
end if
|
||||
set position of item "$APP_NAME" of container window to {200, 190}
|
||||
set position of item "Applications" of container window to {600, 190}
|
||||
close
|
||||
open
|
||||
update without registering applications
|
||||
delay 1
|
||||
close
|
||||
end tell
|
||||
end tell
|
||||
APPLESCRIPT
|
||||
|
||||
echo " Window layout configured"
|
||||
|
||||
# Ensure Finder updates are flushed
|
||||
sync
|
||||
|
||||
# ── Step 6: Finalise ─────────────────────────────────────────────────────────
|
||||
hdiutil detach "$MOUNT_DIR" -quiet
|
||||
hdiutil convert "$TEMP_DMG" -format UDZO -imagekey zlib-level=9 -o "$OUTPUT_DMG"
|
||||
|
||||
echo "==> DMG created: $OUTPUT_DMG"
|
||||
echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)"
|
||||
@@ -0,0 +1,91 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate the DMG background image with a centered drag-to-Applications arrow.
|
||||
|
||||
The output is a 1600×740 retina PNG (2× for 800×400 logical window).
|
||||
Icons are positioned at (200, 190) and (600, 190) in logical coordinates;
|
||||
the arrow is drawn centered between them.
|
||||
|
||||
Usage:
|
||||
python3 generate-background.py [output.png]
|
||||
|
||||
If no output path is given, overwrites the bundled background.png in-place.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
# Retina dimensions (2× logical 800×400)
|
||||
WIDTH = 1600
|
||||
HEIGHT = 740
|
||||
|
||||
# Icon positions in logical coords → retina coords
|
||||
# App icon at (200, 190), Applications at (600, 190)
|
||||
APP_X = 200 * 2 # 400
|
||||
APPS_X = 600 * 2 # 1200
|
||||
ICON_Y = 190 * 2 # 380
|
||||
|
||||
# Arrow drawn between icons, slightly above icon center
|
||||
ARROW_START_X = APP_X + 160 # past the icon
|
||||
ARROW_END_X = APPS_X - 160 # before the Applications icon
|
||||
ARROW_Y = ICON_Y # same height as icons
|
||||
ARROW_RISE = 120 # upward arc height
|
||||
|
||||
|
||||
def draw_arrow(draw: ImageDraw.ImageDraw) -> None:
|
||||
"""Draw a hand-drawn-style curved arrow from app icon toward Applications."""
|
||||
color = (30, 30, 30)
|
||||
line_width = 8
|
||||
|
||||
# Compute bezier curve points for a gentle upward arc
|
||||
points: list[tuple[float, float]] = []
|
||||
steps = 80
|
||||
for i in range(steps + 1):
|
||||
t = i / steps
|
||||
# Quadratic bezier: start → control → end
|
||||
cx = (ARROW_START_X + ARROW_END_X) / 2
|
||||
cy = ARROW_Y - ARROW_RISE
|
||||
x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X
|
||||
y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y
|
||||
points.append((x, y))
|
||||
|
||||
# Draw the curve as connected line segments
|
||||
for i in range(len(points) - 1):
|
||||
draw.line([points[i], points[i + 1]], fill=color, width=line_width)
|
||||
|
||||
# Arrowhead at the end
|
||||
end_x, end_y = points[-1]
|
||||
# Direction from second-to-last to last point
|
||||
prev_x, prev_y = points[-3]
|
||||
angle = math.atan2(end_y - prev_y, end_x - prev_x)
|
||||
head_len = 36
|
||||
head_angle = math.radians(25)
|
||||
|
||||
left_x = end_x - head_len * math.cos(angle - head_angle)
|
||||
left_y = end_y - head_len * math.sin(angle - head_angle)
|
||||
right_x = end_x - head_len * math.cos(angle + head_angle)
|
||||
right_y = end_y - head_len * math.sin(angle + head_angle)
|
||||
|
||||
draw.polygon(
|
||||
[(end_x, end_y), (left_x, left_y), (right_x, right_y)],
|
||||
fill=color,
|
||||
)
|
||||
|
||||
|
||||
def generate_background(output_path: str) -> None:
|
||||
"""Generate a white DMG background with a centered arrow."""
|
||||
img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255))
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw_arrow(draw)
|
||||
img.save(output_path, "PNG")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
default_output = str(Path(__file__).parent / "background.png")
|
||||
out = sys.argv[1] if len(sys.argv) >= 2 else default_output
|
||||
generate_background(out)
|
||||
print(f"Background image written to {out}")
|
||||
+52
-52
@@ -1,35 +1,35 @@
|
||||
[project]
|
||||
name = "exo"
|
||||
version = "0.3.0"
|
||||
version = "0.3.68"
|
||||
description = "Exo"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"aiofiles>=24.1.0",
|
||||
"aiohttp>=3.12.14",
|
||||
"types-aiofiles>=24.1.0.20250708",
|
||||
"pydantic>=2.11.7",
|
||||
"fastapi>=0.116.1",
|
||||
"filelock>=3.18.0",
|
||||
"rustworkx>=0.17.1",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"psutil>=7.0.0",
|
||||
"loguru>=0.7.3",
|
||||
"exo_pyo3_bindings", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"mlx; sys_platform == 'darwin'",
|
||||
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
|
||||
"mlx-lm==0.30.7",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"hypercorn>=0.18.0",
|
||||
"openai-harmony>=0.0.8",
|
||||
"httpx>=0.28.1",
|
||||
"tomlkit>=0.14.0",
|
||||
"pillow>=11.0,<12.0", # compatibility with mflux
|
||||
"mflux==0.15.5",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
"aiofiles>=24.1.0",
|
||||
"aiohttp>=3.12.14",
|
||||
"types-aiofiles>=24.1.0.20250708",
|
||||
"pydantic>=2.11.7",
|
||||
"fastapi>=0.116.1",
|
||||
"filelock>=3.18.0",
|
||||
"rustworkx>=0.17.1",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"psutil>=7.0.0",
|
||||
"loguru>=0.7.3",
|
||||
"exo_pyo3_bindings", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"mlx; sys_platform == 'darwin'",
|
||||
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
|
||||
"mlx-lm",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"hypercorn>=0.18.0",
|
||||
"openai-harmony>=0.0.8",
|
||||
"httpx>=0.28.1",
|
||||
"tomlkit>=0.14.0",
|
||||
"pillow>=11.0,<12.0", # compatibility with mflux
|
||||
"mflux==0.15.5",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -38,12 +38,12 @@ exo = "exo.main:main"
|
||||
# dependencies only required for development
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"basedpyright>=1.29.0",
|
||||
"pyinstaller>=6.17.0",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
"pytest-env",
|
||||
"ruff>=0.11.13",
|
||||
"basedpyright>=1.29.0",
|
||||
"pyinstaller>=6.17.0",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
"pytest-env",
|
||||
"ruff>=0.11.13",
|
||||
]
|
||||
|
||||
# mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
|
||||
@@ -57,15 +57,12 @@ dev = [
|
||||
###
|
||||
|
||||
[tool.uv.workspace]
|
||||
members = [
|
||||
"rust/exo_pyo3_bindings",
|
||||
"bench",
|
||||
]
|
||||
members = ["rust/exo_pyo3_bindings", "bench"]
|
||||
|
||||
[tool.uv.sources]
|
||||
exo_pyo3_bindings = { workspace = true }
|
||||
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
|
||||
#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
|
||||
mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
|
||||
# Uncomment to use local mlx/mlx-lm development versions:
|
||||
# mlx = { path = "/Users/Shared/mlx", editable=true }
|
||||
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
|
||||
@@ -95,7 +92,15 @@ reportUnnecessaryTypeIgnoreComment = "error"
|
||||
pythonVersion = "3.13"
|
||||
pythonPlatform = "Darwin"
|
||||
|
||||
exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
|
||||
exclude = [
|
||||
"**/.venv",
|
||||
"**/venv",
|
||||
"**/__pycache__",
|
||||
"**/exo_scripts",
|
||||
"**/.direnv",
|
||||
"**/rust",
|
||||
"**/.github",
|
||||
]
|
||||
stubPath = ".mlx_typings"
|
||||
|
||||
[[tool.basedpyright.executionEnvironments]]
|
||||
@@ -109,17 +114,18 @@ root = "src"
|
||||
[tool.uv]
|
||||
required-version = ">=0.8.6"
|
||||
prerelease = "allow"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
]
|
||||
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
|
||||
|
||||
###
|
||||
# ruff configuration
|
||||
###
|
||||
|
||||
[tool.ruff]
|
||||
extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
|
||||
extend-exclude = [
|
||||
"shared/protobufs/**",
|
||||
"*mlx_typings/**",
|
||||
"rust/exo_pyo3_bindings/**",
|
||||
]
|
||||
|
||||
[tool.ruff.lint]
|
||||
extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
|
||||
@@ -127,13 +133,7 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = "."
|
||||
asyncio_mode = "auto"
|
||||
markers = [
|
||||
"slow: marks tests as slow (deselected by default)"
|
||||
]
|
||||
env = [
|
||||
"EXO_TESTS=1"
|
||||
]
|
||||
markers = ["slow: marks tests as slow (deselected by default)"]
|
||||
env = ["EXO_TESTS=1"]
|
||||
addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
|
||||
filterwarnings = [
|
||||
"ignore:builtin type Swig:DeprecationWarning",
|
||||
]
|
||||
filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]
|
||||
|
||||
@@ -3,6 +3,10 @@ n_layers = 48
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3 Coder Next"
|
||||
capabilities = ["text", "code"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 45644286500
|
||||
|
||||
@@ -3,6 +3,10 @@ n_layers = 48
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "5bit"
|
||||
base_model = "Qwen3 Coder Next"
|
||||
capabilities = ["text", "code"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 57657697020
|
||||
|
||||
@@ -3,6 +3,10 @@ n_layers = 48
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "6bit"
|
||||
base_model = "Qwen3 Coder Next"
|
||||
capabilities = ["text", "code"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 68899327465
|
||||
|
||||
@@ -3,6 +3,10 @@ n_layers = 48
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3 Coder Next"
|
||||
capabilities = ["text", "code"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 89357758772
|
||||
|
||||
@@ -3,6 +3,10 @@ n_layers = 48
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "bf16"
|
||||
base_model = "Qwen3 Coder Next"
|
||||
capabilities = ["text", "code"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 157548627945
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 69593314272
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "6bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 100120675296
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 130648036320
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "bf16"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 245125640160
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-27B-4bit"
|
||||
n_layers = 64
|
||||
hidden_size = 5120
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 27B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 16054266848
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-27B-8bit"
|
||||
n_layers = 64
|
||||
hidden_size = 5120
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 27B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 29500943328
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
|
||||
n_layers = 24
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 2B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 2662787264
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
|
||||
n_layers = 40
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 35B A3B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 20391405152
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
|
||||
n_layers = 40
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 35B A3B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 37721130592
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 223860768352
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "6bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 322946674272
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 422032580192
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-9B-4bit"
|
||||
n_layers = 32
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 9B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 5950062560
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-9B-8bit"
|
||||
n_layers = 32
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 9B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 10426433504
|
||||
@@ -5,7 +5,6 @@ edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
path = "src/lib.rs"
|
||||
name = "exo_pyo3_bindings"
|
||||
|
||||
@@ -26,37 +25,33 @@ networking = { workspace = true }
|
||||
|
||||
# interop
|
||||
pyo3 = { version = "0.27.2", features = [
|
||||
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
|
||||
# "nightly", # enables better-supported GIL integration
|
||||
"experimental-async", # async support in #[pyfunction] & #[pymethods]
|
||||
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
|
||||
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
|
||||
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
|
||||
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
|
||||
# "nightly", # enables better-supported GIL integration
|
||||
"experimental-async", # async support in #[pyfunction] & #[pymethods]
|
||||
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
|
||||
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
|
||||
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
|
||||
|
||||
# integrations with other libraries
|
||||
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
|
||||
# "ordered-float", "rust_decimal", "smallvec",
|
||||
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
|
||||
# integrations with other libraries
|
||||
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
|
||||
# "ordered-float", "rust_decimal", "smallvec",
|
||||
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
|
||||
] }
|
||||
pyo3-stub-gen = { version = "0.17.2" }
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = [
|
||||
"attributes",
|
||||
"tokio-runtime",
|
||||
"testing",
|
||||
] }
|
||||
pyo3-log = "0.13.2"
|
||||
|
||||
# macro dependencies
|
||||
extend = { workspace = true }
|
||||
delegate = { workspace = true }
|
||||
|
||||
# async runtime
|
||||
tokio = { workspace = true, features = ["full", "tracing"] }
|
||||
futures-lite = { workspace = true }
|
||||
|
||||
# utility dependencies
|
||||
util = { workspace = true }
|
||||
|
||||
# Tracing
|
||||
log = { workspace = true }
|
||||
env_logger = "0.11"
|
||||
|
||||
# Networking
|
||||
libp2p = { workspace = true, features = ["full"] }
|
||||
pin-project = "1.1.10"
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
TODO: do something here....
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import builtins
|
||||
import enum
|
||||
import typing
|
||||
|
||||
@typing.final
|
||||
@@ -11,138 +10,39 @@ class AllQueuesFullError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class ConnectionUpdate:
|
||||
@property
|
||||
def update_type(self) -> ConnectionUpdateType:
|
||||
r"""
|
||||
Whether this is a connection or disconnection event
|
||||
"""
|
||||
@property
|
||||
def peer_id(self) -> PeerId:
|
||||
r"""
|
||||
Identity of the peer that we have connected to or disconnected from.
|
||||
"""
|
||||
@property
|
||||
def remote_ipv4(self) -> builtins.str:
|
||||
r"""
|
||||
Remote connection's IPv4 address.
|
||||
"""
|
||||
@property
|
||||
def remote_tcp_port(self) -> builtins.int:
|
||||
r"""
|
||||
Remote connection's TCP port.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class Keypair:
|
||||
r"""
|
||||
Identity keypair of a node.
|
||||
"""
|
||||
@staticmethod
|
||||
def generate_ed25519() -> Keypair:
|
||||
def generate() -> Keypair:
|
||||
r"""
|
||||
Generate a new Ed25519 keypair.
|
||||
"""
|
||||
@staticmethod
|
||||
def generate_ecdsa() -> Keypair:
|
||||
def from_bytes(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Generate a new ECDSA keypair.
|
||||
"""
|
||||
@staticmethod
|
||||
def generate_secp256k1() -> Keypair:
|
||||
r"""
|
||||
Generate a new Secp256k1 keypair.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_protobuf_encoding(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Decode a private key from a protobuf structure and parse it as a `Keypair`.
|
||||
"""
|
||||
@staticmethod
|
||||
def rsa_from_pkcs8(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
|
||||
format (i.e. unencrypted) as defined in [RFC5208].
|
||||
|
||||
[RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
|
||||
"""
|
||||
@staticmethod
|
||||
def secp256k1_from_der(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
|
||||
structure as defined in [RFC5915].
|
||||
|
||||
[RFC5915]: https://tools.ietf.org/html/rfc5915
|
||||
"""
|
||||
@staticmethod
|
||||
def ed25519_from_bytes(bytes: bytes) -> Keypair: ...
|
||||
def to_protobuf_encoding(self) -> bytes:
|
||||
r"""
|
||||
Encode a private key as protobuf structure.
|
||||
"""
|
||||
def to_peer_id(self) -> PeerId:
|
||||
r"""
|
||||
Convert the `Keypair` into the corresponding `PeerId`.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class Multiaddr:
|
||||
r"""
|
||||
Representation of a Multiaddr.
|
||||
"""
|
||||
@staticmethod
|
||||
def empty() -> Multiaddr:
|
||||
r"""
|
||||
Create a new, empty multiaddress.
|
||||
"""
|
||||
@staticmethod
|
||||
def with_capacity(n: builtins.int) -> Multiaddr:
|
||||
r"""
|
||||
Create a new, empty multiaddress with the given capacity.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_bytes(bytes: bytes) -> Multiaddr:
|
||||
r"""
|
||||
Parse a `Multiaddr` value from its byte slice representation.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_string(string: builtins.str) -> Multiaddr:
|
||||
r"""
|
||||
Parse a `Multiaddr` value from its string representation.
|
||||
"""
|
||||
def len(self) -> builtins.int:
|
||||
r"""
|
||||
Return the length in bytes of this multiaddress.
|
||||
"""
|
||||
def is_empty(self) -> builtins.bool:
|
||||
r"""
|
||||
Returns true if the length of this multiaddress is 0.
|
||||
Construct an Ed25519 keypair from secret key bytes
|
||||
"""
|
||||
def to_bytes(self) -> bytes:
|
||||
r"""
|
||||
Return a copy of this [`Multiaddr`]'s byte representation.
|
||||
Get the secret key bytes underlying the keypair
|
||||
"""
|
||||
def to_string(self) -> builtins.str:
|
||||
def to_node_id(self) -> builtins.str:
|
||||
r"""
|
||||
Convert a Multiaddr to a string.
|
||||
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class MessageTooLargeError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
|
||||
async def connection_update_recv(self) -> ConnectionUpdate:
|
||||
r"""
|
||||
Receives the next `ConnectionUpdate` from networking.
|
||||
"""
|
||||
async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
|
||||
r"""
|
||||
Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
|
||||
|
||||
For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
|
||||
For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
|
||||
will sleep until a `ConnectionUpdate`s is sent.
|
||||
"""
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
@@ -161,18 +61,7 @@ class NetworkingHandle:
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
|
||||
r"""
|
||||
Receives the next message from the `GossipSub` network.
|
||||
"""
|
||||
async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
|
||||
r"""
|
||||
Receives at most `limit` messages from the `GossipSub` network and returns them.
|
||||
|
||||
For `limit = 0`, an empty collection of messages will be returned immediately.
|
||||
For `limit > 0`, if there are no messages in the channel's queue this method
|
||||
will sleep until a message is sent.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
@typing.final
|
||||
class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
@@ -180,42 +69,26 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class PeerId:
|
||||
r"""
|
||||
Identifier of a peer of the network.
|
||||
class PyFromSwarm:
|
||||
@typing.final
|
||||
class Connection(PyFromSwarm):
|
||||
__match_args__ = ("peer_id", "connected",)
|
||||
@property
|
||||
def peer_id(self) -> builtins.str: ...
|
||||
@property
|
||||
def connected(self) -> builtins.bool: ...
|
||||
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
|
||||
|
||||
The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
|
||||
as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
|
||||
"""
|
||||
@staticmethod
|
||||
def random() -> PeerId:
|
||||
r"""
|
||||
Generates a random peer ID from a cryptographically secure PRNG.
|
||||
|
||||
This is useful for randomly walking on a DHT, or for testing purposes.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_bytes(bytes: bytes) -> PeerId:
|
||||
r"""
|
||||
Parses a `PeerId` from bytes.
|
||||
"""
|
||||
def to_bytes(self) -> bytes:
|
||||
r"""
|
||||
Returns a raw bytes representation of this `PeerId`.
|
||||
"""
|
||||
def to_base58(self) -> builtins.str:
|
||||
r"""
|
||||
Returns a base-58 encoded string of this `PeerId`.
|
||||
"""
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class ConnectionUpdateType(enum.Enum):
|
||||
r"""
|
||||
Connection or disconnection event discriminant type.
|
||||
"""
|
||||
Connected = ...
|
||||
Disconnected = ...
|
||||
@typing.final
|
||||
class Message(PyFromSwarm):
|
||||
__match_args__ = ("origin", "topic", "data",)
|
||||
@property
|
||||
def origin(self) -> builtins.str: ...
|
||||
@property
|
||||
def topic(self) -> builtins.str: ...
|
||||
@property
|
||||
def data(self) -> bytes: ...
|
||||
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
|
||||
|
||||
...
|
||||
|
||||
|
||||
@@ -4,21 +4,18 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "exo_pyo3_bindings"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
authors = [
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
|
||||
{ name = "Evan Quiney", email = "evanev7@gmail.com" },
|
||||
]
|
||||
requires-python = ">=3.13"
|
||||
dependencies = []
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"exo_pyo3_bindings",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
]
|
||||
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
|
||||
|
||||
[tool.maturin]
|
||||
#purelib = true
|
||||
|
||||
@@ -1,37 +1,22 @@
|
||||
//! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
//!
|
||||
|
||||
use pin_project::pin_project;
|
||||
//! See: <https://pyo3.rs/v0.27.2/async-await.html#detaching-from-the-interpreter-across-await>
|
||||
use pyo3::prelude::*;
|
||||
use std::{
|
||||
future::Future,
|
||||
pin::Pin,
|
||||
pin::{Pin, pin},
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
#[pin_project]
|
||||
#[repr(transparent)]
|
||||
pub(crate) struct AllowThreads<F>(#[pin] F);
|
||||
|
||||
impl<F> AllowThreads<F>
|
||||
where
|
||||
Self: Future,
|
||||
{
|
||||
pub fn new(f: F) -> Self {
|
||||
Self(f)
|
||||
}
|
||||
}
|
||||
pub struct AllowThreads<F>(pub(crate) F);
|
||||
|
||||
impl<F> Future for AllowThreads<F>
|
||||
where
|
||||
F: Future + Send,
|
||||
F: Future + Unpin + Send,
|
||||
F::Output: Send,
|
||||
{
|
||||
type Output = F::Output;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
let waker = cx.waker();
|
||||
Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker))))
|
||||
Python::attach(|py| py.detach(|| pin!(&mut self.0).poll(&mut Context::from_waker(waker))))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
use crate::ext::ResultExt as _;
|
||||
use libp2p::PeerId;
|
||||
use libp2p::identity::Keypair;
|
||||
use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::types::{PyBytes, PyBytesMethods as _};
|
||||
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
|
||||
@@ -18,142 +16,32 @@ pub struct PyKeypair(pub Keypair);
|
||||
impl PyKeypair {
|
||||
/// Generate a new Ed25519 keypair.
|
||||
#[staticmethod]
|
||||
fn generate_ed25519() -> Self {
|
||||
fn generate() -> Self {
|
||||
Self(Keypair::generate_ed25519())
|
||||
}
|
||||
|
||||
/// Generate a new ECDSA keypair.
|
||||
/// Construct an Ed25519 keypair from secret key bytes
|
||||
#[staticmethod]
|
||||
fn generate_ecdsa() -> Self {
|
||||
Self(Keypair::generate_ecdsa())
|
||||
}
|
||||
|
||||
/// Generate a new Secp256k1 keypair.
|
||||
#[staticmethod]
|
||||
fn generate_secp256k1() -> Self {
|
||||
Self(Keypair::generate_secp256k1())
|
||||
}
|
||||
|
||||
/// Decode a private key from a protobuf structure and parse it as a `Keypair`.
|
||||
#[staticmethod]
|
||||
fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?))
|
||||
}
|
||||
|
||||
/// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
|
||||
/// format (i.e. unencrypted) as defined in [RFC5208].
|
||||
///
|
||||
/// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
|
||||
#[staticmethod]
|
||||
fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let mut bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?))
|
||||
}
|
||||
|
||||
/// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
|
||||
/// structure as defined in [RFC5915].
|
||||
///
|
||||
/// [RFC5915]: https://tools.ietf.org/html/rfc5915
|
||||
#[staticmethod]
|
||||
fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let mut bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let mut bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
|
||||
}
|
||||
|
||||
/// Encode a private key as protobuf structure.
|
||||
fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
|
||||
let bytes = self.0.to_protobuf_encoding().pyerr()?;
|
||||
/// Get the secret key bytes underlying the keypair
|
||||
fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
|
||||
let bytes = self
|
||||
.0
|
||||
.clone()
|
||||
.try_into_ed25519()
|
||||
.pyerr()?
|
||||
.secret()
|
||||
.as_ref()
|
||||
.to_vec();
|
||||
Ok(PyBytes::new(py, &bytes))
|
||||
}
|
||||
|
||||
/// Convert the `Keypair` into the corresponding `PeerId`.
|
||||
fn to_peer_id(&self) -> PyPeerId {
|
||||
PyPeerId(self.0.public().to_peer_id())
|
||||
}
|
||||
|
||||
// /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
|
||||
// #[gen_stub(skip)]
|
||||
// #[new]
|
||||
// fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
// Self::from_protobuf_encoding(bytes)
|
||||
// }
|
||||
//
|
||||
// #[gen_stub(skip)]
|
||||
// fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
|
||||
// *self = Self::from_protobuf_encoding(state)?;
|
||||
// Ok(())
|
||||
// }
|
||||
//
|
||||
// #[gen_stub(skip)]
|
||||
// fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
|
||||
// self.to_protobuf_encoding(py)
|
||||
// }
|
||||
//
|
||||
// #[gen_stub(skip)]
|
||||
// pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
|
||||
// Ok((self.to_protobuf_encoding(py)?,))
|
||||
// }
|
||||
}
|
||||
|
||||
/// Identifier of a peer of the network.
|
||||
///
|
||||
/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
|
||||
/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "PeerId", frozen)]
|
||||
#[derive(Debug, Clone)]
|
||||
#[repr(transparent)]
|
||||
pub struct PyPeerId(pub PeerId);
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
#[allow(clippy::needless_pass_by_value)]
|
||||
impl PyPeerId {
|
||||
/// Generates a random peer ID from a cryptographically secure PRNG.
|
||||
///
|
||||
/// This is useful for randomly walking on a DHT, or for testing purposes.
|
||||
#[staticmethod]
|
||||
fn random() -> Self {
|
||||
Self(PeerId::random())
|
||||
}
|
||||
|
||||
/// Parses a `PeerId` from bytes.
|
||||
#[staticmethod]
|
||||
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
|
||||
let bytes = Vec::from(bytes.as_bytes());
|
||||
Ok(Self(PeerId::from_bytes(&bytes).pyerr()?))
|
||||
}
|
||||
|
||||
/// Returns a raw bytes representation of this `PeerId`.
|
||||
fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
|
||||
let bytes = self.0.to_bytes();
|
||||
PyBytes::new(py, &bytes)
|
||||
}
|
||||
|
||||
/// Returns a base-58 encoded string of this `PeerId`.
|
||||
fn to_base58(&self) -> String {
|
||||
self.0.to_base58()
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("PeerId({})", self.to_base58())
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
self.to_base58()
|
||||
/// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
|
||||
fn to_node_id(&self) -> String {
|
||||
self.0.public().to_peer_id().to_base58()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyKeypair>()?;
|
||||
m.add_class::<PyPeerId>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1,170 +1,5 @@
|
||||
//! TODO: crate documentation
|
||||
//!
|
||||
//! this is here as a placeholder documentation
|
||||
//!
|
||||
//!
|
||||
|
||||
mod allow_threading;
|
||||
mod ident;
|
||||
mod networking;
|
||||
|
||||
use crate::ident::ident_submodule;
|
||||
use crate::networking::networking_submodule;
|
||||
use pyo3::prelude::PyModule;
|
||||
use pyo3::{Bound, PyResult, pyclass, pymodule};
|
||||
use pyo3_stub_gen::define_stub_info_gatherer;
|
||||
|
||||
/// Namespace for all the constants used by this crate.
|
||||
pub(crate) mod r#const {
|
||||
pub const MPSC_CHANNEL_SIZE: usize = 1024;
|
||||
}
|
||||
|
||||
/// Namespace for crate-wide extension traits/methods
|
||||
pub(crate) mod ext {
|
||||
use crate::allow_threading::AllowThreads;
|
||||
use extend::ext;
|
||||
use pyo3::exceptions::{PyConnectionError, PyRuntimeError};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Py, PyErr, PyResult, Python};
|
||||
use tokio::runtime::Runtime;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::mpsc::error::TryRecvError;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
#[ext(pub, name = ByteArrayExt)]
|
||||
impl [u8] {
|
||||
fn pybytes(&self) -> Py<PyBytes> {
|
||||
Python::attach(|py| PyBytes::new(py, self).unbind())
|
||||
}
|
||||
}
|
||||
|
||||
#[ext(pub, name = ResultExt)]
|
||||
impl<T, E> Result<T, E>
|
||||
where
|
||||
E: ToString,
|
||||
{
|
||||
fn pyerr(self) -> PyResult<T> {
|
||||
self.map_err(|e| PyRuntimeError::new_err(e.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait FutureExt: Future + Sized {
|
||||
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
fn allow_threads_py(self) -> AllowThreads<Self>
|
||||
where
|
||||
AllowThreads<Self>: Future,
|
||||
{
|
||||
AllowThreads::new(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Future> FutureExt for T {}
|
||||
|
||||
#[ext(pub, name = PyErrExt)]
|
||||
impl PyErr {
|
||||
fn receiver_channel_closed() -> Self {
|
||||
PyConnectionError::new_err("Receiver channel closed unexpectedly")
|
||||
}
|
||||
}
|
||||
|
||||
#[ext(pub, name = PyResultExt)]
|
||||
impl<T> PyResult<T> {
|
||||
fn write_unraisable(self) -> Option<T> {
|
||||
Python::attach(|py| self.write_unraisable_with(py))
|
||||
}
|
||||
|
||||
fn write_unraisable_with(self, py: Python<'_>) -> Option<T> {
|
||||
match self {
|
||||
Ok(v) => Some(v),
|
||||
Err(e) => {
|
||||
// write error back to python
|
||||
e.write_unraisable(py, None);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[ext(pub, name = TokioRuntimeExt)]
|
||||
impl Runtime {
|
||||
fn spawn_with_scope<F>(&self, py: Python<'_>, future: F) -> PyResult<JoinHandle<F::Output>>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?;
|
||||
Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future)))
|
||||
}
|
||||
}
|
||||
|
||||
#[ext(pub, name = TokioMpscSenderExt)]
|
||||
impl<T> mpsc::Sender<T> {
|
||||
/// Sends a value, waiting until there is capacity.
|
||||
///
|
||||
/// A successful send occurs when it is determined that the other end of the
|
||||
/// channel has not hung up already. An unsuccessful send would be one where
|
||||
/// the corresponding receiver has already been closed.
|
||||
async fn send_py(&self, value: T) -> PyResult<()> {
|
||||
self.send(value)
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())
|
||||
}
|
||||
}
|
||||
|
||||
#[ext(pub, name = TokioMpscReceiverExt)]
|
||||
impl<T> mpsc::Receiver<T> {
|
||||
/// Receives the next value for this receiver.
|
||||
async fn recv_py(&mut self) -> PyResult<T> {
|
||||
self.recv().await.ok_or_else(PyErr::receiver_channel_closed)
|
||||
}
|
||||
|
||||
/// Receives at most `limit` values for this receiver and returns them.
|
||||
///
|
||||
/// For `limit = 0`, an empty collection of messages will be returned immediately.
|
||||
/// For `limit > 0`, if there are no messages in the channel's queue this method
|
||||
/// will sleep until a message is sent.
|
||||
async fn recv_many_py(&mut self, limit: usize) -> PyResult<Vec<T>> {
|
||||
// get updates from receiver channel
|
||||
let mut updates = Vec::with_capacity(limit);
|
||||
let received = self.recv_many(&mut updates, limit).await;
|
||||
|
||||
// if we received zero items, then the channel was unexpectedly closed
|
||||
if limit != 0 && received == 0 {
|
||||
return Err(PyErr::receiver_channel_closed());
|
||||
}
|
||||
|
||||
Ok(updates)
|
||||
}
|
||||
|
||||
/// Tries to receive the next value for this receiver.
|
||||
fn try_recv_py(&mut self) -> PyResult<Option<T>> {
|
||||
match self.try_recv() {
|
||||
Ok(v) => Ok(Some(v)),
|
||||
Err(TryRecvError::Empty) => Ok(None),
|
||||
Err(TryRecvError::Disconnected) => Err(PyErr::receiver_channel_closed()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A Python module implemented in Rust. The name of this function must match
|
||||
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
|
||||
/// import the module.
|
||||
#[pymodule(name = "exo_pyo3_bindings")]
|
||||
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
// install logger
|
||||
pyo3_log::init();
|
||||
|
||||
// TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
|
||||
// work with maturin, where the types generate correctly, in the right folder, without
|
||||
// too many importing issues...
|
||||
ident_submodule(m)?;
|
||||
networking_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
// TODO: ...
|
||||
|
||||
Ok(())
|
||||
}
|
||||
mod allow_threading;
|
||||
|
||||
define_stub_info_gatherer!(stub_info);
|
||||
|
||||
@@ -1,571 +0,0 @@
|
||||
#![allow(
|
||||
clippy::multiple_inherent_impl,
|
||||
clippy::unnecessary_wraps,
|
||||
clippy::unused_self,
|
||||
clippy::needless_pass_by_value
|
||||
)]
|
||||
|
||||
use crate::r#const::MPSC_CHANNEL_SIZE;
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
|
||||
use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
|
||||
use crate::ident::{PyKeypair, PyPeerId};
|
||||
use crate::pyclass;
|
||||
use libp2p::futures::StreamExt as _;
|
||||
use libp2p::gossipsub;
|
||||
use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
|
||||
use libp2p::swarm::SwarmEvent;
|
||||
use networking::discovery;
|
||||
use networking::swarm::create_swarm;
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods as _};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
|
||||
use std::net::IpAddr;
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
|
||||
mod exception {
|
||||
use pyo3::types::PyTuple;
|
||||
use pyo3::{exceptions::PyException, prelude::*};
|
||||
use pyo3_stub_gen::derive::*;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
|
||||
pub struct PyNoPeersSubscribedToTopicError {}
|
||||
|
||||
impl PyNoPeersSubscribedToTopicError {
|
||||
const MSG: &'static str = "\
|
||||
No peers are currently subscribed to receive messages on this topic. \
|
||||
Wait for peers to subscribe or check your network connectivity.";
|
||||
|
||||
/// Creates a new [ `PyErr` ] of this type.
|
||||
///
|
||||
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
|
||||
pub(crate) fn new_err() -> PyErr {
|
||||
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyNoPeersSubscribedToTopicError {
|
||||
#[new]
|
||||
#[pyo3(signature = (*args))]
|
||||
#[allow(unused_variables)]
|
||||
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("PeerId(\"{}\")", Self::MSG)
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
Self::MSG.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
|
||||
pub struct PyAllQueuesFullError {}
|
||||
|
||||
impl PyAllQueuesFullError {
|
||||
const MSG: &'static str =
|
||||
"All libp2p peers are unresponsive, resend the message or reconnect.";
|
||||
|
||||
/// Creates a new [ `PyErr` ] of this type.
|
||||
///
|
||||
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
|
||||
pub(crate) fn new_err() -> PyErr {
|
||||
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyAllQueuesFullError {
|
||||
#[new]
|
||||
#[pyo3(signature = (*args))]
|
||||
#[allow(unused_variables)]
|
||||
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("PeerId(\"{}\")", Self::MSG)
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
Self::MSG.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Connection or disconnection event discriminant type.
|
||||
#[gen_stub_pyclass_enum]
|
||||
#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum PyConnectionUpdateType {
|
||||
Connected = 0,
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, name = "ConnectionUpdate")]
|
||||
#[derive(Debug, Clone)]
|
||||
struct PyConnectionUpdate {
|
||||
/// Whether this is a connection or disconnection event
|
||||
#[pyo3(get)]
|
||||
update_type: PyConnectionUpdateType,
|
||||
|
||||
/// Identity of the peer that we have connected to or disconnected from.
|
||||
#[pyo3(get)]
|
||||
peer_id: PyPeerId,
|
||||
|
||||
/// Remote connection's IPv4 address.
|
||||
#[pyo3(get)]
|
||||
remote_ipv4: String,
|
||||
|
||||
/// Remote connection's TCP port.
|
||||
#[pyo3(get)]
|
||||
remote_tcp_port: u16,
|
||||
}
|
||||
|
||||
enum ToTask {
|
||||
GossipsubSubscribe {
|
||||
topic: String,
|
||||
result_tx: oneshot::Sender<PyResult<bool>>,
|
||||
},
|
||||
GossipsubUnsubscribe {
|
||||
topic: String,
|
||||
result_tx: oneshot::Sender<bool>,
|
||||
},
|
||||
GossipsubPublish {
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
result_tx: oneshot::Sender<PyResult<MessageId>>,
|
||||
},
|
||||
}
|
||||
|
||||
#[allow(clippy::enum_glob_use)]
|
||||
async fn networking_task(
|
||||
mut swarm: networking::swarm::Swarm,
|
||||
mut to_task_rx: mpsc::Receiver<ToTask>,
|
||||
connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
|
||||
gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
|
||||
) {
|
||||
use SwarmEvent::*;
|
||||
use ToTask::*;
|
||||
use networking::swarm::BehaviourEvent::*;
|
||||
|
||||
log::info!("RUST: networking task started");
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
message = to_task_rx.recv() => {
|
||||
// handle closed channel
|
||||
let Some(message) = message else {
|
||||
log::info!("RUST: channel closed");
|
||||
break;
|
||||
};
|
||||
|
||||
// dispatch incoming messages
|
||||
match message {
|
||||
GossipsubSubscribe { topic, result_tx } => {
|
||||
// try to subscribe
|
||||
let result = swarm.behaviour_mut()
|
||||
.gossipsub.subscribe(&IdentTopic::new(topic));
|
||||
|
||||
// send response oneshot
|
||||
if let Err(e) = result_tx.send(result.pyerr()) {
|
||||
log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
GossipsubUnsubscribe { topic, result_tx } => {
|
||||
// try to unsubscribe from the topic
|
||||
let result = swarm.behaviour_mut()
|
||||
.gossipsub.unsubscribe(&IdentTopic::new(topic));
|
||||
|
||||
// send response oneshot (or exit if connection closed)
|
||||
if let Err(e) = result_tx.send(result) {
|
||||
log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
GossipsubPublish { topic, data, result_tx } => {
|
||||
// try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
|
||||
let result = swarm.behaviour_mut().gossipsub.publish(
|
||||
IdentTopic::new(topic), data);
|
||||
let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
|
||||
Err(exception::PyNoPeersSubscribedToTopicError::new_err())
|
||||
} else if let Err(PublishError::AllQueuesFull(_)) = result {
|
||||
Err(exception::PyAllQueuesFullError::new_err())
|
||||
} else {
|
||||
result.pyerr()
|
||||
};
|
||||
|
||||
// send response oneshot (or exit if connection closed)
|
||||
if let Err(e) = result_tx.send(pyresult) {
|
||||
log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// architectural solution to this problem:
|
||||
// create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
|
||||
// -> it will emmit TRUE connected/disconnected events consumable elsewhere
|
||||
//
|
||||
// gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
|
||||
// then for actual communication it will dial those peers if need-be
|
||||
swarm_event = swarm.select_next_some() => {
|
||||
match swarm_event {
|
||||
Behaviour(Gossipsub(gossipsub::Event::Message {
|
||||
message: Message {
|
||||
topic,
|
||||
data,
|
||||
..
|
||||
},
|
||||
..
|
||||
})) => {
|
||||
// topic-ID is just the topic hash!!! (since we used identity hasher)
|
||||
let message = (topic.into_string(), data);
|
||||
|
||||
// send incoming message to channel (or exit if connection closed)
|
||||
if let Err(e) = gossipsub_message_tx.send(message).await {
|
||||
log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
|
||||
// grab IPv4 string
|
||||
let remote_ipv4 = match remote_ip {
|
||||
IpAddr::V4(ip) => ip.to_string(),
|
||||
IpAddr::V6(ip) => {
|
||||
log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// send connection event to channel (or exit if connection closed)
|
||||
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
|
||||
update_type: PyConnectionUpdateType::Connected,
|
||||
peer_id: PyPeerId(peer_id),
|
||||
remote_ipv4,
|
||||
remote_tcp_port,
|
||||
}).await {
|
||||
log::error!("RUST: could not send connection update since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
|
||||
// grab IPv4 string
|
||||
let remote_ipv4 = match remote_ip {
|
||||
IpAddr::V4(ip) => ip.to_string(),
|
||||
IpAddr::V6(ip) => {
|
||||
log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// send disconnection event to channel (or exit if connection closed)
|
||||
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
|
||||
update_type: PyConnectionUpdateType::Disconnected,
|
||||
peer_id: PyPeerId(peer_id),
|
||||
remote_ipv4,
|
||||
remote_tcp_port,
|
||||
}).await {
|
||||
log::error!("RUST: could not send connection update since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
e => {
|
||||
log::info!("RUST: other event {e:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
log::info!("RUST: networking task stopped");
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "NetworkingHandle")]
|
||||
#[derive(Debug)]
|
||||
struct PyNetworkingHandle {
|
||||
// channels
|
||||
to_task_tx: Option<mpsc::Sender<ToTask>>,
|
||||
connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
|
||||
gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
|
||||
}
|
||||
|
||||
impl Drop for PyNetworkingHandle {
|
||||
fn drop(&mut self) {
|
||||
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
|
||||
// to ensure that the networking task is done BEFORE exiting the clear function...
|
||||
// but this may require GIL?? and it may not be safe to call GIL here??
|
||||
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::expect_used)]
|
||||
impl PyNetworkingHandle {
|
||||
fn new(
|
||||
to_task_tx: mpsc::Sender<ToTask>,
|
||||
connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
|
||||
gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
to_task_tx: Some(to_task_tx),
|
||||
connection_update_rx: Mutex::new(connection_update_rx),
|
||||
gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
|
||||
}
|
||||
}
|
||||
|
||||
const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
|
||||
self.to_task_tx
|
||||
.as_ref()
|
||||
.expect("The sender should only be None after de-initialization.")
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyNetworkingHandle {
|
||||
// NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
|
||||
// immediately beforehand to release the interpreter.
|
||||
// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
|
||||
|
||||
// ---- Lifecycle management methods ----
|
||||
|
||||
#[new]
|
||||
fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
|
||||
use pyo3_async_runtimes::tokio::get_runtime;
|
||||
|
||||
// create communication channels
|
||||
let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
|
||||
// get identity
|
||||
let identity = identity.borrow().0.clone();
|
||||
|
||||
// create networking swarm (within tokio context!! or it crashes)
|
||||
let swarm = get_runtime()
|
||||
.block_on(async { create_swarm(identity) })
|
||||
.pyerr()?;
|
||||
|
||||
// spawn tokio task running the networking logic
|
||||
get_runtime().spawn(async move {
|
||||
networking_task(
|
||||
swarm,
|
||||
to_task_rx,
|
||||
connection_update_tx,
|
||||
gossipsub_message_tx,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
Ok(Self::new(
|
||||
to_task_tx,
|
||||
connection_update_rx,
|
||||
gossipsub_message_rx,
|
||||
))
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(()) // This is needed purely so `__clear__` can work
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
fn __clear__(&mut self) {
|
||||
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
|
||||
// to ensure that the networking task is done BEFORE exiting the clear function...
|
||||
// but this may require GIL?? and it may not be safe to call GIL here??
|
||||
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
|
||||
}
|
||||
|
||||
// ---- Connection update receiver methods ----
|
||||
|
||||
/// Receives the next `ConnectionUpdate` from networking.
|
||||
async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
|
||||
self.connection_update_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_py()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
}
|
||||
|
||||
/// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
|
||||
///
|
||||
/// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
|
||||
/// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
|
||||
/// will sleep until a `ConnectionUpdate`s is sent.
|
||||
async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
|
||||
self.connection_update_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_many_py(limit)
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
}
|
||||
|
||||
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
|
||||
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
|
||||
// so things don't randomly block
|
||||
// /// Tries to receive the next `ConnectionUpdate` from networking.
|
||||
// fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
|
||||
// self.connection_update_rx.blocking_lock().try_recv_py()
|
||||
// }
|
||||
//
|
||||
// /// Checks if the `ConnectionUpdate` channel is empty.
|
||||
// fn connection_update_is_empty(&self) -> bool {
|
||||
// self.connection_update_rx.blocking_lock().is_empty()
|
||||
// }
|
||||
//
|
||||
// /// Returns the number of `ConnectionUpdate`s in the channel.
|
||||
// fn connection_update_len(&self) -> usize {
|
||||
// self.connection_update_rx.blocking_lock().len()
|
||||
// }
|
||||
|
||||
// ---- Gossipsub management methods ----
|
||||
|
||||
/// Subscribe to a `GossipSub` topic.
|
||||
///
|
||||
/// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to subscribe
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubSubscribe {
|
||||
topic,
|
||||
result_tx: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
|
||||
// wait for response & return any errors
|
||||
rx.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())?
|
||||
}
|
||||
|
||||
/// Unsubscribes from a `GossipSub` topic.
|
||||
///
|
||||
/// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to unsubscribe
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubUnsubscribe {
|
||||
topic,
|
||||
result_tx: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
|
||||
// wait for response & convert any errors
|
||||
rx.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())
|
||||
}
|
||||
|
||||
/// Publishes a message with multiple topics to the `GossipSub` network.
|
||||
///
|
||||
/// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to subscribe
|
||||
let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubPublish {
|
||||
topic,
|
||||
data,
|
||||
result_tx: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
|
||||
// wait for response & return any errors => ignore messageID for now!!!
|
||||
let _ = rx
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---- Gossipsub message receiver methods ----
|
||||
|
||||
/// Receives the next message from the `GossipSub` network.
|
||||
async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
|
||||
self.gossipsub_message_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_py()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map(|(t, d)| (t, d.pybytes()))
|
||||
}
|
||||
|
||||
/// Receives at most `limit` messages from the `GossipSub` network and returns them.
|
||||
///
|
||||
/// For `limit = 0`, an empty collection of messages will be returned immediately.
|
||||
/// For `limit > 0`, if there are no messages in the channel's queue this method
|
||||
/// will sleep until a message is sent.
|
||||
async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
|
||||
Ok(self
|
||||
.gossipsub_message_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_many_py(limit)
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|(t, d)| (t, d.pybytes()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
|
||||
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
|
||||
// so things don't randomly block
|
||||
// /// Tries to receive the next message from the `GossipSub` network.
|
||||
// fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
|
||||
// Ok(self
|
||||
// .gossipsub_message_rx
|
||||
// .blocking_lock()
|
||||
// .try_recv_py()?
|
||||
// .map(|(t, d)| (t, d.pybytes())))
|
||||
// }
|
||||
//
|
||||
// /// Checks if the `GossipSub` message channel is empty.
|
||||
// fn gossipsub_is_empty(&self) -> bool {
|
||||
// self.gossipsub_message_rx.blocking_lock().is_empty()
|
||||
// }
|
||||
//
|
||||
// /// Returns the number of `GossipSub` messages in the channel.
|
||||
// fn gossipsub_len(&self) -> usize {
|
||||
// self.gossipsub_message_rx.blocking_lock().len()
|
||||
// }
|
||||
}
|
||||
|
||||
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
|
||||
m.add_class::<exception::PyAllQueuesFullError>()?;
|
||||
|
||||
m.add_class::<PyConnectionUpdateType>()?;
|
||||
m.add_class::<PyConnectionUpdate>()?;
|
||||
m.add_class::<PyConnectionUpdateType>()?;
|
||||
m.add_class::<PyNetworkingHandle>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,54 +0,0 @@
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use core::mem::drop;
|
||||
use core::option::Option::Some;
|
||||
use core::time::Duration;
|
||||
use tokio;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_drop_channel() {
|
||||
struct Ping;
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<Ping>(10);
|
||||
|
||||
let _ = tokio::spawn(async move {
|
||||
println!("TASK: entered");
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
result = rx.recv() => {
|
||||
match result {
|
||||
Some(_) => {
|
||||
println!("TASK: pinged");
|
||||
}
|
||||
None => {
|
||||
println!("TASK: closing channel");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
|
||||
println!("TASK: heartbeat");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("TASK: exited");
|
||||
});
|
||||
|
||||
let tx2 = tx.clone();
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
|
||||
tx.send(Ping).await.expect("Should not fail");
|
||||
drop(tx);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
|
||||
tx2.send(Ping).await.expect("Should not fail");
|
||||
drop(tx2);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
|
||||
}
|
||||
}
|
||||
@@ -1,16 +1,20 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
|
||||
from exo_pyo3_bindings import (
|
||||
Keypair,
|
||||
NetworkingHandle,
|
||||
NoPeersSubscribedToTopicError,
|
||||
PyFromSwarm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleep_on_multiple_items() -> None:
|
||||
print("PYTHON: starting handle")
|
||||
h = NetworkingHandle(Keypair.generate_ed25519())
|
||||
h = NetworkingHandle(Keypair.generate())
|
||||
|
||||
ct = asyncio.create_task(_await_cons(h))
|
||||
mt = asyncio.create_task(_await_msg(h))
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
|
||||
# sleep for 4 ticks
|
||||
for i in range(4):
|
||||
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
|
||||
print("caught it", e)
|
||||
|
||||
|
||||
async def _await_cons(h: NetworkingHandle):
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
c = await h.connection_update_recv()
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
|
||||
|
||||
async def _await_msg(h: NetworkingHandle):
|
||||
while True:
|
||||
m = await h.gossipsub_recv()
|
||||
print(f"PYTHON: message: {m}")
|
||||
event = await h.recv()
|
||||
match event:
|
||||
case PyFromSwarm.Connection() as c:
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
case PyFromSwarm.Message() as m:
|
||||
print(f"PYTHON: message: {m}")
|
||||
|
||||
@@ -5,7 +5,6 @@ edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
name = "networking"
|
||||
path = "src/lib.rs"
|
||||
|
||||
@@ -13,25 +12,22 @@ path = "src/lib.rs"
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
# datastructures
|
||||
either = { workspace = true }
|
||||
|
||||
# macro dependencies
|
||||
extend = { workspace = true }
|
||||
delegate = { workspace = true }
|
||||
|
||||
# async
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
futures-lite = { workspace = true }
|
||||
futures-timer = { workspace = true }
|
||||
|
||||
# utility dependencies
|
||||
util = { workspace = true }
|
||||
tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
|
||||
keccak-const = { workspace = true }
|
||||
pin-project = "1.1.10"
|
||||
|
||||
# tracing/logging
|
||||
log = { workspace = true }
|
||||
iroh = "0.96.1"
|
||||
iroh-gossip = "0.96.0"
|
||||
iroh-blobs = "0.98.0"
|
||||
iroh-docs = "0.96.0"
|
||||
|
||||
# networking
|
||||
libp2p = { workspace = true, features = ["full"] }
|
||||
|
||||
@@ -1,7 +1,3 @@
|
||||
use futures_lite::StreamExt;
|
||||
use libp2p::{gossipsub, identity, swarm::SwarmEvent};
|
||||
use networking::{discovery, swarm};
|
||||
use tokio::{io, io::AsyncBufReadExt as _, select};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use tracing_subscriber::filter::LevelFilter;
|
||||
|
||||
@@ -10,65 +6,4 @@ async fn main() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
|
||||
.try_init();
|
||||
|
||||
// Configure swarm
|
||||
let mut swarm =
|
||||
swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
|
||||
|
||||
// Create a Gossipsub topic & subscribe
|
||||
let topic = gossipsub::IdentTopic::new("test-net");
|
||||
swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.subscribe(&topic)
|
||||
.expect("Subscribing to topic failed");
|
||||
|
||||
// Read full lines from stdin
|
||||
let mut stdin = io::BufReader::new(io::stdin()).lines();
|
||||
println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
|
||||
|
||||
// Kick it off
|
||||
loop {
|
||||
select! {
|
||||
// on gossipsub outgoing
|
||||
Ok(Some(line)) = stdin.next_line() => {
|
||||
if let Err(e) = swarm
|
||||
.behaviour_mut().gossipsub
|
||||
.publish(topic.clone(), line.as_bytes()) {
|
||||
println!("Publish error: {e:?}");
|
||||
}
|
||||
}
|
||||
event = swarm.next() => match event {
|
||||
// on gossipsub incoming
|
||||
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
|
||||
propagation_source: peer_id,
|
||||
message_id: id,
|
||||
message,
|
||||
}))) => println!(
|
||||
"\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
|
||||
String::from_utf8_lossy(&message.data),
|
||||
),
|
||||
|
||||
// on discovery
|
||||
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
|
||||
discovery::Event::ConnectionEstablished {
|
||||
peer_id, connection_id, remote_ip, remote_tcp_port
|
||||
} => {
|
||||
println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
|
||||
}
|
||||
discovery::Event::ConnectionClosed {
|
||||
peer_id, connection_id, remote_ip, remote_tcp_port
|
||||
} => {
|
||||
eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
|
||||
}
|
||||
}
|
||||
|
||||
// ignore outgoing errors: those are normal
|
||||
e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
|
||||
|
||||
// otherwise log any other event
|
||||
e => { log::info!("Other event {e:?}"); }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,44 +0,0 @@
|
||||
https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609
|
||||
^raw sockets distributed -> `<net/ndrv.h>` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html
|
||||
--> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets.
|
||||
--> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received.
|
||||
--> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools.
|
||||
--> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html
|
||||
----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET
|
||||
|
||||
https://stackoverflow.com/questions/17169298/af-packet-on-osx
|
||||
^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only
|
||||
^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing
|
||||
|
||||
https://www.unix.com/man_page/mojave/4/ip/
|
||||
^OS X manpages for IP
|
||||
|
||||
https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts
|
||||
^driver kit, system extensions & kexts for macOS
|
||||
|
||||
----
|
||||
|
||||
To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network.
|
||||
--> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a
|
||||
----> here is a guide on how to set up thunderbolt-ethernet on linux
|
||||
----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS
|
||||
|
||||
https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7
|
||||
^GPT discussion about making socket interface
|
||||
|
||||
https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,??
|
||||
https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2
|
||||
^GPT discussion about accessing TB on MacOS low level interactions
|
||||
|
||||
--------------------------------
|
||||
|
||||
https://www.intel.com/content/www/us/en/support/articles/000098893/software.html
|
||||
^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge
|
||||
|
||||
|
||||
---------------------------------
|
||||
|
||||
https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/
|
||||
-->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet??
|
||||
-->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!!
|
||||
-->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c
|
||||
@@ -1,382 +0,0 @@
|
||||
use crate::ext::MultiaddrExt;
|
||||
use delegate::delegate;
|
||||
use either::Either;
|
||||
use futures_lite::FutureExt;
|
||||
use futures_timer::Delay;
|
||||
use libp2p::core::transport::PortUse;
|
||||
use libp2p::core::{ConnectedPoint, Endpoint};
|
||||
use libp2p::swarm::behaviour::ConnectionEstablished;
|
||||
use libp2p::swarm::dial_opts::DialOpts;
|
||||
use libp2p::swarm::{
|
||||
CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
|
||||
ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
|
||||
THandlerOutEvent, ToSwarm, dummy,
|
||||
};
|
||||
use libp2p::{Multiaddr, PeerId, identity, mdns};
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::convert::Infallible;
|
||||
use std::io;
|
||||
use std::net::IpAddr;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
use util::wakerdeque::WakerDeque;
|
||||
|
||||
const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
|
||||
|
||||
mod managed {
|
||||
use libp2p::swarm::NetworkBehaviour;
|
||||
use libp2p::{identity, mdns, ping};
|
||||
use std::io;
|
||||
use std::time::Duration;
|
||||
|
||||
const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
|
||||
const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
|
||||
const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
|
||||
const PING_INTERVAL: Duration = Duration::from_millis(2_500);
|
||||
|
||||
#[derive(NetworkBehaviour)]
|
||||
pub struct Behaviour {
|
||||
mdns: mdns::tokio::Behaviour,
|
||||
ping: ping::Behaviour,
|
||||
}
|
||||
|
||||
impl Behaviour {
|
||||
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
|
||||
Ok(Self {
|
||||
mdns: mdns_behaviour(keypair)?,
|
||||
ping: ping_behaviour(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
|
||||
use mdns::{Config, tokio};
|
||||
|
||||
// mDNS config => enable IPv6
|
||||
let mdns_config = Config {
|
||||
ttl: MDNS_RECORD_TTL,
|
||||
query_interval: MDNS_QUERY_INTERVAL,
|
||||
|
||||
// enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
|
||||
Ok(mdns_behaviour?)
|
||||
}
|
||||
|
||||
fn ping_behaviour() -> ping::Behaviour {
|
||||
ping::Behaviour::new(
|
||||
ping::Config::new()
|
||||
.with_timeout(PING_TIMEOUT)
|
||||
.with_interval(PING_INTERVAL),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Events for when a listening connection is truly established and truly closed.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Event {
|
||||
ConnectionEstablished {
|
||||
peer_id: PeerId,
|
||||
connection_id: ConnectionId,
|
||||
remote_ip: IpAddr,
|
||||
remote_tcp_port: u16,
|
||||
},
|
||||
ConnectionClosed {
|
||||
peer_id: PeerId,
|
||||
connection_id: ConnectionId,
|
||||
remote_ip: IpAddr,
|
||||
remote_tcp_port: u16,
|
||||
},
|
||||
}
|
||||
|
||||
/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
|
||||
///
|
||||
/// The behaviour operates as such:
|
||||
/// 1) All true (listening) connections/disconnections are tracked, emitting corresponding events
|
||||
/// to the swarm.
|
||||
/// 1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
|
||||
/// immediately, and expired but connected peers are disconnected from immediately.
|
||||
/// 2) Every fixed interval: discovered but not connected peers are dialed, and expired but
|
||||
/// connected peers are disconnected from.
|
||||
pub struct Behaviour {
|
||||
// state-tracking for managed behaviors & mDNS-discovered peers
|
||||
managed: managed::Behaviour,
|
||||
mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
|
||||
|
||||
retry_delay: Delay, // retry interval
|
||||
|
||||
// pending events to emmit => waker-backed Deque to control polling
|
||||
pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
|
||||
}
|
||||
|
||||
impl Behaviour {
|
||||
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
|
||||
Ok(Self {
|
||||
managed: managed::Behaviour::new(keypair)?,
|
||||
mdns_discovered: HashMap::new(),
|
||||
retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
|
||||
pending_events: WakerDeque::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
|
||||
self.pending_events.push_back(ToSwarm::Dial {
|
||||
opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
|
||||
})
|
||||
}
|
||||
|
||||
fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
|
||||
// push front to make this IMMEDIATE
|
||||
self.pending_events.push_front(ToSwarm::CloseConnection {
|
||||
peer_id,
|
||||
connection: CloseConnection::One(connection),
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
|
||||
for (p, ma) in peers {
|
||||
self.dial(p, ma.clone()); // always connect
|
||||
|
||||
// get peer's multi-addresses or insert if missing
|
||||
let Some(mas) = self.mdns_discovered.get_mut(&p) else {
|
||||
self.mdns_discovered.insert(p, BTreeSet::from([ma]));
|
||||
continue;
|
||||
};
|
||||
|
||||
// multiaddress should never already be present - else something has gone wrong
|
||||
let is_new_addr = mas.insert(ma);
|
||||
assert!(is_new_addr, "cannot discover a discovered peer");
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
|
||||
for (p, ma) in peers {
|
||||
// at this point, we *must* have the peer
|
||||
let mas = self
|
||||
.mdns_discovered
|
||||
.get_mut(&p)
|
||||
.expect("nonexistent peer cannot expire");
|
||||
|
||||
// at this point, we *must* have the multiaddress
|
||||
let was_present = mas.remove(&ma);
|
||||
assert!(was_present, "nonexistent multiaddress cannot expire");
|
||||
|
||||
// if empty, remove the peer-id entirely
|
||||
if mas.is_empty() {
|
||||
self.mdns_discovered.remove(&p);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn on_connection_established(
|
||||
&mut self,
|
||||
peer_id: PeerId,
|
||||
connection_id: ConnectionId,
|
||||
remote_ip: IpAddr,
|
||||
remote_tcp_port: u16,
|
||||
) {
|
||||
// send out connected event
|
||||
self.pending_events
|
||||
.push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
|
||||
peer_id,
|
||||
connection_id,
|
||||
remote_ip,
|
||||
remote_tcp_port,
|
||||
}));
|
||||
}
|
||||
|
||||
fn on_connection_closed(
|
||||
&mut self,
|
||||
peer_id: PeerId,
|
||||
connection_id: ConnectionId,
|
||||
remote_ip: IpAddr,
|
||||
remote_tcp_port: u16,
|
||||
) {
|
||||
// send out disconnected event
|
||||
self.pending_events
|
||||
.push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
|
||||
peer_id,
|
||||
connection_id,
|
||||
remote_ip,
|
||||
remote_tcp_port,
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
impl NetworkBehaviour for Behaviour {
|
||||
type ConnectionHandler =
|
||||
ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
|
||||
type ToSwarm = Event;
|
||||
|
||||
// simply delegate to underlying mDNS behaviour
|
||||
|
||||
delegate! {
|
||||
to self.managed {
|
||||
fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
|
||||
fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_established_inbound_connection(
|
||||
&mut self,
|
||||
connection_id: ConnectionId,
|
||||
peer: PeerId,
|
||||
local_addr: &Multiaddr,
|
||||
remote_addr: &Multiaddr,
|
||||
) -> Result<THandler<Self>, ConnectionDenied> {
|
||||
Ok(ConnectionHandler::select(
|
||||
dummy::ConnectionHandler,
|
||||
self.managed.handle_established_inbound_connection(
|
||||
connection_id,
|
||||
peer,
|
||||
local_addr,
|
||||
remote_addr,
|
||||
)?,
|
||||
))
|
||||
}
|
||||
|
||||
#[allow(clippy::needless_question_mark)]
|
||||
fn handle_established_outbound_connection(
|
||||
&mut self,
|
||||
connection_id: ConnectionId,
|
||||
peer: PeerId,
|
||||
addr: &Multiaddr,
|
||||
role_override: Endpoint,
|
||||
port_use: PortUse,
|
||||
) -> Result<THandler<Self>, ConnectionDenied> {
|
||||
Ok(ConnectionHandler::select(
|
||||
dummy::ConnectionHandler,
|
||||
self.managed.handle_established_outbound_connection(
|
||||
connection_id,
|
||||
peer,
|
||||
addr,
|
||||
role_override,
|
||||
port_use,
|
||||
)?,
|
||||
))
|
||||
}
|
||||
|
||||
fn on_connection_handler_event(
|
||||
&mut self,
|
||||
peer_id: PeerId,
|
||||
connection_id: ConnectionId,
|
||||
event: THandlerOutEvent<Self>,
|
||||
) {
|
||||
match event {
|
||||
Either::Left(ev) => libp2p::core::util::unreachable(ev),
|
||||
Either::Right(ev) => {
|
||||
self.managed
|
||||
.on_connection_handler_event(peer_id, connection_id, ev)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// hook into these methods to drive behavior
|
||||
|
||||
fn on_swarm_event(&mut self, event: FromSwarm) {
|
||||
self.managed.on_swarm_event(event); // let mDNS handle swarm events
|
||||
|
||||
// handle swarm events to update internal state:
|
||||
match event {
|
||||
FromSwarm::ConnectionEstablished(ConnectionEstablished {
|
||||
peer_id,
|
||||
connection_id,
|
||||
endpoint,
|
||||
..
|
||||
}) => {
|
||||
let remote_address = match endpoint {
|
||||
ConnectedPoint::Dialer { address, .. } => address,
|
||||
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
|
||||
};
|
||||
|
||||
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
|
||||
// handle connection established event which is filtered correctly
|
||||
self.on_connection_established(peer_id, connection_id, ip, port)
|
||||
}
|
||||
}
|
||||
FromSwarm::ConnectionClosed(ConnectionClosed {
|
||||
peer_id,
|
||||
connection_id,
|
||||
endpoint,
|
||||
..
|
||||
}) => {
|
||||
let remote_address = match endpoint {
|
||||
ConnectedPoint::Dialer { address, .. } => address,
|
||||
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
|
||||
};
|
||||
|
||||
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
|
||||
// handle connection closed event which is filtered correctly
|
||||
self.on_connection_closed(peer_id, connection_id, ip, port)
|
||||
}
|
||||
}
|
||||
|
||||
// since we are running TCP/IP transport layer, we are assuming that
|
||||
// no address changes can occur, hence encountering one is a fatal error
|
||||
FromSwarm::AddressChange(a) => {
|
||||
unreachable!("unhandlable: address change encountered: {:?}", a)
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
|
||||
// delegate to managed behaviors for any behaviors they need to perform
|
||||
match self.managed.poll(cx) {
|
||||
Poll::Ready(ToSwarm::GenerateEvent(e)) => {
|
||||
match e {
|
||||
// handle discovered and expired events from mDNS
|
||||
managed::BehaviourEvent::Mdns(e) => match e.clone() {
|
||||
mdns::Event::Discovered(peers) => {
|
||||
self.handle_mdns_discovered(peers);
|
||||
}
|
||||
mdns::Event::Expired(peers) => {
|
||||
self.handle_mdns_expired(peers);
|
||||
}
|
||||
},
|
||||
|
||||
// handle ping events => if error then disconnect
|
||||
managed::BehaviourEvent::Ping(e) => {
|
||||
if let Err(_) = e.result {
|
||||
self.close_connection(e.peer, e.connection.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// since we just consumed an event, we should immediately wake just in case
|
||||
// there are more events to come where that came from
|
||||
cx.waker().wake_by_ref();
|
||||
}
|
||||
|
||||
// forward any other mDNS event to the swarm or its connection handler(s)
|
||||
Poll::Ready(e) => {
|
||||
return Poll::Ready(
|
||||
e.map_out(|_| unreachable!("events returning to swarm already handled"))
|
||||
.map_in(Either::Right),
|
||||
);
|
||||
}
|
||||
|
||||
Poll::Pending => {}
|
||||
}
|
||||
|
||||
// retry connecting to all mDNS peers periodically (fails safely if already connected)
|
||||
if self.retry_delay.poll(cx).is_ready() {
|
||||
for (p, mas) in self.mdns_discovered.clone() {
|
||||
for ma in mas {
|
||||
self.dial(p, ma)
|
||||
}
|
||||
}
|
||||
self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
|
||||
}
|
||||
|
||||
// send out any pending events from our own service
|
||||
if let Some(e) = self.pending_events.pop_front(cx) {
|
||||
return Poll::Ready(e.map_in(Either::Left));
|
||||
}
|
||||
|
||||
// wait for pending events
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
//! TODO: crate documentation
|
||||
//!
|
||||
//! this is here as a placeholder documentation
|
||||
//!
|
||||
//!
|
||||
pub mod discovery;
|
||||
pub mod swarm;
|
||||
|
||||
/// Namespace for all the type/trait aliases used by this crate.
|
||||
pub(crate) mod alias {
|
||||
use std::error::Error;
|
||||
|
||||
pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
|
||||
pub type AnyResult<T> = Result<T, AnyError>;
|
||||
}
|
||||
|
||||
/// Namespace for crate-wide extension traits/methods
|
||||
pub(crate) mod ext {
|
||||
use extend::ext;
|
||||
use libp2p::Multiaddr;
|
||||
use libp2p::multiaddr::Protocol;
|
||||
use std::net::IpAddr;
|
||||
|
||||
#[ext(pub, name = MultiaddrExt)]
|
||||
impl Multiaddr {
|
||||
/// If the multiaddress corresponds to a TCP address, extracts it
|
||||
fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
|
||||
let mut ps = self.into_iter();
|
||||
let ip = if let Some(p) = ps.next() {
|
||||
match p {
|
||||
Protocol::Ip4(ip) => IpAddr::V4(ip),
|
||||
Protocol::Ip6(ip) => IpAddr::V6(ip),
|
||||
_ => return None,
|
||||
}
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
let Some(Protocol::Tcp(port)) = ps.next() else {
|
||||
return None;
|
||||
};
|
||||
Some((ip, port))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
use crate::alias;
|
||||
use crate::swarm::transport::tcp_transport;
|
||||
pub use behaviour::{Behaviour, BehaviourEvent};
|
||||
use libp2p::{SwarmBuilder, identity};
|
||||
|
||||
pub type Swarm = libp2p::Swarm<Behaviour>;
|
||||
|
||||
/// The current version of the network: this prevents devices running different versions of the
|
||||
/// software from interacting with each other.
|
||||
///
|
||||
/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
|
||||
/// even be, and how to inject the right version into this config/initialization. E.g. should
|
||||
/// this be passed in as a parameter? What about rapidly changing versions in debug builds?
|
||||
/// this is all VERY very hard to figure out and needs to be mulled over as a team.
|
||||
pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
|
||||
pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
|
||||
|
||||
/// Create and configure a swarm which listens to all ports on OS
|
||||
pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
|
||||
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
|
||||
.with_tokio()
|
||||
.with_other_transport(tcp_transport)?
|
||||
.with_behaviour(Behaviour::new)?
|
||||
.build();
|
||||
|
||||
// Listen on all interfaces and whatever port the OS assigns
|
||||
swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
|
||||
Ok(swarm)
|
||||
}
|
||||
|
||||
mod transport {
|
||||
use crate::alias;
|
||||
use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
|
||||
use futures_lite::{AsyncRead, AsyncWrite};
|
||||
use keccak_const::Sha3_256;
|
||||
use libp2p::core::muxing;
|
||||
use libp2p::core::transport::Boxed;
|
||||
use libp2p::pnet::{PnetError, PnetOutput};
|
||||
use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
|
||||
use std::{env, sync::LazyLock};
|
||||
|
||||
/// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
|
||||
/// See [`pnet_upgrade`] for more.
|
||||
static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
|
||||
let builder = Sha3_256::new().update(b"exo_discovery_network");
|
||||
|
||||
if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
|
||||
let bytes = var.into_bytes();
|
||||
builder.update(&bytes)
|
||||
} else {
|
||||
builder.update(NETWORK_VERSION)
|
||||
}
|
||||
.finalize()
|
||||
});
|
||||
|
||||
/// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
|
||||
/// also different-versioned instances of this same network.
|
||||
/// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
|
||||
async fn pnet_upgrade<TSocket>(
|
||||
socket: TSocket,
|
||||
_: impl Sized,
|
||||
) -> Result<PnetOutput<TSocket>, PnetError>
|
||||
where
|
||||
TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
|
||||
{
|
||||
use pnet::{PnetConfig, PreSharedKey};
|
||||
PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
|
||||
.handshake(socket)
|
||||
.await
|
||||
}
|
||||
|
||||
/// TCP/IP transport layer configuration.
|
||||
pub fn tcp_transport(
|
||||
keypair: &identity::Keypair,
|
||||
) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
|
||||
use libp2p::{
|
||||
core::upgrade::Version,
|
||||
tcp::{Config, tokio},
|
||||
};
|
||||
|
||||
// `TCP_NODELAY` enabled => avoid latency
|
||||
let tcp_config = Config::default().nodelay(true);
|
||||
|
||||
// V1 + lazy flushing => 0-RTT negotiation
|
||||
let upgrade_version = Version::V1Lazy;
|
||||
|
||||
// Noise is faster than TLS + we don't care much for security
|
||||
let noise_config = noise::Config::new(keypair)?;
|
||||
|
||||
// Use default Yamux config for multiplexing
|
||||
let yamux_config = yamux::Config::default();
|
||||
|
||||
// Create new Tokio-driven TCP/IP transport layer
|
||||
let base_transport = tokio::Transport::new(tcp_config)
|
||||
.and_then(pnet_upgrade)
|
||||
.upgrade(upgrade_version)
|
||||
.authenticate(noise_config)
|
||||
.multiplex(yamux_config);
|
||||
|
||||
// Return boxed transport (to flatten complex type)
|
||||
Ok(base_transport.boxed())
|
||||
}
|
||||
}
|
||||
|
||||
mod behaviour {
|
||||
use crate::{alias, discovery};
|
||||
use libp2p::swarm::NetworkBehaviour;
|
||||
use libp2p::{gossipsub, identity};
|
||||
|
||||
/// Behavior of the Swarm which composes all desired behaviors:
|
||||
/// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
|
||||
#[derive(NetworkBehaviour)]
|
||||
pub struct Behaviour {
|
||||
pub discovery: discovery::Behaviour,
|
||||
pub gossipsub: gossipsub::Behaviour,
|
||||
}
|
||||
|
||||
impl Behaviour {
|
||||
pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
|
||||
Ok(Self {
|
||||
discovery: discovery::Behaviour::new(keypair)?,
|
||||
gossipsub: gossipsub_behaviour(keypair),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
|
||||
use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
|
||||
|
||||
// build a gossipsub network behaviour
|
||||
// => signed message authenticity + strict validation mode means the message-ID is
|
||||
// automatically provided by gossipsub w/out needing to provide custom message-ID function
|
||||
gossipsub::Behaviour::new(
|
||||
MessageAuthenticity::Signed(keypair.clone()),
|
||||
ConfigBuilder::default()
|
||||
.max_transmit_size(1024 * 1024)
|
||||
.validation_mode(ValidationMode::Strict)
|
||||
.build()
|
||||
.expect("the configuration should always be valid"),
|
||||
)
|
||||
.expect("creating gossipsub behavior should always work")
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
// maybe this will hold test in the future...??
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn does_nothing() {}
|
||||
}
|
||||
@@ -1,15 +0,0 @@
|
||||
[package]
|
||||
name = "util"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
name = "util"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
@@ -1 +0,0 @@
|
||||
pub mod wakerdeque;
|
||||
@@ -1,55 +0,0 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::fmt::{Debug, Formatter};
|
||||
use std::task::{Context, Waker};
|
||||
|
||||
/// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods,
|
||||
/// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods.
|
||||
pub struct WakerDeque<T> {
|
||||
waker: Option<Waker>,
|
||||
deque: VecDeque<T>,
|
||||
}
|
||||
|
||||
impl<T: Debug> Debug for WakerDeque<T> {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
self.deque.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> WakerDeque<T> {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
waker: None,
|
||||
deque: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn update(&mut self, cx: &mut Context<'_>) {
|
||||
self.waker = Some(cx.waker().clone());
|
||||
}
|
||||
|
||||
fn wake(&mut self) {
|
||||
let Some(ref mut w) = self.waker else { return };
|
||||
w.wake_by_ref();
|
||||
self.waker = None;
|
||||
}
|
||||
|
||||
pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option<T> {
|
||||
self.update(cx);
|
||||
self.deque.pop_front()
|
||||
}
|
||||
|
||||
pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option<T> {
|
||||
self.update(cx);
|
||||
self.deque.pop_back()
|
||||
}
|
||||
|
||||
pub fn push_front(&mut self, value: T) {
|
||||
self.wake();
|
||||
self.deque.push_front(value);
|
||||
}
|
||||
|
||||
pub fn push_back(&mut self, value: T) {
|
||||
self.wake();
|
||||
self.deque.push_back(value);
|
||||
}
|
||||
}
|
||||
@@ -1,31 +1,28 @@
|
||||
import asyncio
|
||||
import socket
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Iterator
|
||||
|
||||
import anyio
|
||||
from anyio import current_time
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
|
||||
from exo.download.download_utils import (
|
||||
RepoDownloadProgress,
|
||||
delete_model,
|
||||
map_repo_download_progress_to_download_progress_data,
|
||||
resolve_model_in_path,
|
||||
)
|
||||
from exo.download.shard_downloader import ShardDownloader
|
||||
from exo.shared.constants import EXO_MODELS_DIR
|
||||
from exo.shared.models.model_cards import ModelId
|
||||
from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
|
||||
from exo.shared.models.model_cards import ModelId, get_model_cards
|
||||
from exo.shared.types.commands import (
|
||||
CancelDownload,
|
||||
DeleteDownload,
|
||||
ForwarderDownloadCommand,
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
ForwarderEvent,
|
||||
NodeDownloadProgress,
|
||||
)
|
||||
from exo.shared.types.worker.downloads import (
|
||||
@@ -35,36 +32,29 @@ from exo.shared.types.worker.downloads import (
|
||||
DownloadPending,
|
||||
DownloadProgress,
|
||||
)
|
||||
from exo.shared.types.worker.shards import ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
@dataclass
|
||||
class DownloadCoordinator:
|
||||
node_id: NodeId
|
||||
session_id: SessionId
|
||||
shard_downloader: ShardDownloader
|
||||
download_command_receiver: Receiver[ForwarderDownloadCommand]
|
||||
local_event_sender: Sender[ForwarderEvent]
|
||||
event_index_counter: Iterator[int]
|
||||
event_sender: Sender[Event]
|
||||
offline: bool = False
|
||||
|
||||
# Local state
|
||||
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
|
||||
active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
|
||||
|
||||
# Internal event channel for forwarding (initialized in __post_init__)
|
||||
event_sender: Sender[Event] = field(init=False)
|
||||
event_receiver: Receiver[Event] = field(init=False)
|
||||
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
# Per-model throttle for download progress events
|
||||
_last_progress_time: dict[ModelId, float] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.event_sender, self.event_receiver = channel[Event]()
|
||||
if self.offline:
|
||||
self.shard_downloader.set_internet_connection(False)
|
||||
self.shard_downloader.on_progress(self._download_progress_callback)
|
||||
|
||||
def _model_dir(self, model_id: ModelId) -> str:
|
||||
@@ -80,7 +70,7 @@ class DownloadCoordinator:
|
||||
completed = DownloadCompleted(
|
||||
shard_metadata=callback_shard,
|
||||
node_id=self.node_id,
|
||||
total_bytes=progress.total_bytes,
|
||||
total=progress.total,
|
||||
model_directory=self._model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
@@ -113,45 +103,16 @@ class DownloadCoordinator:
|
||||
logger.info(
|
||||
f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
|
||||
)
|
||||
if not self.offline:
|
||||
self._test_internet_connection()
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._command_processor)
|
||||
tg.start_soon(self._forward_events)
|
||||
tg.start_soon(self._emit_existing_download_progress)
|
||||
if not self.offline:
|
||||
tg.start_soon(self._check_internet_connection)
|
||||
|
||||
def _test_internet_connection(self) -> None:
|
||||
# Try multiple endpoints since some ISPs/networks block specific IPs
|
||||
for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
|
||||
try:
|
||||
socket.create_connection((host, 443), timeout=3).close()
|
||||
self.shard_downloader.set_internet_connection(True)
|
||||
logger.debug(f"Internet connectivity: True (via {host})")
|
||||
return
|
||||
except OSError:
|
||||
continue
|
||||
self.shard_downloader.set_internet_connection(False)
|
||||
logger.debug("Internet connectivity: False")
|
||||
|
||||
async def _check_internet_connection(self) -> None:
|
||||
first_connection = True
|
||||
while True:
|
||||
await asyncio.sleep(10)
|
||||
|
||||
# Assume that internet connection is set to False on 443 errors.
|
||||
if self.shard_downloader.internet_connection:
|
||||
continue
|
||||
|
||||
self._test_internet_connection()
|
||||
|
||||
if first_connection and self.shard_downloader.internet_connection:
|
||||
first_connection = False
|
||||
self._tg.start_soon(self._emit_existing_download_progress)
|
||||
try:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._command_processor)
|
||||
tg.start_soon(self._emit_existing_download_progress)
|
||||
finally:
|
||||
for task in self.active_downloads.values():
|
||||
task.cancel()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._tg.cancel_scope.cancel()
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _command_processor(self) -> None:
|
||||
with self.download_command_receiver as commands:
|
||||
@@ -185,6 +146,25 @@ class DownloadCoordinator:
|
||||
)
|
||||
return
|
||||
|
||||
# Check EXO_MODELS_PATH for pre-downloaded models
|
||||
found_path = resolve_model_in_path(model_id)
|
||||
if found_path is not None:
|
||||
logger.info(
|
||||
f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}"
|
||||
)
|
||||
completed = DownloadCompleted(
|
||||
shard_metadata=shard,
|
||||
node_id=self.node_id,
|
||||
total=shard.model_card.storage_size,
|
||||
model_directory=str(found_path),
|
||||
read_only=True,
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=completed)
|
||||
)
|
||||
return
|
||||
|
||||
# Emit pending status
|
||||
progress = DownloadPending(
|
||||
shard_metadata=shard,
|
||||
@@ -203,7 +183,7 @@ class DownloadCoordinator:
|
||||
completed = DownloadCompleted(
|
||||
shard_metadata=shard,
|
||||
node_id=self.node_id,
|
||||
total_bytes=initial_progress.total_bytes,
|
||||
total=initial_progress.total,
|
||||
model_directory=self._model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = completed
|
||||
@@ -269,6 +249,15 @@ class DownloadCoordinator:
|
||||
self.active_downloads[model_id] = task
|
||||
|
||||
async def _delete_download(self, model_id: ModelId) -> None:
|
||||
# Protect read-only models (from EXO_MODELS_PATH) from deletion
|
||||
if model_id in self.download_status:
|
||||
current = self.download_status[model_id]
|
||||
if isinstance(current, DownloadCompleted) and current.read_only:
|
||||
logger.warning(
|
||||
f"Refusing to delete read-only model {model_id} (from EXO_MODELS_PATH)"
|
||||
)
|
||||
return
|
||||
|
||||
# Cancel if active
|
||||
if model_id in self.active_downloads:
|
||||
logger.info(f"Cancelling active download for {model_id} before deletion")
|
||||
@@ -297,21 +286,6 @@ class DownloadCoordinator:
|
||||
)
|
||||
del self.download_status[model_id]
|
||||
|
||||
async def _forward_events(self) -> None:
|
||||
with self.event_receiver as events:
|
||||
async for event in events:
|
||||
idx = next(self.event_index_counter)
|
||||
fe = ForwarderEvent(
|
||||
origin_idx=idx,
|
||||
origin=self.node_id,
|
||||
session=self.session_id,
|
||||
event=event,
|
||||
)
|
||||
logger.debug(
|
||||
f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
|
||||
)
|
||||
await self.local_event_sender.send(fe)
|
||||
|
||||
async def _emit_existing_download_progress(self) -> None:
|
||||
try:
|
||||
while True:
|
||||
@@ -332,19 +306,21 @@ class DownloadCoordinator:
|
||||
status: DownloadProgress = DownloadCompleted(
|
||||
node_id=self.node_id,
|
||||
shard_metadata=progress.shard,
|
||||
total_bytes=progress.total_bytes,
|
||||
total=progress.total,
|
||||
model_directory=self._model_dir(
|
||||
progress.shard.model_card.model_id
|
||||
),
|
||||
)
|
||||
elif progress.status in ["in_progress", "not_started"]:
|
||||
if progress.downloaded_bytes_this_session.in_bytes == 0:
|
||||
if progress.downloaded_this_session.in_bytes == 0:
|
||||
status = DownloadPending(
|
||||
node_id=self.node_id,
|
||||
shard_metadata=progress.shard,
|
||||
model_directory=self._model_dir(
|
||||
progress.shard.model_card.model_id
|
||||
),
|
||||
downloaded=progress.downloaded,
|
||||
total=progress.total,
|
||||
)
|
||||
else:
|
||||
status = DownloadOngoing(
|
||||
@@ -364,6 +340,39 @@ class DownloadCoordinator:
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=status)
|
||||
)
|
||||
# Scan EXO_MODELS_PATH for pre-downloaded models
|
||||
if EXO_MODELS_PATH is not None:
|
||||
for card in await get_model_cards():
|
||||
mid = card.model_id
|
||||
if mid in self.active_downloads:
|
||||
continue
|
||||
if isinstance(
|
||||
self.download_status.get(mid),
|
||||
(DownloadCompleted, DownloadOngoing, DownloadFailed),
|
||||
):
|
||||
continue
|
||||
found = resolve_model_in_path(mid)
|
||||
if found is not None:
|
||||
path_shard = PipelineShardMetadata(
|
||||
model_card=card,
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
start_layer=0,
|
||||
end_layer=card.n_layers,
|
||||
n_layers=card.n_layers,
|
||||
)
|
||||
path_completed: DownloadProgress = DownloadCompleted(
|
||||
node_id=self.node_id,
|
||||
shard_metadata=path_shard,
|
||||
total=card.storage_size,
|
||||
model_directory=str(found),
|
||||
read_only=True,
|
||||
)
|
||||
self.download_status[mid] = path_completed
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=path_completed)
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"DownloadCoordinator: Done emitting existing download progress."
|
||||
)
|
||||
|
||||
@@ -20,7 +20,6 @@ from huggingface_hub import (
|
||||
)
|
||||
from loguru import logger
|
||||
from pydantic import (
|
||||
DirectoryPath,
|
||||
TypeAdapter,
|
||||
)
|
||||
|
||||
@@ -31,7 +30,7 @@ from exo.download.huggingface_utils import (
|
||||
get_hf_endpoint,
|
||||
get_hf_token,
|
||||
)
|
||||
from exo.shared.constants import EXO_MODELS_DIR
|
||||
from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
|
||||
from exo.shared.models.model_cards import ModelTask
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.memory import Memory
|
||||
@@ -80,9 +79,9 @@ def map_repo_file_download_progress_to_download_progress_data(
|
||||
repo_file_download_progress: RepoFileDownloadProgress,
|
||||
) -> DownloadProgressData:
|
||||
return DownloadProgressData(
|
||||
downloaded_bytes=repo_file_download_progress.downloaded,
|
||||
downloaded_bytes_this_session=repo_file_download_progress.downloaded_this_session,
|
||||
total_bytes=repo_file_download_progress.total,
|
||||
downloaded=repo_file_download_progress.downloaded,
|
||||
downloaded_this_session=repo_file_download_progress.downloaded_this_session,
|
||||
total=repo_file_download_progress.total,
|
||||
completed_files=1 if repo_file_download_progress.status == "complete" else 0,
|
||||
total_files=1,
|
||||
speed=repo_file_download_progress.speed,
|
||||
@@ -95,9 +94,9 @@ def map_repo_download_progress_to_download_progress_data(
|
||||
repo_download_progress: RepoDownloadProgress,
|
||||
) -> DownloadProgressData:
|
||||
return DownloadProgressData(
|
||||
total_bytes=repo_download_progress.total_bytes,
|
||||
downloaded_bytes=repo_download_progress.downloaded_bytes,
|
||||
downloaded_bytes_this_session=repo_download_progress.downloaded_bytes_this_session,
|
||||
total=repo_download_progress.total,
|
||||
downloaded=repo_download_progress.downloaded,
|
||||
downloaded_this_session=repo_download_progress.downloaded_this_session,
|
||||
completed_files=repo_download_progress.completed_files,
|
||||
total_files=repo_download_progress.total_files,
|
||||
speed=repo_download_progress.overall_speed,
|
||||
@@ -111,7 +110,27 @@ def map_repo_download_progress_to_download_progress_data(
|
||||
)
|
||||
|
||||
|
||||
def build_model_path(model_id: ModelId) -> DirectoryPath:
|
||||
def resolve_model_in_path(model_id: ModelId) -> Path | None:
|
||||
"""Search EXO_MODELS_PATH directories for a pre-existing model.
|
||||
|
||||
Checks each directory for the normalized name (org--model). A candidate
|
||||
is only returned if ``is_model_directory_complete`` confirms all weight
|
||||
files are present.
|
||||
"""
|
||||
if EXO_MODELS_PATH is None:
|
||||
return None
|
||||
normalized = model_id.normalize()
|
||||
for search_dir in EXO_MODELS_PATH:
|
||||
candidate = search_dir / normalized
|
||||
if candidate.is_dir() and is_model_directory_complete(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def build_model_path(model_id: ModelId) -> Path:
|
||||
found = resolve_model_in_path(model_id)
|
||||
if found is not None:
|
||||
return found
|
||||
return EXO_MODELS_DIR / model_id.normalize()
|
||||
|
||||
|
||||
@@ -142,7 +161,7 @@ async def delete_model(model_id: ModelId) -> bool:
|
||||
|
||||
|
||||
async def seed_models(seed_dir: str | Path):
|
||||
"""Move model in resources folder of app to .cache/huggingface/hub"""
|
||||
"""Move models from resources folder to EXO_MODELS_DIR."""
|
||||
source_dir = Path(seed_dir)
|
||||
dest_dir = await ensure_models_dir()
|
||||
for path in source_dir.iterdir():
|
||||
@@ -158,6 +177,72 @@ async def seed_models(seed_dir: str | Path):
|
||||
logger.error(traceback.format_exc())
|
||||
|
||||
|
||||
def _scan_model_directory(
|
||||
model_dir: Path, recursive: bool = False
|
||||
) -> list[FileListEntry] | None:
|
||||
"""Scan a local model directory and build a file list.
|
||||
|
||||
Requires at least one ``*.safetensors.index.json``. Every weight file
|
||||
referenced by the index that is missing on disk gets ``size=None``.
|
||||
"""
|
||||
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
|
||||
if not index_files:
|
||||
return None
|
||||
|
||||
entries_by_path: dict[str, FileListEntry] = {}
|
||||
|
||||
if recursive:
|
||||
for dirpath, _, filenames in os.walk(model_dir):
|
||||
for filename in filenames:
|
||||
if filename.endswith(".partial"):
|
||||
continue
|
||||
full_path = Path(dirpath) / filename
|
||||
rel_path = str(full_path.relative_to(model_dir))
|
||||
entries_by_path[rel_path] = FileListEntry(
|
||||
type="file",
|
||||
path=rel_path,
|
||||
size=full_path.stat().st_size,
|
||||
)
|
||||
else:
|
||||
for item in model_dir.iterdir():
|
||||
if item.is_file() and not item.name.endswith(".partial"):
|
||||
entries_by_path[item.name] = FileListEntry(
|
||||
type="file",
|
||||
path=item.name,
|
||||
size=item.stat().st_size,
|
||||
)
|
||||
|
||||
# Add expected weight files from index that haven't been downloaded yet
|
||||
for index_file in index_files:
|
||||
try:
|
||||
index_data = ModelSafetensorsIndex.model_validate_json(
|
||||
index_file.read_text()
|
||||
)
|
||||
relative_dir = index_file.parent.relative_to(model_dir)
|
||||
for filename in set(index_data.weight_map.values()):
|
||||
rel_path = (
|
||||
str(relative_dir / filename)
|
||||
if relative_dir != Path(".")
|
||||
else filename
|
||||
)
|
||||
if rel_path not in entries_by_path:
|
||||
entries_by_path[rel_path] = FileListEntry(
|
||||
type="file",
|
||||
path=rel_path,
|
||||
size=None,
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return list(entries_by_path.values())
|
||||
|
||||
|
||||
def is_model_directory_complete(model_dir: Path) -> bool:
|
||||
"""Check if a model directory contains all required weight files."""
|
||||
file_list = _scan_model_directory(model_dir, recursive=True)
|
||||
return file_list is not None and all(f.size is not None for f in file_list)
|
||||
|
||||
|
||||
async def _build_file_list_from_local_directory(
|
||||
model_id: ModelId,
|
||||
recursive: bool = False,
|
||||
@@ -172,59 +257,7 @@ async def _build_file_list_from_local_directory(
|
||||
if not await aios.path.exists(model_dir):
|
||||
return None
|
||||
|
||||
def _scan() -> list[FileListEntry] | None:
|
||||
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
|
||||
if not index_files:
|
||||
return None
|
||||
|
||||
entries_by_path: dict[str, FileListEntry] = {}
|
||||
|
||||
if recursive:
|
||||
for dirpath, _, filenames in os.walk(model_dir):
|
||||
for filename in filenames:
|
||||
if filename.endswith(".partial"):
|
||||
continue
|
||||
full_path = Path(dirpath) / filename
|
||||
rel_path = str(full_path.relative_to(model_dir))
|
||||
entries_by_path[rel_path] = FileListEntry(
|
||||
type="file",
|
||||
path=rel_path,
|
||||
size=full_path.stat().st_size,
|
||||
)
|
||||
else:
|
||||
for item in model_dir.iterdir():
|
||||
if item.is_file() and not item.name.endswith(".partial"):
|
||||
entries_by_path[item.name] = FileListEntry(
|
||||
type="file",
|
||||
path=item.name,
|
||||
size=item.stat().st_size,
|
||||
)
|
||||
|
||||
# Add expected weight files from index that haven't been downloaded yet
|
||||
for index_file in index_files:
|
||||
try:
|
||||
index_data = ModelSafetensorsIndex.model_validate_json(
|
||||
index_file.read_text()
|
||||
)
|
||||
relative_dir = index_file.parent.relative_to(model_dir)
|
||||
for filename in set(index_data.weight_map.values()):
|
||||
rel_path = (
|
||||
str(relative_dir / filename)
|
||||
if relative_dir != Path(".")
|
||||
else filename
|
||||
)
|
||||
if rel_path not in entries_by_path:
|
||||
entries_by_path[rel_path] = FileListEntry(
|
||||
type="file",
|
||||
path=rel_path,
|
||||
size=None,
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return list(entries_by_path.values())
|
||||
|
||||
file_list = await asyncio.to_thread(_scan)
|
||||
file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
|
||||
if not file_list:
|
||||
return None
|
||||
return file_list
|
||||
@@ -281,9 +314,13 @@ async def fetch_file_list_with_cache(
|
||||
_fetched_file_lists_this_session.add(cache_key)
|
||||
return file_list
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning(
|
||||
"Ran into exception when fetching file list from HF."
|
||||
)
|
||||
|
||||
if await aios.path.exists(cache_file):
|
||||
logger.warning(
|
||||
f"No internet and no cached file list for {model_id} - using local file list"
|
||||
f"No cached file list for {model_id} - using local file list"
|
||||
)
|
||||
async with aiofiles.open(cache_file, "r") as f:
|
||||
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
|
||||
@@ -578,19 +615,20 @@ def calculate_repo_progress(
|
||||
file_progress: dict[str, RepoFileDownloadProgress],
|
||||
all_start_time: float,
|
||||
) -> RepoDownloadProgress:
|
||||
all_total_bytes = sum((p.total.in_bytes for p in file_progress.values()), 0)
|
||||
all_downloaded_bytes = sum(
|
||||
(p.downloaded.in_bytes for p in file_progress.values()), 0
|
||||
all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0))
|
||||
all_downloaded = sum(
|
||||
(p.downloaded for p in file_progress.values()), Memory.from_bytes(0)
|
||||
)
|
||||
all_downloaded_bytes_this_session = sum(
|
||||
(p.downloaded_this_session.in_bytes for p in file_progress.values()), 0
|
||||
all_downloaded_this_session = sum(
|
||||
(p.downloaded_this_session for p in file_progress.values()),
|
||||
Memory.from_bytes(0),
|
||||
)
|
||||
elapsed_time = time.time() - all_start_time
|
||||
all_speed = (
|
||||
all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0
|
||||
all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0
|
||||
)
|
||||
all_eta = (
|
||||
timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed)
|
||||
timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed)
|
||||
if all_speed > 0
|
||||
else timedelta(seconds=0)
|
||||
)
|
||||
@@ -609,11 +647,9 @@ def calculate_repo_progress(
|
||||
[p for p in file_progress.values() if p.downloaded == p.total]
|
||||
),
|
||||
total_files=len(file_progress),
|
||||
downloaded_bytes=Memory.from_bytes(all_downloaded_bytes),
|
||||
downloaded_bytes_this_session=Memory.from_bytes(
|
||||
all_downloaded_bytes_this_session
|
||||
),
|
||||
total_bytes=Memory.from_bytes(all_total_bytes),
|
||||
downloaded=all_downloaded,
|
||||
downloaded_this_session=all_downloaded_this_session,
|
||||
total=all_total,
|
||||
overall_speed=all_speed,
|
||||
overall_eta=all_eta,
|
||||
status=status,
|
||||
@@ -791,6 +827,7 @@ async def download_shard(
|
||||
|
||||
for file in filtered_file_list:
|
||||
downloaded_bytes = await get_downloaded_size(target_dir / file.path)
|
||||
final_file_exists = await aios.path.exists(target_dir / file.path)
|
||||
file_progress[file.path] = RepoFileDownloadProgress(
|
||||
repo_id=shard.model_card.model_id,
|
||||
repo_revision=revision,
|
||||
@@ -800,7 +837,9 @@ async def download_shard(
|
||||
total=Memory.from_bytes(file.size or 0),
|
||||
speed=0,
|
||||
eta=timedelta(0),
|
||||
status="complete" if downloaded_bytes == file.size else "not_started",
|
||||
status="complete"
|
||||
if final_file_exists and downloaded_bytes == file.size
|
||||
else "not_started",
|
||||
start_time=time.time(),
|
||||
)
|
||||
|
||||
|
||||
@@ -15,9 +15,13 @@ from exo.shared.types.worker.shards import (
|
||||
)
|
||||
|
||||
|
||||
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
|
||||
def exo_shard_downloader(
|
||||
max_parallel_downloads: int = 8, offline: bool = False
|
||||
) -> ShardDownloader:
|
||||
return SingletonShardDownloader(
|
||||
CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
|
||||
CachedShardDownloader(
|
||||
ResumableShardDownloader(max_parallel_downloads, offline=offline)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -50,10 +54,6 @@ class SingletonShardDownloader(ShardDownloader):
|
||||
self.shard_downloader = shard_downloader
|
||||
self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
self.shard_downloader.set_internet_connection(value)
|
||||
|
||||
def on_progress(
|
||||
self,
|
||||
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
|
||||
@@ -90,10 +90,6 @@ class CachedShardDownloader(ShardDownloader):
|
||||
self.shard_downloader = shard_downloader
|
||||
self.cache: dict[tuple[str, ShardMetadata], Path] = {}
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
self.shard_downloader.set_internet_connection(value)
|
||||
|
||||
def on_progress(
|
||||
self,
|
||||
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
|
||||
@@ -123,8 +119,9 @@ class CachedShardDownloader(ShardDownloader):
|
||||
|
||||
|
||||
class ResumableShardDownloader(ShardDownloader):
|
||||
def __init__(self, max_parallel_downloads: int = 8):
|
||||
def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
|
||||
self.max_parallel_downloads = max_parallel_downloads
|
||||
self.offline = offline
|
||||
self.on_progress_callbacks: list[
|
||||
Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
|
||||
] = []
|
||||
@@ -151,8 +148,7 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
self.on_progress_wrapper,
|
||||
max_parallel_downloads=self.max_parallel_downloads,
|
||||
allow_patterns=allow_patterns,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
return target_dir
|
||||
|
||||
@@ -168,8 +164,7 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
shard,
|
||||
self.on_progress_wrapper,
|
||||
skip_download=True,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
|
||||
semaphore = asyncio.Semaphore(self.max_parallel_downloads)
|
||||
@@ -198,7 +193,6 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
shard,
|
||||
self.on_progress_wrapper,
|
||||
skip_download=True,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
return progress
|
||||
|
||||
@@ -16,11 +16,6 @@ from exo.shared.types.worker.shards import (
|
||||
|
||||
# TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
|
||||
class ShardDownloader(ABC):
|
||||
internet_connection: bool = False
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
|
||||
@abstractmethod
|
||||
async def ensure_shard(
|
||||
self, shard: ShardMetadata, config_only: bool = False
|
||||
@@ -107,9 +102,9 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
|
||||
),
|
||||
completed_files=0,
|
||||
total_files=0,
|
||||
downloaded_bytes=Memory.from_bytes(0),
|
||||
downloaded_bytes_this_session=Memory.from_bytes(0),
|
||||
total_bytes=Memory.from_bytes(0),
|
||||
downloaded=Memory.from_bytes(0),
|
||||
downloaded_this_session=Memory.from_bytes(0),
|
||||
total=Memory.from_bytes(0),
|
||||
overall_speed=0,
|
||||
overall_eta=timedelta(seconds=0),
|
||||
status="complete",
|
||||
|
||||
+53
-42
@@ -1,14 +1,12 @@
|
||||
import argparse
|
||||
import itertools
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import resource
|
||||
import signal
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Iterator, Self
|
||||
from typing import Self
|
||||
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
from pydantic import PositiveInt
|
||||
|
||||
@@ -17,6 +15,7 @@ from exo.download.coordinator import DownloadCoordinator
|
||||
from exo.download.impl_shard_downloader import exo_shard_downloader
|
||||
from exo.master.api import API # TODO: should API be in master?
|
||||
from exo.master.main import Master
|
||||
from exo.routing.event_router import EventRouter
|
||||
from exo.routing.router import Router, get_node_id_keypair
|
||||
from exo.shared.constants import EXO_LOG
|
||||
from exo.shared.election import Election, ElectionResult
|
||||
@@ -24,12 +23,14 @@ from exo.shared.logging import logger_cleanup, logger_setup
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.utils.channels import Receiver, channel
|
||||
from exo.utils.pydantic_ext import CamelCaseModel
|
||||
from exo.utils.task_group import TaskGroup
|
||||
from exo.worker.main import Worker
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
router: Router
|
||||
event_router: EventRouter
|
||||
download_coordinator: DownloadCoordinator | None
|
||||
worker: Worker | None
|
||||
election: Election # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
|
||||
@@ -38,14 +39,13 @@ class Node:
|
||||
api: API | None
|
||||
|
||||
node_id: NodeId
|
||||
event_index_counter: Iterator[int]
|
||||
offline: bool
|
||||
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> "Self":
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
keypair = get_node_id_keypair()
|
||||
node_id = NodeId(keypair.to_peer_id().to_base58())
|
||||
node_id = NodeId(keypair.to_node_id())
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
router = Router.create(keypair)
|
||||
await router.register_topic(topics.GLOBAL_EVENTS)
|
||||
@@ -54,21 +54,22 @@ class Node:
|
||||
await router.register_topic(topics.ELECTION_MESSAGES)
|
||||
await router.register_topic(topics.CONNECTION_MESSAGES)
|
||||
await router.register_topic(topics.DOWNLOAD_COMMANDS)
|
||||
event_router = EventRouter(
|
||||
session_id,
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
external_outbound=router.sender(topics.LOCAL_EVENTS),
|
||||
external_inbound=router.receiver(topics.GLOBAL_EVENTS),
|
||||
)
|
||||
|
||||
logger.info(f"Starting node {node_id}")
|
||||
|
||||
# Create shared event index counter for Worker and DownloadCoordinator
|
||||
event_index_counter = itertools.count()
|
||||
|
||||
# Create DownloadCoordinator (unless --no-downloads)
|
||||
if not args.no_downloads:
|
||||
download_coordinator = DownloadCoordinator(
|
||||
node_id,
|
||||
session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(offline=args.offline),
|
||||
event_sender=event_router.sender(),
|
||||
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
|
||||
local_event_sender=router.sender(topics.LOCAL_EVENTS),
|
||||
event_index_counter=event_index_counter,
|
||||
offline=args.offline,
|
||||
)
|
||||
else:
|
||||
@@ -77,9 +78,8 @@ class Node:
|
||||
if args.spawn_api:
|
||||
api = API(
|
||||
node_id,
|
||||
session_id,
|
||||
port=args.api_port,
|
||||
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
event_receiver=event_router.receiver(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
election_receiver=router.receiver(topics.ELECTION_MESSAGES),
|
||||
@@ -90,12 +90,10 @@ class Node:
|
||||
if not args.no_worker:
|
||||
worker = Worker(
|
||||
node_id,
|
||||
session_id,
|
||||
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
local_event_sender=router.sender(topics.LOCAL_EVENTS),
|
||||
event_receiver=event_router.receiver(),
|
||||
event_sender=event_router.sender(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
event_index_counter=event_index_counter,
|
||||
)
|
||||
else:
|
||||
worker = None
|
||||
@@ -104,6 +102,7 @@ class Node:
|
||||
master = Master(
|
||||
node_id,
|
||||
session_id,
|
||||
event_sender=event_router.sender(),
|
||||
global_event_sender=router.sender(topics.GLOBAL_EVENTS),
|
||||
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
|
||||
command_receiver=router.receiver(topics.COMMANDS),
|
||||
@@ -126,6 +125,7 @@ class Node:
|
||||
|
||||
return cls(
|
||||
router,
|
||||
event_router,
|
||||
download_coordinator,
|
||||
worker,
|
||||
election,
|
||||
@@ -133,7 +133,6 @@ class Node:
|
||||
master,
|
||||
api,
|
||||
node_id,
|
||||
event_index_counter,
|
||||
args.offline,
|
||||
)
|
||||
|
||||
@@ -142,6 +141,7 @@ class Node:
|
||||
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
|
||||
signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
|
||||
tg.start_soon(self.router.run)
|
||||
tg.start_soon(self.event_router.run)
|
||||
tg.start_soon(self.election.run)
|
||||
if self.download_coordinator:
|
||||
tg.start_soon(self.download_coordinator.run)
|
||||
@@ -155,11 +155,11 @@ class Node:
|
||||
|
||||
def shutdown(self):
|
||||
# if this is our second call to shutdown, just sys.exit
|
||||
if self._tg.cancel_scope.cancel_called:
|
||||
if self._tg.cancel_called():
|
||||
import sys
|
||||
|
||||
sys.exit(1)
|
||||
self._tg.cancel_scope.cancel()
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _elect_loop(self):
|
||||
with self.election_result_receiver as results:
|
||||
@@ -176,6 +176,17 @@ class Node:
|
||||
# - Shutdown and re-create the worker
|
||||
# - Shut down and re-create the API
|
||||
|
||||
if result.is_new_master:
|
||||
await anyio.sleep(0)
|
||||
self.event_router.shutdown()
|
||||
self.event_router = EventRouter(
|
||||
result.session_id,
|
||||
self.router.sender(topics.COMMANDS),
|
||||
self.router.receiver(topics.GLOBAL_EVENTS),
|
||||
self.router.sender(topics.LOCAL_EVENTS),
|
||||
)
|
||||
self._tg.start_soon(self.event_router.run)
|
||||
|
||||
if (
|
||||
result.session_id.master_node_id == self.node_id
|
||||
and self.master is not None
|
||||
@@ -189,6 +200,7 @@ class Node:
|
||||
self.master = Master(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
event_sender=self.event_router.sender(),
|
||||
global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
|
||||
local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
|
||||
command_receiver=self.router.receiver(topics.COMMANDS),
|
||||
@@ -211,20 +223,15 @@ class Node:
|
||||
f"Node {result.session_id.master_node_id} elected master"
|
||||
)
|
||||
if result.is_new_master:
|
||||
await anyio.sleep(0)
|
||||
# Fresh counter for new session (buffer expects indices from 0)
|
||||
self.event_index_counter = itertools.count()
|
||||
if self.download_coordinator:
|
||||
self.download_coordinator.shutdown()
|
||||
self.download_coordinator = DownloadCoordinator(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(offline=self.offline),
|
||||
event_sender=self.event_router.sender(),
|
||||
download_command_receiver=self.router.receiver(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
|
||||
event_index_counter=self.event_index_counter,
|
||||
offline=self.offline,
|
||||
)
|
||||
self._tg.start_soon(self.download_coordinator.run)
|
||||
@@ -233,20 +240,16 @@ class Node:
|
||||
# TODO: add profiling etc to resource monitor
|
||||
self.worker = Worker(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
global_event_receiver=self.router.receiver(
|
||||
topics.GLOBAL_EVENTS
|
||||
),
|
||||
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
|
||||
event_receiver=self.event_router.receiver(),
|
||||
event_sender=self.event_router.sender(),
|
||||
command_sender=self.router.sender(topics.COMMANDS),
|
||||
download_command_sender=self.router.sender(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
event_index_counter=self.event_index_counter,
|
||||
)
|
||||
self._tg.start_soon(self.worker.run)
|
||||
if self.api:
|
||||
self.api.reset(result.session_id, result.won_clock)
|
||||
self.api.reset(result.won_clock, self.event_router.receiver())
|
||||
else:
|
||||
if self.api:
|
||||
self.api.unpause(result.won_clock)
|
||||
@@ -258,7 +261,7 @@ def main():
|
||||
target = min(max(soft, 65535), hard)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
|
||||
|
||||
mp.set_start_method("spawn")
|
||||
mp.set_start_method("spawn", force=True)
|
||||
# TODO: Refactor the current verbosity system
|
||||
logger_setup(EXO_LOG, args.verbosity)
|
||||
logger.info("Starting EXO")
|
||||
@@ -276,9 +279,16 @@ def main():
|
||||
logger.info("FAST_SYNCH forced OFF")
|
||||
|
||||
node = anyio.run(Node.create, args)
|
||||
anyio.run(node.run)
|
||||
logger.info("EXO Shutdown complete")
|
||||
logger_cleanup()
|
||||
try:
|
||||
anyio.run(node.run)
|
||||
except BaseException as exception:
|
||||
logger.opt(exception=exception).critical(
|
||||
"EXO terminated due to unhandled exception"
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
logger.info("EXO Shutdown complete")
|
||||
logger_cleanup()
|
||||
|
||||
|
||||
class Args(CamelCaseModel):
|
||||
@@ -289,7 +299,7 @@ class Args(CamelCaseModel):
|
||||
tb_only: bool = False
|
||||
no_worker: bool = False
|
||||
no_downloads: bool = False
|
||||
offline: bool = False
|
||||
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
|
||||
fast_synch: bool | None = None # None = auto, True = force on, False = force off
|
||||
|
||||
@classmethod
|
||||
@@ -340,6 +350,7 @@ class Args(CamelCaseModel):
|
||||
parser.add_argument(
|
||||
"--offline",
|
||||
action="store_true",
|
||||
default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
|
||||
help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
|
||||
)
|
||||
fast_synch_group = parser.add_mutually_exclusive_group()
|
||||
|
||||
@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
|
||||
ToolCallChunk,
|
||||
)
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
TextGenerationTaskParams,
|
||||
resolve_reasoning_params,
|
||||
)
|
||||
|
||||
|
||||
def chat_request_to_text_generation(
|
||||
@@ -59,7 +63,11 @@ def chat_request_to_text_generation(
|
||||
chat_template_messages.append({"role": "system", "content": content})
|
||||
else:
|
||||
# Skip messages with no meaningful content
|
||||
if msg.content is None and msg.thinking is None and msg.tool_calls is None:
|
||||
if (
|
||||
msg.content is None
|
||||
and msg.reasoning_content is None
|
||||
and msg.tool_calls is None
|
||||
):
|
||||
continue
|
||||
|
||||
if msg.role in ("user", "assistant", "developer"):
|
||||
@@ -71,6 +79,10 @@ def chat_request_to_text_generation(
|
||||
dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
|
||||
chat_template_messages.append(dumped)
|
||||
|
||||
resolved_effort, resolved_thinking = resolve_reasoning_params(
|
||||
request.reasoning_effort, request.enable_thinking
|
||||
)
|
||||
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_messages
|
||||
@@ -85,7 +97,8 @@ def chat_request_to_text_generation(
|
||||
seed=request.seed,
|
||||
stream=request.stream,
|
||||
tools=request.tools,
|
||||
enable_thinking=request.enable_thinking,
|
||||
reasoning_effort=resolved_effort,
|
||||
enable_thinking=resolved_thinking,
|
||||
chat_template_messages=chat_template_messages
|
||||
if chat_template_messages
|
||||
else None,
|
||||
@@ -111,6 +124,11 @@ def chunk_to_response(
|
||||
]
|
||||
)
|
||||
|
||||
if chunk.is_thinking:
|
||||
delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text)
|
||||
else:
|
||||
delta = ChatCompletionMessage(role="assistant", content=chunk.text)
|
||||
|
||||
return ChatCompletionResponse(
|
||||
id=command_id,
|
||||
created=int(time.time()),
|
||||
@@ -118,7 +136,7 @@ def chunk_to_response(
|
||||
choices=[
|
||||
StreamingChoiceResponse(
|
||||
index=0,
|
||||
delta=ChatCompletionMessage(role="assistant", content=chunk.text),
|
||||
delta=delta,
|
||||
logprobs=logprobs,
|
||||
finish_reason=chunk.finish_reason,
|
||||
)
|
||||
@@ -208,6 +226,7 @@ async def collect_chat_response(
|
||||
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
|
||||
"""Collect all token chunks and return a single ChatCompletionResponse."""
|
||||
text_parts: list[str] = []
|
||||
thinking_parts: list[str] = []
|
||||
tool_calls: list[ToolCall] = []
|
||||
logprobs_content: list[LogprobsContentItem] = []
|
||||
model: str | None = None
|
||||
@@ -228,7 +247,10 @@ async def collect_chat_response(
|
||||
if model is None:
|
||||
model = chunk.model
|
||||
last_usage = chunk.usage or last_usage
|
||||
text_parts.append(chunk.text)
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
else:
|
||||
text_parts.append(chunk.text)
|
||||
if chunk.logprob is not None:
|
||||
logprobs_content.append(
|
||||
LogprobsContentItem(
|
||||
@@ -258,6 +280,7 @@ async def collect_chat_response(
|
||||
raise ValueError(error_message)
|
||||
|
||||
combined_text = "".join(text_parts)
|
||||
combined_thinking = "".join(thinking_parts) if thinking_parts else None
|
||||
assert model is not None
|
||||
|
||||
yield ChatCompletionResponse(
|
||||
@@ -270,6 +293,7 @@ async def collect_chat_response(
|
||||
message=ChatCompletionMessage(
|
||||
role="assistant",
|
||||
content=combined_text,
|
||||
reasoning_content=combined_thinking,
|
||||
tool_calls=tool_calls if tool_calls else None,
|
||||
),
|
||||
logprobs=Logprobs(content=logprobs_content)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Claude Messages API adapter for converting requests/responses."""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
@@ -28,6 +29,8 @@ from exo.shared.types.claude_api import (
|
||||
ClaudeStopReason,
|
||||
ClaudeTextBlock,
|
||||
ClaudeTextDelta,
|
||||
ClaudeThinkingBlock,
|
||||
ClaudeThinkingDelta,
|
||||
ClaudeToolResultBlock,
|
||||
ClaudeToolUseBlock,
|
||||
ClaudeUsage,
|
||||
@@ -61,6 +64,22 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
|
||||
return "".join(sub_block.text for sub_block in block.content)
|
||||
|
||||
|
||||
# Matches "x-anthropic-billing-header: ...;" (with optional trailing newline)
|
||||
# or similar telemetry headers that change every request and break KV prefix caching.
|
||||
_VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE)
|
||||
|
||||
|
||||
def _strip_volatile_headers(text: str) -> str:
|
||||
"""Remove Anthropic billing/telemetry headers from system prompt text.
|
||||
|
||||
Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...;
|
||||
cc_entrypoint=...; cch=...;' that contain per-request content hashes. These
|
||||
change every request and break KV prefix caching (the prefix diverges at ~20
|
||||
tokens instead of matching thousands of conversation tokens).
|
||||
"""
|
||||
return _VOLATILE_HEADER_RE.sub("", text)
|
||||
|
||||
|
||||
def claude_request_to_text_generation(
|
||||
request: ClaudeMessagesRequest,
|
||||
) -> TextGenerationTaskParams:
|
||||
@@ -73,6 +92,8 @@ def claude_request_to_text_generation(
|
||||
instructions = request.system
|
||||
else:
|
||||
instructions = "".join(block.text for block in request.system)
|
||||
|
||||
instructions = _strip_volatile_headers(instructions)
|
||||
chat_template_messages.append({"role": "system", "content": instructions})
|
||||
|
||||
# Convert messages to input
|
||||
@@ -85,12 +106,15 @@ def claude_request_to_text_generation(
|
||||
|
||||
# Process structured content blocks
|
||||
text_parts: list[str] = []
|
||||
thinking_parts: list[str] = []
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
tool_results: list[ClaudeToolResultBlock] = []
|
||||
|
||||
for block in msg.content:
|
||||
if isinstance(block, ClaudeTextBlock):
|
||||
text_parts.append(block.text)
|
||||
elif isinstance(block, ClaudeThinkingBlock):
|
||||
thinking_parts.append(block.thinking)
|
||||
elif isinstance(block, ClaudeToolUseBlock):
|
||||
tool_calls.append(
|
||||
{
|
||||
@@ -106,6 +130,7 @@ def claude_request_to_text_generation(
|
||||
tool_results.append(block)
|
||||
|
||||
content = "".join(text_parts)
|
||||
reasoning_content = "".join(thinking_parts) if thinking_parts else None
|
||||
|
||||
# Build InputMessage from text content
|
||||
if msg.role in ("user", "assistant"):
|
||||
@@ -113,9 +138,14 @@ def claude_request_to_text_generation(
|
||||
|
||||
# Build chat_template_messages preserving tool structure
|
||||
if tool_calls:
|
||||
chat_template_messages.append(
|
||||
{"role": "assistant", "content": content, "tool_calls": tool_calls}
|
||||
)
|
||||
chat_msg: dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": tool_calls,
|
||||
}
|
||||
if reasoning_content:
|
||||
chat_msg["reasoning_content"] = reasoning_content
|
||||
chat_template_messages.append(chat_msg)
|
||||
elif tool_results:
|
||||
for tr in tool_results:
|
||||
chat_template_messages.append(
|
||||
@@ -126,7 +156,10 @@ def claude_request_to_text_generation(
|
||||
}
|
||||
)
|
||||
else:
|
||||
chat_template_messages.append({"role": msg.role, "content": content})
|
||||
chat_msg = {"role": msg.role, "content": content}
|
||||
if reasoning_content:
|
||||
chat_msg["reasoning_content"] = reasoning_content
|
||||
chat_template_messages.append(chat_msg)
|
||||
|
||||
# Convert Claude tool definitions to OpenAI-style function tools
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
@@ -143,6 +176,10 @@ def claude_request_to_text_generation(
|
||||
for tool in request.tools
|
||||
]
|
||||
|
||||
enable_thinking: bool | None = None
|
||||
if request.thinking is not None:
|
||||
enable_thinking = request.thinking.type in ("enabled", "adaptive")
|
||||
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_messages
|
||||
@@ -156,6 +193,7 @@ def claude_request_to_text_generation(
|
||||
stop=request.stop_sequences,
|
||||
stream=request.stream,
|
||||
tools=tools,
|
||||
enable_thinking=enable_thinking,
|
||||
chat_template_messages=chat_template_messages
|
||||
if chat_template_messages
|
||||
else None,
|
||||
@@ -173,6 +211,7 @@ async def collect_claude_response(
|
||||
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
|
||||
"""Collect all token chunks and return a single ClaudeMessagesResponse."""
|
||||
text_parts: list[str] = []
|
||||
thinking_parts: list[str] = []
|
||||
tool_use_blocks: list[ClaudeToolUseBlock] = []
|
||||
stop_reason: ClaudeStopReason | None = None
|
||||
last_usage: Usage | None = None
|
||||
@@ -200,7 +239,10 @@ async def collect_claude_response(
|
||||
stop_reason = "tool_use"
|
||||
continue
|
||||
|
||||
text_parts.append(chunk.text)
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
else:
|
||||
text_parts.append(chunk.text)
|
||||
|
||||
if chunk.finish_reason is not None:
|
||||
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
|
||||
@@ -209,9 +251,12 @@ async def collect_claude_response(
|
||||
raise ValueError(error_message)
|
||||
|
||||
combined_text = "".join(text_parts)
|
||||
combined_thinking = "".join(thinking_parts)
|
||||
|
||||
# Build content blocks
|
||||
content: list[ClaudeContentBlock] = []
|
||||
if combined_thinking:
|
||||
content.append(ClaudeThinkingBlock(thinking=combined_thinking))
|
||||
if combined_text:
|
||||
content.append(ClaudeTextBlock(text=combined_text))
|
||||
content.extend(tool_use_blocks)
|
||||
@@ -256,16 +301,16 @@ async def generate_claude_stream(
|
||||
start_event = ClaudeMessageStartEvent(message=initial_message)
|
||||
yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
|
||||
|
||||
# content_block_start for text block at index 0
|
||||
block_start = ClaudeContentBlockStartEvent(
|
||||
index=0, content_block=ClaudeTextBlock(text="")
|
||||
)
|
||||
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
|
||||
|
||||
output_tokens = 0
|
||||
stop_reason: ClaudeStopReason | None = None
|
||||
last_usage: Usage | None = None
|
||||
next_block_index = 1 # text block is 0, tool blocks start at 1
|
||||
next_block_index = 0
|
||||
|
||||
# Track whether we've started thinking/text blocks
|
||||
thinking_block_started = False
|
||||
thinking_block_index = -1
|
||||
text_block_started = False
|
||||
text_block_index = -1
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
if isinstance(chunk, PrefillProgressChunk):
|
||||
@@ -310,12 +355,45 @@ async def generate_claude_stream(
|
||||
|
||||
output_tokens += 1 # Count each chunk as one token
|
||||
|
||||
# content_block_delta
|
||||
delta_event = ClaudeContentBlockDeltaEvent(
|
||||
index=0,
|
||||
delta=ClaudeTextDelta(text=chunk.text),
|
||||
)
|
||||
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
|
||||
if chunk.is_thinking:
|
||||
# Start thinking block on first thinking token
|
||||
if not thinking_block_started:
|
||||
thinking_block_started = True
|
||||
thinking_block_index = next_block_index
|
||||
next_block_index += 1
|
||||
block_start = ClaudeContentBlockStartEvent(
|
||||
index=thinking_block_index,
|
||||
content_block=ClaudeThinkingBlock(thinking=""),
|
||||
)
|
||||
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
|
||||
|
||||
delta_event = ClaudeContentBlockDeltaEvent(
|
||||
index=thinking_block_index,
|
||||
delta=ClaudeThinkingDelta(thinking=chunk.text),
|
||||
)
|
||||
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
|
||||
else:
|
||||
# Close thinking block when transitioning to text
|
||||
if thinking_block_started and text_block_index == -1:
|
||||
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
|
||||
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
|
||||
|
||||
# Start text block on first text token
|
||||
if not text_block_started:
|
||||
text_block_started = True
|
||||
text_block_index = next_block_index
|
||||
next_block_index += 1
|
||||
block_start = ClaudeContentBlockStartEvent(
|
||||
index=text_block_index,
|
||||
content_block=ClaudeTextBlock(text=""),
|
||||
)
|
||||
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
|
||||
|
||||
delta_event = ClaudeContentBlockDeltaEvent(
|
||||
index=text_block_index,
|
||||
delta=ClaudeTextDelta(text=chunk.text),
|
||||
)
|
||||
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
|
||||
|
||||
if chunk.finish_reason is not None:
|
||||
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
|
||||
@@ -324,9 +402,22 @@ async def generate_claude_stream(
|
||||
if last_usage is not None:
|
||||
output_tokens = last_usage.completion_tokens
|
||||
|
||||
# content_block_stop for text block
|
||||
block_stop = ClaudeContentBlockStopEvent(index=0)
|
||||
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
|
||||
# Close any open blocks
|
||||
if thinking_block_started and text_block_index == -1:
|
||||
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
|
||||
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
|
||||
|
||||
if text_block_started:
|
||||
block_stop = ClaudeContentBlockStopEvent(index=text_block_index)
|
||||
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
|
||||
|
||||
if not thinking_block_started and not text_block_started:
|
||||
empty_start = ClaudeContentBlockStartEvent(
|
||||
index=0, content_block=ClaudeTextBlock(text="")
|
||||
)
|
||||
yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n"
|
||||
empty_stop = ClaudeContentBlockStopEvent(index=0)
|
||||
yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n"
|
||||
|
||||
# message_delta
|
||||
message_delta = ClaudeMessageDeltaEvent(
|
||||
|
||||
@@ -0,0 +1,456 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
from exo.shared.types.chunks import (
|
||||
ErrorChunk,
|
||||
PrefillProgressChunk,
|
||||
TokenChunk,
|
||||
ToolCallChunk,
|
||||
)
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.ollama_api import (
|
||||
OllamaChatRequest,
|
||||
OllamaChatResponse,
|
||||
OllamaDoneReason,
|
||||
OllamaGenerateRequest,
|
||||
OllamaGenerateResponse,
|
||||
OllamaMessage,
|
||||
OllamaToolCall,
|
||||
OllamaToolFunction,
|
||||
)
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
|
||||
|
||||
def _map_done_reason(
|
||||
finish_reason: str | None,
|
||||
) -> OllamaDoneReason | None:
|
||||
if finish_reason is None:
|
||||
return None
|
||||
if finish_reason == "stop":
|
||||
return "stop"
|
||||
if finish_reason == "length":
|
||||
return "length"
|
||||
if finish_reason in ("tool_calls", "function_call"):
|
||||
return "tool_call"
|
||||
if finish_reason == "error":
|
||||
return "error"
|
||||
return "stop"
|
||||
|
||||
|
||||
def _try_parse_json(value: str) -> dict[str, Any] | str:
|
||||
try:
|
||||
return json.loads(value) # type: ignore
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
|
||||
|
||||
def _build_tool_calls(chunk: ToolCallChunk) -> list[OllamaToolCall]:
|
||||
tool_calls: list[OllamaToolCall] = []
|
||||
for index, tool in enumerate(chunk.tool_calls):
|
||||
# tool.arguments is always str; try to parse as JSON dict for Ollama format
|
||||
arguments: dict[str, Any] | str = _try_parse_json(tool.arguments)
|
||||
tool_calls.append(
|
||||
OllamaToolCall(
|
||||
id=tool.id,
|
||||
type="function",
|
||||
function=OllamaToolFunction(
|
||||
name=tool.name, arguments=arguments, index=index
|
||||
),
|
||||
)
|
||||
)
|
||||
return tool_calls
|
||||
|
||||
|
||||
def _get_usage(
|
||||
chunk: TokenChunk | ToolCallChunk,
|
||||
) -> tuple[int | None, int | None]:
|
||||
"""Extract (prompt_eval_count, eval_count) from a chunk."""
|
||||
if chunk.usage is not None:
|
||||
return (chunk.usage.prompt_tokens, chunk.usage.completion_tokens)
|
||||
if chunk.stats is not None:
|
||||
return (chunk.stats.prompt_tokens, chunk.stats.generation_tokens)
|
||||
return (None, None)
|
||||
|
||||
|
||||
def ollama_request_to_text_generation(
|
||||
request: OllamaChatRequest,
|
||||
) -> TextGenerationTaskParams:
|
||||
"""Convert Ollama chat request to exo's internal text generation format."""
|
||||
instructions: str | None = None
|
||||
input_messages: list[InputMessage] = []
|
||||
chat_template_messages: list[dict[str, Any]] = []
|
||||
tool_message_index = 0
|
||||
|
||||
for msg in request.messages:
|
||||
content = msg.content or ""
|
||||
|
||||
if msg.role == "system":
|
||||
if instructions is None:
|
||||
instructions = content
|
||||
else:
|
||||
instructions = f"{instructions}\n{content}"
|
||||
chat_template_messages.append({"role": "system", "content": content})
|
||||
continue
|
||||
|
||||
if msg.role in ("user", "assistant") and (
|
||||
msg.content is not None or msg.thinking is not None or msg.tool_calls
|
||||
):
|
||||
input_messages.append(InputMessage(role=msg.role, content=content))
|
||||
|
||||
dumped: dict[str, Any] = {"role": msg.role, "content": content}
|
||||
if msg.thinking is not None:
|
||||
dumped["thinking"] = msg.thinking
|
||||
if msg.tool_calls is not None:
|
||||
tool_calls_list: list[dict[str, Any]] = []
|
||||
for tc in msg.tool_calls:
|
||||
function: dict[str, Any] = {
|
||||
"name": tc.function.name,
|
||||
"arguments": (
|
||||
json.dumps(tc.function.arguments)
|
||||
if isinstance(tc.function.arguments, dict)
|
||||
else tc.function.arguments
|
||||
),
|
||||
}
|
||||
if tc.function.index is not None:
|
||||
function["index"] = tc.function.index
|
||||
tool_call: dict[str, Any] = {"function": function}
|
||||
if tc.id is not None:
|
||||
tool_call["id"] = tc.id
|
||||
if tc.type is not None:
|
||||
tool_call["type"] = tc.type
|
||||
tool_calls_list.append(tool_call)
|
||||
dumped["tool_calls"] = tool_calls_list
|
||||
if msg.name is not None:
|
||||
dumped["name"] = msg.name
|
||||
if msg.role == "tool":
|
||||
tool_message_index += 1
|
||||
tool_call_id = msg.tool_name or msg.name or f"tool_{tool_message_index}"
|
||||
dumped["tool_call_id"] = tool_call_id
|
||||
if msg.tool_name is not None:
|
||||
dumped["tool_name"] = msg.tool_name
|
||||
chat_template_messages.append(dumped)
|
||||
|
||||
options = request.options
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_messages
|
||||
if input_messages
|
||||
else [InputMessage(role="user", content="")],
|
||||
instructions=instructions,
|
||||
max_output_tokens=options.num_predict if options else None,
|
||||
temperature=options.temperature if options else None,
|
||||
top_p=options.top_p if options else None,
|
||||
top_k=options.top_k if options else None,
|
||||
stop=options.stop if options else None,
|
||||
seed=options.seed if options else None,
|
||||
stream=request.stream,
|
||||
tools=request.tools,
|
||||
enable_thinking=request.think,
|
||||
chat_template_messages=chat_template_messages
|
||||
if chat_template_messages
|
||||
else None,
|
||||
)
|
||||
|
||||
|
||||
async def generate_ollama_chat_stream(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Generate streaming responses in Ollama format (newline-delimited JSON)."""
|
||||
thinking_parts: list[str] = []
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
match chunk:
|
||||
case PrefillProgressChunk():
|
||||
continue
|
||||
|
||||
case ErrorChunk():
|
||||
error_response = OllamaChatResponse(
|
||||
model=str(chunk.model),
|
||||
message=OllamaMessage(
|
||||
role="assistant", content=chunk.error_message
|
||||
),
|
||||
done=True,
|
||||
done_reason="error",
|
||||
)
|
||||
yield f"{error_response.model_dump_json(exclude_none=True)}\n"
|
||||
return
|
||||
|
||||
case ToolCallChunk():
|
||||
prompt_eval, eval_count = _get_usage(chunk)
|
||||
response = OllamaChatResponse(
|
||||
model=str(chunk.model),
|
||||
message=OllamaMessage(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls=_build_tool_calls(chunk),
|
||||
thinking="".join(thinking_parts) if thinking_parts else None,
|
||||
),
|
||||
done=True,
|
||||
done_reason="tool_call",
|
||||
prompt_eval_count=prompt_eval,
|
||||
eval_count=eval_count,
|
||||
)
|
||||
yield f"{response.model_dump_json(exclude_none=True)}\n"
|
||||
return
|
||||
|
||||
case TokenChunk():
|
||||
done = chunk.finish_reason is not None
|
||||
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
response = OllamaChatResponse(
|
||||
model=str(chunk.model),
|
||||
message=OllamaMessage(
|
||||
role="assistant", content="", thinking=chunk.text
|
||||
),
|
||||
done=False,
|
||||
)
|
||||
yield f"{response.model_dump_json(exclude_none=True)}\n"
|
||||
elif done:
|
||||
prompt_eval, eval_count = _get_usage(chunk)
|
||||
response = OllamaChatResponse(
|
||||
model=str(chunk.model),
|
||||
message=OllamaMessage(
|
||||
role="assistant",
|
||||
content=chunk.text,
|
||||
),
|
||||
done=True,
|
||||
done_reason=_map_done_reason(chunk.finish_reason),
|
||||
prompt_eval_count=prompt_eval,
|
||||
eval_count=eval_count,
|
||||
)
|
||||
yield f"{response.model_dump_json(exclude_none=True)}\n"
|
||||
else:
|
||||
response = OllamaChatResponse(
|
||||
model=str(chunk.model),
|
||||
message=OllamaMessage(role="assistant", content=chunk.text),
|
||||
done=False,
|
||||
)
|
||||
yield f"{response.model_dump_json(exclude_none=True)}\n"
|
||||
|
||||
if done:
|
||||
return
|
||||
|
||||
|
||||
async def collect_ollama_chat_response(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
) -> AsyncGenerator[str]:
|
||||
"""Collect streaming chunks into a single non-streaming Ollama response.
|
||||
|
||||
Returns an AsyncGenerator[str] (single yield) for consistency with FastAPI
|
||||
StreamingResponse cancellation handling.
|
||||
"""
|
||||
text_parts: list[str] = []
|
||||
thinking_parts: list[str] = []
|
||||
tool_calls: list[OllamaToolCall] = []
|
||||
model: str | None = None
|
||||
finish_reason: str | None = None
|
||||
prompt_eval_count: int | None = None
|
||||
eval_count: int | None = None
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
match chunk:
|
||||
case PrefillProgressChunk():
|
||||
continue
|
||||
|
||||
case ErrorChunk():
|
||||
raise ValueError(chunk.error_message or "Internal server error")
|
||||
|
||||
case TokenChunk():
|
||||
if model is None:
|
||||
model = str(chunk.model)
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
else:
|
||||
text_parts.append(chunk.text)
|
||||
if chunk.finish_reason is not None:
|
||||
finish_reason = chunk.finish_reason
|
||||
prompt_eval_count, eval_count = _get_usage(chunk)
|
||||
|
||||
case ToolCallChunk():
|
||||
if model is None:
|
||||
model = str(chunk.model)
|
||||
tool_calls.extend(_build_tool_calls(chunk))
|
||||
finish_reason = chunk.finish_reason
|
||||
prompt_eval_count, eval_count = _get_usage(chunk)
|
||||
|
||||
combined_text = "".join(text_parts)
|
||||
combined_thinking = "".join(thinking_parts) if thinking_parts else None
|
||||
assert model is not None
|
||||
|
||||
yield OllamaChatResponse(
|
||||
model=model,
|
||||
message=OllamaMessage(
|
||||
role="assistant",
|
||||
content=combined_text,
|
||||
thinking=combined_thinking,
|
||||
tool_calls=tool_calls if tool_calls else None,
|
||||
),
|
||||
done=True,
|
||||
done_reason=_map_done_reason(finish_reason),
|
||||
prompt_eval_count=prompt_eval_count,
|
||||
eval_count=eval_count,
|
||||
).model_dump_json(exclude_none=True)
|
||||
return
|
||||
|
||||
|
||||
# ── /api/generate ──
|
||||
|
||||
|
||||
def ollama_generate_request_to_text_generation(
|
||||
request: OllamaGenerateRequest,
|
||||
) -> TextGenerationTaskParams:
|
||||
"""Convert Ollama generate request to exo's internal text generation format."""
|
||||
chat_template_messages: list[dict[str, Any]] = []
|
||||
if request.system:
|
||||
chat_template_messages.append({"role": "system", "content": request.system})
|
||||
chat_template_messages.append({"role": "user", "content": request.prompt})
|
||||
|
||||
options = request.options
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=[InputMessage(role="user", content=request.prompt)],
|
||||
instructions=request.system,
|
||||
max_output_tokens=options.num_predict if options else None,
|
||||
temperature=options.temperature if options else None,
|
||||
top_p=options.top_p if options else None,
|
||||
top_k=options.top_k if options else None,
|
||||
stop=options.stop if options else None,
|
||||
seed=options.seed if options else None,
|
||||
stream=request.stream,
|
||||
enable_thinking=request.think,
|
||||
chat_template_messages=chat_template_messages
|
||||
if chat_template_messages
|
||||
else None,
|
||||
)
|
||||
|
||||
|
||||
async def generate_ollama_generate_stream(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Generate streaming responses for /api/generate in Ollama NDJSON format."""
|
||||
thinking_parts: list[str] = []
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
match chunk:
|
||||
case PrefillProgressChunk():
|
||||
continue
|
||||
|
||||
case ErrorChunk():
|
||||
resp = OllamaGenerateResponse(
|
||||
model=str(chunk.model),
|
||||
response="",
|
||||
done=True,
|
||||
done_reason="error",
|
||||
)
|
||||
yield f"{resp.model_dump_json(exclude_none=True)}\n"
|
||||
return
|
||||
|
||||
case ToolCallChunk():
|
||||
# generate endpoint doesn't support tools; emit as done
|
||||
prompt_eval, eval_count = _get_usage(chunk)
|
||||
resp = OllamaGenerateResponse(
|
||||
model=str(chunk.model),
|
||||
response="",
|
||||
done=True,
|
||||
done_reason="stop",
|
||||
prompt_eval_count=prompt_eval,
|
||||
eval_count=eval_count,
|
||||
)
|
||||
yield f"{resp.model_dump_json(exclude_none=True)}\n"
|
||||
return
|
||||
|
||||
case TokenChunk():
|
||||
done = chunk.finish_reason is not None
|
||||
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
resp = OllamaGenerateResponse(
|
||||
model=str(chunk.model),
|
||||
response="",
|
||||
thinking=chunk.text,
|
||||
done=False,
|
||||
)
|
||||
yield f"{resp.model_dump_json(exclude_none=True)}\n"
|
||||
elif done:
|
||||
prompt_eval, eval_count = _get_usage(chunk)
|
||||
resp = OllamaGenerateResponse(
|
||||
model=str(chunk.model),
|
||||
response=chunk.text,
|
||||
done=True,
|
||||
done_reason=_map_done_reason(chunk.finish_reason),
|
||||
prompt_eval_count=prompt_eval,
|
||||
eval_count=eval_count,
|
||||
)
|
||||
yield f"{resp.model_dump_json(exclude_none=True)}\n"
|
||||
else:
|
||||
resp = OllamaGenerateResponse(
|
||||
model=str(chunk.model),
|
||||
response=chunk.text,
|
||||
done=False,
|
||||
)
|
||||
yield f"{resp.model_dump_json(exclude_none=True)}\n"
|
||||
|
||||
if done:
|
||||
return
|
||||
|
||||
|
||||
async def collect_ollama_generate_response(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
) -> AsyncGenerator[str]:
|
||||
"""Collect chunks into a single non-streaming /api/generate response."""
|
||||
text_parts: list[str] = []
|
||||
thinking_parts: list[str] = []
|
||||
model: str | None = None
|
||||
finish_reason: str | None = None
|
||||
prompt_eval_count: int | None = None
|
||||
eval_count: int | None = None
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
match chunk:
|
||||
case PrefillProgressChunk():
|
||||
continue
|
||||
case ErrorChunk():
|
||||
raise ValueError(chunk.error_message or "Internal server error")
|
||||
case TokenChunk():
|
||||
if model is None:
|
||||
model = str(chunk.model)
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
else:
|
||||
text_parts.append(chunk.text)
|
||||
if chunk.finish_reason is not None:
|
||||
finish_reason = chunk.finish_reason
|
||||
prompt_eval_count, eval_count = _get_usage(chunk)
|
||||
case ToolCallChunk():
|
||||
if model is None:
|
||||
model = str(chunk.model)
|
||||
finish_reason = chunk.finish_reason
|
||||
prompt_eval_count, eval_count = _get_usage(chunk)
|
||||
|
||||
assert model is not None
|
||||
yield OllamaGenerateResponse(
|
||||
model=model,
|
||||
response="".join(text_parts),
|
||||
thinking="".join(thinking_parts) if thinking_parts else None,
|
||||
done=True,
|
||||
done_reason=_map_done_reason(finish_reason),
|
||||
prompt_eval_count=prompt_eval_count,
|
||||
eval_count=eval_count,
|
||||
).model_dump_json(exclude_none=True)
|
||||
return
|
||||
@@ -29,6 +29,12 @@ from exo.shared.types.openai_responses import (
|
||||
ResponseOutputItemAddedEvent,
|
||||
ResponseOutputItemDoneEvent,
|
||||
ResponseOutputText,
|
||||
ResponseReasoningItem,
|
||||
ResponseReasoningSummaryPartAddedEvent,
|
||||
ResponseReasoningSummaryPartDoneEvent,
|
||||
ResponseReasoningSummaryText,
|
||||
ResponseReasoningSummaryTextDeltaEvent,
|
||||
ResponseReasoningSummaryTextDoneEvent,
|
||||
ResponsesRequest,
|
||||
ResponsesResponse,
|
||||
ResponsesStreamEvent,
|
||||
@@ -36,7 +42,11 @@ from exo.shared.types.openai_responses import (
|
||||
ResponseTextDoneEvent,
|
||||
ResponseUsage,
|
||||
)
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
TextGenerationTaskParams,
|
||||
resolve_reasoning_params,
|
||||
)
|
||||
|
||||
|
||||
def _format_sse(event: ResponsesStreamEvent) -> str:
|
||||
@@ -113,6 +123,11 @@ def responses_request_to_text_generation(
|
||||
)
|
||||
built_chat_template = chat_template_messages if chat_template_messages else None
|
||||
|
||||
effort_from_reasoning = request.reasoning.effort if request.reasoning else None
|
||||
resolved_effort, resolved_thinking = resolve_reasoning_params(
|
||||
effort_from_reasoning, request.enable_thinking
|
||||
)
|
||||
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_value,
|
||||
@@ -126,6 +141,8 @@ def responses_request_to_text_generation(
|
||||
stop=request.stop,
|
||||
seed=request.seed,
|
||||
chat_template_messages=built_chat_template or request.chat_template_messages,
|
||||
reasoning_effort=resolved_effort,
|
||||
enable_thinking=resolved_thinking,
|
||||
)
|
||||
|
||||
|
||||
@@ -141,7 +158,9 @@ async def collect_responses_response(
|
||||
"""Collect all token chunks and return a single ResponsesResponse."""
|
||||
response_id = f"resp_{command_id}"
|
||||
item_id = f"item_{command_id}"
|
||||
reasoning_id = f"rs_{command_id}"
|
||||
accumulated_text = ""
|
||||
thinking_parts: list[str] = []
|
||||
function_call_items: list[ResponseFunctionCallItem] = []
|
||||
last_usage: Usage | None = None
|
||||
error_message: str | None = None
|
||||
@@ -168,6 +187,10 @@ async def collect_responses_response(
|
||||
)
|
||||
continue
|
||||
|
||||
if chunk.is_thinking:
|
||||
thinking_parts.append(chunk.text)
|
||||
continue
|
||||
|
||||
accumulated_text += chunk.text
|
||||
|
||||
if error_message is not None:
|
||||
@@ -182,13 +205,21 @@ async def collect_responses_response(
|
||||
total_tokens=last_usage.total_tokens,
|
||||
)
|
||||
|
||||
output: list[ResponseItem] = [
|
||||
output: list[ResponseItem] = []
|
||||
if thinking_parts:
|
||||
output.append(
|
||||
ResponseReasoningItem(
|
||||
id=reasoning_id,
|
||||
summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))],
|
||||
)
|
||||
)
|
||||
output.append(
|
||||
ResponseMessageItem(
|
||||
id=item_id,
|
||||
content=[ResponseOutputText(text=accumulated_text)],
|
||||
status="completed",
|
||||
)
|
||||
]
|
||||
)
|
||||
output.extend(function_call_items)
|
||||
|
||||
yield ResponsesResponse(
|
||||
@@ -212,6 +243,7 @@ async def generate_responses_stream(
|
||||
"""Generate OpenAI Responses API streaming events from TokenChunks."""
|
||||
response_id = f"resp_{command_id}"
|
||||
item_id = f"item_{command_id}"
|
||||
reasoning_id = f"rs_{command_id}"
|
||||
seq = count(1)
|
||||
|
||||
# response.created
|
||||
@@ -233,32 +265,17 @@ async def generate_responses_stream(
|
||||
)
|
||||
yield _format_sse(in_progress_event)
|
||||
|
||||
# response.output_item.added
|
||||
initial_item = ResponseMessageItem(
|
||||
id=item_id,
|
||||
content=[ResponseOutputText(text="")],
|
||||
status="in_progress",
|
||||
)
|
||||
item_added = ResponseOutputItemAddedEvent(
|
||||
sequence_number=next(seq), output_index=0, item=initial_item
|
||||
)
|
||||
yield _format_sse(item_added)
|
||||
|
||||
# response.content_part.added
|
||||
initial_part = ResponseOutputText(text="")
|
||||
part_added = ResponseContentPartAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
part=initial_part,
|
||||
)
|
||||
yield _format_sse(part_added)
|
||||
|
||||
accumulated_text = ""
|
||||
accumulated_thinking = ""
|
||||
function_call_items: list[ResponseFunctionCallItem] = []
|
||||
last_usage: Usage | None = None
|
||||
next_output_index = 1 # message item is at 0
|
||||
next_output_index = 0
|
||||
|
||||
# Track dynamic block creation
|
||||
reasoning_started = False
|
||||
reasoning_output_index = -1
|
||||
message_started = False
|
||||
message_output_index = -1
|
||||
|
||||
async for chunk in chunk_stream:
|
||||
if isinstance(chunk, PrefillProgressChunk):
|
||||
@@ -327,23 +344,184 @@ async def generate_responses_stream(
|
||||
next_output_index += 1
|
||||
continue
|
||||
|
||||
if chunk.is_thinking:
|
||||
# Start reasoning block on first thinking token
|
||||
if not reasoning_started:
|
||||
reasoning_started = True
|
||||
reasoning_output_index = next_output_index
|
||||
next_output_index += 1
|
||||
|
||||
# response.output_item.added for reasoning
|
||||
reasoning_item = ResponseReasoningItem(
|
||||
id=reasoning_id,
|
||||
summary=[],
|
||||
status="in_progress",
|
||||
)
|
||||
rs_added = ResponseOutputItemAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
output_index=reasoning_output_index,
|
||||
item=reasoning_item,
|
||||
)
|
||||
yield _format_sse(rs_added)
|
||||
|
||||
# response.reasoning_summary_part.added
|
||||
part_added = ResponseReasoningSummaryPartAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
part=ResponseReasoningSummaryText(text=""),
|
||||
)
|
||||
yield _format_sse(part_added)
|
||||
|
||||
accumulated_thinking += chunk.text
|
||||
|
||||
# response.reasoning_summary_text.delta
|
||||
rs_delta = ResponseReasoningSummaryTextDeltaEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
delta=chunk.text,
|
||||
)
|
||||
yield _format_sse(rs_delta)
|
||||
continue
|
||||
|
||||
# Close reasoning block when transitioning to text
|
||||
if reasoning_started and not message_started:
|
||||
# response.reasoning_summary_text.done
|
||||
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
text=accumulated_thinking,
|
||||
)
|
||||
yield _format_sse(rs_text_done)
|
||||
|
||||
# response.reasoning_summary_part.done
|
||||
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
part=ResponseReasoningSummaryText(text=accumulated_thinking),
|
||||
)
|
||||
yield _format_sse(rs_part_done)
|
||||
|
||||
# response.output_item.done for reasoning
|
||||
rs_item_done = ResponseOutputItemDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
output_index=reasoning_output_index,
|
||||
item=ResponseReasoningItem(
|
||||
id=reasoning_id,
|
||||
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
|
||||
),
|
||||
)
|
||||
yield _format_sse(rs_item_done)
|
||||
|
||||
# Start message block on first text token
|
||||
if not message_started:
|
||||
message_started = True
|
||||
message_output_index = next_output_index
|
||||
next_output_index += 1
|
||||
|
||||
initial_item = ResponseMessageItem(
|
||||
id=item_id,
|
||||
content=[ResponseOutputText(text="")],
|
||||
status="in_progress",
|
||||
)
|
||||
item_added = ResponseOutputItemAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
output_index=message_output_index,
|
||||
item=initial_item,
|
||||
)
|
||||
yield _format_sse(item_added)
|
||||
|
||||
initial_part = ResponseOutputText(text="")
|
||||
part_added = ResponseContentPartAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=message_output_index,
|
||||
content_index=0,
|
||||
part=initial_part,
|
||||
)
|
||||
yield _format_sse(part_added)
|
||||
|
||||
accumulated_text += chunk.text
|
||||
|
||||
# response.output_text.delta
|
||||
delta_event = ResponseTextDeltaEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=0,
|
||||
output_index=message_output_index,
|
||||
content_index=0,
|
||||
delta=chunk.text,
|
||||
)
|
||||
yield _format_sse(delta_event)
|
||||
|
||||
# Close reasoning block if it was never followed by text
|
||||
if reasoning_started and not message_started:
|
||||
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
text=accumulated_thinking,
|
||||
)
|
||||
yield _format_sse(rs_text_done)
|
||||
|
||||
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=reasoning_id,
|
||||
output_index=reasoning_output_index,
|
||||
summary_index=0,
|
||||
part=ResponseReasoningSummaryText(text=accumulated_thinking),
|
||||
)
|
||||
yield _format_sse(rs_part_done)
|
||||
|
||||
rs_item_done = ResponseOutputItemDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
output_index=reasoning_output_index,
|
||||
item=ResponseReasoningItem(
|
||||
id=reasoning_id,
|
||||
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
|
||||
),
|
||||
)
|
||||
yield _format_sse(rs_item_done)
|
||||
|
||||
# If no message block was started, create one now (empty text)
|
||||
if not message_started:
|
||||
message_output_index = next_output_index
|
||||
next_output_index += 1
|
||||
|
||||
initial_item = ResponseMessageItem(
|
||||
id=item_id,
|
||||
content=[ResponseOutputText(text="")],
|
||||
status="in_progress",
|
||||
)
|
||||
item_added = ResponseOutputItemAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
output_index=message_output_index,
|
||||
item=initial_item,
|
||||
)
|
||||
yield _format_sse(item_added)
|
||||
|
||||
initial_part = ResponseOutputText(text="")
|
||||
part_added_evt = ResponseContentPartAddedEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=message_output_index,
|
||||
content_index=0,
|
||||
part=initial_part,
|
||||
)
|
||||
yield _format_sse(part_added_evt)
|
||||
|
||||
# response.output_text.done
|
||||
text_done = ResponseTextDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=0,
|
||||
output_index=message_output_index,
|
||||
content_index=0,
|
||||
text=accumulated_text,
|
||||
)
|
||||
@@ -354,7 +532,7 @@ async def generate_responses_stream(
|
||||
part_done = ResponseContentPartDoneEvent(
|
||||
sequence_number=next(seq),
|
||||
item_id=item_id,
|
||||
output_index=0,
|
||||
output_index=message_output_index,
|
||||
content_index=0,
|
||||
part=final_part,
|
||||
)
|
||||
@@ -367,7 +545,9 @@ async def generate_responses_stream(
|
||||
status="completed",
|
||||
)
|
||||
item_done = ResponseOutputItemDoneEvent(
|
||||
sequence_number=next(seq), output_index=0, item=final_message_item
|
||||
sequence_number=next(seq),
|
||||
output_index=message_output_index,
|
||||
item=final_message_item,
|
||||
)
|
||||
yield _format_sse(item_done)
|
||||
|
||||
@@ -381,7 +561,15 @@ async def generate_responses_stream(
|
||||
)
|
||||
|
||||
# response.completed
|
||||
output: list[ResponseItem] = [final_message_item]
|
||||
output: list[ResponseItem] = []
|
||||
if reasoning_started:
|
||||
output.append(
|
||||
ResponseReasoningItem(
|
||||
id=reasoning_id,
|
||||
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
|
||||
)
|
||||
)
|
||||
output.append(final_message_item)
|
||||
output.extend(function_call_items)
|
||||
final_response = ResponsesResponse(
|
||||
id=response_id,
|
||||
|
||||
+291
-72
@@ -11,8 +11,7 @@ from typing import Annotated, Literal, cast
|
||||
from uuid import uuid4
|
||||
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, create_task_group
|
||||
from anyio.abc import TaskGroup
|
||||
from anyio import BrokenResourceError
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||||
@@ -32,6 +31,14 @@ from exo.master.adapters.claude import (
|
||||
collect_claude_response,
|
||||
generate_claude_stream,
|
||||
)
|
||||
from exo.master.adapters.ollama import (
|
||||
collect_ollama_chat_response,
|
||||
collect_ollama_generate_response,
|
||||
generate_ollama_chat_stream,
|
||||
generate_ollama_generate_stream,
|
||||
ollama_generate_request_to_text_generation,
|
||||
ollama_request_to_text_generation,
|
||||
)
|
||||
from exo.master.adapters.responses import (
|
||||
collect_responses_response,
|
||||
generate_responses_stream,
|
||||
@@ -43,6 +50,7 @@ from exo.master.placement import place_instance as get_instance_placements
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.constants import (
|
||||
DASHBOARD_DIR,
|
||||
EXO_CACHE_HOME,
|
||||
EXO_EVENT_LOG_DIR,
|
||||
EXO_IMAGE_CACHE_DIR,
|
||||
EXO_MAX_CHUNK_SIZE,
|
||||
@@ -65,6 +73,7 @@ from exo.shared.types.api import (
|
||||
BenchChatCompletionResponse,
|
||||
BenchImageGenerationResponse,
|
||||
BenchImageGenerationTaskParams,
|
||||
CancelCommandResponse,
|
||||
ChatCompletionChoice,
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionRequest,
|
||||
@@ -73,6 +82,8 @@ from exo.shared.types.api import (
|
||||
CreateInstanceResponse,
|
||||
DeleteDownloadResponse,
|
||||
DeleteInstanceResponse,
|
||||
DeleteTracesRequest,
|
||||
DeleteTracesResponse,
|
||||
ErrorInfo,
|
||||
ErrorResponse,
|
||||
FinishReason,
|
||||
@@ -132,16 +143,27 @@ from exo.shared.types.commands import (
|
||||
TaskFinished,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
Event,
|
||||
ForwarderEvent,
|
||||
IndexedEvent,
|
||||
PrefillProgress,
|
||||
TracesMerged,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.ollama_api import (
|
||||
OllamaChatRequest,
|
||||
OllamaChatResponse,
|
||||
OllamaGenerateRequest,
|
||||
OllamaGenerateResponse,
|
||||
OllamaModelDetails,
|
||||
OllamaModelTag,
|
||||
OllamaPsModel,
|
||||
OllamaPsResponse,
|
||||
OllamaShowRequest,
|
||||
OllamaShowResponse,
|
||||
OllamaTagsResponse,
|
||||
)
|
||||
from exo.shared.types.openai_responses import (
|
||||
ResponsesRequest,
|
||||
ResponsesResponse,
|
||||
@@ -152,9 +174,10 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
|
||||
from exo.shared.types.worker.shards import Sharding
|
||||
from exo.utils.banner import print_startup_banner
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.event_buffer import OrderedBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
|
||||
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
|
||||
|
||||
|
||||
def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
|
||||
@@ -174,11 +197,9 @@ class API:
|
||||
def __init__(
|
||||
self,
|
||||
node_id: NodeId,
|
||||
session_id: SessionId,
|
||||
*,
|
||||
port: int,
|
||||
# Ideally this would be a MasterForwarderEvent but type system says no :(
|
||||
global_event_receiver: Receiver[ForwarderEvent],
|
||||
event_receiver: Receiver[IndexedEvent],
|
||||
command_sender: Sender[ForwarderCommand],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
# This lets us pause the API if an election is running
|
||||
@@ -186,13 +207,12 @@ class API:
|
||||
) -> None:
|
||||
self.state = State()
|
||||
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
|
||||
self._system_id = SystemId()
|
||||
self.command_sender = command_sender
|
||||
self.download_command_sender = download_command_sender
|
||||
self.global_event_receiver = global_event_receiver
|
||||
self.event_receiver = event_receiver
|
||||
self.election_receiver = election_receiver
|
||||
self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
|
||||
self.node_id: NodeId = node_id
|
||||
self.session_id: SessionId = session_id
|
||||
self.last_completed_election: int = 0
|
||||
self.port = port
|
||||
|
||||
@@ -230,18 +250,20 @@ class API:
|
||||
CommandId, Sender[ImageChunk | ErrorChunk]
|
||||
] = {}
|
||||
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
|
||||
self._tg: TaskGroup | None = None
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
|
||||
def reset(self, new_session_id: SessionId, result_clock: int):
|
||||
def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
|
||||
logger.info("Resetting API State")
|
||||
self._event_log.close()
|
||||
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
|
||||
self.state = State()
|
||||
self.session_id = new_session_id
|
||||
self.event_buffer = OrderedBuffer[Event]()
|
||||
self._system_id = SystemId()
|
||||
self._text_generation_queues = {}
|
||||
self._image_generation_queues = {}
|
||||
self.unpause(result_clock)
|
||||
self.event_receiver.close()
|
||||
self.event_receiver = event_receiver
|
||||
self._tg.start_soon(self._apply_state)
|
||||
|
||||
def unpause(self, result_clock: int):
|
||||
logger.info("Unpausing API")
|
||||
@@ -301,14 +323,33 @@ class API:
|
||||
self.app.get("/images/{image_id}")(self.get_image)
|
||||
self.app.post("/v1/messages", response_model=None)(self.claude_messages)
|
||||
self.app.post("/v1/responses", response_model=None)(self.openai_responses)
|
||||
self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
|
||||
|
||||
# Ollama API
|
||||
self.app.head("/ollama/")(self.ollama_version)
|
||||
self.app.head("/ollama/api/version")(self.ollama_version)
|
||||
self.app.post("/ollama/api/chat", response_model=None)(self.ollama_chat)
|
||||
self.app.post("/ollama/api/api/chat", response_model=None)(self.ollama_chat)
|
||||
self.app.post("/ollama/api/v1/chat", response_model=None)(self.ollama_chat)
|
||||
self.app.post("/ollama/api/generate", response_model=None)(self.ollama_generate)
|
||||
self.app.get("/ollama/api/tags")(self.ollama_tags)
|
||||
self.app.get("/ollama/api/api/tags")(self.ollama_tags)
|
||||
self.app.get("/ollama/api/v1/tags")(self.ollama_tags)
|
||||
self.app.post("/ollama/api/show")(self.ollama_show)
|
||||
self.app.get("/ollama/api/ps")(self.ollama_ps)
|
||||
self.app.get("/ollama/api/version")(self.ollama_version)
|
||||
|
||||
self.app.get("/state")(lambda: self.state)
|
||||
self.app.get("/events")(self.stream_events)
|
||||
self.app.post("/download/start")(self.start_download)
|
||||
self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
|
||||
self.app.get("/v1/traces")(self.list_traces)
|
||||
self.app.post("/v1/traces/delete")(self.delete_traces)
|
||||
self.app.get("/v1/traces/{task_id}")(self.get_trace)
|
||||
self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
|
||||
self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
|
||||
self.app.get("/onboarding")(self.get_onboarding)
|
||||
self.app.post("/onboarding")(self.complete_onboarding)
|
||||
|
||||
async def place_instance(self, payload: PlaceInstanceParams):
|
||||
command = PlaceInstance(
|
||||
@@ -528,6 +569,25 @@ class API:
|
||||
instance_id=instance_id,
|
||||
)
|
||||
|
||||
async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
|
||||
"""Cancel an active command by closing its stream and notifying workers."""
|
||||
sender = self._text_generation_queues.get(
|
||||
command_id
|
||||
) or self._image_generation_queues.get(command_id)
|
||||
if sender is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Command not found or already completed",
|
||||
)
|
||||
|
||||
await self._send(TaskCancelled(cancelled_command_id=command_id))
|
||||
sender.close()
|
||||
|
||||
return CancelCommandResponse(
|
||||
message="Command cancelled.",
|
||||
command_id=command_id,
|
||||
)
|
||||
|
||||
async def _token_chunk_stream(
|
||||
self, command_id: CommandId
|
||||
) -> AsyncGenerator[
|
||||
@@ -554,7 +614,7 @@ class API:
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self.node_id, command=command)
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
@@ -902,7 +962,7 @@ class API:
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self.node_id, command=command)
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
@@ -988,7 +1048,7 @@ class API:
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self.node_id, command=command)
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
@@ -1294,6 +1354,163 @@ class API:
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
async def _ollama_root(self) -> JSONResponse:
|
||||
"""Respond to HEAD / from Ollama CLI connectivity checks."""
|
||||
return JSONResponse(content="Ollama is running")
|
||||
|
||||
async def ollama_chat(
|
||||
self, request: Request
|
||||
) -> OllamaChatResponse | StreamingResponse:
|
||||
"""Ollama Chat API — accepts JSON regardless of Content-Type."""
|
||||
body = await request.body()
|
||||
payload = OllamaChatRequest.model_validate_json(body)
|
||||
task_params = ollama_request_to_text_generation(payload)
|
||||
resolved_model = await self._resolve_and_validate_text_model(
|
||||
ModelId(task_params.model)
|
||||
)
|
||||
task_params = task_params.model_copy(update={"model": resolved_model})
|
||||
|
||||
command = TextGeneration(task_params=task_params)
|
||||
await self._send(command)
|
||||
|
||||
if payload.stream:
|
||||
return StreamingResponse(
|
||||
generate_ollama_chat_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "close",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
else:
|
||||
return StreamingResponse(
|
||||
collect_ollama_chat_response(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
async def ollama_generate(
|
||||
self, request: Request
|
||||
) -> OllamaGenerateResponse | StreamingResponse:
|
||||
"""Ollama Generate API — accepts JSON regardless of Content-Type."""
|
||||
body = await request.body()
|
||||
payload = OllamaGenerateRequest.model_validate_json(body)
|
||||
task_params = ollama_generate_request_to_text_generation(payload)
|
||||
resolved_model = await self._resolve_and_validate_text_model(
|
||||
ModelId(task_params.model)
|
||||
)
|
||||
task_params = task_params.model_copy(update={"model": resolved_model})
|
||||
|
||||
command = TextGeneration(task_params=task_params)
|
||||
await self._send(command)
|
||||
|
||||
if payload.stream:
|
||||
return StreamingResponse(
|
||||
generate_ollama_generate_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "close",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
else:
|
||||
return StreamingResponse(
|
||||
collect_ollama_generate_response(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
|
||||
async def ollama_tags(self) -> OllamaTagsResponse:
|
||||
"""Returns list of models in Ollama tags format. We return the downloaded ones only."""
|
||||
|
||||
def none_if_empty(value: str) -> str | None:
|
||||
return value or None
|
||||
|
||||
downloaded_model_ids: set[str] = set()
|
||||
for node_downloads in self.state.downloads.values():
|
||||
for dl in node_downloads:
|
||||
if isinstance(dl, DownloadCompleted):
|
||||
downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
|
||||
|
||||
cards = [
|
||||
c for c in await get_model_cards() if c.model_id in downloaded_model_ids
|
||||
]
|
||||
|
||||
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
return OllamaTagsResponse(
|
||||
models=[
|
||||
OllamaModelTag(
|
||||
name=str(card.model_id),
|
||||
model=str(card.model_id),
|
||||
modified_at=now,
|
||||
size=card.storage_size.in_bytes,
|
||||
digest="sha256:000000000000",
|
||||
details=OllamaModelDetails(
|
||||
family=none_if_empty(card.family),
|
||||
quantization_level=none_if_empty(card.quantization),
|
||||
),
|
||||
)
|
||||
for card in cards
|
||||
]
|
||||
)
|
||||
|
||||
async def ollama_show(self, request: Request) -> OllamaShowResponse:
|
||||
"""Returns model information in Ollama show format."""
|
||||
body = await request.body()
|
||||
payload = OllamaShowRequest.model_validate_json(body)
|
||||
model_name = payload.name or payload.model
|
||||
if not model_name:
|
||||
raise HTTPException(status_code=400, detail="name or model is required")
|
||||
try:
|
||||
card = await ModelCard.load(ModelId(model_name))
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Model not found: {model_name}"
|
||||
) from exc
|
||||
|
||||
return OllamaShowResponse(
|
||||
modelfile=f"FROM {card.model_id}",
|
||||
template="{{ .Prompt }}",
|
||||
details=OllamaModelDetails(
|
||||
family=card.family or None,
|
||||
quantization_level=card.quantization or None,
|
||||
),
|
||||
)
|
||||
|
||||
async def ollama_ps(self) -> OllamaPsResponse:
|
||||
"""Returns list of running models (active instances)."""
|
||||
models: list[OllamaPsModel] = []
|
||||
seen: set[str] = set()
|
||||
for instance in self.state.instances.values():
|
||||
model_id = str(instance.shard_assignments.model_id)
|
||||
if model_id in seen:
|
||||
continue
|
||||
seen.add(model_id)
|
||||
models.append(
|
||||
OllamaPsModel(
|
||||
name=model_id,
|
||||
model=model_id,
|
||||
size=0,
|
||||
)
|
||||
)
|
||||
return OllamaPsResponse(models=models)
|
||||
|
||||
async def ollama_version(self) -> dict[str, str]:
|
||||
"""Returns version information for Ollama API compatibility."""
|
||||
return {"version": "exo v1.0"}
|
||||
|
||||
def _calculate_total_available_memory(self) -> Memory:
|
||||
"""Calculate total available memory across all nodes in bytes."""
|
||||
total_available = Memory()
|
||||
@@ -1323,7 +1540,7 @@ class API:
|
||||
name=card.model_id.short(),
|
||||
description="",
|
||||
tags=[],
|
||||
storage_size_megabytes=int(card.storage_size.in_mb),
|
||||
storage_size_megabytes=card.storage_size.in_mb,
|
||||
supports_tensor=card.supports_tensor,
|
||||
tasks=[task.value for task in card.tasks],
|
||||
is_custom=is_custom_card(card.model_id),
|
||||
@@ -1394,8 +1611,7 @@ class API:
|
||||
shutdown_ev = anyio.Event()
|
||||
|
||||
try:
|
||||
async with create_task_group() as tg:
|
||||
self._tg = tg
|
||||
async with self._tg as tg:
|
||||
logger.info("Starting API")
|
||||
tg.start_soon(self._apply_state)
|
||||
tg.start_soon(self._pause_on_new_election)
|
||||
@@ -1410,7 +1626,7 @@ class API:
|
||||
finally:
|
||||
self._event_log.close()
|
||||
self.command_sender.close()
|
||||
self.global_event_receiver.close()
|
||||
self.event_receiver.close()
|
||||
|
||||
async def run_api(self, ev: anyio.Event):
|
||||
cfg = Config()
|
||||
@@ -1427,52 +1643,31 @@ class API:
|
||||
)
|
||||
|
||||
async def _apply_state(self):
|
||||
with self.global_event_receiver as events:
|
||||
async for f_event in events:
|
||||
if f_event.origin != self.session_id.master_node_id:
|
||||
continue
|
||||
self.event_buffer.ingest(f_event.origin_idx, f_event.event)
|
||||
for idx, event in self.event_buffer.drain_indexed():
|
||||
self._event_log.append(event)
|
||||
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
|
||||
with self.event_receiver as events:
|
||||
async for i_event in events:
|
||||
self._event_log.append(i_event.event)
|
||||
self.state = apply(self.state, i_event)
|
||||
event = i_event.event
|
||||
|
||||
if isinstance(event, ChunkGenerated):
|
||||
if queue := self._image_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._image_generation_queues.pop(
|
||||
event.command_id, None
|
||||
)
|
||||
if queue := self._text_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert not isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
|
||||
elif isinstance(event, PrefillProgress):
|
||||
if queue := self._text_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
try:
|
||||
await queue.send(
|
||||
PrefillProgressChunk(
|
||||
model=event.model,
|
||||
processed_tokens=event.processed_tokens,
|
||||
total_tokens=event.total_tokens,
|
||||
)
|
||||
)
|
||||
except BrokenResourceError:
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
|
||||
if isinstance(event, TracesMerged):
|
||||
self._save_merged_trace(event)
|
||||
if isinstance(event, ChunkGenerated):
|
||||
if queue := self._image_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._image_generation_queues.pop(event.command_id, None)
|
||||
if queue := self._text_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert not isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
if isinstance(event, TracesMerged):
|
||||
self._save_merged_trace(event)
|
||||
|
||||
def _save_merged_trace(self, event: TracesMerged) -> None:
|
||||
traces = [
|
||||
@@ -1508,12 +1703,12 @@ class API:
|
||||
while self.paused:
|
||||
await self.paused_ev.wait()
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self.node_id, command=command)
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
|
||||
async def _send_download(self, command: DownloadCommand):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(origin=self.node_id, command=command)
|
||||
ForwarderDownloadCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
|
||||
async def start_download(
|
||||
@@ -1536,8 +1731,12 @@ class API:
|
||||
await self._send_download(command)
|
||||
return DeleteDownloadResponse(command_id=command.command_id)
|
||||
|
||||
def _get_trace_path(self, task_id: str) -> Path:
|
||||
return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
|
||||
@staticmethod
|
||||
def _get_trace_path(task_id: str) -> Path:
|
||||
trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
|
||||
if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
|
||||
raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
|
||||
return trace_path
|
||||
|
||||
async def list_traces(self) -> TraceListResponse:
|
||||
traces: list[TraceListItem] = []
|
||||
@@ -1635,3 +1834,23 @@ class API:
|
||||
media_type="application/json",
|
||||
filename=f"trace_{task_id}.json",
|
||||
)
|
||||
|
||||
async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
|
||||
deleted: list[str] = []
|
||||
not_found: list[str] = []
|
||||
for task_id in request.task_ids:
|
||||
trace_path = self._get_trace_path(task_id)
|
||||
if trace_path.exists():
|
||||
trace_path.unlink()
|
||||
deleted.append(task_id)
|
||||
else:
|
||||
not_found.append(task_id)
|
||||
return DeleteTracesResponse(deleted=deleted, not_found=not_found)
|
||||
|
||||
async def get_onboarding(self) -> JSONResponse:
|
||||
return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
|
||||
|
||||
async def complete_onboarding(self) -> JSONResponse:
|
||||
ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
ONBOARDING_COMPLETE_FILE.write_text("true")
|
||||
return JSONResponse({"completed": True})
|
||||
|
||||
+31
-47
@@ -1,7 +1,6 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
|
||||
from exo.master.event_log import DiskEventLog
|
||||
@@ -29,13 +28,14 @@ from exo.shared.types.commands import (
|
||||
TestCommand,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.common import CommandId, NodeId, SessionId
|
||||
from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
ForwarderEvent,
|
||||
GlobalForwarderEvent,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceDeleted,
|
||||
LocalForwarderEvent,
|
||||
NodeGatheredInfo,
|
||||
NodeTimedOut,
|
||||
TaskCreated,
|
||||
@@ -60,8 +60,9 @@ from exo.shared.types.tasks import (
|
||||
TextGeneration as TextGenerationTask,
|
||||
)
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.event_buffer import MultiSourceBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
class Master:
|
||||
@@ -71,26 +72,23 @@ class Master:
|
||||
session_id: SessionId,
|
||||
*,
|
||||
command_receiver: Receiver[ForwarderCommand],
|
||||
local_event_receiver: Receiver[ForwarderEvent],
|
||||
global_event_sender: Sender[ForwarderEvent],
|
||||
event_sender: Sender[Event],
|
||||
local_event_receiver: Receiver[LocalForwarderEvent],
|
||||
global_event_sender: Sender[GlobalForwarderEvent],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
):
|
||||
self.state = State()
|
||||
self._tg: TaskGroup = anyio.create_task_group()
|
||||
self.node_id = node_id
|
||||
self.session_id = session_id
|
||||
self.state = State()
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
self.command_task_mapping: dict[CommandId, TaskId] = {}
|
||||
self.command_receiver = command_receiver
|
||||
self.local_event_receiver = local_event_receiver
|
||||
self.global_event_sender = global_event_sender
|
||||
self.download_command_sender = download_command_sender
|
||||
send, recv = channel[Event]()
|
||||
self.event_sender: Sender[Event] = send
|
||||
self._loopback_event_receiver: Receiver[Event] = recv
|
||||
self._loopback_event_sender: Sender[ForwarderEvent] = (
|
||||
local_event_receiver.clone_sender()
|
||||
)
|
||||
self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
|
||||
self.event_sender = event_sender
|
||||
self._system_id = SystemId()
|
||||
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
|
||||
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
|
||||
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
|
||||
self._expected_ranks: dict[TaskId, set[int]] = {}
|
||||
@@ -102,19 +100,16 @@ class Master:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._event_processor)
|
||||
tg.start_soon(self._command_processor)
|
||||
tg.start_soon(self._loopback_processor)
|
||||
tg.start_soon(self._plan)
|
||||
finally:
|
||||
self._event_log.close()
|
||||
self.global_event_sender.close()
|
||||
self.local_event_receiver.close()
|
||||
self.command_receiver.close()
|
||||
self._loopback_event_sender.close()
|
||||
self._loopback_event_receiver.close()
|
||||
|
||||
async def shutdown(self):
|
||||
logger.info("Stopping Master")
|
||||
self._tg.cancel_scope.cancel()
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _command_processor(self) -> None:
|
||||
with self.command_receiver as commands:
|
||||
@@ -288,7 +283,7 @@ class Master:
|
||||
):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=self.node_id, command=cmd
|
||||
origin=self._system_id, command=cmd
|
||||
)
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
@@ -333,17 +328,22 @@ class Master:
|
||||
task_id=task_id,
|
||||
)
|
||||
)
|
||||
case TaskFinished():
|
||||
generated_events.append(
|
||||
TaskDeleted(
|
||||
task_id=self.command_task_mapping[
|
||||
command.finished_command_id
|
||||
]
|
||||
else:
|
||||
logger.warning(
|
||||
f"Nonexistent command {command.cancelled_command_id} cancelled"
|
||||
)
|
||||
)
|
||||
self.command_task_mapping.pop(
|
||||
command.finished_command_id, None
|
||||
)
|
||||
case TaskFinished():
|
||||
if (
|
||||
task_id := self.command_task_mapping.pop(
|
||||
command.finished_command_id, None
|
||||
)
|
||||
) is not None:
|
||||
generated_events.append(TaskDeleted(task_id=task_id))
|
||||
else:
|
||||
logger.warning(
|
||||
f"Finished command {command.finished_command_id} finished"
|
||||
)
|
||||
|
||||
case RequestEventLog():
|
||||
# We should just be able to send everything, since other buffers will ignore old messages
|
||||
# rate limit to 1000 at a time
|
||||
@@ -407,27 +407,11 @@ class Master:
|
||||
self._event_log.append(event)
|
||||
await self._send_event(indexed)
|
||||
|
||||
async def _loopback_processor(self) -> None:
|
||||
# this would ideally not be necessary.
|
||||
# this is WAY less hacky than how I was working around this before
|
||||
local_index = 0
|
||||
with self._loopback_event_receiver as events:
|
||||
async for event in events:
|
||||
await self._loopback_event_sender.send(
|
||||
ForwarderEvent(
|
||||
origin=NodeId(f"master_{self.node_id}"),
|
||||
origin_idx=local_index,
|
||||
session=self.session_id,
|
||||
event=event,
|
||||
)
|
||||
)
|
||||
local_index += 1
|
||||
|
||||
# This function is re-entrant, take care!
|
||||
async def _send_event(self, event: IndexedEvent):
|
||||
# Convenience method since this line is ugly
|
||||
await self.global_event_sender.send(
|
||||
ForwarderEvent(
|
||||
GlobalForwarderEvent(
|
||||
origin=self.node_id,
|
||||
origin_idx=event.idx,
|
||||
session=self.session_id,
|
||||
|
||||
@@ -112,7 +112,11 @@ def place_instance(
|
||||
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
|
||||
]
|
||||
|
||||
if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
|
||||
if command.instance_meta == InstanceMeta.MlxJaccl:
|
||||
if not smallest_rdma_cycles:
|
||||
raise ValueError(
|
||||
"Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
|
||||
)
|
||||
smallest_cycles = smallest_rdma_cycles
|
||||
|
||||
cycles_with_leaf_nodes: list[Cycle] = [
|
||||
@@ -129,6 +133,11 @@ def place_instance(
|
||||
),
|
||||
)
|
||||
|
||||
# Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
|
||||
if len(selected_cycle) == 1:
|
||||
command.instance_meta = InstanceMeta.MlxRing
|
||||
command.sharding = Sharding.Pipeline
|
||||
|
||||
shard_assignments = get_shard_assignments(
|
||||
command.model_card, selected_cycle, command.sharding, node_memory
|
||||
)
|
||||
@@ -138,18 +147,29 @@ def place_instance(
|
||||
instance_id = InstanceId()
|
||||
target_instances = dict(deepcopy(current_instances))
|
||||
|
||||
if len(selected_cycle) == 1:
|
||||
command.instance_meta = InstanceMeta.MlxRing
|
||||
|
||||
# TODO: Single node instances
|
||||
match command.instance_meta:
|
||||
case InstanceMeta.MlxJaccl:
|
||||
# TODO(evan): shard assignments should contain information about ranks, this is ugly
|
||||
def get_device_rank(node_id: NodeId) -> int:
|
||||
runner_id = shard_assignments.node_to_runner[node_id]
|
||||
shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
|
||||
assert shard_metadata is not None
|
||||
return shard_metadata.device_rank
|
||||
|
||||
zero_node_ids = [
|
||||
node_id
|
||||
for node_id in selected_cycle.node_ids
|
||||
if get_device_rank(node_id) == 0
|
||||
]
|
||||
assert len(zero_node_ids) == 1
|
||||
coordinator_node_id = zero_node_ids[0]
|
||||
|
||||
mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
|
||||
[node_id for node_id in selected_cycle],
|
||||
cycle_digraph,
|
||||
)
|
||||
mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
|
||||
coordinator=selected_cycle.node_ids[0],
|
||||
coordinator=coordinator_node_id,
|
||||
coordinator_port=random_ephemeral_port(),
|
||||
cycle_digraph=cycle_digraph,
|
||||
node_network=node_network,
|
||||
|
||||
@@ -102,22 +102,21 @@ def _allocate_and_validate_layers(
|
||||
layer_allocations = allocate_layers_proportionally(
|
||||
total_layers=model_card.n_layers,
|
||||
memory_fractions=[
|
||||
node_memory[node_id].ram_available.in_bytes / total_memory.in_bytes
|
||||
for node_id in node_ids
|
||||
node_memory[node_id].ram_available / total_memory for node_id in node_ids
|
||||
],
|
||||
)
|
||||
|
||||
total_storage_bytes = model_card.storage_size.in_bytes
|
||||
total_storage = model_card.storage_size
|
||||
total_layers = model_card.n_layers
|
||||
for i, node_id in enumerate(node_ids):
|
||||
node_layers = layer_allocations[i]
|
||||
required_memory = (total_storage_bytes * node_layers) // total_layers
|
||||
available_memory = node_memory[node_id].ram_available.in_bytes
|
||||
required_memory = (total_storage * node_layers) // total_layers
|
||||
available_memory = node_memory[node_id].ram_available
|
||||
if required_memory > available_memory:
|
||||
raise ValueError(
|
||||
f"Node {i} ({node_id}) has insufficient memory: "
|
||||
f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
|
||||
f"but only has {available_memory / (1024**3):.2f} GB available"
|
||||
f"requires {required_memory.in_gb:.2f} GB for {node_layers} layers, "
|
||||
f"but only has {available_memory.in_gb:.2f} GB available"
|
||||
)
|
||||
|
||||
return layer_allocations
|
||||
@@ -342,6 +341,7 @@ def _find_ip_prioritised(
|
||||
other_node_id: NodeId,
|
||||
cycle_digraph: Topology,
|
||||
node_network: Mapping[NodeId, NodeNetworkInfo],
|
||||
ring: bool,
|
||||
) -> str | None:
|
||||
"""Find an IP address between nodes with prioritization.
|
||||
|
||||
@@ -354,13 +354,27 @@ def _find_ip_prioritised(
|
||||
ip_to_type = {
|
||||
iface.ip_address: iface.interface_type for iface in other_network.interfaces
|
||||
}
|
||||
priority = {
|
||||
"ethernet": 0,
|
||||
"wifi": 1,
|
||||
"unknown": 2,
|
||||
"maybe_ethernet": 3,
|
||||
"thunderbolt": 4,
|
||||
}
|
||||
|
||||
# Ring should prioritise fastest connection. As a best-effort, we prioritise TB.
|
||||
# TODO: Profile and get actual connection speeds.
|
||||
if ring:
|
||||
priority = {
|
||||
"thunderbolt": 0,
|
||||
"maybe_ethernet": 1,
|
||||
"ethernet": 2,
|
||||
"wifi": 3,
|
||||
"unknown": 4,
|
||||
}
|
||||
|
||||
# RDMA prefers ethernet coordinator
|
||||
else:
|
||||
priority = {
|
||||
"ethernet": 0,
|
||||
"wifi": 1,
|
||||
"unknown": 2,
|
||||
"maybe_ethernet": 3,
|
||||
"thunderbolt": 4,
|
||||
}
|
||||
return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
|
||||
|
||||
|
||||
@@ -400,7 +414,7 @@ def get_mlx_ring_hosts_by_node(
|
||||
continue
|
||||
|
||||
connection_ip = _find_ip_prioritised(
|
||||
node_id, other_node_id, cycle_digraph, node_network
|
||||
node_id, other_node_id, cycle_digraph, node_network, ring=True
|
||||
)
|
||||
if connection_ip is None:
|
||||
raise ValueError(
|
||||
@@ -431,7 +445,9 @@ def get_mlx_jaccl_coordinators(
|
||||
if n == coordinator:
|
||||
return "0.0.0.0"
|
||||
|
||||
ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
|
||||
ip = _find_ip_prioritised(
|
||||
n, coordinator, cycle_digraph, node_network, ring=False
|
||||
)
|
||||
if ip is not None:
|
||||
return ip
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user