diff --git a/README.md b/README.md index bf31a2c..fecbaf8 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 2067d79..bd7225d 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -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 @@ -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 @@ -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() @@ -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 { @@ -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: "") @@ -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. /// @@ -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 @@ -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)") @@ -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 { @@ -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 @@ -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() diff --git a/Sources/SwiftLM/StreamingDirectory.swift b/Sources/SwiftLM/StreamingDirectory.swift new file mode 100644 index 0000000..dab5c3e --- /dev/null +++ b/Sources/SwiftLM/StreamingDirectory.swift @@ -0,0 +1,228 @@ +// StreamingDirectory.swift — the directory `--stream-experts` loads a Hub model from. +// +// Streaming targets one directory, and a directory load doesn't fill in missing +// files, so that copy must be complete. A local copy whose files show it is complete +// is used as is; otherwise the Hub listing decides. + +import Foundation +import Hub +import MLXInferenceCore + +/// Files SwiftLM downloads for a model: weights, configs, tokenizer and templates. +let modelDownloadPatterns = ["*.safetensors", "*.json", "*.jinja"] + +/// A `--stream-experts` download finished without the files the loader needs. +struct ModelDownloadIncomplete: Error, CustomStringConvertible { + let modelId: String + let directory: URL + var description: String { + "Download of \(modelId) is incomplete (\(directory.path)). Check the network, or delete the directory and retry." + } +} + +/// The repository has no `tokenizer.json`, which the tokenizer loader requires. +struct ModelMissingTokenizer: Error, CustomStringConvertible { + let modelId: String + var description: String { + "\(modelId) has no tokenizer.json, which SwiftLM needs to load it." + } +} + +/// The Hub answered 401/404 and there is no usable local copy. +struct ModelNotOnHub: Error, CustomStringConvertible { + let modelId: String + let underlying: Error + var description: String { + "The Hub has no accessible repository \(modelId) (\(underlying)), and no usable local copy was found. Check the id, or set HF_TOKEN for a private repository." + } +} + +/// The Hub couldn't be used and there is no usable local copy. +struct ModelUnavailableOffline: Error, CustomStringConvertible { + let modelId: String + let underlying: Error + var description: String { + "Could not get \(modelId) from the Hub (\(underlying)), and no usable local copy was found." + } +} + +/// Whether a local copy's weights are known to be complete. +enum LocalWeights: Equatable { + case complete + case incomplete + /// A layout the file names can't vouch for (`weights.00.safetensors`, say); + /// the Hub listing decides. + case unverified +} + +/// Whether `name` in `directory` exists and is non-empty, following symlinks +/// (Hub snapshots link to blobs, and a dangling link counts as missing). +func isPresentFile(_ name: String, in directory: URL) -> Bool { + let url = directory.appendingPathComponent(name) + let fm = FileManager.default + guard fm.fileExists(atPath: url.path) else { return false } + let attributes = try? fm.attributesOfItem(atPath: url.resolvingSymlinksInPath().path) + return ((attributes?[.size] as? Int) ?? 0) > 0 +} + +/// `stem-00002-of-00004.safetensors` as (stem, 2, 4); nil for other names. +func shardNumber(_ name: String) -> (stem: String, index: Int, count: Int)? { + guard name.hasSuffix(".safetensors") else { return nil } + let base = name.dropLast(".safetensors".count) + guard let of = base.range(of: "-of-", options: .backwards) else { return nil } + let countText = base[of.upperBound...] + let head = base[.. LocalWeights { + let indexed = indexedWeightFiles(in: directory) + if !indexed.isEmpty { + let missing = indexed.filter { !isPresentFile($0, in: directory) } + if missing.isEmpty { return .complete } + // A stale index names top-level shards the repo doesn't ship; a missing file + // in a subdirectory (optiq/optiq_vision.safetensors) is really missing. + if missing.contains(where: { $0.contains("/") }) { return .incomplete } + } + let top = ((try? FileManager.default.contentsOfDirectory(atPath: directory.path)) ?? []) + .filter { $0.hasSuffix(".safetensors") } + let chosen = ["model", "weight"] + .map { prefix in top.filter { $0.hasPrefix(prefix) } } + .first { !$0.isEmpty } ?? top + guard !chosen.isEmpty else { return .incomplete } + + // `stem-00002-of-00004.safetensors`: every shard 1...4 must be there. + var groups = [String: (count: Int, have: Set)]() + for name in chosen { + guard let (stem, i, n) = shardNumber(name) else { continue } + let key = "\(stem)/\(n)" + var group = groups[key] ?? (n, []) + if isPresentFile(name, in: directory) { group.have.insert(i) } + groups[key] = group + } + if !groups.isEmpty { + let whole = groups.values.allSatisfy { $0.count > 0 && $0.have == Set(1 ... $0.count) } + return whole ? .complete : .incomplete + } + if chosen == ["model.safetensors"] { + return isPresentFile("model.safetensors", in: directory) ? .complete : .incomplete + } + return .unverified +} + +/// The files `model.safetensors.index.json` in `directory` names, if any. +func indexedWeightFiles(in directory: URL) -> Set { + let index = directory.appendingPathComponent("model.safetensors.index.json") + guard let data = try? Data(contentsOf: index), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let weightMap = json["weight_map"] as? [String: String] + else { return [] } + return Set(weightMap.values) +} + +/// Listed files `directory` must have: every top-level one, plus any nested one its +/// index names (weights in a subdirectory, such as OptiQ's vision tower). +func requiredListedFiles(_ files: [String], in directory: URL) -> [String] { + let indexed = indexedWeightFiles(in: directory) + return files.filter { !$0.contains("/") || indexed.contains($0) } +} + +/// Whether `directory` has the non-weight files the loader reads. swift-transformers +/// requires `tokenizer.json`; `tokenizer_config.json` is optional. +func hasModelConfigAndTokenizer(in directory: URL) -> Bool { + isPresentFile("config.json", in: directory) && isPresentFile("tokenizer.json", in: directory) +} + +/// Whether a Hub error means the repository isn't available to this user. +private func isNotOnHub(_ error: Error) -> Bool { + guard let hubError = error as? Hub.HubClientError else { return false } + switch hubError { + case .authorizationRequired, .resourceNotFound, .fileNotFound: return true + case .httpStatusCode(let code): return code == 401 || code == 404 + default: return false + } +} + +/// The directory `--stream-experts` loads and streams `modelId` from, and whether that +/// copy is complete (a dense model's partial copy is returned for the loader to finish). +/// +/// A local copy (`candidate`, then the loader's Application Support copy) whose files +/// show it is complete is used without contacting the Hub. Otherwise the Hub listing +/// decides: a copy with every top-level listed file is used, or the model is downloaded +/// into the Application Support copy. If the Hub can't be used, a copy whose layout +/// its names can't verify is used with a warning. +func resolveStreamingDirectory( + modelId: String, candidate: URL?, hub: HubApi +) async throws -> (directory: URL, complete: Bool) { + let fm = FileManager.default + let localRepo = hub.localRepoLocation(Hub.Repo(id: modelId)) + // Present while a download into localRepo is unfinished; that copy isn't trusted. + let inProgress = localRepo.appendingPathComponent(".swiftlm-download-in-progress") + let copies = [candidate, localRepo].compactMap { $0 }.filter { + fm.fileExists(atPath: $0.path) + && !($0 == localRepo && fm.fileExists(atPath: inProgress.path)) + } + let usable = copies.filter(hasModelConfigAndTokenizer) + if let dir = usable.first(where: { localWeightState(in: $0) == .complete }) { + return (dir, true) + } + let unverified = usable.filter { localWeightState(in: $0) == .unverified } + + let files: [String] + do { + files = try await hub.getFilenames(from: modelId, matching: modelDownloadPatterns) + } catch { + if let dir = unverified.first { + print("[SwiftLM] ⚠️ Could not check \(modelId) with the Hub (\(error)); using \(dir.path) as is.") + return (dir, true) + } + if isNotOnHub(error) { throw ModelNotOnHub(modelId: modelId, underlying: error) } + throw ModelUnavailableOffline(modelId: modelId, underlying: error) + } + guard files.contains("tokenizer.json") else { throw ModelMissingTokenizer(modelId: modelId) } + if let dir = usable.first(where: { dir in + requiredListedFiles(files, in: dir).allSatisfy { isPresentFile($0, in: dir) } + }) { + return (dir, true) + } + + // A dense model won't stream: leave completing it to the loader instead of + // downloading everything here first. + let allCopies = [candidate, localRepo].compactMap { $0 } + if let dir = allCopies.first(where: { isPresentFile("config.json", in: $0) }), + let profile = ModelProfiler.profile(modelDirectory: dir, modelId: modelId), !profile.isMoE + { + return (dir, false) + } + for dir in allCopies where fm.fileExists(atPath: dir.path) { + let reason = dir == localRepo && fm.fileExists(atPath: inProgress.path) + ? "its download didn't finish" : "it is missing files" + print("[SwiftLM] \(dir.path) can't be used as is (\(reason)).") + } + + print("[SwiftLM] --stream-experts: downloading \(modelId) before loading...") + try fm.createDirectory(at: localRepo, withIntermediateDirectories: true) + fm.createFile(atPath: inProgress.path, contents: nil) + let tracker = ProgressTracker(modelId: modelId) + defer { tracker.finish() } + let snapshot = try await hub.snapshot(from: modelId, matching: modelDownloadPatterns) { + tracker.printProgress($0) + } + tracker.finish() + guard requiredListedFiles(files, in: snapshot).allSatisfy({ isPresentFile($0, in: snapshot) }), + hasModelConfigAndTokenizer(in: snapshot) + else { + throw ModelDownloadIncomplete(modelId: modelId, directory: snapshot) + } + try? fm.removeItem(at: inProgress) + return (snapshot, true) +} diff --git a/tests/SwiftLMTests/StreamingDirectoryTests.swift b/tests/SwiftLMTests/StreamingDirectoryTests.swift new file mode 100644 index 0000000..778a597 --- /dev/null +++ b/tests/SwiftLMTests/StreamingDirectoryTests.swift @@ -0,0 +1,123 @@ +import Foundation +import XCTest + +@testable import SwiftLM + +/// `localWeightState` must judge a copy the way the loader reads it, so a partial +/// download is never served and a complete copy with a stale index is not rejected. +final class StreamingDirectoryTests: XCTestCase { + private var dir: URL! + + override func setUpWithError() throws { + dir = FileManager.default.temporaryDirectory + .appendingPathComponent("streaming-dir-\(UUID().uuidString)") + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + try write("config.json") + try write("tokenizer.json") + } + + override func tearDownWithError() throws { + try? FileManager.default.removeItem(at: dir) + } + + private func write(_ name: String, _ text: String = "{}") throws { + try Data(text.utf8).write(to: dir.appendingPathComponent(name)) + } + + private func shards(_ indices: [Int], of count: Int) throws { + for i in indices { + try write(String(format: "model-%05d-of-%05d.safetensors", i, count), "weights") + } + } + + /// An index naming shards the repository doesn't ship (carried over from the source). + private func staleIndex(naming count: Int) throws { + let map = Dictionary(uniqueKeysWithValues: (1 ... count).map { + ("layer\($0).weight", String(format: "model-%05d-of-%05d.safetensors", $0, count)) + }) + let json = try JSONSerialization.data(withJSONObject: ["weight_map": map]) + try json.write(to: dir.appendingPathComponent("model.safetensors.index.json")) + } + + func testPartialCopyWithoutIndexIsIncomplete() throws { + try shards([1, 3], of: 3) + XCTAssertEqual(localWeightState(in: dir), .incomplete) + } + + func testAllShardsWithoutIndexAreComplete() throws { + try shards([1, 2, 3], of: 3) + XCTAssertEqual(localWeightState(in: dir), .complete) + } + + func testStaleIndexWithEveryShippedShardIsComplete() throws { + try staleIndex(naming: 13) + try shards([1, 2, 3, 4], of: 4) + XCTAssertEqual(localWeightState(in: dir), .complete) + } + + func testStaleIndexWithAMissingShardIsIncomplete() throws { + try staleIndex(naming: 13) + try shards([1, 2, 4], of: 4) + XCTAssertEqual(localWeightState(in: dir), .incomplete) + } + + /// OptiQ VLMs index their vision tower in optiq/; a download of only model*.safetensors + /// has every top-level shard but not that file. + private func indexWithNestedVision() throws { + var map = Dictionary(uniqueKeysWithValues: (1 ... 2).map { + ("layer\($0).weight", String(format: "model-%05d-of-%05d.safetensors", $0, 2)) + }) + map["vision_tower.weight"] = "optiq/optiq_vision.safetensors" + let json = try JSONSerialization.data(withJSONObject: ["weight_map": map]) + try json.write(to: dir.appendingPathComponent("model.safetensors.index.json")) + } + + func testIndexedFileMissingFromSubdirectoryIsIncomplete() throws { + try indexWithNestedVision() + try shards([1, 2], of: 2) + XCTAssertEqual(localWeightState(in: dir), .incomplete) + } + + func testIndexedSubdirectoryFilePresentIsComplete() throws { + try indexWithNestedVision() + try shards([1, 2], of: 2) + try FileManager.default.createDirectory( + at: dir.appendingPathComponent("optiq"), withIntermediateDirectories: true) + try write("optiq/optiq_vision.safetensors", "weights") + XCTAssertEqual(localWeightState(in: dir), .complete) + } + + func testRequiredListedFilesIncludeIndexedNestedFiles() throws { + try indexWithNestedVision() + let listed = ["config.json", "model-00001-of-00002.safetensors", + "optiq/optiq_vision.safetensors", "optiq/mtp.safetensors"] + XCTAssertEqual( + Set(requiredListedFiles(listed, in: dir)), + ["config.json", "model-00001-of-00002.safetensors", "optiq/optiq_vision.safetensors"]) + } + + func testSingleModelFileIsComplete() throws { + try write("model.safetensors", "weights") + XCTAssertEqual(localWeightState(in: dir), .complete) + } + + func testUnconventionalLayoutIsUnverified() throws { + try write("weights.00.safetensors", "weights") + XCTAssertEqual(localWeightState(in: dir), .unverified) + } + + func testDanglingSymlinkIsMissing() throws { + try FileManager.default.createSymbolicLink( + at: dir.appendingPathComponent("model.safetensors"), + withDestinationURL: dir.appendingPathComponent("gone.bin")) + XCTAssertFalse(isPresentFile("model.safetensors", in: dir)) + XCTAssertEqual(localWeightState(in: dir), .incomplete) + } + + func testTokenizerJSONIsRequired() throws { + XCTAssertTrue(hasModelConfigAndTokenizer(in: dir)) + try FileManager.default.removeItem(at: dir.appendingPathComponent("tokenizer.json")) + try write("tokenizer_config.json") + XCTAssertFalse(hasModelConfigAndTokenizer(in: dir)) + } +} diff --git a/tests/test-server.sh b/tests/test-server.sh index 95b0939..36067b7 100755 --- a/tests/test-server.sh +++ b/tests/test-server.sh @@ -35,6 +35,14 @@ cleanup() { kill -9 "$SERVER_PID" 2>/dev/null || true wait "$SERVER_PID" 2>/dev/null || true fi + if [ -n "${TERM_SERVER_PID:-}" ]; then + kill -9 "$TERM_SERVER_PID" 2>/dev/null || true + wait "$TERM_SERVER_PID" 2>/dev/null || true + fi + if [ -n "${TERM_CURL_PID:-}" ]; then + kill "$TERM_CURL_PID" 2>/dev/null || true + fi + rm -f "${TERM_LOG:-}" "${TERM_STREAM:-}" if [ -n "${CORS_SERVER_PID:-}" ]; then log "Stopping CORS server (PID $CORS_SERVER_PID)" kill -9 "$CORS_SERVER_PID" 2>/dev/null || true @@ -1203,6 +1211,81 @@ fi rm -f /tmp/mlx_ttfb_prompt.txt /tmp/mlx_ttfb_body.json /tmp/mlx_ttfb_stream.txt +# ── Test 39: SIGTERM during generation exits cleanly ───────────────── +# A shutdown request mid-stream must exit 0 with exiting{reason:"requested"} on its +# own line (exit() used to race an in-flight GPU eval and crash, #197). A crash +# shows up as a nonzero status; crash reports are written too late to check here. +log "Test 39: SIGTERM during a streaming generation" +TERM_PORT=$((PORT + 3)) +TERM_LOG=$(mktemp) +TERM_STREAM=$(mktemp) +"$BINARY" --model "$MODEL" --port "$TERM_PORT" --host "$HOST" > "$TERM_LOG" 2>&1 & +TERM_SERVER_PID=$! +TERM_READY=false +for i in $(seq 1 120); do + if curl -sf "http://${HOST}:${TERM_PORT}/health" >/dev/null 2>&1; then + TERM_READY=true + break + fi + kill -0 "$TERM_SERVER_PID" 2>/dev/null || break + sleep 1 +done + +if [ "$TERM_READY" != true ]; then + fail "SIGTERM mid-stream: server did not become ready" + tail -40 "$TERM_LOG" || true + cp "$TERM_LOG" /tmp/SwiftLM-test-sigterm.log 2>/dev/null || true + kill -9 "$TERM_SERVER_PID" 2>/dev/null || true + wait "$TERM_SERVER_PID" 2>/dev/null || true + unset TERM_SERVER_PID +else + curl -sN -X POST "http://${HOST}:${TERM_PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"$MODEL\",\"stream\":true,\"max_tokens\":2048,\"messages\":[{\"role\":\"user\",\"content\":\"Write a very long story about a lighthouse keeper.\"}]}" \ + > "$TERM_STREAM" 2>/dev/null & + TERM_CURL_PID=$! + for i in $(seq 1 200); do + grep -q '"content"' "$TERM_STREAM" 2>/dev/null && break + sleep 0.1 + done + kill -TERM "$TERM_SERVER_PID" 2>/dev/null || true + + # Bounded wait: a shutdown hang must fail here, not at the job timeout. + TERM_HUNG=false + for i in $(seq 1 300); do + kill -0 "$TERM_SERVER_PID" 2>/dev/null || break + sleep 0.1 + done + if kill -0 "$TERM_SERVER_PID" 2>/dev/null; then + TERM_HUNG=true + kill -9 "$TERM_SERVER_PID" 2>/dev/null || true + fi + TERM_STATUS=0 + wait "$TERM_SERVER_PID" 2>/dev/null || TERM_STATUS=$? + unset TERM_SERVER_PID + kill "$TERM_CURL_PID" 2>/dev/null || true + wait "$TERM_CURL_PID" 2>/dev/null || true + unset TERM_CURL_PID + + # Each line on its own: a line that is not exactly one JSON object (an echoed + # token, or the event with a token glued on) must not count. + TERM_REASON=$(jq -rR 'fromjson? | select(type == "object" and .event == "exiting") | .reason' \ + "$TERM_LOG" 2>/dev/null | tail -1 || true) + if [ "$TERM_HUNG" = true ]; then + fail "SIGTERM mid-stream: server still running 30s after SIGTERM" + tail -40 "$TERM_LOG" || true + cp "$TERM_LOG" /tmp/SwiftLM-test-sigterm.log 2>/dev/null || true + elif [ "$TERM_STATUS" -eq 0 ] && [ "$TERM_REASON" = "requested" ]; then + pass "SIGTERM mid-stream: exit 0, exiting{requested} on its own line" + else + fail "SIGTERM mid-stream: status=$TERM_STATUS reason='${TERM_REASON}'" + tail -40 "$TERM_LOG" || true + cp "$TERM_LOG" /tmp/SwiftLM-test-sigterm.log 2>/dev/null || true + fi +fi +rm -f "$TERM_LOG" "$TERM_STREAM" +unset TERM_LOG TERM_STREAM + # ── Results ────────────────────────────────────────────────────────── echo "" log "═══════════════════════════════════════"