From d9bc69cc8c9a56f5e57c974780353f1c3796a263 Mon Sep 17 00:00:00 2001 From: Tao Jia Date: Thu, 23 Jul 2026 10:58:46 -0700 Subject: [PATCH] Polish a few APIs, method names, and remove unnecessarily vending public extension --- .../VanillaDecodingStrategy.swift | 4 ++-- .../LanguageModel/CoreAIRunner.swift | 4 ++-- .../Profiling/PerformanceMetrics.swift | 22 +++++++++---------- .../Profiling/Timing.swift | 8 +++++-- .../Sources/CoreAIShared/Logger/Logger.swift | 6 +---- .../ImageSegmentationRunnerMain.swift | 2 +- .../Tools/llm-runner/LLMRunnerMain.swift | 12 +++++----- .../PerformanceMetricsTests.swift | 10 ++++----- 8 files changed, 34 insertions(+), 34 deletions(-) diff --git a/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift b/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift index 9d40942..b6f447b 100644 --- a/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift +++ b/swift/Sources/CoreAILanguageModels/DecodingStrategies/VanillaDecodingStrategy.swift @@ -164,7 +164,7 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence { if !flushed { let count = generatedTokenCount // Metrics recording is async; fire-and-forget on the abandon path. - Task { await PerformanceMetrics.shared.setGeneratedTokenCount(count) } + Task { await PerformanceMetrics.shared.recordGeneratedTokens(count) } } } @@ -283,7 +283,7 @@ extension VanillaDecodingStrategy.VanillaDecodedSequence { trailing = nil } - await PerformanceMetrics.shared.setGeneratedTokenCount(generatedTokenCount) + await PerformanceMetrics.shared.recordGeneratedTokens(generatedTokenCount) return trailing } } diff --git a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift index 9fe857e..67d1247 100644 --- a/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift +++ b/swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift @@ -32,7 +32,7 @@ public struct CoreAIRunner { kvCacheStrategy: KVCacheStrategy = .auto ) throws { self.init( - from: try LanguageBundle(at: url), + bundle: try LanguageBundle(at: url), variant: variant, kvCacheStrategy: kvCacheStrategy ) @@ -40,7 +40,7 @@ public struct CoreAIRunner { /// Creates a runner from a LanguageBundle. public init( - from bundle: LanguageBundle, + bundle: LanguageBundle, variant: String? = nil, kvCacheStrategy: KVCacheStrategy = .auto ) { diff --git a/swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift b/swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift index ce3c9a6..7c7f663 100644 --- a/swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift +++ b/swift/Sources/CoreAILanguageModels/Profiling/PerformanceMetrics.swift @@ -28,9 +28,8 @@ public final class PerformanceMetrics { private var startInstant: ContinuousClock.Instant? private var endInstant: ContinuousClock.Instant? - private var promptTokenCount: Int = 0 - private var generatedTokenCount: Int = 0 - private var totalTokenCount: Int = 0 + public private(set) var promptTokenCount: Int = 0 + public private(set) var generatedTokenCount: Int = 0 /// The shared PerformanceMetrics instance for production use. public static let shared = PerformanceMetrics() @@ -61,17 +60,19 @@ public final class PerformanceMetrics { // MARK: - Token Counting - public func setPromptTokenCount(_ count: Int) { - promptTokenCount = count + /// Total number of prompt and generated tokens. + public var totalTokenCount: Int { + promptTokenCount + generatedTokenCount } - public func setGeneratedTokenCount(_ count: Int) { - generatedTokenCount = count - totalTokenCount = promptTokenCount + generatedTokenCount + /// Records the number of tokens in the prompt. + public func recordPromptTokens(_ count: Int) { + promptTokenCount = count } - public var getGeneratedTokenCount: Int { - return generatedTokenCount + /// Records the number of tokens produced during generation. + public func recordGeneratedTokens(_ count: Int) { + generatedTokenCount = count } // MARK: - Computed Metrics (from StatsStorage) @@ -236,7 +237,6 @@ public final class PerformanceMetrics { endInstant = nil promptTokenCount = 0 generatedTokenCount = 0 - totalTokenCount = 0 // Also reset StatsStorage since this is a full reset StatsStorage.shared.reset() } diff --git a/swift/Sources/CoreAILanguageModels/Profiling/Timing.swift b/swift/Sources/CoreAILanguageModels/Profiling/Timing.swift index 626a2db..f2e83cf 100644 --- a/swift/Sources/CoreAILanguageModels/Profiling/Timing.swift +++ b/swift/Sources/CoreAILanguageModels/Profiling/Timing.swift @@ -9,6 +9,10 @@ import Foundation /// Extension to convert Duration to common time units. /// +/// - Note: This is an internal implementation detail. It is intentionally not +/// `public`: vending members on a standard-library type we don't own would +/// pollute `Duration`'s API surface for every client of this library. +/// /// Example usage: /// ```swift /// let start = ContinuousClock.now @@ -18,13 +22,13 @@ import Foundation /// ``` extension Duration { /// Duration in seconds as a Double. - public var inSeconds: Double { + var inSeconds: Double { let (secs, attoseconds) = self.components return Double(secs) + Double(attoseconds) / 1e18 } /// Duration in milliseconds as a Double. - public var inMilliseconds: Double { + var inMilliseconds: Double { inSeconds * 1000.0 } } diff --git a/swift/Sources/CoreAIShared/Logger/Logger.swift b/swift/Sources/CoreAIShared/Logger/Logger.swift index 51a65f2..8bde66e 100644 --- a/swift/Sources/CoreAIShared/Logger/Logger.swift +++ b/swift/Sources/CoreAIShared/Logger/Logger.swift @@ -11,7 +11,7 @@ import Synchronization public struct CLILogger { private static let _level = Atomic(0) - static var level: Int { + public static var level: Int { get { _level.load(ordering: .acquiring) } @@ -21,10 +21,6 @@ public struct CLILogger { } } - public static func setLevel(to level: Int) { - Self.level = level - } - /// Performs logging if enabled for the requested level. /// - Parameters: /// - message: The message to log. diff --git a/swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift b/swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift index cd00218..aadfaa6 100644 --- a/swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift +++ b/swift/Sources/Tools/image-segmenter/ImageSegmentationRunnerMain.swift @@ -157,7 +157,7 @@ struct ImageSegmenterCLI: AsyncParsableCommand { throw ValidationError("--image is required.") } if verbose { - CLILogger.setLevel(to: 1) + CLILogger.level = 1 } if verbose { print("Creating image segmenter...") } diff --git a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift index 55b2914..74f0ccc 100644 --- a/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift +++ b/swift/Sources/Tools/llm-runner/LLMRunnerMain.swift @@ -221,7 +221,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { func run() async throws { let verboseLevel = max(self.verboseLevel ?? 0, verbose ? 1 : 0) - CLILogger.setLevel(to: verboseLevel) + CLILogger.level = verboseLevel let resolver = ModelPaths() let resolvedPath = try validateAndResolveModelPath(resolver: resolver) @@ -308,7 +308,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { // Set up verbose logging environment variable first let verboseLevel = max(self.verboseLevel ?? 0, verbose ? 1 : 0) - CLILogger.setLevel(to: verboseLevel) + CLILogger.level = verboseLevel // Bridge hidden CLI overrides to environment variables read by the Core AI engine if let b = bucketSize { @@ -530,7 +530,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { CLILogger.log("Text generator built successfully", component: "Main") // Apply chat template and count tokens for metrics - await PerformanceMetrics.shared.setPromptTokenCount(promptTokens.count) + await PerformanceMetrics.shared.recordPromptTokens(promptTokens.count) CLILogger.log("Generating text...", component: "Main") CLILogger.log("Input: \(displayPrompt)", component: "Main") @@ -632,7 +632,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { InstrumentsProfiler.endDecoding(signpostID: decodingID) // Generated token count is already set by the decoding strategy - let generatedTokenCount = await PerformanceMetrics.shared.getGeneratedTokenCount + let generatedTokenCount = await PerformanceMetrics.shared.generatedTokenCount InstrumentsProfiler.endInference(generatedTokens: generatedTokenCount, signpostID: inferenceID) // End overall timing now that core inference is complete @@ -924,7 +924,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { let inferenceID = InstrumentsProfiler.beginInference( promptTokens: vlmTokens.count, maxTokens: maxTokens) - await PerformanceMetrics.shared.setPromptTokenCount(vlmTokens.count) + await PerformanceMetrics.shared.recordPromptTokens(vlmTokens.count) let tokenStream = try await vlmEngine.generate( with: embeddedInput, @@ -993,7 +993,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable { // Record generation stats InstrumentsProfiler.endInference( generatedTokens: generatedTokens.count, signpostID: inferenceID) - await PerformanceMetrics.shared.setGeneratedTokenCount(generatedTokens.count) + await PerformanceMetrics.shared.recordGeneratedTokens(generatedTokens.count) await PerformanceMetrics.shared.endOverallTiming() await PerformanceMetrics.shared.printSummary(verbose: CLILogger.isVerbose) diff --git a/swift/Tests/LanguageModelsTests/PerformanceMetricsTests.swift b/swift/Tests/LanguageModelsTests/PerformanceMetricsTests.swift index 1873c08..0ef7766 100644 --- a/swift/Tests/LanguageModelsTests/PerformanceMetricsTests.swift +++ b/swift/Tests/LanguageModelsTests/PerformanceMetricsTests.swift @@ -37,12 +37,12 @@ struct PerformanceMetricsTests { metrics.reset() metrics.startOverallTiming() - metrics.setPromptTokenCount(100) - metrics.setGeneratedTokenCount(50) + metrics.recordPromptTokens(100) + metrics.recordGeneratedTokens(50) metrics.reset() #expect(metrics.totalTime == 0) - #expect(metrics.getGeneratedTokenCount == 0) + #expect(metrics.generatedTokenCount == 0) #expect(metrics.modelLoadTime == 0) } @@ -51,8 +51,8 @@ struct PerformanceMetricsTests { let metrics = PerformanceMetrics.shared metrics.reset() - metrics.setPromptTokenCount(100) - metrics.setGeneratedTokenCount(50) + metrics.recordPromptTokens(100) + metrics.recordGeneratedTokens(50) #expect(metrics.promptThroughput == 0) #expect(metrics.generationThroughput == 0)