feat(memory): SKAINET_MEMORY_DEBUG — allocation-site tagging, use-after-close with the closing site, adapter log, leak reports (SKEEP-003 P2, S1.6)

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 <init> (JLsk/ainet/lang/memory/ScopeKind;JLsk/ainet/lang/tensor/TensorId;Ljava/lang/String;Ljava/lang/String;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public synthetic fun <init> (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 <init> ()V
 	public fun <init> (Lsk/ainet/lang/memory/trace/TraceSink;Ljava/lang/String;)V

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)
+}

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<Long, Entry>()
+    private val peakByScope = HashMap<ScopeKind, Long>()
+    private val liveByScope = HashMap<ScopeKind, Long>()
+
+    // --- 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<Entry>) -> 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<StorageId>): List<Entry> {
+        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<Entry> = 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?

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
             }
 

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<StorageClosedException> { 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<MemoryDebug.Entry> = 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)))
+    }
+}

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