mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-09 03:51:22 -04:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
343d5bc6d4 | ||
|
|
0e6a56baee | ||
|
|
f7bdef9f08 | ||
|
|
f792bd5d52 | ||
|
|
9c6ff4ce95 | ||
|
|
b26268dfaf | ||
|
|
8dae3ecb9a | ||
|
|
fb12b403ea | ||
|
|
1606e63816 |
No files matched your search
@@ -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"
|
||||
|
||||
@@ -3435,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"
|
||||
@@ -4822,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"
|
||||
@@ -4968,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
|
||||
|
||||
+1
-13
@@ -254,7 +254,6 @@ class API:
|
||||
self.node_id: NodeId = node_id
|
||||
self.last_completed_election: int = 0
|
||||
self.port = port
|
||||
self._sent_image_hashes: set[str] = set()
|
||||
|
||||
self.paused: bool = False
|
||||
self.paused_ev: anyio.Event = anyio.Event()
|
||||
@@ -304,7 +303,6 @@ class API:
|
||||
self.event_receiver.close()
|
||||
self.event_receiver = event_receiver
|
||||
self._tg.start_soon(self._apply_state)
|
||||
self._sent_image_hashes = set()
|
||||
|
||||
def unpause(self, result_clock: int):
|
||||
logger.info("Unpausing API")
|
||||
@@ -826,18 +824,8 @@ class API:
|
||||
)
|
||||
command = TextGeneration(task_params=task_params)
|
||||
|
||||
new_images: list[tuple[int, str]] = []
|
||||
for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
|
||||
if h not in self._sent_image_hashes:
|
||||
self._sent_image_hashes.add(h)
|
||||
new_images.append((idx, img))
|
||||
|
||||
if not new_images:
|
||||
await self._send(command)
|
||||
return command
|
||||
|
||||
all_chunks: list[tuple[int, str]] = []
|
||||
for img_idx, img_data in new_images:
|
||||
for img_idx, img_data in enumerate(images):
|
||||
for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
|
||||
all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -80,6 +80,9 @@ class EventRouter:
|
||||
def shutdown(self) -> None:
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
def set_buffer_start(self, idx: int) -> None:
|
||||
self.event_buffer.fast_forward_to(idx)
|
||||
|
||||
async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
|
||||
idx = 0
|
||||
with recv as events:
|
||||
@@ -95,7 +98,6 @@ class EventRouter:
|
||||
self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
|
||||
|
||||
async def _run_ext_in(self):
|
||||
buf = OrderedBuffer[Event]()
|
||||
with self.external_inbound as events:
|
||||
async for event in events:
|
||||
if event.session != self.session_id:
|
||||
@@ -103,12 +105,12 @@ class EventRouter:
|
||||
if event.origin != self.session_id.master_node_id:
|
||||
continue
|
||||
|
||||
buf.ingest(event.origin_idx, event.event)
|
||||
self.event_buffer.ingest(event.origin_idx, event.event)
|
||||
event_id = event.event.event_id
|
||||
if event_id in self.out_for_delivery:
|
||||
self.out_for_delivery.pop(event_id)
|
||||
|
||||
drained = buf.drain_indexed()
|
||||
drained = self.event_buffer.drain_indexed()
|
||||
if drained:
|
||||
self._nack_attempts = 0
|
||||
if self._nack_cancel_scope:
|
||||
@@ -119,7 +121,9 @@ class EventRouter:
|
||||
or self._nack_cancel_scope.cancel_called
|
||||
):
|
||||
# Request the next index.
|
||||
self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
|
||||
self._tg.start_soon(
|
||||
self._nack_request, self.event_buffer.next_idx_to_release
|
||||
)
|
||||
continue
|
||||
|
||||
for idx, event in drained:
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Reassembles a snapshot from a stream of `SnapshotChunk`s.
|
||||
|
||||
A receiver belongs to one node; it ignores chunks addressed to other
|
||||
requesters and chunks from prior sessions. Once a transfer's chunks have
|
||||
all been collected and the SHA-256 checks out, the snapshot is decoded into
|
||||
a `State`. Concurrent transfers (for the same requester) are tolerated:
|
||||
each is keyed by `transfer_id`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from dataclasses import dataclass, field
|
||||
from typing import final
|
||||
|
||||
import zstandard
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
|
||||
from exo.shared.types.state import State
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class _Assembly:
|
||||
"""Partial state for one in-flight snapshot transfer."""
|
||||
|
||||
total_chunks: int
|
||||
sha256_hex: str
|
||||
schema_version: int
|
||||
last_event_applied_idx: int
|
||||
chunks: dict[int, bytes] = field(default_factory=dict)
|
||||
|
||||
def is_complete(self) -> bool:
|
||||
return len(self.chunks) == self.total_chunks
|
||||
|
||||
def assemble(self) -> bytes:
|
||||
return b"".join(self.chunks[i] for i in range(self.total_chunks))
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReceivedSnapshot:
|
||||
last_event_applied_idx: int
|
||||
state: State
|
||||
|
||||
|
||||
class SnapshotReceiver:
|
||||
"""Filters and reassembles inbound chunks into a `ReceivedSnapshot`.
|
||||
|
||||
Stateless w.r.t. delivery: callers feed `SnapshotChunk`s in via `ingest`
|
||||
and check the return value for completion.
|
||||
"""
|
||||
|
||||
def __init__(self, my_node_id: NodeId, session_id: SessionId) -> None:
|
||||
self._my_node_id = my_node_id
|
||||
self._session_id = session_id
|
||||
self._assemblies: dict[SnapshotTransferId, _Assembly] = {}
|
||||
|
||||
def ingest(self, chunk: SnapshotChunk) -> ReceivedSnapshot | None:
|
||||
"""Absorb a chunk; return the snapshot once a transfer completes.
|
||||
|
||||
Returns None for partial transfers, mismatched recipients, stale
|
||||
sessions, version mismatches, or corrupt payloads.
|
||||
"""
|
||||
if chunk.requester_node_id != self._my_node_id:
|
||||
return None
|
||||
if chunk.session_id != self._session_id:
|
||||
return None
|
||||
|
||||
existing = self._assemblies.get(chunk.transfer_id)
|
||||
if existing is None:
|
||||
existing = _Assembly(
|
||||
total_chunks=chunk.total_chunks,
|
||||
sha256_hex=chunk.sha256_hex,
|
||||
schema_version=chunk.schema_version,
|
||||
last_event_applied_idx=chunk.last_event_applied_idx,
|
||||
)
|
||||
self._assemblies[chunk.transfer_id] = existing
|
||||
existing.chunks[chunk.chunk_index] = chunk.data
|
||||
|
||||
if not existing.is_complete():
|
||||
return None
|
||||
|
||||
# Transfer complete — finalise and remove from the in-flight map.
|
||||
del self._assemblies[chunk.transfer_id]
|
||||
body = existing.assemble()
|
||||
if hashlib.sha256(body).hexdigest() != existing.sha256_hex:
|
||||
logger.warning(f"Snapshot {chunk.transfer_id} failed checksum; discarding")
|
||||
return None
|
||||
try:
|
||||
decompressed = zstandard.ZstdDecompressor().decompress(body)
|
||||
state = State.model_validate_json(decompressed.decode("utf-8"))
|
||||
except (zstandard.ZstdError, ValueError) as e:
|
||||
logger.opt(exception=e).warning(
|
||||
f"Snapshot {chunk.transfer_id} could not be decoded; discarding"
|
||||
)
|
||||
return None
|
||||
if state.schema_version != existing.schema_version:
|
||||
# Should not happen — the master writes schema_version into both
|
||||
# the chunk meta and the State payload — but treat it as corrupt.
|
||||
logger.warning(
|
||||
f"Snapshot {chunk.transfer_id} schema version mismatch "
|
||||
f"(chunk={existing.schema_version}, state={state.schema_version})"
|
||||
)
|
||||
return None
|
||||
return ReceivedSnapshot(
|
||||
last_event_applied_idx=existing.last_event_applied_idx, state=state
|
||||
)
|
||||
@@ -141,3 +141,28 @@ async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]):
|
||||
assert [e[0] for e in drained] == [2]
|
||||
assert buffer.next_idx_to_release == 3
|
||||
assert 4 in buffer.store
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fast_forward_discards_buffered_stale_events(
|
||||
buffer: OrderedBuffer[Event],
|
||||
):
|
||||
buffer.ingest(*make_indexed_event(0))
|
||||
buffer.ingest(*make_indexed_event(2))
|
||||
buffer.ingest(*make_indexed_event(4))
|
||||
|
||||
buffer.fast_forward_to(3)
|
||||
|
||||
assert buffer.next_idx_to_release == 3
|
||||
assert set(buffer.store) == {4}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fast_forward_only_moves_forward(buffer: OrderedBuffer[Event]):
|
||||
buffer.ingest(*make_indexed_event(0))
|
||||
buffer.ingest(*make_indexed_event(1))
|
||||
buffer.drain()
|
||||
|
||||
buffer.fast_forward_to(1)
|
||||
|
||||
assert buffer.next_idx_to_release == 2
|
||||
@@ -0,0 +1,151 @@
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
import zstandard
|
||||
|
||||
from exo.routing.snapshot_receiver import SnapshotReceiver
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
|
||||
from exo.shared.types.state import State
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session_id() -> SessionId:
|
||||
return SessionId(master_node_id=NodeId("master"), election_clock=0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def my_node() -> NodeId:
|
||||
return NodeId("worker-1")
|
||||
|
||||
|
||||
def _encode(state: State) -> bytes:
|
||||
return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
|
||||
|
||||
|
||||
def _make_chunks(
|
||||
body: bytes,
|
||||
*,
|
||||
chunk_size: int,
|
||||
requester_node_id: NodeId,
|
||||
session_id: SessionId,
|
||||
state: State,
|
||||
transfer_id: SnapshotTransferId | None = None,
|
||||
) -> list[SnapshotChunk]:
|
||||
sha256 = hashlib.sha256(body).hexdigest()
|
||||
transfer_id = transfer_id or SnapshotTransferId()
|
||||
pieces = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] or [
|
||||
b""
|
||||
]
|
||||
return [
|
||||
SnapshotChunk.from_data(
|
||||
data=piece,
|
||||
transfer_id=transfer_id,
|
||||
requester_node_id=requester_node_id,
|
||||
session_id=session_id,
|
||||
schema_version=state.schema_version,
|
||||
last_event_applied_idx=state.last_event_applied_idx,
|
||||
chunk_index=i,
|
||||
total_chunks=len(pieces),
|
||||
sha256_hex=sha256,
|
||||
)
|
||||
for i, piece in enumerate(pieces)
|
||||
]
|
||||
|
||||
|
||||
def test_completes_on_full_transfer(my_node: NodeId, session_id: SessionId):
|
||||
state = State(last_event_applied_idx=42)
|
||||
chunks = _make_chunks(
|
||||
_encode(state),
|
||||
chunk_size=64,
|
||||
requester_node_id=my_node,
|
||||
session_id=session_id,
|
||||
state=state,
|
||||
)
|
||||
receiver = SnapshotReceiver(my_node, session_id)
|
||||
|
||||
received = None
|
||||
for chunk in chunks:
|
||||
received = receiver.ingest(chunk)
|
||||
assert received is not None
|
||||
assert received.last_event_applied_idx == 42
|
||||
assert received.state.last_event_applied_idx == 42
|
||||
|
||||
|
||||
def test_handles_out_of_order_chunks(my_node: NodeId, session_id: SessionId):
|
||||
state = State(last_event_applied_idx=99)
|
||||
chunks = _make_chunks(
|
||||
_encode(state),
|
||||
chunk_size=32,
|
||||
requester_node_id=my_node,
|
||||
session_id=session_id,
|
||||
state=state,
|
||||
)
|
||||
receiver = SnapshotReceiver(my_node, session_id)
|
||||
|
||||
# Reverse them.
|
||||
received = None
|
||||
for chunk in reversed(chunks):
|
||||
received = receiver.ingest(chunk)
|
||||
assert received is not None
|
||||
assert received.last_event_applied_idx == 99
|
||||
|
||||
|
||||
def test_ignores_chunks_for_other_recipients(my_node: NodeId, session_id: SessionId):
|
||||
state = State(last_event_applied_idx=1)
|
||||
other = NodeId("worker-2")
|
||||
chunks = _make_chunks(
|
||||
_encode(state),
|
||||
chunk_size=64,
|
||||
requester_node_id=other,
|
||||
session_id=session_id,
|
||||
state=state,
|
||||
)
|
||||
receiver = SnapshotReceiver(my_node, session_id)
|
||||
for chunk in chunks:
|
||||
assert receiver.ingest(chunk) is None
|
||||
|
||||
|
||||
def test_ignores_chunks_from_stale_session(my_node: NodeId, session_id: SessionId):
|
||||
state = State(last_event_applied_idx=1)
|
||||
other_session = SessionId(master_node_id=NodeId("other-master"), election_clock=99)
|
||||
chunks = _make_chunks(
|
||||
_encode(state),
|
||||
chunk_size=64,
|
||||
requester_node_id=my_node,
|
||||
session_id=other_session,
|
||||
state=state,
|
||||
)
|
||||
receiver = SnapshotReceiver(my_node, session_id)
|
||||
for chunk in chunks:
|
||||
assert receiver.ingest(chunk) is None
|
||||
|
||||
|
||||
def test_discards_on_checksum_mismatch(my_node: NodeId, session_id: SessionId):
|
||||
state = State(last_event_applied_idx=1)
|
||||
chunks = _make_chunks(
|
||||
_encode(state),
|
||||
chunk_size=64,
|
||||
requester_node_id=my_node,
|
||||
session_id=session_id,
|
||||
state=state,
|
||||
)
|
||||
# Corrupt the last byte of the last chunk.
|
||||
original = chunks[-1]
|
||||
chunks[-1] = SnapshotChunk.from_data(
|
||||
data=original.data + b"\x00garbage",
|
||||
transfer_id=original.transfer_id,
|
||||
requester_node_id=original.requester_node_id,
|
||||
session_id=original.session_id,
|
||||
schema_version=original.schema_version,
|
||||
last_event_applied_idx=original.last_event_applied_idx,
|
||||
chunk_index=original.chunk_index,
|
||||
total_chunks=original.total_chunks,
|
||||
sha256_hex=original.sha256_hex,
|
||||
)
|
||||
|
||||
receiver = SnapshotReceiver(my_node, session_id)
|
||||
received = None
|
||||
for chunk in chunks:
|
||||
received = receiver.ingest(chunk)
|
||||
assert received is None
|
||||
+32
-2
@@ -40,7 +40,14 @@ from exo.shared.types.profiling import (
|
||||
ThunderboltBridgeStatus,
|
||||
)
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.tasks import Task, TaskId, TaskStatus
|
||||
from exo.shared.types.tasks import (
|
||||
ImageEdits,
|
||||
ImageGeneration,
|
||||
Task,
|
||||
TaskId,
|
||||
TaskStatus,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.topology import Connection, RDMAConnection
|
||||
from exo.shared.types.worker.downloads import DownloadProgress
|
||||
from exo.shared.types.worker.instances import Instance, InstanceId
|
||||
@@ -72,7 +79,6 @@ def event_apply(event: Event, state: State) -> State:
|
||||
TestEvent()
|
||||
| ChunkGenerated()
|
||||
| TaskAcknowledged()
|
||||
| InputChunkReceived()
|
||||
| TracesCollected()
|
||||
| TracesMerged()
|
||||
| CustomModelCardAdded()
|
||||
@@ -93,6 +99,8 @@ def event_apply(event: Event, state: State) -> State:
|
||||
return apply_runner_status_updated(event, state)
|
||||
case TaskCreated():
|
||||
return apply_task_created(event, state)
|
||||
case InputChunkReceived():
|
||||
return apply_input_chunk_received(event, state)
|
||||
case TaskDeleted():
|
||||
return apply_task_deleted(event, state)
|
||||
case TaskFailed():
|
||||
@@ -157,10 +165,32 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
|
||||
return state.model_copy(update={"tasks": new_tasks})
|
||||
|
||||
|
||||
def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
|
||||
command_chunks = {
|
||||
**state.input_chunks.get(event.command_id, {}),
|
||||
event.chunk.chunk_index: event.chunk,
|
||||
}
|
||||
return state.model_copy(
|
||||
update={
|
||||
"input_chunks": {**state.input_chunks, event.command_id: command_chunks}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def apply_task_deleted(event: TaskDeleted, state: State) -> State:
|
||||
task = state.tasks.get(event.task_id)
|
||||
new_tasks: Mapping[TaskId, Task] = {
|
||||
tid: task for tid, task in state.tasks.items() if tid != event.task_id
|
||||
}
|
||||
if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
|
||||
new_input_chunks = {
|
||||
command_id: chunks
|
||||
for command_id, chunks in state.input_chunks.items()
|
||||
if command_id != task.command_id
|
||||
}
|
||||
return state.model_copy(
|
||||
update={"tasks": new_tasks, "input_chunks": new_input_chunks}
|
||||
)
|
||||
return state.model_copy(update={"tasks": new_tasks})
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.models.model_cards import ModelId
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.events import (
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
TaskCreated,
|
||||
TaskDeleted,
|
||||
)
|
||||
from exo.shared.types.state import State
|
||||
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
InputMessageContent,
|
||||
TextGenerationTaskParams,
|
||||
)
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
|
||||
|
||||
def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
|
||||
command_id = CommandId("command-1")
|
||||
chunk = InputImageChunk(
|
||||
model=ModelId("mlx-community/test-model"),
|
||||
command_id=command_id,
|
||||
data="abc",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
image_index=0,
|
||||
)
|
||||
|
||||
state = apply(
|
||||
State(),
|
||||
IndexedEvent(
|
||||
idx=0,
|
||||
event=InputChunkReceived(command_id=command_id, chunk=chunk),
|
||||
),
|
||||
)
|
||||
|
||||
assert state.input_chunks == {command_id: {0: chunk}}
|
||||
|
||||
|
||||
def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
|
||||
command_id = CommandId("command-1")
|
||||
task_id = TaskId("task-1")
|
||||
chunk = InputImageChunk(
|
||||
model=ModelId("mlx-community/test-model"),
|
||||
command_id=command_id,
|
||||
data="abc",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
image_index=0,
|
||||
)
|
||||
task = TextGeneration(
|
||||
task_id=task_id,
|
||||
instance_id=InstanceId("instance-1"),
|
||||
task_status=TaskStatus.Pending,
|
||||
command_id=command_id,
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=ModelId("mlx-community/test-model"),
|
||||
input=[
|
||||
InputMessage(role="user", content=InputMessageContent("hello")),
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
state = State()
|
||||
state = apply(
|
||||
state,
|
||||
IndexedEvent(
|
||||
idx=0,
|
||||
event=InputChunkReceived(command_id=command_id, chunk=chunk),
|
||||
),
|
||||
)
|
||||
state = apply(
|
||||
state,
|
||||
IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
|
||||
)
|
||||
state = apply(
|
||||
state,
|
||||
IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
|
||||
)
|
||||
|
||||
assert state.tasks == {}
|
||||
assert state.input_chunks == {}
|
||||
@@ -25,6 +25,7 @@ def test_state_serialization_roundtrip() -> None:
|
||||
json_repr = state.model_dump_json()
|
||||
restored_state = State.model_validate_json(json_repr)
|
||||
|
||||
assert restored_state.schema_version == state.schema_version
|
||||
assert (
|
||||
state.topology.to_snapshot().nodes
|
||||
== restored_state.topology.to_snapshot().nodes
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Wire types for snapshot transfer between master and a joining node.
|
||||
|
||||
Snapshots can be tens of MB; the gossipsub message ceiling is around 1 MB.
|
||||
We slice the compressed snapshot body into chunks and publish each chunk on
|
||||
the SNAPSHOT_RESPONSES topic. The receiver collects chunks for its own
|
||||
`requester_node_id`, validates the SHA-256 of the reassembled body, and
|
||||
materialises the State.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
from exo.shared.types.common import Id, NodeId, SessionId
|
||||
from exo.utils.pydantic_ext import FrozenModel
|
||||
|
||||
|
||||
class SnapshotTransferId(Id):
|
||||
"""Identifies a single snapshot transfer (one master response to one
|
||||
`RequestSnapshot`). Distinct transfers may interleave; the id lets
|
||||
receivers keep them apart."""
|
||||
|
||||
|
||||
class SnapshotChunk(FrozenModel):
|
||||
"""One slice of a snapshot in flight.
|
||||
|
||||
`data_b64` carries a base64-encoded slice of the zstd-compressed JSON
|
||||
dump of State. Concatenating the *decoded* bytes of all chunks for a
|
||||
`transfer_id` in order of `chunk_index` yields the full compressed
|
||||
body; `sha256_hex` is the SHA-256 of that decoded blob.
|
||||
|
||||
We use base64 explicitly because the topic layer JSON-encodes messages,
|
||||
and JSON can't carry raw binary. Helpers `from_data` / `data` keep the
|
||||
base64 detail at the boundaries.
|
||||
"""
|
||||
|
||||
transfer_id: SnapshotTransferId
|
||||
requester_node_id: NodeId
|
||||
session_id: SessionId
|
||||
schema_version: int
|
||||
last_event_applied_idx: int
|
||||
chunk_index: int
|
||||
total_chunks: int
|
||||
sha256_hex: str
|
||||
data_b64: str
|
||||
|
||||
@classmethod
|
||||
def from_data(cls, *, data: bytes, **kwargs: object) -> "SnapshotChunk":
|
||||
return cls(data_b64=base64.b64encode(data).decode("ascii"), **kwargs) # pyright: ignore[reportArgumentType]
|
||||
|
||||
@property
|
||||
def data(self) -> bytes:
|
||||
return base64.b64decode(self.data_b64)
|
||||
|
||||
|
||||
__all__ = ["SnapshotChunk", "SnapshotTransferId"]
|
||||
@@ -6,7 +6,8 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
|
||||
from pydantic.alias_generators import to_camel
|
||||
|
||||
from exo.shared.topology import Topology, TopologySnapshot
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.types.common import CommandId, NodeId
|
||||
from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
|
||||
from exo.shared.types.profiling import (
|
||||
DiskUsage,
|
||||
@@ -41,10 +42,16 @@ class State(FrozenModel):
|
||||
strict=True,
|
||||
arbitrary_types_allowed=True,
|
||||
)
|
||||
# Bump when a State change makes older snapshots unsafe to restore.
|
||||
schema_version: int = Field(default=1, ge=1)
|
||||
|
||||
instances: Mapping[InstanceId, Instance] = {}
|
||||
runners: Mapping[RunnerId, RunnerStatus] = {}
|
||||
downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
|
||||
tasks: Mapping[TaskId, Task] = {}
|
||||
# Durable request input chunks for active image requests. Workers rebuild
|
||||
# local image caches from this state instead of reading events directly.
|
||||
input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
|
||||
last_seen: Mapping[NodeId, datetime] = {}
|
||||
topology: Topology = Field(default_factory=Topology)
|
||||
last_event_applied_idx: int = Field(default=-1, ge=-1)
|
||||
|
||||
@@ -47,6 +47,18 @@ class OrderedBuffer[T]:
|
||||
logger.trace(f"Releasing event {ret}")
|
||||
return ret
|
||||
|
||||
def fast_forward_to(self, idx: int) -> None:
|
||||
"""Skip every event before idx.
|
||||
|
||||
Snapshot restore uses this after applying state that already includes
|
||||
events before idx. Any buffered or future event below idx is stale.
|
||||
"""
|
||||
if idx <= self.next_idx_to_release:
|
||||
return
|
||||
self.next_idx_to_release = idx
|
||||
for stale_idx in [i for i in self.store if i < idx]:
|
||||
del self.store[stale_idx]
|
||||
|
||||
|
||||
class MultiSourceBuffer[SourceId, T]:
|
||||
"""
|
||||
|
||||
+30
-43
@@ -24,7 +24,6 @@ from exo.shared.types.events import (
|
||||
CustomModelCardDeleted,
|
||||
Event,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceDeleted,
|
||||
NodeDownloadProgress,
|
||||
NodeGatheredInfo,
|
||||
@@ -141,37 +140,6 @@ class Worker:
|
||||
if isinstance(event, InstanceDeleted):
|
||||
self._instance_backoff.reset(event.instance_id)
|
||||
|
||||
# Buffer input image chunks for image editing
|
||||
if isinstance(event, InputChunkReceived):
|
||||
cmd_id = event.command_id
|
||||
if cmd_id not in self.input_chunk_buffer:
|
||||
self.input_chunk_buffer[cmd_id] = {}
|
||||
self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
|
||||
|
||||
self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
|
||||
event.chunk
|
||||
)
|
||||
|
||||
if (
|
||||
len(self.input_chunk_buffer[cmd_id])
|
||||
== self.input_chunk_counts[cmd_id]
|
||||
):
|
||||
per_image: defaultdict[int, list[InputImageChunk]] = (
|
||||
defaultdict(list)
|
||||
)
|
||||
for chunk in self.input_chunk_buffer[cmd_id].values():
|
||||
per_image[chunk.image_index].append(chunk)
|
||||
for chunks_for_image in per_image.values():
|
||||
sorted_chunks = sorted(
|
||||
chunks_for_image, key=lambda c: c.chunk_index
|
||||
)
|
||||
img = Base64Image("".join(c.data for c in sorted_chunks))
|
||||
self.image_cache[
|
||||
Base64ImageHash(
|
||||
hashlib.sha256(img.encode("ascii")).hexdigest()
|
||||
)
|
||||
] = img
|
||||
|
||||
if isinstance(event, CustomModelCardAdded):
|
||||
await event.model_card.save_to_custom_dir()
|
||||
add_to_card_cache(event.model_card)
|
||||
@@ -179,6 +147,35 @@ class Worker:
|
||||
if isinstance(event, CustomModelCardDeleted):
|
||||
await delete_custom_card(event.model_id)
|
||||
|
||||
self._sync_input_views_from_state()
|
||||
|
||||
def _sync_input_views_from_state(self) -> None:
|
||||
self.input_chunk_buffer = {
|
||||
command_id: dict(chunks)
|
||||
for command_id, chunks in self.state.input_chunks.items()
|
||||
}
|
||||
self.input_chunk_counts = {
|
||||
command_id: next(iter(chunks.values())).total_chunks
|
||||
for command_id, chunks in self.input_chunk_buffer.items()
|
||||
if chunks
|
||||
}
|
||||
|
||||
self.image_cache = {}
|
||||
for command_id, chunks in self.input_chunk_buffer.items():
|
||||
expected_chunks = self.input_chunk_counts.get(command_id)
|
||||
if expected_chunks is None or len(chunks) != expected_chunks:
|
||||
continue
|
||||
|
||||
per_image: defaultdict[int, list[InputImageChunk]] = defaultdict(list)
|
||||
for chunk in chunks.values():
|
||||
per_image[chunk.image_index].append(chunk)
|
||||
for chunks_for_image in per_image.values():
|
||||
sorted_chunks = sorted(chunks_for_image, key=lambda c: c.chunk_index)
|
||||
image = Base64Image("".join(chunk.data for chunk in sorted_chunks))
|
||||
self.image_cache[
|
||||
Base64ImageHash(hashlib.sha256(image.encode("ascii")).hexdigest())
|
||||
] = image
|
||||
|
||||
async def plan_step(self):
|
||||
while True:
|
||||
await anyio.sleep(0.1)
|
||||
@@ -189,7 +186,7 @@ class Worker:
|
||||
self.state.instances,
|
||||
self.state.runners,
|
||||
self.state.tasks,
|
||||
self.input_chunk_buffer,
|
||||
self.state.input_chunks,
|
||||
self.image_cache,
|
||||
self._instance_backoff,
|
||||
self._download_backoff,
|
||||
@@ -321,15 +318,9 @@ class Worker:
|
||||
advanced_params=task.task_params.advanced_params,
|
||||
),
|
||||
)
|
||||
# Cleanup buffers
|
||||
if cmd_id in self.input_chunk_buffer:
|
||||
del self.input_chunk_buffer[cmd_id]
|
||||
if cmd_id in self.input_chunk_counts:
|
||||
del self.input_chunk_counts[cmd_id]
|
||||
await self._start_runner_task(modified_task)
|
||||
|
||||
case TextGeneration() if task.task_params.image_hashes:
|
||||
cmd_id = task.command_id
|
||||
resolved_images = [
|
||||
self.image_cache[h]
|
||||
for _, h in sorted(task.task_params.image_hashes.items())
|
||||
@@ -341,10 +332,6 @@ class Worker:
|
||||
)
|
||||
}
|
||||
)
|
||||
if cmd_id in self.input_chunk_buffer:
|
||||
del self.input_chunk_buffer[cmd_id]
|
||||
if cmd_id in self.input_chunk_counts:
|
||||
del self.input_chunk_counts[cmd_id]
|
||||
await self._start_runner_task(modified_task)
|
||||
case LoadModel(instance_id=instance_id):
|
||||
if (instance := self.state.instances.get(instance_id)) is not None:
|
||||
|
||||
Reference in new issue
Block a user