From c646c408553578a01119cbaf46d80e64757913aa Mon Sep 17 00:00:00 2001 From: Demetrios Agourakis Date: Fri, 11 Sep 2026 08:31:14 -0300 Subject: [PATCH 1/3] fix(streaming): stop layer discovery on IndexError too install_streaming_experts(num_layers=None) counts layers by probing spec.block_of(model, n) with increasing n until it fails, but it only catches AttributeError. Layers live in a list (MoESpec.block_of indexes "layers are plain lists in most families"), so running off the end raises IndexError and the call crashes instead of returning. The shipped engines always pass num_layers, which is why this never showed. Adds test_install_discovers_layer_count: a three-layer list-based model (one MoE layer, two dense) installed with num_layers=None. It raised IndexError before this change; now it returns [twin, None, None] and the twin is swapped into the MoE block. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0146Hf9MXz9hHfRnq3MNXXWX --- python/src/edge0/streaming/install.py | 2 +- python/tests/test_streaming_math.py | 30 +++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/python/src/edge0/streaming/install.py b/python/src/edge0/streaming/install.py index 37b0e7e..acc3585 100644 --- a/python/src/edge0/streaming/install.py +++ b/python/src/edge0/streaming/install.py @@ -41,7 +41,7 @@ def install_streaming_experts( while True: try: spec.block_of(model, n) - except AttributeError: + except (AttributeError, IndexError): # past the last layer break n += 1 if n == 0: diff --git a/python/tests/test_streaming_math.py b/python/tests/test_streaming_math.py index eeb0f19..660d289 100644 --- a/python/tests/test_streaming_math.py +++ b/python/tests/test_streaming_math.py @@ -344,3 +344,33 @@ def test_double_buffered_swap(layer): ref2 = lay(x, mx.array([second], dtype=mx.int32)) assert mx.allclose(out1, ref1).item() assert mx.allclose(out2, ref2).item() + + +def test_install_discovers_layer_count(tmp_path): + """install_streaming_experts(num_layers=None) probes block_path with + increasing layer indices until it stops resolving. Layers live in a + list, so running off the end raises IndexError, not AttributeError.""" + from types import SimpleNamespace + + from edge0.streaming.install import install_streaming_experts + + path = tmp_path / "w.safetensors" + _write_shard(path, fuse_gu=False) + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + quant=QuantSpec(bits=4, group_size=64), + layout=WeightLayout.SEPARATE, + key_template="layers.{layer}.mlp.switch_mlp", + block_path="layers.{layer}.mlp", + ) + moe = SimpleNamespace(switch_mlp=object()) + model = SimpleNamespace(layers=[SimpleNamespace(mlp=moe), + SimpleNamespace(mlp=SimpleNamespace()), + SimpleNamespace(mlp=SimpleNamespace())]) + twins = install_streaming_experts( + model, [SafetensorsMmap(str(path))], spec, options=_options()) + assert len(twins) == 3 + assert isinstance(twins[0], StreamingSwitchGLU) + assert twins[1] is None and twins[2] is None # dense layers + assert moe.switch_mlp is twins[0] + twins[0].close() From 85678c30cce88022924228166627b43eeeb6155a Mon Sep 17 00:00:00 2001 From: Demetrios Agourakis Date: Tue, 22 Sep 2026 17:11:29 -0300 Subject: [PATCH 2/3] fix(moe): isolate layer-index errors from other numeric segments in discovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two coverage nits from the PR review of the previous commit: - The AttributeError half of the widened except was untested: added test_layer_exists_stops_on_attribute_error, using an attribute-based layer container (no list, no __getitem__) to reach it independently of the IndexError case. - Catching (AttributeError, IndexError) around the whole block_of call could mask an IndexError raised by a different, non-layer-index numeric segment further down block_path (e.g. a fixed expert-slot index) — a real bug, not end-of-list. No shipped spec hits this, but it would have silently under-counted layers instead of surfacing the break. Fixes it by adding MoESpec.layer_exists(model, layer), which resolves block_path segment-by-segment and only treats an AttributeError/ IndexError as "past the last layer" when it comes from the segment templated by {layer} itself; any other segment's error propagates. install_streaming_experts's discovery loop now calls this instead of wrapping spec.block_of directly. test_layer_exists_reraises_unrelated_index_error covers the propagation case. Co-Authored-By: Claude Sonnet 5 --- python/src/edge0/moe/spec.py | 29 ++++++++++++++++++++ python/src/edge0/streaming/install.py | 6 +---- python/tests/test_streaming_math.py | 38 +++++++++++++++++++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/python/src/edge0/moe/spec.py b/python/src/edge0/moe/spec.py index ced4ac4..a39d403 100644 --- a/python/src/edge0/moe/spec.py +++ b/python/src/edge0/moe/spec.py @@ -110,6 +110,35 @@ def block_of(self, model, layer: int): obj = getattr(obj, part) return obj + def layer_exists(self, model, layer: int) -> bool: + """True if ``layer`` resolves at block_path's layer-index segment. + + Used by layer-count discovery (``install_streaming_experts`` with + ``num_layers=None``) to find where the layer list ends. Only an + ``AttributeError``/``IndexError`` raised while resolving the + segment templated by ``{layer}`` itself means "past the last + layer"; the same errors raised by a *different* segment further + down ``block_path`` (e.g. a fixed expert-slot index) indicate a + bug in the spec or model and are re-raised rather than read as + end-of-list. + """ + template_parts = self.block_path.split(".") + layer_pos = next( + i for i, p in enumerate(template_parts) if "{layer}" in p) + obj = model + for i, raw_part in enumerate(template_parts): + part = raw_part.format(layer=layer) + try: + if part.isdigit(): + obj = obj[int(part)] + else: + obj = getattr(obj, part) + except (AttributeError, IndexError): + if i == layer_pos: + return False + raise + return True + def layer_of(self, model, layer: int): """Resolve the decoder layer object (the block's owner). diff --git a/python/src/edge0/streaming/install.py b/python/src/edge0/streaming/install.py index acc3585..c0c84de 100644 --- a/python/src/edge0/streaming/install.py +++ b/python/src/edge0/streaming/install.py @@ -38,11 +38,7 @@ def install_streaming_experts( """ if num_layers is None: n = 0 - while True: - try: - spec.block_of(model, n) - except (AttributeError, IndexError): # past the last layer - break + while spec.layer_exists(model, n): n += 1 if n == 0: raise ValueError( diff --git a/python/tests/test_streaming_math.py b/python/tests/test_streaming_math.py index 660d289..b8aae01 100644 --- a/python/tests/test_streaming_math.py +++ b/python/tests/test_streaming_math.py @@ -374,3 +374,41 @@ def test_install_discovers_layer_count(tmp_path): assert twins[1] is None and twins[2] is None # dense layers assert moe.switch_mlp is twins[0] twins[0].close() + + +def test_layer_exists_stops_on_attribute_error(): + """The AttributeError half of ``layer_exists``'s except clause is + reachable independently of IndexError: a family whose layer container + is attribute-based (no list, no ``__getitem__``) runs off the end via + a plain missing attribute, not an out-of-range index.""" + from types import SimpleNamespace + + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + key_template="layer_{layer}.mlp.switch_mlp", + block_path="layer_{layer}", + ) + model = SimpleNamespace(layer_0=object(), layer_1=object()) + assert spec.layer_exists(model, 0) is True + assert spec.layer_exists(model, 1) is True + assert spec.layer_exists(model, 2) is False # no `layer_2` attribute + + +def test_layer_exists_reraises_unrelated_index_error(): + """An IndexError from a segment *other* than the ``{layer}`` slot is a + real bug (a malformed block_path or a broken model), not end-of-list, + and must propagate instead of being read as "past the last layer" -- + otherwise install_streaming_experts(num_layers=None) would silently + under-count layers instead of surfacing the break.""" + from types import SimpleNamespace + + spec = MoESpec( + num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, + key_template="layers.{layer}.experts.9.switch_mlp", + block_path="layers.{layer}.experts.9", + ) + # `layer=0` is in range for `layers`, but the fixed trailing index `9` + # is out of range for `experts` -- unrelated to layer-count discovery. + model = SimpleNamespace(layers=[SimpleNamespace(experts=[object()])]) + with pytest.raises(IndexError): + spec.layer_exists(model, 0) From a6ad384550aabc1f387a447a5410f10aaf44224c Mon Sep 17 00:00:00 2001 From: Demetrios Agourakis Date: Sun, 4 Oct 2026 09:28:34 -0300 Subject: [PATCH 3/3] test(streaming): cover layer discovery without tensor calculations --- python/tests/test_moe_spec.py | 37 +++++++++++++++++++++ python/tests/test_streaming_install.py | 46 ++++++++++++++++++++++++++ python/tests/test_streaming_math.py | 38 --------------------- 3 files changed, 83 insertions(+), 38 deletions(-) create mode 100644 python/tests/test_streaming_install.py diff --git a/python/tests/test_moe_spec.py b/python/tests/test_moe_spec.py index b3dd18f..34a0f13 100644 --- a/python/tests/test_moe_spec.py +++ b/python/tests/test_moe_spec.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest from edge0.moe.spec import (MoESpec, QuantSpec, RouterKind, WeightLayout) @@ -59,6 +61,41 @@ def test_block_of_digit_segments(): assert s.block_of(m, 0) is m.language_model.model.layers[0].mlp.switch_mlp +def test_layer_exists_stops_at_end_of_list(): + spec = _spec() + model = _FakeModel(n=2) + assert spec.layer_exists(model, 0) is True + assert spec.layer_exists(model, 1) is True + assert spec.layer_exists(model, 2) is False + assert spec.layer_exists(_FakeModel(n=0), 0) is False + + +def test_layer_exists_stops_on_attribute_error(): + spec = _spec(block_path="layer_{layer}") + model = SimpleNamespace(layer_0=object(), layer_1=object()) + assert spec.layer_exists(model, 0) is True + assert spec.layer_exists(model, 1) is True + assert spec.layer_exists(model, 2) is False + + +def test_layer_exists_reraises_unrelated_index_error(): + spec = _spec(block_path="layers.{layer}.experts.9") + model = SimpleNamespace(layers=[SimpleNamespace(experts=[object()])]) + with pytest.raises(IndexError): + spec.layer_exists(model, 0) + + +@pytest.mark.parametrize("block_path", [ + "missing.layers.{layer}.mlp", + "layers.{layer}.missing", +]) +def test_layer_exists_reraises_unrelated_attribute_error(block_path): + spec = _spec(block_path=block_path) + model = SimpleNamespace(layers=[SimpleNamespace(mlp=object())]) + with pytest.raises(AttributeError): + spec.layer_exists(model, 0) + + def test_layer_of_defaults_from_block_path(): s = _spec() m = _FakeModel() diff --git a/python/tests/test_streaming_install.py b/python/tests/test_streaming_install.py new file mode 100644 index 0000000..da88de3 --- /dev/null +++ b/python/tests/test_streaming_install.py @@ -0,0 +1,46 @@ +"""Layer discovery and installation without tensor calculations.""" + +from types import SimpleNamespace + +import pytest + +from edge0.moe.spec import MoESpec +from edge0.streaming import install + + +def _spec(): + return MoESpec( + num_experts=8, top_k=2, intermediate_size=64, + block_path="layers.{layer}.mlp", + ) + + +def test_install_discovers_mixed_layers(monkeypatch): + twin = object() + resident = object() + moe = SimpleNamespace(switch_mlp=resident) + dense = SimpleNamespace() + model = SimpleNamespace(layers=[ + SimpleNamespace(mlp=moe), SimpleNamespace(mlp=dense), + ]) + calls = [] + + def make_twin(shards, layer, spec, **kwargs): + calls.append((shards, layer, spec)) + return twin + + monkeypatch.setattr(install, "StreamingSwitchGLU", make_twin) + spec = _spec() + twins = install.install_streaming_experts(model, [], spec) + + assert twins == [twin, None] + assert calls == [([], 0, spec)] + assert moe.switch_mlp is twin + assert moe._edge0_resident_switch is resident + assert not hasattr(dense, "switch_mlp") + + +def test_install_rejects_empty_layer_list(): + model = SimpleNamespace(layers=[]) + with pytest.raises(ValueError, match="resolves no layer 0"): + install.install_streaming_experts(model, [], _spec()) diff --git a/python/tests/test_streaming_math.py b/python/tests/test_streaming_math.py index b8aae01..660d289 100644 --- a/python/tests/test_streaming_math.py +++ b/python/tests/test_streaming_math.py @@ -374,41 +374,3 @@ def test_install_discovers_layer_count(tmp_path): assert twins[1] is None and twins[2] is None # dense layers assert moe.switch_mlp is twins[0] twins[0].close() - - -def test_layer_exists_stops_on_attribute_error(): - """The AttributeError half of ``layer_exists``'s except clause is - reachable independently of IndexError: a family whose layer container - is attribute-based (no list, no ``__getitem__``) runs off the end via - a plain missing attribute, not an out-of-range index.""" - from types import SimpleNamespace - - spec = MoESpec( - num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, - key_template="layer_{layer}.mlp.switch_mlp", - block_path="layer_{layer}", - ) - model = SimpleNamespace(layer_0=object(), layer_1=object()) - assert spec.layer_exists(model, 0) is True - assert spec.layer_exists(model, 1) is True - assert spec.layer_exists(model, 2) is False # no `layer_2` attribute - - -def test_layer_exists_reraises_unrelated_index_error(): - """An IndexError from a segment *other* than the ``{layer}`` slot is a - real bug (a malformed block_path or a broken model), not end-of-list, - and must propagate instead of being read as "past the last layer" -- - otherwise install_streaming_experts(num_layers=None) would silently - under-count layers instead of surfacing the break.""" - from types import SimpleNamespace - - spec = MoESpec( - num_experts=N_EXPERTS, top_k=4, intermediate_size=INTER, - key_template="layers.{layer}.experts.9.switch_mlp", - block_path="layers.{layer}.experts.9", - ) - # `layer=0` is in range for `layers`, but the fixed trailing index `9` - # is out of range for `experts` -- unrelated to layer-count discovery. - model = SimpleNamespace(layers=[SimpleNamespace(experts=[object()])]) - with pytest.raises(IndexError): - spec.layer_exists(model, 0)