64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
"""Tiny host-import/lifetime receipts from the pinned MTPLX runtime; no models."""
|
|
import gc
|
|
import json
|
|
import mmap
|
|
import os
|
|
import weakref
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
|
|
def main():
|
|
assert mx.__version__ == "0.32.2"
|
|
page = os.sysconf("SC_PAGE_SIZE")
|
|
mx.set_cache_limit(4 * page)
|
|
for name, offset, size in (
|
|
("aligned", 0, page),
|
|
("unaligned", 1, page),
|
|
("short", 0, 17),
|
|
("empty", 0, 0),
|
|
):
|
|
mx.clear_cache()
|
|
baseline = mx.get_active_memory()
|
|
owner = mmap.mmap(-1, 2 * page)
|
|
source = np.ndarray((size,), dtype=np.uint8, buffer=owner, offset=offset)
|
|
source[:] = 37
|
|
pointer = source.__array_interface__["data"][0]
|
|
weak = weakref.ref(source)
|
|
# The codec calls this constructor: array.cpp passes copy=true to
|
|
# create_array, even though nd_array_to_mlx has a different default.
|
|
value = mx.array(source)
|
|
mx.eval(value)
|
|
exported = np.asarray(value)
|
|
adopted = bool(size and pointer == exported.__array_interface__["data"][0])
|
|
assert (exported == 37).all()
|
|
del exported, source
|
|
gc.collect()
|
|
retained = weak() is not None
|
|
assert not adopted and not retained
|
|
active = mx.get_active_memory() - baseline
|
|
view = value[:]
|
|
mx.eval(view)
|
|
del value
|
|
gc.collect()
|
|
retained_by_view = weak() is not None
|
|
assert not retained_by_view
|
|
assert (np.asarray(view) == 37).all()
|
|
del view
|
|
gc.collect()
|
|
mx.synchronize()
|
|
released = weak() is None
|
|
assert released
|
|
assert mx.get_active_memory() == baseline
|
|
cached = mx.get_cache_memory()
|
|
owner.close()
|
|
print(json.dumps(dict(event="host_import", name=name, page=page, size=size,
|
|
adopted=adopted, active=active, retained=retained,
|
|
retained_by_view=retained_by_view, released=released,
|
|
cached_after_release=cached)), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|