mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-09 20:10:19 -04:00
Compare commits
13
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f9f4ee0575 | ||
|
|
c002f8111d | ||
|
|
fa57131374 | ||
|
|
414132ae9c | ||
|
|
edef8004f8 | ||
|
|
a0c00f9dfd | ||
|
|
89d20c1888 | ||
|
|
dbcceaa50c | ||
|
|
9c6ff4ce95 | ||
|
|
b26268dfaf | ||
|
|
8dae3ecb9a | ||
|
|
fb12b403ea | ||
|
|
1606e63816 |
No files matched your search
@@ -40,3 +40,4 @@ bench/**/*.json
|
||||
tmp/models
|
||||
/build/exo
|
||||
/.claude/skills
|
||||
/.claude
|
||||
@@ -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"
|
||||
|
||||
@@ -15,9 +15,8 @@ from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import httpx
|
||||
from harness import (
|
||||
ExoClient,
|
||||
ExoHttpError,
|
||||
from exo_tools.client import ExoClient, ExoHttpError
|
||||
from exo_tools.harness import (
|
||||
add_common_instance_args,
|
||||
capture_cluster_snapshot,
|
||||
instance_id_from_instance,
|
||||
|
||||
+2
-3
@@ -30,9 +30,8 @@ from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Any
|
||||
|
||||
from harness import (
|
||||
ExoClient,
|
||||
ExoHttpError,
|
||||
from exo_tools.client import ExoClient, ExoHttpError
|
||||
from exo_tools.harness import (
|
||||
add_common_instance_args,
|
||||
capture_cluster_snapshot,
|
||||
find_existing_instance,
|
||||
|
||||
+2
-3
@@ -42,9 +42,8 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from harness import (
|
||||
ExoClient,
|
||||
ExoHttpError,
|
||||
from exo_tools.client import ExoClient, ExoHttpError
|
||||
from exo_tools.harness import (
|
||||
add_common_instance_args,
|
||||
capture_cluster_snapshot,
|
||||
find_existing_instance,
|
||||
|
||||
@@ -35,9 +35,8 @@ from exo_bench import (
|
||||
load_tokenizer_for_bench,
|
||||
parse_int_list,
|
||||
)
|
||||
from harness import (
|
||||
ExoClient,
|
||||
ExoHttpError,
|
||||
from exo_tools.client import ExoClient, ExoHttpError
|
||||
from exo_tools.harness import (
|
||||
add_common_instance_args,
|
||||
instance_id_from_instance,
|
||||
node_ids_from_instance,
|
||||
|
||||
@@ -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
|
||||
|
||||
+12
-4
@@ -40,6 +40,7 @@ exo = "exo.main:main"
|
||||
dev = [
|
||||
"basedpyright>=1.29.0",
|
||||
"pyinstaller>=6.17.0",
|
||||
"playwright>=1.52.0",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
"pytest-env",
|
||||
@@ -75,11 +76,11 @@ cuda13 = [
|
||||
###
|
||||
|
||||
[tool.uv.workspace]
|
||||
members = ["rust/exo_pyo3_bindings", "bench"]
|
||||
members = ["rust/exo_pyo3_bindings", "bench", "tools"]
|
||||
|
||||
[tool.uv.sources]
|
||||
exo-pyo3-bindings = { workspace = true }
|
||||
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
|
||||
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "leo/order-distributed-ops", marker = "sys_platform == 'darwin'" }
|
||||
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
|
||||
torch = [
|
||||
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
|
||||
@@ -112,7 +113,7 @@ build-backend = "uv_build"
|
||||
###
|
||||
|
||||
[tool.basedpyright]
|
||||
include = ["src", "bench"]
|
||||
include = ["src", "bench", "tools"]
|
||||
typeCheckingMode = "strict"
|
||||
failOnWarnings = true
|
||||
|
||||
@@ -146,6 +147,13 @@ reportMissingModuleSource = false
|
||||
[[tool.basedpyright.executionEnvironments]]
|
||||
root = "src"
|
||||
|
||||
[[tool.basedpyright.executionEnvironments]]
|
||||
root = "bench"
|
||||
extraPaths = ["tools/src"]
|
||||
|
||||
[[tool.basedpyright.executionEnvironments]]
|
||||
root = "tools/src"
|
||||
|
||||
|
||||
###
|
||||
# uv configuration
|
||||
@@ -220,5 +228,5 @@ pythonpath = "."
|
||||
asyncio_mode = "auto"
|
||||
markers = ["slow: marks tests as slow (deselected by default)"]
|
||||
env = ["EXO_TESTS=1"]
|
||||
addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
|
||||
addopts = "-m 'not slow' --ignore=tests"
|
||||
filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]
|
||||
+12
-13
@@ -133,12 +133,10 @@ 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 (
|
||||
@@ -481,6 +479,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
|
||||
@@ -544,6 +543,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:
|
||||
@@ -1633,17 +1633,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())
|
||||
@@ -1656,8 +1655,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
|
||||
@@ -1720,7 +1719,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()
|
||||
@@ -1771,7 +1770,7 @@ class API:
|
||||
|
||||
# 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,
|
||||
@@ -1787,7 +1786,7 @@ 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")
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -365,6 +365,7 @@ class Master:
|
||||
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
|
||||
|
||||
@@ -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,
|
||||
@@ -105,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))
|
||||
@@ -166,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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
+50
-3
@@ -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,
|
||||
@@ -65,6 +66,18 @@ from exo.utils.info_gatherer.info_gatherer import (
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
@@ -75,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():
|
||||
@@ -397,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,
|
||||
@@ -409,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():
|
||||
@@ -432,6 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
|
||||
**state.node_rdma_ctl,
|
||||
event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
|
||||
}
|
||||
# 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)
|
||||
|
||||
@@ -447,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})
|
||||
@@ -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"
|
||||
@@ -196,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
|
||||
@@ -233,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
|
||||
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
@@ -65,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
|
||||
|
||||
|
||||
@@ -19,19 +19,21 @@ class PowerSampler:
|
||||
):
|
||||
self._get_node_system = get_node_system
|
||||
self._interval = interval
|
||||
self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
|
||||
defaultdict(list)
|
||||
)
|
||||
self._samples: defaultdict[
|
||||
NodeId, list[tuple[float, SystemPerformanceProfile]]
|
||||
] = defaultdict(list)
|
||||
self._start_time: float | None = None
|
||||
self._stopped = False
|
||||
|
||||
def _take_sample(self) -> None:
|
||||
def _take_sample(self, t_rel: float | None = None) -> None:
|
||||
assert self._start_time is not None
|
||||
ts = t_rel if t_rel is not None else time.perf_counter() - self._start_time
|
||||
for node_id, profile in self._get_node_system().items():
|
||||
self._samples[node_id].append(profile)
|
||||
self._samples[node_id].append((ts, profile))
|
||||
|
||||
async def run(self) -> None:
|
||||
self._start_time = time.perf_counter()
|
||||
self._take_sample()
|
||||
self._take_sample(t_rel=0.0)
|
||||
while not self._stopped:
|
||||
await anyio.sleep(self._interval)
|
||||
self._take_sample()
|
||||
@@ -39,26 +41,51 @@ class PowerSampler:
|
||||
def result(self) -> PowerUsage:
|
||||
self._stopped = True
|
||||
assert self._start_time is not None, "result() called before run()"
|
||||
self._take_sample()
|
||||
elapsed = time.perf_counter() - self._start_time
|
||||
self._take_sample(t_rel=elapsed)
|
||||
|
||||
node_stats: list[NodePowerStats] = []
|
||||
for node_id, profiles in self._samples.items():
|
||||
n = len(profiles)
|
||||
total_energy_j = 0.0
|
||||
for node_id, ts_profiles in self._samples.items():
|
||||
n = len(ts_profiles)
|
||||
if n == 0:
|
||||
continue
|
||||
node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
|
||||
avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
|
||||
total_energy_j += node_energy_j
|
||||
node_stats.append(
|
||||
NodePowerStats(
|
||||
node_id=node_id,
|
||||
samples=n,
|
||||
avg_sys_power=sum(p.sys_power for p in profiles) / n,
|
||||
avg_sys_power=avg_power_w,
|
||||
)
|
||||
)
|
||||
|
||||
total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
|
||||
total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
|
||||
return PowerUsage(
|
||||
elapsed_seconds=elapsed,
|
||||
nodes=node_stats,
|
||||
total_avg_sys_power_watts=total_avg_sys,
|
||||
total_energy_joules=total_avg_sys * elapsed,
|
||||
total_avg_sys_power_watts=total_avg_sys_w,
|
||||
total_energy_joules=total_energy_j,
|
||||
)
|
||||
|
||||
|
||||
def trapezoidal_energy(
|
||||
ts_profiles: list[tuple[float, SystemPerformanceProfile]],
|
||||
elapsed: float,
|
||||
) -> float:
|
||||
"""Integrate sys_power(t) over the sample window using the trapezoidal rule.
|
||||
First sample is anchored at t=0 and last at t=elapsed (set by `run` /
|
||||
`result`), so the integral spans the full request interval. Falls back to
|
||||
power * elapsed when only one sample exists (constant-power assumption)."""
|
||||
if len(ts_profiles) == 1:
|
||||
return ts_profiles[0][1].sys_power * elapsed
|
||||
energy_j = 0.0
|
||||
for i in range(1, len(ts_profiles)):
|
||||
t_prev, p_prev = ts_profiles[i - 1]
|
||||
t_cur, p_cur = ts_profiles[i]
|
||||
dt = t_cur - t_prev
|
||||
if dt <= 0:
|
||||
continue
|
||||
energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
|
||||
return energy_j
|
||||
@@ -111,6 +111,36 @@ async def test_empty_state() -> None:
|
||||
assert result.total_energy_joules == 0.0
|
||||
|
||||
|
||||
def test_trapezoidal_unit_dt_weighting() -> None:
|
||||
"""Pure unit test on the integration helper. Crafted samples where the
|
||||
arithmetic mean is wildly wrong vs the time-weighted result."""
|
||||
from exo.utils.power_sampler import trapezoidal_energy
|
||||
|
||||
# 5 s window. Power = 10 W for the first 4.9 s, then 100 W for the last 0.1 s.
|
||||
# Three samples: t=0 W=10, t=4.9 W=10, t=5.0 W=100.
|
||||
samples = [
|
||||
(0.0, _make_profile(10.0)),
|
||||
(4.9, _make_profile(10.0)),
|
||||
(5.0, _make_profile(100.0)),
|
||||
]
|
||||
energy = trapezoidal_energy(samples, elapsed=5.0)
|
||||
# (10+10)/2 * 4.9 + (10+100)/2 * 0.1 = 49 + 5.5 = 54.5 J
|
||||
assert abs(energy - 54.5) < 1e-9
|
||||
avg = energy / 5.0 # 10.9 W
|
||||
# Arithmetic mean of the three samples would be (10+10+100)/3 ≈ 40 W.
|
||||
# Trapezoidal correctly weights each segment by its dt.
|
||||
assert abs(avg - 10.9) < 1e-9
|
||||
|
||||
|
||||
def test_trapezoidal_unit_single_sample() -> None:
|
||||
"""One sample: no window to integrate over, so fall back to constant power
|
||||
over the elapsed duration."""
|
||||
from exo.utils.power_sampler import trapezoidal_energy
|
||||
|
||||
samples = [(0.0, _make_profile(42.0))]
|
||||
assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
|
||||
|
||||
|
||||
async def test_result_stops_sampling() -> None:
|
||||
"""Calling result() should stop the sampler's run loop."""
|
||||
state: dict[NodeId, SystemPerformanceProfile] = {
|
||||
|
||||
@@ -143,6 +143,7 @@ class ImageEngine(Engine):
|
||||
Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
|
||||
) = field(init=False, default=None)
|
||||
queue: deque[ImageTask] = field(init=False, default_factory=deque)
|
||||
_cancelled_tasks: set[TaskId] = field(init=False, default_factory=set)
|
||||
|
||||
def warmup(self) -> None:
|
||||
image = warmup_image_generator(model=self.image_model)
|
||||
@@ -168,7 +169,11 @@ class ImageEngine(Engine):
|
||||
task = self.queue.popleft()
|
||||
self.current_gen = self._run_image_task(task.task_id, task.task_params)
|
||||
resp = next(self.current_gen, None)
|
||||
return (resp,) if resp is not None else ()
|
||||
return (
|
||||
(resp,)
|
||||
if resp is not None and _is_primary_output_node(self.shard_metadata)
|
||||
else ()
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
with contextlib.suppress(NameError, AttributeError):
|
||||
|
||||
@@ -133,9 +133,7 @@ class PipelineFirstLayer(CustomMlxLayer):
|
||||
if self.r != 0:
|
||||
# We want to avoid GPU timeout errors by evalling the distributed operation
|
||||
# so that it stays on CPU, which does not have a timeout.
|
||||
mx.eval(x)
|
||||
x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
|
||||
mx.eval(x)
|
||||
return self.original_layer(x, *args, **kwargs)
|
||||
|
||||
|
||||
@@ -162,10 +160,6 @@ class PipelineLastLayer(CustomMlxLayer):
|
||||
|
||||
output: mx.array = self.original_layer(x, *args, **kwargs)
|
||||
|
||||
# Eval layer output to materialize it before send — this splits the graph
|
||||
# so the send is isolated and the receiving rank's recv can complete.
|
||||
mx.eval(output)
|
||||
|
||||
if self.r != self.s - 1:
|
||||
if self.queue_sends:
|
||||
_pending_prefill_sends.append(
|
||||
@@ -181,7 +175,6 @@ class PipelineLastLayer(CustomMlxLayer):
|
||||
_cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore
|
||||
if hasattr(_cache, "keys"): # pyright: ignore[reportAny]
|
||||
_cache.keys = mx.depends(_cache.keys, output) # type: ignore
|
||||
mx.eval(output)
|
||||
if cache is not None and hasattr(_cache, "keys"): # type: ignore
|
||||
mx.eval(_cache.keys) # type: ignore
|
||||
|
||||
@@ -189,7 +182,6 @@ class PipelineLastLayer(CustomMlxLayer):
|
||||
output = mx.distributed.all_gather(output, group=self.group)[
|
||||
-output.shape[0] :
|
||||
]
|
||||
mx.eval(output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
+14
-9
@@ -10,7 +10,7 @@ from exo.api.types import ImageEditsTaskParams
|
||||
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
|
||||
from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
|
||||
from exo.shared.models.model_cards import ModelId, card_cache
|
||||
from exo.shared.types.chunks import InputImageChunk
|
||||
from exo.shared.types.commands import (
|
||||
DeleteInstance,
|
||||
@@ -20,8 +20,6 @@ from exo.shared.types.commands import (
|
||||
)
|
||||
from exo.shared.types.common import CommandId, NodeId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
CustomModelCardAdded,
|
||||
CustomModelCardDeleted,
|
||||
Event,
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
@@ -110,6 +108,8 @@ class Worker:
|
||||
tg.start_soon(self.plan_step)
|
||||
tg.start_soon(self._event_applier)
|
||||
tg.start_soon(self._poll_connection_updates)
|
||||
tg.start_soon(self._reconcile_custom_cards)
|
||||
|
||||
finally:
|
||||
# Actual shutdown code - waits for all tasks to complete before executing.
|
||||
logger.info("Stopping Worker")
|
||||
@@ -151,7 +151,6 @@ class Worker:
|
||||
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]
|
||||
@@ -172,12 +171,18 @@ class Worker:
|
||||
)
|
||||
] = img
|
||||
|
||||
if isinstance(event, CustomModelCardAdded):
|
||||
await event.model_card.save_to_custom_dir()
|
||||
add_to_card_cache(event.model_card)
|
||||
async def _reconcile_custom_cards(self) -> None:
|
||||
while True:
|
||||
await anyio.sleep(1)
|
||||
target = dict(self.state.custom_model_cards)
|
||||
for model_id, card in target.items():
|
||||
if card_cache.get(model_id) == card:
|
||||
continue
|
||||
await card_cache.save(card)
|
||||
|
||||
if isinstance(event, CustomModelCardDeleted):
|
||||
await delete_custom_card(event.model_id)
|
||||
for card in await card_cache.list_all():
|
||||
if card.model_id not in target:
|
||||
await card_cache.pop(card.model_id)
|
||||
|
||||
async def plan_step(self):
|
||||
while True:
|
||||
|
||||
@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
|
||||
def agree_on_tasks(self) -> None:
|
||||
"""Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
|
||||
agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
|
||||
self._queue.extend(task for task in self._maybe_queue if task in agreed)
|
||||
self._maybe_queue = [task for task in self._maybe_queue if task in different]
|
||||
# Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
|
||||
# rank enqueues tasks in the same order, preventing TP collective deadlocks.
|
||||
self._queue.extend(agreed)
|
||||
self._maybe_queue = list(different)
|
||||
|
||||
def agree_on_cancellations(self) -> None:
|
||||
"""Agree between all ranks about which tasks to cancel."""
|
||||
@@ -197,9 +199,14 @@ class SequentialGenerator(Engine):
|
||||
self._active = None
|
||||
raise
|
||||
|
||||
return itertools.chain(
|
||||
output,
|
||||
map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
|
||||
return filter(
|
||||
lambda chunk: (
|
||||
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
|
||||
),
|
||||
itertools.chain(
|
||||
output,
|
||||
map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
|
||||
),
|
||||
)
|
||||
|
||||
def _start_next(self) -> None:
|
||||
@@ -368,8 +375,10 @@ class BatchGenerator(Engine):
|
||||
def agree_on_tasks(self) -> None:
|
||||
"""Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
|
||||
agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
|
||||
self._queue.extend(task for task in self._maybe_queue if task in agreed)
|
||||
self._maybe_queue = [task for task in self._maybe_queue if task in different]
|
||||
# Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
|
||||
# rank enqueues tasks in the same order, preventing TP collective deadlocks.
|
||||
self._queue.extend(agreed)
|
||||
self._maybe_queue = list(different)
|
||||
|
||||
def agree_on_cancellations(self) -> None:
|
||||
"""Agree between all ranks about which tasks to cancel."""
|
||||
@@ -449,7 +458,12 @@ class BatchGenerator(Engine):
|
||||
output.append((task.task_id, FinishedResponse()))
|
||||
del self._active_tasks[uid]
|
||||
|
||||
return itertools.chain(output, self._apply_cancellations())
|
||||
return filter(
|
||||
lambda chunk: (
|
||||
not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
|
||||
),
|
||||
itertools.chain(output, self._apply_cancellations()),
|
||||
)
|
||||
|
||||
def _apply_cancellations(
|
||||
self,
|
||||
|
||||
@@ -390,5 +390,5 @@ class Runner:
|
||||
chunk: Chunk,
|
||||
command_id: CommandId,
|
||||
):
|
||||
if self.device_rank == 0:
|
||||
self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
|
||||
assert isinstance(self.generator, Engine)
|
||||
self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
|
||||
@@ -16,7 +16,7 @@ from exo.download.download_utils import (
|
||||
fetch_file_list_with_cache,
|
||||
resolve_model_dir,
|
||||
)
|
||||
from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
|
||||
from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
|
||||
from exo.worker.engines.mlx.utils_mlx import (
|
||||
get_eos_token_ids_for_model,
|
||||
load_tokenizer_for_model_id,
|
||||
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
|
||||
"""Get a representative sample of models to test."""
|
||||
# Pick one model from each family to test
|
||||
families: dict[str, ModelCard] = {}
|
||||
for card in asyncio.run(get_model_cards()):
|
||||
for card in asyncio.run(card_cache.list_all()):
|
||||
# Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
|
||||
parts = card.model_id.short().split("-")
|
||||
family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
|
||||
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
|
||||
async def test_kimi_tokenizer_specifically():
|
||||
"""Test Kimi tokenizer with its specific patches and quirks."""
|
||||
kimi_models = [
|
||||
card for card in await get_model_cards() if "kimi" in card.model_id.lower()
|
||||
card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
|
||||
]
|
||||
|
||||
if not kimi_models:
|
||||
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
|
||||
|
||||
glm_model_cards = [
|
||||
card
|
||||
for card in await get_model_cards()
|
||||
for card in await card_cache.list_all()
|
||||
if contains(card, "glm")
|
||||
and not contains(card, "-5")
|
||||
and not contains(card, "4.7")
|
||||
|
||||
File renamed without changes.
@@ -0,0 +1,181 @@
|
||||
# type: ignore
|
||||
"""Pytest configuration for marker-driven exo integration tests.
|
||||
|
||||
Test authors declare requirements via markers:
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt='a2a')
|
||||
@pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
|
||||
sharding='tensor', comm='jaccl')
|
||||
def test_jaccl_inference(session):
|
||||
resp = session.chat('What is 2+2?')
|
||||
assert '4' in resp
|
||||
|
||||
Clusters are cached by `ClusterSpec`; tests with the same cluster_spec
|
||||
share a deployment. Each test places its own instance (matching its
|
||||
`@pytest.mark.instance`), and instances are cleaned up after the test.
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/ -v
|
||||
uv run pytest tests/ -v --hosts s2,s4,s9,s10
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from exo_tools.cluster import ClusterInfo, EcoSession
|
||||
from exo_tools.harness import cleanup_all_instances, place_instance
|
||||
|
||||
from .framework import (
|
||||
ClusterSpec,
|
||||
Session,
|
||||
parse_cluster_marker,
|
||||
parse_instance_marker,
|
||||
)
|
||||
|
||||
# Single eco session for the entire test process.
|
||||
eco = EcoSession(user_prefix="test")
|
||||
|
||||
# Cluster cache keyed by ClusterSpec — tests with the same spec share a deployment.
|
||||
# Cleared at session teardown.
|
||||
_cluster_cache: dict[ClusterSpec, ClusterInfo] = {}
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption(
|
||||
"--hosts",
|
||||
default=None,
|
||||
help="Comma-separated list of hosts (e.g. s2,s4,s9,s10). "
|
||||
"Overrides constraint-based reservation.",
|
||||
)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Register custom markers."""
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"cluster(count=N, thunderbolt=Thunderbolt|None, min_memory=GB, chip=PATTERN): "
|
||||
"declare cluster requirements for a test",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"instance(model_id, sharding=Sharding, comm=Comm, min_nodes=N): "
|
||||
"declare instance placement for a test",
|
||||
)
|
||||
|
||||
|
||||
def pytest_report_header(config):
|
||||
"""Show the eco user and hosts for this test session."""
|
||||
hosts = config.getoption("--hosts")
|
||||
lines = [f"eco user: {eco.user}"]
|
||||
if hosts:
|
||||
lines.append(f"hosts override: {hosts}")
|
||||
return lines
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def _host_pool(request) -> list[str] | None:
|
||||
raw = request.config.getoption("--hosts")
|
||||
if raw:
|
||||
return [h.strip() for h in raw.split(",") if h.strip()]
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session(request, _host_pool) -> Session:
|
||||
"""Per-test fixture providing a Session matching the test's markers.
|
||||
|
||||
Reads @pytest.mark.cluster and @pytest.mark.instance from the test, deploys
|
||||
a matching cluster (cached across tests with the same spec), places the
|
||||
model, and yields a Session for the test to interact with. Cleans up the
|
||||
instance after the test, and invalidates the cluster cache if the test
|
||||
left nodes disconnected.
|
||||
"""
|
||||
cluster_marker = request.node.get_closest_marker("cluster")
|
||||
instance_marker = request.node.get_closest_marker("instance")
|
||||
|
||||
cluster_spec = parse_cluster_marker(cluster_marker)
|
||||
instance_spec = parse_instance_marker(instance_marker)
|
||||
|
||||
# Deploy or reuse a cluster matching the spec
|
||||
cluster = _cluster_cache.get(cluster_spec)
|
||||
if cluster is None:
|
||||
if _host_pool:
|
||||
cluster = eco.start_deploy(
|
||||
hosts=_host_pool[: cluster_spec.count], wait=True
|
||||
)
|
||||
else:
|
||||
cluster = eco.start_deploy(
|
||||
count=cluster_spec.count,
|
||||
thunderbolt=cluster_spec.thunderbolt,
|
||||
chip=cluster_spec.chip,
|
||||
min_memory_gb=cluster_spec.min_memory_gb,
|
||||
wait=True,
|
||||
)
|
||||
_cluster_cache[cluster_spec] = cluster
|
||||
|
||||
# Place an instance for this test if the test specified one
|
||||
instance_id = None
|
||||
if instance_spec is not None:
|
||||
client = cluster.make_client()
|
||||
instance_id = place_instance(
|
||||
client,
|
||||
instance_spec.model_id,
|
||||
sharding=instance_spec.sharding,
|
||||
comm=instance_spec.comm,
|
||||
min_nodes=instance_spec.min_nodes,
|
||||
)
|
||||
|
||||
sess = Session(
|
||||
cluster=cluster,
|
||||
eco=eco,
|
||||
instance_spec=instance_spec,
|
||||
instance_id=instance_id,
|
||||
)
|
||||
|
||||
yield sess
|
||||
|
||||
# ---- Teardown ----
|
||||
|
||||
# If the test left nodes disconnected, invalidate the cluster cache and
|
||||
# stop the cluster so the next test deploys fresh.
|
||||
if sess._stopped_hosts:
|
||||
_cluster_cache.pop(cluster_spec, None)
|
||||
with contextlib.suppress(Exception):
|
||||
eco.stop(sess.cluster.hosts)
|
||||
return
|
||||
|
||||
# Otherwise, clean up any instances created during the test
|
||||
with contextlib.suppress(Exception):
|
||||
cleanup_all_instances(sess.client)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session-level teardown — stop all cached clusters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def _teardown_clusters():
|
||||
yield
|
||||
for cluster in _cluster_cache.values():
|
||||
with contextlib.suppress(Exception):
|
||||
eco.stop(cluster.hosts)
|
||||
_cluster_cache.clear()
|
||||
|
||||
|
||||
def pytest_runtest_makereport(item, call):
|
||||
"""Attach cluster logs to the test report when a test fails."""
|
||||
if call.when != "call" or call.excinfo is None:
|
||||
return
|
||||
|
||||
sess = item.funcargs.get("session")
|
||||
if sess is None:
|
||||
return
|
||||
try:
|
||||
logs = eco.logs(sess.cluster.hosts, lines=200)
|
||||
item.add_report_section("call", "Cluster Logs", json.dumps(logs, indent=2))
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Marker-driven test framework for exo integration tests.
|
||||
|
||||
Test authors declare requirements via markers:
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt='a2a')
|
||||
@pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
|
||||
sharding='tensor', comm='jaccl')
|
||||
def test_jaccl_inference(session):
|
||||
resp = session.chat('What is 2+2?')
|
||||
assert '4' in resp
|
||||
|
||||
The `session` fixture reads the markers, deploys the cluster, places the
|
||||
instance, and provides a `Session` object. All cluster/instance orchestration
|
||||
lives in `exo_tools.harness`; this module is purely the pytest-facing layer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from exo_tools.client import ExoClient
|
||||
from exo_tools.cluster import (
|
||||
Chip,
|
||||
ClusterInfo,
|
||||
EcoSession,
|
||||
Thunderbolt,
|
||||
make_client_from_url,
|
||||
)
|
||||
from exo_tools.harness import Comm, Sharding
|
||||
|
||||
from exo.api.types.api import (
|
||||
ChatCompletionChoice,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
)
|
||||
|
||||
DEFAULT_MODEL = "mlx-community/Llama-3.2-1B-Instruct-4bit"
|
||||
|
||||
|
||||
def _extract_content(resp: ChatCompletionResponse) -> str:
|
||||
"""Extract plain-text content from a non-streaming chat completion."""
|
||||
choice = resp.choices[0]
|
||||
if not isinstance(choice, ChatCompletionChoice):
|
||||
raise RuntimeError(
|
||||
f"Expected non-streaming choice, got {type(choice).__name__}"
|
||||
)
|
||||
content = choice.message.content
|
||||
if not isinstance(content, str):
|
||||
raise RuntimeError(f"Expected string content, got {type(content).__name__}")
|
||||
return content
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClusterSpec:
|
||||
count: int = 1
|
||||
thunderbolt: Thunderbolt | None = None
|
||||
min_memory_gb: float | None = None
|
||||
chip: Chip | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InstanceSpec:
|
||||
model_id: str
|
||||
sharding: Sharding = Sharding.PIPELINE
|
||||
comm: Comm = Comm.RING
|
||||
min_nodes: int = 1
|
||||
|
||||
|
||||
def parse_cluster_marker(marker) -> ClusterSpec:
|
||||
if marker is None:
|
||||
return ClusterSpec()
|
||||
return ClusterSpec(
|
||||
count=marker.kwargs.get("count", 1),
|
||||
thunderbolt=marker.kwargs.get("thunderbolt"),
|
||||
min_memory_gb=marker.kwargs.get("min_memory"),
|
||||
chip=marker.kwargs.get("chip"),
|
||||
)
|
||||
|
||||
|
||||
def parse_instance_marker(marker) -> InstanceSpec | None:
|
||||
if marker is None:
|
||||
return None
|
||||
if not marker.args:
|
||||
raise ValueError(
|
||||
"@pytest.mark.instance requires a positional model_id argument"
|
||||
)
|
||||
return InstanceSpec(
|
||||
model_id=marker.args[0],
|
||||
sharding=marker.kwargs.get("sharding", Sharding.PIPELINE),
|
||||
comm=marker.kwargs.get("comm", Comm.RING),
|
||||
min_nodes=marker.kwargs.get("min_nodes", 1),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
cluster: ClusterInfo
|
||||
eco: EcoSession
|
||||
instance_spec: InstanceSpec | None = None
|
||||
instance_id: str | None = None
|
||||
_stopped_hosts: set[str] = field(default_factory=set)
|
||||
|
||||
@property
|
||||
def client(self) -> ExoClient:
|
||||
for host in self.cluster.hosts:
|
||||
if host not in self._stopped_hosts:
|
||||
return make_client_from_url(self.cluster.api_endpoints[host])
|
||||
return self.cluster.make_client()
|
||||
|
||||
@property
|
||||
def state(self) -> dict[str, Any]:
|
||||
return self.client.request_json("GET", "/state") or {}
|
||||
|
||||
@property
|
||||
def instances(self) -> dict[str, Any]:
|
||||
return self.state.get("instances", {})
|
||||
|
||||
# ---- Inference ----
|
||||
|
||||
def chat(self, prompt: str, max_tokens: int = 100) -> str:
|
||||
resp = self.chat_raw(prompt, max_tokens=max_tokens)
|
||||
return _extract_content(resp)
|
||||
|
||||
def chat_raw(self, prompt: str, **kwargs: Any) -> ChatCompletionResponse:
|
||||
if not self.instance_spec:
|
||||
raise RuntimeError(
|
||||
"No instance placed; add @pytest.mark.instance to the test"
|
||||
)
|
||||
max_tokens = kwargs.pop("max_tokens", 100)
|
||||
request = ChatCompletionRequest.model_validate(
|
||||
{
|
||||
"model": self.instance_spec.model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": max_tokens,
|
||||
**kwargs,
|
||||
}
|
||||
)
|
||||
return self._post_chat(request)
|
||||
|
||||
def multi_turn(self, messages: list[dict[str, str]], max_tokens: int = 100) -> str:
|
||||
if not self.instance_spec:
|
||||
raise RuntimeError(
|
||||
"No instance placed; add @pytest.mark.instance to the test"
|
||||
)
|
||||
request = ChatCompletionRequest.model_validate(
|
||||
{
|
||||
"model": self.instance_spec.model_id,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
)
|
||||
return _extract_content(self._post_chat(request))
|
||||
|
||||
def _post_chat(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
|
||||
raw = self.client.request_json(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
body=request.model_dump(exclude_none=True),
|
||||
)
|
||||
return ChatCompletionResponse.model_validate(raw)
|
||||
|
||||
def disconnect_node(self, index: int) -> None:
|
||||
"""Stop exo on a node and wait for the cluster to observe the disconnect."""
|
||||
host = self.cluster.hosts[index]
|
||||
self.eco.stop([host], keep=True)
|
||||
self._stopped_hosts.add(host)
|
||||
|
||||
def reconnect_node(self, index: int) -> None:
|
||||
"""Restart a previously disconnected node into the existing namespace."""
|
||||
host = self.cluster.hosts[index]
|
||||
self.eco.start_hosts([host], namespace=self.cluster.namespace)
|
||||
self._stopped_hosts.discard(host)
|
||||
|
||||
def wait_ready(
|
||||
self, expected_nodes: int | None = None, timeout: float = 60
|
||||
) -> None:
|
||||
"""Wait until the cluster has exactly `expected_nodes` visible and reporting memory.
|
||||
|
||||
Defaults to the count of non-stopped hosts. Use this after
|
||||
`disconnect_node` / `reconnect_node` to wait for the cluster to settle.
|
||||
"""
|
||||
if expected_nodes is None:
|
||||
expected_nodes = len(self.cluster.hosts) - len(self._stopped_hosts)
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
state = self.state
|
||||
identities = len(state.get("nodeIdentities", {}))
|
||||
memory = len(state.get("nodeMemory", {}))
|
||||
if identities == expected_nodes and memory == expected_nodes:
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2.0)
|
||||
raise TimeoutError(
|
||||
f"Cluster did not reach exactly {expected_nodes} ready nodes within {timeout}s"
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
# type: ignore
|
||||
"""Single-node integration tests.
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/test_1node.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
from exo_tools.harness import is_model_downloaded, place_instance
|
||||
|
||||
from .framework import DEFAULT_MODEL, InstanceSpec
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=1)
|
||||
@pytest.mark.instance(DEFAULT_MODEL)
|
||||
def test_place_instance_and_chat(session):
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=1)
|
||||
@pytest.mark.instance(DEFAULT_MODEL)
|
||||
def test_chat_multiple_turns(session):
|
||||
first_reply = session.chat("What is 2 + 2?")
|
||||
assert len(first_reply) > 0
|
||||
|
||||
second_reply = session.multi_turn(
|
||||
[
|
||||
{"role": "user", "content": "What is 2 + 2?"},
|
||||
{"role": "assistant", "content": first_reply},
|
||||
{"role": "user", "content": "Now multiply that by 3."},
|
||||
]
|
||||
)
|
||||
assert len(second_reply) > 0
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=1)
|
||||
@pytest.mark.instance(DEFAULT_MODEL)
|
||||
def test_delete_instance(session):
|
||||
from exo_tools.harness import wait_for_instance_gone
|
||||
|
||||
session.client.request_json("DELETE", f"/instance/{session.instance_id}")
|
||||
wait_for_instance_gone(session.client, session.instance_id, timeout=30.0)
|
||||
assert len(session.instances) == 0, (
|
||||
f"Expected no instances, found {len(session.instances)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=1)
|
||||
def test_download_from_scratch(session):
|
||||
"""Ensure the model is not on the cluster, then place an instance to
|
||||
trigger a fresh download and verify inference.
|
||||
"""
|
||||
node_id = next(iter(session.state.get("nodeIdentities", {})))
|
||||
|
||||
# Delete any existing download — the API call is idempotent
|
||||
session.client.request_json("DELETE", f"/download/{node_id}/{DEFAULT_MODEL}")
|
||||
|
||||
# Poll until the model is gone (it may already be gone)
|
||||
deadline = time.time() + 60.0
|
||||
while time.time() < deadline:
|
||||
if not is_model_downloaded(session.client, DEFAULT_MODEL):
|
||||
break
|
||||
time.sleep(2.0)
|
||||
else:
|
||||
raise AssertionError(f"Expected {DEFAULT_MODEL} to be deleted from cluster")
|
||||
|
||||
place_instance(session.client, DEFAULT_MODEL, timeout=900.0)
|
||||
session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL)
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
@@ -0,0 +1,49 @@
|
||||
# type: ignore
|
||||
"""Two-node integration tests (ring + jaccl parallelism).
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/test_2node.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from exo_tools.cluster import Thunderbolt
|
||||
from exo_tools.harness import Comm, Sharding
|
||||
|
||||
from .framework import DEFAULT_MODEL
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=2
|
||||
)
|
||||
def test_2node_jaccl(session):
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=2
|
||||
)
|
||||
def test_2node_ring(session):
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=2
|
||||
)
|
||||
def test_2node_jaccl_multi_turn(session):
|
||||
first = session.chat("What is the capital of France?")
|
||||
assert len(first) > 0
|
||||
second = session.multi_turn(
|
||||
[
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
{"role": "assistant", "content": first},
|
||||
{"role": "user", "content": "What country is it in?"},
|
||||
]
|
||||
)
|
||||
assert len(second) > 0
|
||||
@@ -0,0 +1,32 @@
|
||||
# type: ignore
|
||||
"""Four-node integration tests.
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/test_4node.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from exo_tools.cluster import Thunderbolt
|
||||
from exo_tools.harness import Comm, Sharding
|
||||
|
||||
from .framework import DEFAULT_MODEL
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=4, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=4
|
||||
)
|
||||
def test_4node_pipeline_ring(session):
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=4, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=4
|
||||
)
|
||||
def test_4node_tensor_jaccl(session):
|
||||
resp = session.chat("Say hello in one sentence.")
|
||||
assert len(resp) > 0
|
||||
@@ -0,0 +1,102 @@
|
||||
# type: ignore
|
||||
"""Dashboard end-to-end tests using Playwright (headless Chromium).
|
||||
|
||||
Prerequisites:
|
||||
uv run playwright install chromium
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/test_dashboard.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
from playwright.sync_api import sync_playwright
|
||||
|
||||
_HAS_PLAYWRIGHT = True
|
||||
except ImportError:
|
||||
_HAS_PLAYWRIGHT = False
|
||||
|
||||
# Check if Chromium is installed by attempting a quick launch
|
||||
_HAS_CHROMIUM = False
|
||||
if _HAS_PLAYWRIGHT:
|
||||
try:
|
||||
with sync_playwright() as p:
|
||||
browser = p.chromium.launch(headless=True)
|
||||
browser.close()
|
||||
_HAS_CHROMIUM = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not _HAS_PLAYWRIGHT or not _HAS_CHROMIUM,
|
||||
reason="playwright or chromium not installed (run: uv run playwright install chromium)",
|
||||
)
|
||||
|
||||
|
||||
def _mark_onboarding_complete(session) -> None:
|
||||
"""Mark onboarding complete on the server so the wizard doesn't auto-launch a model."""
|
||||
with contextlib.suppress(Exception):
|
||||
session.client.request_json("POST", "/onboarding")
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=1)
|
||||
def test_dashboard_chat_inference(session):
|
||||
"""Full UI flow: open dashboard, pick a model, send a chat, verify response.
|
||||
|
||||
The instance is created via the dashboard UI (model picker → chat send
|
||||
triggers the dashboard's auto-launch flow), not via @pytest.mark.instance.
|
||||
"""
|
||||
_mark_onboarding_complete(session)
|
||||
|
||||
with sync_playwright() as p:
|
||||
browser = p.chromium.launch(headless=True)
|
||||
page = browser.new_page(viewport={"width": 1280, "height": 800})
|
||||
page.goto(session.cluster.api_url, wait_until="networkidle")
|
||||
page.wait_for_timeout(3000)
|
||||
page.screenshot(path="/tmp/dashboard_initial.png")
|
||||
|
||||
# Open the model picker by clicking the "SELECT MODEL" button
|
||||
page.get_by_text("SELECT MODEL", exact=False).first.click()
|
||||
page.wait_for_timeout(1000)
|
||||
page.screenshot(path="/tmp/dashboard_picker_open.png")
|
||||
|
||||
# Search for the model — uses the model id substring; the picker
|
||||
# matches against name/id so "Llama-3.2-1B" filters to the small Llama.
|
||||
search_input = page.locator('input[placeholder*="Search models"]').first
|
||||
search_input.fill("Llama-3.2-1B")
|
||||
page.wait_for_timeout(1500)
|
||||
page.screenshot(path="/tmp/dashboard_picker_search.png")
|
||||
|
||||
# Click the only matching result. The picker shows the model's
|
||||
# display name (e.g. "Llama 3.2 1B") which differs from the model_id.
|
||||
# We click the first visible button-like row in the result list.
|
||||
page.get_by_text("Llama 3.2 1B", exact=False).first.click()
|
||||
page.wait_for_timeout(1500)
|
||||
page.screenshot(path="/tmp/dashboard_model_selected.png")
|
||||
|
||||
# Type a chat message — sending triggers the dashboard's auto-launch
|
||||
# flow: it picks an optimal placement for the selected model and POSTs
|
||||
# to /instance, then sends the chat once the runner is ready.
|
||||
chat_input = page.locator("textarea").first
|
||||
chat_input.fill("Say hello")
|
||||
chat_input.press("Enter")
|
||||
page.screenshot(path="/tmp/dashboard_chat_sent.png")
|
||||
|
||||
# Wait for the instance to launch and respond. Generous timeout
|
||||
# because this includes model placement + load + generation.
|
||||
page.wait_for_timeout(60000)
|
||||
page.screenshot(path="/tmp/dashboard_after_chat.png")
|
||||
|
||||
# Verify an instance was created and the chat got a response
|
||||
instances = session.client.request_json("GET", "/state").get("instances", {})
|
||||
assert len(instances) > 0, "Expected the dashboard to have created an instance"
|
||||
|
||||
body_text = page.text_content("body") or ""
|
||||
assert len(body_text) > 0
|
||||
|
||||
browser.close()
|
||||
@@ -0,0 +1,56 @@
|
||||
# type: ignore
|
||||
"""Resilience tests: disconnect/reconnect nodes and verify cluster recovery.
|
||||
|
||||
Run with:
|
||||
uv run pytest tests/test_resilience.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from exo_tools.cluster import Thunderbolt
|
||||
from exo_tools.harness import Comm, Sharding, cleanup_all_instances, place_instance
|
||||
|
||||
from .framework import DEFAULT_MODEL, InstanceSpec
|
||||
|
||||
|
||||
@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
|
||||
@pytest.mark.instance(
|
||||
DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=2
|
||||
)
|
||||
def test_node_recovery(session):
|
||||
"""Full disconnect/reconnect cycle.
|
||||
|
||||
1. Place a 2-node instance, verify inference
|
||||
2. Disconnect one node
|
||||
3. Place a 1-node instance on remaining node, verify inference
|
||||
4. Reconnect the stopped node, wait for the cluster to reform
|
||||
5. Place a 2-node instance again, verify inference
|
||||
"""
|
||||
# --- Phase 1: 2-node inference ---
|
||||
resp = session.chat("Hello")
|
||||
assert len(resp) > 0
|
||||
|
||||
# --- Phase 2: disconnect one node ---
|
||||
session.disconnect_node(1)
|
||||
session.wait_ready(60)
|
||||
|
||||
# Clean up the now-broken 2-node instance
|
||||
cleanup_all_instances(session.client)
|
||||
|
||||
# --- Phase 3: 1-node inference on the remaining node ---
|
||||
place_instance(session.client, DEFAULT_MODEL, min_nodes=1)
|
||||
session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL, min_nodes=1)
|
||||
resp = session.chat("Hello")
|
||||
assert len(resp) > 0
|
||||
|
||||
# --- Phase 4: reconnect and restore 2-node cluster ---
|
||||
cleanup_all_instances(session.client)
|
||||
session.reconnect_node(1)
|
||||
session.wait_ready(60)
|
||||
|
||||
# --- Phase 5: 2-node inference again ---
|
||||
place_instance(session.client, DEFAULT_MODEL, min_nodes=2)
|
||||
session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL, min_nodes=2)
|
||||
resp = session.chat("Hello again")
|
||||
assert len(resp) > 0
|
||||
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
@@ -0,0 +1,10 @@
|
||||
[project]
|
||||
name = "exo-tools"
|
||||
version = "0.1.0"
|
||||
description = "Shared tooling for interacting with exo clusters"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = ["loguru>=0.7.3"]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
Whitespace-only changes.
@@ -0,0 +1,117 @@
|
||||
# type: ignore
|
||||
"""HTTP client for the exo API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import http.client
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
|
||||
class ExoHttpError(RuntimeError):
|
||||
def __init__(self, status: int, reason: str, body_preview: str):
|
||||
super().__init__(f"HTTP {status} {reason}: {body_preview}")
|
||||
self.status = status
|
||||
|
||||
|
||||
class ExoClient:
|
||||
def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout_s = timeout_s
|
||||
|
||||
def request_json(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
body: dict[str, Any] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> Any:
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
if params:
|
||||
path = path + "?" + urlencode(params)
|
||||
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
payload: bytes | None = None
|
||||
hdrs: dict[str, str] = {"Accept": "application/json"}
|
||||
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
if headers:
|
||||
hdrs.update(headers)
|
||||
|
||||
conn.request(method.upper(), path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
raw = resp.read()
|
||||
text = raw.decode("utf-8", errors="replace") if raw else ""
|
||||
|
||||
if resp.status >= 400:
|
||||
raise ExoHttpError(resp.status, resp.reason, text[:300])
|
||||
|
||||
if not text:
|
||||
return None
|
||||
return json.loads(text)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self.request_json("POST", "/bench/chat/completions", body=payload)
|
||||
|
||||
def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
|
||||
"""POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
|
||||
payload = {**payload, "stream": True}
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
conn.request(
|
||||
"POST",
|
||||
"/bench/chat/completions",
|
||||
body=data,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
)
|
||||
resp = conn.getresponse()
|
||||
if resp.status >= 400:
|
||||
raw = resp.read().decode("utf-8", errors="replace")
|
||||
raise ExoHttpError(resp.status, resp.reason, raw[:300])
|
||||
for line in resp:
|
||||
yield line.decode("utf-8", errors="replace")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_state_path(self, path: str) -> Any:
|
||||
try:
|
||||
return self.request_json("GET", f"/state/{path}")
|
||||
except ExoHttpError as e:
|
||||
if e.status == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
def get_instance(self, instance_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"instances/{instance_id}")
|
||||
|
||||
def get_runner(self, runner_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"runners/{runner_id}")
|
||||
|
||||
def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
|
||||
return self.get_state_path(f"downloads/{node_id}")
|
||||
|
||||
def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeDisk/{node_id}")
|
||||
|
||||
def get_node_system(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeSystem/{node_id}")
|
||||
|
||||
def get_node_identities(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("nodeIdentities")
|
||||
|
||||
def get_topology(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("topology")
|
||||
@@ -0,0 +1,243 @@
|
||||
# type: ignore
|
||||
"""Cluster lifecycle management via eco.
|
||||
|
||||
Provides subprocess wrappers for eco commands (deploy, stop, start, release,
|
||||
logs, exec) and a ClusterInfo dataclass. Reusable by integration tests,
|
||||
bench, eval, and CI workflows.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
from .client import ExoClient
|
||||
|
||||
|
||||
class Thunderbolt(str, Enum):
|
||||
A2A = "a2a" # all-to-all (eco --tb-a2a)
|
||||
RING = "ring" # ring topology (eco --tb-ring)
|
||||
|
||||
|
||||
class Chip(str, Enum):
|
||||
M1 = "M1"
|
||||
M1_PRO = "M1 Pro"
|
||||
M1_MAX = "M1 Max"
|
||||
M1_ULTRA = "M1 Ultra"
|
||||
M2 = "M2"
|
||||
M2_PRO = "M2 Pro"
|
||||
M2_MAX = "M2 Max"
|
||||
M2_ULTRA = "M2 Ultra"
|
||||
M3 = "M3"
|
||||
M3_PRO = "M3 Pro"
|
||||
M3_MAX = "M3 Max"
|
||||
M3_ULTRA = "M3 Ultra"
|
||||
M4 = "M4"
|
||||
M4_PRO = "M4 Pro"
|
||||
M4_MAX = "M4 Max"
|
||||
M4_ULTRA = "M4 Ultra"
|
||||
|
||||
|
||||
logger = logging.getLogger("exo_tools.cluster")
|
||||
|
||||
# When set, deploy from a GitHub branch/tag instead of local source (rsync).
|
||||
_EXO_REF = os.environ.get("EXO_REF")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClusterInfo:
|
||||
"""Holds the result of an `eco start --deploy` invocation."""
|
||||
|
||||
hosts: list[str]
|
||||
namespace: str
|
||||
api_endpoints: dict[str, str] # host -> url
|
||||
api_url: str # primary endpoint for ExoClient
|
||||
|
||||
primary_host: str = ""
|
||||
_host: str = field(init=False, repr=False, default="")
|
||||
_port: int = field(init=False, repr=False, default=52415)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not self.primary_host:
|
||||
self.primary_host = self.hosts[0]
|
||||
url = self.api_url.replace("http://", "").replace("https://", "")
|
||||
parts = url.split(":")
|
||||
self._host = parts[0]
|
||||
self._port = int(parts[1]) if len(parts) > 1 else 52415
|
||||
|
||||
def make_client(self, timeout_s: float = 7200.0) -> ExoClient:
|
||||
return ExoClient(self._host, self._port, timeout_s=timeout_s)
|
||||
|
||||
|
||||
class EcoSession:
|
||||
"""Manages an eco session with a unique user and automatic cleanup.
|
||||
|
||||
Usage:
|
||||
session = EcoSession(user_prefix="test")
|
||||
cluster = session.start_deploy(count=2, thunderbolt=True)
|
||||
...
|
||||
session.stop_all() # or let atexit handle it
|
||||
|
||||
The session registers atexit and signal handlers to ensure cleanup
|
||||
on normal exit, uncaught exceptions, SIGTERM, and SIGHUP. SIGINT
|
||||
is left unhandled so KeyboardInterrupt propagates normally.
|
||||
"""
|
||||
|
||||
def __init__(self, user_prefix: str = "test") -> None:
|
||||
self._session_id = uuid.uuid4().hex[:8]
|
||||
self.user = f"{user_prefix}-{self._session_id}"
|
||||
self._env = {**os.environ, "USER": self.user}
|
||||
|
||||
# Register cleanup handlers
|
||||
atexit.register(self.stop_all)
|
||||
for sig in (signal.SIGTERM, signal.SIGHUP):
|
||||
signal.signal(sig, self._signal_handler)
|
||||
|
||||
def _signal_handler(self, signum: int, _frame: object) -> None:
|
||||
self.stop_all()
|
||||
raise SystemExit(128 + signum)
|
||||
|
||||
def stop_all(self) -> None:
|
||||
"""Stop all clusters and release all reservations for this session."""
|
||||
with contextlib.suppress(Exception):
|
||||
subprocess.run(
|
||||
["eco", "stop"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
env=self._env,
|
||||
)
|
||||
|
||||
def _run(
|
||||
self, args: list[str], *, check: bool = True, timeout: int = 120
|
||||
) -> subprocess.CompletedProcess[str]:
|
||||
"""Run an eco command as this session's user.
|
||||
|
||||
stdout is captured (JSON output), stderr is passed through to the
|
||||
console so eco's progress messages are visible.
|
||||
"""
|
||||
logger.info(f"eco: {' '.join(args)}")
|
||||
return subprocess.run(
|
||||
args,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=None,
|
||||
text=True,
|
||||
check=check,
|
||||
timeout=timeout,
|
||||
env=self._env,
|
||||
)
|
||||
|
||||
def start_deploy(
|
||||
self,
|
||||
hosts: list[str] | None = None,
|
||||
*,
|
||||
count: int | None = None,
|
||||
thunderbolt: Thunderbolt | None = None,
|
||||
chip: Chip | None = None,
|
||||
min_memory_gb: float | None = None,
|
||||
wait: bool = True,
|
||||
ref: str | None = _EXO_REF,
|
||||
timeout: int = 600,
|
||||
) -> ClusterInfo:
|
||||
"""Start and deploy exo on a set of hosts via eco.
|
||||
|
||||
By default, deploys from local source via rsync. Set EXO_REF
|
||||
or pass ref= to deploy from a GitHub branch/tag instead (for CI).
|
||||
"""
|
||||
cmd: list[str] = ["eco", "--json", "start", "--deploy"]
|
||||
if hosts:
|
||||
cmd.extend(hosts)
|
||||
if count is not None:
|
||||
cmd.extend(["--count", str(count)])
|
||||
if thunderbolt is not None:
|
||||
cmd.append(f"--tb-{thunderbolt.value}")
|
||||
if chip is not None:
|
||||
cmd.extend(["--chip", chip.value])
|
||||
if min_memory_gb is not None:
|
||||
cmd.extend(["--min-memory", str(min_memory_gb)])
|
||||
if wait:
|
||||
cmd.append("--wait")
|
||||
if ref:
|
||||
cmd.extend(["--ref", ref])
|
||||
|
||||
result = self._run(cmd, timeout=timeout)
|
||||
data = json.loads(result.stdout)["data"]
|
||||
endpoints: dict[str, str] = data["api_endpoints"]
|
||||
primary_host = data["hosts"][0]
|
||||
|
||||
return ClusterInfo(
|
||||
hosts=data["hosts"],
|
||||
namespace=data["namespace"],
|
||||
api_endpoints=endpoints,
|
||||
api_url=endpoints[primary_host],
|
||||
primary_host=primary_host,
|
||||
)
|
||||
|
||||
def stop(self, hosts: list[str], *, keep: bool = False, timeout: int = 120) -> None:
|
||||
"""Stop exo on the given hosts. If keep=True, keep the reservation."""
|
||||
cmd: list[str] = ["eco", "stop"]
|
||||
cmd.extend(hosts)
|
||||
if keep:
|
||||
cmd.append("--keep")
|
||||
self._run(cmd, timeout=timeout)
|
||||
|
||||
def start_hosts(
|
||||
self, hosts: list[str], *, namespace: str, timeout: int = 300
|
||||
) -> None:
|
||||
"""Start (previously stopped) hosts back into an existing namespace."""
|
||||
cmd: list[str] = ["eco", "--json", "start"]
|
||||
cmd.extend(hosts)
|
||||
cmd.extend(["--namespace", namespace])
|
||||
self._run(cmd, timeout=timeout)
|
||||
|
||||
def release(self, hosts: list[str], timeout: int = 120) -> None:
|
||||
"""Release hosts from the reservation."""
|
||||
cmd: list[str] = ["eco", "release"]
|
||||
cmd.extend(hosts)
|
||||
self._run(cmd, timeout=timeout)
|
||||
|
||||
def logs(
|
||||
self, hosts: list[str], lines: int = 500, timeout: int = 60
|
||||
) -> dict[str, list[str]]:
|
||||
"""Fetch recent logs from cluster hosts."""
|
||||
cmd: list[str] = ["eco", "--json", "logs"]
|
||||
cmd.extend(hosts)
|
||||
cmd.extend(["-n", str(lines), "--raw"])
|
||||
result = self._run(cmd, check=False, timeout=timeout)
|
||||
if result.returncode != 0:
|
||||
return {"_error": [result.stderr]}
|
||||
try:
|
||||
return json.loads(result.stdout)
|
||||
except json.JSONDecodeError:
|
||||
return {"_raw": result.stdout.splitlines()}
|
||||
|
||||
def exec(self, hosts: list[str], command: str, timeout: int = 120) -> str:
|
||||
"""Run an arbitrary command on the given hosts via eco."""
|
||||
cmd: list[str] = ["eco", "exec"]
|
||||
cmd.extend(hosts)
|
||||
cmd.append("--")
|
||||
cmd.extend(command.split())
|
||||
result = self._run(cmd, check=False, timeout=timeout)
|
||||
return result.stdout
|
||||
|
||||
|
||||
def make_client(cluster: ClusterInfo, timeout_s: float = 7200.0) -> ExoClient:
|
||||
"""Create an ExoClient from a ClusterInfo."""
|
||||
return cluster.make_client(timeout_s=timeout_s)
|
||||
|
||||
|
||||
def make_client_from_url(url: str, timeout_s: float = 7200.0) -> ExoClient:
|
||||
"""Create an ExoClient from a URL string like 'http://host:port'."""
|
||||
url_clean = url.replace("http://", "").replace("https://", "")
|
||||
parts = url_clean.split(":")
|
||||
host = parts[0]
|
||||
port = int(parts[1]) if len(parts) > 1 else 52415
|
||||
return ExoClient(host, port, timeout_s=timeout_s)
|
||||
@@ -1,129 +1,39 @@
|
||||
# type: ignore
|
||||
"""Instance lifecycle helpers for exo clusters.
|
||||
|
||||
Provides utilities for placing instances, waiting for readiness,
|
||||
managing downloads, filtering placements, and common CLI arguments.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import http.client
|
||||
import json
|
||||
import contextlib
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from .client import ExoClient, ExoHttpError
|
||||
|
||||
|
||||
class Sharding(str, Enum):
|
||||
PIPELINE = "Pipeline" # layers split across nodes
|
||||
TENSOR = "Tensor" # layers split within (across nodes)
|
||||
|
||||
|
||||
class Comm(str, Enum):
|
||||
RING = "MlxRing" # ring all-reduce over network
|
||||
JACCL = "MlxJaccl" # RDMA over Thunderbolt
|
||||
|
||||
|
||||
_SETTLE_INITIAL_BACKOFF_S = 1.0
|
||||
_SETTLE_MAX_BACKOFF_S = 60.0
|
||||
_SETTLE_BACKOFF_MULTIPLIER = 2.0
|
||||
|
||||
|
||||
class ExoHttpError(RuntimeError):
|
||||
def __init__(self, status: int, reason: str, body_preview: str):
|
||||
super().__init__(f"HTTP {status} {reason}: {body_preview}")
|
||||
self.status = status
|
||||
|
||||
|
||||
class ExoClient:
|
||||
def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout_s = timeout_s
|
||||
|
||||
def request_json(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
body: dict[str, Any] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> Any:
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
if params:
|
||||
path = path + "?" + urlencode(params)
|
||||
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
payload: bytes | None = None
|
||||
hdrs: dict[str, str] = {"Accept": "application/json"}
|
||||
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
if headers:
|
||||
hdrs.update(headers)
|
||||
|
||||
conn.request(method.upper(), path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
raw = resp.read()
|
||||
text = raw.decode("utf-8", errors="replace") if raw else ""
|
||||
|
||||
if resp.status >= 400:
|
||||
raise ExoHttpError(resp.status, resp.reason, text[:300])
|
||||
|
||||
if not text:
|
||||
return None
|
||||
return json.loads(text)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self.request_json("POST", "/bench/chat/completions", body=payload)
|
||||
|
||||
def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
|
||||
"""POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
|
||||
payload = {**payload, "stream": True}
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
conn.request(
|
||||
"POST",
|
||||
"/bench/chat/completions",
|
||||
body=data,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
)
|
||||
resp = conn.getresponse()
|
||||
if resp.status >= 400:
|
||||
raw = resp.read().decode("utf-8", errors="replace")
|
||||
raise ExoHttpError(resp.status, resp.reason, raw[:300])
|
||||
for line in resp:
|
||||
yield line.decode("utf-8", errors="replace")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_state_path(self, path: str) -> Any:
|
||||
try:
|
||||
return self.request_json("GET", f"/state/{path}")
|
||||
except ExoHttpError as e:
|
||||
if e.status == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
def get_instance(self, instance_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"instances/{instance_id}")
|
||||
|
||||
def get_runner(self, runner_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"runners/{runner_id}")
|
||||
|
||||
def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
|
||||
return self.get_state_path(f"downloads/{node_id}")
|
||||
|
||||
def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeDisk/{node_id}")
|
||||
|
||||
def get_node_system(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeSystem/{node_id}")
|
||||
|
||||
def get_node_identities(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("nodeIdentities")
|
||||
|
||||
def get_topology(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("topology")
|
||||
|
||||
|
||||
def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
|
||||
if len(instance) != 1:
|
||||
raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
|
||||
@@ -555,7 +465,6 @@ def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
|
||||
except Exception:
|
||||
return None
|
||||
for inst_id, inst in state.get("instances", {}).items():
|
||||
# Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
|
||||
for _inst_type, inner in inst.items():
|
||||
if not isinstance(inner, dict):
|
||||
continue
|
||||
@@ -623,3 +532,112 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
|
||||
action="store_true",
|
||||
help="Reuse an existing running instance for this model instead of creating a new one.",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cluster/instance orchestration helpers (used by tests, bench, eval)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def get_instance_ids(client: ExoClient) -> set[str]:
|
||||
"""Return the set of current instance IDs from cluster state."""
|
||||
state = client.request_json("GET", "/state") or {}
|
||||
result: set[str] = set()
|
||||
for instance in state.get("instances", {}).values():
|
||||
with contextlib.suppress(Exception):
|
||||
result.add(instance_id_from_instance(instance))
|
||||
return result
|
||||
|
||||
|
||||
def wait_for_cluster_ready(
|
||||
client: ExoClient, expected_nodes: int = 1, timeout: float = 120.0
|
||||
) -> None:
|
||||
"""Wait until the cluster has all expected nodes visible and reporting memory.
|
||||
|
||||
Placement requires nodeMemory for all nodes in a cycle. This polls until
|
||||
both nodeIdentities and nodeMemory have at least `expected_nodes` entries.
|
||||
"""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
state = client.request_json("GET", "/state") or {}
|
||||
if (
|
||||
len(state.get("nodeIdentities", {})) >= expected_nodes
|
||||
and len(state.get("nodeMemory", {})) >= expected_nodes
|
||||
):
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(1.0)
|
||||
raise TimeoutError(f"Cluster not ready: expected {expected_nodes} nodes")
|
||||
|
||||
|
||||
def place_instance(
|
||||
client: ExoClient,
|
||||
model_id: str,
|
||||
*,
|
||||
sharding: Sharding = Sharding.PIPELINE,
|
||||
comm: Comm = Comm.RING,
|
||||
min_nodes: int = 1,
|
||||
timeout: float = 600.0,
|
||||
placement_retries: int = 10,
|
||||
placement_retry_delay: float = 10.0,
|
||||
) -> str:
|
||||
"""Place an instance and wait for it to be ready. Returns the instance_id.
|
||||
|
||||
The /place_instance API returns a command_id, but instances are stored
|
||||
under a separately-generated instance_id. This polls cluster state for the
|
||||
new instance, retrying placement if the cluster is still settling.
|
||||
"""
|
||||
wait_for_cluster_ready(client, expected_nodes=min_nodes)
|
||||
|
||||
body = {
|
||||
"model_id": model_id,
|
||||
"sharding": sharding.value,
|
||||
"instance_meta": comm.value,
|
||||
"min_nodes": min_nodes,
|
||||
}
|
||||
|
||||
instance_id: str | None = None
|
||||
for attempt in range(placement_retries):
|
||||
before_ids = get_instance_ids(client)
|
||||
client.request_json("POST", "/place_instance", body=body)
|
||||
|
||||
poll_deadline = time.time() + 30.0
|
||||
while time.time() < poll_deadline:
|
||||
new_ids = get_instance_ids(client) - before_ids
|
||||
if new_ids:
|
||||
instance_id = next(iter(new_ids))
|
||||
break
|
||||
time.sleep(1.0)
|
||||
|
||||
if instance_id is not None:
|
||||
break
|
||||
|
||||
if attempt < placement_retries - 1:
|
||||
time.sleep(placement_retry_delay)
|
||||
|
||||
if instance_id is None:
|
||||
raise TimeoutError(
|
||||
f"Placement failed after {placement_retries} attempts "
|
||||
f"({sharding.value}/{comm.value} for {model_id})"
|
||||
)
|
||||
|
||||
wait_for_instance_ready(client, instance_id, timeout=timeout)
|
||||
return instance_id
|
||||
|
||||
|
||||
def cleanup_all_instances(client: ExoClient) -> None:
|
||||
"""Remove all running instances from the cluster."""
|
||||
state = client.request_json("GET", "/state") or {}
|
||||
for instance in state.get("instances", {}).values():
|
||||
with contextlib.suppress(Exception):
|
||||
iid = instance_id_from_instance(instance)
|
||||
client.request_json("DELETE", f"/instance/{iid}")
|
||||
wait_for_instance_gone(client, iid, timeout=30.0)
|
||||
|
||||
|
||||
def is_model_downloaded(client: ExoClient, model_id: str) -> bool:
|
||||
response = client.request_json("GET", "/models", params={"status": "downloaded"})
|
||||
data = (response or {}).get("data", [])
|
||||
return all(model.get("id") == model_id for model in data)
|
||||
@@ -23,10 +23,11 @@ members = [
|
||||
"exo",
|
||||
"exo-bench",
|
||||
"exo-pyo3-bindings",
|
||||
"exo-tools",
|
||||
]
|
||||
constraints = [{ name = "transformers", specifier = ">=5.6.2" }]
|
||||
overrides = [
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops" },
|
||||
{ name = "mlx", marker = "sys_platform == 'linux'", specifier = "==0.31.1" },
|
||||
]
|
||||
|
||||
@@ -394,7 +395,7 @@ dependencies = [
|
||||
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-vlm", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
@@ -416,21 +417,21 @@ build = [
|
||||
]
|
||||
cpu = [
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cpu') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-cpu", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
|
||||
]
|
||||
cuda12 = [
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-cuda-12", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
|
||||
]
|
||||
cuda13 = [
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "(sys_platform == 'darwin' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-cuda-13", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'linux'" },
|
||||
{ name = "mlx-vlm", marker = "sys_platform == 'linux'" },
|
||||
@@ -439,6 +440,7 @@ cuda13 = [
|
||||
[package.dev-dependencies]
|
||||
dev = [
|
||||
{ name = "basedpyright", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "playwright", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "pyinstaller", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "pytest-asyncio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
@@ -459,7 +461,7 @@ requires-dist = [
|
||||
{ name = "hypercorn", specifier = ">=0.18.0" },
|
||||
{ name = "loguru", specifier = ">=0.7.3" },
|
||||
{ name = "mflux", marker = "sys_platform == 'darwin'", specifier = "==0.17.2" },
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
|
||||
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops" },
|
||||
{ name = "mlx", marker = "sys_platform == 'linux' and extra == 'cpu'", specifier = "==0.31.1" },
|
||||
{ name = "mlx", marker = "sys_platform == 'linux' and extra == 'cuda12'", specifier = "==0.31.1" },
|
||||
{ name = "mlx", marker = "sys_platform == 'linux' and extra == 'cuda13'", specifier = "==0.31.1" },
|
||||
@@ -498,6 +500,7 @@ provides-extras = ["build", "cpu", "cuda12", "cuda13"]
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
{ name = "basedpyright", specifier = ">=1.29.0" },
|
||||
{ name = "playwright", specifier = ">=1.52.0" },
|
||||
{ name = "pyinstaller", specifier = ">=6.17.0" },
|
||||
{ name = "pytest", specifier = ">=8.4.0" },
|
||||
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
|
||||
@@ -561,6 +564,17 @@ dev = [
|
||||
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "exo-tools"
|
||||
version = "0.1.0"
|
||||
source = { editable = "tools" }
|
||||
dependencies = [
|
||||
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [{ name = "loguru", specifier = ">=0.7.3" }]
|
||||
|
||||
[[package]]
|
||||
name = "fastapi"
|
||||
version = "0.128.0"
|
||||
@@ -669,6 +683,24 @@ http = [
|
||||
{ name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "greenlet"
|
||||
version = "3.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/3c/3f/dbf99fb14bfeb88c28f16729215478c0e265cacd6dc22270c8f31bb6892f/greenlet-3.5.0.tar.gz", hash = "sha256:d419647372241bc68e957bf38d5c1f98852155e4146bd1e4121adea81f4f01e4", size = 196995, upload-time = "2026-04-27T13:37:15.544Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/58/fc576f99037ce19c5aa16628e4c3226b6d1419f72a62c79f5f40576e6eb3/greenlet-3.5.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:5a5ed18de6a0f6cc7087f1563f6bd93fc7df1c19165ca01e9bde5a5dc281d106", size = 285066, upload-time = "2026-04-27T12:23:05.033Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4a/ba/b28ddbe6bfad6a8ac196ef0e8cff37bc65b79735995b9e410923fffeeb70/greenlet-3.5.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a717fbc46d8a354fa675f7c1e813485b6ba3885f9bef0cd56e5ba27d758ff5b", size = 604414, upload-time = "2026-04-27T12:52:42.358Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/09/06/4b69f8f0b67603a8be2790e55107a190b376f2627fe0eaf5695d85ffb3cd/greenlet-3.5.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ddc090c5c1792b10246a78e8c2163ebbe04cf877f9d785c230a7b27b39ad038e", size = 617349, upload-time = "2026-04-27T12:59:43.32Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6a/15/a643b4ecd09969e30b8a150d5919960caae0abe4f5af75ab040b1ab85e78/greenlet-3.5.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4964101b8585c144cbda5532b1aa644255126c08a265dae90c16e7a0e63aaa9d", size = 623234, upload-time = "2026-04-27T13:02:40.611Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8a/17/a3918541fd0ddefe024a69de6d16aa7b46d36ac19562adaa63c7fa180eff/greenlet-3.5.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2094acd54b272cb6eae8c03dd87b3fa1820a4cef18d6889c378d503500a1dc13", size = 613927, upload-time = "2026-04-27T12:25:30.28Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/77/18/3b13d5ef1275b0ffaf933b05efa21408ac4ca95823c7411d79682e4fdcff/greenlet-3.5.0-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:7022615368890680e67b9965d33f5773aade330d5343bbe25560135aaa849eae", size = 425243, upload-time = "2026-04-27T13:05:15.689Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ee/e1/bd0af6213c7dd33175d8a462d4c1fe1175124ebed4855bc1475a5b5242c2/greenlet-3.5.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5e05ba267789ea87b5a155cf0e810b1ab88bf18e9e8740813945ceb8ee4350ba", size = 1570893, upload-time = "2026-04-27T12:53:29.483Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9b/2a/0789702f864f5382cb476b93d7a9c823c10472658102ccd65f415747d2e2/greenlet-3.5.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:0ecec963079cd58cbd14723582384f11f166fd58883c15dcbfb342e0bc9b5846", size = 1636060, upload-time = "2026-04-27T12:25:28.845Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b2/8f/22bf9df92bbff0eb07842b60f7e63bf7675a9742df628437a9f02d09137f/greenlet-3.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:728d9667d8f2f586644b748dbd9bb67e50d6a9381767d1357714ea6825bb3bf5", size = 238740, upload-time = "2026-04-27T12:24:01.341Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b6/b7/9c5c3d653bd4ff614277c049ac676422e2c557db47b4fe43e6313fc005dc/greenlet-3.5.0-cp313-cp313-win_arm64.whl", hash = "sha256:47422135b1d308c14b2c6e758beedb1acd33bb91679f5670edf77bf46244722b", size = 235525, upload-time = "2026-04-27T12:23:12.308Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h11"
|
||||
version = "0.16.0"
|
||||
@@ -1213,7 +1245,7 @@ dependencies = [
|
||||
{ name = "hf-transfer", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "matplotlib", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "numpy", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "opencv-python", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "piexif", marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
@@ -1263,8 +1295,8 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "mlx"
|
||||
version = "0.32.0.dev20260427+cc3f3e60"
|
||||
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }
|
||||
version = "0.32.0.dev20260508+df6a7891"
|
||||
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }
|
||||
resolution-markers = [
|
||||
"sys_platform == 'darwin'",
|
||||
]
|
||||
@@ -1315,7 +1347,7 @@ source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Fdeepseek-v4
|
||||
dependencies = [
|
||||
{ name = "jinja2", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "numpy", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "protobuf", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "pyyaml", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
@@ -1332,7 +1364,7 @@ dependencies = [
|
||||
{ name = "fastapi", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "miniaudio", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.31.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260427+cc3f3e60", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#cc3f3e60be1289506125f2fa19b73b05aa770df8" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx", version = "0.32.0.dev20260508+df6a7891", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Forder-distributed-ops#df6a7891226d3262b88fe21a918e4223b552b819" }, marker = "sys_platform == 'darwin' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "numpy", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "opencv-python", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and extra == 'extra-3-exo-cpu') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda12') or (sys_platform == 'linux' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
@@ -1768,6 +1800,25 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cb/28/3bfe2fa5a7b9c46fe7e13c97bda14c895fb10fa2ebf1d0abb90e0cea7ee1/platformdirs-4.5.1-py3-none-any.whl", hash = "sha256:d03afa3963c806a9bed9d5125c8f4cb2fdaf74a55ab60e5d59b3fde758104d31", size = 18731, upload-time = "2025-12-05T13:52:56.823Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "playwright"
|
||||
version = "1.58.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "greenlet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
{ name = "pyee", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/c9/9c6061d5703267f1baae6a4647bfd1862e386fbfdb97d889f6f6ae9e3f64/playwright-1.58.0-py3-none-macosx_10_13_x86_64.whl", hash = "sha256:96e3204aac292ee639edbfdef6298b4be2ea0a55a16b7068df91adac077cc606", size = 42251098, upload-time = "2026-01-30T15:09:24.028Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e0/40/59d34a756e02f8c670f0fee987d46f7ee53d05447d43cd114ca015cb168c/playwright-1.58.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:70c763694739d28df71ed578b9c8202bb83e8fe8fb9268c04dd13afe36301f71", size = 41039625, upload-time = "2026-01-30T15:09:27.558Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e1/ee/3ce6209c9c74a650aac9028c621f357a34ea5cd4d950700f8e2c4b7fe2c4/playwright-1.58.0-py3-none-macosx_11_0_universal2.whl", hash = "sha256:185e0132578733d02802dfddfbbc35f42be23a45ff49ccae5081f25952238117", size = 42251098, upload-time = "2026-01-30T15:09:30.461Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f1/af/009958cbf23fac551a940d34e3206e6c7eed2b8c940d0c3afd1feb0b0589/playwright-1.58.0-py3-none-manylinux1_x86_64.whl", hash = "sha256:c95568ba1eda83812598c1dc9be60b4406dffd60b149bc1536180ad108723d6b", size = 46235268, upload-time = "2026-01-30T15:09:33.787Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d9/a6/0e66ad04b6d3440dae73efb39540c5685c5fc95b17c8b29340b62abbd952/playwright-1.58.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f9999948f1ab541d98812de25e3a8c410776aa516d948807140aff797b4bffa", size = 45964214, upload-time = "2026-01-30T15:09:36.751Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0e/4b/236e60ab9f6d62ed0fd32150d61f1f494cefbf02304c0061e78ed80c1c32/playwright-1.58.0-py3-none-win32.whl", hash = "sha256:1e03be090e75a0fabbdaeab65ce17c308c425d879fa48bb1d7986f96bfad0b99", size = 36815998, upload-time = "2026-01-30T15:09:39.627Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/41/f8/5ec599c5e59d2f2f336a05b4f318e733077cd5044f24adb6f86900c3e6a7/playwright-1.58.0-py3-none-win_amd64.whl", hash = "sha256:a2bf639d0ce33b3ba38de777e08697b0d8f3dc07ab6802e4ac53fb65e3907af8", size = 36816005, upload-time = "2026-01-30T15:09:42.449Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/c4/cc0229fea55c87d6c9c67fe44a21e2cd28d1d558a5478ed4d617e9fb0c93/playwright-1.58.0-py3-none-win_arm64.whl", hash = "sha256:32ffe5c303901a13a0ecab91d1c3f74baf73b84f4bedbb6b935f5bc11cc98e1b", size = 33085919, upload-time = "2026-01-30T15:09:45.71Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
@@ -1941,6 +1992,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/73/7d/f2f9db34af103bea3e09735bb40b021788a5e834c81eedb541991badf8f5/pydantic_core-2.41.5-cp313-cp313-win_arm64.whl", hash = "sha256:3f84d5c1b4ab906093bdc1ff10484838aca54ef08de4afa9de0f5f14d69639cd", size = 1981005, upload-time = "2025-11-04T13:40:54.734Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyee"
|
||||
version = "13.0.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda12') or (extra == 'extra-3-exo-cpu' and extra == 'extra-3-exo-cuda13') or (extra == 'extra-3-exo-cuda12' and extra == 'extra-3-exo-cuda13')" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8b/04/e7c1fe4dc78a6fdbfd6c337b1c3732ff543b8a397683ab38378447baa331/pyee-13.0.1.tar.gz", hash = "sha256:0b931f7c14535667ed4c7e0d531716368715e860b988770fc7eb8578d1f67fc8", size = 31655, upload-time = "2026-02-14T21:12:28.044Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/c4/b4d4827c93ef43c01f599ef31453ccc1c132b353284fc6c87d535c233129/pyee-13.0.1-py3-none-any.whl", hash = "sha256:af2f8fede4171ef667dfded53f96e2ed0d6e6bd7ee3bb46437f77e3b57689228", size = 15659, upload-time = "2026-02-14T21:12:26.263Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pygments"
|
||||
version = "2.19.2"
|
||||
|
||||
Reference in new issue
Block a user