Skip to content

fix(ssd): reuse only finished prefetches; SIGTERM mid-stream test; doc fixes #198

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
solderzzc merged 7 commits into from
Sep 28, 2026
Merged

fix(ssd): reuse only finished prefetches; SIGTERM mid-stream test; doc fixes #198

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0a66dbb
fix(ssd): reuse only finished prefetches; directory load after the Mo…
solderzzc Sep 26, 2026
0f27cbc
docs+test: emitEvent doc placement, SIGTERM mid-stream test
solderzzc Sep 27, 2026
3284dc5
fix(ssd): review fixes for #198
solderzzc Sep 28, 2026
4547bce
fix(ssd): decide local-copy completeness from the Hub's file list
solderzzc Sep 28, 2026
f96e756
fix(ssd): use loadable local copies first; consult the Hub only to do…
solderzzc Sep 28, 2026
cdb16e6
fix(ssd): judge local copies the way the loader reads them
solderzzc Sep 28, 2026
ca32d58
fix(ssd): a missing indexed file in a subdirectory means incomplete
solderzzc Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ Every needle check passed. `--mtp` with the bf16 assistant (`gemma-4-26B-A4B-it-
| ~9.8K | 858 / 43.4 | 20.4 GB | 401 / 12.7 | 5.6 GB |
| 40.8K | 615 / 36.1 | 21.5 GB | 336 / 12.0 | 5.8 GB |

> ⚠️ **`--stream-experts` crashes on quantized MoE models in releases b769 and b773** (`broadcast_shapes … (N,8,8,D)` on the first request). Reproduced on M6 with Qwen3.6-35B-A3B and Gemma 4 26B-A4B. The mlx-swift-lm upstream sync in #167 broke the SSD path. Earlier versions of this table were measured before that sync and were never re-checked afterwards. Fixed in SharpAI/mlx-swift-lm#69 and #71; the table above was re-measured with those fixes.
> ⚠️ **`--stream-experts` crashes on quantized MoE models in releases b769 and b773** (a `broadcast_shapes` fatal error on the first request, e.g. `(263,8,8,2048) and (263,8,1)` with top-k 8). Reproduced on M6 with Qwen3.6-35B-A3B and Gemma 4 26B-A4B. The mlx-swift-lm upstream sync in #167 broke the SSD path. Earlier versions of this table were measured before that sync and were never re-checked afterwards. Fixed in SharpAI/mlx-swift-lm#69 and #71; the table above was re-measured with those fixes.

### Qwen3.8-27B-4bit (dense)

Expand Down
173 changes: 90 additions & 83 deletions Sources/SwiftLM/Server.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ final class ProgressTracker {
private var lastUpdate: TimeInterval = 0
private var lastBytes: Int64 = 0
private var speedStr = "0.0 MB/s"
private var progress: Progress?

init(modelId: String) {
self.modelId = modelId
Expand All @@ -291,6 +292,11 @@ final class ProgressTracker {
isDone = true
trackingTask?.cancel()
if Self.barOpen {
// The Task may be cancelled before its last frame; show where it ended.
if let progress {
print(frame(fraction: progress.fractionCompleted).padding(
toLength: 100, withPad: " ", startingAt: 0), terminator: "")
}
print("")
fflush(stdout)
Self.barOpen = false
Expand Down Expand Up @@ -320,7 +326,37 @@ final class ProgressTracker {
return sumDir(modelHubDir) + sumDir(downloadDir)
}

/// One `\r` frame of the download bar at `fraction`.
private func frame(fraction: Double) -> String {
let pct = Int(fraction * 100)
var completedMB = String(format: "%.1f", Double(self.lastBytes) / 1_048_576)
var totalMB = "???"
if fraction > 0.001 {
let extrapolated = (Double(self.lastBytes) / fraction) / 1_048_576.0
totalMB = String(format: "%.1f", extrapolated)
} else if fraction == 0.0 {
completedMB = "0.0"
}

let barLength = 20
let completedBars = min(barLength, Int(fraction * Double(barLength)))
let emptyBars = max(0, barLength - completedBars)

var bars = ""
if completedBars > 0 {
bars += String(repeating: "=", count: completedBars - 1) + ">"
}
bars += String(repeating: " ", count: emptyBars)

let pctStr = String(format: "%3d%%", pct)
let spinner = self.spinnerFrames[self.frameIndex]
let speedText = "| Speed: \(self.speedStr)"

return String(format: "\r[SwiftLM] Download: [%@] %@ %@ (%@ MB / %@ MB) %@", bars, pctStr, spinner, completedMB, totalMB, speedText)
}

func printProgress(_ progress: Progress) {
self.progress = progress
if trackingTask == nil {
lastUpdate = Date().timeIntervalSince1970
lastBytes = getDownloadedBytes()
Expand All @@ -329,7 +365,6 @@ final class ProgressTracker {
while !self.isDone && !Task.isCancelled {
let now = Date().timeIntervalSince1970
let fraction = progress.fractionCompleted
let pct = Int(fraction * 100)

let interval = now - self.lastUpdate
if interval >= 0.25 {
Expand All @@ -348,30 +383,7 @@ final class ProgressTracker {
self.lastUpdate = now
}

var completedMB = String(format: "%.1f", Double(self.lastBytes) / 1_048_576)
var totalMB = "???"
if fraction > 0.001 {
let extrapolated = (Double(self.lastBytes) / fraction) / 1_048_576.0
totalMB = String(format: "%.1f", extrapolated)
} else if fraction == 0.0 {
completedMB = "0.0"
}

let barLength = 20
let completedBars = min(barLength, Int(fraction * Double(barLength)))
let emptyBars = max(0, barLength - completedBars)

var bars = ""
if completedBars > 0 {
bars += String(repeating: "=", count: completedBars - 1) + ">"
}
bars += String(repeating: " ", count: emptyBars)

let pctStr = String(format: "%3d%%", pct)
let spinner = self.spinnerFrames[self.frameIndex]
let speedText = "| Speed: \(self.speedStr)"

let msg = String(format: "\r[SwiftLM] Download: [%@] %@ %@ (%@ MB / %@ MB) %@", bars, pctStr, spinner, completedMB, totalMB, speedText)
let msg = self.frame(fraction: fraction)

if self.isDone { break }
print(msg.padding(toLength: 100, withPad: " ", startingAt: 0), terminator: "")
Expand Down Expand Up @@ -408,6 +420,26 @@ final class ProgressTracker {
/// for a protocol-critical signal the daemon is meant to positively rely
/// on (not just infer), a dropped emit with zero diagnostic trail would be
/// hard to ever notice. Logs to stderr on failure instead.
/// `preface` lines are written in the same `print` as the JSON, so no generation
/// token can land between them and push the JSON off the start of its line.
func emitEvent(_ payload: [String: Any], preface: String? = nil) {
guard let data = try? JSONSerialization.data(withJSONObject: payload),
let json = String(data: data, encoding: .utf8)
else {
FileHandle.standardError.write(
Data("[SwiftLM] failed to encode event for stdout: \(payload)\n".utf8))
return
}
var out = ""
if ProgressTracker.barOpen {
out += "\n"
ProgressTracker.barOpen = false
}
if let preface { out += preface + "\n" }
print(out + json)
fflush(stdout)
}

/// Ends the process after a requested shutdown (SIGTERM/SIGINT) without running
/// C++ static destructors or `atexit` handlers.
///
Expand All @@ -423,22 +455,6 @@ func exitAfterShutdownRequest() -> Never {
Darwin._exit(0)
}

func emitEvent(_ payload: [String: Any]) {
guard let data = try? JSONSerialization.data(withJSONObject: payload),
let json = String(data: data, encoding: .utf8)
else {
FileHandle.standardError.write(
Data("[SwiftLM] failed to encode event for stdout: \(payload)\n".utf8))
return
}
if ProgressTracker.barOpen {
print("")
ProgressTracker.barOpen = false
}
print(json)
fflush(stdout)
}

/// Where `run()` is when it throws, so `classifyExitReason` can produce a more
/// precise `reason` than "something failed" and `detail` can name which load
/// stalled. Set immediately before each fallible stage begins; read only if
Expand Down Expand Up @@ -651,6 +667,14 @@ struct MLXServer: AsyncParsableCommand {
let modelId = model

// ── Load model ──
// Same root the loader's HubApi uses further down.
let cliHub = HubApi(
downloadBase: URL.applicationSupportDirectory
.appendingPathComponent("MLX", isDirectory: true)
.appendingPathComponent("HuggingFace", isDirectory: true))
let isHubId = !ModelStorage.isLocalDirectoryPath(modelId)
&& !FileManager.default.fileExists(atPath: modelId)

var modelConfig: ModelConfiguration
if ModelStorage.isLocalDirectoryPath(modelId) {
print("[SwiftLM] Loading from local directory: \(modelId)")
Expand Down Expand Up @@ -681,46 +705,20 @@ struct MLXServer: AsyncParsableCommand {
var modelDirectory =
ModelStorage.validatedContentDirectory(for: modelId)
?? resolveModelDirectory(modelId: modelId)
// resolveModelDirectory doesn't check the weights are there; streaming must not be
// activated for an empty or partial snapshot the loader won't read.
if self.streamExperts, let dir = modelDirectory,
!ModelStorage.validateLocalModelDirectory(dir)
var modelDirectoryComplete = false
// --info doesn't download; don't profile a copy that is clearly partial.
if self.streamExperts, self.info, isHubId, let dir = modelDirectory,
localWeightState(in: dir) == .incomplete
{
modelDirectory = nil
}
if self.streamExperts, !self.info, modelDirectory == nil,
!FileManager.default.fileExists(atPath: modelId)
{
// Streaming must be activated for the directory the loader reads, so resolve it
// before loading. Same hub root as the loader below, which reuses these files.
let hub = HubApi(
downloadBase: URL.applicationSupportDirectory
.appendingPathComponent("MLX", isDirectory: true)
.appendingPathComponent("HuggingFace", isDirectory: true))
let localRepo = hub.localRepoLocation(Hub.Repo(id: modelId))
// Every shard must be present: an interrupted download (or the config.json
// the architecture probe fetches) would otherwise plan with a partial size.
if FileManager.default.fileExists(atPath: localRepo.path),
ModelStorage.validateLocalModelDirectory(localRepo)
{
modelDirectory = localRepo
} else {
// First run. A failed download is a model problem, not a binary one.
phase = .architectureProbe
print("[SwiftLM] --stream-experts: downloading \(modelId) before loading...")
let prefetchTracker = ProgressTracker(modelId: modelId)
defer { prefetchTracker.finish() }
modelDirectory = try await hub.snapshot(
from: modelId, matching: ["*.safetensors", "*.json", "*.jinja"]
) { progress in
prefetchTracker.printProgress(progress)
}
}
}
// Streaming is activated for `modelDirectory`, and only a load of that exact
// directory streams. Load from it, or a different lookup could pick another copy.
if self.streamExperts, let dir = modelDirectory {
modelConfig = ModelConfiguration(directory: dir)
if self.streamExperts, !self.info, isHubId {
// A Hub or download failure here is a model problem, not a binary one.
phase = .architectureProbe
let resolved = try await resolveStreamingDirectory(
modelId: modelId, candidate: modelDirectory, hub: cliHub)
modelDirectory = resolved.directory
modelDirectoryComplete = resolved.complete
}
var mainModelProfile: ModelProfile? = nil
if self.streamExperts, let dir = modelDirectory {
Expand All @@ -739,6 +737,13 @@ struct MLXServer: AsyncParsableCommand {
}
}

// Streaming is activated for `modelDirectory`, and only a load of that exact
// directory streams. Load from it, or a different lookup could pick another copy.
// A complete copy also loads from here after the MoE check turns streaming off.
if let dir = modelDirectory, self.streamExperts || modelDirectoryComplete {
modelConfig = ModelConfiguration(directory: dir)
}

// Inject streaming flag into config to bypass eval(model) if requested
if self.streamExperts {
modelConfig.lazyLoad = true
Expand Down Expand Up @@ -1605,20 +1610,22 @@ struct MLXServer: AsyncParsableCommand {
// is already gone by the time a shutdown signal arrives (e.g. the
// daemon itself already crashed), the exiting-event print()/fflush
// below can raise SIGPIPE — whose default disposition kills this
// process via signal instead of reaching Darwin.exit(0), producing
// process via signal instead of reaching exitAfterShutdownRequest(), producing
// exactly the ambiguous "was this a crash?" signature this feature
// exists to eliminate. Ignore SIGPIPE so a closed pipe surfaces as
// an ordinary EPIPE write error instead.
signal(SIGPIPE, SIG_IGN)

shutdownSource.setEventHandler {
print("\n[SwiftLM] Received SIGTERM, shutting down gracefully...")
emitEvent(["event": "exiting", "reason": "requested"])
emitEvent(
["event": "exiting", "reason": "requested"],
preface: "\n[SwiftLM] Received SIGTERM, shutting down gracefully...")
exitAfterShutdownRequest()
}
interruptSource.setEventHandler {
print("\n[SwiftLM] Received SIGINT, shutting down gracefully...")
emitEvent(["event": "exiting", "reason": "requested"])
emitEvent(
["event": "exiting", "reason": "requested"],
preface: "\n[SwiftLM] Received SIGINT, shutting down gracefully...")
exitAfterShutdownRequest()
}
shutdownSource.resume()
Expand Down
Loading
Loading