Add: materialize planned comm endpoint regions - #1770
Conversation
- Add private RegionInstance materialization for supported L3 host to L2 AICPU plans - Validate unsupported W2 plans with typed refusals instead of fallback behavior - Reuse worker-chip cleanup for instance close and rollback paths - Cover shape validation, delegation refusals, access delegation, and cleanup failures
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR adds validated single-owner region materialization, lifecycle states, payload and counter access, rollback handling, worker cleanup integration, control-context checks, and comprehensive unit tests. ChangesRegion materialization
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant EndpointRegistry
participant NativeWorker
participant RegionInstance
Worker->>EndpointRegistry: Resolve endpoint records
Worker->>NativeWorker: Create worker-chip region
NativeWorker-->>RegionInstance: Return adopted region
RegionInstance->>RegionInstance: Enter live state
Worker->>RegionInstance: Close or roll back region
RegionInstance->>NativeWorker: Release mappings and child region
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (7)
tests/ut/py/test_worker/test_comm_region.py (4)
42-54: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind the layout once so the plan and the context share one object.
layout or ce.RegionLayoutSpec(...)is evaluated twice. When the caller passes no layout, line 47 builds oneRegionLayoutSpecfor the plan and line 53 builds a second, separate one forMaterializationContext.layout. The two calls also use different argument styles. A future field with a non-equal default would let the planned layout and the context layout diverge without any test failing.♻️ Proposed fix
def _context(worker: Worker, members, topology, layout=None) -> MaterializationContext: + layout = layout or ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=128) registry = worker._get_endpoint_registry() resolved = registry.resolve_region_spec(members, topology) - plan = ce.BackendResolver(registry, worker._get_region_access_service()).plan( - resolved, - layout or ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=128), - ) + plan = ce.BackendResolver(registry, worker._get_region_access_service()).plan(resolved, layout) return MaterializationContext( worker=worker, registry=registry, plan=plan, - layout=layout or ce.RegionLayoutSpec(64, 128), + layout=layout, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_worker/test_comm_region.py` around lines 42 - 54, Update _context to bind the effective layout once before calling BackendResolver.plan, then reuse that same object for both the plan and MaterializationContext.layout. Preserve the caller-provided layout and the existing default values while eliminating the duplicate RegionLayoutSpec construction and argument-style difference.
109-166: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd refusal tests for the remote-endpoint and device-peer consumer plans.
The PR objectives list remote endpoints and device-peer consumers as explicit refusals. The current tests cover unsupported plan, non-
VMM_WINDOWbacking, direct-map adapter, AICore provider, L4 delegation, and extra consumers. No test drives a remote endpoint member or aDEVICE_PEER/DEVICE_VMM_PEER_IMPORTconsumer attachment throughvalidate_single_owner_region_shape.The device-peer service entry is already registered in this test at lines 151-159, so a device-peer consumer case is cheap to add.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_worker/test_comm_region.py` around lines 109 - 166, Add refusal coverage in the shape-validation tests for both a remote-endpoint member and a device-peer consumer attachment. Build contexts that route each case through validate_single_owner_region_shape, reuse the existing DEVICE_PEER/DEVICE_VMM_PEER_IMPORT service registration, and assert RefusalReason.UNSUPPORTED_MEMBER_SHAPE or the appropriate unsupported-consumer refusal reason.
219-231: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the repeated region-materialization setup into a fixture.
Lines 220-230, 268-278, 302-312, 340-347, and 361-371 repeat the same five-step setup. Only
fail_mapping_closediffers. A fixture that returns(worker, calls, fake_region)removes about 40 duplicated lines and keeps the call-ordering assertions as the only per-test content.♻️ Proposed fixture
`@pytest.fixture` def region_worker(monkeypatch): def _build(*, fail_mapping_close=False, device_ids=(8, 9)): worker = _l3(device_ids=device_ids) calls: list[tuple] = [] fake_region = _FakeRegion(calls, fail_mapping_close=fail_mapping_close) worker._worker = _FakeNativeWorker(calls) def create_region(worker_id, payload_bytes, counter_bytes): calls.append(("create", worker_id, payload_bytes, counter_bytes)) worker._live_worker_chip_regions.append(fake_region) return fake_region monkeypatch.setattr(worker, "_create_worker_chip_region", create_region) return worker, calls, fake_region return _build🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_worker/test_comm_region.py` around lines 219 - 231, Extract the repeated worker, calls, fake-region, native-worker, and region-creation setup from the affected tests into a region_worker fixture returning a builder with fail_mapping_close and device_ids options. Replace each duplicated setup with the fixture result, preserving each test’s existing call-ordering assertions and using fail_mapping_close only where needed.
267-298: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for a failed
rollback().This test covers only a successful rollback. No test drives
rollback()toROLLBACK_FAILED. Reuse_FakeRegion(calls, fail_mapping_close=True)and make the release fail, then assert the state, the cached error replay on a secondrollback(), and the behavior of a followingclose().That case exercises the terminal-state gap I flagged in
python/simpler/comm_region.pylines 133-170.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/py/test_worker/test_comm_region.py` around lines 267 - 298, Add a test alongside test_live_region_instance_rollback_reuses_single_region_cleanup that uses _FakeRegion(calls, fail_mapping_close=True) and configures release to fail, then assert rollback() transitions the instance to ROLLBACK_FAILED, preserves and re-raises the cached error on a second rollback(), and verify the subsequent close() behavior. Reuse the existing call-tracking and reservation setup to cover the rollback failure path in the region instance implementation.python/simpler/comm_region.py (1)
278-283: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winReject duplicate attachment members explicitly.
attachmentsis built as a dict keyed byattachment.member. A duplicated member silently overwrites the earlier entry, and the set-equality check still passes. The overwritten attachment is never validated. Add a length check so a malformed plan is refused rather than partially validated.♻️ Proposed guard
attachments = {attachment.member: attachment for attachment in part.attachments} - if set(attachments) != {provider.identity, consumer.identity}: + if len(attachments) != len(part.attachments) or set(attachments) != {provider.identity, consumer.identity}: raise MaterializationRefusal( RefusalReason.UNSUPPORTED_ATTACHMENT, "Part attachments must match exactly the provider and host consumer", )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/comm_region.py` around lines 278 - 283, In the attachment validation flow, explicitly reject duplicate members before relying on the `attachments` dictionary. Compare the number of original `part.attachments` entries with the number of keys in `attachments`, and raise `MaterializationRefusal` with `RefusalReason.UNSUPPORTED_ATTACHMENT` when they differ; preserve the existing exact provider/consumer membership check.python/simpler/worker.py (2)
5797-5803: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the shared plan-building steps from
_plan_regionand_materialize_region_instance.Lines 5797-5800 repeat
_plan_region(lines 5785-5788) exactly: registry fetch,resolve_region_spec,BackendResolver,plan. Extract a private helper so a future change to the planning inputs stays in one place.♻️ Proposed refactor
+ def _build_region_plan(self, members, topology: SingleOwner, layout_summary: RegionLayoutSpec): + registry = self._get_endpoint_registry() + resolved = registry.resolve_region_spec(members, topology) + resolver = BackendResolver(registry, self._get_region_access_service()) + return registry, resolver.plan(resolved, layout_summary)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 5797 - 5803, Extract the repeated registry lookup, region resolution, BackendResolver construction, and plan creation from `_plan_region` and `_materialize_region_instance` into a shared private helper. Update both methods to call that helper with their existing inputs and preserve their current materialization and return behavior.
7811-7850: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePreserve the second cleanup error instead of dropping it.
region_errorscan hold both the mapping-close error and the release error. Every exit path raisesregion_errors[0]and discards the rest. When the mapping close fails first, the release failure disappears from the traceback, and the poison message names only the first error. Attach the remaining errors so both debts stay diagnosable.♻️ Proposed change
if region_errors: if poison_on_error: self._record_unreclaimable( f"close_worker_chip_region: region {region.region_id} on worker {region._worker_id} could not be " "fully reclaimed; no further work is admitted", region_errors[0], ) + for extra in region_errors[1:]: + if region_errors[0].__cause__ is None: + region_errors[0].__cause__ = extra raise region_errors[0]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/worker.py` around lines 7811 - 7850, Update the cleanup error handling around region_errors and release_error so any additional cleanup failures are attached to the primary exception before every raise, preserving both mapping-close and worker-release errors in the traceback and poison reporting. Ensure all exit paths, including the early poison_on_error branch and final raise, retain the secondary errors rather than discarding them.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/simpler/comm_region.py`:
- Around line 133-153: Define one module-level set containing
RegionInstanceState.CLOSE_FAILED and RegionInstanceState.ROLLBACK_FAILED, then
update RegionInstance.close() at python/simpler/comm_region.py:133-153 and
RegionInstance.rollback() at python/simpler/comm_region.py:154-170 to use
membership in that set when returning the cached cleanup error. This must
prevent either method from retrying cleanup after either terminal-failure state
and preserve the recorded cause.
In `@python/simpler/worker.py`:
- Around line 7859-7874: Update the nested retire_tracking function to replace
both recursive retry paths with a bounded loop that advances next_tracking_list
before retrying, including when catching BaseException. Preserve error
collection in errors, ensure each iteration makes progress, and allow
asynchronous KeyboardInterrupt/BaseException signals to propagate instead of
immediately retrying teardown.
In `@tests/ut/py/test_worker/test_comm_region.py`:
- Line 57: Add postponed annotation evaluation at the top of the test module so
the _accepted_context parameter annotation using Worker | None remains
import-compatible with Python 3.9; alternatively, replace that union with
Optional[Worker].
---
Nitpick comments:
In `@python/simpler/comm_region.py`:
- Around line 278-283: In the attachment validation flow, explicitly reject
duplicate members before relying on the `attachments` dictionary. Compare the
number of original `part.attachments` entries with the number of keys in
`attachments`, and raise `MaterializationRefusal` with
`RefusalReason.UNSUPPORTED_ATTACHMENT` when they differ; preserve the existing
exact provider/consumer membership check.
In `@python/simpler/worker.py`:
- Around line 5797-5803: Extract the repeated registry lookup, region
resolution, BackendResolver construction, and plan creation from `_plan_region`
and `_materialize_region_instance` into a shared private helper. Update both
methods to call that helper with their existing inputs and preserve their
current materialization and return behavior.
- Around line 7811-7850: Update the cleanup error handling around region_errors
and release_error so any additional cleanup failures are attached to the primary
exception before every raise, preserving both mapping-close and worker-release
errors in the traceback and poison reporting. Ensure all exit paths, including
the early poison_on_error branch and final raise, retain the secondary errors
rather than discarding them.
In `@tests/ut/py/test_worker/test_comm_region.py`:
- Around line 42-54: Update _context to bind the effective layout once before
calling BackendResolver.plan, then reuse that same object for both the plan and
MaterializationContext.layout. Preserve the caller-provided layout and the
existing default values while eliminating the duplicate RegionLayoutSpec
construction and argument-style difference.
- Around line 109-166: Add refusal coverage in the shape-validation tests for
both a remote-endpoint member and a device-peer consumer attachment. Build
contexts that route each case through validate_single_owner_region_shape, reuse
the existing DEVICE_PEER/DEVICE_VMM_PEER_IMPORT service registration, and assert
RefusalReason.UNSUPPORTED_MEMBER_SHAPE or the appropriate unsupported-consumer
refusal reason.
- Around line 219-231: Extract the repeated worker, calls, fake-region,
native-worker, and region-creation setup from the affected tests into a
region_worker fixture returning a builder with fail_mapping_close and device_ids
options. Replace each duplicated setup with the fixture result, preserving each
test’s existing call-ordering assertions and using fail_mapping_close only where
needed.
- Around line 267-298: Add a test alongside
test_live_region_instance_rollback_reuses_single_region_cleanup that uses
_FakeRegion(calls, fail_mapping_close=True) and configures release to fail, then
assert rollback() transitions the instance to ROLLBACK_FAILED, preserves and
re-raises the cached error on a second rollback(), and verify the subsequent
close() behavior. Reuse the existing call-tracking and reservation setup to
cover the rollback failure path in the region instance implementation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d62dcc46-10bb-4cfc-a1a6-bc6b51093ffb
📒 Files selected for processing (4)
python/simpler/comm_endpoints.pypython/simpler/comm_region.pypython/simpler/worker.pytests/ut/py/test_worker/test_comm_region.py
- Require RegionInstance access and cleanup to run inside an active orchestration/control context - Reject stale or foreign materialization registries before creating worker-chip regions - Poison worker admission on instance cleanup failure without retrying a released chip region - Cover counter notify/wait delegation and context/refusal cleanup cases
7aa41ba to
d655f14
Compare
- Reuse terminal cleanup failure state across close() and rollback() - Reject duplicate attachment members and remote endpoint materialization - Cover cleanup failure, rollback failure, and BaseException tracking paths in UTs
ChaoWao
left a comment
ChaoWao
left a comment
There was a problem hiding this comment.
Review
Reviewed against merge-base 92eebce9 (3 commits, +979/−43 across 4 files). CI 19/19 green.
Summary: the design is clean and the narrow scope is the right call. Validation is separated from side effects so validate_single_owner_region_shape() is unit-testable on its own; all 9 refusal classes listed in the PR body are implemented and individually tested; existing gates (_operation_lease / _control_admission / _record_unreclaimable) are reused rather than reinvented; no native ABI is touched. The extraction of _close_worker_chip_region so that run-scoped batch cleanup and instance-scoped single cleanup share one body is a genuine improvement.
One blocking finding, reproduced empirically against this branch's code, plus a handful of smaller items.
Must fix
1. RegionInstance.close() releases the same region twice when materialization happened inside a run's orch callback
RegionInstance holds only _worker, so close() / rollback() call
self._worker._close_worker_chip_region(self._region, poison_on_error=True)leaving resources at its default None. _retire_worker_chip_region_tracking(region, None) therefore builds tracking_lists = [self._live_worker_chip_regions] only (worker.py:7869-7871).
But _create_worker_chip_region registers the region into two lists when self._building_run_resources is not None (worker.py:7709-7713), and _building_run_resources is non-None for exactly the duration of the orch callback (worker.py:10072 + _callback_run) — which is precisely the context _require_region_control_context blesses as valid for using the instance (worker.py:9936-9943).
Probe against this branch, simulating the callback + resources environment that Worker.submit() sets up (fake region/native worker mirroring the real tracking behaviour, including WorkerHostRegionMapping.close()'s idempotence):
state after materialize: RegionInstanceState.LIVE
tracked in _live_worker_chip_regions: 1
tracked in resources.worker_chip_regions: 1
state after close: RegionInstanceState.CLOSED
AFTER close():
_live_worker_chip_regions : 0
resources.worker_chip_regions : 1 <-- should be 0
Then the run's ordered cleanup runs (submit()'s finalization cursor, worker.py:10189):
calls: [('create',1,64,128), ('mapping_close',False), ('release',1,42), ('expire',),
('mapping_close',True), ('release',1,42), ('expire',)]
RESULT: control_worker_chip_region_release called 2 time(s) for region 42
The host mapping side is harmless — WorkerHostRegionMapping.close() has a closed idempotence flag. The chip-side release is unguarded: _close_worker_chip_region checks neither region._expired nor region._released. So either the child errors on an unknown region id, which fails ordered cleanup and poisons the worker (a run that should have succeeded goes red), or the child silently accepts it and — if the id has been reused — a second release frees somebody else's region.
Same mechanism, second route: materialize in a callback → submit a task → close() is refused by has_submitted_task → run cleanup releases the region while the RegionInstance is still LIVE → the user closes it afterwards inside a reservation → second release again.
Not reachable from production code today (_materialize_region_instance is private and only the new tests call it, all outside a callback). But _require_region_control_context documents the callback context as first-class, so the first W4/W5 consumer will hit it. Three possible fixes:
- capture
worker._building_run_resourcesat materialize time on theRegionInstanceand pass it back to_close_worker_chip_region(closest to current structure); - have
_retire_worker_chip_region_trackingalways sweep both lists — theresourceslist is an additional registration, not a replacement; - or make
_materialize_region_instancerefuse outright whenself._building_run_resources is not None, deferring that shape to W5 — but then the callback branch of_require_region_control_contextcontradicts it and should go in the same change.
Whichever you pick, this needs a regression test in callback context: the current fixture's fake create_region only appends to _live_worker_chip_regions, so it can never observe this.
Should fix
2. close() after a successful rollback() silently rewrites ROLLED_BACK to CLOSED
comm_region.py:150-152: if self._state in (PLANNED, ROLLED_BACK): self._state = CLOSED; return. For the W5 transaction semantics this object exists to serve, "was rolled back" and "was closed normally" are different outcomes, and here the later caller wins. PLANNED → CLOSED is right (no resource to speak of); ROLLED_BACK → CLOSED should be a no-op.
3. The relocated comment in _close_worker_chip_region no longer covers the code beneath it
worker.py:7838-7840 describes only the resources is None (whole-tree close, retain for journal replay) vs resources is not None (end-of-run, expire) dichotomy. The third case this PR introduces — poison_on_error=True, the RegionInstance path, where resources is always None yet the region is not retained — isn't mentioned. Neither is the asymmetric guard right below it (worker.py:7842: tracking is retained only when the chip release failed, not when the mapping close failed), which is the least obvious invariant in the function. Per .claude/rules/comments.md, what belongs here is the present-tense fact: who retains tracking, and why.
4. _LOCAL_L2_PATH_RE duplicates the existing endpoint-path vocabulary
comm_region.py:36 introduces ^L3/L2\[(?P<worker_id>[0-9]+)\]$, but comm_endpoints.py already has parse_endpoint_path(path, root_level=...) → ParsedEndpointPath.segments (carrying level / index), and EndpointRegistry.root_level is public. The new regex also hardcodes root level 3 — correct today only because the caller already pinned level == 3, but it's a second spelling of one concept. Using parse_endpoint_path and asserting the segment shape gets the same result from the existing vocabulary.
5. The chip-id bound check is a second implementation of _validate_worker_chip_id
comm_region.py:216-217 reads ctx.worker._config.get("device_ids", ()) and recomputes worker_id >= len(device_ids), which Worker._validate_worker_chip_id (worker.py:7523-7530) already does with the same bound. Wanting a MaterializationRefusal rather than a ValueError is a fair motive, but call the helper and translate the exception rather than copying the bound. Incidentally worker_id < 0 is dead — the regex only matches [0-9]+.
6. test_retire_worker_chip_region_tracking_does_not_retry_base_exception cannot fail for the reason it claims
It asserts tracking.assignments == 1 while installing only one tracking list. With a single list, assignments is 1 whether or not retry logic exists. Distinguishing the named behaviour (don't re-attempt the same list, but do continue to the next one) needs two lists.
Consider
-
Unwired speculative surface.
RegionInstanceState.CONSUMER_ATTACHEDis never assigned or read.RegionInstance.generation(module-levelitertools.count, so process-global rather than per-worker),diagnostic_label(eagerly formatted on every materialization),planandlayouthave no readers. If these are for W5, adding them with W5 keeps them covered by tests. -
Private state mutated from outside the class.
materialize_region_instanceassignsinstance._statethree times and callsinstance._adopt_worker_chip_region/instance._rollback_after_failed_materialization(the latter a one-line alias forrollback()); a classmethod onRegionInstancewould own this. Related semantics:_state = OWNER_CREATEDis set before_create_worker_chip_regionis called, so on failure the state claims the owner created something that does not exist. It is safe today only because rollback keys off_region is None, not off the state. -
raise errors[0]now raises a different exception. Previously, when mapping and release both failed withresourcesnon-None,errorswas ordered[expire_exc, mapping_exc, release_exc, …]anderrors[0]was the expire error. Nowregion_errorsis in occurrence order and the mapping error is raised with the rest chained on__cause__. The new behaviour is better (first real failure rather than the bookkeeping one), but it is a behaviour change inside something framed as a refactor and worth a line in the PR body. -
# noqa: PLR0912onvalidate_single_owner_region_shape(max-branches = 15). There's precedent in the file (_create_worker_chip_regioncarries the same), but this one is pure validation and decomposes naturally along the refusal families — splitting beats suppressing. -
Inconsistent error taxonomy: the module defines a
MaterializationErrorhierarchy, butRegionInstance._ensure_live()(comm_region.py:186-190) raises a bareRuntimeError. -
EndpointRegistry.record_for()has no docstring while its neighbourowner_endpoint()does, and it is newly public surface.
Traceability
Every stated goal traces to a design choice except one:
| Stated goal | Assessment |
|---|---|
plan → live internal RegionInstance |
✅ |
stays private, no public create_region |
✅ |
| rollback-capable object for W5 | |
| a place for W4 to hang publish/observe | ✅ |
| refuse all 9 shapes without a materializer | ✅ all covered and tested |
| "region tracking is retired only after cleanup debts are attempted" | ❌ true for _live_worker_chip_regions; the run-scoped list is never retired from the RegionInstance path — finding #1 |
| no C++ binding / wire ABI change | ✅ |
| — | ➕ EndpointRegistry.record_for() is new public surface not covered by any stated goal |
Two things the PR body doesn't answer and the next contributor will have to re-derive: why cleanup was extracted into _close_worker_chip_region with two optional parameters rather than giving RegionInstance its own cleanup closure, and why RegionInstance isn't a context manager given that close() requires a control context anyway (with is the natural shape for that constraint).
Verdict
The only hard blocker is #1. It is unreachable from production code today, but the callback context it breaks in is already documented as valid, so it is what the next W4/W5 consumer will hit — better fixed now than discovered then.
Minimal closing loop: fix #1 (I'd lean toward capturing resources at materialize time) + a callback-context regression test + the #2 no-op + the #3 comment. #4/#5 vocabulary de-duplication and #6's test strengthening are fine either here or as follow-ups.
- Capture run cleanup resources on RegionInstance so callback-created regions retire both tracking lists when closed directly. - Treat already-expired worker-chip regions as cleanup-only handles, and mark regions released after chip release succeeds. - Reuse endpoint path parsing and worker-chip id validation in the private materializer. - Add callback cleanup and state-machine regressions for the W3.5 review fixes.
8b5d30b to
646d20f
Compare
Worker-chip cleanup must only skip native release for already-expired regions. A handle-level free marker can predate cleanup, so using it as the release guard drops the actual child release. Mark regions released after a successful native release when the handle supports it, while keeping unit-test fake regions compatible.
|
Thanks @ChaoWao for the detailed review. The blocker is fixed with option A from the review: Fix commits:
Must Fix1. Callback-context materialization can double-release a worker-chip regionFixed.
self._worker._close_worker_chip_region(
self._region,
self._cleanup_resources,
poison_on_error=True,
)For a region created inside an orchestration callback, instance-level
Two callback-context regressions were added with
The fake create path in the test fixture now mirrors production tracking: when The follow-up CI fix in Should Fix2.
|
ChaoWao
left a comment
ChaoWao
left a comment
There was a problem hiding this comment.
Re-review — head 6ea63054 (2 new commits on top of bffea8b5)
The blocker is fixed. Verified against this branch's code with the same probe that surfaced it — the run-scoped tracking list is now retired and the chip sees exactly one release:
state after close: RegionInstanceState.CLOSED
_live_worker_chip_regions : 0
resources.worker_chip_regions : 0
Then the run's ordered cleanup runs (submit()'s finalization cursor):
calls unchanged — no second release
RESULT: control_worker_chip_region_release called 1 time(s) for region 42
Both routes I described are now pinned by tests: test_callback_region_close_before_submit_retired_from_run_cleanup covers close-inside-callback, and test_callback_region_run_cleanup_then_later_close_is_idempotent covers the submit-then-close-after-run-cleanup path via the new expired guard. CI 19/19 green.
Also fixed: #2 (ROLLED_BACK is now a no-op for close(), with an assertion), #3 (comment rewritten as a present-tense fact about who retains tracking and why), #4 (parse_endpoint_path replaces the local regex, plus a non-L2-child-path test), #5 (_validate_worker_chip_id is reused, with test_shape_validation_does_not_translate_worker_readiness_errors deliberately pinning that readiness RuntimeErrors are not translated into refusals — a reasonable call, and good that it's explicit), #6 (two tracking lists now, asserting both the failing one is attempted once and the succeeding one drains), #11 (MaterializationError), #12 (docstring). 8 of 12 items closed.
One thing worth crediting that the commit messages undersell: the free() call after a successful chip release closes a real hole in the retained-for-replay window. Before this PR, when the mapping close failed but the chip release succeeded, the region stayed tracked with neither _expired nor _released set — _ensure_live() passed, so a later payload_write could write into a region whose chip-side ownership had already been handed back. Verified on the merge-base:
merge-base 92eebce9 : region._expired: False region._released: False <- _ensure_live() passes
head 6ea63054 : region._expired: False region._released: True <- refused
Remaining, none blocking
a. Production code carrying getattr guards for attributes its only real type always has. worker.py:7827:
expired = bool(getattr(region, "expired", getattr(region, "_expired", False)))
...
free = getattr(region, "free", None)
if callable(free):
free()WorkerChipOrchRegion has both an expired property (worker_chip_orch_comm.py:325) and free() (:366), and it is the only type that ever reaches _live_worker_chip_regions / resources.worker_chip_regions — _create_worker_chip_region is the sole producer. Commit 6ea63054 says as much: "while keeping unit-test fake regions compatible." That's the test doubles dictating the shape of production code. Giving _FakeRegion an expired property (it already grew a free()) lets both sites become plain region.expired / region.free(), and then a fake that drifts out of shape fails loudly instead of silently taking the False branch — which for the expired guard would mean silently double-releasing, the exact bug we just fixed.
b. The replay double-release is still open (pre-existing — I verified it reproduces identically on merge-base 92eebce9, so this PR neither introduced nor is obliged to fix it). Whole-tree close where the mapping close fails and the chip release succeeds retains the region by design; on replay, expired is still False, so the release runs a second time:
pass 1 (whole-tree close): mapping_close raises, release, free -> retained, _expired=False
pass 2 (journal replay) : mapping_close no-op, release, free, expire
RESULT: control_worker_chip_region_release called 2 time(s) for region 42
_released can't be the guard, for exactly the reason 6ea63054's message gives — a user-called free() is indistinguishable from a cleanup-owned one. So closing this needs a marker cleanup owns (e.g. _chip_release_committed), separate from the handle-level _released. Fine as a follow-up; worth an issue so it isn't rediscovered by the next person who reads this function.
c. self._cleanup_resources = getattr(ctx.worker, "_building_run_resources", None) (comm_region.py:110) — same defensive-getattr note as (a); Worker always has the attribute (worker.py:4158). Separately, a RegionInstance outliving its run now pins that run's whole _RunResources (domain lists, remote slot refs) alive. Harmless today because closing retires from a list that is already empty, but if the intent is only "retire from whatever lists tracked me," a weakref or just the list object would express it without the lifetime tie.
d. Still open from the first pass: #7 (CONSUMER_ATTACHED is still never assigned or read; generation / diagnostic_label / plan / layout still have no readers), #8 (materialize_region_instance still assigns instance._state from outside the class, and still sets OWNER_CREATED before the create call), #10 (validate_single_owner_region_shape kept its # noqa: PLR0912 and grew from 12 to 16 branches / 65 lines — the path-parsing rewrite added three more; splitting along the refusal families is now more clearly the cheaper option than the suppression).
e. Redundant provider-path checks. comm_region.py:229-234 checks root.level != 3 or root.index is not None, but parse_endpoint_path(..., root_level=ctx.registry.root_level) already enforces both (comm_endpoints.py:121-123), and root_level is 3 here because the earlier worker.level != 3 gate and provider.path.startswith("L3/") both already passed. Three checks for one condition. Keeping len(segments) != 2 and the child.level != 2 or child.index is None half is what actually carries weight.
f. The PR body is now stale. It still describes only the original three commits — no mention of the callback-cleanup fix, the expired-guard idempotence, or the six new tests, and the Tests list omits all of them. It also still doesn't note the behaviour change I flagged as #9 (raise errors[0] now surfaces the mapping error rather than the expire error, with the rest chained on __cause__ — an improvement, but a behaviour change inside something framed as a refactor). Worth regenerating before merge.
Verdict
Approve with nits. The blocking defect is fixed, the fix is the right one structurally (the instance owns the resources it was created under rather than the cleanup function guessing), and it comes with regression tests that actually exercise the callback context the original tests couldn't reach. Everything left is small: (a) is the one I'd fix in this PR since it guards the bug we just closed, (f) is a two-minute regeneration, and (b) deserves an issue rather than a fix here.
This PR handles the #1770 follow-up items. Done here: - Tighten worker-chip region cleanup shape: - replace defensive `getattr(region, "expired", ...)` with direct `region.expired` - replace defensive `getattr(region, "free", None)` / callable check with direct `region.free()` - update test fake regions to expose the same cleanup surface as `WorkerChipOrchRegion` - Fix replay double-release: - add cleanup-owned chip release tracking separate from user-facing `_released` - mark chip release committed only after `control_worker_chip_region_release(...)` succeeds - allow close replay to retry native host mapping close without releasing the same chip region twice - add a whole-tree close regression test for mapping-close failure plus successful chip release Already handled in #1822: - W4 region sync/access refactor onto neutral region-native primitives - `WorkerChipOrchRegion` compatibility facade over W4 region access pieces - neutral native mapped-region backend and `_region_*` binding surface - redundant provider root check cleanup in `validate_single_owner_region_shape` Not done here: - post-W4 `comm_region.py` hygiene - public `create_region(...)` - W5 delegated transactions - queue/ring/mailbox/freelist templates - Symmetric/OpenSHMEM topology
2 participants
Summary
This PR builds on #1696 by adding the private Python materialization layer for
the first planned communication-region shape.
#1696 introduced the endpoint model: endpoint selectors, endpoint registries,
SingleOwnerbackend planning, typedBackendPlans, and explicit unsupportedplans. This PR takes one supported plan shape from that layer and turns it into
a live internal
RegionInstancebacked by the existing worker-chip regionnative path.
In the work breakdown, this corresponds to the internal region
instance/materializer step. It is intentionally private: it gives later W4/W5
work a real plan-to-resource object without exposing a public
create_region(...)API yet.Why This PR Exists
In the work breakdown, this PR is the internal region
instance/materializer step. This PR sits between #1696's endpoint/backend planner
and the later W4/W5 work.
W4 is the sync-primitive normalization step. It promotes the current
worker-chip counter operations (
notify,test,wait) into region-levelpublish/observe APIs, with arch-specific cache-maintenance call sites hidden
behind the region abstraction. W4 needs a real internal region object so those
sync operations are no longer tied directly to the old worker-chip API shape.
W5 is the delegated transaction step. It lets a higher-level worker declare a
region over endpoint members, delegates execution to the owning L3/chip side,
and gives the operation transaction semantics: partial-failure rollback,
membership freeze, epoch/generation tracking, and an all-member READY barrier.
W5 needs a rollback-capable internal region instance before it can safely build
delegated create/commit/rollback orchestration.
Before this PR, #1696 could describe a valid region plan, and the worker already
had a native worker-chip region backend, but there was no internal layer that
bound those two together. This PR adds that layer and its lifecycle/error
behavior without exposing a public
create_region(...)API yet.Supported Materialization Shape
This PR supports exactly the current production shape:
HOST_CPUconsumerL3/L2[i] DEVICE_AICPUproviderSingleOwnertopologyVMM_WINDOWpayload and counter backingOWNER_DELEGATED_COPY / HOST_VMM_COPYThis maps to the existing
Worker._create_worker_chip_region(...)backend anddoes not add a new native backend, C++ binding, or wire ABI.
Why The Scope Is Deliberately Narrow
The current work-breakdown conclusion is that production region materialization
must stay on
VMM_WINDOW + HOST_VMM_COPYfor now. Earlier direct host-map planswere rejected by W0/W1 validation: host CPU mappings alone do not prove that
AICPU can observe host writes through the region device address. As a result,
direct-map and host-mapped profiles remain planning vocabulary, not production
materialization paths.
This PR therefore refuses every shape that would require a materializer we do
not yet have, including:
Those refusals are intentional. They keep this PR as a stable base for W4/W5
instead of smuggling W4 sync semantics or W5 delegated transactions into this
PR.
Region Lifecycle
The private
RegionInstancetracks planned, live, closed, rolled-back, andcleanup-failed states. Cleanup is routed through the worker's existing ordered
cleanup/admission model:
close()androllback()share cached terminal cleanup errorsThis gives W5 a rollback-capable internal object and gives W4 a single place to
hang future region-level publish/observe operations.
Non-Goals
This PR does not:
Worker.create_region(...)orOrchestrator.create_region(...)RegionInstanceas public APITests
BaseException