""" .. _l-example-plot-align-external-data-streaming: Benchmark streaming vs in-memory alignment of external data =========================================================== This example benchmarks two ways to re-align tensor offsets in an ONNX model that stores its weights as external data: 1. **Streaming approach** — :func:`onnx_light.onnx.align_external_data_streaming` rewrites a ``(src.onnx, src.data)`` pair into a new ``(dst.onnx, dst.data)`` pair with every tensor offset aligned to a power-of-two boundary, **without ever loading the weights in RAM**. Only the proto metadata is parsed; tensor bytes are streamed file-to-file (via ``splice(2)`` on Linux for zero-copy in kernel space). 2. **Usual in-memory approach** — :func:`onnx_light.onnx.load` followed by ``SerializeToFile`` with :attr:`onnx_light.onnx.SerializeOptions.alignment` set to the same boundary. This loads the full set of weights into memory before re-writing them. The example reports wall-clock time and peak process-resident memory usage for each approach. At equal alignment the two outputs are byte-equivalent for the tensor payloads — only the peak memory footprint differs. .. note:: Peak memory is sampled from the process's resident set size (RSS), not from :mod:`tracemalloc`. ``tracemalloc`` only tracks the Python allocator (``pymalloc``), so it misses the C++-owned buffers where onnx-light stores tensor ``raw_data`` — making both approaches appear similarly small and hiding the very overhead this example is meant to illustrate. RSS captures both Python and C++ allocations. """ import gc import os import shutil import threading import time import matplotlib.patches as mpatches import matplotlib.pyplot as plt import numpy as np import pandas import onnx_light.onnx as onnxl import onnx_light.onnx.helper as oh import onnx_light.onnx.numpy_helper as onh # %% # Build a synthetic model # ----------------------- # # A handful of large initializers is enough to make the in-memory peak # noticeably larger than the streaming peak. The example runs with a # smaller model under ``UNITTEST_GOING=1`` to keep CI fast. N_INIT = 8 if os.environ.get("UNITTEST_GOING") == "1" else 24 DIM = 256 if os.environ.get("UNITTEST_GOING") == "1" else 1024 ALIGNMENT = 4096 def make_model(n_init: int = N_INIT, dim: int = DIM) -> onnxl.ModelProto: """Creates a synthetic ONNX model with *n_init* float32 Gemm weights.""" initializers = [] nodes = [] inputs = [oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [None, dim])] prev = "X" for i in range(n_init): weight_name = f"W{i}" out_name = f"Y{i}" w = np.random.randn(dim, dim).astype(np.float32) initializers.append(onh.from_array(w, name=weight_name)) nodes.append(oh.make_node("Gemm", [prev, weight_name], [out_name], transB=1)) prev = out_name outputs = [oh.make_tensor_value_info(prev, onnxl.TensorProto.FLOAT, [None, dim])] graph = oh.make_graph(nodes, "bench_graph", inputs, outputs, initializer=initializers) return oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=9) out_dir = "temp_plot_align_external_data_streaming" os.makedirs(out_dir, exist_ok=True) model = make_model() total_weight_bytes = sum(int(np.prod(init.dims)) * 4 for init in model.graph.initializer) print(f"Number of initializers : {len(model.graph.initializer)}") print(f"Total weight bytes : {total_weight_bytes / 2 ** 20:.2f} MB") print(f"Alignment : {ALIGNMENT} bytes") # %% # Save the source two-file model # ------------------------------ # # Every benchmark below starts from the same on-disk source pair # ``(src.onnx, src.data)``. src_onnx = os.path.join(out_dir, "src.onnx") src_data = src_onnx + ".data" def _save_source(proto: onnxl.ModelProto) -> None: """Writes the synthetic model to a two-file ``(.onnx, .data)`` pair.""" onnxl.save(proto, src_onnx, location=src_data) _save_source(model) print(f"Source two-file model : {src_onnx} + {src_data}") print(f"Source weights size : {os.path.getsize(src_data) / 2 ** 20:.2f} MB") # Release the in-memory model now that the source files are on disk. # Keeping ``model`` alive would inflate the RSS baseline for both benchmarks, # making the streaming variant appear to use more memory than it actually does. # After this point all benchmarks start from a clean low-memory baseline. del model gc.collect() # %% # Helpers # ------- # # Each benchmark is wrapped in a function that returns the elapsed wall # time and the peak process-resident memory (RSS) observed while the # function runs. RSS is sampled from ``/proc/self/statm`` (Linux) on a # background thread; on platforms where ``/proc`` is unavailable the # sampler falls back to :func:`resource.getrusage` (Unix) or reports # ``0``. Unlike :mod:`tracemalloc`, RSS includes the C++-owned buffers # where onnx-light stores tensor ``raw_data`` during an in-memory load # — which is exactly the overhead the streaming variant avoids. _PAGE_SIZE = os.sysconf("SC_PAGE_SIZE") if hasattr(os, "sysconf") else 4096 def _read_rss_bytes() -> int: """Returns the current process resident set size in bytes (0 if unknown).""" try: with open("/proc/self/statm", encoding="ascii") as stream: return int(stream.read().split()[1]) * _PAGE_SIZE except OSError: pass try: import resource ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # Linux reports KiB, macOS reports bytes; both branches are >= our usage scale. return ru * 1024 if ru < 2**40 else ru except (ImportError, OSError): return 0 class _PeakRssSampler: """Samples process RSS on a background thread and records the peak observed.""" def __init__(self, interval: float = 0.001): self._interval = interval self._stop = threading.Event() self._thread: threading.Thread | None = None self.baseline = 0 self.peak = 0 def __enter__(self): self.baseline = _read_rss_bytes() self.peak = self.baseline self._thread = threading.Thread(target=self._run, daemon=True) self._thread.start() return self def __exit__(self, *_exc) -> None: self._stop.set() if self._thread is not None: self._thread.join() cur = _read_rss_bytes() if cur > self.peak: self.peak = cur def _run(self) -> None: while not self._stop.wait(self._interval): cur = _read_rss_bytes() if cur > self.peak: self.peak = cur def measure(name: str, fn) -> dict: """Runs *fn* and returns elapsed time (s) and peak RSS delta (MB).""" with _PeakRssSampler() as sampler: t0 = time.perf_counter() fn() elapsed = time.perf_counter() - t0 peak_mb = max(sampler.peak - sampler.baseline, 0) / 2**20 print(f"{name:<40} time={elapsed * 1e3:8.1f} ms peak={peak_mb:7.2f} MB") return {"name": name, "time_s": elapsed, "peak_mb": peak_mb} def _flush(path: str) -> None: """fsync a file so timings include write-back to disk.""" with open(path, "r+b") as stream: stream.flush() os.fsync(stream.fileno()) # %% # Approach 1 — streaming alignment # -------------------------------- # # The source ``.onnx`` is parsed with ``skip_raw_data=True``, so only the # metadata (a few KB) reaches RAM; tensor bytes are streamed # file-to-file in chunks of ``chunk_size`` bytes, zero-padding to the # requested alignment between tensors. streaming_onnx = os.path.join(out_dir, "streaming.onnx") streaming_data = os.path.join(out_dir, "streaming.data") def run_streaming() -> None: onnxl.align_external_data_streaming( src_onnx_path=src_onnx, dst_onnx_path=streaming_onnx, dst_weights_path=streaming_data, alignment=ALIGNMENT, chunk_size=4 * 1024 * 1024, ) _flush(streaming_data) _flush(streaming_onnx) # %% # Approach 2 — in-memory load + re-save with alignment # ---------------------------------------------------- # # The "usual" path: load the full model (weights and all) into RAM, then # re-serialize it with :attr:`SerializeOptions.alignment` set to the # same boundary. inmem_onnx = os.path.join(out_dir, "inmem.onnx") inmem_data = os.path.join(out_dir, "inmem.data") def run_in_memory() -> None: loaded = onnxl.load(src_onnx, load_external_data=True, num_threads=1) opts = onnxl.SerializeOptions() opts.alignment = ALIGNMENT # Disable honouring existing external_data.location entries so every # tensor is written into the destination weights file passed below. opts.use_external_data_location = False loaded.SerializeToFile(inmem_onnx, opts, inmem_data) _flush(inmem_data) _flush(inmem_onnx) # %% # Run the benchmarks # ------------------ results = [] results.append(measure("align/streaming", run_streaming)) results.append(measure("align/in-memory", run_in_memory)) # %% # Verify byte-equivalence of the aligned outputs # ---------------------------------------------- # # Both approaches must produce the same tensor payloads at the same # aligned offsets. We compare each tensor's bytes read from both # destination weights files. streaming_loaded = onnxl.load(streaming_onnx, load_external_data=False) inmem_loaded = onnxl.load(inmem_onnx, load_external_data=False) with open(streaming_data, "rb") as f_stream, open(inmem_data, "rb") as f_inmem: for s_init, m_init in zip(streaming_loaded.graph.initializer, inmem_loaded.graph.initializer): s_meta = {e.key: e.value for e in s_init.external_data} m_meta = {e.key: e.value for e in m_init.external_data} s_off, s_len = int(s_meta["offset"]), int(s_meta["length"]) m_off, m_len = int(m_meta["offset"]), int(m_meta["length"]) assert s_off % ALIGNMENT == 0, f"streaming offset {s_off} not aligned" assert m_off % ALIGNMENT == 0, f"in-memory offset {m_off} not aligned" assert s_len == m_len, f"length mismatch: {s_len} vs {m_len}" f_stream.seek(s_off) f_inmem.seek(m_off) assert f_stream.read(s_len) == f_inmem.read(m_len), f"payload mismatch on {s_init.name}" print("Both approaches produced byte-equivalent aligned tensor payloads.") # %% # Results # ------- df = pandas.DataFrame(results).set_index("name").sort_index() print(df) # %% # Plot # ---- # # The streaming variant trades a slightly different I/O pattern for an # **O(metadata) + chunk_size** memory footprint, independent of the # total weights size. The in-memory variant peaks proportionally to # the model's weight payload. fig, (ax_time, ax_mem) = plt.subplots(1, 2, figsize=(12, 5)) row_names = df.index.tolist() colors = ["darkorange" if "streaming" in n else "steelblue" for n in row_names] ax_time.barh(row_names, df["time_s"], color=colors) ax_time.set_title("Wall time (s, lower is better)") ax_time.set_xlabel("seconds") ax_time.grid(axis="x") ax_mem.barh(row_names, df["peak_mb"], color=colors) ax_mem.set_title("Peak RSS delta (MB, lower is better)") ax_mem.set_xlabel("MB (resident set size)") ax_mem.grid(axis="x") fig.suptitle( f"Alignment={ALIGNMENT}B weights={total_weight_bytes / 2 ** 20:.1f} MB" f" #initializers={N_INIT}" ) fig.legend( handles=[ mpatches.Patch(color="darkorange", label="streaming"), mpatches.Patch(color="steelblue", label="in-memory"), ], loc="lower center", ncol=2, ) fig.tight_layout(rect=(0, 0.05, 1, 1)) fig.savefig("plot_align_external_data_streaming.png") # %% # Cleanup # ------- shutil.rmtree(out_dir, ignore_errors=True)