diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index a79c61926..78158a9c8 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -491,7 +491,7 @@ def malloc(self, worker_id: int, size: int) -> int: return int(self._o.malloc(wid, sz)) with self._worker._child_prov_lock: ptr = int(self._o.malloc(wid, sz)) - self._worker._child_prov_record_malloc(wid, ptr) + self._worker._child_prov_record_malloc(wid, ptr, sz) return ptr def free(self, worker_id: int, ptr: int) -> None: @@ -514,13 +514,13 @@ def free(self, worker_id: int, ptr: int) -> None: def copy_to(self, worker_id: int, dst: int, src: int, size: int) -> None: """Copy *size* bytes from host *src* to worker *dst*.""" - wid, d = int(worker_id), int(dst) + wid, d, nbytes = int(worker_id), int(dst), int(size) if self._worker is None: - self._o.copy_to(wid, d, int(src), int(size)) + self._o.copy_to(wid, d, int(src), nbytes) return with self._worker._child_prov_lock: - self._worker._child_prov_require_live(wid, d, api="copy_to") - self._o.copy_to(wid, d, int(src), int(size)) + self._worker._child_prov_require_copy_range(wid, d, nbytes, api="copy_to") + self._o.copy_to(wid, d, int(src), nbytes) def copy_from(self, worker_id: int, dst: int, src: int, size: int) -> None: """Copy *size* bytes from worker *src* to host *dst*.""" diff --git a/python/simpler/worker.py b/python/simpler/worker.py index b5d4be7b6..0744ed6d6 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -455,24 +455,28 @@ class _ChildProvEntry: Typed rather than a bare presence bit because the same ``(worker_id, ptr)`` can carry more than one role at once: a ``malloc`` base and a CommDomain window / carved buffer pointer can legally alias the same device address. - The key is live while ``malloc_owned or domain_allocation_ids``; only an - exact ``malloc`` base is ``free``-able, while a domain pointer is revoked by - its domain's release. Interior pointers are never recorded, so a pointer - that merely lands inside a live allocation has no entry and is rejected. + Each role records its accessible extent from the exact base. Only a malloc + base is ``free``-able, while a domain pointer is revoked by its domain's + release. Interior pointers are never recorded as independent bases. """ - __slots__ = ("malloc_owned", "domain_allocation_ids") + __slots__ = ("domain_allocation_nbytes", "malloc_nbytes") def __init__(self) -> None: - self.malloc_owned: bool = False - self.domain_allocation_ids: set[int] = set() + self.malloc_nbytes: int | None = None + self.domain_allocation_nbytes: dict[int, int] = {} def is_live(self) -> bool: """True iff this entry still carries a role. A role-less entry is dead — live checks are fail-closed on this, never on key presence alone, so an entry momentarily left empty (e.g. an interrupted revoke) never re-authorizes a freed pointer.""" - return self.malloc_owned or bool(self.domain_allocation_ids) + return self.malloc_nbytes is not None or bool(self.domain_allocation_nbytes) + + def live_nbytes(self) -> int: + """Largest accessible extent carried by a live role at this base.""" + malloc_nbytes = self.malloc_nbytes if self.malloc_nbytes is not None else 0 + return max(malloc_nbytes, max(self.domain_allocation_nbytes.values(), default=0)) @dataclass @@ -5340,9 +5344,19 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm # not by the deferred marker — so the deferred window stays dispatchable. with self._child_prov_lock: for chip_idx, ctx in contexts.items(): - self._child_prov_record_domain(chip_idx, int(ctx.local_window_base), allocation_id) - for buf_ptr in ctx.buffer_ptrs.values(): - self._child_prov_record_domain(chip_idx, int(buf_ptr), allocation_id) + self._child_prov_record_domain( + chip_idx, + int(ctx.local_window_base), + allocation_id, + int(ctx.actual_window_size), + ) + for buffer in buffers: + self._child_prov_record_domain( + chip_idx, + int(ctx.buffer_ptrs[buffer.name]), + allocation_id, + int(buffer.nbytes), + ) return handle def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResources) -> None: @@ -5605,17 +5619,15 @@ def _release_all_live_domains(self, resources: _RunResources | None = None) -> N # successful native alloc; revoke before the native free. # ------------------------------------------------------------------ - def _child_prov_record_malloc(self, worker_id: int, ptr: int) -> None: + def _child_prov_record_malloc(self, worker_id: int, ptr: int, nbytes: int) -> None: """Mark ``(worker_id, ptr)`` as a live malloc base (after a successful malloc).""" entry = self._child_alloc_prov.get((worker_id, ptr)) if entry is None: - # Fully initialise the role BEFORE inserting, so the dict never holds - # a role-less (dead) entry even if an async unwind lands here. entry = _ChildProvEntry() - entry.malloc_owned = True + entry.malloc_nbytes = nbytes self._child_alloc_prov[(worker_id, ptr)] = entry else: - entry.malloc_owned = True + entry.malloc_nbytes = nbytes def _child_prov_require_malloc_base(self, worker_id: int, ptr: int, *, api: str) -> None: """Require ``(worker_id, ptr)`` to be an exact live malloc base (freeable). @@ -5625,7 +5637,7 @@ def _child_prov_require_malloc_base(self, worker_id: int, ptr: int, *, api: str) by ``free``). """ entry = self._child_alloc_prov.get((worker_id, ptr)) - if entry is None or not entry.malloc_owned: + if entry is None or entry.malloc_nbytes is None: raise ValueError( f"Worker.{api}: device pointer 0x{ptr:x} is not a live malloc base on worker " f"{worker_id} (wrong worker, already-freed/stale, an interior pointer, or a " @@ -5639,10 +5651,10 @@ def _child_prov_clear_malloc(self, worker_id: int, ptr: int) -> None: entry = self._child_alloc_prov.get(key) if entry is None: return - if entry.domain_allocation_ids: - entry.malloc_owned = False # still live via a domain — keep the entry + if entry.domain_allocation_nbytes: + entry.malloc_nbytes = None else: - del self._child_alloc_prov[key] # last role — delete directly, no empty state + del self._child_alloc_prov[key] def _child_prov_require_live(self, worker_id: int, ptr: int, *, api: str) -> None: """Require ``(worker_id, ptr)`` to be a live child pointer (malloc or domain).""" @@ -5653,25 +5665,58 @@ def _child_prov_require_live(self, worker_id: int, ptr: int, *, api: str) -> Non f"{worker_id} (wrong worker, freed/stale, or an interior pointer)" ) - def _child_prov_record_domain(self, worker_id: int, ptr: int, allocation_id: int) -> None: - """Record a CommDomain window / buffer pointer at exact ``(worker_id, ptr)``.""" + def _child_prov_require_copy_range( + self, + worker_id: int, + ptr: int, + nbytes: int, + *, + api: str, + ) -> None: + """Require a copy range to be contained in a live child allocation.""" + if nbytes < 0: + raise ValueError(f"Worker.{api}: size must be non-negative, got {nbytes}") + exact = self._child_alloc_prov.get((worker_id, ptr)) + if exact is not None and exact.is_live() and nbytes <= exact.live_nbytes(): + return + for (owner, base), entry in self._child_alloc_prov.items(): + if owner != worker_id or ptr <= base or not entry.is_live(): + continue + live_nbytes = entry.live_nbytes() + offset = ptr - base + if offset <= live_nbytes and nbytes <= live_nbytes - offset: + return + raise ValueError( + f"Worker.{api}: device range [0x{ptr:x}, 0x{ptr + nbytes:x}) is not a live allocation " + f"on worker {worker_id} (wrong worker, freed/stale, or out of bounds)" + ) + + def _child_prov_record_domain( + self, + worker_id: int, + ptr: int, + allocation_id: int, + nbytes: int, + ) -> None: + """Record a CommDomain window or buffer base and its accessible extent.""" entry = self._child_alloc_prov.get((worker_id, ptr)) if entry is None: entry = _ChildProvEntry() self._child_alloc_prov[(worker_id, ptr)] = entry - entry.domain_allocation_ids.add(allocation_id) + prior_nbytes = entry.domain_allocation_nbytes.get(allocation_id, 0) + entry.domain_allocation_nbytes[allocation_id] = max(prior_nbytes, nbytes) def _child_prov_drop_domain(self, allocation_id: int) -> None: """Drop every pointer recorded by a CommDomain allocation (at the start of its physical release, before the backend free — see _release_domain_now).""" for key in list(self._child_alloc_prov): entry = self._child_alloc_prov[key] - if allocation_id not in entry.domain_allocation_ids: + if allocation_id not in entry.domain_allocation_nbytes: continue - if entry.malloc_owned or len(entry.domain_allocation_ids) > 1: - entry.domain_allocation_ids.discard(allocation_id) # other roles remain + if entry.malloc_nbytes is not None or len(entry.domain_allocation_nbytes) > 1: + del entry.domain_allocation_nbytes[allocation_id] else: - del self._child_alloc_prov[key] # last role — delete directly, no empty state + del self._child_alloc_prov[key] @staticmethod def _child_ptrs_in_args(args: Any) -> list[tuple[int, int]]: @@ -5738,7 +5783,7 @@ def malloc(self, size: int, worker_id: int = 0) -> int: # provenance is keyed on the canonical worker 0. with self._child_prov_lock: ptr = self._chip_worker.malloc(size) - self._child_prov_record_malloc(0, int(ptr)) + self._child_prov_record_malloc(0, int(ptr), int(size)) return ptr self._check_chip_worker_id(worker_id) assert self._orch is not None @@ -5767,7 +5812,7 @@ def copy_to(self, dst: int, src: int, size: int, worker_id: int = 0) -> None: if self.level == 2: assert self._chip_worker is not None with self._child_prov_lock: - self._child_prov_require_live(0, int(dst), api="copy_to") + self._child_prov_require_copy_range(0, int(dst), int(size), api="copy_to") self._chip_worker.copy_to(dst, src, size) return self._check_chip_worker_id(worker_id) diff --git a/tests/ut/py/test_worker/test_child_addr_guard.py b/tests/ut/py/test_worker/test_child_addr_guard.py index 5196443dc..c2e748a0c 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -40,9 +40,9 @@ def _child_args(ptr: int, *, n: int = 16) -> TaskArgs: return args -def _record_malloc(w: Worker, worker_id: int, ptr: int) -> None: +def _record_malloc(w: Worker, worker_id: int, ptr: int, nbytes: int = 64) -> None: with w._child_prov_lock: - w._child_prov_record_malloc(worker_id, ptr) + w._child_prov_record_malloc(worker_id, ptr, nbytes) # ---------------------------------------------------------------------------- @@ -54,7 +54,7 @@ class TestProvenanceTable: def test_malloc_base_is_live_then_freed(self): w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x1000) + w._child_prov_record_malloc(0, 0x1000, 64) w._child_prov_require_live(0, 0x1000, api="copy_to") # no raise w._child_prov_require_malloc_base(0, 0x1000, api="free") # no raise w._child_prov_clear_malloc(0, 0x1000) @@ -69,7 +69,7 @@ def test_free_of_unknown_pointer_rejected(self): def test_double_free_stale_before_reuse_rejected(self): w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x1000) + w._child_prov_record_malloc(0, 0x1000, 64) w._child_prov_clear_malloc(0, 0x1000) with pytest.raises(ValueError, match="already-freed/stale"): w._child_prov_require_malloc_base(0, 0x1000, api="free") @@ -77,7 +77,7 @@ def test_double_free_stale_before_reuse_rejected(self): def test_interior_pointer_is_not_a_live_entry(self): w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x1000) + w._child_prov_record_malloc(0, 0x1000, 64) # A pointer physically inside the allocation has no exact entry. with pytest.raises(ValueError, match="interior"): w._child_prov_require_live(0, 0x1008, api="copy_to") @@ -87,8 +87,8 @@ def test_same_numeric_va_on_two_workers_is_independent(self): # chips' identical numeric addresses independent. w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x4000) - w._child_prov_record_malloc(1, 0x4000) + w._child_prov_record_malloc(0, 0x4000, 64) + w._child_prov_record_malloc(1, 0x4000, 64) w._child_prov_clear_malloc(0, 0x4000) with pytest.raises(ValueError, match="not a live allocation"): w._child_prov_require_live(0, 0x4000, api="copy_to") @@ -97,7 +97,7 @@ def test_same_numeric_va_on_two_workers_is_independent(self): def test_domain_pointer_is_not_freeable_but_is_dispatchable(self): w = _l3() with w._child_prov_lock: - w._child_prov_record_domain(0, 0x8000, allocation_id=7) + w._child_prov_record_domain(0, 0x8000, allocation_id=7, nbytes=64) # usable for copy / dispatch w._child_prov_require_live(0, 0x8000, api="copy_to") # but not free-able — domains are revoked by release, not free() @@ -113,10 +113,13 @@ def test_malloc_and_domain_alias_same_pointer(self): # the domain removes it. w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x9000) - w._child_prov_record_domain(0, 0x9000, allocation_id=3) + w._child_prov_record_malloc(0, 0x9000, 128) + w._child_prov_record_domain(0, 0x9000, allocation_id=3, nbytes=64) + w._child_prov_require_copy_range(0, 0x9060, 32, api="copy_to") w._child_prov_clear_malloc(0, 0x9000) w._child_prov_require_live(0, 0x9000, api="copy_to") # still live via domain + with pytest.raises(ValueError, match="not a live allocation"): + w._child_prov_require_copy_range(0, 0x9060, 32, api="copy_to") w._child_prov_drop_domain(3) with pytest.raises(ValueError, match="not a live allocation"): w._child_prov_require_live(0, 0x9000, api="copy_to") @@ -141,10 +144,10 @@ def test_drop_last_role_deletes_entry_no_empty_state(self): # live one without the fail-closed guard). w = _l3() with w._child_prov_lock: - w._child_prov_record_domain(0, 0x5000, allocation_id=1) + w._child_prov_record_domain(0, 0x5000, allocation_id=1, nbytes=64) w._child_prov_drop_domain(1) assert (0, 0x5000) not in w._child_alloc_prov - w._child_prov_record_malloc(0, 0x6000) + w._child_prov_record_malloc(0, 0x6000, 64) w._child_prov_clear_malloc(0, 0x6000) assert (0, 0x6000) not in w._child_alloc_prov @@ -155,9 +158,9 @@ def test_strict_aba_is_explicitly_not_covered(self): # handles; this guard only catches stale-BEFORE-reuse. w = _l3() with w._child_prov_lock: - w._child_prov_record_malloc(0, 0x1000) + w._child_prov_record_malloc(0, 0x1000, 64) w._child_prov_clear_malloc(0, 0x1000) - w._child_prov_record_malloc(0, 0x1000) # device handed back the same VA + w._child_prov_record_malloc(0, 0x1000, 64) # device handed back the same VA w._child_prov_require_live(0, 0x1000, api="copy_to") # NOT rejected — by design @@ -220,6 +223,48 @@ def test_copy_to_requires_live_device_dst(self): o.copy_to(0, 0x2000, 0xABCD, 64) fake.copy_to.assert_called_once() + def test_copy_to_accepts_interior_range(self): + w = _l3() + fake = MagicMock() + fake.malloc.return_value = 0x2000 + o = Orchestrator(fake, w) + with pytest.raises(ValueError, match="not a live allocation"): + o.copy_to(0, 0x2020, 0xABCD, 16) + fake.copy_to.assert_not_called() + o.malloc(0, 64) + o.copy_to(0, 0x2020, 0xABCD, 16) + fake.copy_to.assert_called_once_with(0, 0x2020, 0xABCD, 16) + + @pytest.mark.parametrize(("dst", "size"), [(0x2030, 17), (0x2041, 0), (0x2000, -1)]) + def test_copy_to_rejects_range_past_allocation(self, dst, size): + w = _l3() + fake = MagicMock() + fake.malloc.return_value = 0x2000 + o = Orchestrator(fake, w) + o.malloc(0, 64) + + with pytest.raises(ValueError, match="not a live allocation|size must be non-negative"): + o.copy_to(0, dst, 0xABCD, size) + + fake.copy_to.assert_not_called() + + def test_copy_to_accepts_comm_domain_window_chunk(self): + w = _l3() + fake = MagicMock() + o = Orchestrator(fake, w) + with w._child_prov_lock: + w._child_prov_record_domain(0, 0x4000, allocation_id=7, nbytes=2 << 20) + # The first carved buffer aliases the window base; recording its + # smaller extent must not narrow the full-window role. + w._child_prov_record_domain(0, 0x4000, allocation_id=7, nbytes=4) + + o.copy_to(0, 0x104000, 0xABCD, 1 << 20) + + fake.copy_to.assert_called_once_with(0, 0x104000, 0xABCD, 1 << 20) + with pytest.raises(ValueError, match="not a live allocation"): + o.copy_to(0, 0x104000, 0xABCD, (1 << 20) + 1) + assert fake.copy_to.call_count == 1 + def test_copy_from_requires_live_device_src(self): w = _l3() fake = MagicMock() @@ -349,7 +394,7 @@ def test_domain_pointer_dispatch_to_owner_then_rejected_after_release(self, _fak w = _l3() w._chip_shms = [object(), object()] with w._child_prov_lock: - w._child_prov_record_domain(0, 0x5000, allocation_id=42) + w._child_prov_record_domain(0, 0x5000, allocation_id=42, nbytes=64) fake = MagicMock() o = Orchestrator(fake, w) o.submit_next_level(object(), _child_args(0x5000), None, worker=0) @@ -404,14 +449,17 @@ def test_l2_free_revokes_before_native_free(self): w.free(0x2000) assert seen["live_at_native"] is False - def test_l2_copy_to_requires_live_dst(self): + def test_l2_copy_to_requires_contained_live_range(self): w, chip = self._l2() with pytest.raises(ValueError, match="not a live allocation"): w.copy_to(0x2000, 0xABCD, 64) chip.copy_to.assert_not_called() w.malloc(64) - w.copy_to(0x2000, 0xABCD, 64) - chip.copy_to.assert_called_once() + w.copy_to(0x2020, 0xABCD, 32) + chip.copy_to.assert_called_once_with(0x2020, 0xABCD, 32) + with pytest.raises(ValueError, match="not a live allocation"): + w.copy_to(0x2020, 0xABCD, 33) + assert chip.copy_to.call_count == 1 # ---------------------------------------------------------------------------- @@ -518,8 +566,8 @@ def _domain_worker(self): w = _l3() w._worker = MagicMock() # non-None so _release_domain_now proceeds with w._child_prov_lock: - w._child_prov_record_domain(0, 0x5000, allocation_id=9) - w._child_prov_record_domain(1, 0x6000, allocation_id=9) + w._child_prov_record_domain(0, 0x5000, allocation_id=9, nbytes=64) + w._child_prov_record_domain(1, 0x6000, allocation_id=9, nbytes=64) handle = SimpleNamespace(name="d", workers=(0, 1), allocation_id=9) return w, handle