From e09e56591de9d73a7e02326fde06d21b1ab80bb1 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 23 Aug 2026 21:19:26 +0200 Subject: [PATCH] =?UTF-8?q?feat(memory):=20SKAINET=5FMEMORY=5FDEBUG=20?= =?UTF-8?q?=E2=80=94=20allocation-site=20tagging,=20use-after-close=20with?= =?UTF-8?q?=20the=20closing=20site,=20adapter=20log,=20leak=20reports=20(S?= =?UTF-8?q?KEEP-003=20P2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Milestone M1 (#1002), PRD M1-F9, SKEEP-003 §4.7 / §8 item 4: the tool that would have found #782 in one run. - MemoryDebug: off by default (SKAINET_MEMORY_DEBUG=1 or -Dskainet.memory.debug=true; overrideEnabled for tests), costing one boolean read per allocation. On, it keeps a bounded ledger — storage id, scope, bytes, origin TensorId, allocation site, closing site — plus per-scope live and high-water bytes. - Breakpoint-like hooks (§4.7): onAllocate, onClose, onAdapter, onLeak. A debugger or a test breaks on "an allocation whose origin matches layers[3].*", "storage #412 closed", or "an adapter larger than X". - Storage records allocations and closes into the ledger and passes the call site into the trace event too; checkAlive()'s message now names where the storage was allocated and closed, which is the difference between "something is closed" and "this weight was freed at model.close(), you are reading it from the KV path". - platformCallSite(): a StackWalker frame outside sk.ainet.lang.memory on JVM/Android, null elsewhere; platformMemoryDebugEnabled() reads the env var and the property where there is one. - MemoryDebugTest: nothing recorded while off; allocations tagged with scope/origin/site; peaks survive frees; use-after-close names the tensor, scope and size; close hook fires once; adapter reporting; leak reporting distinguishes an escaped storage from a closed one; report() renders peaks and live storages. 97/97 memory tests. Closes #1026 Co-Authored-By: Claude Fable 5 --- .../api/jvm/skainet-lang-core.api | 53 ++++++ .../ainet/lang/memory/MemoryDebug.android.kt | 14 ++ .../sk/ainet/lang/memory/MemoryDebug.kt | 153 ++++++++++++++++++ .../kotlin/sk/ainet/lang/memory/Storage.kt | 19 ++- .../sk/ainet/lang/memory/MemoryDebugTest.kt | 109 +++++++++++++ .../sk/ainet/lang/memory/MemoryDebug.other.kt | 9 ++ .../sk/ainet/lang/memory/MemoryDebug.jvm.kt | 14 ++ .../sk/ainet/lang/memory/MemoryDebug.other.kt | 9 ++ .../sk/ainet/lang/memory/MemoryDebug.other.kt | 9 ++ .../sk/ainet/lang/memory/MemoryDebug.other.kt | 9 ++ 10 files changed, 395 insertions(+), 3 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/MemoryDebug.android.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/MemoryDebug.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/MemoryDebugTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/jsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt create mode 100644 skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/MemoryDebug.jvm.kt create mode 100644 skainet-lang/skainet-lang-core/src/nativeMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt create mode 100644 skainet-lang/skainet-lang-core/src/wasmJsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt create mode 100644 skainet-lang/skainet-lang-core/src/wasmWasiMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 0452ff880..c75821fad 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -836,6 +836,59 @@ public final class sk/ainet/lang/memory/MappedFileStorage$Companion { public static synthetic fun map$default (Lsk/ainet/lang/memory/MappedFileStorage$Companion;Ljava/nio/file/Path;JJLsk/ainet/lang/memory/ScopeKind;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/memory/trace/TraceSink;ILjava/lang/Object;)Lsk/ainet/lang/memory/MappedFileStorage; } +public final class sk/ainet/lang/memory/MemoryDebug { + public static final field ENV Ljava/lang/String; + public static final field INSTANCE Lsk/ainet/lang/memory/MemoryDebug; + public static final field PROPERTY Ljava/lang/String; + public final fun describeClosed-3jxYLUQ (J)Ljava/lang/String; + public final fun entry-3jxYLUQ (J)Lsk/ainet/lang/memory/MemoryDebug$Entry; + public final fun getOnAdapter ()Lkotlin/jvm/functions/Function3; + public final fun getOnAllocate ()Lkotlin/jvm/functions/Function1; + public final fun getOnClose ()Lkotlin/jvm/functions/Function1; + public final fun getOnLeak ()Lkotlin/jvm/functions/Function1; + public final fun getOverrideEnabled ()Ljava/lang/Boolean; + public final fun isEnabled ()Z + public final fun liveBytes (Lsk/ainet/lang/memory/ScopeKind;)J + public final fun liveEntries ()Ljava/util/List; + public final fun peakBytes (Lsk/ainet/lang/memory/ScopeKind;)J + public final fun recordAdapter (Ljava/lang/String;JLsk/ainet/lang/tensor/TensorId;)V + public static synthetic fun recordAdapter$default (Lsk/ainet/lang/memory/MemoryDebug;Ljava/lang/String;JLsk/ainet/lang/tensor/TensorId;ILjava/lang/Object;)V + public final fun report ()Ljava/lang/String; + public final fun reportLeaks (Ljava/util/List;)Ljava/util/List; + public final fun reset ()V + public final fun setOnAdapter (Lkotlin/jvm/functions/Function3;)V + public final fun setOnAllocate (Lkotlin/jvm/functions/Function1;)V + public final fun setOnClose (Lkotlin/jvm/functions/Function1;)V + public final fun setOnLeak (Lkotlin/jvm/functions/Function1;)V + public final fun setOverrideEnabled (Ljava/lang/Boolean;)V +} + +public final class sk/ainet/lang/memory/MemoryDebug$Entry { + public synthetic fun (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Ljava/lang/String;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public synthetic fun (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Ljava/lang/String;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun component1-TPZW6QE ()J + public final fun component2 ()Lsk/ainet/lang/memory/ScopeKind; + public final fun component3 ()J + public final fun component4 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component5 ()Ljava/lang/String; + public final fun component6 ()Ljava/lang/String; + public final fun component7 ()Z + public final fun copy-IscwgJM (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Ljava/lang/String;Z)Lsk/ainet/lang/memory/MemoryDebug$Entry; + public static synthetic fun copy-IscwgJM$default (Lsk/ainet/lang/memory/MemoryDebug$Entry;JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Ljava/lang/String;ZILjava/lang/Object;)Lsk/ainet/lang/memory/MemoryDebug$Entry; + public fun equals (Ljava/lang/Object;)Z + public final fun getAllocationSite ()Ljava/lang/String; + public final fun getBytes ()J + public final fun getClosed ()Z + public final fun getClosingSite ()Ljava/lang/String; + public final fun getOrigin ()Lsk/ainet/lang/tensor/TensorId; + public final fun getScope ()Lsk/ainet/lang/memory/ScopeKind; + public final fun getStorageId-TPZW6QE ()J + public fun hashCode ()I + public final fun setClosed (Z)V + public final fun setClosingSite (Ljava/lang/String;)V + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/memory/ModelScope : sk/ainet/lang/memory/Scope { public fun ()V public fun (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;)V diff --git a/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/MemoryDebug.android.kt b/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/MemoryDebug.android.kt new file mode 100644 index 000000000..033d6f69c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/androidMain/kotlin/sk/ainet/lang/memory/MemoryDebug.android.kt @@ -0,0 +1,14 @@ +package sk.ainet.lang.memory + +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = + System.getProperty(MemoryDebug.PROPERTY) == "true" || System.getenv(MemoryDebug.ENV).let { it == "1" || it == "true" } + +/** The first frame outside `sk.ainet.lang.memory` — where the allocation was actually asked for. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = StackWalker.getInstance().walk { frames -> + frames.filter { !it.className.startsWith("sk.ainet.lang.memory") } + .findFirst() + .map { "${it.className}.${it.methodName}(${it.fileName}:${it.lineNumber})" } + .orElse(null) +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/MemoryDebug.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/MemoryDebug.kt new file mode 100644 index 000000000..798a5ee63 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/MemoryDebug.kt @@ -0,0 +1,153 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.TensorId + +/** + * Debug mode for the memory model (SKEEP-003 §4.7, §8 item 4, PRD M1-F9): the switch that turns + * `SKAINET_MEMORY_DEBUG=1` (or `-Dskainet.memory.debug=true`) into allocation-site tagging, + * use-after-close reporting with the closing site, adapter logging and per-scope high-water marks. + * + * *The tool that would have found #782 in one run.* Off by default and costing one boolean read + * per allocation; on, it keeps a bounded ledger of live storages plus hooks a debugger or a test + * can break on. + */ +@ExperimentalMemoryApi +public object MemoryDebug { + + /** Environment variable and system property that enable debug mode. */ + public const val ENV: String = "SKAINET_MEMORY_DEBUG" + public const val PROPERTY: String = "skainet.memory.debug" + + /** Explicit override; `null` means "read the platform setting". Tests set this. */ + public var overrideEnabled: Boolean? = null + + /** Whether debug mode is on. */ + public val isEnabled: Boolean get() = overrideEnabled ?: platformMemoryDebugEnabled() + + /** What a storage's ledger entry remembers. */ + public data class Entry( + val storageId: StorageId, + val scope: ScopeKind, + val bytes: Long, + val origin: TensorId?, + val allocationSite: String?, + var closingSite: String? = null, + var closed: Boolean = false, + ) + + private val ledger = LinkedHashMap() + private val peakByScope = HashMap() + private val liveByScope = HashMap() + + // --- hooks (SKEEP-003 §4.7 "breakpoint-like hooks") --- + + /** Called on every allocation while debug mode is on; throw from it to break in a debugger. */ + public var onAllocate: ((Entry) -> Unit)? = null + + /** Called when a storage is closed. */ + public var onClose: ((Entry) -> Unit)? = null + + /** Called when the dispatcher inserts an adapter; `bytes` is what it allocated. */ + public var onAdapter: ((kind: String, bytes: Long, target: TensorId?) -> Unit)? = null + + /** Called when a `Forward` scope is reset with storages still referenced from outside — a leak. */ + public var onLeak: ((List) -> Unit)? = null + + // --- recording (called by Storage / Scope / the dispatcher; no-ops when disabled) --- + + internal fun recordAllocation(id: StorageId, scope: ScopeKind, bytes: Long, origin: TensorId?, site: String?) { + if (!isEnabled) return + val e = Entry(id, scope, bytes, origin, site) + ledger[id.value] = e + val live = (liveByScope[scope] ?: 0L) + bytes + liveByScope[scope] = live + if (live > (peakByScope[scope] ?: 0L)) peakByScope[scope] = live + onAllocate?.invoke(e) + } + + internal fun recordClose(id: StorageId, site: String?) { + if (!isEnabled) return + val e = ledger[id.value] ?: return + if (e.closed) return + e.closed = true; e.closingSite = site + liveByScope[e.scope] = ((liveByScope[e.scope] ?: 0L) - e.bytes).coerceAtLeast(0L) + onClose?.invoke(e) + } + + /** Report an adapter insertion (kind, bytes, target) — the "silent dequant budget" of #782. */ + public fun recordAdapter(kind: String, bytes: Long, target: TensorId? = null) { + if (!isEnabled) return + onAdapter?.invoke(kind, bytes, target) + } + + /** Report storages that outlived a `Forward` reset; the hook fires and the entries are returned. */ + public fun reportLeaks(stillReferenced: List): List { + if (!isEnabled || stillReferenced.isEmpty()) return emptyList() + val entries = stillReferenced.mapNotNull { ledger[it.value] }.filterNot { it.closed } + if (entries.isNotEmpty()) onLeak?.invoke(entries) + return entries + } + + // --- inspection --- + + /** The ledger entry of a storage, if debug mode saw it. */ + public fun entry(id: StorageId): Entry? = ledger[id.value] + + /** Storages allocated and not yet closed. */ + public fun liveEntries(): List = ledger.values.filterNot { it.closed } + + /** Per-scope high-water mark since the last [reset]. */ + public fun peakBytes(scope: ScopeKind): Long = peakByScope[scope] ?: 0L + + /** Live bytes per scope right now. */ + public fun liveBytes(scope: ScopeKind): Long = liveByScope[scope] ?: 0L + + /** + * A use-after-close message that names the closing site — the difference between "something is + * closed" and "this weight was freed at model.close(), you are reading it from the KV path". + */ + public fun describeClosed(id: StorageId): String { + val e = ledger[id.value] ?: return "storage $id (no debug record; set $ENV=1 before allocating)" + return buildString { + append("storage "); append(id) + e.origin?.let { append(" ("); append(it.canonical); append(')') } + append(" of "); append(e.bytes); append(" B in "); append(e.scope.name.lowercase()); append(" scope") + e.allocationSite?.let { append("\n allocated at "); append(it) } + e.closingSite?.let { append("\n closed at "); append(it) } + } + } + + /** Everything the ledger knows, for a debugger watch window or a failing test. */ + public fun report(): String = buildString { + append("SKaiNET memory debug — "); append(ledger.size); append(" storages seen, ") + append(liveEntries().size); append(" live\n") + for (scope in ScopeKind.entries) { + val peak = peakBytes(scope); val live = liveBytes(scope) + if (peak == 0L && live == 0L) continue + append(" "); append(scope.name.lowercase().padEnd(8)) + append("live "); append(live); append(" B, peak "); append(peak); append(" B\n") + } + for (e in liveEntries().take(MAX_REPORTED)) { + append(" live "); append(e.storageId); append(' ') + append(e.origin?.canonical ?: "—"); append(' '); append(e.bytes); append(" B") + e.allocationSite?.let { append(" @ "); append(it) } + append('\n') + } + if (liveEntries().size > MAX_REPORTED) { append(" … "); append(liveEntries().size - MAX_REPORTED); append(" more\n") } + } + + /** Forget everything (between tests, or between models). */ + public fun reset() { + ledger.clear(); peakByScope.clear(); liveByScope.clear() + } + + private const val MAX_REPORTED: Int = 20 +} + +/** Platform reading of [MemoryDebug.ENV] / [MemoryDebug.PROPERTY]. */ +@ExperimentalMemoryApi +internal expect fun platformMemoryDebugEnabled(): Boolean + +/** The current call site, when the platform can cheaply produce one (JVM: a stack frame). */ +@ExperimentalMemoryApi +internal expect fun platformCallSite(): String? diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt index 0173d79b6..168a4279b 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Storage.kt @@ -94,7 +94,13 @@ public sealed class Storage : AutoCloseable { public abstract val isMutable: Boolean /** Throws [StorageClosedException] if this storage is no longer alive. Called by every accessor. */ - public fun checkAlive() { if (!isAlive) throw StorageClosedException(id, debugOrigin) } + public fun checkAlive() { + if (isAlive) return + // In debug mode the ledger knows where this storage was allocated and closed, which is the + // difference between "something is closed" and "this weight was freed at model.close()". + val detail = if (MemoryDebug.isEnabled) "\n" + MemoryDebug.describeClosed(id) else "" + throw StorageClosedException(id, debugOrigin, "Storage $id${debugOrigin?.let { " (" + it.canonical + ")" } ?: ""} is closed$detail") + } /** * Close: an [Owner.Owned] storage releases its bytes (exactly once); an [Owner.Borrowed] storage @@ -105,7 +111,10 @@ public sealed class Storage : AutoCloseable { if (closed) return closed = true onClose() - if (sink.isEnabled && owner !is Owner.Alias) sink.emit(TraceEvent.Free(id.value, scope, sizeBytes)) + if (owner !is Owner.Alias) { + if (sink.isEnabled) sink.emit(TraceEvent.Free(id.value, scope, sizeBytes)) + MemoryDebug.recordClose(id, if (MemoryDebug.isEnabled) platformCallSite() else null) + } } /** Release platform resources (owned storage only); default nothing. */ @@ -155,7 +164,11 @@ public sealed class Storage : AutoCloseable { private fun create(floats: FloatArray?, ints: IntArray?, bytes: ByteArray?, offset: Int, count: Int, owner: Owner, origin: TensorId?, sink: TraceSink, mutable: Boolean): Heap { val eb = if (bytes != null) 1 else 4 val s = Heap(StorageId.next(), floats, ints, bytes, offset, count.toLong() * eb, owner, origin, sink, mutable) - if (sink.isEnabled && owner is Owner.Owned) sink.emit(TraceEvent.Allocation(s.id.value, owner.scope, s.sizeBytes, origin)) + if (owner is Owner.Owned) { + val site = if (MemoryDebug.isEnabled) platformCallSite() else null + if (sink.isEnabled) sink.emit(TraceEvent.Allocation(s.id.value, owner.scope, s.sizeBytes, origin, site)) + MemoryDebug.recordAllocation(s.id, owner.scope, s.sizeBytes, origin, site) + } return s } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/MemoryDebugTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/MemoryDebugTest.kt new file mode 100644 index 000000000..c66e8f624 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/memory/MemoryDebugTest.kt @@ -0,0 +1,109 @@ +package sk.ainet.lang.memory + +import sk.ainet.lang.tensor.TensorId +import kotlin.test.AfterTest +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** SKEEP-003 §4.7 / §8 item 4, PRD M1-F9: the debug mode that would have found #782 in one run. */ +@OptIn(ExperimentalMemoryApi::class) +class MemoryDebugTest { + + @BeforeTest fun on() { MemoryDebug.overrideEnabled = true; MemoryDebug.reset() } + + @AfterTest fun off() { + MemoryDebug.onAllocate = null; MemoryDebug.onClose = null; MemoryDebug.onAdapter = null; MemoryDebug.onLeak = null + MemoryDebug.reset(); MemoryDebug.overrideEnabled = null + } + + @Test + fun offByDefaultAndCostsNothing() { + MemoryDebug.overrideEnabled = null + // the platform default is off unless the env/property says otherwise + val before = MemoryDebug.liveEntries().size + Storage.Heap.floats(4).close() + assertEquals(before, MemoryDebug.liveEntries().size, "nothing is recorded while debug mode is off") + } + + @Test + fun allocationsAreTaggedWithScopeOriginAndSite() { + val id = TensorId.parse("model.layers[3].mlp.down_proj.weight") + var seen: MemoryDebug.Entry? = null + MemoryDebug.onAllocate = { seen = it } + val s = Storage.Heap.floats(64, ScopeKind.MODEL, origin = id) + val e = assertNotNull(seen) + assertEquals(s.id, e.storageId); assertEquals(ScopeKind.MODEL, e.scope); assertEquals(256L, e.bytes) + assertEquals(id, e.origin); assertFalse(e.closed) + assertEquals(e, MemoryDebug.entry(s.id)) + assertEquals(256L, MemoryDebug.liveBytes(ScopeKind.MODEL)) + assertEquals(256L, MemoryDebug.peakBytes(ScopeKind.MODEL)) + assertTrue(MemoryDebug.liveEntries().any { it.storageId == s.id }) + s.close() + assertTrue(assertNotNull(MemoryDebug.entry(s.id)).closed) + assertEquals(0L, MemoryDebug.liveBytes(ScopeKind.MODEL)) + assertEquals(256L, MemoryDebug.peakBytes(ScopeKind.MODEL), "the peak survives the free") + } + + @Test + fun useAfterCloseNamesTheStorageAndItsOrigin() { + val id = TensorId.parse("model.embed_tokens.weight") + val s = Storage.Heap.floats(8, ScopeKind.MODEL, origin = id) + s.close() + val ex = assertFailsWith { s.checkAlive() } + val msg = assertNotNull(ex.message) + assertTrue(msg.contains("model.embed_tokens.weight"), msg) + assertTrue(msg.contains("model scope"), msg) + assertTrue(msg.contains("32 B"), msg) + } + + @Test + fun closeHookFiresOnce() { + var closes = 0 + MemoryDebug.onClose = { closes++ } + val s = Storage.Heap.bytes(16) + s.close(); s.close() + assertEquals(1, closes) + } + + @Test + fun adapterInsertionsAreReported() { + var kind: String? = null; var bytes = 0L; var target: TensorId? = null + MemoryDebug.onAdapter = { k, b, t -> kind = k; bytes = b; target = t } + MemoryDebug.recordAdapter("dequantize", 96L * 1024 * 1024, TensorId.parse("model.layers[3].mlp.down_proj.weight")) + assertEquals("dequantize", kind); assertEquals(96L * 1024 * 1024, bytes) + assertEquals("model.layers[3].mlp.down_proj.weight", assertNotNull(target).canonical) + } + + @Test + fun leakReportingNamesTheStoragesThatOutlivedAReset() { + var leaked: List = emptyList() + MemoryDebug.onLeak = { leaked = it } + val escaped = Storage.Heap.floats(4, ScopeKind.FORWARD, origin = TensorId.parse("model.act#step=1")) + val closed = Storage.Heap.floats(4, ScopeKind.FORWARD) + closed.close() + val reported = MemoryDebug.reportLeaks(listOf(escaped.id, closed.id)) + assertEquals(1, reported.size, "only the still-open storage is a leak") + assertEquals(escaped.id, reported.single().storageId) + assertEquals(reported, leaked) + assertTrue(MemoryDebug.reportLeaks(emptyList()).isEmpty()) + } + + @Test + fun theReportShowsPeaksAndLiveStorages() { + Storage.Heap.floats(32, ScopeKind.MODEL, origin = TensorId.parse("model.w")) + val tmp = Storage.Heap.floats(8, ScopeKind.FORWARD); tmp.close() + val text = MemoryDebug.report() + assertTrue(text.contains("storages seen")); assertTrue(text.contains("model")) + assertTrue(text.contains("model.w"), text) + assertTrue(text.contains("peak")) + MemoryDebug.reset() + assertEquals(0L, MemoryDebug.peakBytes(ScopeKind.MODEL)) + assertNull(MemoryDebug.entry(StorageId(1))) + } +} diff --git a/skainet-lang/skainet-lang-core/src/jsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt b/skainet-lang/skainet-lang-core/src/jsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt new file mode 100644 index 000000000..4325b0d66 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt @@ -0,0 +1,9 @@ +package sk.ainet.lang.memory + +/** No property store here: debug mode is opt-in through [MemoryDebug.overrideEnabled]. */ +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = false + +/** No cheap stack walk on this target. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = null diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/MemoryDebug.jvm.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/MemoryDebug.jvm.kt new file mode 100644 index 000000000..033d6f69c --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/lang/memory/MemoryDebug.jvm.kt @@ -0,0 +1,14 @@ +package sk.ainet.lang.memory + +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = + System.getProperty(MemoryDebug.PROPERTY) == "true" || System.getenv(MemoryDebug.ENV).let { it == "1" || it == "true" } + +/** The first frame outside `sk.ainet.lang.memory` — where the allocation was actually asked for. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = StackWalker.getInstance().walk { frames -> + frames.filter { !it.className.startsWith("sk.ainet.lang.memory") } + .findFirst() + .map { "${it.className}.${it.methodName}(${it.fileName}:${it.lineNumber})" } + .orElse(null) +} diff --git a/skainet-lang/skainet-lang-core/src/nativeMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt b/skainet-lang/skainet-lang-core/src/nativeMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt new file mode 100644 index 000000000..4325b0d66 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/nativeMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt @@ -0,0 +1,9 @@ +package sk.ainet.lang.memory + +/** No property store here: debug mode is opt-in through [MemoryDebug.overrideEnabled]. */ +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = false + +/** No cheap stack walk on this target. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = null diff --git a/skainet-lang/skainet-lang-core/src/wasmJsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt b/skainet-lang/skainet-lang-core/src/wasmJsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt new file mode 100644 index 000000000..4325b0d66 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/wasmJsMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt @@ -0,0 +1,9 @@ +package sk.ainet.lang.memory + +/** No property store here: debug mode is opt-in through [MemoryDebug.overrideEnabled]. */ +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = false + +/** No cheap stack walk on this target. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = null diff --git a/skainet-lang/skainet-lang-core/src/wasmWasiMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt b/skainet-lang/skainet-lang-core/src/wasmWasiMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt new file mode 100644 index 000000000..4325b0d66 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/wasmWasiMain/kotlin/sk/ainet/lang/memory/MemoryDebug.other.kt @@ -0,0 +1,9 @@ +package sk.ainet.lang.memory + +/** No property store here: debug mode is opt-in through [MemoryDebug.overrideEnabled]. */ +@ExperimentalMemoryApi +internal actual fun platformMemoryDebugEnabled(): Boolean = false + +/** No cheap stack walk on this target. */ +@ExperimentalMemoryApi +internal actual fun platformCallSite(): String? = null