diff --git a/README.md b/README.md index 40b5711..d46eaca 100644 --- a/README.md +++ b/README.md @@ -223,11 +223,7 @@ val handle = wildEdge.registerModel("my-model", ModelInfo( )) handle.trackLoad(durationMs = loadMs, accelerator = Accelerator.CPU, coldStart = true) - -val start = System.currentTimeMillis() -val output = model.run(input) -handle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - +val output = handle.trackInference { model.run(input) } handle.trackUnload() ``` @@ -260,19 +256,8 @@ Group related inferences so the server can reconstruct the full pipeline: ```kotlin wildEdge.trace("user-query") { trace -> - trace.span("embed") { - val start = System.currentTimeMillis() - val embedding = embedModel.run(input) - embedHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - embedding - } - - trace.span("classify") { - val start = System.currentTimeMillis() - val label = classifyModel.run(embedding) - classifyHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - label - } + val embedding = trace.span("embed") { embedHandle.trackInference { embedModel.run(input) } } + trace.span("classify") { classifyHandle.trackInference { classifyModel.run(embedding) } } } ``` @@ -318,12 +303,7 @@ Declare your field against `WildEdgeClient` and inject `WildEdgeClient.noop()` i class InferenceService(private val wildEdge: WildEdgeClient) { private val handle = wildEdge.registerModel("my-model", ...) - fun run(input: ByteArray): Result { - val start = System.currentTimeMillis() - val result = model.run(input) - handle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - return result - } + fun run(input: ByteArray): Result = handle.trackInference { model.run(input) } } // In tests: diff --git a/examples/tracing/TracingExample.kt b/examples/tracing/TracingExample.kt index 679ecb8..f48afc4 100644 --- a/examples/tracing/TracingExample.kt +++ b/examples/tracing/TracingExample.kt @@ -2,6 +2,7 @@ package examples.tracing import dev.wildedge.sdk.ModelInfo import dev.wildedge.sdk.WildEdge +import dev.wildedge.sdk.trackInference // Assumes WildEdge.init() has already run (manifest meta-data or Application.onCreate()). class TracingExample { @@ -15,19 +16,8 @@ class TracingExample { // Nested span() calls set parent_span_id so the server can reconstruct the tree. fun runPipeline(input: ByteArray) { wildEdge.trace("pipeline") { trace -> - val embedding = trace.span("embed") { - val start = System.currentTimeMillis() - val result = runEmbedding(input) - embedHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - result - } - - trace.span("classify") { - val start = System.currentTimeMillis() - val result = runClassification(embedding) - classifyHandle.trackInference(durationMs = (System.currentTimeMillis() - start).toInt()) - result - } + val embedding = trace.span("embed") { embedHandle.trackInference { runEmbedding(input) } } + trace.span("classify") { classifyHandle.trackInference { runClassification(embedding) } } } } diff --git a/wildedge/src/main/kotlin/dev/wildedge/sdk/ModelHandle+Coroutines.kt b/wildedge/src/main/kotlin/dev/wildedge/sdk/ModelHandle+Coroutines.kt index 1735c30..4fc6a57 100644 --- a/wildedge/src/main/kotlin/dev/wildedge/sdk/ModelHandle+Coroutines.kt +++ b/wildedge/src/main/kotlin/dev/wildedge/sdk/ModelHandle+Coroutines.kt @@ -6,6 +6,53 @@ import dev.wildedge.sdk.events.TextInputMeta import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow +/** + * Times a blocking inference block and emits a tracking event on completion or error. + * + * Usage: + * ``` + * val result = handle.trackInference { model.run(input) } + * ``` + */ +@Suppress("TooGenericExceptionCaught") +fun ModelHandle.trackInference( + inputModality: InputModality? = null, + outputModality: OutputModality? = null, + inputMeta: Map? = null, + outputMeta: Map? = null, + outputMetaExtractor: ((T) -> Map?)? = null, + traceId: String? = null, + parentSpanId: String? = null, + block: () -> T, +): T { + val start = System.currentTimeMillis() + return try { + val result = block() + trackInference( + durationMs = (System.currentTimeMillis() - start).toInt(), + inputModality = inputModality, + outputModality = outputModality, + inputMeta = inputMeta, + outputMeta = outputMetaExtractor?.invoke(result) ?: outputMeta, + traceId = traceId, + parentSpanId = parentSpanId, + ) + result + } catch (e: Exception) { + trackInference( + durationMs = (System.currentTimeMillis() - start).toInt(), + inputModality = inputModality, + outputModality = outputModality, + inputMeta = inputMeta, + success = false, + errorCode = e.javaClass.simpleName, + traceId = traceId, + parentSpanId = parentSpanId, + ) + throw e + } +} + /** * Times a suspend inference block and emits a tracking event on completion or error. * diff --git a/wildedge/src/test/kotlin/dev/wildedge/sdk/CoroutinesTrackingTest.kt b/wildedge/src/test/kotlin/dev/wildedge/sdk/CoroutinesTrackingTest.kt index 7f692ba..014d66b 100644 --- a/wildedge/src/test/kotlin/dev/wildedge/sdk/CoroutinesTrackingTest.kt +++ b/wildedge/src/test/kotlin/dev/wildedge/sdk/CoroutinesTrackingTest.kt @@ -18,6 +18,58 @@ class CoroutinesTrackingTest { return handle to events } + // --- trackInference (block) --- + + @Test fun blockTrackingEmitsInferenceEvent() { + val (handle, events) = captureHandle() + handle.trackInference { "result" } + assertEquals(1, events.count { it["event_type"] == "inference" }) + } + + @Test fun blockTrackingReturnsBlockResult() { + val (handle, _) = captureHandle() + val result = handle.trackInference { 42 } + assertEquals(42, result) + } + + @Test fun blockTrackingRecordsSuccessTrue() { + val (handle, events) = captureHandle() + handle.trackInference { Unit } + val inference = events.first { it["event_type"] == "inference" } + @Suppress("UNCHECKED_CAST") + assertEquals(true, (inference["inference"] as Map)["success"]) + } + + @Test fun blockTrackingRecordsFailureOnException() { + val (handle, events) = captureHandle() + runCatching { handle.trackInference { throw IllegalStateException("boom") } } + val inference = events.first { it["event_type"] == "inference" } + + @Suppress("UNCHECKED_CAST") + val inf = inference["inference"] as Map + assertEquals(false, inf["success"]) + assertEquals("IllegalStateException", inf["error_code"]) + } + + @Test fun blockTrackingRethrowsException() { + val (handle, _) = captureHandle() + val ex = runCatching { + handle.trackInference { throw RuntimeException("rethrow me") } + }.exceptionOrNull() + assertNotNull(ex) + assertEquals("rethrow me", ex!!.message) + } + + @Test fun blockTrackingOutputMetaExtractorReceivesResult() { + val (handle, events) = captureHandle() + handle.trackInference(outputMetaExtractor = { r: String -> mapOf("extracted" to r) }) { "hello" } + val inference = events.first { it["event_type"] == "inference" } + + @Suppress("UNCHECKED_CAST") + val outputMeta = (inference["inference"] as Map)["output_meta"] as Map? + assertEquals("hello", outputMeta?.get("extracted")) + } + // --- trackSuspendInference --- @Test fun suspendTrackingEmitsInferenceEvent() = runBlocking {