mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-09 03:51:22 -04:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Evan
|
ded7840499 | ||
|
Evan
|
7ed3eaa617 | ||
|
Evan
|
e113d42ba1 | ||
|
Evan
|
c5245dd87e | ||
|
Alex Cheema
|
edef8004f8 | ||
|
Alex CheemaandClaude Opus 4.7
|
a0c00f9dfd | ||
|
Drifter4242
|
89d20c1888 | ||
|
Evan Quineyandciaranbor
|
dbcceaa50c | ||
|
Sam BradburyandSam Bradbury
|
9c6ff4ce95 | ||
|
ecohash-coandJordan Miller
|
b26268dfaf | ||
|
ciaranbor
|
8dae3ecb9a | ||
|
Alex CheemaandClaude Opus 4.7
|
fb12b403ea | ||
|
Alex CheemaandClaude Opus 4.7
|
1606e63816 |
No files matched your search
@@ -1 +1,8 @@
|
||||
use flake
|
||||
|
||||
# creates .venv if doesn't exist and loads its environment
|
||||
export VIRTUAL_ENV=".venv"
|
||||
if ! [ -d "./$VIRTUAL_ENV" ]; then
|
||||
uv venv
|
||||
fi
|
||||
layout python
|
||||
@@ -191,10 +191,13 @@ class RotatingKVCache(_BaseCache):
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@property
|
||||
def meta_state(self) -> tuple[str, ...]: ...
|
||||
def meta_state(self): # -> tuple[str, ...]:
|
||||
...
|
||||
@meta_state.setter
|
||||
def meta_state(self, v: tuple[str, ...]) -> None: ...
|
||||
def is_trimmable(self) -> bool: ...
|
||||
def meta_state(self, v): # -> None:
|
||||
...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
Generated
+2338
-2254
File diff suppressed because it is too large.
Load diff
+16
-15
@@ -1,6 +1,6 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
|
||||
members = ["rust/exo_net", "rust/networking"]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.0.1"
|
||||
@@ -20,30 +20,31 @@ opt-level = 3
|
||||
[workspace.dependencies]
|
||||
## Crate members as common dependencies
|
||||
networking = { path = "rust/networking" }
|
||||
util = { path = "rust/util" }
|
||||
|
||||
# Macro dependecies
|
||||
extend = "1.2"
|
||||
delegate = "0.13"
|
||||
|
||||
# Utility dependencies
|
||||
keccak-const = "0.2"
|
||||
|
||||
# Async dependencies
|
||||
async-stream = "0.3"
|
||||
tokio = "1.46"
|
||||
futures-lite = "2.6.1"
|
||||
futures-timer = "3.0"
|
||||
|
||||
# Data structures
|
||||
either = "1.15"
|
||||
|
||||
# Tracing/logging
|
||||
log = "0.4"
|
||||
env_logger = "0.11.10"
|
||||
|
||||
# networking
|
||||
libp2p = "0.56"
|
||||
libp2p-tcp = "0.44"
|
||||
zenoh = "=1.9.0"
|
||||
async-stream = "0.3.6"
|
||||
netwatcher = "0.6.0"
|
||||
parking_lot = "0.12.5"
|
||||
pin-project = "1.1.10"
|
||||
pyo3 = "0.27.2"
|
||||
pyo3-async-runtimes = "0.27.0"
|
||||
pyo3-log = "0.13.2"
|
||||
pyo3-stub-gen = "0.22.2"
|
||||
rand = "0.10.1"
|
||||
serde_json = "1.0.149"
|
||||
tracing = "0.1.44"
|
||||
zenoh-plugin-storage-manager = { version = "=1.9.0", default-features = false }
|
||||
zenoh-plugin-trait = "=1.9.0"
|
||||
|
||||
[workspace.lints.rust]
|
||||
static_mut_refs = "warn" # Or use "warn" instead of deny
|
||||
|
||||
@@ -16,22 +16,13 @@ struct ContentView: View {
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var settingsWindowController: SettingsWindowController
|
||||
@EnvironmentObject private var bugReportWindowController: BugReportWindowController
|
||||
@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
|
||||
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 = ""
|
||||
@@ -294,6 +285,13 @@ struct ContentView: View {
|
||||
) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
HoverButton(
|
||||
title: "Share Bug Report…",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "ladybug"
|
||||
) {
|
||||
bugReportWindowController.open()
|
||||
}
|
||||
.padding(.bottom, 8)
|
||||
HoverButton(title: "Quit", tint: .secondary) {
|
||||
controller.stop()
|
||||
@@ -477,40 +475,6 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var debugSection: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
HoverButton(
|
||||
title: "Debug Info",
|
||||
tint: .primary,
|
||||
trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
|
||||
small: true
|
||||
) {
|
||||
showDebugInfo.toggle()
|
||||
}
|
||||
if showDebugInfo {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("Version: \(buildTag)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
Text("Commit: \(buildCommit)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
Text(thunderboltStatusText)
|
||||
.font(.caption2)
|
||||
.foregroundColor(thunderboltStatusColor)
|
||||
clusterThunderboltBridgeView
|
||||
interfaceIpList
|
||||
rdmaStatusView
|
||||
sendBugReportButton
|
||||
.padding(.top, 6)
|
||||
}
|
||||
.padding(.leading, 8)
|
||||
.transition(.opacity)
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.25), value: showDebugInfo)
|
||||
}
|
||||
|
||||
private var rdmaStatusView: some View {
|
||||
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
@@ -559,127 +523,6 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
switch bugReportPhase {
|
||||
case .idle:
|
||||
Button {
|
||||
bugReportPhase = .prompting
|
||||
bugReportUserDescription = ""
|
||||
} label: {
|
||||
HStack {
|
||||
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))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
|
||||
case .prompting:
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
Text("Tell us what went wrong (optional)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
Text(
|
||||
"A quick description of what you were doing and what happened helps us track down the bug for you."
|
||||
)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.opacity(0.8)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
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.06))
|
||||
)
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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> {
|
||||
Binding(
|
||||
get: {
|
||||
@@ -720,61 +563,6 @@ struct ContentView: View {
|
||||
)
|
||||
}
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportPhase = .sending("Collecting logs...")
|
||||
let service = BugReportService()
|
||||
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
do {
|
||||
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 {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
let alert = NSAlert()
|
||||
alert.messageText = "Uninstall EXO"
|
||||
@@ -857,13 +645,6 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
private struct HoverButton: View {
|
||||
|
||||
@@ -22,6 +22,7 @@ struct EXOApp: App {
|
||||
@StateObject private var updater: SparkleUpdater
|
||||
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@StateObject private var settingsWindowController: SettingsWindowController
|
||||
@StateObject private var bugReportWindowController: BugReportWindowController
|
||||
private let terminationObserver: TerminationObserver
|
||||
private let firstLaunchPopout = FirstLaunchPopout()
|
||||
private let ciContext = CIContext(options: nil)
|
||||
@@ -46,6 +47,7 @@ struct EXOApp: App {
|
||||
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
|
||||
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
|
||||
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
|
||||
_bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
|
||||
enableLaunchAtLoginIfNeeded()
|
||||
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
|
||||
NetworkSetupHelper.promptAndInstallIfNeeded()
|
||||
@@ -66,6 +68,7 @@ struct EXOApp: App {
|
||||
.environmentObject(updater)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(settingsWindowController)
|
||||
.environmentObject(bugReportWindowController)
|
||||
} label: {
|
||||
menuBarIcon
|
||||
.onReceive(controller.$isFirstLaunchReady) { ready in
|
||||
|
||||
@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
|
||||
|
||||
init(
|
||||
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
|
||||
session: URLSession = .shared
|
||||
session: URLSession = ClusterStateService.makeNonCachingSession()
|
||||
) {
|
||||
self.baseURL = baseURL
|
||||
self.endpoint = baseURL.appendingPathComponent("state")
|
||||
@@ -27,6 +27,23 @@ final class ClusterStateService: ObservableObject {
|
||||
self.decoder = decoder
|
||||
}
|
||||
|
||||
/// `URLSession.shared` carries an on-disk `URLCache` that persists every
|
||||
/// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
|
||||
/// at 2 Hz from `startPolling`, so leaving the shared cache attached
|
||||
/// dirties ~500–620 KB/sec of file-backed memory and trips macOS's
|
||||
/// per-process `disk writes` resource limit (microstackshot reports
|
||||
/// observed on M3 Ultra producing GBs of cached responses per hour).
|
||||
/// Cluster-state polling responses are time-sensitive and small; they
|
||||
/// gain nothing from being cached on disk. Use an ephemeral session
|
||||
/// with `urlCache = nil` so neither response bodies nor metadata
|
||||
/// touch disk.
|
||||
private static func makeNonCachingSession() -> URLSession {
|
||||
let config = URLSessionConfiguration.ephemeral
|
||||
config.urlCache = nil
|
||||
config.requestCachePolicy = .reloadIgnoringLocalCacheData
|
||||
return URLSession(configuration: config)
|
||||
}
|
||||
|
||||
func startPolling(interval: TimeInterval = 0.5) {
|
||||
stopPolling()
|
||||
Task {
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Manages a standalone window for the bug-report flow.
|
||||
/// Ensures only one instance exists and brings it to front on repeated opens.
|
||||
@MainActor
|
||||
final class BugReportWindowController: ObservableObject {
|
||||
private var window: NSWindow?
|
||||
|
||||
func open() {
|
||||
if let existing = window, existing.isVisible {
|
||||
existing.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate()
|
||||
return
|
||||
}
|
||||
|
||||
let view = BugReportView(onDismiss: { [weak self] in
|
||||
self?.window?.close()
|
||||
})
|
||||
|
||||
let hostingController = NSHostingController(rootView: view)
|
||||
hostingController.sizingOptions = [.preferredContentSize, .minSize]
|
||||
|
||||
let newWindow = NSWindow(contentViewController: hostingController)
|
||||
newWindow.styleMask = [.titled, .closable, .resizable]
|
||||
newWindow.title = "Send a Bug Report"
|
||||
newWindow.center()
|
||||
newWindow.setFrameAutosaveName("ExoBugReportWindow")
|
||||
newWindow.isReleasedWhenClosed = false
|
||||
newWindow.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate()
|
||||
|
||||
window = newWindow
|
||||
}
|
||||
}
|
||||
|
||||
private struct BugReportView: View {
|
||||
fileprivate enum Phase: Equatable {
|
||||
case prompting
|
||||
case sending(String)
|
||||
case success(String)
|
||||
case failure(String)
|
||||
}
|
||||
|
||||
let onDismiss: () -> Void
|
||||
|
||||
@State private var phase: Phase = .prompting
|
||||
@State private var userDescription: String = ""
|
||||
@FocusState private var descriptionFocused: Bool
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
switch phase {
|
||||
case .prompting:
|
||||
promptingView
|
||||
case .sending(let message):
|
||||
sendingView(message: message)
|
||||
case .success(let message):
|
||||
successView(message: message)
|
||||
case .failure(let message):
|
||||
failureView(message: message)
|
||||
}
|
||||
}
|
||||
.padding(16)
|
||||
.frame(minWidth: 380)
|
||||
.animation(.easeInOut(duration: 0.2), value: phase)
|
||||
.onAppear { descriptionFocused = true }
|
||||
}
|
||||
|
||||
private var promptingView: some View {
|
||||
VStack(alignment: .leading, spacing: 8) {
|
||||
Text("Description (optional)")
|
||||
.font(.subheadline)
|
||||
.foregroundColor(.secondary)
|
||||
ZStack(alignment: .topLeading) {
|
||||
if userDescription.isEmpty {
|
||||
Text("What were you doing when it broke?")
|
||||
.font(.body)
|
||||
.foregroundColor(Color(nsColor: .placeholderTextColor))
|
||||
.padding(.horizontal, 10)
|
||||
.padding(.vertical, 8)
|
||||
.allowsHitTesting(false)
|
||||
}
|
||||
TextEditor(text: $userDescription)
|
||||
.font(.body)
|
||||
.scrollContentBackground(.hidden)
|
||||
.padding(4)
|
||||
.frame(height: 72)
|
||||
.focused($descriptionFocused)
|
||||
}
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color(nsColor: .textBackgroundColor))
|
||||
)
|
||||
.overlay(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
|
||||
)
|
||||
|
||||
Text("Diagnostic logs will be uploaded with your report.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Cancel") { onDismiss() }
|
||||
.keyboardShortcut(.cancelAction)
|
||||
Button("Send") {
|
||||
Task { await send() }
|
||||
}
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
.padding(.top, 4)
|
||||
}
|
||||
}
|
||||
|
||||
private func sendingView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(spacing: 10) {
|
||||
ProgressView().controlSize(.small)
|
||||
Text(message)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Cancel") { onDismiss() }
|
||||
.keyboardShortcut(.cancelAction)
|
||||
.disabled(true)
|
||||
Button("Send") {}
|
||||
.disabled(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func successView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(alignment: .top, spacing: 10) {
|
||||
Image(systemName: "checkmark.circle.fill")
|
||||
.foregroundColor(.green)
|
||||
.font(.title2)
|
||||
Text(message)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
HStack {
|
||||
Button {
|
||||
openGitHubIssue()
|
||||
} label: {
|
||||
HStack(spacing: 4) {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
Text("Open GitHub Issue")
|
||||
}
|
||||
}
|
||||
Spacer()
|
||||
Button("Done") { onDismiss() }
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func failureView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(alignment: .top, spacing: 10) {
|
||||
Image(systemName: "exclamationmark.triangle.fill")
|
||||
.foregroundColor(.orange)
|
||||
.font(.title2)
|
||||
Text(message)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Try Again") {
|
||||
phase = .prompting
|
||||
}
|
||||
Button("Close") { onDismiss() }
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func send() async {
|
||||
phase = .sending("Collecting logs and uploading…")
|
||||
let service = BugReportService()
|
||||
let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
do {
|
||||
let outcome = try await service.sendReport(
|
||||
isManual: true,
|
||||
userDescription: description.isEmpty ? nil : description
|
||||
)
|
||||
if outcome.success {
|
||||
phase = .success(outcome.message)
|
||||
} else {
|
||||
phase = .failure(outcome.message)
|
||||
}
|
||||
} catch {
|
||||
phase = .failure(error.localizedDescription)
|
||||
}
|
||||
}
|
||||
|
||||
private func openGitHubIssue() {
|
||||
let description = userDescription.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)
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
@@ -21,8 +21,6 @@ struct SettingsView: View {
|
||||
@State private var pendingReadOnlyModelsDirs: String = ""
|
||||
@State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
|
||||
@State private var needsRestart = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
@State private var uninstallInProgress = false
|
||||
|
||||
var body: some View {
|
||||
@@ -202,8 +200,6 @@ struct SettingsView: View {
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
rdmaStatusView
|
||||
}
|
||||
|
||||
sendBugReportButton
|
||||
}
|
||||
|
||||
Section("Danger Zone") {
|
||||
@@ -504,50 +500,8 @@ struct SettingsView: View {
|
||||
}
|
||||
}
|
||||
|
||||
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"
|
||||
|
||||
+1
-3
@@ -589,9 +589,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
|
||||
help="Only consider placements using >= this many nodes.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--instance-meta",
|
||||
choices=["ring", "jaccl", "vllm", "both"],
|
||||
default="both",
|
||||
"--instance-meta", choices=["ring", "jaccl", "both"], default="both"
|
||||
)
|
||||
ap.add_argument(
|
||||
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
|
||||
|
||||
@@ -12,24 +12,24 @@ timeout = 7200.0
|
||||
settle_timeout = 60.0
|
||||
|
||||
# Workload
|
||||
pp = [4096, 8192]
|
||||
tg = [128]
|
||||
pp = [4096]
|
||||
tg = [512]
|
||||
repeat = 1
|
||||
warmup = 0
|
||||
|
||||
json_out = "bench/prefill_decode_results.json"
|
||||
|
||||
[prefill]
|
||||
model = "sakamakismile/Qwen3.6-27B-NVFP4"
|
||||
node = "gx10-de89"
|
||||
instance_meta = "vllm"
|
||||
model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
|
||||
node = "mike"
|
||||
instance_meta = "ring"
|
||||
sharding = "pipeline"
|
||||
min_nodes = 1
|
||||
max_nodes = 1
|
||||
|
||||
[decode]
|
||||
model = "mlx-community/Qwen3.6-27B-4bit"
|
||||
node = "Ryuichi’s MacBook Pro"
|
||||
model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
|
||||
node = "james"
|
||||
instance_meta = "ring"
|
||||
sharding = "pipeline"
|
||||
min_nodes = 1
|
||||
|
||||
@@ -31,7 +31,6 @@ from typing import Any
|
||||
|
||||
from exo_bench import (
|
||||
PromptSizer,
|
||||
SystemMetricsSampler,
|
||||
format_peak_memory,
|
||||
load_tokenizer_for_bench,
|
||||
parse_int_list,
|
||||
@@ -279,7 +278,6 @@ def _run_phase(
|
||||
warmup: int,
|
||||
repeat: int,
|
||||
common_meta: dict[str, Any],
|
||||
sampler: SystemMetricsSampler | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
logger.info(f"=== phase: {label} (model={model_id}) ===")
|
||||
rows: list[dict[str, Any]] = []
|
||||
@@ -290,13 +288,10 @@ def _run_phase(
|
||||
for pp, tg in pp_tg_pairs:
|
||||
logger.info(f"--- {label}: pp={pp} tg={tg} ---")
|
||||
runs: list[dict[str, Any]] = []
|
||||
inference_windows: list[tuple[float, float]] = []
|
||||
for r in range(repeat):
|
||||
time.sleep(2)
|
||||
try:
|
||||
inf_t0 = time.monotonic()
|
||||
row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
|
||||
inference_windows.append((inf_t0, time.monotonic()))
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
continue
|
||||
@@ -320,26 +315,11 @@ def _run_phase(
|
||||
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
|
||||
peak = mean(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
|
||||
avg_elapsed = mean(x["elapsed_s"] for x in runs)
|
||||
energy_str = ""
|
||||
if sampler is not None and inference_windows:
|
||||
joules = sum(
|
||||
sampler.energy_between(t0, t1) for t0, t1 in inference_windows
|
||||
)
|
||||
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
|
||||
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0.0
|
||||
energy_per_run = joules / len(runs) if runs else 0.0
|
||||
energy_str = (
|
||||
f" energy={joules:.1f}J ({avg_watts:.1f}W avg over "
|
||||
f"{inf_seconds:.1f}s inference, {energy_per_run:.1f}J/run)"
|
||||
)
|
||||
for run_row, (t0, t1) in zip(runs, inference_windows, strict=False):
|
||||
run_row["energy_joules"] = sampler.energy_between(t0, t1)
|
||||
run_row["inference_window_s"] = t1 - t0
|
||||
logger.info(
|
||||
f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
|
||||
f"prompt_tokens={ptok} gen_tokens={gtok} "
|
||||
f"peak_memory={format_peak_memory(peak)} "
|
||||
f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
|
||||
f"avg_elapsed={avg_elapsed:.2f}s"
|
||||
)
|
||||
time.sleep(2)
|
||||
return rows
|
||||
@@ -352,36 +332,14 @@ def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, fl
|
||||
grouped.setdefault(key, []).append(r)
|
||||
out: dict[tuple[int, int], dict[str, float]] = {}
|
||||
for key, runs in grouped.items():
|
||||
energy_runs = [x.get("energy_joules") for x in runs if "energy_joules" in x]
|
||||
window_runs = [
|
||||
x.get("inference_window_s") for x in runs if "inference_window_s" in x
|
||||
]
|
||||
out[key] = {
|
||||
"prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
|
||||
"gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
|
||||
"elapsed_s": mean(x["elapsed_s"] for x in runs),
|
||||
"prompt_tokens": mean(x["stats"]["prompt_tokens"] for x in runs),
|
||||
"gen_tokens": mean(x["stats"]["generation_tokens"] for x in runs),
|
||||
"energy_j": mean(energy_runs) if energy_runs else 0.0,
|
||||
"inference_window_s": mean(window_runs) if window_runs else 0.0,
|
||||
}
|
||||
return out
|
||||
|
||||
|
||||
def _normalised_seconds(summary: dict[str, float], pp: int, tg: int) -> float | None:
|
||||
"""Wall-clock time implied by reported tps for the *configured* pp/tg.
|
||||
|
||||
elapsed_s is not comparable across phases when models EOS at different
|
||||
lengths. This formula reconstructs "what would this phase take to do
|
||||
pp prompt tokens + tg generation tokens" using its own reported rates.
|
||||
"""
|
||||
p_tps = summary.get("prompt_tps", 0.0)
|
||||
g_tps = summary.get("gen_tps", 0.0)
|
||||
if p_tps <= 0 or g_tps <= 0:
|
||||
return None
|
||||
return pp / p_tps + tg / g_tps
|
||||
|
||||
|
||||
def _print_diff(
|
||||
disagg_rows: list[dict[str, Any]],
|
||||
decode_alone_rows: list[dict[str, Any]],
|
||||
@@ -392,17 +350,14 @@ def _print_diff(
|
||||
prefill_alone = _summarise(prefill_alone_rows)
|
||||
keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
|
||||
|
||||
width = 110
|
||||
width = 64
|
||||
for key in sorted(keys):
|
||||
pp, tg = key
|
||||
logger.info("─" * width)
|
||||
logger.info(f" pp={pp} tg={tg}")
|
||||
logger.info("─" * width)
|
||||
logger.info(
|
||||
f" {'phase':<16} {'elapsed':>9} {'norm':>9} "
|
||||
f"{'prompt_tps':>11} {'gen_tps':>8} "
|
||||
f"{'p_tok':>6} {'g_tok':>6} "
|
||||
f"{'energy':>9} {'avg_W':>7}"
|
||||
f" {'phase':<16} {'elapsed':>10} {'prompt_tps':>11} {'gen_tps':>9}"
|
||||
)
|
||||
for label, summary in (
|
||||
("disaggregated", disagg.get(key)),
|
||||
@@ -410,51 +365,26 @@ def _print_diff(
|
||||
("prefill_alone", prefill_alone.get(key)),
|
||||
):
|
||||
if summary is None:
|
||||
logger.info(
|
||||
f" {label:<16} {'—':>9} {'—':>9} "
|
||||
f"{'—':>11} {'—':>8} {'—':>6} {'—':>6} "
|
||||
f"{'—':>9} {'—':>7}"
|
||||
)
|
||||
logger.info(f" {label:<16} {'—':>10} {'—':>11} {'—':>9}")
|
||||
continue
|
||||
norm = _normalised_seconds(summary, pp, tg)
|
||||
norm_str = f"{norm:>8.2f}s" if norm is not None else f"{'—':>9}"
|
||||
energy = summary.get("energy_j", 0.0)
|
||||
window = summary.get("inference_window_s", 0.0)
|
||||
energy_str = f"{energy:>8.1f}J" if energy > 0 else f"{'—':>9}"
|
||||
avg_w = energy / window if window > 0 else 0.0
|
||||
avg_w_str = f"{avg_w:>6.1f}W" if avg_w > 0 else f"{'—':>7}"
|
||||
logger.info(
|
||||
f" {label:<16} "
|
||||
f"{summary['elapsed_s']:>8.2f}s "
|
||||
f"{norm_str} "
|
||||
f"{summary['elapsed_s']:>9.2f}s "
|
||||
f"{summary['prompt_tps']:>11.1f} "
|
||||
f"{summary['gen_tps']:>8.2f} "
|
||||
f"{summary['prompt_tokens']:>6.0f} "
|
||||
f"{summary['gen_tokens']:>6.0f} "
|
||||
f"{energy_str} "
|
||||
f"{avg_w_str}"
|
||||
f"{summary['gen_tps']:>9.2f}"
|
||||
)
|
||||
|
||||
d = disagg.get(key)
|
||||
da = decode_alone.get(key)
|
||||
pa = prefill_alone.get(key)
|
||||
d_norm = _normalised_seconds(d, pp, tg) if d else None
|
||||
if d_norm and da:
|
||||
da_norm = _normalised_seconds(da, pp, tg)
|
||||
if da_norm:
|
||||
logger.info(
|
||||
f" norm speedup vs decode_alone: {da_norm / d_norm:.2f}x "
|
||||
f"(prefill {d['prompt_tps'] / da['prompt_tps']:.2f}x, "
|
||||
f"decode {d['gen_tps'] / da['gen_tps']:.2f}x)"
|
||||
)
|
||||
if d_norm and pa:
|
||||
pa_norm = _normalised_seconds(pa, pp, tg)
|
||||
if pa_norm:
|
||||
logger.info(
|
||||
f" norm speedup vs prefill_alone: {pa_norm / d_norm:.2f}x "
|
||||
f"(prefill {d['prompt_tps'] / pa['prompt_tps']:.2f}x, "
|
||||
f"decode {d['gen_tps'] / pa['gen_tps']:.2f}x)"
|
||||
)
|
||||
if d and da and d["elapsed_s"] > 0:
|
||||
logger.info(
|
||||
f" speedup vs decode_alone: {da['elapsed_s'] / d['elapsed_s']:.2f}x"
|
||||
)
|
||||
if d and pa and d["elapsed_s"] > 0:
|
||||
logger.info(
|
||||
f" speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
|
||||
)
|
||||
logger.info("─" * width)
|
||||
|
||||
|
||||
@@ -752,16 +682,6 @@ def main() -> int:
|
||||
link_id = ""
|
||||
prefill_alive = False
|
||||
decode_alive = False
|
||||
sampler_nodes = sorted(
|
||||
{
|
||||
*node_ids_from_instance(prefill_instance),
|
||||
*node_ids_from_instance(decode_instance),
|
||||
}
|
||||
)
|
||||
sampler = SystemMetricsSampler(
|
||||
ExoClient(args.host, args.port, timeout_s=30), sampler_nodes
|
||||
)
|
||||
sampler.start()
|
||||
try:
|
||||
logger.info("Creating prefill instance...")
|
||||
client.request_json("POST", "/instance", body={"instance": prefill_instance})
|
||||
@@ -780,7 +700,6 @@ def main() -> int:
|
||||
warmup=args.warmup,
|
||||
repeat=args.repeat,
|
||||
common_meta=common_meta,
|
||||
sampler=sampler,
|
||||
)
|
||||
all_rows.extend(prefill_alone_rows)
|
||||
|
||||
@@ -810,7 +729,6 @@ def main() -> int:
|
||||
warmup=args.warmup,
|
||||
repeat=args.repeat,
|
||||
common_meta=common_meta,
|
||||
sampler=sampler,
|
||||
)
|
||||
all_rows.extend(disagg_rows)
|
||||
|
||||
@@ -835,13 +753,11 @@ def main() -> int:
|
||||
warmup=args.warmup,
|
||||
repeat=args.repeat,
|
||||
common_meta=common_meta,
|
||||
sampler=sampler,
|
||||
)
|
||||
all_rows.extend(decode_alone_rows)
|
||||
|
||||
_print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
|
||||
finally:
|
||||
sampler.stop()
|
||||
with contextlib.suppress(ExoHttpError):
|
||||
if link_id:
|
||||
_delete_instance_link(client, link_id)
|
||||
|
||||
@@ -202,7 +202,6 @@
|
||||
let instanceType: string | null = null;
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
*/
|
||||
|
||||
interface Props {
|
||||
/** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
|
||||
/** "macbook pro" | "mac studio" | "mac mini" etc. */
|
||||
deviceType: string;
|
||||
/** Center X coordinate in SVG space */
|
||||
cx: number;
|
||||
@@ -38,43 +38,10 @@
|
||||
const LOGO_NATIVE_WIDTH = 814;
|
||||
const LOGO_NATIVE_HEIGHT = 1000;
|
||||
|
||||
// NVIDIA logo SVG path
|
||||
const NVIDIA_LOGO_PATH =
|
||||
"M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
|
||||
|
||||
const wireColor = "rgba(179,179,179,0.8)";
|
||||
const strokeWidth = 1.5;
|
||||
|
||||
const modelLower = $derived(deviceType.toLowerCase());
|
||||
const isSpark = $derived(
|
||||
modelLower.includes("dgx") || modelLower.includes("gx10"),
|
||||
);
|
||||
const isLinux = $derived(!isSpark && modelLower.startsWith("linux"));
|
||||
const isLinuxLaptop = $derived(isLinux && modelLower.includes("laptop"));
|
||||
|
||||
// ── DGX Spark dimensions ──
|
||||
const dgxW = $derived(size * 1.55);
|
||||
const dgxH = $derived(size * 0.58);
|
||||
const dgxX = $derived(cx - dgxW / 2);
|
||||
const dgxY = $derived(cy - dgxH / 2);
|
||||
const dgxChassisX = $derived(dgxX - dgxW * 0.03);
|
||||
const dgxChassisW = $derived(dgxW * 1.05);
|
||||
const dgxHandleW = $derived(dgxW * 0.27);
|
||||
const dgxHandleGap = $derived(dgxH * 0.05);
|
||||
const dgxHandleH = $derived(dgxH - dgxHandleGap * 2);
|
||||
const dgxHandleY = $derived(dgxY + dgxHandleGap);
|
||||
const dgxInnerHandleW = $derived(dgxW * 0.12);
|
||||
const dgxInnerHandleH = $derived(dgxHandleH - dgxH * 0.06);
|
||||
const dgxLeftHandleX = $derived(dgxX + 4);
|
||||
const dgxRightHandleX = $derived(dgxX + dgxW - dgxHandleW - 4);
|
||||
const dgxClipId = $derived(`di-dgx-${uid}`);
|
||||
const dgxTextureId = $derived(`di-dgx-tex-${uid}`);
|
||||
|
||||
// ── Linux Desktop dimensions (reuses Mac Studio proportions) ──
|
||||
const linuxDesktopClipId = $derived(`di-linux-desktop-${uid}`);
|
||||
|
||||
// ── Linux Laptop dimensions (reuses MacBook proportions) ──
|
||||
const linuxScreenClipId = $derived(`di-linux-screen-${uid}`);
|
||||
|
||||
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
|
||||
const studioW = $derived(size * 1.25);
|
||||
@@ -147,264 +114,7 @@
|
||||
const studioClipId = $derived(`di-studio-${uid}`);
|
||||
</script>
|
||||
|
||||
{#if isSpark}
|
||||
<!-- DGX Spark -->
|
||||
<defs>
|
||||
<clipPath id={dgxClipId}>
|
||||
<rect x={dgxX} y={dgxY} width={dgxW} height={dgxH} rx="3" />
|
||||
</clipPath>
|
||||
<pattern
|
||||
id={dgxTextureId}
|
||||
patternUnits="userSpaceOnUse"
|
||||
width="8"
|
||||
height="8"
|
||||
>
|
||||
<rect width="8" height="8" fill="#6f6248" />
|
||||
<circle cx="2" cy="2" r="1" fill="#5a4f3b" opacity="0.5" />
|
||||
<circle cx="6" cy="6" r="1" fill="#4a4232" opacity="0.45" />
|
||||
</pattern>
|
||||
</defs>
|
||||
|
||||
<!-- Main body -->
|
||||
<rect
|
||||
x={dgxChassisX}
|
||||
y={dgxY}
|
||||
width={dgxChassisW}
|
||||
height={dgxH}
|
||||
rx="3"
|
||||
fill="url(#{dgxTextureId})"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
|
||||
<!-- Side border accents -->
|
||||
<rect
|
||||
x={dgxChassisX}
|
||||
y={dgxY}
|
||||
width={dgxW * 0.02}
|
||||
height={dgxH}
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
<rect
|
||||
x={dgxChassisX + dgxChassisW - dgxW * 0.02}
|
||||
y={dgxY}
|
||||
width={dgxW * 0.02}
|
||||
height={dgxH}
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
|
||||
<!-- Memory fill -->
|
||||
{#if ramPercent > 0}
|
||||
<rect
|
||||
x={dgxX}
|
||||
y={dgxY + dgxH - (ramPercent / 100) * dgxH}
|
||||
width={dgxW}
|
||||
height={(ramPercent / 100) * dgxH}
|
||||
fill="rgba(255,215,0,0.45)"
|
||||
clip-path="url(#{dgxClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Left handle -->
|
||||
<rect
|
||||
x={dgxLeftHandleX}
|
||||
y={dgxHandleY}
|
||||
width={dgxHandleW}
|
||||
height={dgxHandleH}
|
||||
rx="2.4"
|
||||
fill="#b3a170"
|
||||
stroke="#403723"
|
||||
stroke-width="0.7"
|
||||
/>
|
||||
<rect
|
||||
x={dgxLeftHandleX + dgxHandleW * 0.06}
|
||||
y={dgxHandleY + dgxH * 0.03}
|
||||
width={dgxInnerHandleW}
|
||||
height={dgxInnerHandleH}
|
||||
rx="1.6"
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
|
||||
<!-- Right handle -->
|
||||
<rect
|
||||
x={dgxRightHandleX}
|
||||
y={dgxHandleY}
|
||||
width={dgxHandleW}
|
||||
height={dgxHandleH}
|
||||
rx="2.4"
|
||||
fill="#b3a170"
|
||||
stroke="#403723"
|
||||
stroke-width="0.7"
|
||||
/>
|
||||
<rect
|
||||
x={dgxRightHandleX + dgxHandleW - dgxInnerHandleW - dgxHandleW * 0.08}
|
||||
y={dgxHandleY + dgxH * 0.03}
|
||||
width={dgxInnerHandleW}
|
||||
height={dgxInnerHandleH}
|
||||
rx="1.6"
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
|
||||
<!-- NVIDIA logo (rotated 90deg on left handle) -->
|
||||
{@const badgeW = dgxW * 0.09}
|
||||
{@const badgeH = dgxHandleH * 0.5}
|
||||
{@const badgeX = dgxLeftHandleX + dgxHandleW - badgeW - dgxHandleW * 0.06}
|
||||
{@const badgeYPos = dgxHandleY + (dgxHandleH - badgeH) / 2}
|
||||
{@const textSz = badgeW * 0.58}
|
||||
{@const logoW = textSz * 1.2}
|
||||
{@const logoH = logoW * (1.438 / 2.174)}
|
||||
{@const ctrX = badgeX + badgeW / 2 - badgeW * 0.03}
|
||||
{@const ctrY = badgeYPos + badgeH / 2}
|
||||
{@const labelGap = badgeW * 0.15}
|
||||
{@const totalW = logoW + labelGap + textSz * 3.6}
|
||||
<g transform="rotate(90 {ctrX} {ctrY})">
|
||||
<svg
|
||||
x={ctrX - totalW / 2}
|
||||
y={ctrY - logoH / 2}
|
||||
width={logoW}
|
||||
height={logoH}
|
||||
viewBox="0 0 2.174 1.438"
|
||||
>
|
||||
<path d={NVIDIA_LOGO_PATH} fill="#76b900" />
|
||||
</svg>
|
||||
<text
|
||||
x={ctrX - totalW / 2 + logoW + labelGap}
|
||||
y={ctrY}
|
||||
text-anchor="start"
|
||||
dominant-baseline="middle"
|
||||
fill="#8a7a56"
|
||||
font-size={textSz}
|
||||
font-family="monospace"
|
||||
font-weight="700">NVIDIA</text
|
||||
>
|
||||
</g>
|
||||
{:else if isLinuxLaptop}
|
||||
<!-- Linux Laptop — MacBook shape with Tux logo -->
|
||||
<defs>
|
||||
<clipPath id={linuxScreenClipId}>
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<rect
|
||||
x={mbScreenX}
|
||||
y={mbY}
|
||||
width={mbScreenW}
|
||||
height={mbScreenH}
|
||||
rx="3"
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
fill="#0a0a12"
|
||||
/>
|
||||
{#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(#{linuxScreenClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Terminal prompt on screen -->
|
||||
<text
|
||||
x={cx}
|
||||
y={mbY + mbScreenH / 2}
|
||||
text-anchor="middle"
|
||||
dominant-baseline="middle"
|
||||
fill="#FFFFFF"
|
||||
opacity="0.9"
|
||||
font-size={mbScreenH * 0.25}
|
||||
font-family="SF Mono, Monaco, monospace"
|
||||
font-weight="700">{">_"}</text
|
||||
>
|
||||
|
||||
<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"
|
||||
/>
|
||||
<rect
|
||||
x={mbKbX}
|
||||
y={mbKbY}
|
||||
width={mbKbW}
|
||||
height={mbKbH}
|
||||
fill="rgba(0,0,0,0.2)"
|
||||
rx="2"
|
||||
/>
|
||||
<rect
|
||||
x={mbTpX}
|
||||
y={mbTpY}
|
||||
width={mbTpW}
|
||||
height={mbTpH}
|
||||
fill="rgba(255,255,255,0.08)"
|
||||
rx="2"
|
||||
/>
|
||||
{:else if isLinux}
|
||||
<!-- Linux Desktop — Mac Studio shape with Tux logo -->
|
||||
<defs>
|
||||
<clipPath id={linuxDesktopClipId}>
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY + studioTopH}
|
||||
width={studioW}
|
||||
height={studioH - studioTopH}
|
||||
rx={studioCorner - 1}
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY}
|
||||
width={studioW}
|
||||
height={studioH}
|
||||
rx={studioCorner}
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
{#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(#{linuxDesktopClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Terminal prompt on front face -->
|
||||
<text
|
||||
x={cx}
|
||||
y={studioY + studioTopH + (studioH - studioTopH) / 2}
|
||||
text-anchor="middle"
|
||||
dominant-baseline="middle"
|
||||
fill="rgba(255,255,255,0.5)"
|
||||
font-size={(studioH - studioTopH) * 0.4}
|
||||
font-family="SF Mono, Monaco, monospace"
|
||||
font-weight="700">{">_"}</text
|
||||
>
|
||||
{:else if modelLower === "mac studio" || modelLower === "mac mini"}
|
||||
{#if modelLower === "mac studio" || modelLower === "mac mini"}
|
||||
<!-- Mac Studio / Mac Mini -->
|
||||
<defs>
|
||||
<clipPath id={studioClipId}>
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
} | null;
|
||||
nodes?: Record<string, NodeInfo>;
|
||||
sharding?: "Pipeline" | "Tensor";
|
||||
runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
|
||||
runtime?: "MlxRing" | "MlxJaccl";
|
||||
onLaunch?: () => void;
|
||||
tags?: string[];
|
||||
apiPreview?: PlacementPreview | null;
|
||||
@@ -168,10 +168,8 @@
|
||||
|
||||
function getDeviceType(
|
||||
name: string,
|
||||
): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
|
||||
): "macbook" | "studio" | "mini" | "unknown" {
|
||||
const lower = name.toLowerCase();
|
||||
if (lower.includes("dgx") || lower.includes("gx10")) return "dgx";
|
||||
if (lower.includes("linux")) return "linux";
|
||||
if (lower.includes("macbook")) return "macbook";
|
||||
if (lower.includes("studio")) return "studio";
|
||||
if (lower.includes("mini")) return "mini";
|
||||
@@ -578,17 +576,13 @@
|
||||
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)."
|
||||
: runtime === "MlxJaccl"
|
||||
? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
|
||||
: "vLLM: NVIDIA CUDA inference engine."}
|
||||
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
|
||||
>
|
||||
{runtime === "MlxRing"
|
||||
? "MLX Ring"
|
||||
: runtime === "MlxJaccl"
|
||||
? "MLX RDMA"
|
||||
: runtime === "Vllm"
|
||||
? "vLLM"
|
||||
: runtime}
|
||||
: runtime}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
@@ -996,81 +990,6 @@
|
||||
/>
|
||||
{/if}
|
||||
</g>
|
||||
{:else if node.deviceType === "dgx"}
|
||||
<!-- DGX Spark icon -->
|
||||
{@const s = node.iconSize}
|
||||
{@const dgxW = s * 1.4}
|
||||
{@const dgxH = s * 0.52}
|
||||
<g transform="translate({-dgxW / 2}, {-dgxH / 2})">
|
||||
<!-- Chassis -->
|
||||
<rect
|
||||
x="0"
|
||||
y="0"
|
||||
width={dgxW}
|
||||
height={dgxH}
|
||||
rx="2"
|
||||
fill="#6f6248"
|
||||
stroke={node.isUsed ? "#FFD700" : "#4B5563"}
|
||||
stroke-width="1.5"
|
||||
/>
|
||||
<!-- Side accents -->
|
||||
<rect
|
||||
x="0"
|
||||
y="0"
|
||||
width={dgxW * 0.02}
|
||||
height={dgxH}
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
<rect
|
||||
x={dgxW - dgxW * 0.02}
|
||||
y="0"
|
||||
width={dgxW * 0.02}
|
||||
height={dgxH}
|
||||
fill="#8a7a56"
|
||||
/>
|
||||
<!-- Left handle -->
|
||||
<rect
|
||||
x={dgxW * 0.04}
|
||||
y={dgxH * 0.08}
|
||||
width={dgxW * 0.22}
|
||||
height={dgxH * 0.84}
|
||||
rx="2"
|
||||
fill="#b3a170"
|
||||
stroke="#403723"
|
||||
stroke-width="0.5"
|
||||
/>
|
||||
<!-- Right handle -->
|
||||
<rect
|
||||
x={dgxW - dgxW * 0.04 - dgxW * 0.22}
|
||||
y={dgxH * 0.08}
|
||||
width={dgxW * 0.22}
|
||||
height={dgxH * 0.84}
|
||||
rx="2"
|
||||
fill="#b3a170"
|
||||
stroke="#403723"
|
||||
stroke-width="0.5"
|
||||
/>
|
||||
<!-- Memory fill -->
|
||||
<rect
|
||||
x="2"
|
||||
y={dgxH - dgxH * (node.currentPercent / 100)}
|
||||
width={dgxW - 4}
|
||||
height={dgxH * (node.currentPercent / 100)}
|
||||
fill="rgba(255,215,0,0.35)"
|
||||
/>
|
||||
{#if node.modelUsageGB > 0 && node.isUsed}
|
||||
<rect
|
||||
x="2"
|
||||
y={dgxH - dgxH * (node.newPercent / 100)}
|
||||
width={dgxW - 4}
|
||||
height={dgxH *
|
||||
((node.newPercent - node.currentPercent) / 100)}
|
||||
fill="#FFD700"
|
||||
filter="url(#memGlow-{filterId})"
|
||||
class="animate-pulse-slow"
|
||||
/>
|
||||
{/if}
|
||||
</g>
|
||||
{:else}
|
||||
<!-- Unknown device - hexagon -->
|
||||
<g
|
||||
|
||||
@@ -9,7 +9,6 @@
|
||||
capabilities?: string[];
|
||||
family?: string;
|
||||
is_custom?: boolean;
|
||||
requires_vllm?: boolean;
|
||||
}
|
||||
|
||||
interface ModelGroup {
|
||||
@@ -20,7 +19,6 @@
|
||||
variants: ModelInfo[];
|
||||
smallestVariant: ModelInfo;
|
||||
hasMultipleVariants: boolean;
|
||||
requiresVllm: boolean;
|
||||
}
|
||||
|
||||
type DownloadAvailability = {
|
||||
@@ -215,14 +213,6 @@
|
||||
<span class="font-mono text-sm text-white truncate">
|
||||
{group.name}
|
||||
</span>
|
||||
{#if group.requiresVllm}
|
||||
<span
|
||||
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
|
||||
title="Requires vLLM runtime"
|
||||
>
|
||||
vLLM
|
||||
</span>
|
||||
{/if}
|
||||
<!-- Capability icons -->
|
||||
{#each group.capabilities.filter((c) => c !== "text") as cap}
|
||||
{#if cap === "thinking"}
|
||||
@@ -533,15 +523,6 @@
|
||||
{variant.quantization || "default"}
|
||||
</span>
|
||||
|
||||
{#if variant.requires_vllm}
|
||||
<span
|
||||
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
|
||||
title="Requires vLLM runtime"
|
||||
>
|
||||
vLLM
|
||||
</span>
|
||||
{/if}
|
||||
|
||||
<!-- Size -->
|
||||
<span
|
||||
class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
|
||||
@@ -647,7 +628,6 @@
|
||||
variants: [variant],
|
||||
smallestVariant: variant,
|
||||
hasMultipleVariants: false,
|
||||
requiresVllm: variant.requires_vllm === true,
|
||||
});
|
||||
}}
|
||||
title="View variant details"
|
||||
|
||||
@@ -22,7 +22,6 @@
|
||||
is_custom?: boolean;
|
||||
tasks?: string[];
|
||||
hugging_face_id?: string;
|
||||
requires_vllm?: boolean;
|
||||
}
|
||||
|
||||
interface ModelGroup {
|
||||
@@ -33,7 +32,6 @@
|
||||
variants: ModelInfo[];
|
||||
smallestVariant: ModelInfo;
|
||||
hasMultipleVariants: boolean;
|
||||
requiresVllm: boolean;
|
||||
}
|
||||
|
||||
interface FilterState {
|
||||
@@ -398,7 +396,6 @@
|
||||
variants: [],
|
||||
smallestVariant: model,
|
||||
hasMultipleVariants: false,
|
||||
requiresVllm: true,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -433,7 +430,6 @@
|
||||
(a.storage_size_megabytes || 0) - (b.storage_size_megabytes || 0),
|
||||
);
|
||||
group.hasMultipleVariants = group.variants.length > 1;
|
||||
group.requiresVllm = group.variants.every((v) => v.requires_vllm);
|
||||
}
|
||||
|
||||
// Convert to array and sort by smallest variant size (biggest first)
|
||||
@@ -591,7 +587,6 @@
|
||||
variants: [model],
|
||||
smallestVariant: model,
|
||||
hasMultipleVariants: false,
|
||||
requiresVllm: model.requires_vllm === true,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -1170,17 +1165,6 @@
|
||||
<span class="text-white/40">Variants:</span>
|
||||
<span class="text-white/70">{infoGroup.variants.length}</span>
|
||||
</div>
|
||||
{#if infoGroup.requiresVllm}
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="text-white/40">Runtime:</span>
|
||||
<span
|
||||
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 tracking-wider uppercase"
|
||||
>
|
||||
vLLM
|
||||
</span>
|
||||
<span class="text-white/40 text-[11px]">required</span>
|
||||
</div>
|
||||
{/if}
|
||||
{#if infoGroup.variants.length > 0}
|
||||
<div class="mt-3 pt-3 border-t border-exo-yellow/10">
|
||||
<span class="text-white/40">Available quantizations:</span>
|
||||
|
||||
@@ -219,7 +219,7 @@
|
||||
Prefill vs Decode
|
||||
</summary>
|
||||
<div class="mt-2 text-white/80 text-sm leading-relaxed">
|
||||
Prefill is the compute-bound pass that consumes the entire prompt and
|
||||
Prefill is the compute-heavy pass that consumes the entire prompt and
|
||||
builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
|
||||
tokens sequentially from that cache. The two phases have very different
|
||||
bottlenecks, so running them on different hardware can be substantially
|
||||
|
||||
@@ -117,10 +117,6 @@
|
||||
const LOGO_NATIVE_WIDTH = 814;
|
||||
const LOGO_NATIVE_HEIGHT = 1000;
|
||||
|
||||
// NVIDIA logo SVG path (from exo-nvidia)
|
||||
const NVIDIA_LOGO_PATH =
|
||||
"M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
|
||||
|
||||
function formatBytes(bytes: number, decimals = 1): string {
|
||||
if (!bytes || bytes === 0) return "0B";
|
||||
const k = 1024;
|
||||
@@ -558,13 +554,6 @@
|
||||
const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
|
||||
|
||||
const modelLower = modelId.toLowerCase();
|
||||
const identity = identitiesData[nodeInfo.id];
|
||||
const nameLower = (friendlyName || "").toLowerCase();
|
||||
const isSpark = modelLower.includes("dgx") || modelLower.includes("gx10");
|
||||
const isLinux =
|
||||
!isSpark &&
|
||||
(modelLower.startsWith("linux") || identity?.osVersion === "Linux");
|
||||
const isLinuxLaptop = isLinux && modelLower.includes("laptop");
|
||||
|
||||
// Check node states for styling
|
||||
const isHighlighted = highlightedNodes.has(nodeInfo.id);
|
||||
@@ -634,382 +623,7 @@
|
||||
`${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
|
||||
);
|
||||
|
||||
if (isSpark) {
|
||||
// NVIDIA DGX Spark — gold chassis with textured front, side handles, and NVIDIA badge
|
||||
iconBaseWidth = nodeRadius * 1.55;
|
||||
iconBaseHeight = nodeRadius * 0.58;
|
||||
const x = nodeInfo.x - iconBaseWidth / 2;
|
||||
const y = nodeInfo.y - iconBaseHeight / 2;
|
||||
const chassisX = x - iconBaseWidth * 0.03;
|
||||
const chassisWidth = iconBaseWidth * 1.05;
|
||||
const cornerRadius = 3;
|
||||
|
||||
const dgxClipId = `dgx-clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
|
||||
defs
|
||||
.append("clipPath")
|
||||
.attr("id", dgxClipId)
|
||||
.append("rect")
|
||||
.attr("x", x)
|
||||
.attr("y", y)
|
||||
.attr("width", iconBaseWidth)
|
||||
.attr("height", iconBaseHeight)
|
||||
.attr("rx", cornerRadius);
|
||||
|
||||
// Chassis texture pattern
|
||||
const textureId = `chassis-texture-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
|
||||
defs
|
||||
.append("pattern")
|
||||
.attr("id", textureId)
|
||||
.attr("patternUnits", "userSpaceOnUse")
|
||||
.attr("width", 8)
|
||||
.attr("height", 8);
|
||||
const texturePattern = defs.select(`#${textureId}`);
|
||||
texturePattern
|
||||
.append("rect")
|
||||
.attr("width", 8)
|
||||
.attr("height", 8)
|
||||
.attr("fill", "#6f6248");
|
||||
texturePattern
|
||||
.append("circle")
|
||||
.attr("cx", 2)
|
||||
.attr("cy", 2)
|
||||
.attr("r", 1)
|
||||
.attr("fill", "#5a4f3b")
|
||||
.attr("opacity", 0.5);
|
||||
texturePattern
|
||||
.append("circle")
|
||||
.attr("cx", 6)
|
||||
.attr("cy", 6)
|
||||
.attr("r", 1)
|
||||
.attr("fill", "#4a4232")
|
||||
.attr("opacity", 0.45);
|
||||
|
||||
// Main body
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("class", "node-outline")
|
||||
.attr("x", chassisX)
|
||||
.attr("y", y)
|
||||
.attr("width", chassisWidth)
|
||||
.attr("height", iconBaseHeight)
|
||||
.attr("rx", cornerRadius)
|
||||
.attr("fill", `url(#${textureId})`)
|
||||
.attr("stroke", wireColor)
|
||||
.attr("stroke-width", strokeWidth);
|
||||
|
||||
// Side border accents
|
||||
const sideThickness = iconBaseWidth * 0.02;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", chassisX)
|
||||
.attr("y", y)
|
||||
.attr("width", sideThickness)
|
||||
.attr("height", iconBaseHeight)
|
||||
.attr("fill", "#8a7a56");
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", chassisX + chassisWidth - sideThickness)
|
||||
.attr("y", y)
|
||||
.attr("width", sideThickness)
|
||||
.attr("height", iconBaseHeight)
|
||||
.attr("fill", "#8a7a56");
|
||||
|
||||
// Memory fill (bottom up)
|
||||
if (ramUsagePercent > 0) {
|
||||
const memFillHeight = (ramUsagePercent / 100) * iconBaseHeight;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", x)
|
||||
.attr("y", y + iconBaseHeight - memFillHeight)
|
||||
.attr("width", iconBaseWidth)
|
||||
.attr("height", memFillHeight)
|
||||
.attr("fill", "rgba(255,215,0,0.45)")
|
||||
.attr("clip-path", `url(#${dgxClipId})`);
|
||||
}
|
||||
|
||||
// Side handles with inner recess
|
||||
const handleWidth = iconBaseWidth * 0.27;
|
||||
const handleGap = iconBaseHeight * 0.05;
|
||||
const handleHeight = iconBaseHeight - handleGap * 2;
|
||||
const handleY = y + handleGap;
|
||||
const innerHandleWidth = iconBaseWidth * 0.12;
|
||||
const innerHandleHeight = handleHeight - iconBaseHeight * 0.06;
|
||||
const leftHandleX = x + 4;
|
||||
const rightHandleX = x + iconBaseWidth - handleWidth - 4;
|
||||
|
||||
// Left handle
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", leftHandleX)
|
||||
.attr("y", handleY)
|
||||
.attr("width", handleWidth)
|
||||
.attr("height", handleHeight)
|
||||
.attr("rx", 2.4)
|
||||
.attr("fill", "#b3a170")
|
||||
.attr("stroke", "#403723")
|
||||
.attr("stroke-width", 0.7);
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", leftHandleX + handleWidth * 0.06)
|
||||
.attr("y", handleY + iconBaseHeight * 0.03)
|
||||
.attr("width", innerHandleWidth)
|
||||
.attr("height", innerHandleHeight)
|
||||
.attr("rx", 1.6)
|
||||
.attr("fill", "#8a7a56");
|
||||
|
||||
// Right handle
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", rightHandleX)
|
||||
.attr("y", handleY)
|
||||
.attr("width", handleWidth)
|
||||
.attr("height", handleHeight)
|
||||
.attr("rx", 2.4)
|
||||
.attr("fill", "#b3a170")
|
||||
.attr("stroke", "#403723")
|
||||
.attr("stroke-width", 0.7);
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr(
|
||||
"x",
|
||||
rightHandleX + handleWidth - innerHandleWidth - handleWidth * 0.08,
|
||||
)
|
||||
.attr("y", handleY + iconBaseHeight * 0.03)
|
||||
.attr("width", innerHandleWidth)
|
||||
.attr("height", innerHandleHeight)
|
||||
.attr("rx", 1.6)
|
||||
.attr("fill", "#8a7a56");
|
||||
|
||||
// NVIDIA logo + text label (rotated 90 deg on left handle)
|
||||
const badgeWidth = iconBaseWidth * 0.09;
|
||||
const badgeHeight = handleHeight * 0.5;
|
||||
const badgeX =
|
||||
leftHandleX + handleWidth - badgeWidth - handleWidth * 0.06;
|
||||
const badgeY = handleY + (handleHeight - badgeHeight) / 2;
|
||||
const textSize = badgeWidth * 0.58;
|
||||
const logoWidth = textSize * 1.2;
|
||||
const logoHeight = logoWidth * (1.438 / 2.174);
|
||||
const centerX = badgeX + badgeWidth / 2 - badgeWidth * 0.03;
|
||||
const centerY = badgeY + badgeHeight / 2;
|
||||
const gap = badgeWidth * 0.15;
|
||||
const totalWidth = logoWidth + gap + textSize * 3.6;
|
||||
|
||||
const labelGroup = nodeG
|
||||
.append("g")
|
||||
.attr("transform", `rotate(90 ${centerX} ${centerY})`);
|
||||
|
||||
labelGroup
|
||||
.append("svg")
|
||||
.attr("x", centerX - totalWidth / 2)
|
||||
.attr("y", centerY - logoHeight / 2)
|
||||
.attr("width", logoWidth)
|
||||
.attr("height", logoHeight)
|
||||
.attr("viewBox", "0 0 2.174 1.438")
|
||||
.append("path")
|
||||
.attr("d", NVIDIA_LOGO_PATH)
|
||||
.attr("fill", "#76b900");
|
||||
|
||||
labelGroup
|
||||
.append("text")
|
||||
.attr("x", centerX - totalWidth / 2 + logoWidth + gap)
|
||||
.attr("y", centerY)
|
||||
.attr("text-anchor", "start")
|
||||
.attr("dominant-baseline", "middle")
|
||||
.attr("fill", "#8a7a56")
|
||||
.attr("font-size", textSize)
|
||||
.attr("font-family", "monospace")
|
||||
.attr("font-weight", "700")
|
||||
.text("NVIDIA");
|
||||
} else if (isLinuxLaptop) {
|
||||
// Linux Laptop — same shape as MacBook but with Tux logo
|
||||
iconBaseWidth = nodeRadius * 1.6;
|
||||
iconBaseHeight = nodeRadius * 1.15;
|
||||
const x = nodeInfo.x - iconBaseWidth / 2;
|
||||
const y = nodeInfo.y - iconBaseHeight / 2;
|
||||
|
||||
const screenHeight = iconBaseHeight * 0.7;
|
||||
const baseHeight = iconBaseHeight * 0.3;
|
||||
const screenWidth = iconBaseWidth * 0.85;
|
||||
const screenX = nodeInfo.x - screenWidth / 2;
|
||||
const screenBezel = 3;
|
||||
|
||||
const linuxScreenClipId = `linux-screen-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
|
||||
defs
|
||||
.append("clipPath")
|
||||
.attr("id", linuxScreenClipId)
|
||||
.append("rect")
|
||||
.attr("x", screenX + screenBezel)
|
||||
.attr("y", y + screenBezel)
|
||||
.attr("width", screenWidth - screenBezel * 2)
|
||||
.attr("height", screenHeight - screenBezel * 2)
|
||||
.attr("rx", 2);
|
||||
|
||||
// Screen outer frame
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("class", "node-outline")
|
||||
.attr("x", screenX)
|
||||
.attr("y", y)
|
||||
.attr("width", screenWidth)
|
||||
.attr("height", screenHeight)
|
||||
.attr("rx", 3)
|
||||
.attr("fill", "#1a1a1a")
|
||||
.attr("stroke", wireColor)
|
||||
.attr("stroke-width", strokeWidth);
|
||||
|
||||
// Screen inner
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", screenX + screenBezel)
|
||||
.attr("y", y + screenBezel)
|
||||
.attr("width", screenWidth - screenBezel * 2)
|
||||
.attr("height", screenHeight - screenBezel * 2)
|
||||
.attr("rx", 2)
|
||||
.attr("fill", "#0a0a12");
|
||||
|
||||
// Memory fill on screen
|
||||
if (ramUsagePercent > 0) {
|
||||
const memFillTotalHeight = screenHeight - screenBezel * 2;
|
||||
const memFillActualHeight =
|
||||
(ramUsagePercent / 100) * memFillTotalHeight;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", screenX + screenBezel)
|
||||
.attr(
|
||||
"y",
|
||||
y + screenBezel + (memFillTotalHeight - memFillActualHeight),
|
||||
)
|
||||
.attr("width", screenWidth - screenBezel * 2)
|
||||
.attr("height", memFillActualHeight)
|
||||
.attr("fill", "rgba(255,215,0,0.85)")
|
||||
.attr("clip-path", `url(#${linuxScreenClipId})`);
|
||||
}
|
||||
|
||||
// Terminal prompt on screen
|
||||
nodeG
|
||||
.append("text")
|
||||
.attr("x", nodeInfo.x)
|
||||
.attr("y", y + screenHeight / 2)
|
||||
.attr("text-anchor", "middle")
|
||||
.attr("dominant-baseline", "middle")
|
||||
.attr("fill", "#FFFFFF")
|
||||
.attr("opacity", 0.9)
|
||||
.attr("font-size", screenHeight * 0.25)
|
||||
.attr("font-family", "SF Mono, Monaco, monospace")
|
||||
.attr("font-weight", "700")
|
||||
.text(">_");
|
||||
|
||||
// Keyboard base (trapezoidal)
|
||||
const baseY = y + screenHeight;
|
||||
const baseTopWidth = screenWidth;
|
||||
const baseBottomWidth = iconBaseWidth;
|
||||
const baseTopX = nodeInfo.x - baseTopWidth / 2;
|
||||
const baseBottomX = nodeInfo.x - baseBottomWidth / 2;
|
||||
|
||||
nodeG
|
||||
.append("path")
|
||||
.attr(
|
||||
"d",
|
||||
`M ${baseTopX} ${baseY} L ${baseTopX + baseTopWidth} ${baseY} L ${baseBottomX + baseBottomWidth} ${baseY + baseHeight} L ${baseBottomX} ${baseY + baseHeight} Z`,
|
||||
)
|
||||
.attr("fill", "#2c2c2c")
|
||||
.attr("stroke", wireColor)
|
||||
.attr("stroke-width", 1);
|
||||
|
||||
// Keyboard area
|
||||
const keyboardX = baseTopX + 6;
|
||||
const keyboardY = baseY + 3;
|
||||
const keyboardWidth = baseTopWidth - 12;
|
||||
const keyboardHeight = baseHeight * 0.55;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", keyboardX)
|
||||
.attr("y", keyboardY)
|
||||
.attr("width", keyboardWidth)
|
||||
.attr("height", keyboardHeight)
|
||||
.attr("fill", "rgba(0,0,0,0.2)")
|
||||
.attr("rx", 2);
|
||||
|
||||
// Trackpad
|
||||
const trackpadWidth = baseTopWidth * 0.4;
|
||||
const trackpadX = nodeInfo.x - trackpadWidth / 2;
|
||||
const trackpadY = baseY + keyboardHeight + 5;
|
||||
const trackpadHeight = baseHeight * 0.3;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", trackpadX)
|
||||
.attr("y", trackpadY)
|
||||
.attr("width", trackpadWidth)
|
||||
.attr("height", trackpadHeight)
|
||||
.attr("fill", "rgba(255,255,255,0.08)")
|
||||
.attr("rx", 2);
|
||||
} else if (isLinux) {
|
||||
// Linux Desktop — same shape as Mac Studio but with Tux logo
|
||||
iconBaseWidth = nodeRadius * 1.25;
|
||||
iconBaseHeight = nodeRadius * 0.85;
|
||||
const x = nodeInfo.x - iconBaseWidth / 2;
|
||||
const y = nodeInfo.y - iconBaseHeight / 2;
|
||||
const cornerRadius = 4;
|
||||
const topSurfaceHeight = iconBaseHeight * 0.15;
|
||||
|
||||
const linuxDesktopClipId = `linux-desktop-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
|
||||
defs
|
||||
.append("clipPath")
|
||||
.attr("id", linuxDesktopClipId)
|
||||
.append("rect")
|
||||
.attr("x", x)
|
||||
.attr("y", y + topSurfaceHeight)
|
||||
.attr("width", iconBaseWidth)
|
||||
.attr("height", iconBaseHeight - topSurfaceHeight)
|
||||
.attr("rx", cornerRadius - 1);
|
||||
|
||||
// Main body
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("class", "node-outline")
|
||||
.attr("x", x)
|
||||
.attr("y", y)
|
||||
.attr("width", iconBaseWidth)
|
||||
.attr("height", iconBaseHeight)
|
||||
.attr("rx", cornerRadius)
|
||||
.attr("fill", "#1a1a1a")
|
||||
.attr("stroke", wireColor)
|
||||
.attr("stroke-width", strokeWidth);
|
||||
|
||||
// Memory fill
|
||||
if (ramUsagePercent > 0) {
|
||||
const memFillTotalHeight = iconBaseHeight - topSurfaceHeight;
|
||||
const memFillActualHeight =
|
||||
(ramUsagePercent / 100) * memFillTotalHeight;
|
||||
nodeG
|
||||
.append("rect")
|
||||
.attr("x", x)
|
||||
.attr(
|
||||
"y",
|
||||
y + topSurfaceHeight + (memFillTotalHeight - memFillActualHeight),
|
||||
)
|
||||
.attr("width", iconBaseWidth)
|
||||
.attr("height", memFillActualHeight)
|
||||
.attr("fill", "rgba(255,215,0,0.75)")
|
||||
.attr("clip-path", `url(#${linuxDesktopClipId})`);
|
||||
}
|
||||
|
||||
// Terminal prompt on front face
|
||||
nodeG
|
||||
.append("text")
|
||||
.attr("x", nodeInfo.x)
|
||||
.attr(
|
||||
"y",
|
||||
y + topSurfaceHeight + (iconBaseHeight - topSurfaceHeight) / 2,
|
||||
)
|
||||
.attr("text-anchor", "middle")
|
||||
.attr("dominant-baseline", "middle")
|
||||
.attr("fill", "rgba(255,255,255,0.5)")
|
||||
.attr("font-size", (iconBaseHeight - topSurfaceHeight) * 0.4)
|
||||
.attr("font-family", "SF Mono, Monaco, monospace")
|
||||
.attr("font-weight", "700")
|
||||
.text(">_");
|
||||
} else if (modelLower === "mac studio") {
|
||||
if (modelLower === "mac studio") {
|
||||
// Mac Studio - classic cube with memory fill
|
||||
iconBaseWidth = nodeRadius * 1.25;
|
||||
iconBaseHeight = nodeRadius * 0.85;
|
||||
@@ -1568,12 +1182,8 @@
|
||||
debugLabelY += debugLineHeight;
|
||||
}
|
||||
|
||||
const dbgIdentity = identitiesData[nodeInfo.id];
|
||||
if (dbgIdentity?.osVersion) {
|
||||
const osLabel =
|
||||
dbgIdentity.osVersion === "Linux"
|
||||
? "Linux"
|
||||
: `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
|
||||
const identity = identitiesData[nodeInfo.id];
|
||||
if (identity?.osVersion) {
|
||||
nodeG
|
||||
.append("text")
|
||||
.attr("x", nodeInfo.x)
|
||||
@@ -1582,7 +1192,9 @@
|
||||
.attr("fill", "rgba(179,179,179,0.7)")
|
||||
.attr("font-size", debugFontSize)
|
||||
.attr("font-family", "SF Mono, Monaco, monospace")
|
||||
.text(osLabel);
|
||||
.text(
|
||||
`macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -65,7 +65,6 @@
|
||||
nodeThunderboltBridge,
|
||||
nodeIdentities,
|
||||
isConnected,
|
||||
featureFlags,
|
||||
type DownloadProgress,
|
||||
type PlacementPreview,
|
||||
} from "$lib/stores/app.svelte";
|
||||
@@ -703,10 +702,7 @@
|
||||
? Object.keys(topologyData()!.nodes).length
|
||||
: 1;
|
||||
const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
|
||||
const instanceType =
|
||||
nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
|
||||
? "MlxRing"
|
||||
: selectedInstanceType;
|
||||
const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
|
||||
try {
|
||||
const placementResponse = await fetch(
|
||||
`/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
|
||||
@@ -787,7 +783,6 @@
|
||||
quantization?: string;
|
||||
base_model?: string;
|
||||
capabilities?: string[];
|
||||
requires_vllm?: boolean;
|
||||
}>
|
||||
>([]);
|
||||
type ModelMemoryFitStatus =
|
||||
@@ -891,7 +886,7 @@
|
||||
}
|
||||
|
||||
let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
|
||||
type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
|
||||
type InstanceMeta = "MlxRing" | "MlxJaccl";
|
||||
|
||||
// Launch defaults persistence
|
||||
const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
|
||||
@@ -937,12 +932,7 @@
|
||||
// Apply sharding and instance type unconditionally
|
||||
selectedSharding = defaults.sharding;
|
||||
selectedInstanceType =
|
||||
defaults.instanceType === "MlxRing"
|
||||
? "MlxRing"
|
||||
: defaults.instanceType === "Vllm"
|
||||
? "Vllm"
|
||||
: "MlxJaccl";
|
||||
userPickedInstanceType = true;
|
||||
defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
|
||||
|
||||
// Apply minNodes if valid (between 1 and maxNodes)
|
||||
if (
|
||||
@@ -964,23 +954,6 @@
|
||||
}
|
||||
|
||||
let selectedInstanceType = $state<InstanceMeta>("MlxRing");
|
||||
let userPickedInstanceType = $state(false);
|
||||
$effect(() => {
|
||||
if (!userPickedInstanceType && featureFlags()["vllm_available"]) {
|
||||
selectedInstanceType = "Vllm";
|
||||
}
|
||||
});
|
||||
const selectedModelRequiresVllm = $derived.by((): boolean => {
|
||||
const id = selectedPreviewModelId();
|
||||
if (!id) return false;
|
||||
const model = models.find((m) => m.id === id);
|
||||
return model?.requires_vllm === true;
|
||||
});
|
||||
$effect(() => {
|
||||
if (selectedModelRequiresVllm) {
|
||||
selectedInstanceType = "Vllm";
|
||||
}
|
||||
});
|
||||
let selectedMinNodes = $state<number>(1);
|
||||
let minNodesInitialized = $state(false);
|
||||
let launchingModelId = $state<string | null>(null);
|
||||
@@ -1173,7 +1146,9 @@
|
||||
}
|
||||
|
||||
const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
|
||||
runtime === selectedInstanceType;
|
||||
selectedInstanceType === "MlxRing"
|
||||
? runtime === "MlxRing"
|
||||
: runtime === "MlxJaccl";
|
||||
|
||||
// Helper to check if a model can be launched (has valid placement with >= minNodes)
|
||||
function canModelFit(modelId: string): boolean {
|
||||
@@ -2088,7 +2063,6 @@
|
||||
let instanceType = "Unknown";
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
|
||||
|
||||
const inst = instance as {
|
||||
shardAssignments?: {
|
||||
@@ -3461,6 +3435,7 @@
|
||||
>
|
||||
<li>Connect nodes with TB5 cables</li>
|
||||
<li>Boot to Recovery (hold power 10s → Options)</li>
|
||||
<li>Open Terminal from the Utilities menu</li>
|
||||
<li>
|
||||
Run
|
||||
<code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
|
||||
@@ -4848,6 +4823,7 @@
|
||||
>
|
||||
<li>Connect nodes with TB5 cables</li>
|
||||
<li>Boot to Recovery (hold power 10s → Options)</li>
|
||||
<li>Open Terminal from the Utilities menu</li>
|
||||
<li>
|
||||
Run
|
||||
<code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
|
||||
@@ -4994,6 +4970,7 @@
|
||||
>
|
||||
<li>Connect nodes with TB5 cables</li>
|
||||
<li>Boot to Recovery (hold power 10s → Options)</li>
|
||||
<li>Open Terminal from the Utilities menu</li>
|
||||
<li>
|
||||
Run
|
||||
<code
|
||||
@@ -5795,18 +5772,14 @@
|
||||
</div>
|
||||
<div class="flex gap-2">
|
||||
<button
|
||||
disabled={selectedModelRequiresVllm}
|
||||
onclick={() => {
|
||||
if (selectedModelRequiresVllm) return;
|
||||
selectedInstanceType = "MlxRing";
|
||||
userPickedInstanceType = true;
|
||||
saveLaunchDefaults();
|
||||
}}
|
||||
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
|
||||
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
|
||||
: selectedInstanceType === 'MlxRing'
|
||||
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
|
||||
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
|
||||
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
|
||||
'MlxRing'
|
||||
? 'bg-transparent text-exo-yellow border-exo-yellow'
|
||||
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
|
||||
>
|
||||
<span
|
||||
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
|
||||
@@ -5822,18 +5795,14 @@
|
||||
TCP/IP
|
||||
</button>
|
||||
<button
|
||||
disabled={selectedModelRequiresVllm}
|
||||
onclick={() => {
|
||||
if (selectedModelRequiresVllm) return;
|
||||
selectedInstanceType = "MlxJaccl";
|
||||
userPickedInstanceType = true;
|
||||
saveLaunchDefaults();
|
||||
}}
|
||||
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
|
||||
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
|
||||
: selectedInstanceType === 'MlxJaccl'
|
||||
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
|
||||
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
|
||||
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
|
||||
'MlxJaccl'
|
||||
? 'bg-transparent text-exo-yellow border-exo-yellow'
|
||||
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
|
||||
>
|
||||
<span
|
||||
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
|
||||
@@ -5848,41 +5817,7 @@
|
||||
</span>
|
||||
RDMA (Fast)
|
||||
</button>
|
||||
{#if featureFlags()["vllm_available"] || selectedModelRequiresVllm}
|
||||
<button
|
||||
onclick={() => {
|
||||
selectedInstanceType = "Vllm";
|
||||
userPickedInstanceType = true;
|
||||
saveLaunchDefaults();
|
||||
}}
|
||||
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
|
||||
'Vllm'
|
||||
? 'bg-transparent text-exo-yellow border-exo-yellow'
|
||||
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
|
||||
>
|
||||
<span
|
||||
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
|
||||
'Vllm'
|
||||
? 'border-exo-yellow'
|
||||
: 'border-exo-medium-gray'}"
|
||||
>
|
||||
{#if selectedInstanceType === "Vllm"}
|
||||
<span
|
||||
class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
|
||||
></span>
|
||||
{/if}
|
||||
</span>
|
||||
vLLM (CUDA)
|
||||
</button>
|
||||
{/if}
|
||||
</div>
|
||||
{#if selectedModelRequiresVllm}
|
||||
<div
|
||||
class="mt-2 text-[11px] font-mono text-orange-300/80"
|
||||
>
|
||||
This model requires vLLM.
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
<!-- Minimum Devices -->
|
||||
|
||||
Generated
+6
-6
@@ -47,11 +47,11 @@
|
||||
"rust-analyzer-src": "rust-analyzer-src"
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1775807984,
|
||||
"narHash": "sha256-Redoe3D9zGN5I9QPHWL9vfMVQBehY1fKsMiRXQ83X3w=",
|
||||
"lastModified": 1777708550,
|
||||
"narHash": "sha256-Qif3UXT0l5OQq8H9pRWt4/ia4gF48MWK2oHKL8uVx8U=",
|
||||
"owner": "nix-community",
|
||||
"repo": "fenix",
|
||||
"rev": "fcf90c0c4d368b2ca917a7afa6d08e98a397e5fd",
|
||||
"rev": "74c1591efaff494756b8d35ebe357c6c2bbdca96",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@@ -218,11 +218,11 @@
|
||||
"rust-analyzer-src": {
|
||||
"flake": false,
|
||||
"locked": {
|
||||
"lastModified": 1775745684,
|
||||
"narHash": "sha256-8MbfLwd60FNa8dRFkjE+G3TT/x21G3Rsplm1bMBQUtU=",
|
||||
"lastModified": 1777639980,
|
||||
"narHash": "sha256-6d7Hdurvbjc5uwJuc0YiK7rZBGj6Gs3uzfBFcTs+xCc=",
|
||||
"owner": "rust-lang",
|
||||
"repo": "rust-analyzer",
|
||||
"rev": "64ddb549bc9a70d011328746fa46a8883f937b6b",
|
||||
"rev": "64cdaeb06f69b6b769a492edd88b022ae88e8ca2",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
|
||||
@@ -110,7 +110,7 @@
|
||||
nixpkgs-fmt.enable = true;
|
||||
ruff-format = {
|
||||
enable = true;
|
||||
excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
|
||||
excludes = [ "rust/exo_net/exo_net.pyi" ];
|
||||
};
|
||||
rustfmt = {
|
||||
enable = true;
|
||||
@@ -146,7 +146,7 @@
|
||||
config.treefmt.build.wrapper
|
||||
|
||||
# PYTHON
|
||||
self'.packages.exo.passthru.evenv
|
||||
#self'.packages.editableVenv
|
||||
uv
|
||||
|
||||
# RUST
|
||||
|
||||
@@ -22,8 +22,8 @@ sync-clean:
|
||||
uv sync --all-packages --force-reinstall --no-cache
|
||||
|
||||
rust-rebuild:
|
||||
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
|
||||
uv sync --reinstall-package exo_pyo3_bindings
|
||||
cargo run --bin stub_gen
|
||||
uv sync --reinstall-package exo_net
|
||||
|
||||
build-dashboard:
|
||||
#!/usr/bin/env bash
|
||||
@@ -40,19 +40,6 @@ build-app: rust-rebuild sync-clean package
|
||||
xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
|
||||
@echo "\nBuild complete. Run with:\n open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
|
||||
|
||||
sync-cuda:
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
uv sync --extra vllm-cuda13 --extra mlx-cpu --no-install-package vllm
|
||||
dest=".venv/lib/python3.13/site-packages"
|
||||
[[ -d $dest/vllm ]] || {
|
||||
nix build .#exo-cuda-13.passthru.evenv
|
||||
# will also grab vllm-0.19.1-distinfo
|
||||
cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
|
||||
chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
|
||||
rm result
|
||||
}
|
||||
|
||||
clean:
|
||||
rm -rf **/__pycache__
|
||||
rm -rf target/
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
diff --git a/setup.py b/setup.py
|
||||
index 6dc2ed028..bdcc6354a 100644
|
||||
--- a/setup.py
|
||||
+++ b/setup.py
|
||||
@@ -18,6 +18,13 @@ from setuptools import Extension, setup
|
||||
from setuptools.command.build_ext import build_ext
|
||||
|
||||
|
||||
+if "NIX_ATTRS_JSON_FILE" in os.environ:
|
||||
+ with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
|
||||
+ NIX_ATTRS = json.load(f)
|
||||
+else:
|
||||
+ NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
|
||||
+
|
||||
+
|
||||
def load_module_from_path(module_name, path):
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
|
||||
cmake_args = [
|
||||
"-DCMAKE_BUILD_TYPE={}".format(cfg),
|
||||
"-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
|
||||
+ *NIX_ATTRS["cmakeFlags"],
|
||||
]
|
||||
|
||||
verbose = envs.VERBOSE
|
||||
+38
-59
@@ -15,18 +15,21 @@ dependencies = [
|
||||
"huggingface-hub>=1.8.0",
|
||||
"psutil>=7.0.0",
|
||||
"loguru>=0.7.3",
|
||||
"exo-pyo3-bindings", # rust bindings
|
||||
"exo-net", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"mlx==0.31.2; sys_platform == 'darwin'",
|
||||
"mlx-lm; sys_platform=='darwin'",
|
||||
"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",
|
||||
"mflux==0.17.2; sys_platform == 'darwin'",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
"mlx-vlm>=0.3.11; sys_platform == 'darwin'",
|
||||
"transformers>=5.6.2",
|
||||
"nvidia-ml-py>=13.595.45",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -45,30 +48,26 @@ dev = [
|
||||
|
||||
[project.optional-dependencies]
|
||||
build = ["nanobind"]
|
||||
mlx-none = ["anyio"]
|
||||
mlx = [
|
||||
"mlx==0.31.2",
|
||||
"mlx-lm",
|
||||
"mlx-vlm>=0.3.11",
|
||||
"mflux==0.17.5",
|
||||
# pinning vllms versions for consistency.
|
||||
"torch==2.10.0; sys_platform == 'darwin'",
|
||||
"torch==2.10.0; sys_platform == 'linux'",
|
||||
"torchaudio==2.10.0; sys_platform == 'darwin'",
|
||||
"torchaudio==2.10.0; sys_platform == 'linux'",
|
||||
"torchvision==0.25.0; sys_platform == 'darwin'",
|
||||
"torchvision==0.25.0; sys_platform == 'linux'",
|
||||
|
||||
cpu = [
|
||||
"mlx==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-cpu==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-lm; sys_platform == 'linux'",
|
||||
"mlx-vlm>=0.3.11; sys_platform== 'linux'",
|
||||
"torch>=2.10.0; sys_platform == 'linux'",
|
||||
]
|
||||
mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
|
||||
mlx-cuda12 = ["exo[mlx]", "mlx-cuda-12==0.31.1; sys_platform == 'linux'"]
|
||||
mlx-cuda13 = ["exo[mlx]", "mlx-cuda-13==0.31.1; sys_platform == 'linux'"]
|
||||
vllm-none = ["anyio"]
|
||||
vllm-cuda13 = [
|
||||
"vllm[cuda13, fastsafetensors]; sys_platform == 'linux'",
|
||||
"torch==2.10.0; sys_platform == 'linux'",
|
||||
"torchaudio==2.10.0; sys_platform == 'linux'",
|
||||
"torchvision==0.25.0; sys_platform == 'linux'",
|
||||
cuda12 = [
|
||||
"mlx==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-cuda-12==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-lm; sys_platform == 'linux'",
|
||||
"mlx-vlm>=0.3.11; sys_platform== 'linux'",
|
||||
"torch>=2.10.0; sys_platform == 'linux'",
|
||||
]
|
||||
cuda13 = [
|
||||
"mlx==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-cuda-13==0.31.1; sys_platform == 'linux'",
|
||||
"mlx-lm; sys_platform == 'linux'",
|
||||
"mlx-vlm>=0.3.11; sys_platform== 'linux'",
|
||||
"torch>=2.10.0; sys_platform == 'linux'",
|
||||
]
|
||||
|
||||
###
|
||||
@@ -76,29 +75,18 @@ vllm-cuda13 = [
|
||||
###
|
||||
|
||||
[tool.uv.workspace]
|
||||
members = ["rust/exo_pyo3_bindings", "bench"]
|
||||
members = ["rust/exo_net", "bench"]
|
||||
|
||||
[tool.uv.sources]
|
||||
exo-pyo3-bindings = { workspace = true }
|
||||
exo-net = { 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/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
|
||||
mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
|
||||
vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
|
||||
torch = [
|
||||
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
|
||||
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
|
||||
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
|
||||
]
|
||||
torchvision = [
|
||||
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
|
||||
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
|
||||
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
|
||||
]
|
||||
torchaudio = [
|
||||
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
|
||||
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
|
||||
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
|
||||
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
|
||||
{ index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
|
||||
{ index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
|
||||
]
|
||||
vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cu130"
|
||||
@@ -106,8 +94,8 @@ url = "https://download.pytorch.org/whl/cu130"
|
||||
explicit = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cu128"
|
||||
url = "https://download.pytorch.org/whl/cu128"
|
||||
name = "pytorch-cu120"
|
||||
url = "https://download.pytorch.org/whl/cu120"
|
||||
explicit = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
@@ -168,19 +156,11 @@ root = "src"
|
||||
required-version = ">=0.8.6"
|
||||
prerelease = "allow"
|
||||
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
|
||||
override-dependencies = ["opencv-python; python_version < '0'"]
|
||||
conflicts = [
|
||||
[
|
||||
{ extra = "mlx-cuda13" },
|
||||
{ extra = "mlx-cuda12" },
|
||||
{ extra = "mlx-cpu" },
|
||||
{ extra = "mlx-none" },
|
||||
],
|
||||
[
|
||||
{ extra = "vllm-cuda13" },
|
||||
{ extra = "mlx-cuda12" },
|
||||
{ extra = "vllm-none" },
|
||||
],
|
||||
conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
|
||||
constraint-dependencies = ["transformers>=5.6.2"]
|
||||
override-dependencies = [
|
||||
"mlx==0.31.1; sys_platform=='linux'",
|
||||
"mlx; sys_platform=='darwin'",
|
||||
]
|
||||
|
||||
[tool.uv.extra-build-dependencies]
|
||||
@@ -195,7 +175,6 @@ mlx = [
|
||||
"ninja",
|
||||
]
|
||||
mlx-lm = ["setuptools"]
|
||||
mflux = ["uv_build"]
|
||||
xgrammar = [
|
||||
"nanobind",
|
||||
"setuptools",
|
||||
|
||||
+26
-207
@@ -10,10 +10,8 @@ let
|
||||
inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
|
||||
inherit (pkgs.config) cudaSupport;
|
||||
inherit (pkgs) cudaPackages;
|
||||
libmlx_source =
|
||||
if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
|
||||
else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
|
||||
else "mlx-cpu";
|
||||
cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
|
||||
libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
|
||||
python = pkgs.python313;
|
||||
cudaLibs = with cudaPackages; [
|
||||
cuda_cudart
|
||||
@@ -37,17 +35,18 @@ let
|
||||
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
|
||||
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
|
||||
# Copy .pyi stub + py.typed marker so basedpyright can find the types.
|
||||
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
|
||||
pname = "exo-pyo3-bindings";
|
||||
exo-net = pkgs.stdenv.mkDerivation {
|
||||
pname = "exo-net";
|
||||
version = "0.1.0";
|
||||
src = self'.packages.exo_pyo3_bindings;
|
||||
src = self'.packages.exo-net;
|
||||
# Install from pre-built wheel
|
||||
nativeBuildInputs = [ final.pyprojectWheelHook ];
|
||||
dontStrip = true;
|
||||
passthru = prev.exo-pyo3-bindings.passthru or { };
|
||||
|
||||
postInstall = ''
|
||||
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
|
||||
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
|
||||
local siteDir=$out/${final.python.sitePackages}/exo_net
|
||||
cp ${inputs.self}/rust/exo_net/exo_net.pyi $siteDir/
|
||||
touch $siteDir/py.typed
|
||||
'';
|
||||
};
|
||||
@@ -115,213 +114,37 @@ let
|
||||
});
|
||||
} // lib.optionalAttrs isLinux {
|
||||
mlx = prev.mlx.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
|
||||
postInstall = ''
|
||||
cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
|
||||
'';
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
} // lib.optionalAttrs cudaSupport {
|
||||
"${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ cudaLibs;
|
||||
buildInputs = old.buildInputs ++ [ cudaLibs ];
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
torch = prev.torch.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
});
|
||||
torchaudio = prev.torchaudio.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
|
||||
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
|
||||
});
|
||||
torchvision = prev.torchvision.overrideAttrs (old: {
|
||||
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
|
||||
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
|
||||
});
|
||||
|
||||
torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
|
||||
buildInputs = old.buildInputs ++ cudaLibs;
|
||||
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
|
||||
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
|
||||
});
|
||||
# Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
|
||||
vllm = prev.vllm.overrideAttrs (old:
|
||||
let
|
||||
cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
|
||||
mkdir -p $out/include
|
||||
ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
|
||||
'';
|
||||
|
||||
cudaRoot = pkgs.symlinkJoin {
|
||||
name = "cuda-merged-exo";
|
||||
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
|
||||
};
|
||||
|
||||
cutlass = pkgs.fetchFromGitHub {
|
||||
name = "cutlass-source";
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
tag = "v4.2.1";
|
||||
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
|
||||
};
|
||||
triton-kernels = pkgs.fetchFromGitHub {
|
||||
owner = "triton-lang";
|
||||
repo = "triton";
|
||||
tag = "v3.6.0";
|
||||
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
|
||||
};
|
||||
|
||||
cutlass-flashmla = pkgs.fetchFromGitHub {
|
||||
owner = "NVIDIA";
|
||||
repo = "cutlass";
|
||||
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
|
||||
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
|
||||
};
|
||||
|
||||
flashmla = pkgs.stdenv.mkDerivation {
|
||||
pname = "flashmla";
|
||||
version = "1.0.0";
|
||||
|
||||
src = pkgs.fetchFromGitHub {
|
||||
name = "FlashMLA-source";
|
||||
owner = "vllm-project";
|
||||
repo = "FlashMLA";
|
||||
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
|
||||
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
|
||||
};
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass-flashmla} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
qutlass = pkgs.fetchFromGitHub {
|
||||
name = "qutlass-source";
|
||||
owner = "IST-DASLab";
|
||||
repo = "qutlass";
|
||||
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
|
||||
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
|
||||
};
|
||||
vllm-flash-attn = pkgs.stdenv.mkDerivation {
|
||||
pname = "vllm-flash-attn";
|
||||
version = "2.7.2.post1";
|
||||
|
||||
src = pkgs.fetchFromGitHub {
|
||||
name = "flash-attention-source";
|
||||
owner = "vllm-project";
|
||||
repo = "flash-attention";
|
||||
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
|
||||
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
|
||||
};
|
||||
|
||||
patches = [
|
||||
(pkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
|
||||
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
|
||||
})
|
||||
(pkgs.fetchpatch {
|
||||
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
|
||||
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
|
||||
})
|
||||
];
|
||||
|
||||
dontConfigure = true;
|
||||
|
||||
buildPhase = ''
|
||||
rm -rf csrc/cutlass
|
||||
ln -sf ${cutlass} csrc/cutlass
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
cp -rva . $out
|
||||
'';
|
||||
};
|
||||
in
|
||||
{
|
||||
patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
pkgs.cmake
|
||||
pkgs.ninja
|
||||
pkgs.autoAddDriverRunpath
|
||||
] ++ lib.optionals cudaSupport [
|
||||
cudaPackages.cuda_nvcc
|
||||
];
|
||||
# TODO: vllm rocm/cpu
|
||||
VLLM_TARGET_DEVICE = "empty";
|
||||
preConfigure = ''
|
||||
export MAX_JOBS="$NIX_BUILD_CORES"
|
||||
'';
|
||||
|
||||
# TODO: vllm non cuda13 support, more arch's, etc.
|
||||
} // lib.optionalAttrs cudaSupport {
|
||||
buildInputs = cudaLibs ++ [ cudaRoot ];
|
||||
|
||||
VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
|
||||
CUDA_HOME = "${cudaRoot}";
|
||||
CUDAToolkit_ROOT = "${cudaRoot}";
|
||||
CUDACXX = "${cudaRoot}/bin/nvcc";
|
||||
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
|
||||
VLLM_TARGET_DEVICE = "cuda";
|
||||
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
|
||||
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
|
||||
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
|
||||
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
|
||||
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
|
||||
CAFFE2_USE_CUDNN = "ON";
|
||||
CAFFE2_USE_CUFILE = "ON";
|
||||
CUTLASS_ENABLE_CUBLAS = "ON";
|
||||
CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
|
||||
|
||||
cmakeFlags = [
|
||||
(lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
|
||||
(lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
|
||||
(lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
|
||||
(lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
|
||||
(lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
|
||||
(lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
|
||||
(lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
|
||||
(lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
|
||||
(lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
|
||||
(lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
|
||||
(lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
|
||||
(lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
|
||||
(lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
|
||||
(lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
|
||||
(lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
|
||||
(lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
|
||||
];
|
||||
});
|
||||
|
||||
} // lib.optionalAttrs (cudaSupport && isx86_64) {
|
||||
numba = prev.numba.overrideAttrs (old: {
|
||||
buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
|
||||
});
|
||||
};
|
||||
pyprojectOverlay = workspace.mkPyprojectOverlay {
|
||||
sourcePreference = "wheel";
|
||||
@@ -342,28 +165,24 @@ let
|
||||
buildSystemsOverlay
|
||||
]
|
||||
);
|
||||
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
|
||||
venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
|
||||
mkApp = text: name: pkgs.writeShellApplication {
|
||||
venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
|
||||
mkApp = cmd: name: pkgs.writeShellApplication {
|
||||
inherit name;
|
||||
text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
|
||||
runtimeEnv = {
|
||||
EXO_DASHBOARD_DIR = self'.packages.dashboard;
|
||||
EXO_RESOURCES_DIR = inputs.self + /resources;
|
||||
};
|
||||
runtimeInputs = [
|
||||
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
|
||||
(venv name)
|
||||
pkgs.nix-gl-host
|
||||
]
|
||||
++ lib.optionals isDarwin [ pkgs.macmon ];
|
||||
passthru = {
|
||||
venv = venv name;
|
||||
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
|
||||
};
|
||||
text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
|
||||
};
|
||||
in
|
||||
{
|
||||
inherit venv;
|
||||
editablePythonSet = pythonSet.overrideScope editableOverlay;
|
||||
mkPythonScript = path: mkApp ''python ${path} "$@"'';
|
||||
mkExo = mkApp ''exo "$@"'';
|
||||
};
|
||||
@@ -373,18 +192,18 @@ in
|
||||
{ self', pkgs, unfreePkgs, lib, ... }:
|
||||
let
|
||||
inherit (pkgs.stdenv.hostPlatform) isLinux;
|
||||
inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
|
||||
inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
|
||||
|
||||
# Virtual environment with dev dependencies for testing
|
||||
testVenv = (mkPythonSet {
|
||||
inherit self' pkgs lib; members = {
|
||||
exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
|
||||
exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
|
||||
};
|
||||
}).venv "exo-test";
|
||||
|
||||
mkBenchScript = (mkPythonSet {
|
||||
inherit self' pkgs lib; members = {
|
||||
exo = [ "mlx-cpu" "vllm-none" ];
|
||||
exo = [ "cpu" ];
|
||||
exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
|
||||
};
|
||||
}).mkPythonScript;
|
||||
@@ -394,12 +213,12 @@ in
|
||||
runtimeInputs = [ pkgs.python313 ];
|
||||
text = ''exec python ${path} "$@"'';
|
||||
};
|
||||
cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
|
||||
cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
|
||||
|
||||
in
|
||||
{
|
||||
packages = {
|
||||
exo = mkExo "exo";
|
||||
editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
|
||||
# for running tests in ci
|
||||
exo-test-env = testVenv;
|
||||
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
|
||||
@@ -408,8 +227,8 @@ in
|
||||
# used by ./tests/run_exo_on.sh
|
||||
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
|
||||
} // lib.optionalAttrs isLinux {
|
||||
exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
|
||||
exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
|
||||
exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
|
||||
exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
|
||||
};
|
||||
|
||||
checks = {
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
model_id = "2imi9/gpt-oss-20B-NVFP4A16-BF16"
|
||||
n_layers = 24
|
||||
hidden_size = 2880
|
||||
num_key_value_heads = 8
|
||||
supports_tensor = false
|
||||
tasks = ["TextGeneration"]
|
||||
family = "gpt-oss"
|
||||
quantization = "nvfp4"
|
||||
base_model = "GPT-OSS 20B"
|
||||
capabilities = ["text", "thinking"]
|
||||
reasoning_dialect = "channel"
|
||||
context_length = 131072
|
||||
requires_vllm = true
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 41829514752
|
||||
|
||||
[sampling_defaults]
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
top_k = 0
|
||||
@@ -1,27 +0,0 @@
|
||||
model_id = "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
n_layers = 48
|
||||
hidden_size = 2048
|
||||
num_key_value_heads = 4
|
||||
supports_tensor = false
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "nvfp4"
|
||||
base_model = "Qwen3 30B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
context_length = 32768
|
||||
requires_vllm = true
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 18087458688
|
||||
|
||||
[sampling_defaults]
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
top_k = 20
|
||||
min_p = 0.0
|
||||
|
||||
[sampling_defaults.non_thinking]
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
top_k = 20
|
||||
min_p = 0.0
|
||||
@@ -1,20 +0,0 @@
|
||||
model_id = "openai/gpt-oss-120b"
|
||||
n_layers = 36
|
||||
hidden_size = 2880
|
||||
num_key_value_heads = 8
|
||||
supports_tensor = false
|
||||
tasks = ["TextGeneration"]
|
||||
family = "gpt-oss"
|
||||
quantization = "mxfp4"
|
||||
base_model = "GPT-OSS 120B"
|
||||
capabilities = ["text", "thinking"]
|
||||
reasoning_dialect = "channel"
|
||||
context_length = 131072
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 65248815744
|
||||
|
||||
[sampling_defaults]
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
top_k = 0
|
||||
@@ -1,32 +0,0 @@
|
||||
model_id = "sakamakismile/Qwen3.6-27B-NVFP4"
|
||||
n_layers = 64
|
||||
hidden_size = 5120
|
||||
num_key_value_heads = 4
|
||||
supports_tensor = false
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "nvfp4"
|
||||
base_model = "Qwen3.6 27B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle", "vision"]
|
||||
reasoning_dialect = "post_last_user"
|
||||
context_length = 262144
|
||||
requires_vllm = true
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 16703361232
|
||||
|
||||
[sampling_defaults]
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
top_k = 20
|
||||
min_p = 0.0
|
||||
repetition_penalty = 1.0
|
||||
presence_penalty = 1.5
|
||||
|
||||
[sampling_defaults.non_thinking]
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
top_k = 20
|
||||
min_p = 0.0
|
||||
repetition_penalty = 1.0
|
||||
presence_penalty = 1.5
|
||||
@@ -0,0 +1,52 @@
|
||||
[package]
|
||||
name = "exo_net"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
path = "src/lib.rs"
|
||||
name = "exo_net"
|
||||
|
||||
# "cdylib" needed to produce shared library for Python to import
|
||||
# "rlib" needed for stub-gen to run
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[[bin]]
|
||||
path = "src/bin/stub_gen.rs"
|
||||
name = "stub_gen"
|
||||
doc = false
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
networking.workspace = true
|
||||
extend.workspace = true
|
||||
|
||||
# interop
|
||||
pyo3 = { workspace = true, features = ["experimental-async"] }
|
||||
pyo3-stub-gen.workspace = true
|
||||
pyo3-async-runtimes = { workspace = true, features = [
|
||||
"attributes",
|
||||
"tokio-runtime",
|
||||
"testing",
|
||||
] }
|
||||
pyo3-log.workspace = true
|
||||
|
||||
# async runtime
|
||||
tokio = { workspace = true, features = ["full", "tracing"] }
|
||||
futures-lite.workspace = true
|
||||
pin-project.workspace = true
|
||||
|
||||
# Tracing
|
||||
log.workspace = true
|
||||
env_logger.workspace = true
|
||||
|
||||
# Networking
|
||||
zenoh.workspace = true
|
||||
rand.workspace = true
|
||||
serde_json.workspace = true
|
||||
parking_lot.workspace = true
|
||||
tracing.workspace = true
|
||||
File renamed without changes.
@@ -0,0 +1,76 @@
|
||||
# This file is automatically generated by pyo3_stub_gen
|
||||
# ruff: noqa: E501, F401, F403, F405
|
||||
|
||||
import builtins
|
||||
import collections.abc
|
||||
import typing
|
||||
|
||||
__all__ = [
|
||||
"NetReceiver",
|
||||
"NetSender",
|
||||
"NetworkingHandle",
|
||||
"PyFromSwarm",
|
||||
"PySession",
|
||||
"StateProxy",
|
||||
]
|
||||
|
||||
@typing.final
|
||||
class NetReceiver:
|
||||
def recv(self) -> collections.abc.Awaitable[bytes | None]: ...
|
||||
|
||||
@typing.final
|
||||
class NetSender:
|
||||
def send(self, data: bytes) -> collections.abc.Awaitable[bool]: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
@staticmethod
|
||||
def new(identity: bytes, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> tuple[NetworkingHandle, PySession]: ...
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
|
||||
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
"""
|
||||
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Unsubscribes from a `GossipSub` topic.
|
||||
|
||||
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
"""
|
||||
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
|
||||
r"""
|
||||
Publishes a message with multiple topics to the `GossipSub` network.
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
class PyFromSwarm:
|
||||
@typing.final
|
||||
class Connection(PyFromSwarm):
|
||||
__match_args__ = ("connected",)
|
||||
@property
|
||||
def connected(self) -> builtins.bool: ...
|
||||
def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
|
||||
|
||||
@typing.final
|
||||
class Message(PyFromSwarm):
|
||||
__match_args__ = ("topic", "data",)
|
||||
@property
|
||||
def topic(self) -> builtins.str: ...
|
||||
@property
|
||||
def data(self) -> bytes: ...
|
||||
def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
|
||||
|
||||
|
||||
@typing.final
|
||||
class PySession:
|
||||
def net_receiver(self, key: builtins.str) -> NetReceiver: ...
|
||||
def net_sender(self, key: builtins.str) -> NetSender: ...
|
||||
def state_proxy(self) -> StateProxy: ...
|
||||
|
||||
@typing.final
|
||||
class StateProxy:
|
||||
def snapshot(self) -> collections.abc.Awaitable[str]: ...
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# This file is automatically generated by pyo3_stub_gen
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import builtins
|
||||
import collections.abc
|
||||
import typing
|
||||
|
||||
@typing.final
|
||||
class NetReceiver: ...
|
||||
|
||||
@typing.final
|
||||
class NetSender: ...
|
||||
|
||||
@typing.final
|
||||
class PySession:
|
||||
@staticmethod
|
||||
def init() -> collections.abc.Awaitable[PySession]: ...
|
||||
def net_receiver(self, key: builtins.str) -> NetReceiver: ...
|
||||
def net_sender(self, key: builtins.str) -> NetSender: ...
|
||||
def state_proxy(self) -> StateProxy: ...
|
||||
|
||||
@typing.final
|
||||
class StateProxy:
|
||||
def snapshot(self) -> collections.abc.Awaitable[str]: ...
|
||||
@@ -3,24 +3,22 @@ requires = ["maturin>=1.0,<2.0"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "exo_pyo3_bindings"
|
||||
version = "0.2.1"
|
||||
name = "exo_net"
|
||||
version = "0.3.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
authors = [
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
|
||||
{ name = "Evan Quiney", email = "evanev7@gmail.com" },
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
|
||||
]
|
||||
requires-python = ">=3.13"
|
||||
dependencies = []
|
||||
|
||||
[dependency-groups]
|
||||
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
|
||||
dev = ["exo-net", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
|
||||
|
||||
[tool.maturin]
|
||||
#purelib = true
|
||||
#python-source = "python"
|
||||
module-name = "exo_pyo3_bindings"
|
||||
module-name = "exo_net"
|
||||
features = ["pyo3/extension-module", "pyo3/experimental-async"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
File renamed without changes.
@@ -2,7 +2,7 @@ use pyo3_stub_gen::Result;
|
||||
|
||||
fn main() -> Result<()> {
|
||||
env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
|
||||
let stub = exo_pyo3_bindings::stub_info()?;
|
||||
let stub = exo_net::stub_info()?;
|
||||
stub.generate()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -5,21 +5,21 @@
|
||||
//!
|
||||
|
||||
mod allow_threading;
|
||||
mod ident;
|
||||
// mod ident;
|
||||
mod networking;
|
||||
mod point_to_point;
|
||||
mod session;
|
||||
mod state;
|
||||
|
||||
use crate::ident::PyKeypair;
|
||||
use crate::networking::networking_submodule;
|
||||
use crate::point_to_point::{NetReceiver, NetSender};
|
||||
use crate::session::PySession;
|
||||
use crate::state::StateProxy;
|
||||
use pyo3::prelude::PyModule;
|
||||
use pyo3::types::PyModuleMethods;
|
||||
use pyo3::{Bound, PyResult, pyclass, pymodule};
|
||||
use pyo3::{Bound, PyResult, 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;
|
||||
@@ -151,7 +151,7 @@ pub(crate) mod ext {
|
||||
/// 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")]
|
||||
#[pymodule(name = "exo_net")]
|
||||
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
// install logger
|
||||
pyo3_log::init();
|
||||
@@ -162,7 +162,12 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
// 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...
|
||||
m.add_class::<PyKeypair>()?;
|
||||
// m.add_class::<PyKeypair>()?;
|
||||
// networking_submodule(m)?;
|
||||
m.add_class::<StateProxy>()?;
|
||||
m.add_class::<PySession>()?;
|
||||
m.add_class::<NetReceiver>()?;
|
||||
m.add_class::<NetSender>()?;
|
||||
networking_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
@@ -0,0 +1,192 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
|
||||
use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
|
||||
use crate::session::PySession;
|
||||
use futures_lite::{Stream, StreamExt as _};
|
||||
use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{
|
||||
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
|
||||
};
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "NetworkingHandle")]
|
||||
struct PyNetworkingHandle {
|
||||
// channels
|
||||
pub to_swarm: mpsc::Sender<ToSwarm>,
|
||||
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass_complex_enum]
|
||||
#[pyclass]
|
||||
enum PyFromSwarm {
|
||||
Connection { connected: bool },
|
||||
Message { topic: String, data: Py<PyBytes> },
|
||||
}
|
||||
impl From<FromSwarm> for PyFromSwarm {
|
||||
fn from(value: FromSwarm) -> Self {
|
||||
match value {
|
||||
FromSwarm::Discovered {} => Self::Connection { connected: true },
|
||||
FromSwarm::Expired {} => Self::Connection { connected: false },
|
||||
FromSwarm::Message { topic, data } => Self::Message {
|
||||
topic: topic,
|
||||
data: data.pybytes(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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 ----
|
||||
|
||||
#[staticmethod]
|
||||
fn new<'py>(
|
||||
identity: Bound<'py, PyBytes>,
|
||||
bootstrap_peers: Vec<String>,
|
||||
listen_port: u16,
|
||||
) -> PyResult<(PyNetworkingHandle, PySession)> {
|
||||
// create communication channels
|
||||
let (to_swarm, from_client) = mpsc::channel(1024);
|
||||
|
||||
// get identity
|
||||
let identity = u128::from_le_bytes(
|
||||
identity
|
||||
.extract::<'_, Vec<u8>>()?
|
||||
.try_into()
|
||||
.map_err(|_| PyValueError::new_err("invalid identity bytes"))?,
|
||||
);
|
||||
|
||||
// create networking swarm (within tokio context!! or it crashes)
|
||||
let swarm = pyo3_async_runtimes::tokio::get_runtime()
|
||||
.block_on(create_swarm(
|
||||
identity,
|
||||
from_client,
|
||||
bootstrap_peers,
|
||||
listen_port,
|
||||
))
|
||||
.pyerr()?;
|
||||
|
||||
let session = swarm.session.z.clone();
|
||||
|
||||
Ok((
|
||||
PyNetworkingHandle {
|
||||
swarm: Arc::new(Mutex::new(swarm.into_stream())),
|
||||
to_swarm,
|
||||
},
|
||||
PySession { session },
|
||||
))
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let swarm = Arc::clone(&self.swarm);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
swarm
|
||||
.try_lock()
|
||||
.map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
|
||||
.next()
|
||||
.await
|
||||
.ok_or(PyErr::receiver_channel_closed())
|
||||
.map(PyFromSwarm::from)
|
||||
})
|
||||
}
|
||||
|
||||
// ---- 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_swarm
|
||||
.send_py(ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_sender: 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())?
|
||||
.pyerr()
|
||||
}
|
||||
|
||||
/// 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_swarm
|
||||
.send_py(ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_sender: 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_swarm
|
||||
.send_py(ToSwarm::Publish {
|
||||
topic,
|
||||
data,
|
||||
result_sender: 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())?
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pyo3_stub_gen::inventory::submit! {
|
||||
gen_methods_from_python! {
|
||||
r#"
|
||||
class PyNetworkingHandle:
|
||||
async def recv() -> PyFromSwarm: ...
|
||||
"#
|
||||
}
|
||||
}
|
||||
|
||||
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyNetworkingHandle>()?;
|
||||
m.add_class::<PyFromSwarm>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use pyo3::exceptions::PyConnectionError;
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::types::PyNone;
|
||||
use pyo3::{BoundObject, prelude::*};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use zenoh::Result;
|
||||
use zenoh::{
|
||||
handlers::FifoChannelHandler,
|
||||
pubsub::{Publisher, Subscriber},
|
||||
sample::Sample,
|
||||
};
|
||||
|
||||
use crate::ext::ByteArrayExt;
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct NetReceiver {
|
||||
pub subscriber: Subscriber<FifoChannelHandler<Sample>>,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl NetReceiver {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[bytes | None]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
assert!(
|
||||
self.subscriber.receiver_count() == 1,
|
||||
"tried to receive twice on the same receiver"
|
||||
);
|
||||
let subscriber = self.subscriber.clone();
|
||||
async move {
|
||||
match subscriber.recv_async().await {
|
||||
Err(_) => {
|
||||
// stream closed;
|
||||
Ok(Python::attach(|py| PyNone::get(py).unbind()).into_any())
|
||||
}
|
||||
Ok(sample) => Ok(sample.payload().to_bytes().to_vec().pybytes().into_any()),
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct NetSender {
|
||||
pub publisher: Arc<Publisher<'static>>,
|
||||
pub first: bool,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl NetSender {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[bool]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn send<'py>(
|
||||
&'py mut self,
|
||||
py: Python<'py>,
|
||||
data: Bound<'py, PyBytes>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let is_first = self.first;
|
||||
self.first = false;
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
let publisher = Arc::clone(&self.publisher);
|
||||
// clone the data so py can have it back
|
||||
let bytes = data.as_bytes().to_vec();
|
||||
async move {
|
||||
if is_first {
|
||||
wait_for_listener(&*publisher)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;
|
||||
}
|
||||
if !publisher
|
||||
.matching_status()
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(e.to_string()))?
|
||||
.matching()
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
publisher
|
||||
.put(&bytes)
|
||||
.await
|
||||
.map_err(|e| PyConnectionError::new_err(e.to_string()))?;
|
||||
Ok(true)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_listener<'a>(publisher: &Publisher<'a>) -> Result<()> {
|
||||
let matcher = publisher.matching_listener().await?;
|
||||
if publisher.matching_status().await?.matching() {
|
||||
return Ok(());
|
||||
}
|
||||
while let Ok(status) = matcher.recv_async().await {
|
||||
if status.matching() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
|
||||
use zenoh::Session;
|
||||
use zenoh::Wait;
|
||||
use zenoh::qos::CongestionControl;
|
||||
|
||||
use crate::{
|
||||
point_to_point::{NetReceiver, NetSender},
|
||||
state::StateProxy,
|
||||
};
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct PySession {
|
||||
pub session: Session,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PySession {
|
||||
/* for now construct with NetworkingHandle
|
||||
#[staticmethod]
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[PySession]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn init<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
Ok(Self {
|
||||
session: networking::open(
|
||||
networking::cfg(rand::random(), 0).expect("default cfg is valid"),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?,
|
||||
})
|
||||
})
|
||||
}
|
||||
*/
|
||||
|
||||
pub fn net_receiver<'py>(&self, key: String) -> PyResult<NetReceiver> {
|
||||
Ok(NetReceiver {
|
||||
subscriber: self
|
||||
.session
|
||||
.declare_subscriber(key)
|
||||
.wait()
|
||||
// C5: key format error
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn net_sender<'py>(&self, key: String) -> PyResult<NetSender> {
|
||||
Ok(NetSender {
|
||||
publisher: Arc::new(
|
||||
self.session
|
||||
.declare_publisher(key)
|
||||
.congestion_control(CongestionControl::Block)
|
||||
.wait()
|
||||
// C5: key format error, could be declaration error
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))?,
|
||||
),
|
||||
first: true,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn state_proxy(&self) -> StateProxy {
|
||||
StateProxy {
|
||||
session: self.session.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
use serde_json::{Map, Value};
|
||||
use zenoh::{Result, Session, sample::SampleFields};
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass]
|
||||
pub struct StateProxy {
|
||||
pub session: Session,
|
||||
}
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl StateProxy {
|
||||
#[gen_stub(override_return_type(
|
||||
type_repr="collections.abc.Awaitable[str]",
|
||||
imports=("collections.abc")
|
||||
))]
|
||||
pub fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, {
|
||||
let session = self.session.clone();
|
||||
async move {
|
||||
Self::_snapshot(session)
|
||||
.await
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
.map(|v| v.to_string())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl StateProxy {
|
||||
async fn _snapshot(session: Session) -> Result<Value> {
|
||||
let q = session.get("storage/mem1/**").await?;
|
||||
|
||||
let mut v = Value::Object(Map::default());
|
||||
|
||||
while let Ok(sample) = q.recv_async().await {
|
||||
let mut cur_v = &mut v;
|
||||
let Ok(sample) = sample.into_result() else {
|
||||
continue;
|
||||
};
|
||||
// skip storage/mem1
|
||||
let SampleFields {
|
||||
payload, key_expr, ..
|
||||
} = sample.into();
|
||||
let mut iter = key_expr.split('/').skip(2).peekable();
|
||||
loop {
|
||||
let Some(p) = iter.next() else {
|
||||
break;
|
||||
};
|
||||
if iter.peek().is_none() {
|
||||
// terminal; write value into json
|
||||
let existing = cur_v
|
||||
.as_object_mut()
|
||||
.expect("path terminated unexpectedly - value stored at some/path and some/path/two")
|
||||
.insert(p.to_owned(), Value::String(payload.try_to_string()?.to_string()));
|
||||
|
||||
if let Some(value) = existing {
|
||||
assert!(value.is_string())
|
||||
// could log, but string overwrites are fine
|
||||
}
|
||||
} else {
|
||||
// non-terminal; ensure key exists in v, then replace cur with that object
|
||||
cur_v = cur_v
|
||||
.as_object_mut()
|
||||
.expect("path terminated unexpectedly - value stored at some/path and some/path/two")
|
||||
.entry(p)
|
||||
.or_insert(Value::Object(Map::default()));
|
||||
assert!(
|
||||
cur_v.is_object(),
|
||||
"path terminated unexpectedly - value stored at some/path and some/path/two"
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(v)
|
||||
}
|
||||
}
|
||||
@@ -1,66 +0,0 @@
|
||||
[package]
|
||||
name = "exo_pyo3_bindings"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
path = "src/lib.rs"
|
||||
name = "exo_pyo3_bindings"
|
||||
|
||||
# "cdylib" needed to produce shared library for Python to import
|
||||
# "rlib" needed for stub-gen to run
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[[bin]]
|
||||
path = "src/bin/stub_gen.rs"
|
||||
name = "stub_gen"
|
||||
doc = false
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
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
|
||||
|
||||
# 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-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,94 +0,0 @@
|
||||
# This file is automatically generated by pyo3_stub_gen
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import builtins
|
||||
import typing
|
||||
|
||||
@typing.final
|
||||
class AllQueuesFullError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class Keypair:
|
||||
r"""
|
||||
Identity keypair of a node.
|
||||
"""
|
||||
@staticmethod
|
||||
def generate() -> Keypair:
|
||||
r"""
|
||||
Generate a new Ed25519 keypair.
|
||||
"""
|
||||
@staticmethod
|
||||
def from_bytes(bytes: bytes) -> Keypair:
|
||||
r"""
|
||||
Construct an Ed25519 keypair from secret key bytes
|
||||
"""
|
||||
def to_bytes(self) -> bytes:
|
||||
r"""
|
||||
Get the secret key bytes underlying the keypair
|
||||
"""
|
||||
def to_node_id(self) -> builtins.str:
|
||||
r"""
|
||||
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, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
|
||||
Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
|
||||
"""
|
||||
async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Unsubscribes from a `GossipSub` topic.
|
||||
|
||||
Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
|
||||
"""
|
||||
async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
|
||||
r"""
|
||||
Publishes a message with multiple topics to the `GossipSub` network.
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
@typing.final
|
||||
class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
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: ...
|
||||
|
||||
@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: ...
|
||||
|
||||
...
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
use crate::ext::ResultExt as _;
|
||||
use libp2p::identity::Keypair;
|
||||
use pyo3::types::{PyBytes, PyBytesMethods as _};
|
||||
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
|
||||
|
||||
/// Identity keypair of a node.
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "Keypair", frozen)]
|
||||
#[repr(transparent)]
|
||||
pub struct PyKeypair(pub Keypair);
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
#[allow(clippy::needless_pass_by_value)]
|
||||
impl PyKeypair {
|
||||
/// Generate a new Ed25519 keypair.
|
||||
#[staticmethod]
|
||||
fn generate() -> Self {
|
||||
Self(Keypair::generate_ed25519())
|
||||
}
|
||||
|
||||
/// Construct an Ed25519 keypair from secret key bytes
|
||||
#[staticmethod]
|
||||
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()?))
|
||||
}
|
||||
|
||||
/// 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` string, which we use as our `NodeId`.
|
||||
fn to_node_id(&self) -> String {
|
||||
self.0.public().to_peer_id().to_base58()
|
||||
}
|
||||
}
|
||||
@@ -1,318 +0,0 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::r#const::MPSC_CHANNEL_SIZE;
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
|
||||
use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
|
||||
use crate::ident::PyKeypair;
|
||||
use crate::networking::exception::{
|
||||
PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
|
||||
};
|
||||
use crate::pyclass;
|
||||
use futures_lite::{Stream, StreamExt as _};
|
||||
use libp2p::gossipsub::PublishError;
|
||||
use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods as _};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{
|
||||
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
|
||||
};
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
|
||||
pub struct PyMessageTooLargeError {}
|
||||
|
||||
impl PyMessageTooLargeError {
|
||||
const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
|
||||
|
||||
pub(crate) fn new_err() -> PyErr {
|
||||
PyErr::new::<Self, _>(())
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyMessageTooLargeError {
|
||||
#[new]
|
||||
#[pyo3(signature = (*args))]
|
||||
#[allow(unused_variables)]
|
||||
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("MessageTooLargeError(\"{}\")", Self::MSG)
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
Self::MSG.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "NetworkingHandle")]
|
||||
struct PyNetworkingHandle {
|
||||
// channels
|
||||
pub to_swarm: mpsc::Sender<ToSwarm>,
|
||||
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass_complex_enum]
|
||||
#[pyclass]
|
||||
enum PyFromSwarm {
|
||||
Connection {
|
||||
peer_id: String,
|
||||
connected: bool,
|
||||
},
|
||||
Message {
|
||||
origin: String,
|
||||
topic: String,
|
||||
data: Py<PyBytes>,
|
||||
},
|
||||
}
|
||||
impl From<FromSwarm> for PyFromSwarm {
|
||||
fn from(value: FromSwarm) -> Self {
|
||||
match value {
|
||||
FromSwarm::Discovered { peer_id } => Self::Connection {
|
||||
peer_id: peer_id.to_base58(),
|
||||
connected: true,
|
||||
},
|
||||
FromSwarm::Expired { peer_id } => Self::Connection {
|
||||
peer_id: peer_id.to_base58(),
|
||||
connected: false,
|
||||
},
|
||||
FromSwarm::Message { from, topic, data } => Self::Message {
|
||||
origin: from.to_base58(),
|
||||
topic: topic,
|
||||
data: data.pybytes(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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]
|
||||
#[pyo3(signature = (identity, bootstrap_peers, listen_port))]
|
||||
fn py_new(
|
||||
identity: Bound<'_, PyKeypair>,
|
||||
bootstrap_peers: Vec<String>,
|
||||
listen_port: u16,
|
||||
) -> PyResult<Self> {
|
||||
// create communication channels
|
||||
let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
|
||||
// get identity
|
||||
let identity = identity.borrow().0.clone();
|
||||
|
||||
// create networking swarm (within tokio context!! or it crashes)
|
||||
let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
|
||||
let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
|
||||
.pyerr()?
|
||||
.into_stream();
|
||||
|
||||
Ok(Self {
|
||||
swarm: Arc::new(Mutex::new(swarm)),
|
||||
to_swarm,
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let swarm = Arc::clone(&self.swarm);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
swarm
|
||||
.try_lock()
|
||||
.map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
|
||||
.next()
|
||||
.await
|
||||
.ok_or(PyErr::receiver_channel_closed())
|
||||
.map(PyFromSwarm::from)
|
||||
})
|
||||
}
|
||||
|
||||
// ---- 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_swarm
|
||||
.send_py(ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_sender: 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())?
|
||||
.pyerr()
|
||||
}
|
||||
|
||||
/// 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_swarm
|
||||
.send_py(ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_sender: 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_swarm
|
||||
.send_py(ToSwarm::Publish {
|
||||
topic,
|
||||
data,
|
||||
result_sender: 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())?
|
||||
.map_err(|e| match e {
|
||||
PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
|
||||
PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
|
||||
PublishError::NoPeersSubscribedToTopic => {
|
||||
PyNoPeersSubscribedToTopicError::new_err()
|
||||
}
|
||||
e => PyRuntimeError::new_err(e.to_string()),
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pyo3_stub_gen::inventory::submit! {
|
||||
gen_methods_from_python! {
|
||||
r#"
|
||||
class PyNetworkingHandle:
|
||||
async def recv() -> PyFromSwarm: ...
|
||||
"#
|
||||
}
|
||||
}
|
||||
|
||||
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
|
||||
m.add_class::<exception::PyAllQueuesFullError>()?;
|
||||
m.add_class::<exception::PyMessageTooLargeError>()?;
|
||||
|
||||
m.add_class::<PyNetworkingHandle>()?;
|
||||
m.add_class::<PyFromSwarm>()?;
|
||||
|
||||
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,36 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
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(), [], 0)
|
||||
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
|
||||
# sleep for 4 ticks
|
||||
for i in range(4):
|
||||
await asyncio.sleep(1)
|
||||
|
||||
try:
|
||||
await h.gossipsub_publish("topic", b"somehting or other")
|
||||
except NoPeersSubscribedToTopicError as e:
|
||||
print("caught it", e)
|
||||
|
||||
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
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}")
|
||||
+13
-36
@@ -1,42 +1,19 @@
|
||||
[package]
|
||||
name = "networking"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
name = "networking"
|
||||
path = "src/lib.rs"
|
||||
[dependencies]
|
||||
async-stream.workspace = true
|
||||
futures-lite.workspace = true
|
||||
netwatcher = { workspace = true, features = ["tokio"] }
|
||||
parking_lot.workspace = true
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
zenoh = { workspace = true, features = ["internal", "plugins", "unstable"] }
|
||||
zenoh-plugin-storage-manager.workspace = true
|
||||
zenoh-plugin-trait.workspace = true
|
||||
rand.workspace = true
|
||||
log.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
# datastructures
|
||||
either = { workspace = true }
|
||||
|
||||
# macro dependencies
|
||||
extend = { workspace = true }
|
||||
delegate = { workspace = true }
|
||||
|
||||
# async
|
||||
async-stream = { workspace = true }
|
||||
futures-lite = { workspace = true }
|
||||
futures-timer = { workspace = true }
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
|
||||
# utility dependencies
|
||||
util = { workspace = true }
|
||||
tracing-subscriber = { version = "0.3.19", features = [
|
||||
"default",
|
||||
"env-filter",
|
||||
] }
|
||||
keccak-const = { workspace = true }
|
||||
|
||||
# tracing/logging
|
||||
log = { workspace = true }
|
||||
|
||||
# networking
|
||||
libp2p = { workspace = true, features = ["full"] }
|
||||
pin-project = "1.1.10"
|
||||
@@ -1,86 +0,0 @@
|
||||
use futures_lite::StreamExt;
|
||||
use libp2p::identity;
|
||||
use networking::swarm;
|
||||
use networking::swarm::{FromSwarm, ToSwarm};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::{io, io::AsyncBufReadExt as _};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use tracing_subscriber::filter::LevelFilter;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
|
||||
.try_init();
|
||||
|
||||
let (to_swarm, from_client) = mpsc::channel(20);
|
||||
|
||||
// Configure swarm
|
||||
let mut swarm = swarm::create_swarm(
|
||||
identity::Keypair::generate_ed25519(),
|
||||
from_client,
|
||||
vec![],
|
||||
0,
|
||||
)
|
||||
.expect("Swarm creation failed")
|
||||
.into_stream();
|
||||
|
||||
// Create a Gossipsub topic & subscribe
|
||||
let (tx, rx) = oneshot::channel();
|
||||
_ = to_swarm
|
||||
.send(ToSwarm::Subscribe {
|
||||
topic: "test-net".to_string(),
|
||||
result_sender: tx,
|
||||
})
|
||||
.await
|
||||
.expect("should send");
|
||||
|
||||
// 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");
|
||||
|
||||
tokio::task::spawn(async move {
|
||||
rx.await
|
||||
.expect("tx not dropped")
|
||||
.expect("subscribe shouldn't fail");
|
||||
loop {
|
||||
if let Ok(Some(line)) = stdin.next_line().await {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
if let Err(e) = to_swarm
|
||||
.send(swarm::ToSwarm::Publish {
|
||||
topic: "test-net".to_string(),
|
||||
data: line.as_bytes().to_vec(),
|
||||
result_sender: tx,
|
||||
})
|
||||
.await
|
||||
{
|
||||
println!("Send error: {e:?}");
|
||||
return;
|
||||
};
|
||||
match rx.await {
|
||||
Ok(Err(e)) => println!("Publish error: {e:?}"),
|
||||
Err(e) => println!("Publish error: {e:?}"),
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Kick it off
|
||||
loop {
|
||||
// on gossipsub outgoing
|
||||
match swarm.next().await {
|
||||
// on gossipsub incoming
|
||||
Some(FromSwarm::Discovered { peer_id }) => {
|
||||
println!("\n\nconnected to {peer_id}\n\n")
|
||||
}
|
||||
Some(FromSwarm::Expired { peer_id }) => {
|
||||
println!("\n\ndisconnected from {peer_id}\n\n")
|
||||
}
|
||||
Some(FromSwarm::Message { from, topic, data }) => {
|
||||
println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
use log::info;
|
||||
use networking;
|
||||
use zenoh::Result;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
zenoh::init_log_from_env_or("info");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 0)?;
|
||||
let session = networking::open(cfg).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
.declare_token(format!("nodes/{}/live", session.z.zid()))
|
||||
.await?;
|
||||
let key_expr = "storage/mem1/name";
|
||||
let payload = "me";
|
||||
|
||||
info!("Putting Data ('{key_expr}': '{payload}')...");
|
||||
session.z.put(key_expr, payload).await?;
|
||||
tokio::signal::ctrl_c().await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
use log::info;
|
||||
use networking;
|
||||
use zenoh::Result;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
zenoh::init_log_from_env_or("info");
|
||||
info!("Opening session...");
|
||||
let cfg = networking::cfg(rand::random(), 0)?;
|
||||
let session = networking::open(cfg).await?;
|
||||
let _tok = session
|
||||
.z
|
||||
.liveliness()
|
||||
.declare_token(format!("nodes/{}/live", session.z.zid()))
|
||||
.await?;
|
||||
let _sub = session
|
||||
.z
|
||||
.liveliness()
|
||||
.declare_subscriber("nodes/*/live")
|
||||
.history(true)
|
||||
.callback(|tok| println!("{tok:?}"))
|
||||
.await?;
|
||||
tokio::signal::ctrl_c().await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -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,390 +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>>,
|
||||
bootstrap_peers: Vec<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, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
|
||||
Ok(Self {
|
||||
managed: managed::Behaviour::new(keypair)?,
|
||||
mdns_discovered: HashMap::new(),
|
||||
bootstrap_peers,
|
||||
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)
|
||||
}
|
||||
}
|
||||
// dial bootstrap peers (for environments where mDNS is unavailable)
|
||||
for addr in &self.bootstrap_peers {
|
||||
self.pending_events.push_back(ToSwarm::Dial {
|
||||
opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
|
||||
})
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
+96
-36
@@ -1,44 +1,104 @@
|
||||
//! TODO: crate documentation
|
||||
//!
|
||||
//! this is here as a placeholder documentation
|
||||
//!
|
||||
//!
|
||||
pub mod discovery;
|
||||
use std::{env, panic, sync::Arc};
|
||||
|
||||
use netwatcher::WatchHandle;
|
||||
use parking_lot::Mutex;
|
||||
use tokio::{sync::mpsc, task::JoinHandle};
|
||||
use zenoh::{Result, Session as ZSession, config::WhatAmI, internal::runtime::Runtime};
|
||||
use zenoh_plugin_storage_manager::StoragesPlugin;
|
||||
use zenoh_plugin_trait::PluginsManager;
|
||||
|
||||
pub use zenoh::{Config, config::ZenohId};
|
||||
|
||||
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>;
|
||||
pub fn cfg(identity: u128, listen_port: u16) -> Result<zenoh::Config> {
|
||||
let namespace = env::var("EXO_ZENOH_NAMESPACE").unwrap_or_else(|_| "exo".to_string());
|
||||
let mut cfg = zenoh::Config::default();
|
||||
// todo: cleanup
|
||||
cfg.insert_json5("id", &format!("\"{identity:x}\""))?;
|
||||
cfg.insert_json5("mode", "\"peer\"")?;
|
||||
cfg.insert_json5("listen/endpoints", &format!("[\"tcp/[::]:{listen_port}\"]"))?;
|
||||
cfg.insert_json5("scouting/multicast/enabled", "true")?;
|
||||
cfg.insert_json5("scouting/multicast/autoconnect", "[]")?;
|
||||
cfg.insert_json5("scouting/gossip/multihop", "true")?;
|
||||
cfg.insert_json5("namespace", &format!("{namespace:?}"))?;
|
||||
cfg.insert_json5("transport/link/tx/batch_size", "9216")?;
|
||||
cfg.insert_json5("timestamping/enabled", "true")?;
|
||||
cfg.insert_json5("plugins/storage_manager/__required__", "true")?;
|
||||
cfg.insert_json5(
|
||||
"plugins/storage_manager/storages/mem1",
|
||||
r#"{
|
||||
key_expr: "storage/mem1/**",
|
||||
strip_prefix: "storage/mem1",
|
||||
volume: "memory",
|
||||
replication: {
|
||||
interval: 2,
|
||||
}
|
||||
}"#,
|
||||
)?;
|
||||
Ok(cfg)
|
||||
}
|
||||
|
||||
/// 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,
|
||||
pub async fn open(cfg: zenoh::Config) -> Result<Session> {
|
||||
let mut plugins = PluginsManager::static_plugins_only();
|
||||
plugins.declare_static_plugin::<StoragesPlugin, _>("storage_manager", true);
|
||||
let mut runtime = zenoh::internal::runtime::RuntimeBuilder::new(cfg)
|
||||
.plugins_manager(plugins)
|
||||
.build()
|
||||
.await?;
|
||||
let z = zenoh::session::init(runtime.clone().into()).await?;
|
||||
runtime.start().await?;
|
||||
let _watch_all_handle = watch_all(runtime).await?;
|
||||
Ok(Session {
|
||||
z,
|
||||
_watch_all_handle,
|
||||
})
|
||||
}
|
||||
async fn watch_all(runtime: Runtime) -> Result<WatchAllHandle> {
|
||||
log::info!("spawning scout");
|
||||
let mut cfg = Config::default();
|
||||
cfg.insert_json5("scouting/multicast/ttl", "3")?;
|
||||
cfg.insert_json5("scouting/multicast/interface", "\"auto\"")?;
|
||||
let mut scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
|
||||
let (send, mut recv) = mpsc::unbounded_channel();
|
||||
let _sync = Arc::new(Mutex::new(netwatcher::watch_interfaces_with_callback(
|
||||
move |u| _ = send.send(u),
|
||||
)?));
|
||||
let _async = tokio::task::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
u = recv.recv() => {
|
||||
if u.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
log::info!("reloading scout");
|
||||
scout = zenoh::scout(WhatAmI::Peer, cfg.clone()).await?;
|
||||
}
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
let Some(Protocol::Tcp(port)) = ps.next() else {
|
||||
return None;
|
||||
};
|
||||
Some((ip, port))
|
||||
hello = scout.recv_async() => {
|
||||
if let Ok(hello) = hello {
|
||||
// todo: auth
|
||||
runtime
|
||||
.connect_peer(&hello.zid().into(), hello.locators())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok(WatchAllHandle { _sync, _async })
|
||||
}
|
||||
|
||||
pub struct Session {
|
||||
pub z: ZSession,
|
||||
_watch_all_handle: WatchAllHandle,
|
||||
}
|
||||
impl Drop for WatchAllHandle {
|
||||
fn drop(&mut self) {
|
||||
self._async.abort();
|
||||
}
|
||||
}
|
||||
|
||||
pub struct WatchAllHandle {
|
||||
_sync: Arc<Mutex<WatchHandle>>,
|
||||
_async: JoinHandle<Result<()>>,
|
||||
}
|
||||
+132
-230
@@ -1,24 +1,20 @@
|
||||
//! Compat shim for the old libp2p code
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::swarm::transport::tcp_transport;
|
||||
use crate::{alias, discovery};
|
||||
pub use behaviour::{Behaviour, BehaviourEvent};
|
||||
use futures_lite::{Stream, StreamExt};
|
||||
use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use futures_lite::Stream;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use zenoh::Result;
|
||||
use zenoh::Session;
|
||||
use zenoh::handlers::FifoChannelHandler;
|
||||
use zenoh::liveliness::LivelinessToken;
|
||||
use zenoh::pubsub::Subscriber;
|
||||
use zenoh::sample::Sample;
|
||||
use zenoh::sample::SampleKind;
|
||||
|
||||
/// 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";
|
||||
|
||||
// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
|
||||
// of the Swarm.
|
||||
#[derive(Debug)]
|
||||
pub enum ToSwarm {
|
||||
Unsubscribe {
|
||||
topic: String,
|
||||
@@ -26,52 +22,66 @@ pub enum ToSwarm {
|
||||
},
|
||||
Subscribe {
|
||||
topic: String,
|
||||
result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
|
||||
result_sender: oneshot::Sender<Result<bool>>,
|
||||
},
|
||||
Publish {
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
|
||||
result_sender: oneshot::Sender<Result<()>>,
|
||||
},
|
||||
}
|
||||
#[derive(Debug)]
|
||||
pub enum FromSwarm {
|
||||
Message {
|
||||
from: PeerId,
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
},
|
||||
Discovered {
|
||||
peer_id: PeerId,
|
||||
},
|
||||
Expired {
|
||||
peer_id: PeerId,
|
||||
},
|
||||
Message { topic: String, data: Vec<u8> },
|
||||
Discovered {},
|
||||
Expired {},
|
||||
}
|
||||
|
||||
pub type Topics = HashMap<String, Subscriber<()>>;
|
||||
pub struct Swarm {
|
||||
swarm: libp2p::Swarm<Behaviour>,
|
||||
pub session: crate::Session,
|
||||
from_client: mpsc::Receiver<ToSwarm>,
|
||||
}
|
||||
|
||||
impl Swarm {
|
||||
pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
|
||||
let Swarm {
|
||||
mut swarm,
|
||||
session,
|
||||
mut from_client,
|
||||
} = self;
|
||||
let stream = async_stream::stream! {
|
||||
let mut session = session;
|
||||
let (mut to_topics, mut from_topics) = mpsc::channel(1024);
|
||||
let mut topics = Topics::new();
|
||||
let Ok((_token, discovery)) = register_liveness(&mut session.z).await else { return; };
|
||||
loop {
|
||||
tokio::select! {
|
||||
msg = from_client.recv() => {
|
||||
let Some(msg) = msg else { break };
|
||||
on_message(&mut swarm, msg);
|
||||
on_message(&mut session.z, &mut topics, &mut to_topics, msg).await;
|
||||
}
|
||||
event = swarm.next() => {
|
||||
let Some(event) = event else { break };
|
||||
if let Some(item) = filter_swarm_event(event) {
|
||||
yield item;
|
||||
event = from_topics.recv() => {
|
||||
if let Some(event) = event {
|
||||
yield event
|
||||
}
|
||||
}
|
||||
token = discovery.recv_async() => {
|
||||
if let Ok(token) = token {
|
||||
let key_expr = token.key_expr().as_str().to_owned();
|
||||
let nid = key_expr.strip_prefix("nodes/").and_then(|s| s.strip_suffix("/live"));
|
||||
yield match token.kind() {
|
||||
SampleKind::Put => {
|
||||
log::info!("discovered: {nid:?}");
|
||||
FromSwarm::Discovered {}
|
||||
}
|
||||
SampleKind::Delete => {
|
||||
log::info!("expired: {nid:?}");
|
||||
FromSwarm::Expired {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
@@ -79,208 +89,100 @@ impl Swarm {
|
||||
}
|
||||
}
|
||||
|
||||
fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
|
||||
match message {
|
||||
ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.subscribe(&gossipsub::IdentTopic::new(topic));
|
||||
_ = result_sender.send(result);
|
||||
}
|
||||
ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.unsubscribe(&gossipsub::IdentTopic::new(topic));
|
||||
_ = result_sender.send(result);
|
||||
}
|
||||
async fn register_liveness(
|
||||
session: &mut Session,
|
||||
) -> Result<(LivelinessToken, Subscriber<FifoChannelHandler<Sample>>)> {
|
||||
let token = session
|
||||
.liveliness()
|
||||
.declare_token(format!("nodes/{}/live", session.zid()))
|
||||
.await?;
|
||||
let sub = session
|
||||
.liveliness()
|
||||
.declare_subscriber("nodes/*/live")
|
||||
.history(true)
|
||||
.await?;
|
||||
Ok((token, sub))
|
||||
}
|
||||
|
||||
async fn on_message(
|
||||
session: &mut Session,
|
||||
topics: &mut Topics,
|
||||
to_topics: &mut mpsc::Sender<FromSwarm>,
|
||||
msg: ToSwarm,
|
||||
) {
|
||||
match msg {
|
||||
ToSwarm::Publish {
|
||||
topic,
|
||||
data,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.publish(gossipsub::IdentTopic::new(topic), data);
|
||||
_ = result_sender.send(result);
|
||||
let res = session.put(format!("topics/{topic}"), data).await;
|
||||
_ = result_sender.send(res);
|
||||
}
|
||||
ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
let Some((_, subscriber)) = topics.remove_entry(&topic) else {
|
||||
_ = result_sender.send(false);
|
||||
return;
|
||||
};
|
||||
_ = subscriber.undeclare().await;
|
||||
_ = result_sender.send(true);
|
||||
}
|
||||
ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
assert!(topic.is_ascii());
|
||||
if topics.contains_key(&topic) {
|
||||
_ = result_sender.send(Ok(false));
|
||||
return;
|
||||
}
|
||||
let subscriber = match session
|
||||
.declare_subscriber(format!("topics/{topic}"))
|
||||
.allowed_origin(zenoh::sample::Locality::Remote)
|
||||
.callback({
|
||||
let sender = to_topics.clone();
|
||||
let topic = topic.clone();
|
||||
move |sample| {
|
||||
if sample.kind() != SampleKind::Put {
|
||||
return;
|
||||
}
|
||||
_ = sender.try_send(FromSwarm::Message {
|
||||
topic: topic.clone(),
|
||||
data: sample.payload().to_bytes().to_vec(),
|
||||
});
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
_ = result_sender.send(Err(e));
|
||||
return;
|
||||
}
|
||||
};
|
||||
assert!(topics.insert(topic, subscriber).is_none());
|
||||
_ = result_sender.send(Ok(true));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
|
||||
match event {
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
|
||||
message:
|
||||
gossipsub::Message {
|
||||
source: Some(peer_id),
|
||||
topic,
|
||||
data,
|
||||
..
|
||||
},
|
||||
..
|
||||
})) => Some(FromSwarm::Message {
|
||||
from: peer_id,
|
||||
topic: topic.into_string(),
|
||||
data,
|
||||
}),
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Discovery(
|
||||
discovery::Event::ConnectionEstablished { peer_id, .. },
|
||||
)) => Some(FromSwarm::Discovered { peer_id }),
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
|
||||
peer_id,
|
||||
..
|
||||
})) => Some(FromSwarm::Expired { peer_id }),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create and configure a swarm.
|
||||
///
|
||||
/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
|
||||
/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
|
||||
pub fn create_swarm(
|
||||
keypair: identity::Keypair,
|
||||
pub async fn create_swarm(
|
||||
identity: u128,
|
||||
from_client: mpsc::Receiver<ToSwarm>,
|
||||
bootstrap_peers: Vec<String>,
|
||||
listen_port: u16,
|
||||
) -> alias::AnyResult<Swarm> {
|
||||
let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
|
||||
.iter()
|
||||
.filter(|s| !s.is_empty())
|
||||
.filter_map(|s| s.parse().ok())
|
||||
.collect();
|
||||
|
||||
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
|
||||
.with_tokio()
|
||||
.with_other_transport(tcp_transport)?
|
||||
.with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
|
||||
.build();
|
||||
|
||||
swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
|
||||
Ok(Swarm { swarm, from_client })
|
||||
}
|
||||
|
||||
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,
|
||||
bootstrap_peers: Vec<libp2p::Multiaddr>,
|
||||
) -> alias::AnyResult<Self> {
|
||||
Ok(Self {
|
||||
discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
|
||||
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(8 * 1024 * 1024)
|
||||
.validation_mode(ValidationMode::Strict)
|
||||
.build()
|
||||
.expect("the configuration should always be valid"),
|
||||
)
|
||||
.expect("creating gossipsub behavior should always work")
|
||||
}
|
||||
) -> Result<Swarm> {
|
||||
// todo: bootstrap
|
||||
if !bootstrap_peers.is_empty() || listen_port != 0 {
|
||||
todo!();
|
||||
}
|
||||
let cfg = crate::cfg(identity, listen_port)?;
|
||||
let session = crate::open(cfg).await?;
|
||||
Ok(Swarm {
|
||||
session,
|
||||
from_client,
|
||||
})
|
||||
}
|
||||
@@ -1,107 +0,0 @@
|
||||
use futures_lite::StreamExt;
|
||||
use networking::swarm::{FromSwarm, create_swarm};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::timeout;
|
||||
|
||||
/// Helper: find a free TCP port.
|
||||
fn free_port() -> u16 {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
listener.local_addr().unwrap().port()
|
||||
}
|
||||
|
||||
/// Two nodes connect via bootstrap peers — no mDNS needed.
|
||||
///
|
||||
/// Node A listens on a fixed port. Node B bootstraps to A's address.
|
||||
/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
|
||||
#[tokio::test]
|
||||
async fn two_nodes_connect_via_bootstrap_peers() {
|
||||
let port_a = free_port();
|
||||
|
||||
// Node A: listens on a known port, no bootstrap peers
|
||||
let keypair_a = libp2p::identity::Keypair::generate_ed25519();
|
||||
let peer_id_a = keypair_a.public().to_peer_id();
|
||||
let (_tx_a, rx_a) = mpsc::channel(16);
|
||||
let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
|
||||
let mut stream_a = swarm_a.into_stream();
|
||||
|
||||
// Node B: bootstraps to A's address
|
||||
let keypair_b = libp2p::identity::Keypair::generate_ed25519();
|
||||
let (_tx_b, rx_b) = mpsc::channel(16);
|
||||
let swarm_b = create_swarm(
|
||||
keypair_b,
|
||||
rx_b,
|
||||
vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
|
||||
0,
|
||||
)
|
||||
.expect("create swarm B");
|
||||
let mut stream_b = swarm_b.into_stream();
|
||||
|
||||
// Wait for B to discover A (connection established)
|
||||
let connected = timeout(Duration::from_secs(10), async {
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(event) = stream_a.next() => {
|
||||
// A will also see B connect, but we check from B's perspective
|
||||
let _ = event;
|
||||
}
|
||||
Some(event) = stream_b.next() => {
|
||||
if let FromSwarm::Discovered { peer_id } = event {
|
||||
if peer_id == peer_id_a {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
connected.is_ok() && connected.unwrap(),
|
||||
"Node B should discover Node A via bootstrap peer"
|
||||
);
|
||||
}
|
||||
|
||||
/// Empty bootstrap peers should work (backward compatible).
|
||||
#[tokio::test]
|
||||
async fn create_swarm_with_empty_bootstrap_peers() {
|
||||
let keypair = libp2p::identity::Keypair::generate_ed25519();
|
||||
let (_tx, rx) = mpsc::channel(16);
|
||||
let swarm = create_swarm(keypair, rx, vec![], 0);
|
||||
assert!(
|
||||
swarm.is_ok(),
|
||||
"create_swarm with no bootstrap peers should succeed"
|
||||
);
|
||||
}
|
||||
|
||||
/// Invalid multiaddr strings are silently filtered out.
|
||||
#[tokio::test]
|
||||
async fn create_swarm_ignores_invalid_bootstrap_addrs() {
|
||||
let keypair = libp2p::identity::Keypair::generate_ed25519();
|
||||
let (_tx, rx) = mpsc::channel(16);
|
||||
let swarm = create_swarm(
|
||||
keypair,
|
||||
rx,
|
||||
vec![
|
||||
"not-a-valid-multiaddr".to_string(),
|
||||
"".to_string(),
|
||||
"/ip4/10.0.0.1/tcp/30000".to_string(), // valid
|
||||
],
|
||||
0,
|
||||
);
|
||||
assert!(
|
||||
swarm.is_ok(),
|
||||
"create_swarm should succeed even with invalid bootstrap addrs"
|
||||
);
|
||||
}
|
||||
|
||||
/// Fixed listen port works correctly.
|
||||
#[tokio::test]
|
||||
async fn create_swarm_with_fixed_port() {
|
||||
let port = free_port();
|
||||
let keypair = libp2p::identity::Keypair::generate_ed25519();
|
||||
let (_tx, rx) = mpsc::channel(16);
|
||||
let swarm = create_swarm(keypair, rx, vec![], port);
|
||||
assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
// maybe this will hold test in the future...??
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn does_nothing() {}
|
||||
}
|
||||
+4
-3
@@ -55,6 +55,7 @@
|
||||
];
|
||||
|
||||
OPENSSL_NO_VENDOR = "1";
|
||||
MATURIN_NO_INSTALL_RUST = "1";
|
||||
|
||||
# Required for pyo3 tests to find libpython
|
||||
LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ];
|
||||
@@ -81,11 +82,11 @@
|
||||
config = {
|
||||
packages = {
|
||||
# Python bindings wheel via maturin
|
||||
exo_pyo3_bindings = craneLib.buildPackage (
|
||||
exo-net = craneLib.buildPackage (
|
||||
commonArgs
|
||||
// {
|
||||
inherit cargoArtifacts;
|
||||
pname = "exo_pyo3_bindings";
|
||||
pname = "exo-net";
|
||||
|
||||
nativeBuildInputs = commonArgs.nativeBuildInputs ++ [
|
||||
pkgs.maturin
|
||||
@@ -95,7 +96,7 @@
|
||||
maturin build \
|
||||
--release \
|
||||
--manylinux off \
|
||||
--manifest-path rust/exo_pyo3_bindings/Cargo.toml \
|
||||
--manifest-path rust/exo_net/Cargo.toml \
|
||||
--features "pyo3/extension-module,pyo3/experimental-async" \
|
||||
--interpreter ${pkgs.python313}/bin/python \
|
||||
--out dist
|
||||
|
||||
@@ -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,222 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
"""Standalone smoke test for VllmEngine.serve_prefill.
|
||||
|
||||
Loads a real vLLM engine, runs serve_prefill against an in-memory buffer
|
||||
twice in a row with the same prompt, and verifies both runs produce a
|
||||
well-formed wire stream (header -> KV chunks -> Done).
|
||||
|
||||
The second run is the regression guard: with vLLM APC enabled this would
|
||||
trip the chunked-prefill + APC + custom kv-connector CUDA assert
|
||||
(`vectorized_gather_kernel: ind >= ind_dim_size`) and the server would
|
||||
close the socket before the Done frame.
|
||||
|
||||
Usage on the Spark (gx10-de89):
|
||||
|
||||
cd /home/larry/exo
|
||||
/nix/store/2b82iz9ac0pxqafrgxmgdkq8sr2hwlx6-exo-cuda-13-venv/bin/python \\
|
||||
scripts/check_serve_prefill.py Qwen/Qwen3-0.6B
|
||||
|
||||
Exits 0 on success, non-zero with a diagnostic on failure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
|
||||
def _ensure_repo_on_path() -> None:
|
||||
repo = Path(__file__).resolve().parent.parent
|
||||
src = repo / "src"
|
||||
if str(src) not in sys.path:
|
||||
sys.path.insert(0, str(src))
|
||||
|
||||
|
||||
_ensure_repo_on_path()
|
||||
|
||||
from exo.shared.types.common import ModelId # noqa: E402
|
||||
from exo.worker.disaggregated.protocol import ( # noqa: E402
|
||||
ArraysState,
|
||||
Done,
|
||||
ErrorMessage,
|
||||
KVChunk,
|
||||
read_header,
|
||||
read_message,
|
||||
)
|
||||
from exo.worker.disaggregated.server import PrefillRequest # noqa: E402
|
||||
|
||||
|
||||
def _make_token_ids(n: int) -> list[int]:
|
||||
return [(i * 1009 + 17) % 30000 + 100 for i in range(n)]
|
||||
|
||||
|
||||
def _decode(
|
||||
payload: bytes,
|
||||
) -> tuple[list[KVChunk], list[ArraysState], Done | None, ErrorMessage | None]:
|
||||
buf = io.BytesIO(payload)
|
||||
_ = read_header(buf)
|
||||
chunks: list[KVChunk] = []
|
||||
arrays: list[ArraysState] = []
|
||||
done: Done | None = None
|
||||
error: ErrorMessage | None = None
|
||||
while True:
|
||||
msg = read_message(buf)
|
||||
if msg is None:
|
||||
break
|
||||
if isinstance(msg, KVChunk):
|
||||
chunks.append(msg)
|
||||
elif isinstance(msg, ArraysState):
|
||||
arrays.append(msg)
|
||||
elif isinstance(msg, Done):
|
||||
done = msg
|
||||
break
|
||||
elif isinstance(msg, ErrorMessage):
|
||||
error = msg
|
||||
break
|
||||
return chunks, arrays, done, error
|
||||
|
||||
|
||||
def _build_engine(model_id: ModelId) -> object:
|
||||
from exo.worker.engines.vllm.engine import VllmEngine
|
||||
from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
|
||||
from exo.worker.engines.vllm.kv_connector import (
|
||||
ExoKVProducerConnector,
|
||||
_patch_gdn_capture,
|
||||
_patch_vllm_for_connector,
|
||||
)
|
||||
|
||||
_patch_vllm_for_connector(ExoKVProducerConnector)
|
||||
_patch_gdn_capture()
|
||||
|
||||
llm_engine, tool_parser = load_vllm_engine(
|
||||
model_id=model_id,
|
||||
trust_remote_code=False,
|
||||
n_layers=1,
|
||||
kv_connector_cls=ExoKVProducerConnector,
|
||||
)
|
||||
gen = VllmBatchEngine(engine=llm_engine, model_id=model_id)
|
||||
|
||||
class _S:
|
||||
def send(self, _: object) -> None: ...
|
||||
|
||||
class _R:
|
||||
def collect(self) -> list[object]:
|
||||
return []
|
||||
|
||||
return VllmEngine(
|
||||
tool_parser=tool_parser,
|
||||
model_id=model_id,
|
||||
cancel_receiver=cast("object", _R()), # pyright: ignore[reportArgumentType]
|
||||
event_sender=cast("object", _S()), # pyright: ignore[reportArgumentType]
|
||||
_gen=gen,
|
||||
max_concurrent_requests=1,
|
||||
)
|
||||
|
||||
|
||||
def _run_one(engine: object, n_tokens: int, label: str) -> int:
|
||||
request = PrefillRequest(
|
||||
request_id=f"check-{label}-{os.getpid()}",
|
||||
model_id="ignored",
|
||||
token_ids=_make_token_ids(n_tokens),
|
||||
start_pos=0,
|
||||
use_prefix_cache=True,
|
||||
)
|
||||
buf = io.BytesIO()
|
||||
engine.serve_prefill(request, buf) # pyright: ignore[reportAttributeAccessIssue]
|
||||
payload = buf.getvalue()
|
||||
if not payload:
|
||||
raise AssertionError(f"{label}: server wrote nothing")
|
||||
|
||||
chunks, arrays, done, error = _decode(payload)
|
||||
if error is not None:
|
||||
raise AssertionError(
|
||||
f"{label}: server returned ErrorMessage [{error.code}]: {error.message}"
|
||||
)
|
||||
if done is None:
|
||||
raise AssertionError(
|
||||
f"{label}: stream did not end with Done "
|
||||
f"({len(chunks)} kv chunks, {len(arrays)} arrays)"
|
||||
)
|
||||
if done.total_tokens <= 0:
|
||||
raise AssertionError(f"{label}: Done reported {done.total_tokens} tokens")
|
||||
if not chunks:
|
||||
raise AssertionError(f"{label}: no KV chunks shipped")
|
||||
|
||||
expected = max(0, n_tokens - 2)
|
||||
if done.total_tokens < expected - 64:
|
||||
raise AssertionError(
|
||||
f"{label}: got {done.total_tokens} tokens, expected ~{expected}"
|
||||
)
|
||||
print(
|
||||
f" [{label}] OK: tokens={done.total_tokens} "
|
||||
f"kv_chunks={len(chunks)} arrays={len(arrays)}"
|
||||
)
|
||||
return done.total_tokens
|
||||
|
||||
|
||||
def main(argv: list[str]) -> int:
|
||||
if len(argv) < 2:
|
||||
print(__doc__)
|
||||
return 2
|
||||
model_id = ModelId(argv[1])
|
||||
|
||||
from exo.download.download_utils import build_model_path
|
||||
|
||||
model_path = build_model_path(model_id)
|
||||
if not model_path.exists():
|
||||
print(f"FAIL: model {model_id} not found at {model_path}")
|
||||
return 1
|
||||
print(f"Loading vLLM engine for {model_id} ({model_path}) ...")
|
||||
|
||||
engine = _build_engine(model_id)
|
||||
failures: list[str] = []
|
||||
try:
|
||||
try:
|
||||
t1 = _run_one(engine, n_tokens=512, label="run1-fresh")
|
||||
except AssertionError as e:
|
||||
failures.append(f"run1: {e}")
|
||||
t1 = 0
|
||||
try:
|
||||
t2 = _run_one(engine, n_tokens=512, label="run2-same-prompt")
|
||||
except AssertionError as e:
|
||||
failures.append(f"run2: {e}")
|
||||
t2 = 0
|
||||
if t1 and t2 and t1 != t2:
|
||||
failures.append(
|
||||
f"run1 returned {t1} tokens but run2 returned {t2} (should match)"
|
||||
)
|
||||
try:
|
||||
ta = _run_one(engine, n_tokens=256, label="run3-shorter")
|
||||
tb = _run_one(engine, n_tokens=768, label="run4-longer")
|
||||
if ta and tb and tb <= ta:
|
||||
failures.append(
|
||||
f"longer prompt should produce more tokens: 256->{ta} 768->{tb}"
|
||||
)
|
||||
except AssertionError as e:
|
||||
failures.append(f"length-variation: {e}")
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
engine.close() # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
if failures:
|
||||
print()
|
||||
print("FAIL")
|
||||
for f in failures:
|
||||
print(f" - {f}")
|
||||
return 1
|
||||
print()
|
||||
print("PASS")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
sys.exit(main(sys.argv))
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
@@ -1,124 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -Eeuo pipefail
|
||||
|
||||
SELF_IP="169.254.100.1"
|
||||
PEER_IP="169.254.100.2"
|
||||
PREFIX="16"
|
||||
IFACE="enP7s7"
|
||||
USE_NM="auto"
|
||||
DRY_RUN=0
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: sudo $(basename "$0") [options]
|
||||
|
||||
Configure a Linux Ethernet interface with a static IPv4 for a host-to-host
|
||||
link to a Mac peer.
|
||||
|
||||
Defaults: this host = ${SELF_IP}/${PREFIX}, peer = ${PEER_IP}, iface = ${IFACE}.
|
||||
|
||||
Options:
|
||||
--iface IFACE Default: ${IFACE}
|
||||
--self-ip IP Default: ${SELF_IP}
|
||||
--peer-ip IP For verification ping. Default: ${PEER_IP}
|
||||
--prefix N Default: ${PREFIX}
|
||||
--no-nm Use 'ip addr' directly (transient, no NetworkManager).
|
||||
--dry-run Print actions without applying.
|
||||
-h, --help Show this help.
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($#)); do
|
||||
case "$1" in
|
||||
--iface)
|
||||
shift
|
||||
IFACE="${1:?}"
|
||||
;;
|
||||
--self-ip)
|
||||
shift
|
||||
SELF_IP="${1:?}"
|
||||
;;
|
||||
--peer-ip)
|
||||
shift
|
||||
PEER_IP="${1:?}"
|
||||
;;
|
||||
--prefix)
|
||||
shift
|
||||
PREFIX="${1:?}"
|
||||
;;
|
||||
--no-nm) USE_NM=no ;;
|
||||
--dry-run) DRY_RUN=1 ;;
|
||||
-h | --help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown arg: $1" >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
[[ $EUID -eq 0 ]] || {
|
||||
echo "Run as root." >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
run() {
|
||||
printf '+'
|
||||
printf ' %q' "$@"
|
||||
printf '\n'
|
||||
((DRY_RUN)) || "$@"
|
||||
}
|
||||
|
||||
ip link show "$IFACE" >/dev/null 2>&1 || {
|
||||
echo "Interface $IFACE does not exist." >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
if [[ $USE_NM == "auto" ]]; then
|
||||
if command -v nmcli >/dev/null 2>&1 && systemctl is-active --quiet NetworkManager 2>/dev/null; then
|
||||
USE_NM=yes
|
||||
else
|
||||
USE_NM=no
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ $USE_NM == "yes" ]]; then
|
||||
CONN="$(nmcli -g GENERAL.CONNECTION device show "$IFACE" 2>/dev/null | head -n1 || true)"
|
||||
if [[ -z $CONN || $CONN == "--" ]]; then
|
||||
CONN="static-${IFACE}"
|
||||
run nmcli connection add type ethernet ifname "$IFACE" con-name "$CONN"
|
||||
fi
|
||||
run nmcli connection modify "$CONN" \
|
||||
connection.interface-name "$IFACE" \
|
||||
connection.autoconnect yes \
|
||||
connection.autoconnect-priority 100 \
|
||||
ipv4.method manual \
|
||||
ipv4.addresses "${SELF_IP}/${PREFIX}" \
|
||||
ipv4.gateway "" \
|
||||
ipv4.dns "" \
|
||||
ipv4.never-default yes \
|
||||
ipv6.method link-local \
|
||||
ipv6.addr-gen-mode stable-privacy
|
||||
run nmcli connection up "$CONN"
|
||||
else
|
||||
run ip link set "$IFACE" up
|
||||
run ip addr flush dev "$IFACE"
|
||||
run ip addr add "${SELF_IP}/${PREFIX}" dev "$IFACE"
|
||||
fi
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
printf '\n'
|
||||
ip -br addr show "$IFACE"
|
||||
printf '\n'
|
||||
if ping -c2 -W2 "$PEER_IP" >/dev/null 2>&1; then
|
||||
echo "OK: $PEER_IP reachable on $IFACE."
|
||||
else
|
||||
echo "WARN: $PEER_IP not reachable yet."
|
||||
echo " Verify the peer is configured (run setup_linklocal_mac.sh on the Mac)."
|
||||
echo " ip neigh show dev $IFACE # check for the peer MAC"
|
||||
fi
|
||||
fi
|
||||
@@ -1,170 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -Eeuo pipefail
|
||||
|
||||
SELF_IP="169.254.100.2"
|
||||
PEER_IP="169.254.100.1"
|
||||
NETMASK="255.255.0.0"
|
||||
IFACE=""
|
||||
DRY_RUN=0
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: sudo $(basename "$0") [options]
|
||||
|
||||
Configure a Mac Ethernet interface with a static IPv4 for a host-to-host link
|
||||
to the DGX/GX10 peer.
|
||||
|
||||
Defaults: this Mac = ${SELF_IP}, peer = ${PEER_IP}, mask = ${NETMASK}.
|
||||
|
||||
Options:
|
||||
--iface IFACE Interface (e.g. en12). Default: auto-detect.
|
||||
--self-ip IP This Mac's address. Default: ${SELF_IP}.
|
||||
--peer-ip IP Peer for verification ping. Default: ${PEER_IP}.
|
||||
--netmask MASK Default: ${NETMASK}.
|
||||
--dry-run Print actions without applying.
|
||||
-h, --help Show this help.
|
||||
EOF
|
||||
}
|
||||
|
||||
while (($#)); do
|
||||
case "$1" in
|
||||
--iface)
|
||||
shift
|
||||
IFACE="${1:?}"
|
||||
;;
|
||||
--self-ip)
|
||||
shift
|
||||
SELF_IP="${1:?}"
|
||||
;;
|
||||
--peer-ip)
|
||||
shift
|
||||
PEER_IP="${1:?}"
|
||||
;;
|
||||
--netmask)
|
||||
shift
|
||||
NETMASK="${1:?}"
|
||||
;;
|
||||
--dry-run) DRY_RUN=1 ;;
|
||||
-h | --help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown arg: $1" >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
[[ $EUID -eq 0 ]] || {
|
||||
echo "Run with sudo." >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
run() {
|
||||
printf '+'
|
||||
printf ' %q' "$@"
|
||||
printf '\n'
|
||||
((DRY_RUN)) || "$@"
|
||||
}
|
||||
|
||||
target_subnet_prefix() {
|
||||
local ip="$1"
|
||||
printf '%s.' "${ip%.*}"
|
||||
}
|
||||
|
||||
iface_score() {
|
||||
local iface="$1" info subnet
|
||||
info="$(ifconfig "$iface" 2>/dev/null || true)"
|
||||
[[ -n $info ]] || {
|
||||
echo 0
|
||||
return
|
||||
}
|
||||
grep -q 'status: active' <<<"$info" || {
|
||||
echo 0
|
||||
return
|
||||
}
|
||||
subnet="$(target_subnet_prefix "$SELF_IP")"
|
||||
if grep -qE "inet ${subnet//./\\.}" <<<"$info"; then
|
||||
echo 100
|
||||
return
|
||||
fi
|
||||
if grep -qE 'inet 169\.254\.' <<<"$info"; then
|
||||
echo 80
|
||||
return
|
||||
fi
|
||||
if ! grep -qE '^[[:space:]]*inet ' <<<"$info"; then
|
||||
echo 60
|
||||
return
|
||||
fi
|
||||
echo 10
|
||||
}
|
||||
|
||||
detect_iface() {
|
||||
local best="" best_score=0 iface score
|
||||
for iface in $(ifconfig -l); do
|
||||
[[ $iface =~ ^en[0-9]+$ ]] || continue
|
||||
score="$(iface_score "$iface")"
|
||||
if ((score > best_score)); then
|
||||
best="$iface"
|
||||
best_score="$score"
|
||||
fi
|
||||
done
|
||||
((best_score >= 60)) || return 1
|
||||
printf '%s\n' "$best"
|
||||
}
|
||||
|
||||
iface_to_service() {
|
||||
local iface="$1" line port=""
|
||||
while IFS= read -r line; do
|
||||
if [[ $line == "Hardware Port: "* ]]; then
|
||||
port="${line#Hardware Port: }"
|
||||
elif [[ $line == "Device: $iface" ]]; then
|
||||
printf '%s\n' "$port"
|
||||
return 0
|
||||
fi
|
||||
done < <(networksetup -listallhardwareports)
|
||||
return 1
|
||||
}
|
||||
|
||||
if [[ -z $IFACE ]]; then
|
||||
IFACE="$(detect_iface || true)"
|
||||
[[ -n $IFACE ]] || {
|
||||
echo "Could not auto-detect a wired interface. Pass --iface enX." >&2
|
||||
echo "Active interfaces:" >&2
|
||||
ifconfig -l | tr ' ' '\n' | grep -E '^en[0-9]+$' | while read -r i; do
|
||||
printf ' %-6s %s\n' "$i" "$(ifconfig "$i" | grep -E 'status:|inet ' | tr '\n' ' ')" >&2
|
||||
done
|
||||
exit 1
|
||||
}
|
||||
echo "Auto-detected interface: $IFACE"
|
||||
fi
|
||||
|
||||
ifconfig "$IFACE" >/dev/null 2>&1 || {
|
||||
echo "Interface $IFACE does not exist." >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
SERVICE="$(iface_to_service "$IFACE" || true)"
|
||||
[[ -n $SERVICE ]] || {
|
||||
echo "No network service maps to $IFACE. Check System Settings -> Network." >&2
|
||||
exit 1
|
||||
}
|
||||
echo "Network service: $SERVICE"
|
||||
|
||||
run networksetup -setmanual "$SERVICE" "$SELF_IP" "$NETMASK" ""
|
||||
|
||||
if ((!DRY_RUN)); then
|
||||
printf '\n'
|
||||
ifconfig "$IFACE" | grep -E 'inet |status:'
|
||||
printf '\n'
|
||||
if ping -c2 -t3 "$PEER_IP" >/dev/null 2>&1; then
|
||||
echo "OK: $PEER_IP reachable on $IFACE."
|
||||
else
|
||||
echo "WARN: $PEER_IP not reachable yet."
|
||||
echo " Verify the peer is configured (run setup_linklocal_dgx.sh on the GX10)."
|
||||
echo " arp -an -i $IFACE # check for the peer MAC"
|
||||
fi
|
||||
fi
|
||||
@@ -20,7 +20,6 @@ from exo.shared.types.chunks import (
|
||||
TokenChunk,
|
||||
ToolCallChunk,
|
||||
)
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.text_generation import (
|
||||
Base64Image,
|
||||
InputMessage,
|
||||
@@ -181,7 +180,6 @@ def ollama_request_to_text_generation(
|
||||
|
||||
|
||||
async def generate_ollama_chat_stream(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
@@ -264,7 +262,6 @@ async def generate_ollama_chat_stream(
|
||||
|
||||
|
||||
async def collect_ollama_chat_response(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
@@ -369,7 +366,6 @@ def ollama_generate_request_to_text_generation(
|
||||
|
||||
|
||||
async def generate_ollama_generate_stream(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
@@ -442,7 +438,6 @@ async def generate_ollama_generate_stream(
|
||||
|
||||
|
||||
async def collect_ollama_generate_response(
|
||||
_command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[
|
||||
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
|
||||
],
|
||||
|
||||
+278
-319
@@ -5,6 +5,7 @@ import json
|
||||
import random
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
@@ -12,7 +13,7 @@ from typing import Annotated, Any, Literal, cast
|
||||
from uuid import uuid4
|
||||
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, ClosedResourceError
|
||||
from exo_net import NetSender, PySession
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||||
@@ -21,6 +22,7 @@ from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType
|
||||
from hypercorn.config import Config
|
||||
from hypercorn.typing import ASGIFramework
|
||||
from loguru import logger
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from exo.api.adapters.chat_completions import (
|
||||
chat_request_to_text_generation,
|
||||
@@ -133,19 +135,19 @@ from exo.shared.constants import (
|
||||
)
|
||||
from exo.shared.election import ElectionMessage
|
||||
from exo.shared.logging import InterceptLogger
|
||||
from exo.shared.models import model_cards
|
||||
from exo.shared.models.model_cards import (
|
||||
ModelCard,
|
||||
ModelId,
|
||||
add_to_card_cache,
|
||||
get_card,
|
||||
get_model_cards,
|
||||
)
|
||||
from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
|
||||
from exo.shared.types.chunks import (
|
||||
ErrorChunk,
|
||||
ImageChunk,
|
||||
Chunk,
|
||||
ImageGenerationChunk,
|
||||
InputImageChunk,
|
||||
PrefillProgressChunk,
|
||||
StatusChunk,
|
||||
TextGenerationChunk,
|
||||
TokenChunk,
|
||||
ToolCallChunk,
|
||||
)
|
||||
@@ -159,7 +161,6 @@ from exo.shared.types.commands import (
|
||||
DeleteInstance,
|
||||
DeleteInstanceLink,
|
||||
DownloadCommand,
|
||||
ForwarderCommand,
|
||||
ForwarderDownloadCommand,
|
||||
ImageEdits,
|
||||
ImageGeneration,
|
||||
@@ -173,7 +174,6 @@ from exo.shared.types.commands import (
|
||||
)
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
Event,
|
||||
IndexedEvent,
|
||||
InstanceDeleted,
|
||||
@@ -199,7 +199,7 @@ from exo.shared.types.worker.downloads import DownloadCompleted
|
||||
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.channels import Receiver, Sender
|
||||
from exo.utils.disk_event_log import DiskEventLog
|
||||
from exo.utils.power_sampler import PowerSampler
|
||||
from exo.utils.task_group import TaskGroup
|
||||
@@ -232,6 +232,95 @@ def _require_disaggregation_enabled() -> None:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transport:
|
||||
session: PySession
|
||||
cancel_scopes: dict[CommandId, anyio.CancelScope] = field(
|
||||
init=False, default_factory=dict
|
||||
)
|
||||
command_sender: NetSender = field(init=False)
|
||||
paused: bool = field(init=False, default=False)
|
||||
paused_ev: anyio.Event = field(init=False, default_factory=anyio.Event)
|
||||
|
||||
def __post_init__(self):
|
||||
# TODO: retire root keyspace
|
||||
self.command_sender = self.session.net_sender("orchestrator")
|
||||
|
||||
async def send_command(self, command: Command) -> bool:
|
||||
while self.paused:
|
||||
await self.paused_ev.wait()
|
||||
return await self.command_sender.send(command.model_dump_json().encode("utf-8"))
|
||||
|
||||
async def stream_text(
|
||||
self,
|
||||
command_id: CommandId,
|
||||
) -> AsyncGenerator[TextGenerationChunk | StatusChunk]:
|
||||
async for chunk in self.stream(command_id):
|
||||
if isinstance(chunk, (TextGenerationChunk | StatusChunk)):
|
||||
yield chunk
|
||||
|
||||
async def stream_images(
|
||||
self,
|
||||
command_id: CommandId,
|
||||
) -> AsyncGenerator[ImageGenerationChunk]:
|
||||
async for chunk in self.stream(command_id):
|
||||
if isinstance(chunk, (ImageGenerationChunk)):
|
||||
yield chunk
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
command_id: CommandId,
|
||||
) -> AsyncGenerator[Chunk,]:
|
||||
try:
|
||||
with anyio.CancelScope() as cs:
|
||||
self.cancel_scopes[command_id] = cs
|
||||
# recv from any node
|
||||
receiver = self.session.net_receiver(
|
||||
f"runners/*/active_tasks/{command_id}/chunks"
|
||||
)
|
||||
while True:
|
||||
data = await receiver.recv()
|
||||
if data is None:
|
||||
logger.warning(
|
||||
"stream terminated early without finish reason EOF"
|
||||
)
|
||||
break
|
||||
yield (
|
||||
chunk := cast(
|
||||
Chunk,
|
||||
TypeAdapter(Chunk).validate_json(
|
||||
data, strict=True, extra="forbid"
|
||||
),
|
||||
)
|
||||
)
|
||||
if (
|
||||
not isinstance(chunk, StatusChunk)
|
||||
and chunk.finish_reason is not None
|
||||
):
|
||||
break
|
||||
except anyio.get_cancelled_exc_class():
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
TaskCancelled(cancelled_command_id=command_id)
|
||||
.model_dump_json()
|
||||
.encode("utf-8")
|
||||
)
|
||||
finally:
|
||||
self.cancel_scopes.pop(command_id, None)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
TaskFinished(finished_command_id=command_id)
|
||||
.model_dump_json()
|
||||
.encode("utf-8")
|
||||
)
|
||||
|
||||
def cancel(self, command_id: CommandId) -> bool:
|
||||
if (cs := self.cancel_scopes.get(command_id, None)) is not None:
|
||||
cs.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class API:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -239,15 +328,14 @@ class API:
|
||||
*,
|
||||
port: int,
|
||||
event_receiver: Receiver[IndexedEvent],
|
||||
command_sender: Sender[ForwarderCommand],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
# This lets us pause the API if an election is running
|
||||
election_receiver: Receiver[ElectionMessage],
|
||||
session: PySession,
|
||||
) -> 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.event_receiver = event_receiver
|
||||
self.election_receiver = election_receiver
|
||||
@@ -256,9 +344,6 @@ class API:
|
||||
self.port = port
|
||||
self._sent_image_hashes: set[str] = set()
|
||||
|
||||
self.paused: bool = False
|
||||
self.paused_ev: anyio.Event = anyio.Event()
|
||||
|
||||
self.app = FastAPI()
|
||||
|
||||
@self.app.middleware("http")
|
||||
@@ -282,13 +367,7 @@ class API:
|
||||
name="dashboard",
|
||||
)
|
||||
|
||||
self._text_generation_queues: dict[
|
||||
CommandId,
|
||||
Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk],
|
||||
] = {}
|
||||
self._image_generation_queues: dict[
|
||||
CommandId, Sender[ImageChunk | ErrorChunk]
|
||||
] = {}
|
||||
self.transport = Transport(session)
|
||||
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
|
||||
@@ -309,9 +388,9 @@ class API:
|
||||
def unpause(self, result_clock: int):
|
||||
logger.info("Unpausing API")
|
||||
self.last_completed_election = result_clock
|
||||
self.paused = False
|
||||
self.paused_ev.set()
|
||||
self.paused_ev = anyio.Event()
|
||||
self.transport.paused = False
|
||||
self.transport.paused_ev.set()
|
||||
self.transport.paused_ev = anyio.Event()
|
||||
|
||||
def _setup_exception_handlers(self) -> None:
|
||||
self.app.exception_handler(HTTPException)(self.http_exception_handler)
|
||||
@@ -426,7 +505,7 @@ class API:
|
||||
instance_meta=payload.instance_meta,
|
||||
min_nodes=payload.min_nodes,
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
|
||||
return CreateInstanceResponse(
|
||||
message="Command received.",
|
||||
@@ -451,7 +530,7 @@ class API:
|
||||
command = CreateInstance(
|
||||
instance=instance,
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
|
||||
return CreateInstanceResponse(
|
||||
message="Command received.",
|
||||
@@ -481,6 +560,7 @@ class API:
|
||||
topology=self.state.topology,
|
||||
current_instances=self.state.instances,
|
||||
download_status=self.state.downloads,
|
||||
node_rdma_ctl=self.state.node_rdma_ctl,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
@@ -526,8 +606,8 @@ class API:
|
||||
)
|
||||
]
|
||||
)
|
||||
if any(self.state.node_vllm.values()):
|
||||
instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
|
||||
# TODO: PDD
|
||||
# instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
|
||||
|
||||
for sharding, instance_meta, min_nodes in instance_combinations:
|
||||
try:
|
||||
@@ -544,6 +624,7 @@ class API:
|
||||
current_instances=self.state.instances,
|
||||
required_nodes=required_nodes,
|
||||
download_status=self.state.downloads,
|
||||
node_rdma_ctl=self.state.node_rdma_ctl,
|
||||
)
|
||||
except ValueError as exc:
|
||||
if (model_card.model_id, sharding, instance_meta, 0) not in seen:
|
||||
@@ -632,7 +713,7 @@ class API:
|
||||
command = DeleteInstance(
|
||||
instance_id=instance_id,
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return DeleteInstanceResponse(
|
||||
message="Command received.",
|
||||
command_id=command.command_id,
|
||||
@@ -640,10 +721,7 @@ class API:
|
||||
)
|
||||
|
||||
async def get_feature_flags(self) -> dict[str, bool]:
|
||||
return {
|
||||
"disaggregation": ENABLE_DISAGGREGATION,
|
||||
"vllm_available": any(self.state.node_vllm.values()),
|
||||
}
|
||||
return {"disaggregation": ENABLE_DISAGGREGATION}
|
||||
|
||||
async def list_instance_links(self) -> list[InstanceLink]:
|
||||
if not ENABLE_DISAGGREGATION:
|
||||
@@ -670,7 +748,7 @@ class API:
|
||||
prefill_instances=list(body.prefill_instances),
|
||||
decode_instances=list(body.decode_instances),
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return InstanceLinkResponse(
|
||||
message="Command received.", command_id=command.command_id
|
||||
)
|
||||
@@ -680,64 +758,24 @@ class API:
|
||||
) -> InstanceLinkResponse:
|
||||
_require_disaggregation_enabled()
|
||||
command = DeleteInstanceLink(link_id=link_id)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return InstanceLinkResponse(
|
||||
message="Command received.", command_id=command.command_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:
|
||||
if self.transport.cancel(command_id):
|
||||
return CancelCommandResponse(
|
||||
message="Command cancelled.",
|
||||
command_id=command_id,
|
||||
)
|
||||
else:
|
||||
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[
|
||||
TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None
|
||||
]:
|
||||
"""Yield chunks for a given command until completion.
|
||||
|
||||
This is the internal low-level stream used by all API adapters.
|
||||
"""
|
||||
try:
|
||||
self._text_generation_queues[command_id], recv = channel[
|
||||
TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk
|
||||
]()
|
||||
|
||||
with recv as token_chunks:
|
||||
async for chunk in token_chunks:
|
||||
yield chunk
|
||||
if isinstance(chunk, PrefillProgressChunk):
|
||||
continue
|
||||
if chunk.finish_reason is not None:
|
||||
break
|
||||
|
||||
except anyio.get_cancelled_exc_class():
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
await self._send(TaskFinished(finished_command_id=command_id))
|
||||
if command_id in self._text_generation_queues:
|
||||
del self._text_generation_queues[command_id]
|
||||
|
||||
async def _collect_text_generation_with_stats(
|
||||
self, command_id: CommandId
|
||||
) -> BenchChatCompletionResponse:
|
||||
@@ -752,7 +790,7 @@ class API:
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(sampler.run)
|
||||
|
||||
async for chunk in self._token_chunk_stream(command_id):
|
||||
async for chunk in self.transport.stream_text(command_id):
|
||||
if isinstance(chunk, PrefillProgressChunk):
|
||||
continue
|
||||
|
||||
@@ -819,7 +857,7 @@ class API:
|
||||
images = task_params.images
|
||||
if not images:
|
||||
command = TextGeneration(task_params=task_params)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return command
|
||||
|
||||
hashes = [hashlib.sha256(img.encode("ascii")).hexdigest() for img in images]
|
||||
@@ -836,7 +874,7 @@ class API:
|
||||
new_images.append((idx, img))
|
||||
|
||||
if not new_images:
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return command
|
||||
|
||||
all_chunks: list[tuple[int, str]] = []
|
||||
@@ -845,7 +883,7 @@ class API:
|
||||
all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
|
||||
|
||||
for global_idx, (img_idx, chunk_data) in enumerate(all_chunks):
|
||||
await self._send(
|
||||
await self.transport.send_command(
|
||||
SendInputChunk(
|
||||
chunk=InputImageChunk(
|
||||
model=task_params.model,
|
||||
@@ -858,7 +896,7 @@ class API:
|
||||
)
|
||||
)
|
||||
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return command
|
||||
|
||||
async def chat_completions(
|
||||
@@ -878,7 +916,7 @@ class API:
|
||||
with_sse_keepalive(
|
||||
generate_chat_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
@@ -892,7 +930,7 @@ class API:
|
||||
return StreamingResponse(
|
||||
collect_chat_response(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -921,7 +959,7 @@ class API:
|
||||
with_sse_keepalive(
|
||||
generate_chat_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
@@ -1027,7 +1065,7 @@ class API:
|
||||
command = ImageGeneration(
|
||||
task_params=payload,
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
|
||||
# Check if streaming is requested
|
||||
if payload.stream and payload.partial_images and payload.partial_images > 0:
|
||||
@@ -1063,105 +1101,85 @@ class API:
|
||||
image_metadata: dict[tuple[int, bool], tuple[int | None, int | None]] = {}
|
||||
images_complete = 0
|
||||
|
||||
try:
|
||||
self._image_generation_queues[command_id], recv = channel[
|
||||
ImageChunk | ErrorChunk
|
||||
]()
|
||||
|
||||
with recv as chunks:
|
||||
async for chunk in chunks:
|
||||
if chunk.finish_reason == "error":
|
||||
error_response = ErrorResponse(
|
||||
error=ErrorInfo(
|
||||
message=chunk.error_message or "Internal server error",
|
||||
type="InternalServerError",
|
||||
code=500,
|
||||
)
|
||||
)
|
||||
yield f"data: {error_response.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return
|
||||
|
||||
key = (chunk.image_index, chunk.is_partial)
|
||||
|
||||
if key not in image_chunks:
|
||||
image_chunks[key] = {}
|
||||
image_total_chunks[key] = chunk.total_chunks
|
||||
image_metadata[key] = (
|
||||
chunk.partial_index,
|
||||
chunk.total_partials,
|
||||
)
|
||||
|
||||
image_chunks[key][chunk.chunk_index] = chunk.data
|
||||
|
||||
# Check if this image is complete
|
||||
if len(image_chunks[key]) == image_total_chunks[key]:
|
||||
full_data = "".join(
|
||||
image_chunks[key][i] for i in range(len(image_chunks[key]))
|
||||
)
|
||||
|
||||
partial_idx, total_partials = image_metadata[key]
|
||||
|
||||
if chunk.is_partial:
|
||||
# Yield partial image event (always use b64_json for partials)
|
||||
event_data = {
|
||||
"type": "partial",
|
||||
"image_index": chunk.image_index,
|
||||
"partial_index": partial_idx,
|
||||
"total_partials": total_partials,
|
||||
"format": str(chunk.format),
|
||||
"data": {
|
||||
"b64_json": full_data
|
||||
if response_format == "b64_json"
|
||||
else None,
|
||||
},
|
||||
}
|
||||
yield f"data: {json.dumps(event_data)}\n\n"
|
||||
else:
|
||||
# Final image
|
||||
if response_format == "url":
|
||||
image_bytes = base64.b64decode(full_data)
|
||||
content_type = _format_to_content_type(chunk.format)
|
||||
stored = self._image_store.store(
|
||||
image_bytes, content_type
|
||||
)
|
||||
url = self._build_image_url(request, stored.image_id)
|
||||
event_data = {
|
||||
"type": "final",
|
||||
"image_index": chunk.image_index,
|
||||
"format": str(chunk.format),
|
||||
"data": {"url": url},
|
||||
}
|
||||
else:
|
||||
event_data = {
|
||||
"type": "final",
|
||||
"image_index": chunk.image_index,
|
||||
"format": str(chunk.format),
|
||||
"data": {"b64_json": full_data},
|
||||
}
|
||||
yield f"data: {json.dumps(event_data)}\n\n"
|
||||
images_complete += 1
|
||||
|
||||
if images_complete >= num_images:
|
||||
yield "data: [DONE]\n\n"
|
||||
break
|
||||
|
||||
# Clean up completed image chunks
|
||||
del image_chunks[key]
|
||||
del image_total_chunks[key]
|
||||
del image_metadata[key]
|
||||
|
||||
except anyio.get_cancelled_exc_class():
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
async for chunk in self.transport.stream_images(command_id):
|
||||
if chunk.finish_reason == "error":
|
||||
error_response = ErrorResponse(
|
||||
error=ErrorInfo(
|
||||
message=chunk.error_message or "Internal server error",
|
||||
type="InternalServerError",
|
||||
code=500,
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
await self._send(TaskFinished(finished_command_id=command_id))
|
||||
if command_id in self._image_generation_queues:
|
||||
del self._image_generation_queues[command_id]
|
||||
yield f"data: {error_response.model_dump_json()}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
return
|
||||
|
||||
key = (chunk.image_index, chunk.is_partial)
|
||||
|
||||
if key not in image_chunks:
|
||||
image_chunks[key] = {}
|
||||
image_total_chunks[key] = chunk.total_chunks
|
||||
image_metadata[key] = (
|
||||
chunk.partial_index,
|
||||
chunk.total_partials,
|
||||
)
|
||||
|
||||
image_chunks[key][chunk.chunk_index] = chunk.data
|
||||
|
||||
# Check if this image is complete
|
||||
if len(image_chunks[key]) == image_total_chunks[key]:
|
||||
full_data = "".join(
|
||||
image_chunks[key][i] for i in range(len(image_chunks[key]))
|
||||
)
|
||||
|
||||
partial_idx, total_partials = image_metadata[key]
|
||||
|
||||
if chunk.is_partial:
|
||||
# Yield partial image event (always use b64_json for partials)
|
||||
event_data = {
|
||||
"type": "partial",
|
||||
"image_index": chunk.image_index,
|
||||
"partial_index": partial_idx,
|
||||
"total_partials": total_partials,
|
||||
"format": str(chunk.format),
|
||||
"data": {
|
||||
"b64_json": full_data
|
||||
if response_format == "b64_json"
|
||||
else None,
|
||||
},
|
||||
}
|
||||
yield f"data: {json.dumps(event_data)}\n\n"
|
||||
else:
|
||||
# Final image
|
||||
if response_format == "url":
|
||||
image_bytes = base64.b64decode(full_data)
|
||||
content_type = _format_to_content_type(chunk.format)
|
||||
stored = self._image_store.store(image_bytes, content_type)
|
||||
url = self._build_image_url(request, stored.image_id)
|
||||
event_data = {
|
||||
"type": "final",
|
||||
"image_index": chunk.image_index,
|
||||
"format": str(chunk.format),
|
||||
"data": {"url": url},
|
||||
}
|
||||
else:
|
||||
event_data = {
|
||||
"type": "final",
|
||||
"image_index": chunk.image_index,
|
||||
"format": str(chunk.format),
|
||||
"data": {"b64_json": full_data},
|
||||
}
|
||||
yield f"data: {json.dumps(event_data)}\n\n"
|
||||
images_complete += 1
|
||||
|
||||
if images_complete >= num_images:
|
||||
yield "data: [DONE]\n\n"
|
||||
break
|
||||
|
||||
# Clean up completed image chunks
|
||||
del image_chunks[key]
|
||||
del image_total_chunks[key]
|
||||
del image_metadata[key]
|
||||
|
||||
async def _collect_image_chunks(
|
||||
self,
|
||||
@@ -1180,74 +1198,55 @@ class API:
|
||||
images_complete = 0
|
||||
stats: ImageGenerationStats | None = None
|
||||
|
||||
try:
|
||||
self._image_generation_queues[command_id], recv = channel[
|
||||
ImageChunk | ErrorChunk
|
||||
]()
|
||||
|
||||
while images_complete < num_images:
|
||||
with recv as chunks:
|
||||
async for chunk in chunks:
|
||||
if chunk.finish_reason == "error":
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=chunk.error_message or "Internal server error",
|
||||
)
|
||||
|
||||
if chunk.is_partial:
|
||||
continue
|
||||
|
||||
if chunk.image_index not in image_chunks:
|
||||
image_chunks[chunk.image_index] = {}
|
||||
image_total_chunks[chunk.image_index] = chunk.total_chunks
|
||||
image_formats[chunk.image_index] = chunk.format
|
||||
|
||||
image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data
|
||||
|
||||
if capture_stats and chunk.stats is not None:
|
||||
stats = chunk.stats
|
||||
|
||||
if (
|
||||
len(image_chunks[chunk.image_index])
|
||||
== image_total_chunks[chunk.image_index]
|
||||
):
|
||||
images_complete += 1
|
||||
|
||||
if images_complete >= num_images:
|
||||
break
|
||||
|
||||
images: list[ImageData] = []
|
||||
for image_idx in range(num_images):
|
||||
chunks_dict = image_chunks[image_idx]
|
||||
full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
|
||||
if response_format == "url" and request is not None:
|
||||
image_bytes = base64.b64decode(full_data)
|
||||
content_type = _format_to_content_type(image_formats.get(image_idx))
|
||||
stored = self._image_store.store(image_bytes, content_type)
|
||||
url = self._build_image_url(request, stored.image_id)
|
||||
images.append(ImageData(b64_json=None, url=url))
|
||||
else:
|
||||
images.append(
|
||||
ImageData(
|
||||
b64_json=full_data
|
||||
if response_format == "b64_json"
|
||||
else None,
|
||||
url=None,
|
||||
)
|
||||
while images_complete < num_images:
|
||||
async for chunk in self.transport.stream_images(command_id):
|
||||
if chunk.finish_reason == "error":
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=chunk.error_message or "Internal server error",
|
||||
)
|
||||
|
||||
return (images, stats if capture_stats else None)
|
||||
except anyio.get_cancelled_exc_class():
|
||||
command = TaskCancelled(cancelled_command_id=command_id)
|
||||
with anyio.CancelScope(shield=True):
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
if chunk.is_partial:
|
||||
continue
|
||||
|
||||
if chunk.image_index not in image_chunks:
|
||||
image_chunks[chunk.image_index] = {}
|
||||
image_total_chunks[chunk.image_index] = chunk.total_chunks
|
||||
image_formats[chunk.image_index] = chunk.format
|
||||
|
||||
image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data
|
||||
|
||||
if capture_stats and chunk.stats is not None:
|
||||
stats = chunk.stats
|
||||
|
||||
if (
|
||||
len(image_chunks[chunk.image_index])
|
||||
== image_total_chunks[chunk.image_index]
|
||||
):
|
||||
images_complete += 1
|
||||
|
||||
if images_complete >= num_images:
|
||||
break
|
||||
|
||||
images: list[ImageData] = []
|
||||
for image_idx in range(num_images):
|
||||
chunks_dict = image_chunks[image_idx]
|
||||
full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict)))
|
||||
if response_format == "url" and request is not None:
|
||||
image_bytes = base64.b64decode(full_data)
|
||||
content_type = _format_to_content_type(image_formats.get(image_idx))
|
||||
stored = self._image_store.store(image_bytes, content_type)
|
||||
url = self._build_image_url(request, stored.image_id)
|
||||
images.append(ImageData(b64_json=None, url=url))
|
||||
else:
|
||||
images.append(
|
||||
ImageData(
|
||||
b64_json=full_data if response_format == "b64_json" else None,
|
||||
url=None,
|
||||
)
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
await self._send(TaskFinished(finished_command_id=command_id))
|
||||
if command_id in self._image_generation_queues:
|
||||
del self._image_generation_queues[command_id]
|
||||
|
||||
return (images, stats if capture_stats else None)
|
||||
|
||||
async def _collect_image_generation(
|
||||
self,
|
||||
@@ -1297,7 +1296,7 @@ class API:
|
||||
command = ImageGeneration(
|
||||
task_params=payload,
|
||||
)
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
|
||||
return await self._collect_image_generation_with_stats(
|
||||
request=request,
|
||||
@@ -1360,7 +1359,7 @@ class API:
|
||||
f"Sending input image: {len(image_data)} bytes in {total_chunks} chunks"
|
||||
)
|
||||
for chunk_index, chunk_data in enumerate(data_chunks):
|
||||
await self._send(
|
||||
await self.transport.send_command(
|
||||
SendInputChunk(
|
||||
chunk=InputImageChunk(
|
||||
model=resolved_model,
|
||||
@@ -1372,7 +1371,7 @@ class API:
|
||||
)
|
||||
)
|
||||
|
||||
await self._send(command)
|
||||
await self.transport.send_command(command)
|
||||
return command
|
||||
|
||||
async def image_edits(
|
||||
@@ -1500,7 +1499,7 @@ class API:
|
||||
generate_claude_stream(
|
||||
command.command_id,
|
||||
payload.model,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
@@ -1515,7 +1514,7 @@ class API:
|
||||
collect_claude_response(
|
||||
command.command_id,
|
||||
payload.model,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -1536,7 +1535,7 @@ class API:
|
||||
generate_responses_stream(
|
||||
command.command_id,
|
||||
payload.model,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
@@ -1552,7 +1551,7 @@ class API:
|
||||
collect_responses_response(
|
||||
command.command_id,
|
||||
payload.model,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -1578,8 +1577,7 @@ class API:
|
||||
if payload.stream:
|
||||
return StreamingResponse(
|
||||
generate_ollama_chat_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
@@ -1591,8 +1589,7 @@ class API:
|
||||
else:
|
||||
return StreamingResponse(
|
||||
collect_ollama_chat_response(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -1614,8 +1611,7 @@ class API:
|
||||
if payload.stream:
|
||||
return StreamingResponse(
|
||||
generate_ollama_generate_stream(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/x-ndjson",
|
||||
headers={
|
||||
@@ -1627,8 +1623,7 @@ class API:
|
||||
else:
|
||||
return StreamingResponse(
|
||||
collect_ollama_generate_response(
|
||||
command.command_id,
|
||||
self._token_chunk_stream(command.command_id),
|
||||
self.transport.stream_text(command.command_id),
|
||||
),
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -1636,17 +1631,16 @@ class API:
|
||||
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()
|
||||
downloaded_model_ids: set[ModelId] = 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
|
||||
c
|
||||
for c in await model_cards.card_cache.list_all()
|
||||
if c.model_id in downloaded_model_ids
|
||||
]
|
||||
|
||||
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
@@ -1659,8 +1653,8 @@ class API:
|
||||
size=card.storage_size.in_bytes,
|
||||
digest="sha256:000000000000",
|
||||
details=OllamaModelDetails(
|
||||
family=none_if_empty(card.family),
|
||||
quantization_level=none_if_empty(card.quantization),
|
||||
family=card.family or None,
|
||||
quantization_level=card.quantization or None,
|
||||
),
|
||||
)
|
||||
for card in cards
|
||||
@@ -1723,7 +1717,7 @@ class API:
|
||||
|
||||
async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
|
||||
"""Returns list of available models, optionally filtered by being downloaded."""
|
||||
cards = await get_model_cards()
|
||||
cards = await model_cards.card_cache.list_all()
|
||||
|
||||
if status == "downloaded":
|
||||
downloaded_model_ids: set[str] = set()
|
||||
@@ -1751,7 +1745,6 @@ class API:
|
||||
capabilities=card.capabilities,
|
||||
reasoning_dialect=card.reasoning_dialect,
|
||||
context_length=card.context_length,
|
||||
requires_vllm=card.requires_vllm,
|
||||
)
|
||||
for card in cards
|
||||
]
|
||||
@@ -1766,16 +1759,13 @@ class API:
|
||||
status_code=400, detail=f"Failed to fetch model: {exc}"
|
||||
) from exc
|
||||
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(
|
||||
origin=self._system_id,
|
||||
command=AddCustomModelCard(model_card=card),
|
||||
)
|
||||
await self.transport.command_sender.send(
|
||||
AddCustomModelCard(model_card=card).model_dump_json().encode("utf-8")
|
||||
)
|
||||
|
||||
# Immediately update the local cache so the subsequent GET /models
|
||||
# returns the new model without waiting for the event round-trip.
|
||||
add_to_card_cache(card)
|
||||
model_cards.card_cache.cc[card.model_id] = card
|
||||
|
||||
return ModelListModel(
|
||||
id=card.model_id,
|
||||
@@ -1791,15 +1781,12 @@ class API:
|
||||
|
||||
async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
|
||||
"""Delete a user-added custom model card and sync deletion across the cluster."""
|
||||
card = get_card(model_id)
|
||||
card = model_cards.card_cache.get(model_id)
|
||||
if card is None or not card.is_custom:
|
||||
raise HTTPException(status_code=404, detail="Custom model card not found")
|
||||
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(
|
||||
origin=self._system_id,
|
||||
command=DeleteCustomModelCard(model_id=model_id),
|
||||
)
|
||||
await self.transport.command_sender.send(
|
||||
DeleteCustomModelCard(model_id=model_id).model_dump_json().encode("utf-8")
|
||||
)
|
||||
|
||||
return JSONResponse(
|
||||
@@ -1864,7 +1851,6 @@ class API:
|
||||
shutdown_ev.set()
|
||||
finally:
|
||||
self._event_log.close()
|
||||
self.command_sender.close()
|
||||
self.event_receiver.close()
|
||||
|
||||
async def run_api(self, ev: anyio.Event):
|
||||
@@ -1888,23 +1874,6 @@ class API:
|
||||
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, ClosedResourceError):
|
||||
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, ClosedResourceError):
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
if isinstance(event, InstanceDeleted):
|
||||
self._close_streams_for_instance(event.instance_id)
|
||||
if isinstance(event, TracesMerged):
|
||||
@@ -1919,10 +1888,7 @@ class API:
|
||||
task, (TextGenerationTask, ImageGenerationTask, ImageEditsTask)
|
||||
):
|
||||
continue
|
||||
if sender := self._text_generation_queues.pop(task.command_id, None):
|
||||
sender.close()
|
||||
if sender := self._image_generation_queues.pop(task.command_id, None):
|
||||
sender.close()
|
||||
self.transport.cancel(task.command_id)
|
||||
|
||||
def _save_merged_trace(self, event: TracesMerged) -> None:
|
||||
traces = [
|
||||
@@ -1943,7 +1909,7 @@ class API:
|
||||
with self.election_receiver as ems:
|
||||
async for message in ems:
|
||||
if message.clock > self.last_completed_election:
|
||||
self.paused = True
|
||||
self.transport.paused = True
|
||||
|
||||
async def _cleanup_expired_images(self):
|
||||
"""Periodically clean up expired images from the store."""
|
||||
@@ -1954,13 +1920,6 @@ class API:
|
||||
if removed > 0:
|
||||
logger.debug(f"Cleaned up {removed} expired images")
|
||||
|
||||
async def _send(self, command: Command):
|
||||
while self.paused:
|
||||
await self.paused_ev.wait()
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(origin=self._system_id, command=command)
|
||||
)
|
||||
|
||||
async def _send_download(self, command: DownloadCommand):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(origin=self._system_id, command=command)
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
# pyright: reportUnusedFunction=false, reportAny=false
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from exo.api.main import API
|
||||
from exo.api.main import API, Transport
|
||||
from exo.shared.types.common import CommandId
|
||||
|
||||
|
||||
@@ -15,9 +15,9 @@ def _make_api() -> Any:
|
||||
app = FastAPI()
|
||||
api = object.__new__(API)
|
||||
api.app = app
|
||||
api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._send = AsyncMock() # pyright: ignore[reportPrivateUsage]
|
||||
api.transport = object.__new__(Transport)
|
||||
api.transport.cancel = AsyncMock()
|
||||
api.transport.send_command = AsyncMock()
|
||||
api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage]
|
||||
app.post("/v1/cancel/{command_id}")(api.cancel_command)
|
||||
return api
|
||||
@@ -43,16 +43,14 @@ def test_cancel_active_text_generation() -> None:
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("text-cmd-123")
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[cid] = sender
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
api.transport.cancel.assert_called_once()
|
||||
api.transport.send_command.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
|
||||
@@ -63,15 +61,13 @@ def test_cancel_active_image_generation() -> None:
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("img-cmd-456")
|
||||
sender = MagicMock()
|
||||
api._image_generation_queues[cid] = sender
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
api.transport.cancel.assert_called_once()
|
||||
api.transport.send_command.assert_called_once()
|
||||
task_cancelled = api.transport.send_command.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
@@ -1,9 +1,10 @@
|
||||
# pyright: reportUnusedFunction=false, reportAny=false
|
||||
"""Tests that InstanceDeleted events close active generation streams."""
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from exo.api.main import API
|
||||
from exo.api.main import API, Transport
|
||||
from exo.api.types import ImageGenerationTaskParams
|
||||
from exo.shared.types.common import CommandId, ModelId
|
||||
from exo.shared.types.state import State
|
||||
@@ -16,12 +17,11 @@ from exo.shared.types.text_generation import (
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def _make_api_with_state(state: State) -> API:
|
||||
def _make_api_with_state(state: State) -> Any:
|
||||
"""Create a minimal API instance with pre-set state."""
|
||||
api = object.__new__(API)
|
||||
api.state = state
|
||||
api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api.transport = object.__new__(Transport)
|
||||
return api
|
||||
|
||||
|
||||
@@ -47,13 +47,10 @@ def test_close_streams_for_deleted_instance() -> None:
|
||||
state = State(tasks={task.task_id: task})
|
||||
api = _make_api_with_state(state)
|
||||
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[command_id] = sender # pyright: ignore[reportPrivateUsage]
|
||||
api._close_streams_for_instance(instance_id)
|
||||
|
||||
api._close_streams_for_instance(instance_id) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
sender.close.assert_called_once()
|
||||
assert command_id not in api._text_generation_queues # pyright: ignore[reportPrivateUsage]
|
||||
api.transport.cancel.assert_called_once()
|
||||
assert api.transport.cancel.call_args[0][0] == command_id
|
||||
|
||||
|
||||
def test_close_streams_ignores_unrelated_instances() -> None:
|
||||
@@ -72,7 +69,6 @@ def test_close_streams_ignores_unrelated_instances() -> None:
|
||||
api._close_streams_for_instance(target_id) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
sender.close.assert_not_called()
|
||||
assert other_cmd in api._text_generation_queues # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
def test_close_streams_for_deleted_instance_image_generation() -> None:
|
||||
|
||||
@@ -49,7 +49,6 @@ class ModelListModel(BaseModel):
|
||||
base_model: str = Field(default="")
|
||||
capabilities: list[str] = Field(default_factory=list)
|
||||
reasoning_dialect: ReasoningDialect = "none"
|
||||
requires_vllm: bool = Field(default=False)
|
||||
|
||||
|
||||
class ModelList(BaseModel):
|
||||
|
||||
@@ -16,7 +16,8 @@ from exo.download.download_utils import (
|
||||
)
|
||||
from exo.download.shard_downloader import ShardDownloader
|
||||
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
|
||||
from exo.shared.models.model_cards import ModelId, get_model_cards
|
||||
from exo.shared.models import model_cards
|
||||
from exo.shared.models.model_cards import ModelId
|
||||
from exo.shared.types.commands import (
|
||||
CancelDownload,
|
||||
DeleteDownload,
|
||||
@@ -422,7 +423,7 @@ class DownloadCoordinator:
|
||||
)
|
||||
# Scan read-only directories for pre-downloaded models
|
||||
if EXO_MODELS_READ_ONLY_DIRS:
|
||||
for card in await get_model_cards():
|
||||
for card in await model_cards.card_cache.list_all():
|
||||
mid = card.model_id
|
||||
if mid in self.active_downloads:
|
||||
continue
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import ssl
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Awaitable
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Callable, Literal
|
||||
@@ -55,6 +56,36 @@ class HuggingFaceAuthenticationError(Exception):
|
||||
class HuggingFaceRateLimitError(Exception):
|
||||
"""429 Huggingface code"""
|
||||
|
||||
def __init__(self, msg: str, retry_after: float | None = None) -> None:
|
||||
super().__init__(msg)
|
||||
self.retry_after = retry_after
|
||||
|
||||
|
||||
def _parse_retry_after(headers: Mapping[str, str]) -> float | None:
|
||||
"""Parse seconds-to-reset from HF's RateLimit header.
|
||||
|
||||
HF sends e.g. ``ratelimit: "api";r=0;t=52`` on 429s; ``t`` is the wait.
|
||||
Returns ``None`` if the header is missing or has no ``t`` field.
|
||||
"""
|
||||
raw = headers.get("RateLimit") or headers.get("ratelimit")
|
||||
if raw is None:
|
||||
return None
|
||||
for part in raw.split(";"):
|
||||
key, _, val = part.strip().partition("=")
|
||||
if key == "t":
|
||||
try:
|
||||
return float(val)
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
# reset window is 5 min
|
||||
_RATE_LIMIT_MAX_SLEEP_SECS = 300.0
|
||||
|
||||
# 24h. Manually clear the cache (or `delete_model`) to force a refresh.
|
||||
_FILE_LIST_CACHE_TTL_SECS = 24 * 60 * 60
|
||||
|
||||
|
||||
async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
|
||||
token = await get_hf_token()
|
||||
@@ -348,9 +379,6 @@ async def _build_file_list_from_local_directory(
|
||||
return None
|
||||
|
||||
|
||||
_fetched_file_lists_this_session: set[str] = set()
|
||||
|
||||
|
||||
async def fetch_file_list_with_cache(
|
||||
model_id: ModelId,
|
||||
revision: str = "main",
|
||||
@@ -360,13 +388,16 @@ async def fetch_file_list_with_cache(
|
||||
) -> list[FileListEntry]:
|
||||
target_dir = await ensure_cache_dir(model_id)
|
||||
cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
|
||||
cache_key = f"{model_id.normalize()}--{revision}"
|
||||
|
||||
if cache_key in _fetched_file_lists_this_session and await aios.path.exists(
|
||||
cache_file
|
||||
):
|
||||
async with aiofiles.open(cache_file, "r") as f:
|
||||
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
|
||||
# cache survives process restarts so cold starts don't re-burst HF
|
||||
if await aios.path.exists(cache_file):
|
||||
try:
|
||||
cache_age = time.time() - (await aios.stat(cache_file)).st_mtime
|
||||
except OSError:
|
||||
cache_age = float("inf")
|
||||
if cache_age < _FILE_LIST_CACHE_TTL_SECS:
|
||||
async with aiofiles.open(cache_file, "r") as f:
|
||||
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
|
||||
|
||||
if skip_internet:
|
||||
if await aios.path.exists(cache_file):
|
||||
@@ -395,7 +426,6 @@ async def fetch_file_list_with_cache(
|
||||
await f.write(
|
||||
TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
|
||||
)
|
||||
_fetched_file_lists_this_session.add(cache_key)
|
||||
return file_list
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning(
|
||||
@@ -426,17 +456,29 @@ async def fetch_file_list_with_retry(
|
||||
recursive: bool = False,
|
||||
on_connection_lost: Callable[[], None] = lambda: None,
|
||||
) -> list[FileListEntry]:
|
||||
n_attempts = 3
|
||||
n_attempts = 5
|
||||
for attempt in range(n_attempts):
|
||||
try:
|
||||
return await _fetch_file_list(model_id, revision, path, recursive)
|
||||
except HuggingFaceAuthenticationError:
|
||||
raise
|
||||
except HuggingFaceRateLimitError as e:
|
||||
if attempt == n_attempts - 1:
|
||||
raise
|
||||
sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
|
||||
sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
|
||||
0, 1
|
||||
)
|
||||
logger.warning(
|
||||
f"Rate limited by HuggingFace fetching file list for {model_id}; "
|
||||
f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
|
||||
)
|
||||
await asyncio.sleep(sleep_for)
|
||||
except Exception as e:
|
||||
on_connection_lost()
|
||||
if attempt == n_attempts - 1:
|
||||
raise e
|
||||
await asyncio.sleep(2.0**attempt)
|
||||
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
|
||||
raise Exception(
|
||||
f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
|
||||
)
|
||||
@@ -447,6 +489,9 @@ async def _fetch_file_list(
|
||||
) -> list[FileListEntry]:
|
||||
api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}"
|
||||
url = f"{api_url}/{path}" if path else api_url
|
||||
# ?recursive=true returns the whole subtree in one request
|
||||
if recursive:
|
||||
url = f"{url}?recursive=true"
|
||||
|
||||
headers = await get_download_headers()
|
||||
async with (
|
||||
@@ -458,7 +503,8 @@ async def _fetch_file_list(
|
||||
raise HuggingFaceAuthenticationError(msg)
|
||||
elif response.status == 429:
|
||||
raise HuggingFaceRateLimitError(
|
||||
f"Couldn't download {model_id} because of HuggingFace rate limit."
|
||||
f"HuggingFace rate limit hit fetching file list for {model_id}",
|
||||
retry_after=_parse_retry_after(response.headers),
|
||||
)
|
||||
elif response.status == 200:
|
||||
data_json = await response.text()
|
||||
@@ -468,10 +514,14 @@ async def _fetch_file_list(
|
||||
if item.type == "file":
|
||||
files.append(FileListEntry.model_validate(item))
|
||||
elif item.type == "directory" and recursive:
|
||||
subfiles = await _fetch_file_list(
|
||||
model_id, revision, item.path, recursive
|
||||
)
|
||||
files.extend(subfiles)
|
||||
# already inlined by ?recursive=true
|
||||
continue
|
||||
if recursive and len(data) >= 1000:
|
||||
# HF tree endpoint paginates at 1000; we don't follow cursors
|
||||
logger.warning(
|
||||
f"File list for {model_id} hit the 1000-entry page cap "
|
||||
"and may be truncated; cursor pagination is not implemented"
|
||||
)
|
||||
return files
|
||||
else:
|
||||
raise Exception(f"Failed to fetch file list: {response.status}")
|
||||
@@ -552,6 +602,11 @@ async def file_meta(
|
||||
if r.status in [401, 403]:
|
||||
msg = await _build_auth_error_message(r.status, model_id)
|
||||
raise HuggingFaceAuthenticationError(msg)
|
||||
if r.status == 429:
|
||||
raise HuggingFaceRateLimitError(
|
||||
f"HuggingFace rate limit hit fetching metadata for {model_id}/{path}",
|
||||
retry_after=_parse_retry_after(r.headers),
|
||||
)
|
||||
content_length = int(
|
||||
r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
|
||||
)
|
||||
@@ -571,7 +626,7 @@ async def download_file_with_retry(
|
||||
on_connection_lost: Callable[[], None] = lambda: None,
|
||||
skip_internet: bool = False,
|
||||
) -> Path:
|
||||
n_attempts = 3
|
||||
n_attempts = 5
|
||||
for attempt in range(n_attempts):
|
||||
try:
|
||||
return await _download_file(
|
||||
@@ -583,12 +638,16 @@ async def download_file_with_retry(
|
||||
raise
|
||||
except HuggingFaceRateLimitError as e:
|
||||
if attempt == n_attempts - 1:
|
||||
raise e
|
||||
logger.error(
|
||||
f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
|
||||
raise
|
||||
sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
|
||||
sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
|
||||
0, 1
|
||||
)
|
||||
logger.error(traceback.format_exc())
|
||||
await asyncio.sleep(2.0**attempt)
|
||||
logger.warning(
|
||||
f"Rate limited by HuggingFace downloading {model_id}/{path}; "
|
||||
f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
|
||||
)
|
||||
await asyncio.sleep(sleep_for)
|
||||
except Exception as e:
|
||||
if attempt == n_attempts - 1:
|
||||
on_connection_lost()
|
||||
@@ -597,7 +656,7 @@ async def download_file_with_retry(
|
||||
f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
|
||||
)
|
||||
logger.error(traceback.format_exc())
|
||||
await asyncio.sleep(2.0**attempt)
|
||||
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
|
||||
raise Exception(
|
||||
f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
|
||||
)
|
||||
@@ -665,6 +724,11 @@ async def _download_file(
|
||||
if r.status in [401, 403]:
|
||||
msg = await _build_auth_error_message(r.status, model_id)
|
||||
raise HuggingFaceAuthenticationError(msg)
|
||||
if r.status == 429:
|
||||
raise HuggingFaceRateLimitError(
|
||||
f"HuggingFace rate limit hit downloading {model_id}/{path}",
|
||||
retry_after=_parse_retry_after(r.headers),
|
||||
)
|
||||
assert r.status in [200, 206], (
|
||||
f"Failed to download {path} from {url}: {r.status}"
|
||||
)
|
||||
|
||||
@@ -11,11 +11,11 @@ from exo.download.download_utils import (
|
||||
download_shard,
|
||||
)
|
||||
from exo.download.shard_downloader import ShardDownloader
|
||||
from exo.shared.models import model_cards
|
||||
from exo.shared.models.model_cards import (
|
||||
ModelCard,
|
||||
ModelId,
|
||||
ModelTask,
|
||||
get_model_cards,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.shards import (
|
||||
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
|
||||
tasks = [
|
||||
create_task(download_with_semaphore(model_card))
|
||||
for model_card in await get_model_cards()
|
||||
for model_card in await model_cards.card_cache.list_all()
|
||||
]
|
||||
|
||||
for task in asyncio.as_completed(tasks):
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
"""Tests for offline/air-gapped mode."""
|
||||
|
||||
import os
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
@@ -231,3 +233,64 @@ class TestFetchFileListOffline:
|
||||
raise FileNotFoundError."""
|
||||
with pytest.raises(FileNotFoundError, match="No internet"):
|
||||
await fetch_file_list_with_cache(model_id, "main", skip_internet=True)
|
||||
|
||||
|
||||
class TestFileListCacheTTL:
|
||||
async def test_uses_fresh_cache_without_fetching(
|
||||
self, model_id: ModelId, temp_models_dir: Path
|
||||
) -> None:
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
cache_dir = temp_models_dir / "caches" / model_id.normalize()
|
||||
await aios.makedirs(cache_dir, exist_ok=True)
|
||||
|
||||
cached_list = [
|
||||
FileListEntry(type="file", path="model.safetensors", size=1000),
|
||||
]
|
||||
cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
|
||||
async with aiofiles.open(cache_file, "w") as f:
|
||||
await f.write(
|
||||
TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode()
|
||||
)
|
||||
|
||||
with patch(
|
||||
"exo.download.download_utils.fetch_file_list_with_retry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_fetch:
|
||||
result = await fetch_file_list_with_cache(model_id, "main")
|
||||
|
||||
assert result == cached_list
|
||||
mock_fetch.assert_not_called()
|
||||
|
||||
async def test_refetches_when_cache_older_than_ttl(
|
||||
self, model_id: ModelId, temp_models_dir: Path
|
||||
) -> None:
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from exo.download.download_utils import (
|
||||
_FILE_LIST_CACHE_TTL_SECS, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
|
||||
cache_dir = temp_models_dir / "caches" / model_id.normalize()
|
||||
await aios.makedirs(cache_dir, exist_ok=True)
|
||||
|
||||
stale_list = [FileListEntry(type="file", path="stale.bin", size=1)]
|
||||
cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
|
||||
async with aiofiles.open(cache_file, "w") as f:
|
||||
await f.write(
|
||||
TypeAdapter(list[FileListEntry]).dump_json(stale_list).decode()
|
||||
)
|
||||
|
||||
old_mtime = time.time() - _FILE_LIST_CACHE_TTL_SECS - 60
|
||||
os.utime(cache_file, (old_mtime, old_mtime))
|
||||
|
||||
fresh_list = [FileListEntry(type="file", path="fresh.bin", size=2)]
|
||||
with patch(
|
||||
"exo.download.download_utils.fetch_file_list_with_retry",
|
||||
new_callable=AsyncMock,
|
||||
return_value=fresh_list,
|
||||
) as mock_fetch:
|
||||
result = await fetch_file_list_with_cache(model_id, "main")
|
||||
|
||||
assert result == fresh_list
|
||||
mock_fetch.assert_called_once()
|
||||
@@ -0,0 +1,355 @@
|
||||
"""Tests for HuggingFace 429 rate-limit handling in download_utils."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import aiofiles.os as aios
|
||||
import pytest
|
||||
|
||||
from exo.download.download_utils import (
|
||||
HuggingFaceRateLimitError,
|
||||
_download_file, # pyright: ignore[reportPrivateUsage]
|
||||
_fetch_file_list, # pyright: ignore[reportPrivateUsage]
|
||||
_parse_retry_after, # pyright: ignore[reportPrivateUsage]
|
||||
download_file_with_retry,
|
||||
fetch_file_list_with_retry,
|
||||
file_meta,
|
||||
)
|
||||
from exo.shared.types.common import ModelId
|
||||
|
||||
# captured from a real HF 429 on 2026-04-30 (header is lowercased by Cloudfront)
|
||||
REAL_HF_429_HEADERS_2026_04_30 = {
|
||||
"ratelimit": '"api";r=0;t=52',
|
||||
"ratelimit-policy": '"fixed window";"api";q=500;w=300',
|
||||
}
|
||||
|
||||
|
||||
class TestParseRetryAfter:
|
||||
def test_parses_documented_format(self) -> None:
|
||||
assert _parse_retry_after({"RateLimit": '"api";r=0;t=243'}) == 243.0
|
||||
|
||||
def test_parses_real_hf_response(self) -> None:
|
||||
assert _parse_retry_after(REAL_HF_429_HEADERS_2026_04_30) == 52.0
|
||||
|
||||
def test_parses_resolvers_bucket(self) -> None:
|
||||
assert _parse_retry_after({"ratelimit": '"resolvers";r=0;t=120'}) == 120.0
|
||||
|
||||
def test_parses_pages_bucket(self) -> None:
|
||||
assert _parse_retry_after({"ratelimit": '"pages";r=0;t=10'}) == 10.0
|
||||
|
||||
def test_returns_none_when_header_missing(self) -> None:
|
||||
assert _parse_retry_after({}) is None
|
||||
|
||||
def test_returns_none_when_only_retry_after_present(self) -> None:
|
||||
assert _parse_retry_after({"Retry-After": "60"}) is None
|
||||
|
||||
def test_returns_none_when_format_unrecognised(self) -> None:
|
||||
assert _parse_retry_after({"ratelimit": "garbage"}) is None
|
||||
|
||||
def test_handles_extra_whitespace(self) -> None:
|
||||
assert _parse_retry_after({"ratelimit": '"api"; r=0; t=42'}) == 42.0
|
||||
|
||||
|
||||
class TestFetchFileListRetry:
|
||||
async def test_uses_retry_after_from_error(self) -> None:
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
|
||||
if not sleeps:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=2.0)
|
||||
return []
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
):
|
||||
result = await fetch_file_list_with_retry(ModelId("test/model"))
|
||||
|
||||
assert result == []
|
||||
assert len(sleeps) == 1
|
||||
assert 2.0 <= sleeps[0] < 3.0 # retry_after + jitter[0,1)
|
||||
|
||||
async def test_falls_back_to_exp_backoff_when_no_retry_after(self) -> None:
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
|
||||
if not sleeps:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=None)
|
||||
return []
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
):
|
||||
await fetch_file_list_with_retry(ModelId("test/model"))
|
||||
|
||||
assert len(sleeps) == 1
|
||||
assert 1.0 <= sleeps[0] < 2.0 # 2**0 + jitter[0,1)
|
||||
|
||||
async def test_caps_sleep_at_max_window(self) -> None:
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
|
||||
if not sleeps:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=10_000.0)
|
||||
return []
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
):
|
||||
await fetch_file_list_with_retry(ModelId("test/model"))
|
||||
|
||||
assert len(sleeps) == 1
|
||||
assert 300.0 <= sleeps[0] < 301.0 # cap + jitter[0,1)
|
||||
|
||||
async def test_retries_up_to_five_times(self) -> None:
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=1.0)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
pytest.raises(HuggingFaceRateLimitError),
|
||||
):
|
||||
await fetch_file_list_with_retry(ModelId("test/model"))
|
||||
|
||||
assert len(sleeps) == 4 # 5 attempts -> 4 sleeps before giving up
|
||||
|
||||
|
||||
class TestDownloadFileRetry:
|
||||
@pytest.fixture
|
||||
async def target_dir(self, tmp_path: Path) -> AsyncIterator[Path]:
|
||||
target = tmp_path / "downloads"
|
||||
await aios.makedirs(target, exist_ok=True)
|
||||
yield target
|
||||
|
||||
async def test_uses_retry_after_from_error(self, target_dir: Path) -> None:
|
||||
sleeps: list[float] = []
|
||||
results: list[Path] = [target_dir / "file.bin"]
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_download(*args: object, **kwargs: object) -> Path:
|
||||
if not sleeps:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=5.0)
|
||||
return results[0]
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._download_file",
|
||||
side_effect=fake_download,
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
):
|
||||
result = await download_file_with_retry(
|
||||
ModelId("test/model"), "main", "file.bin", target_dir
|
||||
)
|
||||
|
||||
assert result == results[0]
|
||||
assert len(sleeps) == 1
|
||||
assert 5.0 <= sleeps[0] < 6.0
|
||||
|
||||
async def test_caps_sleep_at_max_window(self, target_dir: Path) -> None:
|
||||
sleeps: list[float] = []
|
||||
results: list[Path] = [target_dir / "file.bin"]
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
async def fake_download(*args: object, **kwargs: object) -> Path:
|
||||
if not sleeps:
|
||||
raise HuggingFaceRateLimitError("rate limited", retry_after=99_999.0)
|
||||
return results[0]
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._download_file",
|
||||
side_effect=fake_download,
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
):
|
||||
await download_file_with_retry(
|
||||
ModelId("test/model"), "main", "file.bin", target_dir
|
||||
)
|
||||
|
||||
assert len(sleeps) == 1
|
||||
assert 300.0 <= sleeps[0] < 301.0
|
||||
|
||||
async def test_retries_up_to_five_times(self, target_dir: Path) -> None:
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils._download_file",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=HuggingFaceRateLimitError("rate limited", retry_after=1.0),
|
||||
),
|
||||
patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
|
||||
pytest.raises(HuggingFaceRateLimitError),
|
||||
):
|
||||
await download_file_with_retry(
|
||||
ModelId("test/model"), "main", "file.bin", target_dir
|
||||
)
|
||||
|
||||
assert len(sleeps) == 4
|
||||
|
||||
|
||||
def _make_mock_session_returning(
|
||||
response_attrs: dict[str, object], method: str = "get"
|
||||
) -> MagicMock:
|
||||
"""Build a MagicMock that mimics ``create_http_session`` returning a
|
||||
response whose ``status`` / ``headers`` are set from ``response_attrs``.
|
||||
|
||||
Mocks the chain ``create_http_session().__aenter__() -> session``, and
|
||||
``session.<method>().__aenter__() -> response``.
|
||||
"""
|
||||
mock_response = MagicMock()
|
||||
for k, v in response_attrs.items():
|
||||
setattr(mock_response, k, v)
|
||||
|
||||
mock_session = MagicMock()
|
||||
method_mock = getattr(mock_session, method) # pyright: ignore[reportAny]
|
||||
method_mock.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=mock_response
|
||||
)
|
||||
method_mock.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=None
|
||||
)
|
||||
|
||||
mock_factory = MagicMock()
|
||||
mock_factory.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=mock_session
|
||||
)
|
||||
mock_factory.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=None
|
||||
)
|
||||
return mock_factory
|
||||
|
||||
|
||||
REAL_HF_429_HEADER_DICT = {"ratelimit": '"api";r=0;t=52'}
|
||||
|
||||
|
||||
class TestRateLimitAtHttpCallSites:
|
||||
"""Verify each HF call site translates an HTTP 429 into a
|
||||
``HuggingFaceRateLimitError`` carrying the parsed ``retry_after``.
|
||||
|
||||
These tests would catch regressions where (a) the 429 branch is
|
||||
deleted, (b) ``_parse_retry_after`` stops being called, or
|
||||
(c) the wrong header object is passed to it.
|
||||
"""
|
||||
|
||||
async def test_fetch_file_list_maps_429_to_rate_limit_error(self) -> None:
|
||||
mock_factory = _make_mock_session_returning(
|
||||
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}
|
||||
)
|
||||
with (
|
||||
patch("exo.download.download_utils.create_http_session", mock_factory),
|
||||
pytest.raises(HuggingFaceRateLimitError) as exc_info,
|
||||
):
|
||||
await _fetch_file_list(ModelId("test/model"), "main")
|
||||
assert exc_info.value.retry_after == 52.0
|
||||
|
||||
async def test_file_meta_maps_429_to_rate_limit_error(self) -> None:
|
||||
mock_factory = _make_mock_session_returning(
|
||||
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}, method="head"
|
||||
)
|
||||
with (
|
||||
patch("exo.download.download_utils.create_http_session", mock_factory),
|
||||
pytest.raises(HuggingFaceRateLimitError) as exc_info,
|
||||
):
|
||||
await file_meta(ModelId("test/model"), "main", "weights.safetensors")
|
||||
assert exc_info.value.retry_after == 52.0
|
||||
|
||||
async def test_file_meta_maps_429_after_307_redirect(self) -> None:
|
||||
"""When the initial HEAD 307s and the redirected HEAD then 429s,
|
||||
the 429 must still surface as ``HuggingFaceRateLimitError``."""
|
||||
# First HEAD -> 307 with a Location header pointing somewhere new.
|
||||
first_response = MagicMock()
|
||||
first_response.status = 307
|
||||
first_response.headers = {"location": "/redirected/url"}
|
||||
|
||||
# Second HEAD (the recursive call) -> 429 with the real-HF header.
|
||||
second_response = MagicMock()
|
||||
second_response.status = 429
|
||||
second_response.headers = REAL_HF_429_HEADER_DICT
|
||||
|
||||
responses = iter([first_response, second_response])
|
||||
|
||||
mock_session = MagicMock()
|
||||
mock_session.head.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
side_effect=lambda: next(responses)
|
||||
)
|
||||
mock_session.head.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=None
|
||||
)
|
||||
|
||||
mock_factory = MagicMock()
|
||||
mock_factory.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=mock_session
|
||||
)
|
||||
mock_factory.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny]
|
||||
return_value=None
|
||||
)
|
||||
|
||||
with (
|
||||
patch("exo.download.download_utils.create_http_session", mock_factory),
|
||||
pytest.raises(HuggingFaceRateLimitError) as exc_info,
|
||||
):
|
||||
await file_meta(ModelId("test/model"), "main", "weights.safetensors")
|
||||
assert exc_info.value.retry_after == 52.0
|
||||
|
||||
async def test_download_file_maps_429_to_rate_limit_error(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
target_dir = tmp_path / "downloads"
|
||||
await aios.makedirs(target_dir, exist_ok=True)
|
||||
# No local file -> _download_file goes straight to file_meta then GET.
|
||||
# We need both calls to succeed enough to reach the GET branch:
|
||||
# - file_meta returns a non-429 (size, etag) so we proceed.
|
||||
# - the GET then 429s.
|
||||
with (
|
||||
patch(
|
||||
"exo.download.download_utils.file_meta",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(100, "abc123"),
|
||||
),
|
||||
patch(
|
||||
"exo.download.download_utils.create_http_session",
|
||||
_make_mock_session_returning(
|
||||
{"status": 429, "headers": REAL_HF_429_HEADER_DICT}
|
||||
),
|
||||
),
|
||||
pytest.raises(HuggingFaceRateLimitError) as exc_info,
|
||||
):
|
||||
await _download_file(
|
||||
ModelId("test/model"), "main", "weights.safetensors", target_dir
|
||||
)
|
||||
assert exc_info.value.retry_after == 52.0
|
||||
+16
-14
@@ -5,8 +5,10 @@ import resource
|
||||
import signal
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Self
|
||||
from uuid import uuid4
|
||||
|
||||
import anyio
|
||||
from exo_net import PySession
|
||||
from loguru import logger
|
||||
from pydantic import PositiveInt
|
||||
|
||||
@@ -16,7 +18,7 @@ from exo.download.coordinator import DownloadCoordinator
|
||||
from exo.download.impl_shard_downloader import exo_shard_downloader
|
||||
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.routing.router import Router
|
||||
from exo.shared.constants import EXO_LOG
|
||||
from exo.shared.election import Election, ElectionResult
|
||||
from exo.shared.logging import logger_cleanup, logger_setup
|
||||
@@ -39,31 +41,31 @@ class Node:
|
||||
api: API | None
|
||||
|
||||
node_id: NodeId
|
||||
session: PySession
|
||||
offline: bool
|
||||
_api_port: int
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
keypair = get_node_id_keypair()
|
||||
node_id = NodeId(keypair.to_node_id())
|
||||
node_id_bytes = uuid4()
|
||||
node_id = NodeId(str(node_id_bytes))
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
router = Router.create(
|
||||
keypair,
|
||||
router, session = Router.create(
|
||||
node_id_bytes.bytes,
|
||||
bootstrap_peers=args.bootstrap_peers,
|
||||
listen_port=args.libp2p_port,
|
||||
)
|
||||
await router.register_topic(topics.GLOBAL_EVENTS)
|
||||
await router.register_topic(topics.LOCAL_EVENTS)
|
||||
await router.register_topic(topics.COMMANDS)
|
||||
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),
|
||||
command_sender=session.net_sender("orchestrator"),
|
||||
)
|
||||
|
||||
logger.info(f"Starting node {node_id}")
|
||||
@@ -85,9 +87,9 @@ class Node:
|
||||
node_id,
|
||||
port=args.api_port,
|
||||
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),
|
||||
session=session,
|
||||
)
|
||||
else:
|
||||
api = None
|
||||
@@ -97,9 +99,9 @@ class Node:
|
||||
node_id,
|
||||
event_receiver=event_router.receiver(),
|
||||
event_sender=event_router.sender(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
api_port=args.api_port,
|
||||
session=session,
|
||||
)
|
||||
else:
|
||||
worker = None
|
||||
@@ -111,8 +113,8 @@ class Node:
|
||||
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),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
command_receiver=session.net_receiver("orchestrator")
|
||||
)
|
||||
|
||||
er_send, er_recv = channel[ElectionResult]()
|
||||
@@ -125,7 +127,6 @@ class Node:
|
||||
election_message_sender=router.sender(topics.ELECTION_MESSAGES),
|
||||
election_message_receiver=router.receiver(topics.ELECTION_MESSAGES),
|
||||
connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
|
||||
command_receiver=router.receiver(topics.COMMANDS),
|
||||
election_result_sender=er_send,
|
||||
)
|
||||
|
||||
@@ -139,6 +140,7 @@ class Node:
|
||||
master,
|
||||
api,
|
||||
node_id,
|
||||
session,
|
||||
args.offline,
|
||||
args.api_port,
|
||||
)
|
||||
@@ -188,7 +190,7 @@ class Node:
|
||||
self.event_router.shutdown()
|
||||
self.event_router = EventRouter(
|
||||
result.session_id,
|
||||
self.router.sender(topics.COMMANDS),
|
||||
self.session.net_sender("orchestrator"),
|
||||
self.router.receiver(topics.GLOBAL_EVENTS),
|
||||
self.router.sender(topics.LOCAL_EVENTS),
|
||||
)
|
||||
@@ -209,10 +211,10 @@ class Node:
|
||||
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),
|
||||
download_command_sender=self.router.sender(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
command_receiver=self.session.net_receiver("orchestrator"),
|
||||
)
|
||||
self._tg.start_soon(self.master.run)
|
||||
elif (
|
||||
@@ -248,11 +250,11 @@ class Node:
|
||||
self.node_id,
|
||||
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
|
||||
),
|
||||
api_port=self._api_port,
|
||||
session=self.session,
|
||||
)
|
||||
self._tg.start_soon(self.worker.run)
|
||||
if self.api:
|
||||
|
||||
+283
-302
@@ -1,7 +1,9 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import anyio
|
||||
from exo_net import NetReceiver
|
||||
from loguru import logger
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from exo.master.placement import (
|
||||
add_instance_to_placements,
|
||||
@@ -15,11 +17,11 @@ from exo.shared.apply import apply
|
||||
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
|
||||
from exo.shared.types.commands import (
|
||||
AddCustomModelCard,
|
||||
Command,
|
||||
CreateInstance,
|
||||
DeleteCustomModelCard,
|
||||
DeleteInstance,
|
||||
DeleteInstanceLink,
|
||||
ForwarderCommand,
|
||||
ForwarderDownloadCommand,
|
||||
ImageEdits,
|
||||
ImageGeneration,
|
||||
@@ -121,7 +123,7 @@ class Master:
|
||||
node_id: NodeId,
|
||||
session_id: SessionId,
|
||||
*,
|
||||
command_receiver: Receiver[ForwarderCommand],
|
||||
command_receiver: NetReceiver, # todo: not this type
|
||||
event_sender: Sender[Event],
|
||||
local_event_receiver: Receiver[LocalForwarderEvent],
|
||||
global_event_sender: Sender[GlobalForwarderEvent],
|
||||
@@ -155,323 +157,302 @@ class Master:
|
||||
self._event_log.close()
|
||||
self.global_event_sender.close()
|
||||
self.local_event_receiver.close()
|
||||
self.command_receiver.close()
|
||||
|
||||
async def shutdown(self):
|
||||
logger.info("Stopping Master")
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _command_processor(self) -> None:
|
||||
with self.command_receiver as commands:
|
||||
async for forwarder_command in commands:
|
||||
try:
|
||||
logger.info(f"Executing command: {forwarder_command.command}")
|
||||
while True:
|
||||
data = await self.command_receiver.recv()
|
||||
if not data:
|
||||
break
|
||||
try:
|
||||
command = TypeAdapter[Command](Command).validate_json(data)
|
||||
logger.info(f"Executing command: {command}")
|
||||
|
||||
generated_events: list[Event] = []
|
||||
command = forwarder_command.command
|
||||
instance_task_counts: dict[InstanceId, int] = {}
|
||||
match command:
|
||||
case TestCommand():
|
||||
pass
|
||||
case TextGeneration():
|
||||
prefill_only: set[InstanceId] = set()
|
||||
for link in self.state.instance_links.values():
|
||||
prefill_only.update(link.prefill_instances)
|
||||
for link in self.state.instance_links.values():
|
||||
prefill_only.difference_update(link.decode_instances)
|
||||
generated_events: list[Event] = []
|
||||
instance_task_counts: dict[InstanceId, int] = {}
|
||||
match command:
|
||||
case TestCommand():
|
||||
pass
|
||||
case TextGeneration():
|
||||
prefill_only: set[InstanceId] = set()
|
||||
for link in self.state.instance_links.values():
|
||||
prefill_only.update(link.prefill_instances)
|
||||
for link in self.state.instance_links.values():
|
||||
prefill_only.difference_update(link.decode_instances)
|
||||
|
||||
# If the user typed a prefill-only model id (e.g.
|
||||
# the vLLM-side producer of a P/D pair), the
|
||||
# candidate decode side is whatever it's linked
|
||||
# to. Expand the requested model id to also
|
||||
# include those linked decode instances.
|
||||
requested_model = command.task_params.model
|
||||
linked_decode_ids: set[InstanceId] = set()
|
||||
for link in self.state.instance_links.values():
|
||||
if any(
|
||||
self.state.instances.get(pid) is not None
|
||||
and self.state.instances[
|
||||
pid
|
||||
].shard_assignments.model_id
|
||||
== requested_model
|
||||
for pid in link.prefill_instances
|
||||
):
|
||||
linked_decode_ids.update(link.decode_instances)
|
||||
|
||||
for instance in self.state.instances.values():
|
||||
model_match = (
|
||||
instance.shard_assignments.model_id
|
||||
== requested_model
|
||||
) or (instance.instance_id in linked_decode_ids)
|
||||
if (
|
||||
model_match
|
||||
and instance.instance_id not in prefill_only
|
||||
):
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
decode_instance_id = available_instance_ids[0]
|
||||
task_id = TaskId()
|
||||
params = command.task_params.model_copy(
|
||||
update={
|
||||
"prefill_endpoint": _prefill_endpoint_for(
|
||||
self.state, decode_instance_id
|
||||
),
|
||||
}
|
||||
)
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=TextGenerationTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=decode_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=params,
|
||||
),
|
||||
)
|
||||
)
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
case ImageGeneration():
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
instance.shard_assignments.model_id
|
||||
== command.task_params.model
|
||||
):
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
task_id = TaskId()
|
||||
selected_instance_id = available_instance_ids[0]
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=ImageGenerationTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=selected_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case ImageEdits():
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
instance.shard_assignments.model_id
|
||||
== command.task_params.model
|
||||
):
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
task_id = TaskId()
|
||||
selected_instance_id = available_instance_ids[0]
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=ImageEditsTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=selected_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case DeleteInstance():
|
||||
placement = delete_instance(command, self.state.instances)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
for cmd in cancel_unnecessary_downloads(
|
||||
placement, self.state.downloads
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
instance.shard_assignments.model_id
|
||||
== command.task_params.model
|
||||
and instance.instance_id not in prefill_only
|
||||
):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=self._system_id, command=cmd
|
||||
)
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case PlaceInstance():
|
||||
placement = place_instance(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
self.state.node_memory,
|
||||
self.state.node_network,
|
||||
download_status=self.state.downloads,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case CreateInstance():
|
||||
placement = add_instance_to_placements(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case SendInputChunk(chunk=chunk):
|
||||
generated_events.append(
|
||||
InputChunkReceived(
|
||||
command_id=chunk.command_id,
|
||||
chunk=chunk,
|
||||
)
|
||||
)
|
||||
case TaskCancelled():
|
||||
if (
|
||||
task_id := self.command_task_mapping.get(
|
||||
command.cancelled_command_id
|
||||
)
|
||||
) is not None:
|
||||
generated_events.append(
|
||||
TaskStatusUpdated(
|
||||
task_status=TaskStatus.Cancelled,
|
||||
task_id=task_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Nonexistent command {command.cancelled_command_id} cancelled"
|
||||
)
|
||||
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"
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
case AddCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardAdded(model_card=command.model_card)
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
case DeleteCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardDeleted(model_id=command.model_id)
|
||||
)
|
||||
case SetInstanceLink():
|
||||
link = InstanceLink(
|
||||
link_id=command.link_id,
|
||||
prefill_instances=list(
|
||||
dict.fromkeys(command.prefill_instances)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
decode_instance_id = available_instance_ids[0]
|
||||
task_id = TaskId()
|
||||
params = command.task_params.model_copy(
|
||||
update={
|
||||
"prefill_endpoint": _prefill_endpoint_for(
|
||||
self.state, decode_instance_id
|
||||
),
|
||||
decode_instances=list(
|
||||
dict.fromkeys(command.decode_instances)
|
||||
}
|
||||
)
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=TextGenerationTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=decode_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=params,
|
||||
),
|
||||
)
|
||||
generated_events.append(InstanceLinkCreated(link=link))
|
||||
case DeleteInstanceLink():
|
||||
generated_events.append(
|
||||
InstanceLinkDeleted(link_id=command.link_id)
|
||||
)
|
||||
case RequestEventLog():
|
||||
# We should just be able to send everything, since other buffers will ignore old messages
|
||||
# rate limit to 1000 at a time
|
||||
end = min(command.since_idx + 1000, len(self._event_log))
|
||||
for i, event in enumerate(
|
||||
self._event_log.read_range(command.since_idx, end),
|
||||
start=command.since_idx,
|
||||
)
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
case ImageGeneration():
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
instance.shard_assignments.model_id
|
||||
== command.task_params.model
|
||||
):
|
||||
await self._send_event(IndexedEvent(idx=i, event=event))
|
||||
for event in generated_events:
|
||||
await self.event_sender.send(event)
|
||||
except ValueError as e:
|
||||
logger.opt(exception=e).warning("Error in command processor")
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
task_id = TaskId()
|
||||
selected_instance_id = available_instance_ids[0]
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=ImageGenerationTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=selected_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case ImageEdits():
|
||||
for instance in self.state.instances.values():
|
||||
if (
|
||||
instance.shard_assignments.model_id
|
||||
== command.task_params.model
|
||||
):
|
||||
in_flight = {TaskStatus.Pending, TaskStatus.Running}
|
||||
task_count = sum(
|
||||
1
|
||||
for task in self.state.tasks.values()
|
||||
if task.instance_id == instance.instance_id
|
||||
and task.task_status in in_flight
|
||||
)
|
||||
instance_task_counts[instance.instance_id] = (
|
||||
task_count
|
||||
)
|
||||
|
||||
if not instance_task_counts:
|
||||
raise ValueError(
|
||||
f"No instance found for model {command.task_params.model}"
|
||||
)
|
||||
|
||||
available_instance_ids = sorted(
|
||||
instance_task_counts.keys(),
|
||||
key=lambda instance_id: instance_task_counts[
|
||||
instance_id
|
||||
],
|
||||
)
|
||||
|
||||
task_id = TaskId()
|
||||
selected_instance_id = available_instance_ids[0]
|
||||
generated_events.append(
|
||||
TaskCreated(
|
||||
task_id=task_id,
|
||||
task=ImageEditsTask(
|
||||
task_id=task_id,
|
||||
command_id=command.command_id,
|
||||
instance_id=selected_instance_id,
|
||||
task_status=TaskStatus.Pending,
|
||||
task_params=command.task_params,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
self.command_task_mapping[command.command_id] = task_id
|
||||
|
||||
if EXO_TRACING_ENABLED:
|
||||
selected_instance = self.state.instances.get(
|
||||
selected_instance_id
|
||||
)
|
||||
if selected_instance:
|
||||
ranks = set(
|
||||
shard.device_rank
|
||||
for shard in selected_instance.shard_assignments.runner_to_shard.values()
|
||||
)
|
||||
self._expected_ranks[task_id] = ranks
|
||||
case DeleteInstance():
|
||||
placement = delete_instance(command, self.state.instances)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
for cmd in cancel_unnecessary_downloads(
|
||||
placement, self.state.downloads
|
||||
):
|
||||
await self.download_command_sender.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=self._system_id, command=cmd
|
||||
)
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case PlaceInstance():
|
||||
placement = place_instance(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
self.state.node_memory,
|
||||
self.state.node_network,
|
||||
download_status=self.state.downloads,
|
||||
node_rdma_ctl=self.state.node_rdma_ctl,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case CreateInstance():
|
||||
placement = add_instance_to_placements(
|
||||
command,
|
||||
self.state.topology,
|
||||
self.state.instances,
|
||||
)
|
||||
transition_events = get_transition_events(
|
||||
self.state.instances, placement, self.state.tasks
|
||||
)
|
||||
generated_events.extend(transition_events)
|
||||
case SendInputChunk(chunk=chunk):
|
||||
generated_events.append(
|
||||
InputChunkReceived(
|
||||
command_id=chunk.command_id,
|
||||
chunk=chunk,
|
||||
)
|
||||
)
|
||||
case TaskCancelled():
|
||||
if (
|
||||
task_id := self.command_task_mapping.get(
|
||||
command.cancelled_command_id
|
||||
)
|
||||
) is not None:
|
||||
generated_events.append(
|
||||
TaskStatusUpdated(
|
||||
task_status=TaskStatus.Cancelled,
|
||||
task_id=task_id,
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Nonexistent command {command.cancelled_command_id} cancelled"
|
||||
)
|
||||
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 AddCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardAdded(model_card=command.model_card)
|
||||
)
|
||||
case DeleteCustomModelCard():
|
||||
generated_events.append(
|
||||
CustomModelCardDeleted(model_id=command.model_id)
|
||||
)
|
||||
case SetInstanceLink():
|
||||
link = InstanceLink(
|
||||
link_id=command.link_id,
|
||||
prefill_instances=list(
|
||||
dict.fromkeys(command.prefill_instances)
|
||||
),
|
||||
decode_instances=list(
|
||||
dict.fromkeys(command.decode_instances)
|
||||
),
|
||||
)
|
||||
generated_events.append(InstanceLinkCreated(link=link))
|
||||
case DeleteInstanceLink():
|
||||
generated_events.append(
|
||||
InstanceLinkDeleted(link_id=command.link_id)
|
||||
)
|
||||
case RequestEventLog():
|
||||
end = len(self._event_log)
|
||||
for i, event in enumerate(
|
||||
self._event_log.read_range(command.since_idx, end),
|
||||
start=command.since_idx,
|
||||
):
|
||||
await self._send_event(IndexedEvent(idx=i, event=event))
|
||||
for event in generated_events:
|
||||
await self.event_sender.send(event)
|
||||
except ValueError as e:
|
||||
logger.opt(exception=e).warning("Error in command processor")
|
||||
|
||||
# These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
|
||||
async def _plan(self) -> None:
|
||||
|
||||
@@ -28,7 +28,7 @@ from exo.shared.types.events import (
|
||||
TaskStatusUpdated,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
|
||||
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
|
||||
from exo.shared.types.tasks import Task, TaskId, TaskStatus
|
||||
from exo.shared.types.worker.downloads import (
|
||||
DownloadCompleted,
|
||||
@@ -43,7 +43,6 @@ from exo.shared.types.worker.instances import (
|
||||
InstanceMeta,
|
||||
MlxJacclInstance,
|
||||
MlxRingInstance,
|
||||
VllmInstance,
|
||||
)
|
||||
from exo.shared.types.worker.shards import Sharding
|
||||
from exo.utils.ports import random_ephemeral_port
|
||||
@@ -106,6 +105,7 @@ def place_instance(
|
||||
node_network: Mapping[NodeId, NodeNetworkInfo],
|
||||
required_nodes: set[NodeId] | None = None,
|
||||
download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
|
||||
node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
|
||||
) -> dict[InstanceId, Instance]:
|
||||
cycles = topology.get_cycles()
|
||||
candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
|
||||
@@ -167,8 +167,18 @@ def place_instance(
|
||||
|
||||
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
|
||||
|
||||
rdma_ctl_status = node_rdma_ctl or {}
|
||||
|
||||
def _all_rdma_ctl_enabled(cycle: Cycle) -> bool:
|
||||
return all(
|
||||
((status := rdma_ctl_status.get(node_id)) is not None and status.enabled)
|
||||
for node_id in cycle
|
||||
)
|
||||
|
||||
smallest_rdma_cycles = [
|
||||
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
|
||||
cycle
|
||||
for cycle in smallest_cycles
|
||||
if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
|
||||
]
|
||||
|
||||
if command.instance_meta == InstanceMeta.MlxJaccl:
|
||||
@@ -203,7 +213,7 @@ def place_instance(
|
||||
)
|
||||
|
||||
# Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
|
||||
if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
|
||||
if len(selected_cycle) == 1:
|
||||
command = command.model_copy(
|
||||
update={
|
||||
"instance_meta": InstanceMeta.MlxRing,
|
||||
@@ -267,11 +277,6 @@ def place_instance(
|
||||
hosts_by_node=hosts_by_node,
|
||||
ephemeral_port=ephemeral_port,
|
||||
)
|
||||
case InstanceMeta.Vllm:
|
||||
target_instances[instance_id] = VllmInstance(
|
||||
instance_id=instance_id,
|
||||
shard_assignments=shard_assignments,
|
||||
)
|
||||
|
||||
return target_instances
|
||||
|
||||
|
||||
@@ -375,13 +375,7 @@ def find_ip_prioritised(
|
||||
"maybe_ethernet": 3,
|
||||
"thunderbolt": 4,
|
||||
}
|
||||
|
||||
def _key(ip: str) -> tuple[int, int]:
|
||||
link_local = 0 if ip.startswith("169.254.") else 1
|
||||
type_pri = priority.get(ip_to_type.get(ip, "unknown"), 2)
|
||||
return (link_local, type_pri)
|
||||
|
||||
return min(ips, key=_key)
|
||||
return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
|
||||
|
||||
|
||||
def get_mlx_ring_hosts_by_node(
|
||||
|
||||
@@ -6,7 +6,6 @@ import pytest
|
||||
from loguru import logger
|
||||
|
||||
from exo.master.main import Master
|
||||
from exo.routing.router import get_node_id_keypair
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.commands import (
|
||||
CommandId,
|
||||
@@ -47,8 +46,7 @@ from exo.utils.channels import channel
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_master():
|
||||
keypair = get_node_id_keypair()
|
||||
node_id = NodeId(keypair.to_node_id())
|
||||
node_id = NodeId("master test")
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
|
||||
ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
|
||||
|
||||
@@ -21,7 +21,11 @@ from exo.shared.types.events import (
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.multiaddr import Multiaddr
|
||||
from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
|
||||
from exo.shared.types.profiling import (
|
||||
NetworkInterfaceInfo,
|
||||
NodeNetworkInfo,
|
||||
NodeRdmaCtlStatus,
|
||||
)
|
||||
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
@@ -439,8 +443,21 @@ def test_tensor_rdma_backend_connectivity_matrix(
|
||||
min_nodes=1,
|
||||
)
|
||||
|
||||
node_rdma_ctl = {
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
node_c: NodeRdmaCtlStatus(enabled=True),
|
||||
}
|
||||
|
||||
# act
|
||||
placements = place_instance(cic, topology, {}, node_memory, node_network)
|
||||
placements = place_instance(
|
||||
cic,
|
||||
topology,
|
||||
{},
|
||||
node_memory,
|
||||
node_network,
|
||||
node_rdma_ctl=node_rdma_ctl,
|
||||
)
|
||||
|
||||
# assert
|
||||
assert len(placements) == 1
|
||||
@@ -482,6 +499,131 @@ def test_tensor_rdma_backend_connectivity_matrix(
|
||||
assert len(ip_part.split(".")) == 4
|
||||
|
||||
|
||||
def _build_three_node_rdma_topology() -> tuple[
|
||||
Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
|
||||
]:
|
||||
topology = Topology()
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
node_c = NodeId()
|
||||
|
||||
ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
|
||||
ethernet_conn = SocketConnection(
|
||||
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
|
||||
)
|
||||
node_network = {
|
||||
node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
|
||||
node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
|
||||
node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
|
||||
}
|
||||
|
||||
for n in (node_a, node_b, node_c):
|
||||
topology.add_node(n)
|
||||
|
||||
rdma_pairs = [
|
||||
(node_a, node_b, 3),
|
||||
(node_b, node_a, 3),
|
||||
(node_b, node_c, 4),
|
||||
(node_c, node_b, 4),
|
||||
(node_a, node_c, 5),
|
||||
(node_c, node_a, 5),
|
||||
]
|
||||
for src, sink, iface in rdma_pairs:
|
||||
topology.add_connection(
|
||||
Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
|
||||
)
|
||||
|
||||
socket_pairs = [
|
||||
(node_a, node_b),
|
||||
(node_b, node_c),
|
||||
(node_c, node_a),
|
||||
(node_a, node_c),
|
||||
(node_b, node_a),
|
||||
(node_c, node_b),
|
||||
]
|
||||
for src, sink in socket_pairs:
|
||||
topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
|
||||
|
||||
return topology, node_a, node_b, node_c, node_network
|
||||
|
||||
|
||||
def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
|
||||
model_card: ModelCard,
|
||||
):
|
||||
# arrange
|
||||
model_card = model_card.model_copy(
|
||||
update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
|
||||
)
|
||||
topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
|
||||
node_memory = {
|
||||
node_a: create_node_memory(500),
|
||||
node_b: create_node_memory(500),
|
||||
node_c: create_node_memory(500),
|
||||
}
|
||||
node_rdma_ctl = {
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
node_c: NodeRdmaCtlStatus(enabled=False),
|
||||
}
|
||||
cic = PlaceInstance(
|
||||
sharding=Sharding.Tensor,
|
||||
instance_meta=InstanceMeta.MlxJaccl,
|
||||
command_id=CommandId(),
|
||||
model_card=model_card,
|
||||
min_nodes=3,
|
||||
)
|
||||
|
||||
# act / assert
|
||||
with pytest.raises(
|
||||
ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
|
||||
):
|
||||
place_instance(
|
||||
cic,
|
||||
topology,
|
||||
{},
|
||||
node_memory,
|
||||
node_network,
|
||||
node_rdma_ctl=node_rdma_ctl,
|
||||
)
|
||||
|
||||
|
||||
def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
|
||||
"""A node with no observed rdma_ctl status must not participate in RDMA placement."""
|
||||
# arrange
|
||||
model_card = model_card.model_copy(
|
||||
update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
|
||||
)
|
||||
topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
|
||||
node_memory = {
|
||||
node_a: create_node_memory(500),
|
||||
node_b: create_node_memory(500),
|
||||
node_c: create_node_memory(500),
|
||||
}
|
||||
# node_c has no rdma_ctl entry at all
|
||||
node_rdma_ctl = {
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
}
|
||||
cic = PlaceInstance(
|
||||
sharding=Sharding.Tensor,
|
||||
instance_meta=InstanceMeta.MlxJaccl,
|
||||
command_id=CommandId(),
|
||||
model_card=model_card,
|
||||
min_nodes=3,
|
||||
)
|
||||
|
||||
# act / assert
|
||||
with pytest.raises(ValueError):
|
||||
place_instance(
|
||||
cic,
|
||||
topology,
|
||||
{},
|
||||
node_memory,
|
||||
node_network,
|
||||
node_rdma_ctl=node_rdma_ctl,
|
||||
)
|
||||
|
||||
|
||||
def _make_task(
|
||||
instance_id: InstanceId,
|
||||
status: TaskStatus = TaskStatus.Running,
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
from exo_pyo3_bindings import PyFromSwarm
|
||||
from exo_net import PyFromSwarm
|
||||
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
|
||||
"""Serialisable types for Connection Updates/Messages"""
|
||||
|
||||
|
||||
class ConnectionMessage(FrozenModel):
|
||||
node_id: NodeId
|
||||
connected: bool
|
||||
|
||||
@classmethod
|
||||
def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
|
||||
return cls(node_id=NodeId(update.peer_id), connected=update.connected)
|
||||
return cls(connected=update.connected)
|
||||
@@ -4,9 +4,10 @@ from random import random
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, ClosedResourceError
|
||||
from anyio.abc import CancelScope
|
||||
from exo_net import NetSender
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.types.commands import ForwarderCommand, RequestEventLog
|
||||
from exo.shared.types.commands import RequestEventLog
|
||||
from exo.shared.types.common import SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
@@ -23,7 +24,7 @@ from exo.utils.task_group import TaskGroup
|
||||
@dataclass
|
||||
class EventRouter:
|
||||
session_id: SessionId
|
||||
command_sender: Sender[ForwarderCommand]
|
||||
command_sender: NetSender
|
||||
external_inbound: Receiver[GlobalForwarderEvent]
|
||||
external_outbound: Sender[LocalForwarderEvent]
|
||||
_system_id: SystemId = field(init=False, default_factory=SystemId)
|
||||
@@ -152,10 +153,9 @@ class EventRouter:
|
||||
f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
|
||||
)
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(
|
||||
origin=self._system_id,
|
||||
command=RequestEventLog(since_idx=since_idx),
|
||||
)
|
||||
RequestEventLog(since_idx=since_idx)
|
||||
.model_dump_json()
|
||||
.encode("utf-8")
|
||||
)
|
||||
finally:
|
||||
if self._nack_cancel_scope is scope:
|
||||
|
||||
+12
-66
@@ -2,8 +2,6 @@ from collections.abc import Sequence
|
||||
from copy import copy
|
||||
from itertools import count
|
||||
from math import inf
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
from anyio import (
|
||||
@@ -12,18 +10,9 @@ from anyio import (
|
||||
move_on_after,
|
||||
sleep_forever,
|
||||
)
|
||||
from exo_pyo3_bindings import (
|
||||
AllQueuesFullError,
|
||||
Keypair,
|
||||
MessageTooLargeError,
|
||||
NetworkingHandle,
|
||||
NoPeersSubscribedToTopicError,
|
||||
PyFromSwarm,
|
||||
)
|
||||
from filelock import FileLock
|
||||
from exo_net import NetworkingHandle, PyFromSwarm, PySession
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.constants import EXO_NODE_ID_KEYPAIR
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
from exo.utils.task_group import TaskGroup
|
||||
@@ -105,13 +94,14 @@ class Router:
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
identity: Keypair,
|
||||
identity: bytes,
|
||||
bootstrap_peers: Sequence[str] = (),
|
||||
listen_port: int = 0,
|
||||
) -> "Router":
|
||||
return cls(
|
||||
handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
|
||||
) -> "tuple[Router, PySession]":
|
||||
handle, session = NetworkingHandle.new(
|
||||
identity, list(bootstrap_peers), listen_port
|
||||
)
|
||||
return cls(handle=handle), session
|
||||
|
||||
def __init__(self, handle: NetworkingHandle):
|
||||
self.topic_routers: dict[str, TopicRouter[FrozenModel]] = {}
|
||||
@@ -191,10 +181,8 @@ class Router:
|
||||
from_swarm = await self._net.recv()
|
||||
logger.debug(from_swarm)
|
||||
match from_swarm:
|
||||
case PyFromSwarm.Message(origin, topic, data):
|
||||
logger.trace(
|
||||
f"Received message on {topic} from {origin} with payload {data}"
|
||||
)
|
||||
case PyFromSwarm.Message(topic, data):
|
||||
logger.trace(f"Received message on {topic} with payload {data}")
|
||||
if topic not in self.topic_routers:
|
||||
logger.warning(
|
||||
f"Received message on unknown or inactive topic {topic}"
|
||||
@@ -225,51 +213,9 @@ class Router:
|
||||
async def _networking_publish(self):
|
||||
with self.networking_receiver as networked_items:
|
||||
async for topic, data in networked_items:
|
||||
try:
|
||||
logger.trace(f"Sending message on {topic} with payload {data}")
|
||||
if len(data) > 1024 * 1024:
|
||||
logger.warning(
|
||||
"Sending overlarge payload, network performance may be temporarily degraded"
|
||||
)
|
||||
await self._net.gossipsub_publish(topic, data)
|
||||
except NoPeersSubscribedToTopicError:
|
||||
pass
|
||||
except AllQueuesFullError:
|
||||
logger.warning(f"All peer queues full, dropping message on {topic}")
|
||||
except MessageTooLargeError:
|
||||
logger.trace(f"Sending message on {topic} with payload {data}")
|
||||
if len(data) > 1024 * 1024:
|
||||
logger.warning(
|
||||
f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
|
||||
"Sending overlarge payload, network performance may be temporarily degraded"
|
||||
)
|
||||
|
||||
|
||||
def get_node_id_keypair(
|
||||
path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR,
|
||||
) -> Keypair:
|
||||
"""
|
||||
Obtains the :class:`Keypair` associated with this node-ID.
|
||||
Obtain the :class:`PeerId` by from it.
|
||||
"""
|
||||
# TODO(evan): bring back node id persistence once we figure out how to deal with duplicates
|
||||
return Keypair.generate()
|
||||
|
||||
def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
|
||||
return Path(str(path) + ".lock")
|
||||
|
||||
# operate with cross-process lock to avoid race conditions
|
||||
with FileLock(lock_path(path)):
|
||||
with open(path, "a+b") as f: # opens in append-mode => starts at EOF
|
||||
# if non-zero EOF, then file exists => use to get node-ID
|
||||
if f.tell() != 0:
|
||||
f.seek(0) # go to start & read protobuf-encoded bytes
|
||||
protobuf_encoded = f.read()
|
||||
|
||||
try: # if decoded successfully, save & return
|
||||
return Keypair.from_bytes(protobuf_encoded)
|
||||
except ValueError as e: # on runtime error, assume corrupt file
|
||||
logger.warning(f"Encountered error when trying to get keypair: {e}")
|
||||
|
||||
# if no valid credentials, create new ones and persist
|
||||
with open(path, "w+b") as f:
|
||||
keypair = Keypair.generate()
|
||||
f.write(keypair.to_bytes())
|
||||
return keypair
|
||||
await self._net.gossipsub_publish(topic, data)
|
||||
+50
-19
@@ -4,7 +4,8 @@ from datetime import datetime
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.models.model_cards import ModelCard
|
||||
from exo.shared.types.common import ModelId, NodeId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
CustomModelCardAdded,
|
||||
@@ -59,14 +60,24 @@ from exo.utils.info_gatherer.info_gatherer import (
|
||||
NodeConfig,
|
||||
NodeDiskUsage,
|
||||
NodeNetworkInterfaces,
|
||||
NvmlMetrics,
|
||||
RdmaCtlStatus,
|
||||
StaticNodeInformation,
|
||||
ThunderboltBridgeInfo,
|
||||
VllmCapability,
|
||||
)
|
||||
|
||||
|
||||
def _is_rdma_ctl_enabled(
|
||||
node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
|
||||
) -> bool:
|
||||
"""A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
|
||||
|
||||
Missing entries default to ``False`` — if we have not yet observed (or the node
|
||||
cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
|
||||
"""
|
||||
status = node_rdma_ctl.get(node_id)
|
||||
return status is not None and status.enabled
|
||||
|
||||
|
||||
def event_apply(event: Event, state: State) -> State:
|
||||
"""Apply an event to state."""
|
||||
match event:
|
||||
@@ -77,10 +88,12 @@ def event_apply(event: Event, state: State) -> State:
|
||||
| InputChunkReceived()
|
||||
| TracesCollected()
|
||||
| TracesMerged()
|
||||
| CustomModelCardAdded()
|
||||
| CustomModelCardDeleted()
|
||||
): # Pass-through events that don't modify state
|
||||
return state
|
||||
case CustomModelCardAdded():
|
||||
return apply_custom_model_card_added(event, state)
|
||||
case CustomModelCardDeleted():
|
||||
return apply_custom_model_card_deleted(event, state)
|
||||
case InstanceCreated():
|
||||
return apply_instance_created(event, state)
|
||||
case InstanceDeleted():
|
||||
@@ -306,9 +319,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
|
||||
node_rdma_ctl = {
|
||||
key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
|
||||
}
|
||||
node_vllm = {
|
||||
key: value for key, value in state.node_vllm.items() if key != event.node_id
|
||||
}
|
||||
# Only recompute cycles if the leaving node had TB bridge enabled
|
||||
leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
|
||||
leaving_node_had_tb_enabled = (
|
||||
@@ -331,7 +341,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
|
||||
"node_thunderbolt": node_thunderbolt,
|
||||
"node_thunderbolt_bridge": node_thunderbolt_bridge,
|
||||
"node_rdma_ctl": node_rdma_ctl,
|
||||
"node_vllm": node_vllm,
|
||||
"thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
|
||||
}
|
||||
)
|
||||
@@ -358,11 +367,6 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
event.node_id: info.system_profile,
|
||||
}
|
||||
update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
|
||||
case NvmlMetrics():
|
||||
update["node_system"] = {
|
||||
**state.node_system,
|
||||
event.node_id: info.system_profile,
|
||||
}
|
||||
case MemoryUsage():
|
||||
update["node_memory"] = {**state.node_memory, event.node_id: info}
|
||||
case NodeDiskUsage():
|
||||
@@ -408,6 +412,9 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
for nid in state.node_thunderbolt
|
||||
for tb_ident in state.node_thunderbolt[nid].interfaces
|
||||
}
|
||||
source_is_rdma_enabled = _is_rdma_ctl_enabled(
|
||||
event.node_id, state.node_rdma_ctl
|
||||
)
|
||||
as_rdma_conns = [
|
||||
Connection(
|
||||
source=event.node_id,
|
||||
@@ -420,6 +427,10 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
for tb_conn in info.conns
|
||||
if tb_conn.source_uuid in conn_map
|
||||
if tb_conn.sink_uuid in conn_map
|
||||
if source_is_rdma_enabled
|
||||
and _is_rdma_ctl_enabled(
|
||||
conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
|
||||
)
|
||||
]
|
||||
topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
|
||||
case ThunderboltBridgeInfo():
|
||||
@@ -443,11 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
**state.node_rdma_ctl,
|
||||
event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
|
||||
}
|
||||
case VllmCapability():
|
||||
update["node_vllm"] = {
|
||||
**state.node_vllm,
|
||||
event.node_id: info.available,
|
||||
}
|
||||
# If RDMA just got disabled on this node, drop any RDMA edges touching it
|
||||
# so placement / topology consumers cannot pick a disabled node for an
|
||||
# RDMA-backed instance. (Edges will repopulate on the next
|
||||
# MacThunderboltConnections poll once both endpoints are enabled again.)
|
||||
if not info.enabled:
|
||||
topology.remove_all_rdma_connections_touching(event.node_id)
|
||||
|
||||
return state.model_copy(update=update)
|
||||
|
||||
@@ -463,3 +475,22 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
|
||||
topology.remove_connection(event.conn)
|
||||
# TODO: Clean up removing the reverse connection
|
||||
return state.model_copy(update={"topology": topology})
|
||||
|
||||
|
||||
def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
|
||||
new_cards: Mapping[ModelId, ModelCard] = {
|
||||
**state.custom_model_cards,
|
||||
event.model_card.model_id: event.model_card,
|
||||
}
|
||||
return state.model_copy(update={"custom_model_cards": new_cards})
|
||||
|
||||
|
||||
def apply_custom_model_card_deleted(
|
||||
event: CustomModelCardDeleted, state: State
|
||||
) -> State:
|
||||
new_cards: Mapping[ModelId, ModelCard] = {
|
||||
model_id: card
|
||||
for model_id, card in state.custom_model_cards.items()
|
||||
if model_id != event.model_id
|
||||
}
|
||||
return state.model_copy(update={"custom_model_cards": new_cards})
|
||||
@@ -9,7 +9,6 @@ from anyio import (
|
||||
from loguru import logger
|
||||
|
||||
from exo.routing.connection_message import ConnectionMessage
|
||||
from exo.shared.types.commands import ForwarderCommand
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
@@ -22,7 +21,6 @@ class ElectionMessage(FrozenModel):
|
||||
clock: int
|
||||
seniority: int
|
||||
proposed_session: SessionId
|
||||
commands_seen: int
|
||||
|
||||
# Could eventually include a list of neighbour nodes for centrality
|
||||
def __lt__(self, other: Self) -> bool:
|
||||
@@ -30,8 +28,6 @@ class ElectionMessage(FrozenModel):
|
||||
return self.clock < other.clock
|
||||
if self.seniority != other.seniority:
|
||||
return self.seniority < other.seniority
|
||||
elif self.commands_seen != other.commands_seen:
|
||||
return self.commands_seen < other.commands_seen
|
||||
else:
|
||||
return (
|
||||
self.proposed_session.master_node_id
|
||||
@@ -54,7 +50,6 @@ class Election:
|
||||
election_message_sender: Sender[ElectionMessage],
|
||||
election_result_sender: Sender[ElectionResult],
|
||||
connection_message_receiver: Receiver[ConnectionMessage],
|
||||
command_receiver: Receiver[ForwarderCommand],
|
||||
is_candidate: bool = True,
|
||||
seniority: int = 0,
|
||||
):
|
||||
@@ -64,7 +59,6 @@ class Election:
|
||||
self.seniority = seniority if is_candidate else -1
|
||||
self.clock = 0
|
||||
self.node_id = node_id
|
||||
self.commands_seen = 0
|
||||
# Every node spawns as master
|
||||
self.current_session: SessionId = SessionId(
|
||||
master_node_id=node_id, election_clock=0
|
||||
@@ -75,7 +69,6 @@ class Election:
|
||||
self._em_receiver = election_message_receiver
|
||||
self._er_sender = election_result_sender
|
||||
self._cm_receiver = connection_message_receiver
|
||||
self._co_receiver = command_receiver
|
||||
|
||||
# Campaign state
|
||||
self._candidates: list[ElectionMessage] = []
|
||||
@@ -89,7 +82,6 @@ class Election:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._election_receiver)
|
||||
tg.start_soon(self._connection_receiver)
|
||||
tg.start_soon(self._command_counter)
|
||||
|
||||
# And start an election immediately, that instantly resolves
|
||||
candidates: list[ElectionMessage] = []
|
||||
@@ -179,11 +171,6 @@ class Election:
|
||||
logger.debug("Campaign started")
|
||||
logger.debug("Connection message added")
|
||||
|
||||
async def _command_counter(self) -> None:
|
||||
with self._co_receiver as commands:
|
||||
async for _command in commands:
|
||||
self.commands_seen += 1
|
||||
|
||||
async def _campaign(
|
||||
self, candidates: list[ElectionMessage], campaign_timeout: float
|
||||
) -> None:
|
||||
@@ -261,5 +248,4 @@ class Election:
|
||||
),
|
||||
clock=c,
|
||||
seniority=self.seniority,
|
||||
commands_seen=self.commands_seen,
|
||||
)
|
||||
@@ -46,7 +46,8 @@ class _InterceptHandler(logging.Handler):
|
||||
def logger_setup(log_file: Path | None, verbosity: int = 0):
|
||||
"""Set up logging for this process - formatting, file handles, verbosity and output"""
|
||||
|
||||
logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
|
||||
logging.getLogger("exo_net").setLevel(logging.INFO)
|
||||
logging.getLogger("networking").setLevel(logging.INFO)
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
|
||||
Path(RESOURCES_DIR) / "image_model_cards",
|
||||
]
|
||||
|
||||
_card_cache: dict[ModelId, "ModelCard"] = {}
|
||||
|
||||
class _CardCache:
|
||||
def __init__(self):
|
||||
self.cc: dict[ModelId, "ModelCard"] = {}
|
||||
|
||||
def get(self, model_id: ModelId) -> "ModelCard | None":
|
||||
return self.cc.get(model_id)
|
||||
|
||||
async def save(self, card: "ModelCard"):
|
||||
self.cc[card.model_id] = card
|
||||
try:
|
||||
await card.save_to_custom_dir()
|
||||
except OSError as e:
|
||||
logger.warning(f"failed to save custom model card ({e.strerror})")
|
||||
|
||||
async def pop(self, model_id: ModelId) -> "ModelCard | None":
|
||||
"""Delete a user-added custom model card. Returns True if deleted."""
|
||||
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
|
||||
try:
|
||||
if await card_path.exists():
|
||||
await card_path.unlink()
|
||||
return self.cc.pop(model_id, None)
|
||||
except OSError as e:
|
||||
logger.warning(f"failed to delete custom model card ({e.strerror})")
|
||||
|
||||
async def list_all(self) -> list["ModelCard"]:
|
||||
if len(self.cc) == 0:
|
||||
await self.refresh()
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
return list(self.cc.values())
|
||||
return [c for c in self.cc.values() if not _is_image_card(c)]
|
||||
|
||||
async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
|
||||
"""Load all TOML model cards from a directory into the cache."""
|
||||
async for toml_file in directory.rglob("*.toml"):
|
||||
try:
|
||||
card = await ModelCard.load_from_path(toml_file)
|
||||
if is_custom:
|
||||
card = card.model_copy(update={"is_custom": True})
|
||||
if self.get(card.model_id) is None:
|
||||
self.cc[card.model_id] = card
|
||||
except (ValidationError, TOMLKitError):
|
||||
pass
|
||||
|
||||
async def refresh(self) -> None:
|
||||
for path in _BUILTIN_CARD_DIRS:
|
||||
await self._load_cards_from_dir(path, is_custom=False)
|
||||
await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
|
||||
|
||||
|
||||
card_cache = _CardCache()
|
||||
|
||||
|
||||
def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
|
||||
@@ -59,42 +109,10 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
|
||||
return None
|
||||
|
||||
|
||||
async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
|
||||
"""Load all TOML model cards from a directory into the cache."""
|
||||
async for toml_file in directory.rglob("*.toml"):
|
||||
try:
|
||||
card = await ModelCard.load_from_path(toml_file)
|
||||
if is_custom:
|
||||
card = card.model_copy(update={"is_custom": True})
|
||||
if card.model_id not in _card_cache:
|
||||
_card_cache[card.model_id] = card
|
||||
except (ValidationError, TOMLKitError):
|
||||
pass
|
||||
|
||||
|
||||
async def _refresh_card_cache() -> None:
|
||||
for path in _BUILTIN_CARD_DIRS:
|
||||
await _load_cards_from_dir(path, is_custom=False)
|
||||
await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
|
||||
|
||||
|
||||
def _is_image_card(card: "ModelCard") -> bool:
|
||||
return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
|
||||
|
||||
|
||||
def get_card(model_id: ModelId) -> "ModelCard | None":
|
||||
"""Look up a single model card from the cache by ID."""
|
||||
return _card_cache.get(model_id)
|
||||
|
||||
|
||||
async def get_model_cards() -> list["ModelCard"]:
|
||||
if len(_card_cache) == 0:
|
||||
await _refresh_card_cache()
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
return list(_card_cache.values())
|
||||
return [c for c in _card_cache.values() if not _is_image_card(c)]
|
||||
|
||||
|
||||
class ModelTask(str, Enum):
|
||||
TextGeneration = "TextGeneration"
|
||||
TextToImage = "TextToImage"
|
||||
@@ -150,7 +168,6 @@ class ModelCard(FrozenModel):
|
||||
context_length: int = 0
|
||||
uses_cfg: bool = False
|
||||
trust_remote_code: bool = True
|
||||
requires_vllm: bool = False
|
||||
is_custom: bool = False
|
||||
vision: VisionCardConfig | None = None
|
||||
sampling_defaults: SamplingDefaults = Field(default_factory=SamplingDefaults)
|
||||
@@ -197,14 +214,13 @@ class ModelCard(FrozenModel):
|
||||
# Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
|
||||
@staticmethod
|
||||
async def load(model_id: ModelId) -> "ModelCard":
|
||||
if model_id not in _card_cache:
|
||||
await _refresh_card_cache()
|
||||
if (mc := _card_cache.get(model_id)) is not None:
|
||||
if card_cache.get(model_id) is None:
|
||||
await card_cache.refresh()
|
||||
if (mc := card_cache.get(model_id)) is not None:
|
||||
return mc
|
||||
|
||||
mc = await ModelCard.fetch_from_hf(model_id)
|
||||
await mc.save_to_custom_dir()
|
||||
_card_cache[model_id] = mc
|
||||
return mc
|
||||
|
||||
@staticmethod
|
||||
@@ -234,21 +250,6 @@ class ModelCard(FrozenModel):
|
||||
)
|
||||
|
||||
|
||||
def add_to_card_cache(card: "ModelCard") -> None:
|
||||
"""Add or update a model card in the in-memory cache."""
|
||||
_card_cache[card.model_id] = card
|
||||
|
||||
|
||||
async def delete_custom_card(model_id: ModelId) -> bool:
|
||||
"""Delete a user-added custom model card. Returns True if deleted."""
|
||||
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
|
||||
if await card_path.exists():
|
||||
await card_path.unlink()
|
||||
_card_cache.pop(model_id, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class ConfigData(BaseModel):
|
||||
model_config = {"extra": "ignore"} # Allow unknown fields
|
||||
|
||||
@@ -350,11 +351,7 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData:
|
||||
|
||||
|
||||
async def fetch_safetensors_size(model_id: ModelId) -> Memory:
|
||||
"""Gets model size from safetensors index or falls back to HF API.
|
||||
|
||||
Single-shard repos don't have a `model.safetensors.index.json`; fall back
|
||||
to the HF API for those.
|
||||
"""
|
||||
"""Gets model size from safetensors index or falls back to HF API."""
|
||||
from exo.download.download_utils import (
|
||||
download_file_with_retry,
|
||||
resolve_model_dir,
|
||||
@@ -362,25 +359,21 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
|
||||
from exo.shared.types.worker.downloads import ModelSafetensorsIndex
|
||||
|
||||
target_dir = await resolve_model_dir(model_id)
|
||||
try:
|
||||
index_path = await download_file_with_retry(
|
||||
model_id,
|
||||
"main",
|
||||
"model.safetensors.index.json",
|
||||
target_dir,
|
||||
lambda curr_bytes, total_bytes, is_renamed: logger.debug(
|
||||
f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
|
||||
),
|
||||
)
|
||||
except FileNotFoundError:
|
||||
index_path = None
|
||||
index_path = await download_file_with_retry(
|
||||
model_id,
|
||||
"main",
|
||||
"model.safetensors.index.json",
|
||||
target_dir,
|
||||
lambda curr_bytes, total_bytes, is_renamed: logger.debug(
|
||||
f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
|
||||
),
|
||||
)
|
||||
async with aiofiles.open(index_path, "r") as f:
|
||||
index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
|
||||
|
||||
if index_path is not None:
|
||||
async with aiofiles.open(index_path, "r") as f:
|
||||
index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
|
||||
metadata = index_data.metadata
|
||||
if metadata is not None and metadata.total_size is not None:
|
||||
return Memory.from_bytes(metadata.total_size)
|
||||
metadata = index_data.metadata
|
||||
if metadata is not None and metadata.total_size is not None:
|
||||
return Memory.from_bytes(metadata.total_size)
|
||||
|
||||
info = model_info(model_id)
|
||||
if info.safetensors is None:
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.events import (
|
||||
CustomModelCardAdded,
|
||||
CustomModelCardDeleted,
|
||||
IndexedEvent,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.state import State
|
||||
|
||||
|
||||
def _model_card(model_id: ModelId) -> ModelCard:
|
||||
return ModelCard(
|
||||
model_id=model_id,
|
||||
n_layers=1,
|
||||
storage_size=Memory.from_bytes(1),
|
||||
hidden_size=1,
|
||||
supports_tensor=True,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
)
|
||||
|
||||
|
||||
def test_custom_model_card_added_is_reduced_into_state() -> None:
|
||||
card = _model_card(ModelId("custom/model"))
|
||||
|
||||
state = apply(
|
||||
State(),
|
||||
IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
|
||||
)
|
||||
|
||||
assert state.custom_model_cards == {card.model_id: card}
|
||||
|
||||
|
||||
def test_custom_model_card_deleted_removes_card_from_state() -> None:
|
||||
card = _model_card(ModelId("custom/model"))
|
||||
state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
|
||||
|
||||
state = apply(
|
||||
state,
|
||||
IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
|
||||
)
|
||||
|
||||
assert state.custom_model_cards == {}
|
||||
@@ -0,0 +1,231 @@
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from exo.shared.apply import apply_node_gathered_info
|
||||
from exo.shared.topology import Topology
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import NodeGatheredInfo
|
||||
from exo.shared.types.profiling import (
|
||||
NodeRdmaCtlStatus,
|
||||
NodeThunderboltInfo,
|
||||
)
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.thunderbolt import ThunderboltConnection, ThunderboltIdentifier
|
||||
from exo.shared.types.topology import RDMAConnection
|
||||
from exo.utils.info_gatherer.info_gatherer import (
|
||||
MacThunderboltConnections,
|
||||
RdmaCtlStatus,
|
||||
)
|
||||
|
||||
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _make_state_with_thunderbolt_idents(
|
||||
*node_ids_and_uuids: tuple[NodeId, str, str],
|
||||
rdma_ctl: dict[NodeId, NodeRdmaCtlStatus] | None = None,
|
||||
) -> State:
|
||||
"""Build a State with Thunderbolt identifiers per node so the apply MacThunderboltConnections
|
||||
case can resolve uuid -> (node, iface)."""
|
||||
node_thunderbolt = {
|
||||
nid: NodeThunderboltInfo(
|
||||
interfaces=[ThunderboltIdentifier(rdma_interface=iface, domain_uuid=uuid)]
|
||||
)
|
||||
for nid, uuid, iface in node_ids_and_uuids
|
||||
}
|
||||
return State(
|
||||
node_thunderbolt=node_thunderbolt,
|
||||
node_rdma_ctl=rdma_ctl or {},
|
||||
)
|
||||
|
||||
|
||||
def _has_rdma_edge(topology: Topology, source: NodeId, sink: NodeId) -> bool:
|
||||
return any(
|
||||
isinstance(edge, RDMAConnection)
|
||||
for edge in topology.get_all_connections_between(source, sink)
|
||||
)
|
||||
|
||||
|
||||
def test_mac_thunderbolt_connections_emits_rdma_when_both_endpoints_enabled():
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
state = _make_state_with_thunderbolt_idents(
|
||||
(node_a, "uuid-a", "rdma_en1"),
|
||||
(node_b, "uuid-b", "rdma_en1"),
|
||||
rdma_ctl={
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
},
|
||||
)
|
||||
|
||||
event = NodeGatheredInfo(
|
||||
node_id=node_a,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
|
||||
),
|
||||
)
|
||||
|
||||
new_state = apply_node_gathered_info(event, state)
|
||||
|
||||
assert _has_rdma_edge(new_state.topology, node_a, node_b)
|
||||
|
||||
|
||||
def test_mac_thunderbolt_connections_skips_rdma_when_source_rdma_ctl_disabled():
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
state = _make_state_with_thunderbolt_idents(
|
||||
(node_a, "uuid-a", "rdma_en1"),
|
||||
(node_b, "uuid-b", "rdma_en1"),
|
||||
rdma_ctl={
|
||||
node_a: NodeRdmaCtlStatus(enabled=False),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
},
|
||||
)
|
||||
|
||||
event = NodeGatheredInfo(
|
||||
node_id=node_a,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
|
||||
),
|
||||
)
|
||||
|
||||
new_state = apply_node_gathered_info(event, state)
|
||||
|
||||
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
|
||||
|
||||
|
||||
def test_mac_thunderbolt_connections_skips_rdma_when_sink_rdma_ctl_disabled():
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
state = _make_state_with_thunderbolt_idents(
|
||||
(node_a, "uuid-a", "rdma_en1"),
|
||||
(node_b, "uuid-b", "rdma_en1"),
|
||||
rdma_ctl={
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=False),
|
||||
},
|
||||
)
|
||||
|
||||
event = NodeGatheredInfo(
|
||||
node_id=node_a,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
|
||||
),
|
||||
)
|
||||
|
||||
new_state = apply_node_gathered_info(event, state)
|
||||
|
||||
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
|
||||
|
||||
|
||||
def test_mac_thunderbolt_connections_skips_rdma_when_rdma_ctl_status_missing():
|
||||
"""Missing rdma_ctl status defaults to not-enabled — node is RDMA-incapable."""
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
state = _make_state_with_thunderbolt_idents(
|
||||
(node_a, "uuid-a", "rdma_en1"),
|
||||
(node_b, "uuid-b", "rdma_en1"),
|
||||
rdma_ctl={
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
# node_b intentionally absent
|
||||
},
|
||||
)
|
||||
|
||||
event = NodeGatheredInfo(
|
||||
node_id=node_a,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
|
||||
),
|
||||
)
|
||||
|
||||
new_state = apply_node_gathered_info(event, state)
|
||||
|
||||
assert not _has_rdma_edge(new_state.topology, node_a, node_b)
|
||||
|
||||
|
||||
def test_rdma_ctl_status_disabled_purges_existing_rdma_edges():
|
||||
"""When a node reports rdma_ctl disabled, all RDMA edges touching it must be removed."""
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
|
||||
# Start with both nodes RDMA-enabled and existing RDMA edges in the topology.
|
||||
state = _make_state_with_thunderbolt_idents(
|
||||
(node_a, "uuid-a", "rdma_en1"),
|
||||
(node_b, "uuid-b", "rdma_en1"),
|
||||
rdma_ctl={
|
||||
node_a: NodeRdmaCtlStatus(enabled=True),
|
||||
node_b: NodeRdmaCtlStatus(enabled=True),
|
||||
},
|
||||
)
|
||||
state = apply_node_gathered_info(
|
||||
NodeGatheredInfo(
|
||||
node_id=node_a,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
|
||||
),
|
||||
),
|
||||
state,
|
||||
)
|
||||
state = apply_node_gathered_info(
|
||||
NodeGatheredInfo(
|
||||
node_id=node_b,
|
||||
when=_now(),
|
||||
info=MacThunderboltConnections(
|
||||
conns=[ThunderboltConnection(source_uuid="uuid-b", sink_uuid="uuid-a")]
|
||||
),
|
||||
),
|
||||
state,
|
||||
)
|
||||
assert _has_rdma_edge(state.topology, node_a, node_b)
|
||||
assert _has_rdma_edge(state.topology, node_b, node_a)
|
||||
|
||||
# Now node_a flips to rdma_ctl disabled — both directions of RDMA edge must drop.
|
||||
state = apply_node_gathered_info(
|
||||
NodeGatheredInfo(
|
||||
node_id=node_a, when=_now(), info=RdmaCtlStatus(enabled=False)
|
||||
),
|
||||
state,
|
||||
)
|
||||
|
||||
assert not _has_rdma_edge(state.topology, node_a, node_b)
|
||||
assert not _has_rdma_edge(state.topology, node_b, node_a)
|
||||
assert state.node_rdma_ctl[node_a].enabled is False
|
||||
|
||||
|
||||
def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
|
||||
"""Purging RDMA edges for a disabled node must not affect non-RDMA edges."""
|
||||
from exo.shared.types.multiaddr import Multiaddr
|
||||
from exo.shared.types.topology import Connection, SocketConnection
|
||||
|
||||
topology = Topology()
|
||||
node_a = NodeId()
|
||||
node_b = NodeId()
|
||||
topology.add_node(node_a)
|
||||
topology.add_node(node_b)
|
||||
topology.add_connection(
|
||||
Connection(
|
||||
source=node_a,
|
||||
sink=node_b,
|
||||
edge=RDMAConnection(
|
||||
source_rdma_iface="rdma_en1", sink_rdma_iface="rdma_en1"
|
||||
),
|
||||
)
|
||||
)
|
||||
socket_edge = SocketConnection(
|
||||
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
|
||||
)
|
||||
topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
|
||||
|
||||
topology.remove_all_rdma_connections_touching(node_a)
|
||||
|
||||
assert not _has_rdma_edge(topology, node_a, node_b)
|
||||
# Socket edge survives.
|
||||
assert any(
|
||||
isinstance(edge, SocketConnection)
|
||||
for edge in topology.get_all_connections_between(node_a, node_b)
|
||||
)
|
||||
@@ -327,7 +327,7 @@ async def test_connection_message_triggers_new_round_broadcast() -> None:
|
||||
tg.start_soon(election.run)
|
||||
|
||||
# Send any connection message object; we close quickly to cancel before result creation
|
||||
await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True))
|
||||
await cm_tx.send(ConnectionMessage(connected=True))
|
||||
|
||||
# Expect a broadcast for the new round at clock=1
|
||||
while True:
|
||||
|
||||
@@ -169,6 +169,22 @@ class Topology:
|
||||
for conn in new_connections:
|
||||
self.add_connection(conn)
|
||||
|
||||
def remove_all_rdma_connections_touching(self, node_id: NodeId) -> None:
|
||||
"""Remove every RDMA edge incident to ``node_id`` (incoming or outgoing)."""
|
||||
if node_id not in self._vertex_indices:
|
||||
return
|
||||
rx_idx = self._vertex_indices[node_id]
|
||||
rdma_edge_idxs = [
|
||||
edge_idx
|
||||
for edge_idx in (
|
||||
*self._graph.out_edge_indices(rx_idx),
|
||||
*self._graph.in_edge_indices(rx_idx),
|
||||
)
|
||||
if isinstance(self._graph.get_edge_data_by_index(edge_idx), RDMAConnection)
|
||||
]
|
||||
for edge_idx in rdma_edge_idxs:
|
||||
self._graph.remove_edge_from_index(edge_idx)
|
||||
|
||||
def remove_connection(self, conn: Connection) -> None:
|
||||
if (
|
||||
conn.source not in self._vertex_indices
|
||||
|
||||
@@ -32,15 +32,21 @@ class TokenChunk(BaseChunk):
|
||||
|
||||
class ErrorChunk(BaseChunk):
|
||||
error_message: str
|
||||
finish_reason: Literal["error"] = "error"
|
||||
|
||||
@property
|
||||
def finish_reason(self) -> Literal["error"]:
|
||||
return "error"
|
||||
|
||||
|
||||
class ToolCallChunk(BaseChunk):
|
||||
tool_calls: list[ToolCallItem]
|
||||
usage: Usage | None
|
||||
finish_reason: Literal["tool_calls"] = "tool_calls"
|
||||
stats: GenerationStats | None = None
|
||||
|
||||
@property
|
||||
def finish_reason(self) -> Literal["tool_calls"]:
|
||||
return "tool_calls"
|
||||
|
||||
|
||||
class ImageChunk(BaseChunk):
|
||||
data: str
|
||||
@@ -84,7 +90,13 @@ class PrefillProgressChunk(BaseChunk):
|
||||
processed_tokens: int
|
||||
total_tokens: int
|
||||
|
||||
@property
|
||||
def finish_reason(self) -> FinishReason | None:
|
||||
return None
|
||||
|
||||
|
||||
StatusChunk = PrefillProgressChunk
|
||||
GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk
|
||||
TextGenerationChunk = TokenChunk | ToolCallChunk | ErrorChunk
|
||||
ImageGenerationChunk = ImageChunk | ErrorChunk
|
||||
Chunk = StatusChunk | GenerationChunk
|
||||
@@ -159,10 +159,10 @@ Event = (
|
||||
| NodeTimedOut
|
||||
| NodeGatheredInfo
|
||||
| NodeDownloadProgress
|
||||
| ChunkGenerated
|
||||
| InputChunkReceived
|
||||
| TopologyEdgeCreated
|
||||
| TopologyEdgeDeleted
|
||||
| ChunkGenerated
|
||||
| InputChunkReceived
|
||||
| TracesCollected
|
||||
| TracesMerged
|
||||
| CustomModelCardAdded
|
||||
|
||||
@@ -5,8 +5,9 @@ from typing import Any, cast
|
||||
from pydantic import ConfigDict, Field, field_serializer, field_validator
|
||||
from pydantic.alias_generators import to_camel
|
||||
|
||||
from exo.shared.models.model_cards import ModelCard
|
||||
from exo.shared.topology import Topology, TopologySnapshot
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.common import ModelId, NodeId
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.profiling import (
|
||||
DiskUsage,
|
||||
@@ -58,7 +59,6 @@ class State(FrozenModel):
|
||||
node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
|
||||
node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
|
||||
node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
|
||||
node_vllm: Mapping[NodeId, bool] = {}
|
||||
|
||||
# Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
|
||||
thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
|
||||
@@ -66,6 +66,9 @@ class State(FrozenModel):
|
||||
instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
|
||||
prefill_server_ports: Mapping[RunnerId, int] = {}
|
||||
|
||||
# User-added model cards. Workers can reconcile their on-disk custom card cache
|
||||
custom_model_cards: Mapping[ModelId, ModelCard] = {}
|
||||
|
||||
@field_serializer("topology", mode="plain")
|
||||
def _encode_topology(self, value: Topology) -> TopologySnapshot:
|
||||
return value.to_snapshot()
|
||||
|
||||
@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
|
||||
prefill_endpoint: str | None = None
|
||||
|
||||
def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
|
||||
from exo.shared.models.model_cards import get_card
|
||||
from exo.shared.models import model_cards
|
||||
|
||||
card = get_card(self.model)
|
||||
card = model_cards.card_cache.get(self.model)
|
||||
if card is None:
|
||||
return self
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@ class InstanceId(Id):
|
||||
class InstanceMeta(str, Enum):
|
||||
MlxRing = "MlxRing"
|
||||
MlxJaccl = "MlxJaccl"
|
||||
Vllm = "Vllm"
|
||||
|
||||
|
||||
class BaseInstance(TaggedModel):
|
||||
@@ -36,12 +35,8 @@ class MlxJacclInstance(BaseInstance):
|
||||
jaccl_coordinators: dict[NodeId, str]
|
||||
|
||||
|
||||
class VllmInstance(BaseInstance):
|
||||
pass
|
||||
|
||||
|
||||
# TODO: Single node instance
|
||||
Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
|
||||
Instance = MlxRingInstance | MlxJacclInstance
|
||||
|
||||
|
||||
class BoundInstance(FrozenModel):
|
||||
|
||||
@@ -25,12 +25,12 @@ def print_startup_banner(port: int) -> None:
|
||||
banner = f"""
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ ║
|
||||
║ ███████╗██╗ ██╗ ██████╗ ██████╗ ██████╗ ██╗ ██╗ ║
|
||||
║ ██╔════╝╚██╗██╔╝██╔═══██╗ ██ ██╔══██╗██╔════╝ ╚██╗██╔╝ ║
|
||||
║ █████╗ ╚███╔╝ ██║ ██║ ██████╗ ██║ ██║██║ ███╗ ╚███╔╝ ║
|
||||
║ ██╔══╝ ██╔██╗ ██║ ██║ ╚═██╔═╝ ██║ ██║██║ ██║ ██╔██╗ ║
|
||||
║ ███████╗██╔╝ ██╗╚██████╔╝ ╚═╝ ██████╔╝╚██████╔╝██╔╝ ██╗ ║
|
||||
║ ╚══════╝╚═╝ ╚═╝ ╚═════╝ ╚═════╝ ╚═════╝ ╚═╝ ╚═╝ ║
|
||||
║ ███████╗██╗ ██╗ ██████╗ ║
|
||||
║ ██╔════╝╚██╗██╔╝██╔═══██╗ ║
|
||||
║ █████╗ ╚███╔╝ ██║ ██║ ║
|
||||
║ ██╔══╝ ██╔██╗ ██║ ██║ ║
|
||||
║ ███████╗██╔╝ ██╗╚██████╔╝ ║
|
||||
║ ╚══════╝╚═╝ ╚═╝ ╚═════╝ ║
|
||||
║ ║
|
||||
║ Distributed AI Inference Cluster ║
|
||||
║ ║
|
||||
@@ -38,7 +38,7 @@ def print_startup_banner(port: int) -> None:
|
||||
|
||||
╔═══════════════════════════════════════════════════════════════════════╗
|
||||
║ ║
|
||||
║ 🌐 Dashboard & API Ready ║
|
||||
║ Dashboard & API Ready ║
|
||||
║ ║
|
||||
║ {dashboard_url}{" " * (69 - len(dashboard_url))}║
|
||||
║ ║
|
||||
|
||||
Loaded 100 of 133 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user