diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index f57a426036..8a4f103a4f 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -31,12 +31,15 @@ #include #include +#include #include #include #include #include #include +#include #include +#include #include #include #include @@ -62,7 +65,6 @@ #include "data_type.h" #include "dma_workspace.h" #include "worker_chip_orch_comm.h" -#include "worker_chip_orch_region_access.h" #include "worker_bind.h" #include "task_args.h" #include "tensor.h" @@ -247,6 +249,18 @@ class AclRuntimeApi { return handle; } + int vmm_unmap(void *va) const { return aclrtUnmapMem(va); } + + int vmm_release_address(void *va) const { return aclrtReleaseMemAddress(va); } + + int vmm_free_physical(void *handle) const { return aclrtFreePhysical(handle); } + + void vmm_zero_bytes(void *va, uint64_t offset, uint64_t nbytes) const { + std::vector zeros(static_cast(nbytes), 0); + void *dst = reinterpret_cast(reinterpret_cast(va) + static_cast(offset)); + memcpy_h2d_with_check(dst, static_cast(nbytes), zeros.data(), static_cast(nbytes)); + } + void vmm_release_collecting(void *va, void *handle, std::string &cleanup_error) const { if (va != nullptr) { int rc = aclrtUnmapMem(va); @@ -314,29 +328,6 @@ AclRuntimeApi &acl_api() { return *api; } -class ChipChildOnboardRegion { -public: - int device_id{-1}; - uint64_t device_addr{0}; - uint64_t mapping_bytes{0}; - uint64_t shareable_handle{0}; - void *vmm_handle{nullptr}; - - void bind_acl_device() const { - if (device_id < 0) { - throw std::runtime_error("L3-L2 onboard child region has no device id"); - } - acl_api().bind_device_with_check(device_id); - } -}; - -struct ChipChildOnboardRegionExport { - uint64_t device_addr{0}; - uint64_t mapping_bytes{0}; - uint64_t shareable_handle{0}; - uint64_t registry_handle{0}; -}; - std::string region_shm_name_for_open(const std::string &token) { if (token.empty()) { throw std::invalid_argument("mapped-region sim backing shm token must be non-empty"); @@ -880,34 +871,6 @@ RegionHandle import_onboard_region( return RegionHandle(handle, std::move(handle_owner_token)); } -class ChipChildOnboardRegionRegistry { -public: - uint64_t emplace(ChipChildOnboardRegion region) { - std::lock_guard lk(mu_); - uint64_t handle = next_handle_++; - regions_.emplace(handle, std::move(region)); - return handle; - } - - std::optional remove(uint64_t handle) { - std::lock_guard lk(mu_); - auto it = regions_.find(handle); - if (it == regions_.end()) { - return std::nullopt; - } - ChipChildOnboardRegion region = std::move(it->second); - regions_.erase(it); - return region; - } - -private: - mutable std::mutex mu_; - std::unordered_map regions_; - uint64_t next_handle_{1}; -}; - -ChipChildOnboardRegionRegistry g_chip_child_onboard_regions; - uint64_t align_vmm_bytes(uint64_t bytes, uint64_t granularity) { if (bytes == 0 || bytes > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("L3-L2 onboard VMM region requires a positive mapping size"); @@ -933,6 +896,691 @@ void append_cleanup_error(std::string &cleanup_error, const std::string &message cleanup_error += message; } +enum class RegionVmmStage { + None = 0, + PhysicalAlloc, + VaReserve, + Map, + SetAccess, + Export, + Unmap, + VaRelease, + PhysicalFree, + ZeroBytes, + BindDevice, +}; + +enum class RegionVmmFailWhen { Before, After }; + +enum class RegionVmmUnmapClass { Success, AlreadyGone, Failed }; + +struct RegionVmmExport { + uint64_t device_addr{0}; + uint64_t mapping_bytes{0}; + uint64_t shareable_handle{0}; + uint64_t registry_handle{0}; +}; + +struct RegionVmmInspect { + bool present{false}; + int device_id{-1}; + uint64_t device_addr{0}; + uint64_t mapping_bytes{0}; + uint64_t shareable_handle{0}; + uint64_t physical_handle{0}; + bool physical_allocated{false}; + bool va_reserved{false}; + bool mapped{false}; + bool access_set{false}; + bool exported{false}; + bool unmap_attempted{false}; + bool unmap_complete{false}; + bool va_release_attempted{false}; + bool va_release_complete{false}; + bool physical_free_attempted{false}; + bool physical_free_complete{false}; + bool release_once_done{false}; + std::string first_cleanup_failure; + std::vector local_cleanup_details; +}; + +struct RegionVmmAllocation { + int device_id{-1}; + void *physical_handle{nullptr}; + void *va{nullptr}; + uint64_t mapping_bytes{0}; + uint64_t shareable_handle{0}; + bool physical_allocated{false}; + bool va_reserved{false}; + bool mapped{false}; + bool access_set{false}; + bool exported{false}; + bool unmap_attempted{false}; + bool unmap_complete{false}; + bool va_release_attempted{false}; + bool va_release_complete{false}; + bool physical_free_attempted{false}; + bool physical_free_complete{false}; + bool release_once_done{false}; + std::string first_cleanup_failure; + std::vector local_cleanup_details; +}; + +struct RegionVmmTestHooks { + bool use_fake{false}; + std::optional fail_stage; + RegionVmmFailWhen fail_when{RegionVmmFailWhen::Before}; + bool unmap_already_gone{false}; + std::vector issued_ops; + uint64_t fake_granularity{64}; + uint64_t next_fake_id{1}; + uint64_t next_shareable{1}; + std::unordered_map> fake_va_bytes; +}; + +std::mutex ®ion_vmm_mu() { + static auto *mu = new std::mutex(); + return *mu; +} + +RegionVmmTestHooks ®ion_vmm_test_hooks() { + static auto *hooks = new RegionVmmTestHooks(); + return *hooks; +} + +std::unordered_map &g_region_vmm_allocations() { + static auto *allocations = new std::unordered_map(); + return *allocations; +} + +uint64_t ®ion_vmm_next_handle() { + static auto *next = new uint64_t(1); + return *next; +} + +const char *region_vmm_stage_name(RegionVmmStage stage) { + switch (stage) { + case RegionVmmStage::PhysicalAlloc: + return "physical_alloc"; + case RegionVmmStage::VaReserve: + return "va_reserve"; + case RegionVmmStage::Map: + return "map"; + case RegionVmmStage::SetAccess: + return "set_access"; + case RegionVmmStage::Export: + return "export"; + case RegionVmmStage::Unmap: + return "unmap"; + case RegionVmmStage::VaRelease: + return "va_release"; + case RegionVmmStage::PhysicalFree: + return "physical_free"; + case RegionVmmStage::ZeroBytes: + return "zero_bytes"; + case RegionVmmStage::BindDevice: + return "bind_device"; + case RegionVmmStage::None: + return "none"; + } + return "none"; +} + +RegionVmmStage region_vmm_parse_stage(const std::string &name) { + if (name == "physical_alloc") { + return RegionVmmStage::PhysicalAlloc; + } + if (name == "va_reserve") { + return RegionVmmStage::VaReserve; + } + if (name == "map") { + return RegionVmmStage::Map; + } + if (name == "set_access") { + return RegionVmmStage::SetAccess; + } + if (name == "export") { + return RegionVmmStage::Export; + } + if (name == "unmap") { + return RegionVmmStage::Unmap; + } + if (name == "va_release") { + return RegionVmmStage::VaRelease; + } + if (name == "physical_free") { + return RegionVmmStage::PhysicalFree; + } + if (name == "zero_bytes") { + return RegionVmmStage::ZeroBytes; + } + if (name == "bind_device") { + return RegionVmmStage::BindDevice; + } + throw std::invalid_argument("unknown region VMM stage: " + name); +} + +void region_vmm_record_issued(const char *op) { region_vmm_test_hooks().issued_ops.emplace_back(op); } + +bool region_vmm_consume_fail(RegionVmmStage stage, RegionVmmFailWhen when) { + RegionVmmTestHooks &hooks = region_vmm_test_hooks(); + if (!hooks.fail_stage.has_value() || *hooks.fail_stage != stage || hooks.fail_when != when) { + return false; + } + hooks.fail_stage.reset(); + return true; +} + +void region_vmm_throw_stage_failure(RegionVmmStage stage, const char *when) { + throw std::runtime_error( + std::string("injected region VMM ") + when + " " + region_vmm_stage_name(stage) + " failure" + ); +} + +void *region_vmm_next_fake_ptr() { + RegionVmmTestHooks &hooks = region_vmm_test_hooks(); + uint64_t id = hooks.next_fake_id++; + return reinterpret_cast(static_cast(id) << 12); +} + +void region_vmm_bind_device(int device_id) { + region_vmm_record_issued("bind_device"); + if (region_vmm_consume_fail(RegionVmmStage::BindDevice, RegionVmmFailWhen::Before)) { + region_vmm_throw_stage_failure(RegionVmmStage::BindDevice, "before"); + } + if (region_vmm_test_hooks().use_fake) { + if (device_id < 0) { + throw std::runtime_error("region VMM fake driver requires a non-negative device id"); + } + return; + } + acl_api().bind_device_with_check(device_id); +} + +uint64_t region_vmm_granularity(int device_id) { + region_vmm_record_issued("granularity"); + if (region_vmm_test_hooks().use_fake) { + return region_vmm_test_hooks().fake_granularity; + } + return acl_api().vmm_granularity_with_check(device_id); +} + +void *region_vmm_malloc_physical(uint64_t bytes, int device_id) { + region_vmm_record_issued("physical_alloc"); + if (region_vmm_test_hooks().use_fake) { + (void)bytes; + (void)device_id; + return region_vmm_next_fake_ptr(); + } + return acl_api().vmm_malloc_physical_with_check(bytes, device_id); +} + +void *region_vmm_reserve(uint64_t bytes) { + region_vmm_record_issued("va_reserve"); + if (region_vmm_test_hooks().use_fake) { + void *va = region_vmm_next_fake_ptr(); + region_vmm_test_hooks().fake_va_bytes[reinterpret_cast(va)] = + std::vector(static_cast(bytes), 0xff); + return va; + } + return acl_api().vmm_reserve_with_check(bytes); +} + +void region_vmm_map(void *va, uint64_t bytes, void *handle) { + region_vmm_record_issued("map"); + if (region_vmm_test_hooks().use_fake) { + (void)va; + (void)bytes; + (void)handle; + return; + } + acl_api().vmm_map_with_check(va, bytes, handle); +} + +void region_vmm_set_access(void *va, uint64_t bytes, int device_id) { + region_vmm_record_issued("set_access"); + if (region_vmm_test_hooks().use_fake) { + (void)va; + (void)bytes; + (void)device_id; + return; + } + acl_api().vmm_set_access_with_check(va, bytes, device_id); +} + +uint64_t region_vmm_export_shareable(void *handle) { + region_vmm_record_issued("export"); + if (region_vmm_test_hooks().use_fake) { + (void)handle; + return region_vmm_test_hooks().next_shareable++; + } + return acl_api().vmm_export_shareable_with_check(handle); +} + +int region_vmm_unmap(void *va) { + region_vmm_record_issued("unmap"); + RegionVmmTestHooks &hooks = region_vmm_test_hooks(); + if (hooks.use_fake) { + (void)va; + if (std::exchange(hooks.unmap_already_gone, false)) { + return 2; + } + return 0; + } + return acl_api().vmm_unmap(va); +} + +int region_vmm_release_address(void *va) { + region_vmm_record_issued("va_release"); + if (region_vmm_test_hooks().use_fake) { + region_vmm_test_hooks().fake_va_bytes.erase(reinterpret_cast(va)); + return 0; + } + return acl_api().vmm_release_address(va); +} + +int region_vmm_free_physical(void *handle) { + region_vmm_record_issued("physical_free"); + if (region_vmm_test_hooks().use_fake) { + (void)handle; + return 0; + } + return acl_api().vmm_free_physical(handle); +} + +void region_vmm_zero_bytes(void *va, uint64_t offset, uint64_t nbytes) { + region_vmm_record_issued("zero_bytes"); + if (region_vmm_test_hooks().use_fake) { + auto &backing = region_vmm_test_hooks().fake_va_bytes; + auto it = backing.find(reinterpret_cast(va)); + if (it == backing.end()) { + throw std::runtime_error("region VMM fake zero_bytes has no reserved VA"); + } + uint64_t end = offset + nbytes; + if (end < offset || end > static_cast(it->second.size())) { + throw std::invalid_argument("region VMM fake zero_bytes range exceeds mapping"); + } + std::fill( + it->second.begin() + static_cast(offset), + it->second.begin() + static_cast(end), static_cast(0) + ); + return; + } + acl_api().vmm_zero_bytes(va, offset, nbytes); +} + +RegionVmmUnmapClass region_vmm_classify_unmap(int rc) { + if (rc == 0) { + return RegionVmmUnmapClass::Success; + } + if (region_vmm_test_hooks().use_fake && rc == 2) { + return RegionVmmUnmapClass::AlreadyGone; + } + return RegionVmmUnmapClass::Failed; +} + +void region_vmm_note_cleanup_failure(RegionVmmAllocation &record, const std::string &detail) { + record.local_cleanup_details.push_back(detail); + if (record.first_cleanup_failure.empty()) { + record.first_cleanup_failure = detail; + } +} + +bool region_vmm_cleanup_complete(const RegionVmmAllocation &record) { + if (record.mapped && !record.unmap_complete) { + return false; + } + if (record.va_reserved && !record.va_release_complete) { + return false; + } + if (record.physical_allocated && !record.physical_free_complete) { + return false; + } + return record.first_cleanup_failure.empty(); +} + +RegionVmmInspect region_vmm_inspect_record(uint64_t handle, const RegionVmmAllocation *record) { + RegionVmmInspect inspect{}; + if (record == nullptr) { + return inspect; + } + inspect.present = true; + inspect.device_id = record->device_id; + inspect.device_addr = reinterpret_cast(record->va); + inspect.mapping_bytes = record->mapping_bytes; + inspect.shareable_handle = record->shareable_handle; + inspect.physical_handle = reinterpret_cast(record->physical_handle); + inspect.physical_allocated = record->physical_allocated; + inspect.va_reserved = record->va_reserved; + inspect.mapped = record->mapped; + inspect.access_set = record->access_set; + inspect.exported = record->exported; + inspect.unmap_attempted = record->unmap_attempted; + inspect.unmap_complete = record->unmap_complete; + inspect.va_release_attempted = record->va_release_attempted; + inspect.va_release_complete = record->va_release_complete; + inspect.physical_free_attempted = record->physical_free_attempted; + inspect.physical_free_complete = record->physical_free_complete; + inspect.release_once_done = record->release_once_done; + inspect.first_cleanup_failure = record->first_cleanup_failure; + inspect.local_cleanup_details = record->local_cleanup_details; + (void)handle; + return inspect; +} + +uint64_t region_vmm_begin(int device_id) { + if (device_id < 0) { + throw std::invalid_argument("region VMM begin requires a non-negative device id"); + } + std::lock_guard lk(region_vmm_mu()); + uint64_t handle = region_vmm_next_handle(); + if (handle == 0) { + throw std::overflow_error("region VMM handle space is exhausted"); + } + RegionVmmAllocation record{}; + record.device_id = device_id; + auto inserted = g_region_vmm_allocations().emplace(handle, record); + if (!inserted.second) { + throw std::overflow_error("region VMM handle space is exhausted"); + } + region_vmm_next_handle() = handle + 1; + return handle; +} + +void region_vmm_run_allocate_stage( + RegionVmmStage stage, const std::function &issue, const std::function &commit +) { + if (region_vmm_consume_fail(stage, RegionVmmFailWhen::Before)) { + region_vmm_throw_stage_failure(stage, "before"); + } + issue(); + commit(); + if (region_vmm_consume_fail(stage, RegionVmmFailWhen::After)) { + region_vmm_throw_stage_failure(stage, "after"); + } +} + +RegionVmmExport region_vmm_allocate_export(uint64_t handle, uint64_t logical_bytes) { + std::lock_guard lk(region_vmm_mu()); + auto it = g_region_vmm_allocations().find(handle); + if (it == g_region_vmm_allocations().end()) { + throw std::runtime_error("region VMM allocate_export handle is unknown"); + } + RegionVmmAllocation &record = it->second; + if (record.physical_allocated || record.va_reserved || record.mapped || record.exported) { + throw std::runtime_error("region VMM allocate_export already has partial or complete state"); + } + region_vmm_bind_device(record.device_id); + uint64_t mapping_bytes = align_vmm_bytes(logical_bytes, region_vmm_granularity(record.device_id)); + record.mapping_bytes = mapping_bytes; + + region_vmm_run_allocate_stage( + RegionVmmStage::PhysicalAlloc, + [&]() { + record.physical_handle = region_vmm_malloc_physical(mapping_bytes, record.device_id); + }, + [&]() { + record.physical_allocated = true; + } + ); + region_vmm_run_allocate_stage( + RegionVmmStage::VaReserve, + [&]() { + record.va = region_vmm_reserve(mapping_bytes); + }, + [&]() { + record.va_reserved = true; + } + ); + region_vmm_run_allocate_stage( + RegionVmmStage::Map, + [&]() { + region_vmm_map(record.va, mapping_bytes, record.physical_handle); + }, + [&]() { + record.mapped = true; + } + ); + region_vmm_run_allocate_stage( + RegionVmmStage::SetAccess, + [&]() { + region_vmm_set_access(record.va, mapping_bytes, record.device_id); + }, + [&]() { + record.access_set = true; + } + ); + region_vmm_run_allocate_stage( + RegionVmmStage::Export, + [&]() { + record.shareable_handle = region_vmm_export_shareable(record.physical_handle); + }, + [&]() { + record.exported = true; + } + ); + + RegionVmmExport exported{}; + exported.device_addr = reinterpret_cast(record.va); + exported.mapping_bytes = record.mapping_bytes; + exported.shareable_handle = record.shareable_handle; + exported.registry_handle = handle; + return exported; +} + +bool region_vmm_run_cleanup_stage( + RegionVmmAllocation &record, RegionVmmStage stage, bool &attempted, bool &complete, + const std::function &issue, const std::function &commit_success +) { + attempted = true; + if (region_vmm_consume_fail(stage, RegionVmmFailWhen::Before)) { + region_vmm_note_cleanup_failure( + record, std::string("injected region VMM before ") + region_vmm_stage_name(stage) + " failure" + ); + return false; + } + int rc = issue(); + if (stage == RegionVmmStage::Unmap) { + RegionVmmUnmapClass classified = region_vmm_classify_unmap(rc); + if (classified == RegionVmmUnmapClass::Failed) { + region_vmm_note_cleanup_failure(record, std::string("unmap failed with code ") + std::to_string(rc)); + return false; + } + commit_success(); + complete = true; + if (region_vmm_consume_fail(stage, RegionVmmFailWhen::After)) { + region_vmm_note_cleanup_failure( + record, std::string("injected region VMM after ") + region_vmm_stage_name(stage) + " failure" + ); + } + return true; + } + if (rc != 0) { + region_vmm_note_cleanup_failure( + record, std::string(region_vmm_stage_name(stage)) + " failed with code " + std::to_string(rc) + ); + return false; + } + commit_success(); + complete = true; + if (region_vmm_consume_fail(stage, RegionVmmFailWhen::After)) { + region_vmm_note_cleanup_failure( + record, std::string("injected region VMM after ") + region_vmm_stage_name(stage) + " failure" + ); + return false; + } + return true; +} + +void region_vmm_release(uint64_t handle) { + std::lock_guard lk(region_vmm_mu()); + auto it = g_region_vmm_allocations().find(handle); + if (it == g_region_vmm_allocations().end()) { + return; + } + RegionVmmAllocation &record = it->second; + if (record.release_once_done) { + if (!record.first_cleanup_failure.empty()) { + throw std::runtime_error(record.first_cleanup_failure); + } + return; + } + record.release_once_done = true; + try { + region_vmm_bind_device(record.device_id); + } catch (const std::exception &exc) { + region_vmm_note_cleanup_failure(record, std::string("bind_device failed: ") + exc.what()); + throw std::runtime_error(record.first_cleanup_failure); + } + + bool unmap_ok = true; + if (record.mapped) { + unmap_ok = region_vmm_run_cleanup_stage( + record, RegionVmmStage::Unmap, record.unmap_attempted, record.unmap_complete, + [&]() { + return region_vmm_unmap(record.va); + }, + [&]() { + record.mapped = false; + } + ); + } + + if (record.mapped && !unmap_ok) { + throw std::runtime_error(record.first_cleanup_failure); + } + + if (record.va_reserved && !record.va_release_attempted) { + region_vmm_run_cleanup_stage( + record, RegionVmmStage::VaRelease, record.va_release_attempted, record.va_release_complete, + [&]() { + return region_vmm_release_address(record.va); + }, + [&]() { + record.va_reserved = false; + record.va = nullptr; + } + ); + } + if (record.physical_allocated && !record.physical_free_attempted) { + region_vmm_run_cleanup_stage( + record, RegionVmmStage::PhysicalFree, record.physical_free_attempted, record.physical_free_complete, + [&]() { + return region_vmm_free_physical(record.physical_handle); + }, + [&]() { + record.physical_allocated = false; + record.physical_handle = nullptr; + } + ); + } + + if (region_vmm_cleanup_complete(record)) { + g_region_vmm_allocations().erase(it); + return; + } + throw std::runtime_error( + record.first_cleanup_failure.empty() ? "region VMM release left uncleared ledger state" : + record.first_cleanup_failure + ); +} + +void region_vmm_zero_range(uint64_t handle, uint64_t offset, uint64_t nbytes) { + std::lock_guard lk(region_vmm_mu()); + auto it = g_region_vmm_allocations().find(handle); + if (it == g_region_vmm_allocations().end()) { + throw std::runtime_error("region VMM zero_bytes handle is unknown"); + } + RegionVmmAllocation &record = it->second; + if (!record.mapped || record.va == nullptr) { + throw std::runtime_error("region VMM zero_bytes requires a mapped VA"); + } + uint64_t end = offset + nbytes; + if (end < offset || end > record.mapping_bytes) { + throw std::invalid_argument("region VMM zero_bytes range exceeds mapping_bytes"); + } + region_vmm_bind_device(record.device_id); + if (region_vmm_consume_fail(RegionVmmStage::ZeroBytes, RegionVmmFailWhen::Before)) { + region_vmm_throw_stage_failure(RegionVmmStage::ZeroBytes, "before"); + } + region_vmm_zero_bytes(record.va, offset, nbytes); + if (region_vmm_consume_fail(RegionVmmStage::ZeroBytes, RegionVmmFailWhen::After)) { + region_vmm_throw_stage_failure(RegionVmmStage::ZeroBytes, "after"); + } +} + +RegionVmmInspect region_vmm_inspect(uint64_t handle) { + std::lock_guard lk(region_vmm_mu()); + auto it = g_region_vmm_allocations().find(handle); + if (it == g_region_vmm_allocations().end()) { + return RegionVmmInspect{}; + } + return region_vmm_inspect_record(handle, &it->second); +} + +std::vector region_vmm_live_handles() { + std::lock_guard lk(region_vmm_mu()); + std::vector handles; + handles.reserve(g_region_vmm_allocations().size()); + for (const auto &entry : g_region_vmm_allocations()) { + handles.push_back(entry.first); + } + return handles; +} + +std::vector region_vmm_test_fake_bytes(uint64_t handle) { + std::lock_guard lk(region_vmm_mu()); + if (!region_vmm_test_hooks().use_fake) { + throw std::runtime_error("region VMM fake bytes require the fake driver"); + } + auto it = g_region_vmm_allocations().find(handle); + if (it == g_region_vmm_allocations().end() || it->second.va == nullptr) { + throw std::runtime_error("region VMM fake bytes require a reserved VA"); + } + auto bytes_it = region_vmm_test_hooks().fake_va_bytes.find(reinterpret_cast(it->second.va)); + if (bytes_it == region_vmm_test_hooks().fake_va_bytes.end()) { + throw std::runtime_error("region VMM fake bytes have no reserved VA backing"); + } + return bytes_it->second; +} + +std::vector region_vmm_issued_ops() { + std::lock_guard lk(region_vmm_mu()); + return region_vmm_test_hooks().issued_ops; +} + +void region_vmm_test_use_fake_driver() { + std::lock_guard lk(region_vmm_mu()); + region_vmm_test_hooks().use_fake = true; +} + +void region_vmm_test_reset_hooks() { + std::lock_guard lk(region_vmm_mu()); + RegionVmmTestHooks &hooks = region_vmm_test_hooks(); + hooks.fail_stage.reset(); + hooks.fail_when = RegionVmmFailWhen::Before; + hooks.unmap_already_gone = false; + hooks.issued_ops.clear(); +} + +void region_vmm_test_fail_stage(const std::string &stage, const std::string &when) { + std::lock_guard lk(region_vmm_mu()); + RegionVmmTestHooks &hooks = region_vmm_test_hooks(); + hooks.fail_stage = region_vmm_parse_stage(stage); + if (when == "before") { + hooks.fail_when = RegionVmmFailWhen::Before; + } else if (when == "after") { + hooks.fail_when = RegionVmmFailWhen::After; + } else { + throw std::invalid_argument("region VMM fail when must be before or after"); + } +} + +void region_vmm_test_set_unmap_already_gone() { + std::lock_guard lk(region_vmm_mu()); + region_vmm_test_hooks().unmap_already_gone = true; +} + // The int wire value of a dtype given either a DataType enumerator or its int value. The nanobind // DataType enum is not arithmetic, so a caller holding one has only `.value`; accept both forms. uint8_t datatype_wire_value(nb::object dtype) { @@ -2509,11 +3157,33 @@ NB_MODULE(_task_interface, m) { "way out. Tags are not preserved (the wire format strips them)." ); - nb::class_(m, "_ChipChildOnboardRegionExport") - .def_ro("device_addr", &ChipChildOnboardRegionExport::device_addr) - .def_ro("mapping_bytes", &ChipChildOnboardRegionExport::mapping_bytes) - .def_ro("shareable_handle", &ChipChildOnboardRegionExport::shareable_handle) - .def_ro("registry_handle", &ChipChildOnboardRegionExport::registry_handle); + nb::class_(m, "_RegionVmmExport") + .def_ro("device_addr", &RegionVmmExport::device_addr) + .def_ro("mapping_bytes", &RegionVmmExport::mapping_bytes) + .def_ro("shareable_handle", &RegionVmmExport::shareable_handle) + .def_ro("registry_handle", &RegionVmmExport::registry_handle); + + nb::class_(m, "_RegionVmmInspect") + .def_ro("present", &RegionVmmInspect::present) + .def_ro("device_id", &RegionVmmInspect::device_id) + .def_ro("device_addr", &RegionVmmInspect::device_addr) + .def_ro("mapping_bytes", &RegionVmmInspect::mapping_bytes) + .def_ro("shareable_handle", &RegionVmmInspect::shareable_handle) + .def_ro("physical_handle", &RegionVmmInspect::physical_handle) + .def_ro("physical_allocated", &RegionVmmInspect::physical_allocated) + .def_ro("va_reserved", &RegionVmmInspect::va_reserved) + .def_ro("mapped", &RegionVmmInspect::mapped) + .def_ro("access_set", &RegionVmmInspect::access_set) + .def_ro("exported", &RegionVmmInspect::exported) + .def_ro("unmap_attempted", &RegionVmmInspect::unmap_attempted) + .def_ro("unmap_complete", &RegionVmmInspect::unmap_complete) + .def_ro("va_release_attempted", &RegionVmmInspect::va_release_attempted) + .def_ro("va_release_complete", &RegionVmmInspect::va_release_complete) + .def_ro("physical_free_attempted", &RegionVmmInspect::physical_free_attempted) + .def_ro("physical_free_complete", &RegionVmmInspect::physical_free_complete) + .def_ro("release_once_done", &RegionVmmInspect::release_once_done) + .def_ro("first_cleanup_failure", &RegionVmmInspect::first_cleanup_failure) + .def_ro("local_cleanup_details", &RegionVmmInspect::local_cleanup_details); nb::class_(m, "_RegionHandle").def("__int__", &RegionHandle::value); @@ -2815,60 +3485,92 @@ NB_MODULE(_task_interface, m) { nb::call_guard(), "Poll one L3 Host-side L3-L2 signal counter until match or timeout." ); m.def( - "_l3_child_onboard_region_create", - [](uint64_t nbytes) -> ChipChildOnboardRegionExport { - if (nbytes == 0 || nbytes > static_cast(std::numeric_limits::max())) { - throw std::invalid_argument("L3-L2 onboard child region requires a positive mapping size"); - } - AclRuntimeApi &api = acl_api(); - int device_id = api.current_device_with_check(); - uint64_t mapping_bytes = align_vmm_bytes(nbytes, api.vmm_granularity_with_check(device_id)); - ChipChildOnboardRegion region{}; - region.device_id = device_id; - region.mapping_bytes = mapping_bytes; - region.vmm_handle = api.vmm_malloc_physical_with_check(mapping_bytes, device_id); - void *mapped_addr = nullptr; - try { - mapped_addr = api.vmm_reserve_with_check(mapping_bytes); - api.vmm_map_with_check(mapped_addr, mapping_bytes, region.vmm_handle); - api.vmm_set_access_with_check(mapped_addr, mapping_bytes, device_id); - region.shareable_handle = api.vmm_export_shareable_with_check(region.vmm_handle); - } catch (...) { - std::string cleanup_error; - api.vmm_release_collecting(mapped_addr, region.vmm_handle, cleanup_error); - throw; - } - region.device_addr = reinterpret_cast(mapped_addr); - uint64_t registry_handle = g_chip_child_onboard_regions.emplace(region); - return ChipChildOnboardRegionExport{ - region.device_addr, - region.mapping_bytes, - region.shareable_handle, - registry_handle, - }; + "_region_vmm_begin", + [](int device_id) -> uint64_t { + return region_vmm_begin(device_id); }, - nb::arg("nbytes"), nb::call_guard(), - "Create and export a child-owned onboard VMM region." + nb::arg("device_id"), nb::call_guard(), + "Insert an empty native VMM owner record with no physical side effect." ); m.def( - "_l3_child_onboard_region_close", - [](uint64_t registry_handle) { - std::optional removed = g_chip_child_onboard_regions.remove(registry_handle); - if (!removed.has_value()) { - return; - } - ChipChildOnboardRegion region = *removed; - region.bind_acl_device(); - std::string cleanup_error; - acl_api().vmm_release_collecting( - reinterpret_cast(static_cast(region.device_addr)), region.vmm_handle, cleanup_error - ); - if (!cleanup_error.empty()) { - throw std::runtime_error(cleanup_error); - } + "_region_vmm_allocate_export", + [](uint64_t handle, uint64_t logical_bytes) -> RegionVmmExport { + return region_vmm_allocate_export(handle, logical_bytes); }, - nb::arg("registry_handle"), nb::call_guard(), "Close a child-owned onboard VMM region." + nb::arg("handle"), nb::arg("logical_bytes"), nb::call_guard(), + "Physically allocate, map, and export one native VMM record without rollback." + ); + m.def( + "_region_vmm_release", + [](uint64_t handle) { + region_vmm_release(handle); + }, + nb::arg("handle"), nb::call_guard(), + "Make one dependency-aware cleanup pass and delete the record only after complete success." + ); + m.def( + "_region_vmm_zero_bytes", + [](uint64_t handle, uint64_t offset, uint64_t nbytes) { + region_vmm_zero_range(handle, offset, nbytes); + }, + nb::arg("handle"), nb::arg("offset"), nb::arg("nbytes"), nb::call_guard(), + "Zero a mapped native VMM byte range." + ); + m.def( + "_region_vmm_inspect", + [](uint64_t handle) -> RegionVmmInspect { + return region_vmm_inspect(handle); + }, + nb::arg("handle"), "Inspect one native VMM ledger record." + ); + m.def( + "_region_vmm_test_use_fake_driver", + []() { + region_vmm_test_use_fake_driver(); + }, + "Route native VMM physical calls through the in-process fake driver." + ); + m.def( + "_region_vmm_test_reset_hooks", + []() { + region_vmm_test_reset_hooks(); + }, + "Clear native VMM fail-injection and issued-op logs without dropping live records." + ); + m.def( + "_region_vmm_test_fail_stage", + [](const std::string &stage, const std::string &when) { + region_vmm_test_fail_stage(stage, when); + }, + nb::arg("stage"), nb::arg("when"), "Fail one named native VMM stage before or after ledger mutation." + ); + m.def( + "_region_vmm_test_set_unmap_already_gone", + []() { + region_vmm_test_set_unmap_already_gone(); + }, + "Make the next fake unmap report a backend-defined already-gone outcome." + ); + m.def( + "_region_vmm_test_issued_ops", + []() -> std::vector { + return region_vmm_issued_ops(); + }, + "Return native VMM driver operations issued since the last hook reset." + ); + m.def( + "_region_vmm_test_live_handles", + []() -> std::vector { + return region_vmm_live_handles(); + }, + "Return currently retained native VMM registry handles." + ); + m.def( + "_region_vmm_test_fake_bytes", + [](uint64_t handle) -> std::vector { + return region_vmm_test_fake_bytes(handle); + }, + nb::arg("handle"), "Return fake-driver VA backing bytes for one native VMM record." ); - bind_worker(m); } diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 6821055a5a..6c148ac321 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -845,14 +845,14 @@ inline void bind_worker(nb::module_ &m) { "Drive one NEXT_LEVEL chip child through CTRL_COMM_INIT (lazy base comm init)." ) .def( - "control_worker_chip_region_create", &Worker::control_worker_chip_region_create, nb::arg("worker_id"), + "control_region_allocate", &Worker::control_region_allocate, nb::arg("worker_id"), nb::arg("request_shm_name"), nb::arg("reply_shm_name"), nb::call_guard(), - "Drive one NEXT_LEVEL chip child through CTRL_WORKER_CHIP_REGION_CREATE." + "Drive one NEXT_LEVEL chip child through CTRL_REGION_ALLOCATE." ) .def( - "control_worker_chip_region_release", &Worker::control_worker_chip_region_release, nb::arg("worker_id"), - nb::arg("region_id"), nb::call_guard(), - "Drive one NEXT_LEVEL chip child through CTRL_WORKER_CHIP_REGION_RELEASE." + "control_region_release", &Worker::control_region_release, nb::arg("worker_id"), + nb::arg("request_shm_name"), nb::arg("reply_shm_name"), nb::call_guard(), + "Drive one NEXT_LEVEL chip child through CTRL_REGION_RELEASE." ); m.attr("DEFAULT_HEAP_RING_SIZE") = static_cast(DEFAULT_HEAP_RING_SIZE); diff --git a/python/simpler/comm_provider.py b/python/simpler/comm_provider.py new file mode 100644 index 0000000000..db86a69f56 --- /dev/null +++ b/python/simpler/comm_provider.py @@ -0,0 +1,1249 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Neutral provider-region values, typed codes, store, and structured results. + +This module is the provider-agent value surface for CPU-NPU comm regions. It +does not import worker-chip compatibility types, Worker, mailbox transport, or +W5a transaction identity. +""" + +from __future__ import annotations + +import ctypes +import errno +import logging +import uuid +from dataclasses import dataclass +from enum import Enum, IntEnum +from multiprocessing.shared_memory import SharedMemory +from typing import Any, Callable, Protocol, TypeVar, Union + +from _task_interface import ( # pyright: ignore[reportMissingImports] + BackendKind, + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_release, + _region_vmm_zero_bytes, +) + +_UINT64_MAX = (1 << 64) - 1 +_INT32_MIN = -(1 << 31) +_INT32_MAX = (1 << 31) - 1 +_COUNTER_LOGICAL_ALIGNMENT = 4 +_COUNTER_BASE_ALIGNMENT = 64 +POSIX_SHM_TOKEN_MAX_BYTES = 32 + + +class RegionPartKind(IntEnum): + INVALID = 0 + PAYLOAD = 1 + COUNTER = 2 + + +class RegionEnvironmentKind(str, Enum): + SIM = "SIM" + ONBOARD = "ONBOARD" + + +class ProviderRegionStoreState(str, Enum): + OPEN = "OPEN" + CLOSING = "CLOSING" + CLOSE_FAILED = "CLOSE_FAILED" + CLOSED = "CLOSED" + + +class ProviderRegionResourceState(str, Enum): + CREATING = "CREATING" + ACTIVE = "ACTIVE" + CLEANUP_PENDING = "CLEANUP_PENDING" + + +class ProviderPartResourceState(str, Enum): + SHELL = "SHELL" + MATERIALIZING = "MATERIALIZING" + READY = "READY" + CLEANUP_PENDING = "CLEANUP_PENDING" + RELEASED = "RELEASED" + + +class ProviderReleaseStatus(IntEnum): + RELEASED = 1 + ALREADY_GONE = 2 + UNKNOWN_RESOURCE = 3 + CLEANUP_INCOMPLETE = 4 + + +class RegionControlErrorKind(IntEnum): + NONE = 0 + BAD_MAGIC_VERSION = 1 + BAD_MESSAGE_SIZE = 2 + INVALID_ENUM_VALUE = 3 + RESERVED_NONZERO = 4 + INVALID_FIELD_VALUE = 5 + STORE_LIFECYCLE = 6 + INTERNAL_INVARIANT = 7 + BACKEND_FAILURE = 8 + + +class RegionOperationKind(IntEnum): + NONE = 0 + MATERIALIZE = 1 + ZERO_BYTES = 2 + DESCRIBE = 3 + LOCAL_VIEW = 4 + RELEASE = 5 + + +class RegionCleanupCause(IntEnum): + NONE = 0 + BACKEND_ERROR = 1 + INTERRUPTED = 2 + BACKEND_STATE_MISMATCH = 3 + + +REGION_PARTS = (RegionPartKind.PAYLOAD, RegionPartKind.COUNTER) +_ALLOCATION_ERROR_KINDS = ( + RegionControlErrorKind.BACKEND_FAILURE, + RegionControlErrorKind.INTERNAL_INVARIANT, +) +_BACKEND_FAILURE_OPERATIONS = ( + RegionOperationKind.MATERIALIZE, + RegionOperationKind.ZERO_BYTES, + RegionOperationKind.DESCRIBE, + RegionOperationKind.LOCAL_VIEW, +) + + +def _require_int(name: str, value: object) -> int: + if type(value) is not int: + raise TypeError(f"{name} must be an int") + return value + + +def _require_bool(name: str, value: object) -> bool: + if type(value) is not bool: + raise TypeError(f"{name} must be a bool") + return value + + +def _require_uint64(name: str, value: object) -> int: + number = _require_int(name, value) + if number < 0 or number > _UINT64_MAX: + raise ValueError(f"{name} overflowed uint64") + return number + + +def _require_positive_uint64(name: str, value: object) -> int: + number = _require_uint64(name, value) + if number < 1: + raise ValueError(f"{name} must be positive") + return number + + +def _require_nonzero_uint64(name: str, value: object) -> int: + return _require_positive_uint64(name, value) + + +def _require_int32(name: str, value: object) -> int: + number = _require_int(name, value) + if number < _INT32_MIN or number > _INT32_MAX: + raise ValueError(f"{name} overflowed int32") + return number + + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +def _require_enum(enum_cls: type[_EnumT], value: object, name: str) -> _EnumT: + try: + return enum_cls(value) + except ValueError as exc: + raise ValueError(f"unknown {name} {value!r}") from exc + + +def _require_backend_kind(value: object) -> BackendKind: + return _require_enum(BackendKind, value, "planned backing kind") + + +def _require_part(value: object, *, allow_invalid: bool = False) -> RegionPartKind: + part = _require_enum(RegionPartKind, value, "region part") + if part is RegionPartKind.INVALID and not allow_invalid: + raise ValueError("region part must be PAYLOAD or COUNTER") + return part + + +def _require_posix_shm_token(value: object) -> str: + if not isinstance(value, str): + raise TypeError("POSIX shm token must be str") + try: + encoded = value.encode("ascii") + except UnicodeEncodeError as exc: + raise ValueError("POSIX shm token must be ASCII") from exc + if not encoded or len(encoded) > POSIX_SHM_TOKEN_MAX_BYTES: + raise ValueError("POSIX shm token length must be in 1..32 bytes") + if b"\x00" in encoded: + raise ValueError("POSIX shm token must not contain NUL") + return value + + +def _checked_add_u64(lhs: int, rhs: int, name: str) -> int: + total = lhs + rhs + if total > _UINT64_MAX: + raise ValueError(f"{name} overflowed uint64") + return total + + +def _require_counter_logical_bytes(value: object) -> int: + logical_bytes = _require_positive_uint64("COUNTER logical_bytes", value) + if logical_bytes % _COUNTER_LOGICAL_ALIGNMENT != 0: + raise ValueError("COUNTER logical_bytes must be a multiple of 4") + return logical_bytes + + +@dataclass(frozen=True) +class DeviceAllocationTarget: + device_id: int + + def __post_init__(self) -> None: + object.__setattr__(self, "device_id", _require_int32("device_id", self.device_id)) + + +@dataclass(frozen=True) +class HostAllocationTarget: + pass + + +AllocationTarget = Union[DeviceAllocationTarget, HostAllocationTarget] + + +@dataclass(frozen=True) +class RegionAllocationContext: + environment_kind: RegionEnvironmentKind + target: AllocationTarget + + def __post_init__(self) -> None: + environment = _require_enum(RegionEnvironmentKind, self.environment_kind, "environment kind") + object.__setattr__(self, "environment_kind", environment) + if not isinstance(self.target, (DeviceAllocationTarget, HostAllocationTarget)): + raise TypeError("allocation target must be DeviceAllocationTarget or HostAllocationTarget") + + +@dataclass(frozen=True) +class RegionPartAllocationSpec: + planned_backing_kind: BackendKind + logical_bytes: int + + def __post_init__(self) -> None: + object.__setattr__(self, "planned_backing_kind", _require_backend_kind(self.planned_backing_kind)) + object.__setattr__(self, "logical_bytes", _require_positive_uint64("logical_bytes", self.logical_bytes)) + + +@dataclass(frozen=True) +class RegionAllocationSpec: + payload: RegionPartAllocationSpec + counter: RegionPartAllocationSpec + + def __post_init__(self) -> None: + if not isinstance(self.payload, RegionPartAllocationSpec): + raise TypeError("payload spec must be RegionPartAllocationSpec") + if not isinstance(self.counter, RegionPartAllocationSpec): + raise TypeError("counter spec must be RegionPartAllocationSpec") + _require_counter_logical_bytes(self.counter.logical_bytes) + + def part(self, kind: RegionPartKind) -> RegionPartAllocationSpec: + part = _require_part(kind) + if part is RegionPartKind.PAYLOAD: + return self.payload + return self.counter + + +@dataclass(frozen=True) +class VmmShareableHandleImport: + device_id: int + shareable_handle: int + + def __post_init__(self) -> None: + object.__setattr__(self, "device_id", _require_int32("device_id", self.device_id)) + object.__setattr__( + self, + "shareable_handle", + _require_uint64("shareable_handle", self.shareable_handle), + ) + + +@dataclass(frozen=True) +class PosixShmImport: + shm_name: str + + def __post_init__(self) -> None: + object.__setattr__(self, "shm_name", _require_posix_shm_token(self.shm_name)) + + +ImportCapability = Union[VmmShareableHandleImport, PosixShmImport] + + +@dataclass(frozen=True) +class RegionPartExportDescriptor: + planned_backing_kind: BackendKind + logical_bytes: int + mapping_bytes: int + import_capability: ImportCapability + + def __post_init__(self) -> None: + object.__setattr__(self, "planned_backing_kind", _require_backend_kind(self.planned_backing_kind)) + logical_bytes = _require_positive_uint64("logical_bytes", self.logical_bytes) + mapping_bytes = _require_uint64("mapping_bytes", self.mapping_bytes) + if mapping_bytes < logical_bytes: + raise ValueError("mapping_bytes must cover logical_bytes") + if not isinstance(self.import_capability, (VmmShareableHandleImport, PosixShmImport)): + raise TypeError("import_capability must be a single VMM or POSIX variant") + object.__setattr__(self, "logical_bytes", logical_bytes) + object.__setattr__(self, "mapping_bytes", mapping_bytes) + + +@dataclass(frozen=True) +class RegionExportDescriptor: + payload: RegionPartExportDescriptor + counter: RegionPartExportDescriptor + + def __post_init__(self) -> None: + if not isinstance(self.payload, RegionPartExportDescriptor): + raise TypeError("payload export descriptor must be RegionPartExportDescriptor") + if not isinstance(self.counter, RegionPartExportDescriptor): + raise TypeError("counter export descriptor must be RegionPartExportDescriptor") + _require_counter_logical_bytes(self.counter.logical_bytes) + + +@dataclass(frozen=True) +class RegionPartLocalView: + part: RegionPartKind + local_base: int + logical_bytes: int + + def __post_init__(self) -> None: + part = _require_part(self.part) + local_base = _require_uint64("local_base", self.local_base) + if part is RegionPartKind.COUNTER: + logical_bytes = _require_counter_logical_bytes(self.logical_bytes) + if local_base % _COUNTER_BASE_ALIGNMENT != 0: + raise ValueError("COUNTER local_base must be 64-byte aligned") + else: + logical_bytes = _require_positive_uint64("logical_bytes", self.logical_bytes) + _checked_add_u64(local_base, logical_bytes, f"{part.name} local span") + object.__setattr__(self, "part", part) + object.__setattr__(self, "local_base", local_base) + object.__setattr__(self, "logical_bytes", logical_bytes) + + @property + def local_end(self) -> int: + return self.local_base + self.logical_bytes + + +def validate_independent_local_views( + payload: RegionPartLocalView, counter: RegionPartLocalView +) -> tuple[RegionPartLocalView, RegionPartLocalView]: + if not isinstance(payload, RegionPartLocalView) or payload.part is not RegionPartKind.PAYLOAD: + raise ValueError("payload local view must use part PAYLOAD") + if not isinstance(counter, RegionPartLocalView) or counter.part is not RegionPartKind.COUNTER: + raise ValueError("counter local view must use part COUNTER") + if payload.local_base < counter.local_end and counter.local_base < payload.local_end: + raise ValueError("PAYLOAD and COUNTER local spans must not overlap") + return payload, counter + + +@dataclass(frozen=True) +class RegionAllocationResult: + provider_resource_id: int + export_descriptor: RegionExportDescriptor + + def __post_init__(self) -> None: + object.__setattr__( + self, + "provider_resource_id", + _require_nonzero_uint64("provider_resource_id", self.provider_resource_id), + ) + if not isinstance(self.export_descriptor, RegionExportDescriptor): + raise TypeError("export_descriptor must be RegionExportDescriptor") + + +@dataclass(frozen=True) +class ProviderCleanupFailure: + part: RegionPartKind + backend_operation: RegionOperationKind + typed_cause: RegionCleanupCause + + def __post_init__(self) -> None: + part = _require_part(self.part) + operation = _require_enum(RegionOperationKind, self.backend_operation, "backend operation") + cause = _require_enum(RegionCleanupCause, self.typed_cause, "cleanup cause") + if operation is RegionOperationKind.NONE: + raise ValueError("cleanup failure requires a nonzero backend operation") + if cause is RegionCleanupCause.NONE: + raise ValueError("cleanup failure requires a nonzero typed cause") + object.__setattr__(self, "part", part) + object.__setattr__(self, "backend_operation", operation) + object.__setattr__(self, "typed_cause", cause) + + +@dataclass(frozen=True) +class ProviderReleaseResult: + provider_resource_id: int + status: ProviderReleaseStatus + failures: tuple[ProviderCleanupFailure, ...] = () + + def __post_init__(self) -> None: + object.__setattr__( + self, + "provider_resource_id", + _require_nonzero_uint64("provider_resource_id", self.provider_resource_id), + ) + status = _require_enum(ProviderReleaseStatus, self.status, "release status") + object.__setattr__(self, "status", status) + failures = tuple(self.failures) + if any(not isinstance(failure, ProviderCleanupFailure) for failure in failures): + raise TypeError("release failures must be ProviderCleanupFailure values") + parts = [failure.part for failure in failures] + if len(parts) != len(set(parts)): + raise ValueError("ProviderReleaseResult allows at most one failure per part") + if status is ProviderReleaseStatus.CLEANUP_INCOMPLETE: + if not failures: + raise ValueError("CLEANUP_INCOMPLETE requires structured per-part failures") + elif failures: + raise ValueError(f"{status.name} must not carry cleanup failures") + object.__setattr__(self, "failures", failures) + + +class RegionProviderError(Exception): + pass + + +class RegionControlError(RegionProviderError): + def __init__( + self, + kind: RegionControlErrorKind | int, + message: str = "", + *, + failed_part: RegionPartKind | int = RegionPartKind.INVALID, + failed_operation: RegionOperationKind | int = RegionOperationKind.NONE, + ) -> None: + control_kind = _require_enum(RegionControlErrorKind, kind, "control error kind") + if control_kind is RegionControlErrorKind.NONE: + raise ValueError("RegionControlError requires a nonzero control kind") + part = _require_part(failed_part, allow_invalid=True) + operation = _require_enum(RegionOperationKind, failed_operation, "failed operation") + self.kind = control_kind + self.failed_part = part + self.failed_operation = operation + self.message = str(message) + super().__init__(self.message or control_kind.name) + + +class RegionAllocationError(RegionProviderError): + def __init__( + self, + *, + provisional_resource_id: int, + control_kind: RegionControlErrorKind, + failed_part: RegionPartKind | int, + failed_operation: RegionOperationKind | int, + cleanup_debt_remaining: bool, + message: str = "", + ) -> None: + resource_id = _require_nonzero_uint64("provisional_resource_id", provisional_resource_id) + kind = _require_enum(RegionControlErrorKind, control_kind, "control error kind") + if kind not in _ALLOCATION_ERROR_KINDS: + raise ValueError("RegionAllocationError control kind must be BACKEND_FAILURE or INTERNAL_INVARIANT") + part = _require_part(failed_part, allow_invalid=True) + operation = _require_enum(RegionOperationKind, failed_operation, "failed operation") + if kind is RegionControlErrorKind.BACKEND_FAILURE and operation not in _BACKEND_FAILURE_OPERATIONS: + raise ValueError("BACKEND_FAILURE requires a create-path operation") + if operation is RegionOperationKind.RELEASE: + raise ValueError("RegionAllocationError must not use RELEASE as the failed operation") + self.provisional_resource_id = resource_id + self.control_kind = kind + self.failed_part = part + self.failed_operation = operation + self.cleanup_debt_remaining = _require_bool("cleanup_debt_remaining", cleanup_debt_remaining) + self.message = str(message) + super().__init__(self.message or kind.name) + + +class RegionPartAllocation(Protocol): + """Owning backend shell for one PAYLOAD or COUNTER allocation.""" + + def materialize(self) -> None: ... + + def mapping_bytes(self) -> int: ... + + def import_capability(self) -> ImportCapability: ... + + def local_base(self) -> int: ... + + def zero_bytes(self, offset: int, nbytes: int) -> None: ... + + def release_once(self) -> ProviderCleanupFailure | None: ... + + +ShellFactory = Callable[..., Any] + + +def _generate_posix_shm_token() -> str: + return _require_posix_shm_token("smp_" + uuid.uuid4().hex[:24]) + + +def _posix_shm_create_name(token: str) -> str: + return token[1:] if token.startswith("/") else token + + +def _unlink_posix_shm_token(token: str) -> None: + import _posixshmem # noqa: PLC0415 + from multiprocessing import resource_tracker # noqa: PLC0415 + + name = token if token.startswith("/") else f"/{token}" + try: + _posixshmem.shm_unlink(name) + except FileNotFoundError: + return + except OSError as exc: + if exc.errno == errno.ENOENT: + return + raise + try: + resource_tracker.unregister(name, "shared_memory") + except Exception as exc: + logging.getLogger("simpler").warning( + "POSIX shm resource-tracker unregister failed: token=%s error=%s", + token, + exc, + ) + + +def _require_shm_buf(shm: SharedMemory) -> memoryview: + buf = shm.buf + if buf is None: + raise RuntimeError("shared memory mapping is missing") + return buf + + +def _shm_local_base(shm: SharedMemory) -> int: + exported = ctypes.c_char.from_buffer(_require_shm_buf(shm)) + try: + return ctypes.addressof(exported) + finally: + del exported + + +class SimPosixShmAllocation: + """One POSIX shm object for a single PAYLOAD or COUNTER part.""" + + def __init__( + self, + context: RegionAllocationContext, + part: RegionPartKind, + spec: RegionPartAllocationSpec, + *, + candidate_name: str | None = None, + shm_cls: Any = SharedMemory, + ) -> None: + del context + self._part = _require_part(part) + if not isinstance(spec, RegionPartAllocationSpec): + raise TypeError("spec must be RegionPartAllocationSpec") + self._spec = spec + self._shm_cls = shm_cls + self.candidate_name = _require_posix_shm_token(candidate_name or _generate_posix_shm_token()) + self._name_ownership_known = False + self._shm_object_installed = False + self._mapping_available = False + self._shm: SharedMemory | None = None + self._local_base: int | None = None + self._close_attempted = False + self._close_complete = False + self._unlink_attempted = False + self._unlink_complete = False + self._release_once_count = 0 + self._first_cleanup_failure: ProviderCleanupFailure | None = None + self.local_cleanup_details: list[tuple[str, BaseException]] = [] + + def materialize(self) -> None: + try: + shm = self._shm_cls( + name=_posix_shm_create_name(self.candidate_name), + create=True, + size=self._spec.logical_bytes, + ) + except FileExistsError as exc: + raise RegionControlError( + RegionControlErrorKind.BACKEND_FAILURE, + "POSIX shm name collision", + failed_part=self._part, + failed_operation=RegionOperationKind.MATERIALIZE, + ) from exc + except BaseException: + self._name_ownership_known = True + raise + self._shm = shm + self._name_ownership_known = True + self._shm_object_installed = True + self._local_base = _shm_local_base(shm) + self._mapping_available = True + + def mapping_bytes(self) -> int: + self._require_mapping("mapping_bytes") + return self._spec.logical_bytes + + def import_capability(self) -> PosixShmImport: + self._require_mapping("import_capability") + return PosixShmImport(shm_name=self.candidate_name) + + def local_base(self) -> int: + self._require_mapping("local_base") + if self._local_base is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "POSIX shm mapping has no local base", + failed_part=self._part, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) + return self._local_base + + def zero_bytes(self, offset: int, nbytes: int) -> None: + self._require_mapping("zero_bytes") + if type(offset) is not int or type(nbytes) is not int: + raise TypeError("zero_bytes offset and nbytes must be int") + if offset < 0 or nbytes < 0: + raise ValueError("zero_bytes offset and nbytes must be non-negative") + end = offset + nbytes + if end > self._spec.logical_bytes: + raise ValueError("zero_bytes range exceeds logical_bytes") + assert self._shm is not None + _require_shm_buf(self._shm)[offset:end] = b"\x00" * nbytes + + def release_once(self) -> ProviderCleanupFailure | None: + self._release_once_count += 1 + if self._release_once_count > 1: + return self._first_cleanup_failure + if self._shm_object_installed and not self._close_attempted: + self._close_attempted = True + try: + if self._shm is not None: + self._shm.close() + self._close_complete = True + except BaseException as exc: + self._record_cleanup_failure("close", exc) + if self._name_ownership_known and not self._unlink_attempted: + self._unlink_attempted = True + try: + self._unlink_owned_name() + self._unlink_complete = True + except BaseException as exc: + self._record_cleanup_failure("unlink", exc) + return self._first_cleanup_failure + + def _unlink_owned_name(self) -> None: + if self._shm is not None: + try: + self._shm.unlink() + return + except FileNotFoundError: + return + except OSError as exc: + if exc.errno == errno.ENOENT: + return + raise + _unlink_posix_shm_token(self.candidate_name) + + def _record_cleanup_failure(self, step: str, exc: BaseException) -> None: + cause = RegionCleanupCause.INTERRUPTED if _interrupt_like(exc) else RegionCleanupCause.BACKEND_ERROR + failure = ProviderCleanupFailure( + part=self._part, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=cause, + ) + self.local_cleanup_details.append((step, exc)) + if self._first_cleanup_failure is None: + self._first_cleanup_failure = failure + + def _require_mapping(self, operation: str) -> None: + if not self._mapping_available or self._shm is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + f"{operation} requires a materialized POSIX shm mapping", + failed_part=self._part, + failed_operation=RegionOperationKind.DESCRIBE + if operation in {"mapping_bytes", "import_capability"} + else RegionOperationKind.LOCAL_VIEW + if operation == "local_base" + else RegionOperationKind.ZERO_BYTES, + ) + + +class VmmAllocation: + """One native VMM owner record for a single PAYLOAD or COUNTER part.""" + + def __init__( + self, + context: RegionAllocationContext, + part: RegionPartKind, + spec: RegionPartAllocationSpec, + ) -> None: + self._part = _require_part(part) + if not isinstance(spec, RegionPartAllocationSpec): + raise TypeError("spec must be RegionPartAllocationSpec") + if not isinstance(context.target, DeviceAllocationTarget): + raise TypeError("VmmAllocation requires DeviceAllocationTarget") + self._spec = spec + self._device_id = context.target.device_id + self._handle: int | None = None + self._device_addr: int | None = None + self._mapping_bytes: int | None = None + self._shareable_handle: int | None = None + self._mapping_available = False + self._release_once_count = 0 + self._first_cleanup_failure: ProviderCleanupFailure | None = None + self.local_cleanup_details: list[tuple[str, BaseException]] = [] + + @property + def registry_handle(self) -> int | None: + return self._handle + + def materialize(self) -> None: + try: + handle = _region_vmm_begin(self._device_id) + self._handle = handle + export = _region_vmm_allocate_export(handle, self._spec.logical_bytes) + except RegionControlError: + raise + except BaseException as exc: + raise RegionControlError( + RegionControlErrorKind.BACKEND_FAILURE, + str(exc) or "VMM materialize failed", + failed_part=self._part, + failed_operation=RegionOperationKind.MATERIALIZE, + ) from exc + self._device_addr = int(export.device_addr) + self._mapping_bytes = int(export.mapping_bytes) + self._shareable_handle = int(export.shareable_handle) + self._mapping_available = True + + def mapping_bytes(self) -> int: + self._require_mapping("mapping_bytes") + assert self._mapping_bytes is not None + return self._mapping_bytes + + def import_capability(self) -> VmmShareableHandleImport: + self._require_mapping("import_capability") + assert self._shareable_handle is not None + return VmmShareableHandleImport(device_id=self._device_id, shareable_handle=self._shareable_handle) + + def local_base(self) -> int: + self._require_mapping("local_base") + if self._device_addr is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "VMM mapping has no local base", + failed_part=self._part, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) + return self._device_addr + + def zero_bytes(self, offset: int, nbytes: int) -> None: + self._require_mapping("zero_bytes") + if type(offset) is not int or type(nbytes) is not int: + raise TypeError("zero_bytes offset and nbytes must be int") + if offset < 0 or nbytes < 0: + raise ValueError("zero_bytes offset and nbytes must be non-negative") + end = offset + nbytes + if end > self._spec.logical_bytes: + raise ValueError("zero_bytes range exceeds logical_bytes") + assert self._handle is not None + try: + _region_vmm_zero_bytes(self._handle, offset, nbytes) + except RegionControlError: + raise + except BaseException as exc: + raise RegionControlError( + RegionControlErrorKind.BACKEND_FAILURE, + str(exc) or "VMM zero_bytes failed", + failed_part=self._part, + failed_operation=RegionOperationKind.ZERO_BYTES, + ) from exc + + def release_once(self) -> ProviderCleanupFailure | None: + self._release_once_count += 1 + if self._release_once_count > 1: + return self._first_cleanup_failure + if self._handle is None: + return None + try: + _region_vmm_release(self._handle) + except BaseException as exc: + self._record_cleanup_failure("release", exc) + return self._first_cleanup_failure + self._handle = None + self._mapping_available = False + return self._first_cleanup_failure + + def _record_cleanup_failure(self, step: str, exc: BaseException) -> None: + cause = RegionCleanupCause.INTERRUPTED if _interrupt_like(exc) else RegionCleanupCause.BACKEND_ERROR + failure = ProviderCleanupFailure( + part=self._part, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=cause, + ) + self.local_cleanup_details.append((step, exc)) + if self._first_cleanup_failure is None: + self._first_cleanup_failure = failure + + def _require_mapping(self, operation: str) -> None: + if not self._mapping_available or self._handle is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + f"{operation} requires a materialized VMM mapping", + failed_part=self._part, + failed_operation=RegionOperationKind.DESCRIBE + if operation in {"mapping_bytes", "import_capability"} + else RegionOperationKind.LOCAL_VIEW + if operation == "local_base" + else RegionOperationKind.ZERO_BYTES, + ) + + +def _closed_part_dispatcher( + context: RegionAllocationContext, + part: RegionPartKind, + spec: RegionPartAllocationSpec, +) -> RegionPartAllocation: + if spec.planned_backing_kind is not BackendKind.VMM_WINDOW or not isinstance( + context.target, DeviceAllocationTarget + ): + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "admitted allocation combination has no provider backend", + failed_part=part, + failed_operation=RegionOperationKind.NONE, + ) + if context.environment_kind is RegionEnvironmentKind.SIM: + return SimPosixShmAllocation(context, part, spec) + if context.environment_kind is RegionEnvironmentKind.ONBOARD: + return VmmAllocation(context, part, spec) + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "admitted allocation combination has no provider backend", + failed_part=part, + failed_operation=RegionOperationKind.NONE, + ) + + +def _initialize_counter_storage(allocation: RegionPartAllocation, logical_bytes: int) -> None: + allocation.zero_bytes(0, logical_bytes) + + +def _store_lifecycle_error(state: ProviderRegionStoreState) -> RegionControlError: + return RegionControlError( + RegionControlErrorKind.STORE_LIFECYCLE, + f"store is {state.name}", + ) + + +def _invalid_resource_id_error(message: str) -> RegionControlError: + return RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, message) + + +def _interrupt_like(exc: BaseException) -> bool: + return isinstance(exc, (KeyboardInterrupt, SystemExit, GeneratorExit)) + + +def _cleanup_failure_from_exception(part: RegionPartKind, exc: BaseException) -> ProviderCleanupFailure: + cause = RegionCleanupCause.INTERRUPTED if _interrupt_like(exc) else RegionCleanupCause.BACKEND_ERROR + return ProviderCleanupFailure( + part=part, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=cause, + ) + + +def _coerce_part_cleanup_failure(kind: RegionPartKind, returned: object) -> ProviderCleanupFailure: + if isinstance(returned, ProviderCleanupFailure): + if returned.part is kind: + return returned + return ProviderCleanupFailure( + part=kind, + backend_operation=returned.backend_operation, + typed_cause=returned.typed_cause, + ) + return ProviderCleanupFailure( + part=kind, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_ERROR, + ) + + +class _AllocateStage: + def __init__(self) -> None: + self.part = RegionPartKind.INVALID + self.operation = RegionOperationKind.NONE + + +class _ProviderPartRecord: + def __init__( + self, + kind: RegionPartKind, + spec: RegionPartAllocationSpec, + allocation: RegionPartAllocation, + ) -> None: + self.kind = kind + self.spec = spec + self.allocation = allocation + self.state = ProviderPartResourceState.SHELL + self.cleanup_failure: ProviderCleanupFailure | None = None + + +class ProviderRegionResource: + """Store-private logical record for one provider resource ID.""" + + def __init__(self, provider_resource_id: int, spec: RegionAllocationSpec) -> None: + self.provider_resource_id = provider_resource_id + self.spec = spec + self.state = ProviderRegionResourceState.CREATING + self.parts: dict[RegionPartKind, _ProviderPartRecord] = {} + self.export_descriptor: RegionExportDescriptor | None = None + self.local_views: dict[RegionPartKind, RegionPartLocalView] = {} + + +class ProviderRegionStore: + """Provider-agent physical ownership authority for one store incarnation.""" + + def __init__( + self, + context: RegionAllocationContext, + *, + _shell_factory: ShellFactory | None = None, + ) -> None: + if not isinstance(context, RegionAllocationContext): + raise TypeError("context must be RegionAllocationContext") + self._context = context + self._state = ProviderRegionStoreState.OPEN + self._next_provider_resource_id = 1 + self._resources: dict[int, ProviderRegionResource] = {} + self._shell_factory = _closed_part_dispatcher if _shell_factory is None else _shell_factory + + @property + def context(self) -> RegionAllocationContext: + return self._context + + @property + def state(self) -> ProviderRegionStoreState: + return self._state + + def allocate_and_export(self, spec: RegionAllocationSpec) -> RegionAllocationResult: + self._require_open() + if not isinstance(spec, RegionAllocationSpec): + raise TypeError("spec must be RegionAllocationSpec") + resource_id = self._burn_id() + resource = ProviderRegionResource(resource_id, spec) + self._resources[resource_id] = resource + stage = _AllocateStage() + try: + for kind in REGION_PARTS: + stage.part = kind + stage.operation = RegionOperationKind.NONE + allocation = self._shell_factory(self._context, kind, spec.part(kind)) + resource.parts[kind] = _ProviderPartRecord(kind, spec.part(kind), allocation) + for kind in REGION_PARTS: + stage.part = kind + stage.operation = RegionOperationKind.MATERIALIZE + part = resource.parts[kind] + part.state = ProviderPartResourceState.MATERIALIZING + part.allocation.materialize() + stage.part = RegionPartKind.COUNTER + stage.operation = RegionOperationKind.ZERO_BYTES + _initialize_counter_storage(resource.parts[RegionPartKind.COUNTER].allocation, spec.counter.logical_bytes) + descriptor = self._freeze_descriptor(resource, stage) + local_views = self._freeze_local_views(resource, stage) + for part in resource.parts.values(): + part.state = ProviderPartResourceState.READY + resource.export_descriptor = descriptor + resource.local_views = local_views + resource.state = ProviderRegionResourceState.ACTIVE + return RegionAllocationResult(resource_id, descriptor) + except BaseException as primary: + if isinstance(primary, RegionControlError): + if primary.failed_part is not RegionPartKind.INVALID: + stage.part = primary.failed_part + if primary.failed_operation is not RegionOperationKind.NONE: + stage.operation = primary.failed_operation + resource.state = ProviderRegionResourceState.CLEANUP_PENDING + complete = self._release_installed_parts(resource) + if complete: + del self._resources[resource_id] + else: + self._state = ProviderRegionStoreState.CLOSE_FAILED + error = self._allocation_error(primary, resource_id, stage.part, stage.operation, debt=not complete) + raise error from primary + + def describe(self, provider_resource_id: int) -> RegionExportDescriptor: + self._require_open() + resource = self._require_active_resource(provider_resource_id) + if resource.export_descriptor is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "ACTIVE resource is missing a frozen export descriptor", + ) + return resource.export_descriptor + + def local_view(self, provider_resource_id: int, part: RegionPartKind | int) -> RegionPartLocalView: + self._require_open() + kind = _require_part(part) + resource = self._require_active_resource(provider_resource_id) + view = resource.local_views.get(kind) + if view is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "ACTIVE resource is missing a frozen local view", + failed_part=kind, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) + return view + + def release(self, provider_resource_id: int) -> ProviderReleaseResult: + self._require_open() + resource_id = self._require_resource_id(provider_resource_id) + classified = self._classify_absent_id(resource_id) + if classified is not None: + return ProviderReleaseResult(provider_resource_id=resource_id, status=classified) + resource = self._resources[resource_id] + resource.state = ProviderRegionResourceState.CLEANUP_PENDING + resource.export_descriptor = None + resource.local_views = {} + complete = self._release_installed_parts(resource) + if complete: + del self._resources[resource_id] + return ProviderReleaseResult( + provider_resource_id=resource_id, + status=ProviderReleaseStatus.RELEASED, + ) + self._state = ProviderRegionStoreState.CLOSE_FAILED + return self._incomplete_result(resource) + + def sweep(self) -> tuple[ProviderReleaseResult, ...]: + if self._state is ProviderRegionStoreState.CLOSED: + return () + self._state = ProviderRegionStoreState.CLOSING + results: list[ProviderReleaseResult] = [] + for resource_id in sorted(self._resources): + resource = self._resources[resource_id] + if resource.state in ( + ProviderRegionResourceState.ACTIVE, + ProviderRegionResourceState.CREATING, + ): + resource.state = ProviderRegionResourceState.CLEANUP_PENDING + resource.export_descriptor = None + resource.local_views = {} + complete = self._release_installed_parts(resource) + if complete: + del self._resources[resource_id] + results.append( + ProviderReleaseResult( + provider_resource_id=resource_id, + status=ProviderReleaseStatus.RELEASED, + ) + ) + else: + results.append(self._incomplete_result(resource)) + else: + results.append(self._incomplete_result(resource)) + self._state = ProviderRegionStoreState.CLOSE_FAILED if self._resources else ProviderRegionStoreState.CLOSED + return tuple(results) + + def _require_open(self) -> None: + if self._state is not ProviderRegionStoreState.OPEN: + raise _store_lifecycle_error(self._state) + + def _burn_id(self) -> int: + resource_id = self._next_provider_resource_id + if resource_id > _UINT64_MAX: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "provider_resource_id space exhausted", + ) + self._next_provider_resource_id = resource_id + 1 + return resource_id + + def _require_resource_id(self, provider_resource_id: object) -> int: + resource_id = _require_uint64("provider_resource_id", provider_resource_id) + if resource_id == 0: + raise _invalid_resource_id_error("provider_resource_id must be nonzero") + return resource_id + + def _classify_absent_id(self, resource_id: int) -> ProviderReleaseStatus | None: + if resource_id in self._resources: + return None + if resource_id < self._next_provider_resource_id: + return ProviderReleaseStatus.ALREADY_GONE + return ProviderReleaseStatus.UNKNOWN_RESOURCE + + def _require_active_resource(self, provider_resource_id: object) -> ProviderRegionResource: + resource_id = self._require_resource_id(provider_resource_id) + resource = self._resources.get(resource_id) + if resource is None: + if resource_id < self._next_provider_resource_id: + raise _invalid_resource_id_error("provider_resource_id is already gone") + raise _invalid_resource_id_error("provider_resource_id is unknown") + if resource.state is not ProviderRegionResourceState.ACTIVE: + raise RegionControlError( + RegionControlErrorKind.STORE_LIFECYCLE, + f"resource is {resource.state.name}", + ) + return resource + + def _freeze_descriptor(self, resource: ProviderRegionResource, stage: _AllocateStage) -> RegionExportDescriptor: + exports: dict[RegionPartKind, RegionPartExportDescriptor] = {} + for kind in REGION_PARTS: + stage.part = kind + stage.operation = RegionOperationKind.DESCRIBE + part = resource.parts[kind] + spec = part.spec + try: + mapping_bytes = part.allocation.mapping_bytes() + import_capability = part.allocation.import_capability() + except RegionControlError: + raise + except Exception as exc: + raise RegionControlError( + RegionControlErrorKind.BACKEND_FAILURE, + str(exc) or "export descriptor facts failed", + failed_part=kind, + failed_operation=RegionOperationKind.DESCRIBE, + ) from exc + try: + exports[kind] = RegionPartExportDescriptor( + planned_backing_kind=spec.planned_backing_kind, + logical_bytes=spec.logical_bytes, + mapping_bytes=mapping_bytes, + import_capability=import_capability, + ) + except Exception as exc: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + str(exc) or "export descriptor construction failed", + failed_part=kind, + failed_operation=RegionOperationKind.DESCRIBE, + ) from exc + return RegionExportDescriptor( + payload=exports[RegionPartKind.PAYLOAD], + counter=exports[RegionPartKind.COUNTER], + ) + + def _freeze_local_views( + self, resource: ProviderRegionResource, stage: _AllocateStage + ) -> dict[RegionPartKind, RegionPartLocalView]: + views: dict[RegionPartKind, RegionPartLocalView] = {} + for kind in REGION_PARTS: + stage.part = kind + stage.operation = RegionOperationKind.LOCAL_VIEW + part = resource.parts[kind] + try: + local_base = part.allocation.local_base() + except RegionControlError: + raise + except Exception as exc: + raise RegionControlError( + RegionControlErrorKind.BACKEND_FAILURE, + str(exc) or "local view facts failed", + failed_part=kind, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) from exc + try: + views[kind] = RegionPartLocalView( + part=kind, + local_base=local_base, + logical_bytes=part.spec.logical_bytes, + ) + except Exception as exc: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + str(exc) or "local view construction failed", + failed_part=kind, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) from exc + try: + validate_independent_local_views(views[RegionPartKind.PAYLOAD], views[RegionPartKind.COUNTER]) + except ValueError as exc: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + str(exc), + failed_part=RegionPartKind.INVALID, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) from exc + return views + + def _release_installed_parts(self, resource: ProviderRegionResource) -> bool: + for kind in REGION_PARTS: + part = resource.parts.get(kind) + if part is None or part.state is ProviderPartResourceState.RELEASED: + continue + part.state = ProviderPartResourceState.CLEANUP_PENDING + try: + returned = part.allocation.release_once() + except BaseException as exc: + part.cleanup_failure = _cleanup_failure_from_exception(kind, exc) + continue + if returned is None: + part.state = ProviderPartResourceState.RELEASED + part.cleanup_failure = None + else: + part.cleanup_failure = _coerce_part_cleanup_failure(kind, returned) + return all(part.state is ProviderPartResourceState.RELEASED for part in resource.parts.values()) + + def _incomplete_result(self, resource: ProviderRegionResource) -> ProviderReleaseResult: + failures: list[ProviderCleanupFailure] = [] + for kind in REGION_PARTS: + part = resource.parts.get(kind) + if part is None or part.state is ProviderPartResourceState.RELEASED: + continue + if part.cleanup_failure is None: + part.cleanup_failure = ProviderCleanupFailure( + part=kind, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_STATE_MISMATCH, + ) + failures.append(part.cleanup_failure) + return ProviderReleaseResult( + provider_resource_id=resource.provider_resource_id, + status=ProviderReleaseStatus.CLEANUP_INCOMPLETE, + failures=tuple(failures), + ) + + def _allocation_error( + self, + primary: BaseException, + resource_id: int, + part: RegionPartKind, + operation: RegionOperationKind, + *, + debt: bool, + ) -> RegionAllocationError: + kind = RegionControlErrorKind.INTERNAL_INVARIANT + if isinstance(primary, RegionControlError): + if primary.kind in _ALLOCATION_ERROR_KINDS: + kind = primary.kind + elif operation in _BACKEND_FAILURE_OPERATIONS: + kind = RegionControlErrorKind.BACKEND_FAILURE + if kind is RegionControlErrorKind.BACKEND_FAILURE and operation not in _BACKEND_FAILURE_OPERATIONS: + kind = RegionControlErrorKind.INTERNAL_INVARIANT + operation = RegionOperationKind.NONE + return RegionAllocationError( + provisional_resource_id=resource_id, + control_kind=kind, + failed_part=part, + failed_operation=operation, + cleanup_debt_remaining=debt, + message=str(primary), + ) diff --git a/python/simpler/comm_provider_control.py b/python/simpler/comm_provider_control.py new file mode 100644 index 0000000000..b6e34677a7 --- /dev/null +++ b/python/simpler/comm_provider_control.py @@ -0,0 +1,973 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Fixed private provider-region control wire, transport handlers, and one-shot clients. + +This module owns the versioned mailbox codec and transport adapters. It does not +import worker-chip compatibility types, W2 planning, or W5a identity. +""" + +from __future__ import annotations + +import logging +import struct +from enum import IntEnum +from multiprocessing.shared_memory import SharedMemory +from typing import Any, Protocol, TypeVar + +from _task_interface import BackendKind # pyright: ignore[reportMissingImports] + +from .comm_provider import ( + POSIX_SHM_TOKEN_MAX_BYTES, + PosixShmImport, + ProviderCleanupFailure, + ProviderRegionStore, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationError, + RegionAllocationResult, + RegionAllocationSpec, + RegionCleanupCause, + RegionControlError, + RegionControlErrorKind, + RegionExportDescriptor, + RegionOperationKind, + RegionPartAllocationSpec, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, + VmmShareableHandleImport, + validate_independent_local_views, +) + +PROVIDER_REGION_CTRL_MAGIC = 0x50524354 +PROVIDER_REGION_CTRL_ABI_MAJOR = 1 +PROVIDER_REGION_CTRL_ABI_MINOR = 0 +PROVIDER_REGION_CTRL_MAGIC_VERSION = ( + (PROVIDER_REGION_CTRL_MAGIC << 32) | (PROVIDER_REGION_CTRL_ABI_MAJOR << 16) | PROVIDER_REGION_CTRL_ABI_MINOR +) + +ALLOCATE_REQUEST_BYTES = 48 +ALLOCATE_REPLY_BYTES = 232 +RELEASE_REQUEST_BYTES = 24 +RELEASE_REPLY_BYTES = 48 +EXPORT_PART_BYTES = 72 +LOCAL_VIEW_BYTES = 24 +CAPABILITY_PAYLOAD_BYTES = 40 +VMM_CAPABILITY_PAYLOAD_BYTES = 16 +ALLOCATE_PAYLOAD_EXPORT_OFFSET = 40 +ALLOCATE_COUNTER_EXPORT_OFFSET = 112 +ALLOCATE_PAYLOAD_VIEW_OFFSET = 184 +ALLOCATE_COUNTER_VIEW_OFFSET = 208 +RELEASE_PAYLOAD_FAILURE_OFFSET = 32 +RELEASE_COUNTER_FAILURE_OFFSET = 40 +COMMIT_TAG_OFFSET = 12 +FAILURE_MASK_PAYLOAD = 1 +FAILURE_MASK_COUNTER = 2 + +_ALLOCATE_REQUEST = struct.Struct(" None: ... + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: ... + + +class RegionControlProtocolError(RegionControlError): + """Committed-reply protocol failure that must not be retried.""" + + +def _as_memoryview(buf: memoryview | bytes | bytearray, size: int, *, writable: bool) -> memoryview: + view = buf if isinstance(buf, memoryview) else memoryview(buf) + # POSIX shm mappings may be larger than the requested create size (page + # rounding). The private wire frame is still exactly `size` bytes and lives + # in the prefix; trailing mapping bytes are not part of the protocol. + if view.nbytes < size: + raise RegionControlError( + RegionControlErrorKind.BAD_MESSAGE_SIZE, + f"buffer must hold a {size}-byte wire frame", + ) + if view.nbytes > size: + view = memoryview(view).cast("B")[:size] + if writable and view.readonly: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "reply buffer must be writable", + ) + return view + + +def _require_zero_span(view: memoryview, offset: int, length: int) -> None: + if any(view[offset : offset + length]): + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "reserved bytes must be zero") + + +def _require_magic(magic: int) -> None: + if magic != PROVIDER_REGION_CTRL_MAGIC_VERSION: + raise RegionControlError(RegionControlErrorKind.BAD_MAGIC_VERSION, "unsupported provider control version") + + +_IntEnumT = TypeVar("_IntEnumT", bound=IntEnum) + + +def _require_shm_buf(shm: SharedMemory) -> memoryview: + buf = shm.buf + if buf is None: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "shared memory mapping is missing", + ) + return buf + + +def _require_enum(enum_cls: type[_IntEnumT], value: int, *, allow_zero: bool = False) -> _IntEnumT: + try: + typed = enum_cls(value) + except ValueError as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_ENUM_VALUE, + f"unknown {enum_cls.__name__} {value}", + ) from exc + if not allow_zero and int(typed) == 0: + raise RegionControlError(RegionControlErrorKind.INVALID_ENUM_VALUE, f"{enum_cls.__name__} must be nonzero") + return typed + + +def _require_backend_kind(value: int) -> BackendKind: + kind = _BACKEND_KINDS.get(int(value)) + if kind is None: + raise RegionControlError( + RegionControlErrorKind.INVALID_ENUM_VALUE, + f"unknown planned backing kind {value}", + ) + return kind + + +def _publish_tag(view: memoryview, tag: int) -> None: + _COMMIT_TAG.pack_into(view, COMMIT_TAG_OFFSET, int(tag)) + + +def _acquire_tag(view: memoryview) -> int: + return int(_COMMIT_TAG.unpack_from(view, COMMIT_TAG_OFFSET)[0]) + + +def encode_allocate_request(buf: memoryview | bytearray, spec: RegionAllocationSpec) -> None: + view = _as_memoryview(buf, ALLOCATE_REQUEST_BYTES, writable=True) + view[:] = b"\x00" * ALLOCATE_REQUEST_BYTES + _ALLOCATE_REQUEST.pack_into( + view, + 0, + PROVIDER_REGION_CTRL_MAGIC_VERSION, + ALLOCATE_REQUEST_BYTES, + int(spec.payload.planned_backing_kind), + int(spec.counter.planned_backing_kind), + 0, + int(spec.payload.logical_bytes), + int(spec.counter.logical_bytes), + 0, + ) + + +def decode_allocate_request(buf: memoryview | bytes | bytearray) -> RegionAllocationSpec: + view = _as_memoryview(buf, ALLOCATE_REQUEST_BYTES, writable=False) + magic, request_bytes, payload_kind, counter_kind, reserved20, payload_bytes, counter_bytes, reserved40 = ( + _ALLOCATE_REQUEST.unpack_from(view, 0) + ) + _require_magic(int(magic)) + if int(request_bytes) != ALLOCATE_REQUEST_BYTES: + raise RegionControlError(RegionControlErrorKind.BAD_MESSAGE_SIZE, "allocate request_bytes must be 48") + if int(reserved20) != 0 or int(reserved40) != 0: + raise RegionControlError( + RegionControlErrorKind.RESERVED_NONZERO, "allocate request reserved fields must be zero" + ) + try: + return RegionAllocationSpec( + payload=RegionPartAllocationSpec( + planned_backing_kind=_require_backend_kind(int(payload_kind)), + logical_bytes=int(payload_bytes), + ), + counter=RegionPartAllocationSpec( + planned_backing_kind=_require_backend_kind(int(counter_kind)), + logical_bytes=int(counter_bytes), + ), + ) + except (TypeError, ValueError) as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + str(exc) or "allocate request field is invalid", + ) from exc + + +def _capability_wire_kind(capability: VmmShareableHandleImport | PosixShmImport) -> ImportCapabilityWireKind: + if isinstance(capability, VmmShareableHandleImport): + return ImportCapabilityWireKind.VMM_SHAREABLE_HANDLE + return ImportCapabilityWireKind.POSIX_SHM + + +def _encode_capability_payload(capability: VmmShareableHandleImport | PosixShmImport) -> tuple[int, bytes]: + payload = bytearray(CAPABILITY_PAYLOAD_BYTES) + if isinstance(capability, VmmShareableHandleImport): + _VMM_CAPABILITY.pack_into(payload, 0, int(capability.device_id), 0, int(capability.shareable_handle)) + return VMM_CAPABILITY_PAYLOAD_BYTES, bytes(payload) + token = capability.shm_name.encode("ascii") + payload[: len(token)] = token + return len(token), bytes(payload) + + +def _decode_capability_payload( + kind: ImportCapabilityWireKind, payload_bytes: int, payload: bytes +) -> VmmShareableHandleImport | PosixShmImport: + if len(payload) != CAPABILITY_PAYLOAD_BYTES: + raise RegionControlError(RegionControlErrorKind.BAD_MESSAGE_SIZE, "capability payload must be 40 bytes") + if payload_bytes < 0 or payload_bytes > CAPABILITY_PAYLOAD_BYTES: + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, "capability payload bytes out of range") + if any(payload[payload_bytes:]): + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "inactive capability bytes must be zero") + if kind is ImportCapabilityWireKind.VMM_SHAREABLE_HANDLE: + if payload_bytes != VMM_CAPABILITY_PAYLOAD_BYTES: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "VMM capability payload_bytes must be 16", + ) + device_id, reserved, shareable = _VMM_CAPABILITY.unpack_from(payload, 0) + if int(reserved) != 0: + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "VMM capability reserved must be zero") + try: + return VmmShareableHandleImport(device_id=int(device_id), shareable_handle=int(shareable)) + except (TypeError, ValueError) as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + str(exc) or "VMM capability fields are invalid", + ) from exc + if kind is ImportCapabilityWireKind.POSIX_SHM: + if payload_bytes < 1 or payload_bytes > POSIX_SHM_TOKEN_MAX_BYTES: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "POSIX shm token length must be in 1..32 bytes", + ) + token = bytes(payload[:payload_bytes]) + if b"\x00" in token: + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, "POSIX shm token must not contain NUL") + try: + return PosixShmImport(shm_name=token.decode("ascii")) + except (TypeError, ValueError, UnicodeDecodeError) as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + str(exc) or "POSIX shm token is invalid", + ) from exc + raise RegionControlError(RegionControlErrorKind.INVALID_ENUM_VALUE, "import capability kind is invalid") + + +def _encode_export_part(view: memoryview, offset: int, part: RegionPartExportDescriptor) -> None: + payload_bytes, payload = _encode_capability_payload(part.import_capability) + _EXPORT_PART.pack_into( + view, + offset, + int(part.planned_backing_kind), + int(_capability_wire_kind(part.import_capability)), + int(part.logical_bytes), + int(part.mapping_bytes), + int(payload_bytes), + 0, + payload, + ) + + +def _decode_export_part(view: memoryview, offset: int) -> RegionPartExportDescriptor: + backing, cap_kind, logical_bytes, mapping_bytes, payload_bytes, reserved, payload = _EXPORT_PART.unpack_from( + view, offset + ) + if int(reserved) != 0: + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "export-part reserved must be zero") + capability = _decode_capability_payload( + _require_enum(ImportCapabilityWireKind, int(cap_kind)), # type: ignore[arg-type] + int(payload_bytes), + bytes(payload), + ) + try: + return RegionPartExportDescriptor( + planned_backing_kind=_require_backend_kind(int(backing)), + logical_bytes=int(logical_bytes), + mapping_bytes=int(mapping_bytes), + import_capability=capability, + ) + except (TypeError, ValueError) as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + str(exc) or "export-part fields are invalid", + ) from exc + + +def _encode_local_view(view: memoryview, offset: int, local_view: RegionPartLocalView) -> None: + _LOCAL_VIEW.pack_into( + view, + offset, + int(local_view.part), + 0, + int(local_view.local_base), + int(local_view.logical_bytes), + ) + + +def _decode_local_view(view: memoryview, offset: int, expected: RegionPartKind) -> RegionPartLocalView: + part, reserved, local_base, logical_bytes = _LOCAL_VIEW.unpack_from(view, offset) + if int(reserved) != 0: + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "local-view reserved must be zero") + typed_part = _require_enum(RegionPartKind, int(part)) + if typed_part is not expected: + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, "local-view part does not match slot") + try: + return RegionPartLocalView(part=expected, local_base=int(local_base), logical_bytes=int(logical_bytes)) + except (TypeError, ValueError) as exc: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + str(exc) or "local-view fields are invalid", + ) from exc + + +def _zero_reply(view: memoryview) -> None: + view[:] = b"\x00" * view.nbytes + + +def _discard_control_shm(shm: Any) -> None: + name = getattr(shm, "name", "") + logger = logging.getLogger("simpler") + try: + shm.close() + except BaseException as exc: # noqa: BLE001 + logger.warning( + "provider control shm cleanup failed: name=%s operation=close error=%s", + name, + exc, + ) + try: + shm.unlink() + except FileNotFoundError: + return + except BaseException as exc: # noqa: BLE001 + logger.warning( + "provider control shm cleanup failed: name=%s operation=unlink error=%s", + name, + exc, + ) + + +def _pack_allocate_header( + view: memoryview, + *, + resource_id: int, + error_kind: int, + failed_part: int, + failed_operation: int, + debt: int, +) -> None: + _ALLOCATE_REPLY_HEADER.pack_into( + view, + 0, + PROVIDER_REGION_CTRL_MAGIC_VERSION, + ALLOCATE_REPLY_BYTES, + 0, + int(resource_id), + int(error_kind), + int(failed_part), + int(failed_operation), + int(debt), + ) + + +def peek_allocate_reply_resource_id(buf: memoryview | bytes | bytearray) -> int: + """Return the SUCCESS resource id, or 0 if the reply is not a committed SUCCESS.""" + try: + view = _as_memoryview(buf, ALLOCATE_REPLY_BYTES, writable=False) + except RegionControlError: + return 0 + tag = _acquire_tag(view) + if int(tag) != int(AllocateReplyTag.SUCCESS): + return 0 + return int(_ALLOCATE_REPLY_HEADER.unpack_from(view, 0)[3]) + + +def encode_allocate_success_reply( + buf: memoryview | bytearray, + result: RegionAllocationResult, + payload_view: RegionPartLocalView, + counter_view: RegionPartLocalView, +) -> None: + view = _as_memoryview(buf, ALLOCATE_REPLY_BYTES, writable=True) + validate_independent_local_views(payload_view, counter_view) + _zero_reply(view) + _pack_allocate_header( + view, resource_id=result.provider_resource_id, error_kind=0, failed_part=0, failed_operation=0, debt=0 + ) + _encode_export_part(view, ALLOCATE_PAYLOAD_EXPORT_OFFSET, result.export_descriptor.payload) + _encode_export_part(view, ALLOCATE_COUNTER_EXPORT_OFFSET, result.export_descriptor.counter) + _encode_local_view(view, ALLOCATE_PAYLOAD_VIEW_OFFSET, payload_view) + _encode_local_view(view, ALLOCATE_COUNTER_VIEW_OFFSET, counter_view) + _publish_tag(view, AllocateReplyTag.SUCCESS) + + +def encode_allocate_request_error_reply(buf: memoryview | bytearray, kind: RegionControlErrorKind) -> None: + view = _as_memoryview(buf, ALLOCATE_REPLY_BYTES, writable=True) + if kind not in _REQUEST_ERROR_KINDS: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "REQUEST_ERROR kind must be a wire syntax code", + ) + _zero_reply(view) + _pack_allocate_header(view, resource_id=0, error_kind=int(kind), failed_part=0, failed_operation=0, debt=0) + _publish_tag(view, AllocateReplyTag.REQUEST_ERROR) + + +def encode_allocate_allocation_error_reply(buf: memoryview | bytearray, error: RegionAllocationError) -> None: + view = _as_memoryview(buf, ALLOCATE_REPLY_BYTES, writable=True) + _zero_reply(view) + _pack_allocate_header( + view, + resource_id=error.provisional_resource_id, + error_kind=int(error.control_kind), + failed_part=int(error.failed_part), + failed_operation=int(error.failed_operation), + debt=1 if error.cleanup_debt_remaining else 0, + ) + _publish_tag(view, AllocateReplyTag.ALLOCATION_ERROR) + + +def decode_allocate_reply( + buf: memoryview | bytes | bytearray, +) -> tuple[ + AllocateReplyTag, + RegionAllocationResult | None, + RegionPartLocalView | None, + RegionPartLocalView | None, + RegionControlError | RegionAllocationError | None, +]: + view = _as_memoryview(buf, ALLOCATE_REPLY_BYTES, writable=False) + tag = _require_enum(AllocateReplyTag, _acquire_tag(view), allow_zero=True) + if tag is AllocateReplyTag.EMPTY: + raise RegionControlProtocolError( + RegionControlErrorKind.INVALID_FIELD_VALUE, "allocate reply is missing commit tag" + ) + magic, reply_bytes, _tag, resource_id, error_kind, failed_part, failed_operation, debt = ( + _ALLOCATE_REPLY_HEADER.unpack_from(view, 0) + ) + _require_magic(int(magic)) + if int(reply_bytes) != ALLOCATE_REPLY_BYTES: + raise RegionControlError(RegionControlErrorKind.BAD_MESSAGE_SIZE, "allocate reply_bytes must be 232") + if tag is AllocateReplyTag.SUCCESS: + if ( + int(resource_id) == 0 + or int(error_kind) != 0 + or int(failed_part) != 0 + or int(failed_operation) != 0 + or int(debt) != 0 + ): + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "SUCCESS requires a resource id and zero error fields", + ) + descriptor = RegionExportDescriptor( + payload=_decode_export_part(view, ALLOCATE_PAYLOAD_EXPORT_OFFSET), + counter=_decode_export_part(view, ALLOCATE_COUNTER_EXPORT_OFFSET), + ) + payload_view = _decode_local_view(view, ALLOCATE_PAYLOAD_VIEW_OFFSET, RegionPartKind.PAYLOAD) + counter_view = _decode_local_view(view, ALLOCATE_COUNTER_VIEW_OFFSET, RegionPartKind.COUNTER) + validate_independent_local_views(payload_view, counter_view) + result = RegionAllocationResult(provider_resource_id=int(resource_id), export_descriptor=descriptor) + return tag, result, payload_view, counter_view, None + _require_zero_span(view, ALLOCATE_PAYLOAD_EXPORT_OFFSET, 2 * EXPORT_PART_BYTES + 2 * LOCAL_VIEW_BYTES) + if tag is AllocateReplyTag.REQUEST_ERROR: + if int(resource_id) != 0 or int(failed_part) != 0 or int(failed_operation) != 0 or int(debt) != 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "REQUEST_ERROR requires resource id zero and inactive error slots", + ) + kind = _require_enum(RegionControlErrorKind, int(error_kind)) + if kind not in _REQUEST_ERROR_KINDS: + raise RegionControlError( + RegionControlErrorKind.INVALID_ENUM_VALUE, + "REQUEST_ERROR kind must be a wire syntax code", + ) + return tag, None, None, None, RegionControlError(kind) # type: ignore[arg-type] + if tag is AllocateReplyTag.ALLOCATION_ERROR: + if int(resource_id) == 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "ALLOCATION_ERROR requires a provisional resource id", + ) + if int(debt) not in (0, 1): + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, "cleanup debt remaining must be 0 or 1" + ) + kind = _require_enum(RegionControlErrorKind, int(error_kind)) + if kind not in _ALLOCATION_ERROR_KINDS: + raise RegionControlError( + RegionControlErrorKind.INVALID_ENUM_VALUE, + "ALLOCATION_ERROR kind must be BACKEND_FAILURE or INTERNAL_INVARIANT", + ) + error = RegionAllocationError( + provisional_resource_id=int(resource_id), + control_kind=kind, # type: ignore[arg-type] + failed_part=int(failed_part), + failed_operation=int(failed_operation), + cleanup_debt_remaining=bool(debt), + ) + return tag, None, None, None, error + raise RegionControlError(RegionControlErrorKind.INVALID_ENUM_VALUE, f"unknown allocate reply tag {int(tag)}") + + +def encode_release_request(buf: memoryview | bytearray, provider_resource_id: int) -> None: + view = _as_memoryview(buf, RELEASE_REQUEST_BYTES, writable=True) + view[:] = b"\x00" * RELEASE_REQUEST_BYTES + _RELEASE_REQUEST.pack_into( + view, + 0, + PROVIDER_REGION_CTRL_MAGIC_VERSION, + RELEASE_REQUEST_BYTES, + 0, + int(provider_resource_id), + ) + + +def decode_release_request(buf: memoryview | bytes | bytearray) -> int: + view = _as_memoryview(buf, RELEASE_REQUEST_BYTES, writable=False) + magic, request_bytes, reserved, resource_id = _RELEASE_REQUEST.unpack_from(view, 0) + _require_magic(int(magic)) + if int(request_bytes) != RELEASE_REQUEST_BYTES: + raise RegionControlError(RegionControlErrorKind.BAD_MESSAGE_SIZE, "release request_bytes must be 24") + if int(reserved) != 0: + raise RegionControlError(RegionControlErrorKind.RESERVED_NONZERO, "release request reserved must be zero") + if int(resource_id) == 0: + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, "provider_resource_id must be nonzero") + return int(resource_id) + + +def _pack_release_reply( + view: memoryview, + *, + resource_id: int, + mask: int, + error_kind: int, + payload_op: int, + payload_cause: int, + counter_op: int, + counter_cause: int, +) -> None: + _RELEASE_REPLY.pack_into( + view, + 0, + PROVIDER_REGION_CTRL_MAGIC_VERSION, + RELEASE_REPLY_BYTES, + 0, + int(resource_id), + int(mask), + int(error_kind), + int(payload_op), + int(payload_cause), + int(counter_op), + int(counter_cause), + ) + + +def encode_release_result_reply(buf: memoryview | bytearray, result: ProviderReleaseResult) -> None: + view = _as_memoryview(buf, RELEASE_REPLY_BYTES, writable=True) + tag = { + ProviderReleaseStatus.RELEASED: ReleaseReplyTag.RELEASED, + ProviderReleaseStatus.ALREADY_GONE: ReleaseReplyTag.ALREADY_GONE, + ProviderReleaseStatus.UNKNOWN_RESOURCE: ReleaseReplyTag.UNKNOWN_RESOURCE, + ProviderReleaseStatus.CLEANUP_INCOMPLETE: ReleaseReplyTag.CLEANUP_INCOMPLETE, + }[result.status] + mask = 0 + payload_op = payload_cause = counter_op = counter_cause = 0 + if result.status is ProviderReleaseStatus.CLEANUP_INCOMPLETE: + for failure in result.failures: + if failure.part is RegionPartKind.PAYLOAD: + mask |= FAILURE_MASK_PAYLOAD + payload_op = int(failure.backend_operation) + payload_cause = int(failure.typed_cause) + else: + mask |= FAILURE_MASK_COUNTER + counter_op = int(failure.backend_operation) + counter_cause = int(failure.typed_cause) + if mask == 0: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "CLEANUP_INCOMPLETE requires a nonzero failure mask", + ) + _zero_reply(view) + _pack_release_reply( + view, + resource_id=result.provider_resource_id, + mask=mask, + error_kind=0, + payload_op=payload_op, + payload_cause=payload_cause, + counter_op=counter_op, + counter_cause=counter_cause, + ) + _publish_tag(view, tag) + + +def encode_release_error_reply( + buf: memoryview | bytearray, kind: RegionControlErrorKind, provider_resource_id: int = 0 +) -> None: + view = _as_memoryview(buf, RELEASE_REPLY_BYTES, writable=True) + if kind not in _RELEASE_ERROR_KINDS: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "RELEASE_ERROR kind must be INVALID_FIELD_VALUE, STORE_LIFECYCLE, or INTERNAL_INVARIANT", + ) + _zero_reply(view) + _pack_release_reply( + view, + resource_id=int(provider_resource_id), + mask=0, + error_kind=int(kind), + payload_op=0, + payload_cause=0, + counter_op=0, + counter_cause=0, + ) + _publish_tag(view, ReleaseReplyTag.RELEASE_ERROR) + + +def _failure_from_entry(part: RegionPartKind, operation: int, cause: int) -> ProviderCleanupFailure: + typed_op = _require_enum(RegionOperationKind, operation) + typed_cause = _require_enum(RegionCleanupCause, cause) + return ProviderCleanupFailure(part=part, backend_operation=typed_op, typed_cause=typed_cause) # type: ignore[arg-type] + + +def _require_zero_release_side_fields( + mask: int, + payload_op: int, + payload_cause: int, + counter_op: int, + counter_cause: int, + message: str, +) -> None: + if mask != 0 or payload_op != 0 or payload_cause != 0 or counter_op != 0 or counter_cause != 0: + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, message) + + +def _decode_cleanup_incomplete_failures( + mask: int, + payload_op: int, + payload_cause: int, + counter_op: int, + counter_cause: int, +) -> list[ProviderCleanupFailure]: + if mask & ~(FAILURE_MASK_PAYLOAD | FAILURE_MASK_COUNTER): + raise RegionControlError(RegionControlErrorKind.INVALID_FIELD_VALUE, "failure mask has unknown bits") + if mask == 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "CLEANUP_INCOMPLETE requires a nonzero failure mask", + ) + failures: list[ProviderCleanupFailure] = [] + if mask & FAILURE_MASK_PAYLOAD: + failures.append(_failure_from_entry(RegionPartKind.PAYLOAD, payload_op, payload_cause)) + elif payload_op != 0 or payload_cause != 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "unselected PAYLOAD failure entry must be zero", + ) + if mask & FAILURE_MASK_COUNTER: + failures.append(_failure_from_entry(RegionPartKind.COUNTER, counter_op, counter_cause)) + elif counter_op != 0 or counter_cause != 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "unselected COUNTER failure entry must be zero", + ) + return failures + + +def decode_release_reply(buf: memoryview | bytes | bytearray) -> ProviderReleaseResult | RegionControlError: + view = _as_memoryview(buf, RELEASE_REPLY_BYTES, writable=False) + tag = _require_enum(ReleaseReplyTag, _acquire_tag(view), allow_zero=True) + if tag is ReleaseReplyTag.EMPTY: + raise RegionControlProtocolError( + RegionControlErrorKind.INVALID_FIELD_VALUE, "release reply is missing commit tag" + ) + ( + magic, + reply_bytes, + _tag, + resource_id, + mask, + error_kind, + payload_op, + payload_cause, + counter_op, + counter_cause, + ) = _RELEASE_REPLY.unpack_from(view, 0) + _require_magic(int(magic)) + if int(reply_bytes) != RELEASE_REPLY_BYTES: + raise RegionControlError(RegionControlErrorKind.BAD_MESSAGE_SIZE, "release reply_bytes must be 48") + if tag is ReleaseReplyTag.RELEASE_ERROR: + _require_zero_release_side_fields( + int(mask), + int(payload_op), + int(payload_cause), + int(counter_op), + int(counter_cause), + "RELEASE_ERROR requires a zero failure mask and entries", + ) + kind = _require_enum(RegionControlErrorKind, int(error_kind)) + if kind not in _RELEASE_ERROR_KINDS: + raise RegionControlError( + RegionControlErrorKind.INVALID_ENUM_VALUE, + "RELEASE_ERROR kind is not allowed", + ) + return RegionControlError(kind) + if int(resource_id) == 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, "release reply resource id must be nonzero" + ) + if int(error_kind) != 0: + raise RegionControlError( + RegionControlErrorKind.INVALID_FIELD_VALUE, + "lifecycle release replies require a zero control error", + ) + status = { + ReleaseReplyTag.RELEASED: ProviderReleaseStatus.RELEASED, + ReleaseReplyTag.ALREADY_GONE: ProviderReleaseStatus.ALREADY_GONE, + ReleaseReplyTag.UNKNOWN_RESOURCE: ProviderReleaseStatus.UNKNOWN_RESOURCE, + ReleaseReplyTag.CLEANUP_INCOMPLETE: ProviderReleaseStatus.CLEANUP_INCOMPLETE, + }.get(tag) + if status is None: + raise RegionControlError(RegionControlErrorKind.INVALID_ENUM_VALUE, f"unknown release reply tag {int(tag)}") + failures: list[ProviderCleanupFailure] = [] + if status is ProviderReleaseStatus.CLEANUP_INCOMPLETE: + failures = _decode_cleanup_incomplete_failures( + int(mask), + int(payload_op), + int(payload_cause), + int(counter_op), + int(counter_cause), + ) + else: + _require_zero_release_side_fields( + int(mask), + int(payload_op), + int(payload_cause), + int(counter_op), + int(counter_cause), + "non-incomplete release replies require zero failure mask and entries", + ) + return ProviderReleaseResult( + provider_resource_id=int(resource_id), + status=status, + failures=tuple(failures), + ) + + +def handle_ctrl_region_allocate(req_buf: memoryview, reply_buf: memoryview, store: ProviderRegionStore) -> None: + reply = _as_memoryview(reply_buf, ALLOCATE_REPLY_BYTES, writable=True) + _zero_reply(reply) + try: + spec = decode_allocate_request(req_buf) + except RegionControlError as exc: + if exc.kind in _REQUEST_ERROR_KINDS: + encode_allocate_request_error_reply(reply, exc.kind) + return + raise + try: + result = store.allocate_and_export(spec) + except RegionAllocationError as exc: + encode_allocate_allocation_error_reply(reply, exc) + return + try: + payload_view = store.local_view(result.provider_resource_id, RegionPartKind.PAYLOAD) + counter_view = store.local_view(result.provider_resource_id, RegionPartKind.COUNTER) + encode_allocate_success_reply(reply, result, payload_view, counter_view) + finally: + if _acquire_tag(reply) != int(AllocateReplyTag.SUCCESS): + store.release(result.provider_resource_id) + + +def handle_ctrl_region_release(req_buf: memoryview, reply_buf: memoryview, store: ProviderRegionStore) -> None: + reply = _as_memoryview(reply_buf, RELEASE_REPLY_BYTES, writable=True) + _zero_reply(reply) + try: + resource_id = decode_release_request(req_buf) + except RegionControlError as exc: + kind = exc.kind if exc.kind in _RELEASE_ERROR_KINDS else RegionControlErrorKind.INVALID_FIELD_VALUE + encode_release_error_reply(reply, kind) + return + try: + result = store.release(resource_id) + except RegionControlError as exc: + if exc.kind in _RELEASE_ERROR_KINDS: + encode_release_error_reply(reply, exc.kind, provider_resource_id=resource_id) + return + raise + encode_release_result_reply(reply, result) + + +class ProviderAllocateClient: + """Mailbox client that publishes one allocate request and waits for a committed reply.""" + + def __init__(self, mailbox: RegionControlMailbox, worker_id: int) -> None: + self._mailbox = mailbox + self._worker_id = int(worker_id) + self.dispatch_started = False + self.committed_resource_id = 0 + + def allocate( + self, spec: RegionAllocationSpec + ) -> tuple[RegionAllocationResult, RegionPartLocalView, RegionPartLocalView]: + self.dispatch_started = False + self.committed_resource_id = 0 + req_shm = SharedMemory(create=True, size=ALLOCATE_REQUEST_BYTES) + reply_shm = SharedMemory(create=True, size=ALLOCATE_REPLY_BYTES) + req_buf = memoryview(_require_shm_buf(req_shm)) + reply_buf = memoryview(_require_shm_buf(reply_shm)) + try: + encode_allocate_request(req_buf, spec) + self.dispatch_started = True + try: + self._mailbox.control_region_allocate(self._worker_id, req_shm.name, reply_shm.name) + except BaseException: + self._observe_committed_success(reply_buf) + raise + tag, result, payload_view, counter_view, error = decode_allocate_reply(reply_buf) + if tag is AllocateReplyTag.SUCCESS: + assert result is not None and payload_view is not None and counter_view is not None + self.committed_resource_id = int(result.provider_resource_id) + return result, payload_view, counter_view + if error is not None: + raise error + raise RegionControlProtocolError( + RegionControlErrorKind.INVALID_FIELD_VALUE, "allocate reply is uncommitted" + ) + finally: + del req_buf + del reply_buf + _discard_control_shm(req_shm) + _discard_control_shm(reply_shm) + + def _observe_committed_success(self, reply_buf: memoryview) -> None: + try: + view = _as_memoryview(reply_buf, ALLOCATE_REPLY_BYTES, writable=False) + if _acquire_tag(view) != int(AllocateReplyTag.SUCCESS): + return + tag, result, _payload, _counter, _error = decode_allocate_reply(reply_buf) + except Exception: + return + if tag is AllocateReplyTag.SUCCESS and result is not None: + self.committed_resource_id = int(result.provider_resource_id) + + +class ProviderReleaseClient: + """One-shot mailbox client for provider release. Transport and decode failures are terminal.""" + + def __init__(self, mailbox: RegionControlMailbox, worker_id: int) -> None: + self._mailbox = mailbox + self._worker_id = int(worker_id) + self._issued = False + self._terminal: BaseException | ProviderReleaseResult | None = None + + def release(self, provider_resource_id: int) -> ProviderReleaseResult: + if self._issued: + if isinstance(self._terminal, ProviderReleaseResult): + return self._terminal + assert self._terminal is not None + raise self._terminal + self._issued = True + try: + result = self._issue(provider_resource_id) + except BaseException as exc: + self._terminal = exc + raise + self._terminal = result + return result + + def _issue(self, provider_resource_id: int) -> ProviderReleaseResult: + req_shm = SharedMemory(create=True, size=RELEASE_REQUEST_BYTES) + reply_shm = SharedMemory(create=True, size=RELEASE_REPLY_BYTES) + req_buf = memoryview(_require_shm_buf(req_shm)) + reply_buf = memoryview(_require_shm_buf(reply_shm)) + try: + encode_release_request(req_buf, provider_resource_id) + try: + self._mailbox.control_region_release(self._worker_id, req_shm.name, reply_shm.name) + except BaseException as exc: + raise RegionControlProtocolError( + RegionControlErrorKind.INTERNAL_INVARIANT, + str(exc) or "release transport failed", + ) from exc + decoded = decode_release_reply(reply_buf) + if isinstance(decoded, RegionControlError): + raise RegionControlProtocolError(decoded.kind, decoded.message or decoded.kind.name) + return decoded + finally: + del req_buf + del reply_buf + _discard_control_shm(req_shm) + _discard_control_shm(reply_shm) diff --git a/python/simpler/comm_region.py b/python/simpler/comm_region.py index 6ac6b954c4..0471a9d876 100644 --- a/python/simpler/comm_region.py +++ b/python/simpler/comm_region.py @@ -22,6 +22,7 @@ _region_counter_notify, _region_counter_test, _region_counter_wait, + _worker_host_mapped_region_close, ) from .buffer import AddressSpace, Buffer @@ -42,6 +43,19 @@ UnsupportedRegionPlan, parse_endpoint_path, ) +from .comm_provider import ( + PosixShmImport, + ProviderReleaseStatus, + RegionAllocationError, + RegionAllocationResult, + RegionAllocationSpec, + RegionPartAllocationSpec, + RegionPartKind, + RegionPartLocalView, + VmmShareableHandleImport, + validate_independent_local_views, +) +from .comm_provider_control import ProviderAllocateClient, ProviderReleaseClient _GENERATION_COUNTER = itertools.count(1) _MAX_SIGNED_CHRONO_TIMEOUT_NS = 2**63 - 1 @@ -127,23 +141,23 @@ def __exit__(self, *_: Any) -> None: class HostVmmCopyAccess: def __init__(self, handle: Any) -> None: - self._handle = int(handle) + self._handle = handle @classmethod def from_mapping(cls, mapping: Any) -> HostVmmCopyAccess: - return cls(getattr(mapping, "handle")) + return cls(getattr(mapping, "handle", mapping)) @property - def handle(self) -> int: + def handle(self) -> Any: return self._handle def write_bytes(self, span: RegionPartSpan, offset: int, host_ptr: int, nbytes: int) -> None: span.validate_range(offset, nbytes, "region byte write") - _host_vmm_copy_to(self._handle, int(span.offset) + int(offset), int(host_ptr), int(nbytes)) + _host_vmm_copy_to(int(self._handle), int(span.offset) + int(offset), int(host_ptr), int(nbytes)) def read_bytes(self, span: RegionPartSpan, offset: int, host_ptr: int, nbytes: int) -> None: span.validate_range(offset, nbytes, "region byte read") - _host_vmm_copy_from(self._handle, int(span.offset) + int(offset), int(host_ptr), int(nbytes)) + _host_vmm_copy_from(int(self._handle), int(span.offset) + int(offset), int(host_ptr), int(nbytes)) class PayloadPart: @@ -246,26 +260,12 @@ def wait(self, cmp_value: int, cmp: WaitCmp, timeout: float) -> int: class RegionInstanceState(str, Enum): - PLANNED = "PLANNED" - OWNER_CREATED = "OWNER_CREATED" - CONSUMER_ATTACHED = "CONSUMER_ATTACHED" LIVE = "LIVE" - ROLLING_BACK = "ROLLING_BACK" - ROLLED_BACK = "ROLLED_BACK" - ROLLBACK_FAILED = "ROLLBACK_FAILED" CLOSING = "CLOSING" CLOSED = "CLOSED" CLOSE_FAILED = "CLOSE_FAILED" -_TERMINAL_FAILURE_STATES = frozenset( - ( - RegionInstanceState.ROLLBACK_FAILED, - RegionInstanceState.CLOSE_FAILED, - ) -) - - class RefusalReason(str, Enum): UNSUPPORTED_PLAN = "UNSUPPORTED_PLAN" NEEDS_DELEGATION = "NEEDS_DELEGATION" @@ -302,6 +302,85 @@ class SingleOwnerRegionShape: worker_id: int +class RegionInstanceRegistry: + """Strong reachability for consumer instances. Not an ID-addressed data path.""" + + def __init__(self) -> None: + self._instances: dict[int, RegionInstance] = {} + self._run_scopes: dict[int, Any] = {} + + def track(self, instance: RegionInstance, run_scope: Any) -> None: + key = id(instance) + if key in self._instances: + raise MaterializationError("region instance is already tracked") + self._instances[key] = instance + self._run_scopes[key] = run_scope + + def close(self, instance: RegionInstance) -> None: + try: + instance._close_owned(poison_on_error=True) + finally: + self._settle(instance) + + def cleanup_run(self, run_scope: Any) -> None: + errors: list[BaseException] = [] + for instance in self._iter_run(run_scope): + if instance._close_attempted or instance._state in ( + RegionInstanceState.CLOSED, + RegionInstanceState.CLOSE_FAILED, + ): + self._settle(instance) + continue + try: + self.close(instance) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if errors: + raise errors[0] + + def sweep(self) -> None: + errors: list[BaseException] = [] + for instance in tuple(self._instances.values()): + if not instance._close_attempted and instance._state not in ( + RegionInstanceState.CLOSED, + RegionInstanceState.CLOSE_FAILED, + ): + try: + instance._close_owned(poison_on_error=True) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + self._retire(id(instance)) + if errors: + raise errors[0] + + def record_data_plane_failure(self, run_scope: Any, resource_id: int, error: BaseException) -> None: + resource_id = int(resource_id) + matches = [ + instance for instance in self._iter_run(run_scope) if int(instance.provider_resource_id) == resource_id + ] + if not matches: + raise MaterializationError(f"no region instance for resource {resource_id}") + if len(matches) > 1: + raise MaterializationError(f"duplicate region instances for resource {resource_id}") + instance = matches[0] + if instance._data_plane_error is None: + instance._data_plane_error = error + + def _iter_run(self, run_scope: Any) -> tuple[RegionInstance, ...]: + return tuple(instance for key, instance in self._instances.items() if self._run_scopes[key] is run_scope) + + def _settle(self, instance: RegionInstance) -> None: + if id(instance) not in self._instances: + return + if instance._retains_cleanup_only_reachability(): + return + self._retire(id(instance)) + + def _retire(self, key: int) -> None: + self._instances.pop(key, None) + self._run_scopes.pop(key, None) + + class RegionInstance: def __init__(self, ctx: MaterializationContext, shape: SingleOwnerRegionShape) -> None: self.plan = ctx.plan @@ -317,126 +396,271 @@ def __init__(self, ctx: MaterializationContext, shape: SingleOwnerRegionShape) - ) self._worker = ctx.worker self._cleanup_resources = getattr(ctx.worker, "_building_run_resources", None) - self._region = None + self._payload_mapping: Any | None = None + self._counter_mapping: Any | None = None self._payload_part: PayloadPart | None = None self._counter_part: CounterPart | None = None - self._state = RegionInstanceState.PLANNED + self._payload_local_view: RegionPartLocalView | None = None + self._counter_local_view: RegionPartLocalView | None = None + self._provider_resource_id = 0 + self._release_client: ProviderReleaseClient | None = None + self._allocate_client: ProviderAllocateClient | None = None + self._state: RegionInstanceState | None = None self._cleanup_error: BaseException | None = None + self._close_attempted = False + self._ever_live = False + self._provider_release_committed = False + self._data_plane_error: BaseException | None = None @property def state(self) -> RegionInstanceState: + if self._state is None: + raise MaterializationError("region instance is not published") return self._state + @property + def provider_resource_id(self) -> int: + return int(self._provider_resource_id) + + @property + def data_plane_error(self) -> BaseException | None: + return self._data_plane_error + + def local_view(self, part: RegionPartKind) -> RegionPartLocalView | None: + if self._state is not RegionInstanceState.LIVE: + return None + selected = RegionPartKind(part) + if selected is RegionPartKind.PAYLOAD: + return self._payload_local_view + if selected is RegionPartKind.COUNTER: + return self._counter_local_view + raise ValueError("region part must be PAYLOAD or COUNTER") + @classmethod def planned(cls, ctx: MaterializationContext, shape: SingleOwnerRegionShape) -> RegionInstance: return cls(ctx, shape) - def _adopt_worker_chip_region(self, region: Any) -> None: - self._region = region - mapping = getattr(region, "_worker_host_mapping", None) - if mapping is None: - return - plan = self.plan - if not isinstance(plan, BackendPlan): - raise MaterializationRefusal(RefusalReason.UNSUPPORTED_PLAN, "materialized region requires a BackendPlan") - payload_access = _select_host_vmm_copy_access(plan.payload, self.provider, self.consumer, mapping) - counter_access = _select_host_vmm_copy_access(plan.counter, self.provider, self.consumer, mapping) - self._payload_part = PayloadPart( - RegionPartSpan(offset=int(mapping.payload_offset), nbytes=int(mapping.payload_bytes)), payload_access - ) - self._counter_part = CounterPart( - RegionPartSpan(offset=int(mapping.counter_offset), nbytes=int(mapping.counter_bytes)), counter_access - ) - def payload_write(self, offset: int, host_buffer: Any, nbytes: int | None = None) -> None: self._ensure_live() self._worker._require_region_control_context("region_instance.payload_write") payload_part = self._payload_part - if payload_part is None: - region = self._region - assert region is not None - region.payload_write(offset, host_buffer, nbytes) - return + assert payload_part is not None payload_part.write(offset, host_buffer, nbytes) def payload_read(self, offset: int, host_buffer: Any, nbytes: int | None = None) -> None: self._ensure_live() self._worker._require_region_control_context("region_instance.payload_read") payload_part = self._payload_part - if payload_part is None: - region = self._region - assert region is not None - region.payload_read(offset, host_buffer, nbytes) - return + assert payload_part is not None payload_part.read(offset, host_buffer, nbytes) def counter(self, offset: int): self._ensure_live() self._worker._require_region_control_context("region_instance.counter") counter_part = self._counter_part - if counter_part is None: - region = self._region - assert region is not None - return region.counter(offset) + assert counter_part is not None return counter_part.counter(offset) def close(self) -> None: if self._state is RegionInstanceState.CLOSED: return - if self._state in _TERMINAL_FAILURE_STATES and self._cleanup_error is not None: + if self._state is RegionInstanceState.CLOSE_FAILED and self._cleanup_error is not None: raise self._cleanup_error - if self._state is RegionInstanceState.ROLLED_BACK: - return - if self._state is RegionInstanceState.PLANNED: + if self._state is None and self._provider_resource_id == 0 and self._payload_mapping is None: self._state = RegionInstanceState.CLOSED + self._worker._region_instance_registry._settle(self) return - if self._region is None: + self._worker._require_region_control_before_submit("region_instance.close") + self._worker._region_instance_registry.close(self) + + def _materialize(self, spec: RegionAllocationSpec) -> None: + plan = self.plan + if not isinstance(plan, BackendPlan): + raise MaterializationRefusal(RefusalReason.UNSUPPORTED_PLAN, "materialized region requires a BackendPlan") + prior = self._worker._consume_worker_host_mapped_cleanup_error("region_instance.materialize") + if prior is not None: + raise prior + native = self._worker._worker + if native is None: + raise RuntimeError("region instance materialize requires Worker.init()") + self._allocate_client = ProviderAllocateClient(native, self.worker_id) + try: + result, payload_view, counter_view = self._allocate_client.allocate(spec) + except RegionAllocationError as exc: + self._provider_resource_id = int(exc.provisional_resource_id) + if exc.cleanup_debt_remaining: + error = self._worker._record_unreclaimable( + f"region instance: allocation left cleanup debt for resource " + f"{exc.provisional_resource_id} on worker {self.worker_id}; no further work is admitted", + exc, + ) + self._fail_terminal(error) + raise error self._state = RegionInstanceState.CLOSED - return - self._worker._require_region_control_context("region_instance.close") - self._state = RegionInstanceState.CLOSING + raise + except BaseException: + committed = int(self._allocate_client.committed_resource_id) + if committed: + self._provider_resource_id = committed + self._release_client = ProviderReleaseClient(native, self.worker_id) + raise + self._provider_resource_id = int(result.provider_resource_id) + self._release_client = ProviderReleaseClient(native, self.worker_id) try: - self._worker._close_worker_chip_region( - self._region, - self._cleanup_resources, - poison_on_error=True, + validate_committed_region_allocation( + plan, + spec, + result, + payload_view, + counter_view, + expected_capability_type=self._worker._provider_import_capability_type(), + expected_device_id=self._worker._provider_import_device_id(self.worker_id), + ) + payload_lease = self._worker._import_region_part_lease( + self.worker_id, self._provider_resource_id, result.export_descriptor.payload + ) + self._payload_mapping = payload_lease + counter_lease = self._worker._import_region_part_lease( + self.worker_id, self._provider_resource_id, result.export_descriptor.counter + ) + self._counter_mapping = counter_lease + self._payload_part = PayloadPart( + RegionPartSpan(offset=0, nbytes=int(spec.payload.logical_bytes)), + _select_host_vmm_copy_access(plan.payload, self.provider, self.consumer, payload_lease), ) + self._counter_part = CounterPart( + RegionPartSpan(offset=0, nbytes=int(spec.counter.logical_bytes)), + _select_host_vmm_copy_access(plan.counter, self.provider, self.consumer, counter_lease), + ) + self._payload_local_view = payload_view + self._counter_local_view = counter_view except BaseException as exc: - self._cleanup_error = exc - self._state = RegionInstanceState.CLOSE_FAILED + self._abort_materialization(exc) raise - self._state = RegionInstanceState.CLOSED + self._ever_live = True + self._state = RegionInstanceState.LIVE + + def _abort_materialization(self, cause: BaseException) -> None: + if isinstance(cause, RegionAllocationError): + poison = bool(cause.cleanup_debt_remaining) + else: + client = self._allocate_client + poison = bool(client is not None and client.dispatch_started) + close_error: BaseException | None = None + try: + self._close_owned(poison_on_error=False) + except BaseException as exc: # noqa: BLE001 + close_error = exc + if not poison: + if close_error is None and self._state is not RegionInstanceState.CLOSE_FAILED: + self._state = RegionInstanceState.CLOSED + elif close_error is not None: + self._cleanup_error = close_error + self._state = RegionInstanceState.CLOSE_FAILED + return + if close_error is not None and close_error.__cause__ is None: + close_error.__cause__ = cause + self._cleanup_error = self._worker._record_unreclaimable( + f"region instance: committed allocation could not be published on worker {self.worker_id}; " + "no further work is admitted", + close_error or cause, + ) + self._state = RegionInstanceState.CLOSE_FAILED + + def _close_mapping_leases(self) -> list[BaseException]: + errors: list[BaseException] = [] + for lease in (self._payload_mapping, self._counter_mapping): + if lease is None: + continue + try: + closer = getattr(lease, "close", None) + if closer is not None: + closer() + else: + _worker_host_mapped_region_close(int(lease)) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + return errors + + def _release_provider_resource(self) -> BaseException | None: + if self._release_client is None or int(self._provider_resource_id) == 0: + return None + try: + result = self._release_client.release(int(self._provider_resource_id)) + except BaseException as exc: # noqa: BLE001 + return exc + if result.status in (ProviderReleaseStatus.RELEASED, ProviderReleaseStatus.ALREADY_GONE): + self._provider_release_committed = True + return None + if result.status is ProviderReleaseStatus.CLEANUP_INCOMPLETE: + return RuntimeError( + f"region instance: provider cleanup incomplete for resource {self._provider_resource_id}" + ) + if result.status is ProviderReleaseStatus.UNKNOWN_RESOURCE: + return RuntimeError(f"region instance: provider resource {self._provider_resource_id} is unknown") + return None - def rollback(self) -> None: - if self._state in (RegionInstanceState.CLOSED, RegionInstanceState.ROLLED_BACK): + def _close_owned(self, *, poison_on_error: bool) -> None: + if self._state is RegionInstanceState.CLOSED: return - if self._state in _TERMINAL_FAILURE_STATES and self._cleanup_error is not None: + if self._state is RegionInstanceState.CLOSE_FAILED and self._cleanup_error is not None: raise self._cleanup_error - if self._region is None: - self._state = RegionInstanceState.ROLLED_BACK + if self._close_attempted: + if self._cleanup_error is not None: + raise self._cleanup_error return - self._worker._require_region_control_context("region_instance.rollback") - self._state = RegionInstanceState.ROLLING_BACK - try: - self._worker._close_worker_chip_region( - self._region, - self._cleanup_resources, - poison_on_error=True, + self._close_attempted = True + if self._state is RegionInstanceState.LIVE: + self._state = RegionInstanceState.CLOSING + errors = self._close_mapping_leases() + release_error = self._release_provider_resource() + if release_error is not None: + errors.append(release_error) + if not errors: + self._state = RegionInstanceState.CLOSED + return + self._cleanup_error = errors[0] + self._state = RegionInstanceState.CLOSE_FAILED + if poison_on_error: + self._worker._record_unreclaimable( + f"region instance: resource {self._provider_resource_id} on worker {self.worker_id} " + "could not be fully reclaimed; no further work is admitted", + self._cleanup_error, ) - except BaseException as exc: - self._cleanup_error = exc - self._state = RegionInstanceState.ROLLBACK_FAILED - raise - self._state = RegionInstanceState.ROLLED_BACK + raise self._cleanup_error - def _rollback_after_failed_materialization(self) -> None: - self.rollback() + def _fail_terminal(self, error: BaseException) -> None: + self._cleanup_error = error + self._state = RegionInstanceState.CLOSE_FAILED + + def _retains_cleanup_only_reachability(self) -> bool: + if self._state is RegionInstanceState.CLOSED: + return False + if self._state is RegionInstanceState.CLOSE_FAILED: + return not (self._ever_live and self._provider_release_committed) + return True def _ensure_live(self) -> None: if self._state is not RegionInstanceState.LIVE: - raise MaterializationError(f"region instance is not live: {self._state.value}") - if self._region is None: - raise MaterializationError("region instance has no adopted worker-chip region") + label = "unpublished" if self._state is None else self._state.value + raise MaterializationError(f"region instance is not live: {label}") + + +def project_region_allocation_spec(plan: object, layout: object) -> RegionAllocationSpec: + if not isinstance(plan, BackendPlan): + raise TypeError("project_region_allocation_spec requires a BackendPlan") + if not isinstance(layout, RegionLayoutSpec): + raise TypeError("project_region_allocation_spec requires a RegionLayoutSpec") + return RegionAllocationSpec( + payload=RegionPartAllocationSpec( + planned_backing_kind=plan.payload.backend_kind, + logical_bytes=int(layout.payload_bytes), + ), + counter=RegionPartAllocationSpec( + planned_backing_kind=plan.counter.backend_kind, + logical_bytes=int(layout.counter_bytes), + ), + ) def validate_single_owner_region_shape(ctx: MaterializationContext) -> SingleOwnerRegionShape: # noqa: PLR0912 @@ -508,22 +732,81 @@ def validate_single_owner_region_shape(ctx: MaterializationContext) -> SingleOwn def materialize_region_instance(ctx: MaterializationContext) -> RegionInstance: shape = validate_single_owner_region_shape(ctx) + spec = project_region_allocation_spec(ctx.plan, ctx.layout) instance = RegionInstance.planned(ctx, shape) + resources = getattr(ctx.worker, "_building_run_resources", None) + ctx.worker._region_instance_registry.track(instance, resources) + if resources is not None: + resources.requires_ordered_cleanup = True try: - instance._state = RegionInstanceState.OWNER_CREATED - region = ctx.worker._create_worker_chip_region( - shape.worker_id, - int(ctx.layout.payload_bytes), - int(ctx.layout.counter_bytes), - ) - instance._adopt_worker_chip_region(region) - instance._state = RegionInstanceState.LIVE + instance._materialize(spec) return instance - except BaseException: - instance._rollback_after_failed_materialization() + except BaseException as exc: + if instance._state is None: + instance._abort_materialization(exc) + ctx.worker._region_instance_registry._settle(instance) raise +def validate_committed_region_allocation( # noqa: PLR0912 + plan: BackendPlan, + spec: RegionAllocationSpec, + result: RegionAllocationResult, + payload_view: RegionPartLocalView, + counter_view: RegionPartLocalView, + *, + expected_capability_type: type, + expected_device_id: int | None, +) -> None: + payload_export = result.export_descriptor.payload + counter_export = result.export_descriptor.counter + if payload_export.planned_backing_kind is not plan.payload.backend_kind: + raise RuntimeError("committed PAYLOAD planned backing does not match the admitted plan") + if counter_export.planned_backing_kind is not plan.counter.backend_kind: + raise RuntimeError("committed COUNTER planned backing does not match the admitted plan") + if payload_export.planned_backing_kind is not spec.payload.planned_backing_kind: + raise RuntimeError("committed PAYLOAD planned backing does not match the admitted spec") + if counter_export.planned_backing_kind is not spec.counter.planned_backing_kind: + raise RuntimeError("committed COUNTER planned backing does not match the admitted spec") + if int(payload_export.logical_bytes) != int(spec.payload.logical_bytes): + raise RuntimeError("committed PAYLOAD logical_bytes do not match the admitted spec") + if int(counter_export.logical_bytes) != int(spec.counter.logical_bytes): + raise RuntimeError("committed COUNTER logical_bytes do not match the admitted spec") + if not isinstance(payload_export.import_capability, expected_capability_type) or not isinstance( + counter_export.import_capability, expected_capability_type + ): + raise RuntimeError("committed import capability does not match the current execution environment") + if expected_capability_type is VmmShareableHandleImport: + payload_cap = payload_export.import_capability + counter_cap = counter_export.import_capability + assert isinstance(payload_cap, VmmShareableHandleImport) + assert isinstance(counter_cap, VmmShareableHandleImport) + if int(payload_cap.shareable_handle) == int(counter_cap.shareable_handle): + raise RuntimeError("committed VMM shareable handles must be distinct") + if expected_device_id is not None and ( + int(payload_cap.device_id) != int(expected_device_id) + or int(counter_cap.device_id) != int(expected_device_id) + ): + raise RuntimeError("committed VMM device_id is outside this worker's device namespace") + elif expected_capability_type is PosixShmImport: + payload_cap = payload_export.import_capability + counter_cap = counter_export.import_capability + assert isinstance(payload_cap, PosixShmImport) + assert isinstance(counter_cap, PosixShmImport) + if payload_cap.shm_name == counter_cap.shm_name: + raise RuntimeError("committed POSIX shm tokens must be distinct") + if payload_view.part is not RegionPartKind.PAYLOAD or counter_view.part is not RegionPartKind.COUNTER: + raise RuntimeError("committed local views must be PAYLOAD then COUNTER") + if int(payload_view.logical_bytes) != int(spec.payload.logical_bytes): + raise RuntimeError("committed PAYLOAD local view bytes do not match the admitted spec") + if int(counter_view.logical_bytes) != int(spec.counter.logical_bytes): + raise RuntimeError("committed COUNTER local view bytes do not match the admitted spec") + try: + validate_independent_local_views(payload_view, counter_view) + except ValueError as exc: + raise RuntimeError(str(exc) or "committed local views are not independent") from exc + + def _record_for(ctx: MaterializationContext, endpoint: Any) -> EndpointRecord: try: return ctx.registry.record_for(endpoint) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index f8333a28e5..1ed690bb07 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -93,14 +93,11 @@ def my_l4_orch(orch, args, config): RUNTIME_ENV_RING_COUNT, WorkerType, _emit_host_span, - _l3_child_onboard_region_close, - _l3_child_onboard_region_create, _mailbox_load_i32, _mailbox_store_i32, _read_control_copy_request, _set_host_span_level_prefix, _worker_host_mapped_region_ack_cleanup_error, - _worker_host_mapped_region_close, _worker_host_mapped_region_import_onboard, _worker_host_mapped_region_import_sim, _worker_host_mapped_region_peek_cleanup_error, @@ -162,7 +159,33 @@ def my_l4_orch(orch, args, config): at, parse_endpoint_path, ) -from .comm_region import MaterializationContext, RegionInstance, materialize_region_instance +from .comm_provider import ( + DeviceAllocationTarget, + PosixShmImport, + ProviderRegionStore, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationContext, + RegionAllocationSpec, + RegionEnvironmentKind, + RegionPartExportDescriptor, + RegionPartKind, + VmmShareableHandleImport, +) +from .comm_provider_control import ( + handle_ctrl_region_allocate, + handle_ctrl_region_release, +) +from .comm_region import ( + MaterializationContext, + MaterializationError, + RegionInstance, + RegionInstanceRegistry, + RegionInstanceState, + materialize_region_instance, + project_region_allocation_spec, + validate_single_owner_region_shape, +) from .global_comm_domain import ( CTRL_GLOBAL_DOMAIN_COPY_FROM, CTRL_GLOBAL_DOMAIN_COPY_TO, @@ -235,22 +258,8 @@ def my_l4_orch(orch, args, config): _Worker, ) from .worker_chip_orch_comm import ( - _CTRL_SHM_TOKEN_BYTES, - _REGION_CREATE_REPLY, - _REGION_CREATE_REPLY_BYTES, - _REGION_CREATE_REQUEST, - _REGION_CREATE_REQUEST_BYTES, - _REGION_LAYOUT_ALIGNMENT, - _REGION_MAGIC_VERSION, WorkerChipOrchRegion, - WorkerChipRegionAccessProfile, - WorkerChipRegionCreateRequest, - WorkerHostRegionMapping, - _align_up, - _checked_add_u64, - decode_region_create_reply, - peek_region_create_reply_region_id, - validate_region_create_reply, + worker_chip_orch_region_desc_from_local_views, ) from .worker_level import WorkerLevel from .worker_level import span_prefix as _span_prefix @@ -539,8 +548,8 @@ def _shm_name(token: str, suffix: str): } -_CTRL_WORKER_CHIP_REGION_CREATE = 16 -_CTRL_WORKER_CHIP_REGION_RELEASE = 17 +_CTRL_REGION_ALLOCATE = 16 +_CTRL_REGION_RELEASE = 17 _CTRL_COMMITTED_DEVICE_MEMORY = 18 # L4-to-local-L3 envelope for the Global CommDomain control protocol. The # enclosed command uses remote_l3_protocol.ControlName; values 18-23 belong to @@ -2334,26 +2343,6 @@ def _handle_ctrl_comm_init(cw: ChipWorker, buf: memoryview) -> None: cw._comm_base_handle_cached = int(handle) -@dataclass -class _HostWorkerChipRegion: - region_id: int - payload_bytes: int - counter_offset: int - counter_bytes: int - total_bytes: int - shm: SharedMemory | None = None - dev_ptr: int = 0 - onboard_handle: int = 0 - - -@dataclass -class _HostWorkerChipRegionStore: - """Per-chip-child registry of live L3-L2 direct regions (loop-local state).""" - - regions: dict[int, _HostWorkerChipRegion] = field(default_factory=dict) - next_region_id: int = 1 - - @dataclass class _L2GlobalDomain: domain_id: int @@ -2373,167 +2362,38 @@ class _L2GlobalDomainStore: domains: dict[int, _L2GlobalDomain] = field(default_factory=dict) -@dataclass(frozen=True) -class _HostWorkerChipRegionReplyMeta: - payload_base: int - backing_name: bytes - access_profile: WorkerChipRegionAccessProfile - mapping_bytes: int - shareable_handle: int - - -def _release_host_worker_chip_region(region: _HostWorkerChipRegion) -> None: - if region.shm is not None: - region.shm.close() - region.shm.unlink() - return - if region.onboard_handle: - _l3_child_onboard_region_close(region.onboard_handle) - - -def _create_sim_worker_chip_region( - request: WorkerChipRegionCreateRequest, region_id: int, counter_offset: int, total_bytes: int -) -> tuple[_HostWorkerChipRegion, _HostWorkerChipRegionReplyMeta]: - shm = SharedMemory(create=True, size=total_bytes) - region = _HostWorkerChipRegion( - region_id=region_id, - payload_bytes=request.payload_bytes, - counter_offset=counter_offset, - counter_bytes=request.counter_bytes, - total_bytes=total_bytes, - shm=shm, - ) - region_buf = cast(memoryview, shm.buf) - region_buf[counter_offset : counter_offset + request.counter_bytes] = b"\x00" * request.counter_bytes - exported = ctypes.c_char.from_buffer(region_buf) +def _handle_ctrl_region_allocate(buf: memoryview, store: ProviderRegionStore) -> None: + request_shm_name = _read_shm_name(buf, _OFF_ARGS) + reply_shm_name = _read_shm_name(buf, _OFF_ARGS + _CTRL_SHM_NAME_BYTES) + req_shm = SharedMemory(name=request_shm_name) + reply_shm = SharedMemory(name=reply_shm_name) + req_buf = cast(memoryview, req_shm.buf) + reply_buf = cast(memoryview, reply_shm.buf) try: - payload_base = ctypes.addressof(exported) + handle_ctrl_region_allocate(req_buf, reply_buf, store) finally: - del exported - del region_buf - backing_name = shm.name.encode("utf-8") - if len(backing_name) >= _CTRL_SHM_TOKEN_BYTES: - raise RuntimeError("CTRL_WORKER_CHIP_REGION_CREATE backing shm token is too long") - meta = _HostWorkerChipRegionReplyMeta( - payload_base=payload_base, - backing_name=backing_name, - access_profile=WorkerChipRegionAccessProfile.SIM_POSIX_SHM, - mapping_bytes=total_bytes, - shareable_handle=0, - ) - return region, meta - - -def _create_onboard_worker_chip_region( - cw: ChipWorker, request: WorkerChipRegionCreateRequest, region_id: int, counter_offset: int, total_bytes: int -) -> tuple[_HostWorkerChipRegion, _HostWorkerChipRegionReplyMeta]: - export = _l3_child_onboard_region_create(total_bytes) - dev_ptr = int(export.device_addr) - region = _HostWorkerChipRegion( - region_id=region_id, - payload_bytes=request.payload_bytes, - counter_offset=counter_offset, - counter_bytes=request.counter_bytes, - total_bytes=total_bytes, - dev_ptr=dev_ptr, - onboard_handle=int(export.registry_handle), - ) - zeros = ctypes.create_string_buffer(request.counter_bytes) - cw.copy_to(dev_ptr + counter_offset, ctypes.addressof(zeros), request.counter_bytes) - meta = _HostWorkerChipRegionReplyMeta( - payload_base=dev_ptr, - backing_name=b"", - access_profile=WorkerChipRegionAccessProfile.ONBOARD_VMM, - mapping_bytes=int(export.mapping_bytes), - shareable_handle=int(export.shareable_handle), - ) - return region, meta + del req_buf + del reply_buf + req_shm.close() + reply_shm.close() -def _handle_ctrl_worker_chip_region_create( - cw: ChipWorker, buf: memoryview, chip_platform: str, store: _HostWorkerChipRegionStore -) -> None: +def _handle_ctrl_region_release(buf: memoryview, store: ProviderRegionStore) -> None: request_shm_name = _read_shm_name(buf, _OFF_ARGS) reply_shm_name = _read_shm_name(buf, _OFF_ARGS + _CTRL_SHM_NAME_BYTES) req_shm = SharedMemory(name=request_shm_name) reply_shm = SharedMemory(name=reply_shm_name) req_buf = cast(memoryview, req_shm.buf) reply_buf = cast(memoryview, reply_shm.buf) - region: _HostWorkerChipRegion | None = None try: - fields = _REGION_CREATE_REQUEST.unpack_from(req_buf, 0) - request = WorkerChipRegionCreateRequest( - magic_version=int(fields[0]), - request_bytes=int(fields[1]), - payload_bytes=int(fields[2]), - counter_bytes=int(fields[3]), - ) - # Reject ABI mismatches loudly. The reply carries this child's own - # magic (not the request echo) so the L3 side can detect version skew. - if request.magic_version != _REGION_MAGIC_VERSION: - raise RuntimeError("CTRL_WORKER_CHIP_REGION_CREATE magic_version mismatch") - if request.request_bytes != _REGION_CREATE_REQUEST_BYTES: - raise RuntimeError("CTRL_WORKER_CHIP_REGION_CREATE request_bytes mismatch") - if request.payload_bytes <= 0: - raise RuntimeError("CTRL_WORKER_CHIP_REGION_CREATE payload_bytes must be positive") - if request.counter_bytes <= 0 or request.counter_bytes % 4 != 0: - raise RuntimeError("CTRL_WORKER_CHIP_REGION_CREATE counter_bytes must be positive and a multiple of 4") - counter_offset = _align_up(request.payload_bytes, _REGION_LAYOUT_ALIGNMENT) - total_bytes = _checked_add_u64(counter_offset, request.counter_bytes) - - region_id = store.next_region_id - store.next_region_id += 1 - if str(chip_platform).endswith("sim"): - region, meta = _create_sim_worker_chip_region(request, region_id, counter_offset, total_bytes) - else: - region, meta = _create_onboard_worker_chip_region(cw, request, region_id, counter_offset, total_bytes) - _REGION_CREATE_REPLY.pack_into( - reply_buf, - 0, - _REGION_MAGIC_VERSION, - region_id, - meta.payload_base, - request.payload_bytes, - meta.payload_base + counter_offset, - request.counter_bytes, - int(meta.access_profile), - 0, - int(getattr(cw, "device_id", -1)), - meta.backing_name + b"\x00" * (_CTRL_SHM_TOKEN_BYTES - len(meta.backing_name)), - meta.mapping_bytes, - meta.shareable_handle, - ) - store.regions[region_id] = region - region = None + handle_ctrl_region_release(req_buf, reply_buf, store) finally: - if region is not None: - try: - _release_host_worker_chip_region(region) - except (BufferError, FileNotFoundError, OSError, RuntimeError): - pass del req_buf del reply_buf req_shm.close() reply_shm.close() -def _handle_ctrl_worker_chip_region_release(buf: memoryview, store: _HostWorkerChipRegionStore) -> None: - region_id = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0] - region = store.regions.pop(int(region_id), None) - if region is None: - return - _release_host_worker_chip_region(region) - - -def _sweep_host_worker_chip_regions(store: _HostWorkerChipRegionStore) -> None: - for region_id in list(store.regions): - region = store.regions.pop(region_id) - try: - _release_host_worker_chip_region(region) - except (BufferError, FileNotFoundError, OSError, RuntimeError): - pass - - def _open_global_domain_payload(buf: memoryview) -> tuple[SharedMemory, memoryview, int]: payload_size = int(struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0]) if payload_size <= 0: @@ -2743,12 +2603,65 @@ def _handle_ctrl_global_domain_copy( def _sweep_l2_global_domains(cw: ChipWorker, store: _L2GlobalDomainStore) -> None: - for domain_id in list(store.domains): + failures: list[tuple[int, BaseException]] = [] + first: BaseException | None = None + for domain_id in sorted(store.domains): store.domains.pop(domain_id, None) try: cw._impl.comm_global_domain_release(int(domain_id)) - except Exception: # noqa: BLE001 - pass + except BaseException as exc: # noqa: BLE001 + if first is None: + first = exc + failures.append((int(domain_id), exc)) + if not failures: + return + error = RuntimeError( + "domain cleanup failed: " + "; ".join(f"domain {domain_id}: {exc}" for domain_id, exc in failures) + ) + error.__cause__ = first + raise error + + +def _provider_sweep_debt_errors(results: tuple[ProviderReleaseResult, ...]) -> list[BaseException]: + errors: list[BaseException] = [] + for result in results: + if result.status is not ProviderReleaseStatus.CLEANUP_INCOMPLETE: + continue + if result.failures: + detail = ", ".join( + f"{failure.part.name} {failure.backend_operation.name} {failure.typed_cause.name}" + for failure in result.failures + ) + else: + detail = "CLEANUP_INCOMPLETE" + errors.append(RuntimeError(f"provider resource {result.provider_resource_id}: {detail}")) + return errors + + +def _teardown_chip_process_resources( + import_registry: ImportRegistry, + cw: ChipWorker, + global_domain_store: _L2GlobalDomainStore, + provider_region_store: ProviderRegionStore, +) -> None: + errors: list[BaseException] = [] + try: + import_registry.close() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + try: + _sweep_l2_global_domains(cw, global_domain_store) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + try: + errors.extend(_provider_sweep_debt_errors(provider_region_store.sweep())) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if not errors: + return + aggregated = RuntimeError("chip process resource teardown failed: " + "; ".join(str(item) for item in errors)) + aggregated.__cause__ = errors[0] + raise aggregated def _handle_ctrl_release_domain(cw: ChipWorker, buf: memoryview) -> None: @@ -2826,7 +2739,13 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de DEVICE_MALLOC/VMM_WINDOW backings this chip may materialize. """ prepared = prepared if prepared is not None else set() - worker_chip_region_store = _HostWorkerChipRegionStore() + environment = RegionEnvironmentKind.SIM if str(chip_platform).endswith("sim") else RegionEnvironmentKind.ONBOARD + provider_region_store = ProviderRegionStore( + RegionAllocationContext( + environment_kind=environment, + target=DeviceAllocationTarget(int(device_id)), + ) + ) import_registry = ImportRegistry(ImportContext(is_host_endpoint=False, owning_chip_instance_id=owner_instance_id)) global_domain_store = _L2GlobalDomainStore() @@ -2987,10 +2906,10 @@ def handle_control( # noqa: PLR0912, PLR0915 -- one branch per control sub-comm _handle_ctrl_release_domain(cw, buf) elif sub_cmd == _CTRL_COMM_INIT: _handle_ctrl_comm_init(cw, buf) - elif sub_cmd == _CTRL_WORKER_CHIP_REGION_CREATE: - _handle_ctrl_worker_chip_region_create(cw, buf, chip_platform, worker_chip_region_store) - elif sub_cmd == _CTRL_WORKER_CHIP_REGION_RELEASE: - _handle_ctrl_worker_chip_region_release(buf, worker_chip_region_store) + elif sub_cmd == _CTRL_REGION_ALLOCATE: + _handle_ctrl_region_allocate(buf, provider_region_store) + elif sub_cmd == _CTRL_REGION_RELEASE: + _handle_ctrl_region_release(buf, provider_region_store) elif sub_cmd == _CTRL_COMMITTED_DEVICE_MEMORY: struct.pack_into("Q", buf, _CTRL_OFF_RESULT, cw.committed_device_memory) elif sub_cmd == _CTRL_IMPORT_RELEASE: @@ -3282,9 +3201,7 @@ def submit_frame(frame: _StagedFrame) -> None: else: _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: - import_registry.close() - _sweep_l2_global_domains(cw, global_domain_store) - _sweep_host_worker_chip_regions(worker_chip_region_store) + _teardown_chip_process_resources(import_registry, cw, global_domain_store, provider_region_store) def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, identity tables, log config, prewarm sizing) must cross the fork as explicit COW args; the child cannot read parent state after os.fork @@ -3853,7 +3770,6 @@ class _RunResources: pending_release_domains: list[CommDomainHandle] = field(default_factory=list) live_global_domains: dict[str, GlobalCommDomainHandle] = field(default_factory=dict) pending_release_global_domains: list[GlobalCommDomainHandle] = field(default_factory=list) - worker_chip_regions: list[Any] = field(default_factory=list) worker_chip_orch_comm_host_buffers: dict[int, int] = field(default_factory=dict) # Every Buffer identity a NEXT_LEVEL dispatch (submit_next_level / _group) sent as a Tensor # arg during this run. release_buffer() checks this set across every not-yet-settled run before @@ -4613,7 +4529,7 @@ def __init__( self._endpoint_registry_epoch: int = 0 self._region_access_service: RegionAccessService | None = None - self._live_worker_chip_regions: list[Any] = [] + self._region_instance_registry = RegionInstanceRegistry() self._worker_chip_orch_comm_host_buffers: dict[int, int] = {} # Live-provenance of child (kind4, device) pointers, keyed on the exact @@ -6456,6 +6372,29 @@ def _materialize_region_instance( MaterializationContext(worker=self, registry=registry, plan=plan, layout=layout_summary) ) + def _admitted_worker_chip_region_context( + self, worker_id: int, payload_bytes: int, counter_bytes: int + ) -> MaterializationContext: + worker_id = int(worker_id) + root_path = _format_worker_path(int(self.level)) + provider_path = _format_worker_path(2, parent_path=root_path, index=worker_id) + provider = at(provider_path, DEVICE_AICPU) + layout = RegionLayoutSpec(payload_bytes=int(payload_bytes), counter_bytes=int(counter_bytes)) + members = (at(root_path, HOST_CPU), provider) + topology = SingleOwner(provider=provider) + registry = self._get_endpoint_registry() + resolved = registry.resolve_region_spec(members, topology) + plan = BackendResolver(registry, self._get_region_access_service()).plan(resolved, layout) + ctx = MaterializationContext(worker=self, registry=registry, plan=plan, layout=layout) + validate_single_owner_region_shape(ctx) + return ctx + + def _project_admitted_worker_chip_region_spec( + self, worker_id: int, payload_bytes: int, counter_bytes: int + ) -> RegionAllocationSpec: + ctx = self._admitted_worker_chip_region_context(int(worker_id), int(payload_bytes), int(counter_bytes)) + return project_region_allocation_spec(ctx.plan, ctx.layout) + def _register_into_snapshot_or_wait(self, reg: _CallableRegistration) -> CallableHandle | None: """Linearize a level>=3 register against the startup epoch. @@ -8271,13 +8210,20 @@ def _poison_worker_chip_region_from_endpoint_error( region_id = int(match.group(1)) if region_id == 0: return False - poisoned = False - regions = self._live_worker_chip_regions if resources is None else resources.worker_chip_regions - for region in regions: - if int(region.region_id) == region_id: - region._poison() - poisoned = True - return poisoned + try: + self._region_instance_registry.record_data_plane_failure(resources, region_id, exc) + except MaterializationError as routing: + self._record_unreclaimable( + f"region instance: data-plane failure routing failed for resource {region_id}; " + "no further work is admitted", + routing, + ) + return True + self._record_unreclaimable( + f"region instance: issued local operation failed for resource {region_id}; no further work is admitted", + exc, + ) + return True def _register_worker_chip_orch_comm_host_buffer(self, handle) -> None: if not isinstance(handle, Buffer): @@ -8368,281 +8314,85 @@ def _consume_worker_host_mapped_cleanup_error(self, api: str) -> RuntimeError | with self._hierarchical_start_cv: return self._consume_worker_host_mapped_cleanup_error_locked(api) - def _create_worker_chip_region(self, worker_id: int, payload_bytes: int, counter_bytes: int): # noqa: PLR0912 + def _import_provider_part(self, export: RegionPartExportDescriptor): + capability = export.import_capability + if isinstance(capability, PosixShmImport): + return _worker_host_mapped_region_import_sim(capability.shm_name, int(export.mapping_bytes), self._owner_id) + if isinstance(capability, VmmShareableHandleImport): + return _worker_host_mapped_region_import_onboard( + int(capability.device_id), + int(capability.shareable_handle), + int(export.mapping_bytes), + self._owner_id, + ) + raise RuntimeError("create_worker_chip_region: unsupported import capability") + + def _provider_import_capability_type(self) -> type: + platform = str(self._config.get("platform", "")) + return PosixShmImport if platform.endswith("sim") else VmmShareableHandleImport + + def _provider_import_device_id(self, worker_id: int) -> int: + device_ids = self._config.get("device_ids", []) + return int(device_ids[int(worker_id)]) + + def _import_region_part_lease(self, worker_id: int, resource_id: int, export: RegionPartExportDescriptor): + return self._import_provider_part(export) + + def _create_worker_chip_region(self, worker_id: int, payload_bytes: int, counter_bytes: int): if payload_bytes <= 0: raise ValueError("create_worker_chip_region: payload_bytes must be positive") if counter_bytes <= 0 or counter_bytes % 4 != 0: raise ValueError("create_worker_chip_region: counter_bytes must be positive and a multiple of 4") self._validate_worker_chip_id(int(worker_id)) - prior_native_cleanup_error = self._consume_worker_host_mapped_cleanup_error("create_worker_chip_region") - if prior_native_cleanup_error is not None: - raise prior_native_cleanup_error resources = self._building_run_resources - req_shm = SharedMemory(create=True, size=_REGION_CREATE_REQUEST_BYTES) - reply_shm = SharedMemory(create=True, size=_REGION_CREATE_REPLY_BYTES) - req_buf = cast(memoryview, req_shm.buf) - reply_buf = cast(memoryview, reply_shm.buf) - region_id = 0 - native_mapping_handle = None - worker_host_mapping = None + instance: RegionInstance | None = None region = None - dispatched = False required_ordered_cleanup_before = resources.requires_ordered_cleanup if resources is not None else False try: - WorkerChipRegionCreateRequest( - magic_version=_REGION_MAGIC_VERSION, - request_bytes=_REGION_CREATE_REQUEST_BYTES, - payload_bytes=int(payload_bytes), - counter_bytes=int(counter_bytes), - ).encode_into(req_buf) - worker = self._worker - assert worker is not None - # From here the chip may have created a region. The create releases - # the GIL, so an interrupt can land after the child committed and - # before the id is read back — the rollback below cannot assume - # "no id" means "nothing exists". - dispatched = True - worker.control_worker_chip_region_create(int(worker_id), req_shm.name, reply_shm.name) - # Peek before decode: decode rejects malformed replies, but the - # child has already created the region and the rollback below - # still needs the id. - region_id = peek_region_create_reply_region_id(reply_buf) - reply = decode_region_create_reply(reply_buf) - platform = str(self._config.get("platform", "")) - expected_access_profile = ( - WorkerChipRegionAccessProfile.SIM_POSIX_SHM - if platform.endswith("sim") - else WorkerChipRegionAccessProfile.ONBOARD_VMM + ctx = self._admitted_worker_chip_region_context(int(worker_id), int(payload_bytes), int(counter_bytes)) + instance = materialize_region_instance(ctx) + payload_view = instance.local_view(RegionPartKind.PAYLOAD) + counter_view = instance.local_view(RegionPartKind.COUNTER) + if payload_view is None or counter_view is None: + raise RuntimeError("create_worker_chip_region: materialized instance is missing local views") + desc = worker_chip_orch_region_desc_from_local_views( + instance.provider_resource_id, payload_view, counter_view ) - counter_offset, total_bytes = validate_region_create_reply(reply, expected_access_profile) - if platform.endswith("sim"): - native_mapping_handle = _worker_host_mapped_region_import_sim( - reply.backing_shm, int(reply.mapping_bytes), self._owner_id - ) - else: - native_mapping_handle = _worker_host_mapped_region_import_onboard( - int(reply.device_id), - int(reply.shareable_handle), - int(reply.mapping_bytes), - self._owner_id, - ) - worker_host_mapping = WorkerHostRegionMapping( - worker_id=int(worker_id), - region_id=region_id, - access_profile=reply.access_profile, - total_bytes=total_bytes, - payload_offset=0, - payload_bytes=int(reply.desc.payload_bytes), - counter_offset=counter_offset, - counter_bytes=int(reply.desc.counter_bytes), - handle=native_mapping_handle, - ) - native_mapping_handle = None - region = WorkerChipOrchRegion(self, int(worker_id), reply.desc, worker_host_mapping) - self._live_worker_chip_regions.append(region) + region = WorkerChipOrchRegion(self, instance, desc) if resources is not None: - resources.worker_chip_regions.append(region) - # Region teardown is mailbox control on its owning chip. resources.requires_ordered_cleanup = True return region - except BaseException as exc: - mapping_cleanup_error: BaseException | None = None - if region is not None: - try: - self._live_worker_chip_regions.remove(region) - except ValueError: - pass - if resources is not None: - try: - resources.worker_chip_regions.remove(region) - except ValueError: - pass - resources.requires_ordered_cleanup = required_ordered_cleanup_before - region._expire() - if worker_host_mapping is not None: - try: - worker_host_mapping.close() - except BaseException as mapping_exc: # noqa: BLE001 - mapping_cleanup_error = mapping_exc - elif native_mapping_handle is not None: - try: - _worker_host_mapped_region_close(int(native_mapping_handle)) - except BaseException as mapping_exc: # noqa: BLE001 - mapping_cleanup_error = mapping_exc - if not region_id: - # The failure may have landed after the child created its - # region. The reply shm is parent-owned and zero-filled and the - # child writes the id as part of committing, so a nonzero id - # here is the child saying the region exists. - with contextlib.suppress(BaseException): - region_id = peek_region_create_reply_region_id(reply_buf) - if not region_id and dispatched and not isinstance(exc, Exception): - # An asynchronous unwind through the create. The child may still - # be finishing a region, and the id it would write lands in a - # reply this frame is about to unlink — leaving something on the - # chip that nothing here can name. An ordinary failure is not - # this case: the child releases its own region before reporting - # one, so a zero id there really does mean nothing exists. - self._record_unreclaimable( - f"create_worker_chip_region: interrupted on worker {int(worker_id)} before the region id was " - "read back; a region may be live on the chip and no further work is admitted", - exc, - ) - if region_id: + except BaseException: + if resources is not None: + resources.requires_ordered_cleanup = required_ordered_cleanup_before + if instance is not None and instance._state is RegionInstanceState.LIVE and not instance._close_attempted: try: - assert self._worker is not None - self._worker.control_worker_chip_region_release(int(worker_id), int(region_id)) - except BaseException as release_exc: # noqa: BLE001 - # The chip created the region and this call could not give - # it back. Nothing else knows the id — it was never tracked - # — so no later cleanup can reclaim it, and a rollback that - # swallows this reports an ordinary failure over a region - # that stays allocated for the worker's life. That is the - # unreclaimed-device-state condition, recorded here directly - # because there is no handle for a fence to fail on. + self._region_instance_registry.close(instance) + except BaseException as close_exc: # noqa: BLE001 raise self._record_unreclaimable( - f"create_worker_chip_region: rollback could not release region {region_id} on worker " - f"{int(worker_id)}; it is leaked and no further work is admitted", - release_exc, + f"create_worker_chip_region: rollback could not close the L3 Host mapping for region " + f"{int(instance.provider_resource_id)} on worker {int(worker_id)}; " + "it is leaked and no further work is admitted", + close_exc, ) deferred_native_cleanup_error = self._consume_worker_host_mapped_cleanup_error( "create_worker_chip_region rollback" ) if deferred_native_cleanup_error is not None: - deferred_exc = deferred_native_cleanup_error.__cause__ or deferred_native_cleanup_error - if mapping_cleanup_error is not None: - combined_exc = RuntimeError( - f"{mapping_cleanup_error}; deferred native cleanup also failed: {deferred_exc}" - ) - combined_exc.__cause__ = mapping_cleanup_error - mapping_cleanup_error = combined_exc - else: - mapping_cleanup_error = deferred_exc - if mapping_cleanup_error is not None: + region_id = int(instance.provider_resource_id) if instance is not None else 0 raise self._record_unreclaimable( f"create_worker_chip_region: rollback could not close the L3 Host mapping for region " f"{region_id} on worker {int(worker_id)}; it is leaked and no further work is admitted", - mapping_cleanup_error, - ) - raise exc - finally: - del req_buf - del reply_buf - for shm in (req_shm, reply_shm): - try: - shm.close() - shm.unlink() - except (BufferError, FileNotFoundError, OSError): - pass - - def _close_worker_chip_region( - self, - region, - resources: _RunResources | None = None, - *, - poison_on_error: bool = False, - ) -> None: - region_errors: list[BaseException] = [] - release_error: BaseException | None = None - - def primary_region_error() -> BaseException: - primary = region_errors[0] - tail = primary - seen = {id(tail)} - while tail.__cause__ is not None and id(tail.__cause__) not in seen: - tail = tail.__cause__ - seen.add(id(tail)) - for extra in region_errors[1:]: - tail.__cause__ = extra - tail = extra - return primary - - expired = bool(region.expired) - if not expired: - try: - region._close_worker_host_mapping() - except BaseException as exc: # noqa: BLE001 - region_errors.append(exc) - try: - if self._worker is not None and not region._chip_release_committed: - self._worker.control_worker_chip_region_release(region._worker_id, region.region_id) - region._chip_release_committed = True - region.free() - except BaseException as exc: # noqa: BLE001 - release_error = exc - region_errors.append(exc) - # A region remains tracked while chip ownership may still be live, so - # whole-tree close can replay it and instance cleanup can poison future - # admission. Once the chip release is committed or the native handle is - # already expired, cleanup may retire every tracking list that owns it. - if region_errors and resources is None and not poison_on_error: - raise primary_region_error() - if poison_on_error and region_errors and release_error is not None: - primary = primary_region_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", - primary, - ) - raise primary - try: - if not expired: - region._expire() - except BaseException as exc: # noqa: BLE001 - region_errors.append(exc) - try: - self._retire_worker_chip_region_tracking(region, resources) - except BaseException as exc: # noqa: BLE001 - region_errors.append(exc) - if region_errors: - primary = primary_region_error() - 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", - primary, + deferred_native_cleanup_error.__cause__ or deferred_native_cleanup_error, ) - raise primary - - def _retire_worker_chip_region_tracking(self, region, resources: _RunResources | None = None) -> None: - tracking_lists = [self._live_worker_chip_regions] - if resources is not None and resources.worker_chip_regions is not self._live_worker_chip_regions: - tracking_lists.insert(0, resources.worker_chip_regions) - errors: list[BaseException] = [] - - for owned_regions in tracking_lists: - try: - owned_regions[:] = [owned for owned in owned_regions if owned is not region] - except BaseException as exc: # noqa: BLE001 - errors.append(exc) - if errors: - raise errors[0] + raise - def _cleanup_worker_chip_regions(self, resources: _RunResources | None = None) -> None: - # A region stays tracked until both ownership debts commit. This makes a - # failed close replayable by the cleanup journal instead of publishing a - # false success after dropping the only reference to the region. - tracked = self._live_worker_chip_regions if resources is None else resources.worker_chip_regions - if not tracked: - return - regions = list(tracked) - errors: list[BaseException] = [] - for region in regions: - try: - self._close_worker_chip_region(region, resources) - except BaseException as exc: # noqa: BLE001 - errors.append(exc) - if errors: - raise errors[0] + def _sweep_region_instances(self) -> None: + self._region_instance_registry.sweep() def _close_worker_chip_orch_comm(self) -> None: - for region in self._live_worker_chip_regions: - try: - region._close_worker_host_mapping() - except RuntimeError: - pass - self._live_worker_chip_regions.clear() self._worker_chip_orch_comm_host_buffers.clear() + self._region_instance_registry.sweep() # ------------------------------------------------------------------ # Dynamic CommDomain allocation (driven by Orchestrator.allocate_domain; @@ -10906,10 +10656,6 @@ def _require_no_ordered_cleanup_failure(self, api: str) -> None: def _require_region_control_context(self, api: str) -> None: frame = _callback_frame_for(self) if frame is not None: - if frame.has_submitted_task: - raise RuntimeError( - f"Worker.{api}: RegionInstance access cannot follow a task submission in the same run" - ) self._require_no_ordered_cleanup_failure(api) return if id(self) in _held_control_reservations(): @@ -10917,6 +10663,14 @@ def _require_region_control_context(self, api: str) -> None: return raise RuntimeError(f"Worker.{api}: RegionInstance access requires an active orchestration/control context") + def _require_region_control_before_submit(self, api: str) -> None: + # close/release travels the mailbox; a ready-queue task already in the + # same run has no defined order against that command. + self._require_region_control_context(api) + frame = _callback_frame_for(self) + if frame is not None and frame.has_submitted_task: + raise RuntimeError(f"Worker.{api}: RegionInstance access cannot follow a task submission in the same run") + def _control_admission(self, api: str): """The direct-control ordering policy for a Worker-level command. @@ -11176,7 +10930,7 @@ def _release_native_run() -> None: ("endpoint_poison", _poison_endpoint), ("remote_slot_refs", lambda: self._release_active_remote_slot_refs(resources)), ("remote_frees", self._flush_pending_remote_frees), - ("worker_chip_regions", lambda: self._cleanup_worker_chip_regions(resources)), + ("region_instances", lambda: self._region_instance_registry.cleanup_run(resources)), ("worker_chip_host_buffers", resources.worker_chip_orch_comm_host_buffers.clear), ( "pending_global_domains", @@ -11337,7 +11091,7 @@ def _has_live_resources(self) -> bool: or bool(self._sub_pids or self._chip_pids or self._next_level_pids) or bool(self._sub_shms or self._chip_shms or self._next_level_shms) or any(group.process is not None or group.ready_dir is not None for group in self._mpi_l3_groups) - or bool(self._live_worker_chip_regions) + or bool(self._region_instance_registry._instances) or bool(self._live_domains) or bool(self._live_global_domains or self._failed_global_domain_releases) or bool(self._global_node_domains) @@ -11361,8 +11115,9 @@ def _describe_live_resources(self) -> str: n_mpi = sum(1 for group in self._mpi_l3_groups if group.process is not None or group.ready_dir is not None) if n_mpi: parts.append(f"{n_mpi} mpirun group(s)") - if self._live_worker_chip_regions: - parts.append(f"{len(self._live_worker_chip_regions)} L3-L2 region(s)") + n_instances = len(self._region_instance_registry._instances) + if n_instances: + parts.append(f"{n_instances} region instance(s)") if self._live_domains: parts.append(f"{len(self._live_domains)} comm domain(s)") if self._live_global_domains or self._failed_global_domain_releases: @@ -11843,7 +11598,7 @@ def _step(fn) -> None: # entries. Any survivor fences the transport teardown below. pre_transport_keys: set[tuple[str, str]] = set() for kind, identity, cleanup in ( - ("region", "all Worker-Chip regions", self._cleanup_worker_chip_regions), + ("region", "all region instances", self._sweep_region_instances), ("domain", "all Global CommDomains", self._release_all_live_global_domains), ("domain", "Global CommDomain nodes", self._release_all_global_domain_nodes), ("domain", "all CommDomains", self._release_all_live_domains), diff --git a/python/simpler/worker_chip_orch_comm.py b/python/simpler/worker_chip_orch_comm.py index cad7da835b..e3cb7dcf37 100644 --- a/python/simpler/worker_chip_orch_comm.py +++ b/python/simpler/worker_chip_orch_comm.py @@ -10,67 +10,25 @@ from __future__ import annotations -import struct from dataclasses import dataclass -from enum import IntEnum from typing import Any -from _task_interface import ( # pyright: ignore[reportMissingImports] - _worker_host_mapped_counter_notify, # noqa: F401 - _worker_host_mapped_counter_test, # noqa: F401 - _worker_host_mapped_counter_wait, # noqa: F401 - _worker_host_mapped_payload_read, # noqa: F401 - _worker_host_mapped_payload_write, # noqa: F401 - _worker_host_mapped_region_close, -) - +from .comm_provider import RegionPartLocalView, validate_independent_local_views from .comm_region import ( - CounterPart, - HostVmmCopyAccess, + MaterializationError, NotifyOp, - PayloadPart, - RegionPartSpan, + RegionInstanceState, SignalTestResult, WaitCmp, ) - -class WorkerChipRegionAccessProfile(IntEnum): - INVALID = 0 - ONBOARD_VMM = 1 - SIM_POSIX_SHM = 2 - - -_DESC = struct.Struct("<6Q") -_WORKER_CHIP_ORCH_REGION_DESC_SCALAR_COUNT = 6 -_CTRL_SHM_TOKEN_BYTES = 32 -_REGION_CREATE_REQUEST = struct.Struct(" int: - value = int(value) - if value < 0 or value > _UINT64_MAX: - raise ValueError("L3-L2 region layout overflowed uint64") - remainder = value % align - bump = 0 if remainder == 0 else align - remainder - result = value + bump - if result > _UINT64_MAX: - raise ValueError("L3-L2 region layout overflowed uint64") - return result - - -def _checked_add_u64(lhs: int, rhs: int) -> int: - result = int(lhs) + int(rhs) - if int(lhs) < 0 or int(rhs) < 0 or result > _UINT64_MAX: - raise ValueError("L3-L2 region layout overflowed uint64") - return result +WORKER_CHIP_ORCH_COMM_MAGIC = 0x4C334C32 +WORKER_CHIP_ORCH_COMM_ABI_MAJOR = 3 +WORKER_CHIP_ORCH_COMM_ABI_MINOR = 0 +_REGION_MAGIC_VERSION = ( + (WORKER_CHIP_ORCH_COMM_MAGIC << 32) | (WORKER_CHIP_ORCH_COMM_ABI_MAJOR << 16) | WORKER_CHIP_ORCH_COMM_ABI_MINOR +) @dataclass(frozen=True) @@ -93,115 +51,20 @@ def scalars(self) -> list[int]: ] -@dataclass(frozen=True) -class WorkerChipRegionCreateRequest: - magic_version: int - request_bytes: int - payload_bytes: int - counter_bytes: int - - def encode_into(self, buf: memoryview, offset: int = 0) -> None: - _REGION_CREATE_REQUEST.pack_into( - buf, - offset, - int(self.magic_version), - int(self.request_bytes), - int(self.payload_bytes), - int(self.counter_bytes), - ) - - -@dataclass(frozen=True) -class WorkerChipRegionCreateReply: - desc: WorkerChipOrchRegionDesc - access_profile: WorkerChipRegionAccessProfile - device_id: int - backing_shm: str - mapping_bytes: int - shareable_handle: int - - -@dataclass -class WorkerHostRegionMapping: - worker_id: int - region_id: int - access_profile: WorkerChipRegionAccessProfile - total_bytes: int - payload_offset: int - payload_bytes: int - counter_offset: int - counter_bytes: int - handle: Any - closed: bool = False - - def close(self) -> None: - if self.closed: - return - _worker_host_mapped_region_close(int(self.handle)) - self.closed = True - - -def decode_region_create_reply(buf: memoryview) -> WorkerChipRegionCreateReply: - fields = _REGION_CREATE_REPLY.unpack_from(buf, 0) - desc = WorkerChipOrchRegionDesc(*[int(v) for v in fields[:6]]) - access_profile = WorkerChipRegionAccessProfile(int(fields[6])) - device_id = int(fields[8]) - backing_shm = bytes(fields[9]).split(b"\x00", 1)[0].decode("utf-8", "strict") - return WorkerChipRegionCreateReply( - desc=desc, - access_profile=access_profile, - device_id=device_id, - backing_shm=backing_shm, - mapping_bytes=int(fields[10]), - shareable_handle=int(fields[11]), +def worker_chip_orch_region_desc_from_local_views( + provider_resource_id: int, payload_view: RegionPartLocalView, counter_view: RegionPartLocalView +) -> WorkerChipOrchRegionDesc: + payload_view, counter_view = validate_independent_local_views(payload_view, counter_view) + return WorkerChipOrchRegionDesc( + magic_version=_REGION_MAGIC_VERSION, + region_id=int(provider_resource_id), + payload_base=int(payload_view.local_base), + payload_bytes=int(payload_view.logical_bytes), + counter_base=int(counter_view.local_base), + counter_bytes=int(counter_view.logical_bytes), ) -def peek_region_create_reply_region_id(buf: memoryview) -> int: - """Raw-unpack just the region_id for the create rollback path. - - Must tolerate malformed replies: decode_region_create_reply raises on - unknown access_profile / non-UTF-8 backing_shm, but the L2 child has - already created the region by then, so the rollback release still needs - the id. Call this BEFORE decode_region_create_reply. - """ - return int(_REGION_CREATE_REPLY.unpack_from(buf, 0)[1]) - - -def validate_region_create_reply( - reply: WorkerChipRegionCreateReply, expected_access_profile: WorkerChipRegionAccessProfile -) -> tuple[int, int]: - desc = reply.desc - if desc.magic_version != _REGION_MAGIC_VERSION: - raise RuntimeError("create_worker_chip_region: reply magic_version is invalid") - if desc.region_id == 0: - raise RuntimeError("create_worker_chip_region: reply region_id must be nonzero") - if reply.access_profile != expected_access_profile: - raise RuntimeError( - f"create_worker_chip_region: reply access_profile must be {expected_access_profile.name.lower()}" - ) - if desc.payload_bytes <= 0: - raise RuntimeError("create_worker_chip_region: reply payload_bytes must be positive") - if desc.counter_bytes <= 0 or desc.counter_bytes % 4 != 0: - raise RuntimeError("create_worker_chip_region: reply counter_bytes must be positive and a multiple of 4") - counter_offset = _align_up(desc.payload_bytes, _REGION_LAYOUT_ALIGNMENT) - total_bytes = _checked_add_u64(counter_offset, desc.counter_bytes) - expected_counter_base = _checked_add_u64(desc.payload_base, counter_offset) - if desc.counter_base != expected_counter_base: - raise RuntimeError("create_worker_chip_region: reply counter_base does not match fixed region layout") - if desc.counter_base % _REGION_LAYOUT_ALIGNMENT != 0: - raise RuntimeError("create_worker_chip_region: reply counter_base must be 64-byte aligned") - if reply.access_profile == WorkerChipRegionAccessProfile.SIM_POSIX_SHM and reply.mapping_bytes != total_bytes: - profile = reply.access_profile.name.lower() - raise RuntimeError(f"create_worker_chip_region: {profile} reply mapping_bytes does not match descriptor layout") - if reply.access_profile == WorkerChipRegionAccessProfile.ONBOARD_VMM: - if reply.mapping_bytes < total_bytes: - raise RuntimeError( - "create_worker_chip_region: onboard_vmm reply mapping_bytes is smaller than descriptor layout" - ) - return counter_offset, total_bytes - - class WorkerChipOrchCounter: def __init__(self, region: WorkerChipOrchRegion, offset: int) -> None: self._region = region @@ -232,34 +95,12 @@ def wait(self, cmp_value: int, cmp: WaitCmp, timeout: float) -> int: class WorkerChipOrchRegion: - def __init__( - self, - owner: Any, - worker_id: int, - desc: WorkerChipOrchRegionDesc, - worker_host_mapping: WorkerHostRegionMapping, - ) -> None: + def __init__(self, owner: Any, instance: Any, desc: WorkerChipOrchRegionDesc) -> None: self._owner = owner - self._worker_id = int(worker_id) + self._instance = instance + self._worker_id = int(instance.worker_id) self._descriptor = desc - self._worker_host_mapping = worker_host_mapping - access = HostVmmCopyAccess.from_mapping(worker_host_mapping) - self._payload_part = PayloadPart( - RegionPartSpan( - offset=int(worker_host_mapping.payload_offset), nbytes=int(worker_host_mapping.payload_bytes) - ), - access, - ) - self._counter_part = CounterPart( - RegionPartSpan( - offset=int(worker_host_mapping.counter_offset), nbytes=int(worker_host_mapping.counter_bytes) - ), - access, - ) self._released = False - self._chip_release_committed = False - self._poisoned = False - self._expired = False @property def descriptor(self) -> WorkerChipOrchRegionDesc: @@ -271,7 +112,10 @@ def region_id(self) -> int: @property def expired(self) -> bool: - return self._expired + try: + return self._instance.state is not RegionInstanceState.LIVE + except MaterializationError: + return True def _validate_host_buffer(self, buffer: Any) -> None: self._owner._validate_worker_chip_orch_comm_host_buffer(buffer) @@ -283,23 +127,23 @@ def descriptor_scalars(self) -> list[int]: def payload_write(self, offset: int, host_buffer: Any, nbytes: int | None = None) -> None: self._ensure_live() try: - self._payload_part.write(offset, host_buffer, nbytes) - except Exception: - self._poison() + self._instance.payload_write(offset, host_buffer, nbytes) + except Exception as exc: + self._remember_data_plane_error(exc) raise def payload_read(self, offset: int, host_buffer: Any, nbytes: int | None = None) -> None: self._ensure_live() try: - self._payload_part.read(offset, host_buffer, nbytes) - except Exception: - self._poison() + self._instance.payload_read(offset, host_buffer, nbytes) + except Exception as exc: + self._remember_data_plane_error(exc) raise def counter(self, offset: int) -> WorkerChipOrchCounter: self._ensure_live() offset = int(offset) - self._counter_part.counter(offset) + self._instance.counter(offset) return WorkerChipOrchCounter(self, offset) def free(self) -> None: @@ -307,47 +151,38 @@ def free(self) -> None: return self._released = True - def _expire(self) -> None: - self._expired = True - - def _poison(self) -> None: - self._poisoned = True + def _remember_data_plane_error(self, error: BaseException) -> None: + if self._instance.data_plane_error is None: + self._instance._data_plane_error = error def _ensure_live(self) -> None: - if self._expired: + if self.expired: raise RuntimeError(f"L3-L2 region {self.region_id} expired after orchestration run") if self._released: raise RuntimeError(f"L3-L2 region {self.region_id} has been released") - if self._poisoned: + if self._instance.data_plane_error is not None: raise RuntimeError(f"L3-L2 region {self.region_id} is poisoned") - def _close_worker_host_mapping(self) -> None: - try: - self._worker_host_mapping.close() - except Exception: - self._poison() - raise - def _direct_counter_notify(self, offset: int, value: int, op: NotifyOp) -> None: try: - self._counter_part.notify(offset, int(value), NotifyOp(op)) - except Exception: - self._poison() + self._instance.counter(offset).notify(int(value), NotifyOp(op)) + except Exception as exc: + self._remember_data_plane_error(exc) raise def _direct_counter_test(self, offset: int, cmp_value: int, cmp: WaitCmp) -> SignalTestResult: try: - return self._counter_part.test(offset, int(cmp_value), WaitCmp(cmp)) - except Exception: - self._poison() + return self._instance.counter(offset).test(int(cmp_value), WaitCmp(cmp)) + except Exception as exc: + self._remember_data_plane_error(exc) raise def _direct_counter_wait(self, offset: int, cmp_value: int, cmp: WaitCmp, timeout_ns: int) -> int: timeout = float(timeout_ns) / 1_000_000_000 try: - return self._counter_part.wait(offset, int(cmp_value), WaitCmp(cmp), timeout) + return self._instance.counter(offset).wait(int(cmp_value), WaitCmp(cmp), timeout) except TimeoutError: raise - except Exception: - self._poison() + except Exception as exc: + self._remember_data_plane_error(exc) raise diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index 629a4beda8..ab741eee3e 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -126,13 +126,12 @@ class Worker { void control_comm_init(int worker_id, const std::string &request_shm_name) { manager_.control_comm_init(worker_id, request_shm_name.c_str()); } - void control_worker_chip_region_create( - int worker_id, const std::string &request_shm_name, const std::string &reply_shm_name - ) { - manager_.control_worker_chip_region_create(worker_id, request_shm_name.c_str(), reply_shm_name.c_str()); + void + control_region_allocate(int worker_id, const std::string &request_shm_name, const std::string &reply_shm_name) { + manager_.control_region_allocate(worker_id, request_shm_name.c_str(), reply_shm_name.c_str()); } - void control_worker_chip_region_release(int worker_id, uint64_t region_id) { - manager_.control_worker_chip_region_release(worker_id, region_id); + void control_region_release(int worker_id, const std::string &request_shm_name, const std::string &reply_shm_name) { + manager_.control_region_release(worker_id, request_shm_name.c_str(), reply_shm_name.c_str()); } ControlResult diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 70be691d45..c473b19a08 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -193,11 +193,11 @@ void WorkerEndpoint::control_alloc_domain(const char *, const char *) { } void WorkerEndpoint::control_release_domain(const char *) { throw_unsupported_control("control_release_domain"); } void WorkerEndpoint::control_comm_init(const char *) { throw_unsupported_control("control_comm_init"); } -void WorkerEndpoint::control_worker_chip_region_create(const char *, const char *) { - throw_unsupported_control("control_worker_chip_region_create"); +void WorkerEndpoint::control_region_allocate(const char *, const char *) { + throw_unsupported_control("control_region_allocate"); } -void WorkerEndpoint::control_worker_chip_region_release(uint64_t) { - throw_unsupported_control("control_worker_chip_region_release"); +void WorkerEndpoint::control_region_release(const char *, const char *) { + throw_unsupported_control("control_region_release"); } void WorkerEndpoint::submit_progress(Ring *, const WorkerDispatch &) { @@ -1421,21 +1421,26 @@ void LocalMailboxEndpoint::control_comm_init(const char *request_shm_name) { run_control_command("control_comm_init"); } -void LocalMailboxEndpoint::control_worker_chip_region_create(const char *request_shm_name, const char *reply_shm_name) { +void LocalMailboxEndpoint::control_region_allocate(const char *request_shm_name, const char *reply_shm_name) { if (!request_shm_name || !*request_shm_name || !reply_shm_name || !*reply_shm_name) { - throw std::runtime_error("control_worker_chip_region_create: request and reply shm names must be non-empty"); + throw std::runtime_error("control_region_allocate: request and reply shm names must be non-empty"); } std::lock_guard lk(mailbox_mu_); - uint64_t sub_cmd = CTRL_WORKER_CHIP_REGION_CREATE; + uint64_t sub_cmd = CTRL_REGION_ALLOCATE; std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); write_shm_name_pair(mbox(), request_shm_name, reply_shm_name); - run_control_command("control_worker_chip_region_create"); + run_control_command("control_region_allocate"); } -void LocalMailboxEndpoint::control_worker_chip_region_release(uint64_t region_id) { +void LocalMailboxEndpoint::control_region_release(const char *request_shm_name, const char *reply_shm_name) { + if (!request_shm_name || !*request_shm_name || !reply_shm_name || !*reply_shm_name) { + throw std::runtime_error("control_region_release: request and reply shm names must be non-empty"); + } std::lock_guard lk(mailbox_mu_); - write_control_args(mbox(), CTRL_WORKER_CHIP_REGION_RELEASE, region_id); - run_control_command("control_worker_chip_region_release"); + uint64_t sub_cmd = CTRL_REGION_RELEASE; + std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); + write_shm_name_pair(mbox(), request_shm_name, reply_shm_name); + run_control_command("control_region_release"); } uint64_t WorkerThread::control_malloc(size_t size) { @@ -1578,14 +1583,14 @@ void WorkerThread::control_comm_init(const char *request_shm_name) { endpoint_->control_comm_init(request_shm_name); } -void WorkerThread::control_worker_chip_region_create(const char *request_shm_name, const char *reply_shm_name) { - if (!endpoint_) throw std::runtime_error("control_worker_chip_region_create: null endpoint"); - endpoint_->control_worker_chip_region_create(request_shm_name, reply_shm_name); +void WorkerThread::control_region_allocate(const char *request_shm_name, const char *reply_shm_name) { + if (!endpoint_) throw std::runtime_error("control_region_allocate: null endpoint"); + endpoint_->control_region_allocate(request_shm_name, reply_shm_name); } -void WorkerThread::control_worker_chip_region_release(uint64_t region_id) { - if (!endpoint_) throw std::runtime_error("control_worker_chip_region_release: null endpoint"); - endpoint_->control_worker_chip_region_release(region_id); +void WorkerThread::control_region_release(const char *request_shm_name, const char *reply_shm_name) { + if (!endpoint_) throw std::runtime_error("control_region_release: null endpoint"); + endpoint_->control_region_release(request_shm_name, reply_shm_name); } bool WorkerManager::any_busy() const { @@ -1730,22 +1735,20 @@ void WorkerManager::control_comm_init(int worker_id, const char *request_shm_nam wt->control_comm_init(request_shm_name); } -void WorkerManager::control_worker_chip_region_create( - int worker_id, const char *request_shm_name, const char *reply_shm_name -) { +void WorkerManager::control_region_allocate(int worker_id, const char *request_shm_name, const char *reply_shm_name) { auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (wt == nullptr) { - throw std::runtime_error("control_worker_chip_region_create: invalid worker_id " + std::to_string(worker_id)); + throw std::runtime_error("control_region_allocate: invalid worker_id " + std::to_string(worker_id)); } - wt->control_worker_chip_region_create(request_shm_name, reply_shm_name); + wt->control_region_allocate(request_shm_name, reply_shm_name); } -void WorkerManager::control_worker_chip_region_release(int worker_id, uint64_t region_id) { +void WorkerManager::control_region_release(int worker_id, const char *request_shm_name, const char *reply_shm_name) { auto *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (wt == nullptr) { - throw std::runtime_error("control_worker_chip_region_release: invalid worker_id " + std::to_string(worker_id)); + throw std::runtime_error("control_region_release: invalid worker_id " + std::to_string(worker_id)); } - wt->control_worker_chip_region_release(region_id); + wt->control_region_release(request_shm_name, reply_shm_name); } ControlResult WorkerManager::control_digest_only( diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 39483fd038..4bb25cf416 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -212,14 +212,14 @@ static constexpr uint64_t CTRL_RELEASE_DOMAIN = 8; static constexpr uint64_t CTRL_COMM_INIT = 9; static constexpr uint64_t CTRL_PY_REGISTER = 10; static constexpr uint64_t CTRL_PY_UNREGISTER = 11; -static constexpr uint64_t CTRL_WORKER_CHIP_REGION_CREATE = 16; -static constexpr uint64_t CTRL_WORKER_CHIP_REGION_RELEASE = 17; +static constexpr uint64_t CTRL_REGION_ALLOCATE = 16; +static constexpr uint64_t CTRL_REGION_RELEASE = 17; // Query a chip child's MemoryAllocator-committed HBM (bytes). The child writes // the value to CTRL_OFF_RESULT; the parent sums across children for L3. static constexpr uint64_t CTRL_COMMITTED_DEVICE_MEMORY = 18; // Control args occupy the base frame's config-sized region: -// offset 16: uint64 arg0 (size for malloc/register; ptr for free; region id for region release) +// offset 16: uint64 arg0 (size for malloc/register; ptr for free) // offset 40: uint64 result (returned ptr from malloc) static constexpr ptrdiff_t CTRL_OFF_ARG0 = 16; static constexpr ptrdiff_t CTRL_OFF_RESULT = 40; @@ -384,8 +384,8 @@ class WorkerEndpoint { virtual void control_alloc_domain(const char *request_shm_name, const char *reply_shm_name); virtual void control_release_domain(const char *request_shm_name); virtual void control_comm_init(const char *request_shm_name); - virtual void control_worker_chip_region_create(const char *request_shm_name, const char *reply_shm_name); - virtual void control_worker_chip_region_release(uint64_t region_id); + virtual void control_region_allocate(const char *request_shm_name, const char *reply_shm_name); + virtual void control_region_release(const char *request_shm_name, const char *reply_shm_name); }; class LocalMailboxEndpoint : public WorkerEndpoint { @@ -445,8 +445,8 @@ class LocalMailboxEndpoint : public WorkerEndpoint { void control_alloc_domain(const char *request_shm_name, const char *reply_shm_name) override; void control_release_domain(const char *request_shm_name) override; void control_comm_init(const char *request_shm_name) override; - void control_worker_chip_region_create(const char *request_shm_name, const char *reply_shm_name) override; - void control_worker_chip_region_release(uint64_t region_id) override; + void control_region_allocate(const char *request_shm_name, const char *reply_shm_name) override; + void control_region_release(const char *request_shm_name, const char *reply_shm_name) override; private: WorkerEndpointCaps caps_; @@ -621,8 +621,8 @@ class WorkerThread { // Lazy comm_init driver — payload shm carries (rank, nranks, rootinfo_path). // Caller dispatches in parallel to every chip; child runs cw.comm_init. void control_comm_init(const char *request_shm_name); - void control_worker_chip_region_create(const char *request_shm_name, const char *reply_shm_name); - void control_worker_chip_region_release(uint64_t region_id); + void control_region_allocate(const char *request_shm_name, const char *reply_shm_name); + void control_region_release(const char *request_shm_name, const char *reply_shm_name); private: enum class SubmitDispatchResult : uint8_t { @@ -721,8 +721,8 @@ class WorkerManager { void control_alloc_domain(int worker_id, const char *request_shm_name, const char *reply_shm_name); void control_release_domain(int worker_id, const char *request_shm_name); void control_comm_init(int worker_id, const char *request_shm_name); - void control_worker_chip_region_create(int worker_id, const char *request_shm_name, const char *reply_shm_name); - void control_worker_chip_region_release(int worker_id, uint64_t region_id); + void control_region_allocate(int worker_id, const char *request_shm_name, const char *reply_shm_name); + void control_region_release(int worker_id, const char *request_shm_name, const char *reply_shm_name); ControlResult control_digest_only(WorkerType type, int worker_id, uint64_t sub_cmd, const uint8_t *digest, double timeout_s); std::vector control_payload( diff --git a/src/common/platform/include/aicpu/region_instance_view.h b/src/common/platform/include/aicpu/region_instance_view.h new file mode 100644 index 0000000000..0e8d71b650 --- /dev/null +++ b/src/common/platform/include/aicpu/region_instance_view.h @@ -0,0 +1,380 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include +#include + +#include "aicpu/cache_maintenance.h" +#include "aicpu/device_time.h" +#include "common/region_instance_semantics.h" + +struct RegionSignalTestResult { + bool matched; + int32_t observed; +}; + +struct RegionPayloadView { + uint64_t local_addr; + uint64_t nbytes; +}; + +enum class RegionViewErrorKind : uint32_t { + NONE = 0, + INVALID_VIEW = 1, + OUT_OF_BOUNDS = 2, + INVALID_ENUM = 3, + TIMEOUT = 4, + ISSUED_FAILURE = 5, +}; + +enum class RegionViewOp : uint32_t { + CONSTRUCT = 1, + PAYLOAD_READ = 2, + PAYLOAD_WRITE = 3, + NOTIFY = 4, + TEST = 5, + WAIT = 6, + COUNTER_ADDR = 7, +}; + +struct RegionViewError { + RegionViewErrorKind kind; + RegionViewOp op; + int32_t observed; + char message[256]; +}; + +struct PayloadLocalOps { + bool invalidate(const void *addr, size_t nbytes) const { + cache_invalidate_range(addr, nbytes); + return true; + } + + void flush(const void *addr, size_t nbytes) const { cache_flush_range(addr, nbytes); } + + bool copy_bytes(void *dst, const void *src, size_t nbytes) const { + if (dst == nullptr || src == nullptr) { + return false; + } + memcpy(dst, src, nbytes); + return true; + } +}; + +struct CounterLocalOps { + bool load_i32(uint64_t addr, int32_t &out) const { + volatile int32_t *ptr = reinterpret_cast(static_cast(addr)); + cache_invalidate_range(reinterpret_cast(static_cast(addr)), sizeof(int32_t)); + out = *ptr; + return true; + } + + bool store_i32(uint64_t addr, int32_t value) const { + volatile int32_t *ptr = reinterpret_cast(static_cast(addr)); + *ptr = value; + cache_flush_range(reinterpret_cast(static_cast(addr)), sizeof(int32_t)); + return true; + } +}; + +struct LocalMemoryOps : PayloadLocalOps, + CounterLocalOps {}; + +struct RegionViewState { + bool sticky_failed; + RegionViewError last_error; +}; + +inline void region_view_copy_message(char *dst, size_t dst_size, const char *message) { + if (dst == nullptr || dst_size == 0) { + return; + } + const char *src = message == nullptr ? "" : message; + size_t n = strnlen(src, dst_size - 1); + memcpy(dst, src, n); + dst[n] = '\0'; +} + +template +class RegionInstanceViewImpl { +public: + class PayloadPart { + public: + bool read(uint64_t offset, uint64_t nbytes, RegionPayloadView &out) { + return view_->payload_read(offset, nbytes, out); + } + + bool write(uint64_t offset, const void *src, uint64_t nbytes) { + return view_->payload_write(offset, src, nbytes); + } + + RegionPartLocalSpan span() const { return view_->payload_span(); } + + private: + friend class RegionInstanceViewImpl; + explicit PayloadPart(RegionInstanceViewImpl *view) : + view_(view) {} + RegionInstanceViewImpl *view_; + }; + + class CounterPart { + public: + bool addr(uint64_t offset, uint64_t &out) { return view_->counter_addr(offset, out); } + + bool notify(uint64_t offset, int32_t value, RegionNotifyOp op) { return view_->notify(offset, value, op); } + + bool test(uint64_t offset, int32_t cmp_value, RegionWaitCmp cmp, RegionSignalTestResult &out) { + return view_->test(offset, cmp_value, cmp, out); + } + + bool wait(uint64_t offset, int32_t cmp_value, RegionWaitCmp cmp, uint64_t timeout_ns, int32_t &observed) { + return view_->wait(offset, cmp_value, cmp, timeout_ns, observed); + } + + bool contains_addr(uint64_t counter_addr) const { return view_->counter_contains(counter_addr); } + + RegionPartLocalSpan span() const { return view_->counter_span(); } + + private: + friend class RegionInstanceViewImpl; + explicit CounterPart(RegionInstanceViewImpl *view) : + view_(view) {} + RegionInstanceViewImpl *view_; + }; + + RegionInstanceViewImpl() { + mark_sticky(RegionViewErrorKind::INVALID_VIEW, RegionViewOp::CONSTRUCT, 0, "uninitialized view"); + } + + RegionInstanceViewImpl(RegionPartLocalSpan payload, RegionPartLocalSpan counter) { + if (!region_validate_independent_spans(payload, counter)) { + mark_sticky( + RegionViewErrorKind::INVALID_VIEW, RegionViewOp::CONSTRUCT, 0, "invalid independent local spans" + ); + return; + } + payload_span_ = payload; + counter_span_ = counter; + } + + bool failed() const { return state_.sticky_failed; } + + const RegionViewError &error() const { return state_.last_error; } + + RegionPartLocalSpan payload_span() const { return payload_span_; } + + RegionPartLocalSpan counter_span() const { return counter_span_; } + + PayloadPart payload() { return PayloadPart(this); } + + CounterPart counter() { return CounterPart(this); } + + bool payload_read(uint64_t offset, uint64_t nbytes, RegionPayloadView &out) { + out = RegionPayloadView{0, 0}; + if (state_.sticky_failed) { + return false; + } + if (!validate_payload_range(RegionViewOp::PAYLOAD_READ, offset, nbytes)) { + return false; + } + uint64_t addr = payload_span_.base + offset; + if (!ops_.invalidate( + reinterpret_cast(static_cast(addr)), static_cast(nbytes) + )) { + mark_sticky( + RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::PAYLOAD_READ, 0, "payload invalidate failed" + ); + return false; + } + out = RegionPayloadView{addr, nbytes}; + return true; + } + + bool payload_write(uint64_t offset, const void *src, uint64_t nbytes) { + if (state_.sticky_failed) { + return false; + } + if (src == nullptr) { + record(RegionViewErrorKind::OUT_OF_BOUNDS, RegionViewOp::PAYLOAD_WRITE, 0, false, "null payload source"); + return false; + } + if (!validate_payload_range(RegionViewOp::PAYLOAD_WRITE, offset, nbytes)) { + return false; + } + void *dst = reinterpret_cast(static_cast(payload_span_.base + offset)); + if (!ops_.copy_bytes(dst, src, static_cast(nbytes))) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::PAYLOAD_WRITE, 0, "payload copy failed"); + return false; + } + ops_.flush(dst, static_cast(nbytes)); + return true; + } + + bool counter_addr(uint64_t offset, uint64_t &out) { + out = 0; + if (state_.sticky_failed) { + return false; + } + uint64_t addr = 0; + if (!resolve_counter_addr(RegionViewOp::COUNTER_ADDR, offset, addr, "counter offset is out of bounds")) { + return false; + } + out = addr; + return true; + } + + bool counter_contains(uint64_t counter_addr) const { + if (state_.sticky_failed) { + return false; + } + return counter_addr_in_span(counter_addr); + } + + bool notify(uint64_t offset, int32_t value, RegionNotifyOp op) { + if (state_.sticky_failed) { + return false; + } + uint64_t addr = 0; + if (!resolve_counter_addr(RegionViewOp::NOTIFY, offset, addr, "invalid counter address")) { + return false; + } + if (!region_valid_notify_op(op)) { + record(RegionViewErrorKind::INVALID_ENUM, RegionViewOp::NOTIFY, 0, false, "invalid notify operation"); + return false; + } + if (op == RegionNotifyOp::Set) { + if (!ops_.store_i32(addr, value)) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::NOTIFY, 0, "counter store failed"); + return false; + } + return true; + } + int32_t observed = 0; + if (!ops_.load_i32(addr, observed)) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::NOTIFY, 0, "counter load failed"); + return false; + } + uint32_t wrapped = static_cast(observed) + static_cast(value); + int32_t next = 0; + memcpy(&next, &wrapped, sizeof(next)); + if (!ops_.store_i32(addr, next)) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::NOTIFY, 0, "counter store failed"); + return false; + } + return true; + } + + bool test(uint64_t offset, int32_t cmp_value, RegionWaitCmp cmp, RegionSignalTestResult &out) { + out = RegionSignalTestResult{false, 0}; + if (state_.sticky_failed) { + return false; + } + uint64_t addr = 0; + if (!resolve_counter_addr(RegionViewOp::TEST, offset, addr, "invalid counter address")) { + return false; + } + if (!region_valid_wait_cmp(cmp)) { + record(RegionViewErrorKind::INVALID_ENUM, RegionViewOp::TEST, 0, false, "invalid wait comparison"); + return false; + } + int32_t observed = 0; + if (!ops_.load_i32(addr, observed)) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::TEST, 0, "counter load failed"); + return false; + } + out = RegionSignalTestResult{region_compare_counter(observed, cmp_value, cmp), observed}; + return true; + } + + bool wait(uint64_t offset, int32_t cmp_value, RegionWaitCmp cmp, uint64_t timeout_ns, int32_t &observed) { + observed = 0; + if (state_.sticky_failed) { + return false; + } + uint64_t addr = 0; + if (!resolve_counter_addr(RegionViewOp::WAIT, offset, addr, "invalid counter address")) { + return false; + } + if (!region_valid_wait_cmp(cmp)) { + record(RegionViewErrorKind::INVALID_ENUM, RegionViewOp::WAIT, 0, false, "invalid wait comparison"); + return false; + } + uint64_t start = device_time_now_ticks(); + uint64_t frequency_hz = device_time_frequency_hz(); + while (true) { + int32_t current = 0; + if (!ops_.load_i32(addr, current)) { + mark_sticky(RegionViewErrorKind::ISSUED_FAILURE, RegionViewOp::WAIT, 0, "counter load failed"); + return false; + } + observed = current; + if (region_compare_counter(current, cmp_value, cmp)) { + return true; + } + uint64_t now = device_time_now_ticks(); + if (timeout_ns == 0 || sys_cnt_elapsed_ns(start, now, frequency_hz) >= timeout_ns) { + record(RegionViewErrorKind::TIMEOUT, RegionViewOp::WAIT, current, false, "wait timed out"); + return false; + } + } + } + +private: + void record(RegionViewErrorKind kind, RegionViewOp op, int32_t observed, bool sticky, const char *message) { + if (state_.sticky_failed) { + return; + } + state_.last_error = RegionViewError{kind, op, observed, ""}; + region_view_copy_message(state_.last_error.message, sizeof(state_.last_error.message), message); + state_.sticky_failed = sticky; + } + + void mark_sticky(RegionViewErrorKind kind, RegionViewOp op, int32_t observed, const char *message) { + record(kind, op, observed, true, message); + } + + bool validate_payload_range(RegionViewOp op, uint64_t offset, uint64_t nbytes) { + if (!region_validate_payload_range(offset, nbytes, payload_span_.logical_bytes)) { + record(RegionViewErrorKind::OUT_OF_BOUNDS, op, 0, false, "payload range is out of bounds"); + return false; + } + return true; + } + + bool counter_addr_in_span(uint64_t counter_addr) const { + return region_counter_addr_in_span(counter_span_, counter_addr); + } + + bool resolve_counter_addr(RegionViewOp op, uint64_t offset, uint64_t &addr, const char *message) { + addr = 0; + if ((offset % REGION_COUNTER_LOGICAL_ALIGNMENT) != 0 || region_add_overflows(counter_span_.base, offset)) { + record(RegionViewErrorKind::OUT_OF_BOUNDS, op, 0, false, message); + return false; + } + uint64_t resolved = counter_span_.base + offset; + if (!counter_addr_in_span(resolved)) { + record(RegionViewErrorKind::OUT_OF_BOUNDS, op, 0, false, message); + return false; + } + addr = resolved; + return true; + } + + RegionPartLocalSpan payload_span_{}; + RegionPartLocalSpan counter_span_{}; + Ops ops_{}; + RegionViewState state_{}; +}; + +using RegionInstanceView = RegionInstanceViewImpl; diff --git a/src/common/platform/include/aicpu/worker_chip_orch_endpoint.h b/src/common/platform/include/aicpu/worker_chip_orch_endpoint.h index b206e83c79..4580e403b9 100644 --- a/src/common/platform/include/aicpu/worker_chip_orch_endpoint.h +++ b/src/common/platform/include/aicpu/worker_chip_orch_endpoint.h @@ -11,12 +11,11 @@ #pragma once +#include #include #include -#include -#include "aicpu/cache_maintenance.h" -#include "aicpu/device_time.h" +#include "aicpu/region_instance_view.h" #include "common/worker_chip_orch_comm.h" struct WorkerChipOrchPayloadView { @@ -30,6 +29,7 @@ enum class WorkerChipEndpointErrorKind : uint32_t { OUT_OF_BOUNDS = 2, SIGNAL_TIMEOUT = 3, SIGNAL_PROTOCOL = 4, + LOCAL_OPERATION = 5, }; enum class WorkerChipEndpointOp : uint32_t { @@ -58,9 +58,62 @@ inline const char *worker_chip_endpoint_op_to_string(WorkerChipEndpointOp op) { return "signal_test"; case WorkerChipEndpointOp::SIGNAL_WAIT: return "signal_wait"; - default: - return "unknown"; } + return "unknown"; +} + +inline bool worker_chip_notify_op_to_region(WorkerChipOrchNotifyOp op, RegionNotifyOp &out) { + switch (op) { + case WorkerChipOrchNotifyOp::Set: + out = RegionNotifyOp::Set; + return true; + case WorkerChipOrchNotifyOp::Add: + out = RegionNotifyOp::Add; + return true; + } + return false; +} + +inline bool worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp cmp, RegionWaitCmp &out) { + switch (cmp) { + case WorkerChipOrchWaitCmp::EQ: + out = RegionWaitCmp::EQ; + return true; + case WorkerChipOrchWaitCmp::NE: + out = RegionWaitCmp::NE; + return true; + case WorkerChipOrchWaitCmp::GT: + out = RegionWaitCmp::GT; + return true; + case WorkerChipOrchWaitCmp::GE: + out = RegionWaitCmp::GE; + return true; + case WorkerChipOrchWaitCmp::LT: + out = RegionWaitCmp::LT; + return true; + case WorkerChipOrchWaitCmp::LE: + out = RegionWaitCmp::LE; + return true; + } + return false; +} + +inline WorkerChipEndpointErrorKind worker_chip_map_view_error_kind(RegionViewErrorKind kind) { + switch (kind) { + case RegionViewErrorKind::INVALID_VIEW: + return WorkerChipEndpointErrorKind::BAD_DESCRIPTOR; + case RegionViewErrorKind::OUT_OF_BOUNDS: + return WorkerChipEndpointErrorKind::OUT_OF_BOUNDS; + case RegionViewErrorKind::INVALID_ENUM: + return WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL; + case RegionViewErrorKind::TIMEOUT: + return WorkerChipEndpointErrorKind::SIGNAL_TIMEOUT; + case RegionViewErrorKind::ISSUED_FAILURE: + return WorkerChipEndpointErrorKind::LOCAL_OPERATION; + case RegionViewErrorKind::NONE: + return WorkerChipEndpointErrorKind::OUT_OF_BOUNDS; + } + return WorkerChipEndpointErrorKind::OUT_OF_BOUNDS; } struct WorkerChipEndpointError { @@ -73,123 +126,103 @@ struct WorkerChipEndpointError { char message[256]; }; -class WorkerChipOrchEndpoint { +template +class WorkerChipOrchEndpointImpl { public: - explicit WorkerChipOrchEndpoint(const WorkerChipOrchRegionDesc &desc) : + explicit WorkerChipOrchEndpointImpl(const WorkerChipOrchRegionDesc &desc) : desc_(desc) { - if (worker_chip_orch_comm::validate_desc(desc_) != WorkerChipOrchCommValidationError::OK) { - set_error( - WorkerChipEndpointErrorKind::BAD_DESCRIPTOR, WorkerChipEndpointOp::INIT, desc_.region_id, 0, 0, - "invalid descriptor" - ); - } + install_view_or_error(); } - WorkerChipOrchEndpoint(const uint64_t *scalars, size_t scalar_count) { + WorkerChipOrchEndpointImpl(const uint64_t *scalars, size_t scalar_count) { WorkerChipOrchCommValidationError error = WorkerChipOrchCommValidationError::OK; if (!worker_chip_orch_comm::decode_desc(scalars, scalar_count, &desc_, &error)) { uint64_t region_id = scalar_count > 1 && scalars != nullptr ? scalars[1] : 0; set_error( WorkerChipEndpointErrorKind::BAD_DESCRIPTOR, WorkerChipEndpointOp::INIT, region_id, 0, 0, - "invalid descriptor scalars" + "invalid compatibility descriptor" ); + return; } + install_view_or_error(); } const WorkerChipEndpointError &error() const { return error_; } const WorkerChipOrchRegionDesc &descriptor() const { return desc_; } + const RegionInstanceViewImpl &view() const { + static const RegionInstanceViewImpl kInvalid; + return view_.has_value() ? *view_ : kInvalid; + } + bool counter_addr(uint64_t offset, uint64_t &out_addr) { out_addr = 0; - if (has_error()) { + if (!ready()) { return false; } - if (worker_chip_orch_comm_add_overflows(desc_.counter_base, offset)) { - set_error( - WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, WorkerChipEndpointOp::COUNTER_ADDR, desc_.region_id, 0, 0, - "counter offset is out of bounds" - ); - return false; + uint64_t reported = 0; + if (!worker_chip_orch_comm_add_overflows(desc_.counter_base, offset)) { + reported = desc_.counter_base + offset; } - uint64_t addr = desc_.counter_base + offset; - if (!validate_counter_addr_for_op( - WorkerChipEndpointOp::COUNTER_ADDR, addr, 0, 0, "counter offset is out of bounds" - )) { + if (!view_->counter_addr(offset, out_addr)) { + adopt_view_error(WorkerChipEndpointOp::COUNTER_ADDR, reported, 0); return false; } - out_addr = addr; return true; } - bool validate_counter_addr(uint64_t counter_addr) const { - return worker_chip_orch_comm::validate_counter_addr(desc_, counter_addr) == - WorkerChipOrchCommValidationError::OK; - } + bool validate_counter_addr(uint64_t counter_addr) const { return ready() && view_->counter_contains(counter_addr); } bool payload_read(uint64_t offset, uint64_t nbytes, WorkerChipOrchPayloadView &out) { out = WorkerChipOrchPayloadView{0, 0}; - if (has_error()) { + if (!ready()) { return false; } - if (!validate_payload_range(WorkerChipEndpointOp::PAYLOAD_READ, offset, nbytes)) { + RegionPayloadView inner{}; + if (!view_->payload_read(offset, nbytes, inner)) { + adopt_view_error(WorkerChipEndpointOp::PAYLOAD_READ, 0, 0); return false; } - uint64_t gm_addr = desc_.payload_base + offset; - cache_invalidate_range( - reinterpret_cast(static_cast(gm_addr)), static_cast(nbytes) - ); - out = WorkerChipOrchPayloadView{gm_addr, nbytes}; + out = WorkerChipOrchPayloadView{inner.local_addr, inner.nbytes}; return true; } bool payload_write(uint64_t offset, const void *src, uint64_t nbytes) { - if (has_error()) { - return false; - } - if (src == nullptr) { - set_error( - WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, WorkerChipEndpointOp::PAYLOAD_WRITE, desc_.region_id, 0, 0, - "null payload source" - ); + if (!ready()) { return false; } - if (!validate_payload_range(WorkerChipEndpointOp::PAYLOAD_WRITE, offset, nbytes)) { + if (!view_->payload_write(offset, src, nbytes)) { + adopt_view_error(WorkerChipEndpointOp::PAYLOAD_WRITE, 0, 0); return false; } - void *dst = reinterpret_cast(static_cast(desc_.payload_base + offset)); - memcpy(dst, src, static_cast(nbytes)); - cache_flush_range(dst, static_cast(nbytes)); return true; } bool signal_notify(uint64_t counter_addr, int32_t value, WorkerChipOrchNotifyOp op) { - if (has_error()) { + if (!ready()) { return false; } - if (!validate_counter_addr_for_op( - WorkerChipEndpointOp::SIGNAL_NOTIFY, counter_addr, value, 0, "invalid counter address" - )) { + uint64_t offset = 0; + if (!absolute_to_counter_offset(counter_addr, offset)) { + set_error( + WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, WorkerChipEndpointOp::SIGNAL_NOTIFY, desc_.region_id, + counter_addr, value, "invalid counter address" + ); return false; } - if (!worker_chip_orch_comm::valid_notify_op(op)) { + RegionNotifyOp region_op{}; + if (!worker_chip_notify_op_to_region(op, region_op)) { set_error( WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL, WorkerChipEndpointOp::SIGNAL_NOTIFY, desc_.region_id, counter_addr, value, "invalid notify operation" ); return false; } - - volatile int32_t *counter = counter_ptr(counter_addr); - if (op == WorkerChipOrchNotifyOp::Set) { - *counter = value; - } else { - cache_invalidate_range( - reinterpret_cast(static_cast(counter_addr)), sizeof(*counter) - ); - *counter = static_cast(*counter + value); + if (!view_->notify(offset, value, region_op)) { + adopt_view_error(WorkerChipEndpointOp::SIGNAL_NOTIFY, counter_addr, value); + return false; } - cache_flush_range(reinterpret_cast(static_cast(counter_addr)), sizeof(*counter)); return true; } @@ -197,24 +230,31 @@ class WorkerChipOrchEndpoint { uint64_t counter_addr, int32_t cmp_value, WorkerChipOrchWaitCmp cmp, WorkerChipOrchSignalTestResult &out ) { out = WorkerChipOrchSignalTestResult{false, 0}; - if (has_error()) { + if (!ready()) { return false; } - if (!validate_counter_addr_for_op( - WorkerChipEndpointOp::SIGNAL_TEST, counter_addr, cmp_value, 0, "invalid counter address" - )) { + uint64_t offset = 0; + if (!absolute_to_counter_offset(counter_addr, offset)) { + set_error( + WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, WorkerChipEndpointOp::SIGNAL_TEST, desc_.region_id, + counter_addr, cmp_value, "invalid counter address" + ); return false; } - if (!worker_chip_orch_comm::valid_wait_cmp(cmp)) { + RegionWaitCmp region_cmp{}; + if (!worker_chip_wait_cmp_to_region(cmp, region_cmp)) { set_error( WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL, WorkerChipEndpointOp::SIGNAL_TEST, desc_.region_id, counter_addr, cmp_value, "invalid wait comparison" ); return false; } - int32_t observed = load_counter(counter_addr); - out = - WorkerChipOrchSignalTestResult{worker_chip_orch_comm::compare_counter(observed, cmp_value, cmp), observed}; + RegionSignalTestResult inner{}; + if (!view_->test(offset, cmp_value, region_cmp, inner)) { + adopt_view_error(WorkerChipEndpointOp::SIGNAL_TEST, counter_addr, cmp_value); + return false; + } + out = WorkerChipOrchSignalTestResult{inner.matched, inner.observed}; return true; } @@ -222,79 +262,74 @@ class WorkerChipOrchEndpoint { uint64_t counter_addr, int32_t cmp_value, WorkerChipOrchWaitCmp cmp, uint64_t timeout, int32_t &observed ) { observed = 0; - if (has_error()) { + if (!ready()) { return false; } - if (!validate_counter_addr_for_op( - WorkerChipEndpointOp::SIGNAL_WAIT, counter_addr, cmp_value, 0, "invalid counter address" - )) { + uint64_t offset = 0; + if (!absolute_to_counter_offset(counter_addr, offset)) { + set_error( + WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, WorkerChipEndpointOp::SIGNAL_WAIT, desc_.region_id, + counter_addr, cmp_value, "invalid counter address" + ); return false; } - if (!worker_chip_orch_comm::valid_wait_cmp(cmp)) { + RegionWaitCmp region_cmp{}; + if (!worker_chip_wait_cmp_to_region(cmp, region_cmp)) { set_error( WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL, WorkerChipEndpointOp::SIGNAL_WAIT, desc_.region_id, counter_addr, cmp_value, "invalid wait comparison" ); return false; } - - uint64_t start = device_time_now_ticks(); - uint64_t frequency_hz = device_time_frequency_hz(); - while (true) { - int32_t current = load_counter(counter_addr); - observed = current; - if (worker_chip_orch_comm::compare_counter(current, cmp_value, cmp)) { - return true; - } - uint64_t now = device_time_now_ticks(); - if (timeout == 0 || sys_cnt_elapsed_ns(start, now, frequency_hz) >= timeout) { - set_error( - WorkerChipEndpointErrorKind::SIGNAL_TIMEOUT, WorkerChipEndpointOp::SIGNAL_WAIT, desc_.region_id, - counter_addr, cmp_value, current, "wait timed out" - ); - return false; - } + if (!view_->wait(offset, cmp_value, region_cmp, timeout, observed)) { + adopt_view_error(WorkerChipEndpointOp::SIGNAL_WAIT, counter_addr, cmp_value); + return false; } + return true; } private: bool has_error() const { return error_.kind != WorkerChipEndpointErrorKind::NONE; } - bool validate_payload_range(WorkerChipEndpointOp op, uint64_t offset, uint64_t nbytes) { - WorkerChipOrchCommValidationError error = - worker_chip_orch_comm::validate_payload_bounds(offset, nbytes, desc_.payload_bytes); - if (error == WorkerChipOrchCommValidationError::OK) { - return true; - } - set_error( - WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, op, desc_.region_id, 0, 0, "payload range is out of bounds" - ); - return false; - } + bool ready() const { return !has_error() && view_.has_value(); } - bool validate_counter_addr_for_op( - WorkerChipEndpointOp op, uint64_t counter_addr, int32_t counter_operand, int32_t observed_counter, - const char *message - ) { - if (worker_chip_orch_comm::validate_counter_addr(desc_, counter_addr) == - WorkerChipOrchCommValidationError::OK) { - return true; + void install_view_or_error() { + if (worker_chip_orch_comm_magic(desc_.magic_version) != WORKER_CHIP_ORCH_COMM_MAGIC || + worker_chip_orch_comm_abi_major(desc_.magic_version) != WORKER_CHIP_ORCH_COMM_ABI_MAJOR || + desc_.region_id == 0) { + set_error( + WorkerChipEndpointErrorKind::BAD_DESCRIPTOR, WorkerChipEndpointOp::INIT, desc_.region_id, 0, 0, + "invalid compatibility descriptor" + ); + return; } - set_error( - WorkerChipEndpointErrorKind::OUT_OF_BOUNDS, op, desc_.region_id, counter_addr, counter_operand, - observed_counter, message + view_.emplace( + RegionPartLocalSpan{desc_.payload_base, desc_.payload_bytes}, + RegionPartLocalSpan{desc_.counter_base, desc_.counter_bytes} ); - return false; + if (view_->failed()) { + set_error( + WorkerChipEndpointErrorKind::BAD_DESCRIPTOR, WorkerChipEndpointOp::INIT, desc_.region_id, 0, 0, + "invalid independent local spans" + ); + } } - static volatile int32_t *counter_ptr(uint64_t counter_addr) { - return reinterpret_cast(static_cast(counter_addr)); + bool absolute_to_counter_offset(uint64_t counter_addr, uint64_t &offset) const { + offset = 0; + if (counter_addr < desc_.counter_base) { + return false; + } + offset = counter_addr - desc_.counter_base; + return true; } - static int32_t load_counter(uint64_t counter_addr) { - volatile int32_t *counter = counter_ptr(counter_addr); - cache_invalidate_range(reinterpret_cast(static_cast(counter_addr)), sizeof(*counter)); - return *counter; + void adopt_view_error(WorkerChipEndpointOp op, uint64_t counter_addr, int32_t counter_operand) { + const RegionViewError &ve = view_->error(); + set_error( + worker_chip_map_view_error_kind(ve.kind), op, desc_.region_id, counter_addr, counter_operand, ve.observed, + ve.message + ); } void set_error( @@ -317,4 +352,7 @@ class WorkerChipOrchEndpoint { WorkerChipOrchRegionDesc desc_{}; WorkerChipEndpointError error_{WorkerChipEndpointErrorKind::NONE, WorkerChipEndpointOp::INIT, 0, 0, 0, 0, ""}; + std::optional> view_{}; }; + +using WorkerChipOrchEndpoint = WorkerChipOrchEndpointImpl; diff --git a/src/common/platform/include/common/region_instance_semantics.h b/src/common/platform/include/common/region_instance_semantics.h new file mode 100644 index 0000000000..615d760d46 --- /dev/null +++ b/src/common/platform/include/common/region_instance_semantics.h @@ -0,0 +1,113 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include + +static constexpr uint64_t REGION_COUNTER_LOGICAL_ALIGNMENT = 4; +static constexpr uint64_t REGION_COUNTER_BASE_ALIGNMENT = 64; + +struct RegionPartLocalSpan { + uint64_t base; + uint64_t logical_bytes; +}; + +enum class RegionNotifyOp : uint32_t { + Set = 0, + Add = 1, +}; + +enum class RegionWaitCmp : uint32_t { + EQ = 0, + NE = 1, + GT = 2, + GE = 3, + LT = 4, + LE = 5, +}; + +inline bool region_add_overflows(uint64_t a, uint64_t b) { +#if defined(__clang__) || defined(__GNUC__) + uint64_t result = 0; + return __builtin_add_overflow(a, b, &result); +#else + return a > UINT64_MAX - b; +#endif +} + +inline bool region_spans_overlap(RegionPartLocalSpan first, RegionPartLocalSpan second) { + if (first.logical_bytes == 0 || second.logical_bytes == 0 || + region_add_overflows(first.base, first.logical_bytes) || + region_add_overflows(second.base, second.logical_bytes)) { + return false; + } + if (first.base < second.base) { + return second.base - first.base < first.logical_bytes; + } + return first.base - second.base < second.logical_bytes; +} + +inline bool region_valid_notify_op(RegionNotifyOp op) { return op == RegionNotifyOp::Set || op == RegionNotifyOp::Add; } + +inline bool region_valid_wait_cmp(RegionWaitCmp cmp) { + return cmp == RegionWaitCmp::EQ || cmp == RegionWaitCmp::NE || cmp == RegionWaitCmp::GT || + cmp == RegionWaitCmp::GE || cmp == RegionWaitCmp::LT || cmp == RegionWaitCmp::LE; +} + +inline bool region_compare_counter(int32_t observed, int32_t cmp_value, RegionWaitCmp cmp) { + switch (cmp) { + case RegionWaitCmp::EQ: + return observed == cmp_value; + case RegionWaitCmp::NE: + return observed != cmp_value; + case RegionWaitCmp::GT: + return observed > cmp_value; + case RegionWaitCmp::GE: + return observed >= cmp_value; + case RegionWaitCmp::LT: + return observed < cmp_value; + case RegionWaitCmp::LE: + return observed <= cmp_value; + } + return false; +} + +inline bool region_validate_payload_span(RegionPartLocalSpan span) { + return span.logical_bytes > 0 && !region_add_overflows(span.base, span.logical_bytes); +} + +inline bool region_validate_payload_range(uint64_t offset, uint64_t nbytes, uint64_t logical_bytes) { + return nbytes > 0 && logical_bytes > 0 && !region_add_overflows(offset, nbytes) && offset + nbytes <= logical_bytes; +} + +inline bool region_validate_counter_span(RegionPartLocalSpan span) { + return span.logical_bytes > 0 && (span.logical_bytes % REGION_COUNTER_LOGICAL_ALIGNMENT) == 0 && + (span.base % REGION_COUNTER_BASE_ALIGNMENT) == 0 && !region_add_overflows(span.base, span.logical_bytes); +} + +inline bool region_validate_independent_spans(RegionPartLocalSpan payload, RegionPartLocalSpan counter) { + return region_validate_payload_span(payload) && region_validate_counter_span(counter) && + !region_spans_overlap(payload, counter); +} + +inline bool region_counter_addr_in_span(RegionPartLocalSpan span, uint64_t counter_addr) { + if ((counter_addr % REGION_COUNTER_LOGICAL_ALIGNMENT) != 0 || + span.logical_bytes < REGION_COUNTER_LOGICAL_ALIGNMENT) { + return false; + } + if (counter_addr < span.base) { + return false; + } + uint64_t max_addr = span.base + span.logical_bytes - REGION_COUNTER_LOGICAL_ALIGNMENT; + return counter_addr <= max_addr; +} diff --git a/src/common/platform/include/common/worker_chip_orch_comm.h b/src/common/platform/include/common/worker_chip_orch_comm.h index 5face680ce..99eb7ff7d3 100644 --- a/src/common/platform/include/common/worker_chip_orch_comm.h +++ b/src/common/platform/include/common/worker_chip_orch_comm.h @@ -15,8 +15,10 @@ #include #include +#include "common/region_instance_semantics.h" + static constexpr uint32_t WORKER_CHIP_ORCH_COMM_MAGIC = 0x4C334C32u; // "WorkerChip" -static constexpr uint16_t WORKER_CHIP_ORCH_COMM_ABI_MAJOR = 2; +static constexpr uint16_t WORKER_CHIP_ORCH_COMM_ABI_MAJOR = 3; static constexpr uint16_t WORKER_CHIP_ORCH_COMM_ABI_MINOR = 0; static constexpr size_t WORKER_CHIP_ORCH_REGION_DESC_SCALAR_COUNT = 6; static constexpr uint64_t WORKER_CHIP_ORCH_COMM_COUNTER_BYTES = sizeof(int32_t); @@ -97,14 +99,7 @@ static inline uint16_t worker_chip_orch_comm_abi_minor(uint64_t magic_version) { return static_cast(magic_version & 0xFFFFu); } -static inline bool worker_chip_orch_comm_add_overflows(uint64_t a, uint64_t b) { -#if defined(__clang__) || defined(__GNUC__) - uint64_t result = 0; - return __builtin_add_overflow(a, b, &result); -#else - return a > UINT64_MAX - b; -#endif -} +static inline bool worker_chip_orch_comm_add_overflows(uint64_t a, uint64_t b) { return region_add_overflows(a, b); } static inline uint64_t worker_chip_orch_comm_add_sat(uint64_t a, uint64_t b) { #if defined(__clang__) || defined(__GNUC__) @@ -135,14 +130,9 @@ static inline bool worker_chip_orch_comm_range_contains(uint64_t base, uint64_t static inline bool worker_chip_orch_comm_ranges_overlap( uint64_t first_base, uint64_t first_size, uint64_t second_base, uint64_t second_size ) { - if (first_size == 0 || second_size == 0 || worker_chip_orch_comm_add_overflows(first_base, first_size) || - worker_chip_orch_comm_add_overflows(second_base, second_size)) { - return false; - } - if (first_base < second_base) { - return second_base - first_base < first_size; - } - return first_base - second_base < second_size; + return region_spans_overlap( + RegionPartLocalSpan{first_base, first_size}, RegionPartLocalSpan{second_base, second_size} + ); } static inline WorkerChipOrchCommValidationError @@ -150,8 +140,7 @@ worker_chip_orch_comm_validate_payload_bounds(uint64_t offset, uint64_t nbytes, if (nbytes == 0 || payload_bytes == 0) { return WorkerChipOrchCommValidationError::BAD_PAYLOAD_RANGE; } - if (worker_chip_orch_comm_add_overflows(offset, nbytes) || - worker_chip_orch_comm_add_sat(offset, nbytes) > payload_bytes) { + if (!region_validate_payload_range(offset, nbytes, payload_bytes)) { return WorkerChipOrchCommValidationError::OUT_OF_BOUNDS; } return WorkerChipOrchCommValidationError::OK; @@ -159,15 +148,9 @@ worker_chip_orch_comm_validate_payload_bounds(uint64_t offset, uint64_t nbytes, static inline WorkerChipOrchCommValidationError worker_chip_orch_comm_validate_counter_range(const WorkerChipOrchRegionDesc &desc) { - if (desc.counter_bytes == 0 || - !worker_chip_orch_comm_is_aligned(desc.counter_base) || - (desc.counter_bytes % WORKER_CHIP_ORCH_COMM_COUNTER_BYTES) != 0 || - worker_chip_orch_comm_add_overflows(desc.counter_base, desc.counter_bytes)) { - return WorkerChipOrchCommValidationError::BAD_COUNTER_RANGE; - } - if (worker_chip_orch_comm_ranges_overlap( - desc.payload_base, desc.payload_bytes, desc.counter_base, desc.counter_bytes - )) { + RegionPartLocalSpan payload{desc.payload_base, desc.payload_bytes}; + RegionPartLocalSpan counter{desc.counter_base, desc.counter_bytes}; + if (!region_validate_counter_span(counter) || region_spans_overlap(payload, counter)) { return WorkerChipOrchCommValidationError::BAD_COUNTER_RANGE; } return WorkerChipOrchCommValidationError::OK; @@ -235,50 +218,70 @@ static inline bool worker_chip_orch_comm_decode_desc( return error == WorkerChipOrchCommValidationError::OK; } -static inline bool worker_chip_orch_comm_valid_notify_op(WorkerChipOrchNotifyOp op) { - return op == WorkerChipOrchNotifyOp::Set || op == WorkerChipOrchNotifyOp::Add; -} - -static inline bool worker_chip_orch_comm_valid_wait_cmp(WorkerChipOrchWaitCmp cmp) { - return cmp == WorkerChipOrchWaitCmp::EQ || cmp == WorkerChipOrchWaitCmp::NE || cmp == WorkerChipOrchWaitCmp::GT || - cmp == WorkerChipOrchWaitCmp::GE || cmp == WorkerChipOrchWaitCmp::LT || cmp == WorkerChipOrchWaitCmp::LE; +static inline bool worker_chip_orch_comm_as_region_notify_op(WorkerChipOrchNotifyOp op, RegionNotifyOp &out) { + switch (op) { + case WorkerChipOrchNotifyOp::Set: + out = RegionNotifyOp::Set; + return true; + case WorkerChipOrchNotifyOp::Add: + out = RegionNotifyOp::Add; + return true; + } + return false; } -static inline bool -worker_chip_orch_comm_compare_counter(int32_t observed, int32_t cmp_value, WorkerChipOrchWaitCmp cmp) { +static inline bool worker_chip_orch_comm_as_region_wait_cmp(WorkerChipOrchWaitCmp cmp, RegionWaitCmp &out) { switch (cmp) { case WorkerChipOrchWaitCmp::EQ: - return observed == cmp_value; + out = RegionWaitCmp::EQ; + return true; case WorkerChipOrchWaitCmp::NE: - return observed != cmp_value; + out = RegionWaitCmp::NE; + return true; case WorkerChipOrchWaitCmp::GT: - return observed > cmp_value; + out = RegionWaitCmp::GT; + return true; case WorkerChipOrchWaitCmp::GE: - return observed >= cmp_value; + out = RegionWaitCmp::GE; + return true; case WorkerChipOrchWaitCmp::LT: - return observed < cmp_value; + out = RegionWaitCmp::LT; + return true; case WorkerChipOrchWaitCmp::LE: - return observed <= cmp_value; - default: + out = RegionWaitCmp::LE; + return true; + } + return false; +} + +static inline bool worker_chip_orch_comm_valid_notify_op(WorkerChipOrchNotifyOp op) { + RegionNotifyOp converted{}; + return worker_chip_orch_comm_as_region_notify_op(op, converted) && region_valid_notify_op(converted); +} + +static inline bool worker_chip_orch_comm_valid_wait_cmp(WorkerChipOrchWaitCmp cmp) { + RegionWaitCmp converted{}; + return worker_chip_orch_comm_as_region_wait_cmp(cmp, converted) && region_valid_wait_cmp(converted); +} + +static inline bool +worker_chip_orch_comm_compare_counter(int32_t observed, int32_t cmp_value, WorkerChipOrchWaitCmp cmp) { + RegionWaitCmp converted{}; + if (!worker_chip_orch_comm_as_region_wait_cmp(cmp, converted)) { return false; } + return region_compare_counter(observed, cmp_value, converted); } static inline WorkerChipOrchCommValidationError worker_chip_orch_comm_validate_counter_addr(const WorkerChipOrchRegionDesc &desc, uint64_t counter_addr) { - // Address validity is 4-byte; wrapper protocols must keep different - // counter writers off the same cache line. if (!worker_chip_orch_comm_is_aligned(counter_addr)) { return WorkerChipOrchCommValidationError::BAD_COUNTER_RANGE; } if (worker_chip_orch_comm_validate_counter_range(desc) != WorkerChipOrchCommValidationError::OK) { return WorkerChipOrchCommValidationError::BAD_COUNTER_RANGE; } - if (desc.counter_bytes < WORKER_CHIP_ORCH_COMM_COUNTER_BYTES) { - return WorkerChipOrchCommValidationError::BAD_COUNTER_RANGE; - } - uint64_t max_counter_addr = desc.counter_base + desc.counter_bytes - WORKER_CHIP_ORCH_COMM_COUNTER_BYTES; - if (counter_addr < desc.counter_base || counter_addr > max_counter_addr) { + if (!region_counter_addr_in_span(RegionPartLocalSpan{desc.counter_base, desc.counter_bytes}, counter_addr)) { return WorkerChipOrchCommValidationError::OUT_OF_BOUNDS; } return WorkerChipOrchCommValidationError::OK; diff --git a/src/common/platform/include/host/worker_chip_orch_region_access.h b/src/common/platform/include/host/worker_chip_orch_region_access.h deleted file mode 100644 index 5fae367953..0000000000 --- a/src/common/platform/include/host/worker_chip_orch_region_access.h +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright (c) PyPTO Contributors. - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of - * CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - * ----------------------------------------------------------------------------------------------------------- - */ - -#pragma once - -#include -#include - -enum class WorkerChipRegionAccessProfile : uint32_t { - INVALID = 0, - ONBOARD_VMM = 1, - SIM_POSIX_SHM = 2, -}; - -struct WorkerHostRegionMappingHandle { - uint64_t id{0}; - WorkerChipRegionAccessProfile profile{WorkerChipRegionAccessProfile::INVALID}; - uint64_t mapping_bytes{0}; -}; - -struct WorkerChipRegionCreateRequest { - uint64_t magic_version; - uint64_t request_bytes; - uint64_t payload_bytes; - uint64_t counter_bytes; -}; - -struct WorkerChipRegionCreateReply { - uint64_t desc[6]; - uint32_t access_profile; - uint32_t reserved; - int32_t device_id; - uint8_t backing_shm[32]; - uint64_t mapping_bytes; - uint64_t shareable_handle; -}; - -static constexpr size_t WORKER_CHIP_REGION_CREATE_REQUEST_BYTES = sizeof(WorkerChipRegionCreateRequest); -static constexpr size_t WORKER_CHIP_REGION_CREATE_REPLY_BYTES = sizeof(WorkerChipRegionCreateReply); diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index eda5096252..c3c2544980 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -649,6 +649,23 @@ target_link_libraries(test_worker_chip_orch_endpoint PRIVATE add_test(NAME test_worker_chip_orch_endpoint COMMAND test_worker_chip_orch_endpoint) set_tests_properties(test_worker_chip_orch_endpoint PROPERTIES LABELS "no_hardware") +add_executable(test_region_instance_view + common/test_region_instance_view.cpp + stubs/test_stubs.cpp +) +target_include_directories(test_region_instance_view PRIVATE + ${GTEST_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/../../../src/a2a3/platform/include + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/include +) +target_link_libraries(test_region_instance_view PRIVATE + ${GTEST_MAIN_LIB} + ${GTEST_LIB} + pthread +) +add_test(NAME test_region_instance_view COMMAND test_region_instance_view) +set_tests_properties(test_region_instance_view PROPERTIES LABELS "no_hardware") + add_common_utils_test(test_runtime_timeout_config common/test_runtime_timeout_config.cpp) target_include_directories(test_runtime_timeout_config PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../../src/a2a3/platform/include diff --git a/tests/ut/cpp/common/test_region_instance_view.cpp b/tests/ut/cpp/common/test_region_instance_view.cpp new file mode 100644 index 0000000000..8d607b2d04 --- /dev/null +++ b/tests/ut/cpp/common/test_region_instance_view.cpp @@ -0,0 +1,284 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include +#include +#include + +#include + +#include "aicpu/region_instance_view.h" + +namespace { + +struct RegionStorage { + alignas(64) std::array payload{}; + alignas(64) std::array counters{}; +}; + +RegionPartLocalSpan payload_span(RegionStorage *storage) { + return RegionPartLocalSpan{reinterpret_cast(storage->payload.data()), storage->payload.size()}; +} + +RegionPartLocalSpan counter_span(RegionStorage *storage) { + return RegionPartLocalSpan{ + reinterpret_cast(storage->counters.data()), storage->counters.size() * sizeof(int32_t) + }; +} + +RegionInstanceView make_view(RegionStorage *storage) { + return RegionInstanceView(payload_span(storage), counter_span(storage)); +} + +TEST(RegionInstanceViewTest, AcceptsIndependentNonAdjacentZeroAndReversedBases) { + alignas(64) std::array first{}; + alignas(64) std::array second{}; + RegionPartLocalSpan payload{reinterpret_cast(first.data()), first.size()}; + RegionPartLocalSpan counter{reinterpret_cast(second.data()), second.size()}; + + RegionInstanceView adjacent_or_not(payload, counter); + ASSERT_FALSE(adjacent_or_not.failed()) << adjacent_or_not.error().message; + EXPECT_EQ(adjacent_or_not.payload_span().base, payload.base); + EXPECT_EQ(adjacent_or_not.counter_span().base, counter.base); + + RegionInstanceView reversed(counter, payload); + ASSERT_FALSE(reversed.failed()) << reversed.error().message; + EXPECT_EQ(reversed.payload_span().base, counter.base); + EXPECT_EQ(reversed.counter_span().base, payload.base); + + RegionInstanceView zero_payload(RegionPartLocalSpan{0, 64}, RegionPartLocalSpan{256, 64}); + ASSERT_FALSE(zero_payload.failed()) << zero_payload.error().message; +} + +TEST(RegionInstanceViewTest, InvalidConstructionIsStickyAndHasNoIdentityCookie) { + static_assert(sizeof(RegionViewError::kind) == sizeof(uint32_t), "error kind is not an identity cookie"); + RegionInstanceView uninitialized; + EXPECT_TRUE(uninitialized.failed()); + EXPECT_EQ(uninitialized.error().kind, RegionViewErrorKind::INVALID_VIEW); + + RegionInstanceView overlap(RegionPartLocalSpan{0, 128}, RegionPartLocalSpan{64, 64}); + EXPECT_TRUE(overlap.failed()); + EXPECT_EQ(overlap.error().kind, RegionViewErrorKind::INVALID_VIEW); + EXPECT_EQ(overlap.error().op, RegionViewOp::CONSTRUCT); + EXPECT_STRNE(overlap.error().message, ""); + + RegionInstanceView unaligned(RegionPartLocalSpan{64, 64}, RegionPartLocalSpan{1, 64}); + EXPECT_TRUE(unaligned.failed()); + EXPECT_EQ(unaligned.error().kind, RegionViewErrorKind::INVALID_VIEW); + + RegionInstanceView odd_bytes(RegionPartLocalSpan{64, 64}, RegionPartLocalSpan{128, 6}); + EXPECT_TRUE(odd_bytes.failed()); + EXPECT_EQ(odd_bytes.error().kind, RegionViewErrorKind::INVALID_VIEW); + + RegionPayloadView view{1, 1}; + EXPECT_FALSE(overlap.payload().write(0, "x", 1)); + EXPECT_FALSE(overlap.payload().read(0, 1, view)); + EXPECT_EQ(view.local_addr, 0u); + EXPECT_EQ(overlap.error().kind, RegionViewErrorKind::INVALID_VIEW); +} + +TEST(RegionInstanceViewTest, RejectsOverflowUnalignedAndNonMultipleCounterSpans) { + EXPECT_FALSE(region_validate_payload_span(RegionPartLocalSpan{UINT64_MAX - 7, 16})); + EXPECT_FALSE(region_validate_counter_span(RegionPartLocalSpan{1, 64})); + EXPECT_FALSE(region_validate_counter_span(RegionPartLocalSpan{64, 6})); + EXPECT_FALSE(region_validate_counter_span(RegionPartLocalSpan{64, 0})); + EXPECT_FALSE(region_validate_independent_spans(RegionPartLocalSpan{0, 128}, RegionPartLocalSpan{64, 64})); +} + +TEST(RegionInstanceViewTest, PayloadWriteAndReadUseOffsetZeroSpans) { + RegionStorage storage{}; + RegionInstanceView view = make_view(&storage); + ASSERT_FALSE(view.failed()) << view.error().message; + EXPECT_EQ(view.payload().span().logical_bytes, storage.payload.size()); + EXPECT_EQ(view.counter().span().logical_bytes, storage.counters.size() * sizeof(int32_t)); + + const uint32_t marker = 0xA5B6C7D8u; + ASSERT_TRUE(view.payload().write(12, &marker, sizeof(marker))) << view.error().message; + uint32_t observed = 0; + std::memcpy(&observed, storage.payload.data() + 12, sizeof(observed)); + EXPECT_EQ(observed, marker); + + RegionPayloadView inner{}; + ASSERT_TRUE(view.payload().read(12, sizeof(marker), inner)) << view.error().message; + EXPECT_EQ(inner.nbytes, sizeof(marker)); + EXPECT_EQ(*reinterpret_cast(static_cast(inner.local_addr)), marker); +} + +TEST(RegionInstanceViewTest, PreIssuePayloadBoundsDoNotPoison) { + RegionStorage storage{}; + RegionInstanceView view = make_view(&storage); + RegionPayloadView inner{0xCAFE, 0xBEEF}; + + EXPECT_FALSE(view.payload().read(120, 16, inner)); + EXPECT_EQ(inner.local_addr, 0u); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::OUT_OF_BOUNDS); + EXPECT_FALSE(view.failed()); + + EXPECT_FALSE(view.payload().write(0, nullptr, 1)); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::OUT_OF_BOUNDS); + EXPECT_FALSE(view.failed()); + + const uint8_t byte = 0x11; + EXPECT_TRUE(view.payload().write(0, &byte, 1)) << view.error().message; + EXPECT_EQ(storage.payload[0], 0x11); +} + +TEST(RegionInstanceViewTest, NotifyTestAndWaitFollowNeutralPoisonRules) { + RegionStorage storage{}; + RegionInstanceView view = make_view(&storage); + uint64_t addr = 0; + ASSERT_TRUE(view.counter().addr(0, addr)) << view.error().message; + EXPECT_TRUE(view.counter().contains_addr(addr)); + + ASSERT_TRUE(view.counter().notify(0, 5, RegionNotifyOp::Set)) << view.error().message; + EXPECT_EQ(storage.counters[0], 5); + ASSERT_TRUE(view.counter().notify(0, -2, RegionNotifyOp::Add)) << view.error().message; + EXPECT_EQ(storage.counters[0], 3); + + RegionSignalTestResult result{}; + storage.counters[1] = 7; + ASSERT_TRUE(view.counter().test(4, 7, RegionWaitCmp::EQ, result)) << view.error().message; + EXPECT_TRUE(result.matched); + EXPECT_EQ(result.observed, 7); + ASSERT_TRUE(view.counter().test(4, 8, RegionWaitCmp::NE, result)) << view.error().message; + EXPECT_TRUE(result.matched); + ASSERT_TRUE(view.counter().test(4, 6, RegionWaitCmp::GT, result)) << view.error().message; + EXPECT_TRUE(result.matched); + ASSERT_TRUE(view.counter().test(4, 7, RegionWaitCmp::GE, result)) << view.error().message; + EXPECT_TRUE(result.matched); + ASSERT_TRUE(view.counter().test(4, 8, RegionWaitCmp::LT, result)) << view.error().message; + EXPECT_TRUE(result.matched); + ASSERT_TRUE(view.counter().test(4, 7, RegionWaitCmp::LE, result)) << view.error().message; + EXPECT_TRUE(result.matched); + ASSERT_TRUE(view.counter().test(4, 8, RegionWaitCmp::EQ, result)) << view.error().message; + EXPECT_FALSE(result.matched); + EXPECT_EQ(result.observed, 7); + EXPECT_FALSE(view.failed()); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::NONE); + + int32_t observed = 0; + EXPECT_FALSE(view.counter().wait(0, 99, RegionWaitCmp::EQ, 1, observed)); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::TIMEOUT); + EXPECT_FALSE(view.failed()); + EXPECT_EQ(observed, 3); + EXPECT_TRUE(view.counter().notify(0, 9, RegionNotifyOp::Set)) << view.error().message; + EXPECT_TRUE(view.counter().wait(0, 8, RegionWaitCmp::GE, 1'000'000, observed)) << view.error().message; + EXPECT_EQ(observed, 9); +} + +TEST(RegionInstanceViewTest, InvalidEnumAndUnalignedOffsetDoNotPoison) { + RegionStorage storage{}; + RegionInstanceView view = make_view(&storage); + + EXPECT_FALSE(view.counter().notify(0, 1, static_cast(99))); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::INVALID_ENUM); + EXPECT_FALSE(view.failed()); + + uint64_t addr = 0xCAFE; + EXPECT_FALSE(view.counter().addr(2, addr)); + EXPECT_EQ(addr, 0u); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::OUT_OF_BOUNDS); + EXPECT_FALSE(view.failed()); + EXPECT_TRUE(view.counter().notify(0, 1, RegionNotifyOp::Set)) << view.error().message; +} + +struct FailingCopyOps : LocalMemoryOps { + bool copy_bytes(void *, const void *, size_t) const { return false; } +}; + +struct FailingLoadOps : LocalMemoryOps { + bool load_i32(uint64_t, int32_t &) const { return false; } +}; + +TEST(RegionInstanceViewTest, IssuedLocalOpsFailureIsSticky) { + RegionStorage storage{}; + RegionInstanceViewImpl view(payload_span(&storage), counter_span(&storage)); + ASSERT_FALSE(view.failed()) << view.error().message; + const uint32_t marker = 1; + EXPECT_FALSE(view.payload().write(0, &marker, sizeof(marker))); + EXPECT_TRUE(view.failed()); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); + EXPECT_EQ(view.error().op, RegionViewOp::PAYLOAD_WRITE); + + RegionPayloadView inner{}; + EXPECT_FALSE(view.payload().read(0, 4, inner)); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); +} + +TEST(RegionInstanceViewTest, DefaultInvalidAndIssuedFailedViewsHaveNoReset) { + RegionInstanceView uninitialized; + EXPECT_TRUE(uninitialized.failed()); + EXPECT_EQ(uninitialized.error().kind, RegionViewErrorKind::INVALID_VIEW); + RegionPayloadView unread{}; + EXPECT_FALSE(uninitialized.payload().read(0, 1, unread)); + EXPECT_EQ(uninitialized.error().kind, RegionViewErrorKind::INVALID_VIEW); + + RegionInstanceView overlap(RegionPartLocalSpan{0, 128}, RegionPartLocalSpan{64, 64}); + EXPECT_TRUE(overlap.failed()); + const uint8_t byte = 1; + EXPECT_FALSE(overlap.payload().write(0, &byte, 1)); + EXPECT_EQ(overlap.error().kind, RegionViewErrorKind::INVALID_VIEW); + EXPECT_EQ(overlap.error().op, RegionViewOp::CONSTRUCT); +} + +TEST(RegionInstanceViewTest, AddUsesDefined32BitWrap) { + RegionStorage storage{}; + RegionInstanceView view = make_view(&storage); + ASSERT_FALSE(view.failed()) << view.error().message; + + ASSERT_TRUE(view.counter().notify(0, INT32_MAX, RegionNotifyOp::Set)) << view.error().message; + ASSERT_TRUE(view.counter().notify(0, 1, RegionNotifyOp::Add)) << view.error().message; + EXPECT_EQ(storage.counters[0], INT32_MIN); + + ASSERT_TRUE(view.counter().notify(0, INT32_MIN, RegionNotifyOp::Set)) << view.error().message; + ASSERT_TRUE(view.counter().notify(0, -1, RegionNotifyOp::Add)) << view.error().message; + EXPECT_EQ(storage.counters[0], INT32_MAX); + + ASSERT_TRUE(view.counter().notify(0, 10, RegionNotifyOp::Set)) << view.error().message; + ASSERT_TRUE(view.counter().notify(0, 4, RegionNotifyOp::Add)) << view.error().message; + EXPECT_EQ(storage.counters[0], 14); + + ASSERT_TRUE(view.counter().notify(0, 10, RegionNotifyOp::Set)) << view.error().message; + ASSERT_TRUE(view.counter().notify(0, -3, RegionNotifyOp::Add)) << view.error().message; + EXPECT_EQ(storage.counters[0], 7); +} + +struct FailingInvalidateOps : LocalMemoryOps { + bool invalidate(const void *, size_t) const { return false; } +}; + +TEST(RegionInstanceViewTest, IssuedPayloadReadFailureIsSticky) { + RegionStorage storage{}; + RegionInstanceViewImpl view(payload_span(&storage), counter_span(&storage)); + ASSERT_FALSE(view.failed()) << view.error().message; + RegionPayloadView inner{}; + EXPECT_FALSE(view.payload().read(0, 4, inner)); + EXPECT_TRUE(view.failed()); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); + EXPECT_EQ(view.error().op, RegionViewOp::PAYLOAD_READ); + EXPECT_FALSE(view.payload().write(0, "x", 1)); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); +} + +TEST(RegionInstanceViewTest, IssuedCounterLoadFailureIsSticky) { + RegionStorage storage{}; + RegionInstanceViewImpl view(payload_span(&storage), counter_span(&storage)); + ASSERT_FALSE(view.failed()) << view.error().message; + EXPECT_FALSE(view.counter().notify(0, 1, RegionNotifyOp::Add)); + EXPECT_TRUE(view.failed()); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); + EXPECT_EQ(view.error().op, RegionViewOp::NOTIFY); + EXPECT_FALSE(view.counter().notify(0, 1, RegionNotifyOp::Set)); + EXPECT_EQ(view.error().kind, RegionViewErrorKind::ISSUED_FAILURE); +} + +} // namespace diff --git a/tests/ut/cpp/common/test_worker_chip_orch_comm.cpp b/tests/ut/cpp/common/test_worker_chip_orch_comm.cpp index d1319f6be4..6c08a671ce 100644 --- a/tests/ut/cpp/common/test_worker_chip_orch_comm.cpp +++ b/tests/ut/cpp/common/test_worker_chip_orch_comm.cpp @@ -10,14 +10,12 @@ */ #include -#include #include #include #include #include "common/worker_chip_orch_comm.h" -#include "host/worker_chip_orch_region_access.h" namespace { @@ -71,6 +69,14 @@ TEST(WorkerChipOrchCommTest, DescriptorRejectsBadMajorVersion) { EXPECT_EQ(error, WorkerChipOrchCommValidationError::BAD_MAGIC_VERSION); } +TEST(WorkerChipOrchCommTest, DescriptorRejectsAbiMajor2) { + WorkerChipOrchRegionDesc desc = valid_desc(); + desc.magic_version = worker_chip_orch_comm_pack_magic_version(WORKER_CHIP_ORCH_COMM_MAGIC, 2, 0); + + WorkerChipOrchCommValidationError error = worker_chip_orch_comm_validate_desc(desc); + EXPECT_EQ(error, WorkerChipOrchCommValidationError::BAD_MAGIC_VERSION); +} + TEST(WorkerChipOrchCommTest, DescriptorKeepsRegionIdZeroAsInvalidSentinel) { WorkerChipOrchRegionDesc desc = valid_desc(); desc.region_id = 0; @@ -196,6 +202,32 @@ TEST(WorkerChipOrchCommTest, CompileTimeAlignmentRequiresPowerOfTwoAbiAlignment) EXPECT_FALSE(worker_chip_orch_comm_is_aligned_runtime(24, 3)); } +TEST(WorkerChipOrchCommTest, NeutralHelpersMatchCompatibilityDelegation) { + EXPECT_EQ(worker_chip_orch_comm_add_overflows(7, 9), region_add_overflows(7, 9)); + EXPECT_EQ(worker_chip_orch_comm_add_overflows(UINT64_MAX - 1, 2), region_add_overflows(UINT64_MAX - 1, 2)); + EXPECT_EQ( + worker_chip_orch_comm_ranges_overlap(0, 64, 32, 64), + region_spans_overlap(RegionPartLocalSpan{0, 64}, RegionPartLocalSpan{32, 64}) + ); + EXPECT_EQ( + worker_chip_orch_comm_ranges_overlap(0, 64, 256, 64), + region_spans_overlap(RegionPartLocalSpan{0, 64}, RegionPartLocalSpan{256, 64}) + ); + + RegionNotifyOp notify{}; + ASSERT_TRUE(worker_chip_orch_comm_as_region_notify_op(WorkerChipOrchNotifyOp::Set, notify)); + EXPECT_TRUE(region_valid_notify_op(notify)); + EXPECT_TRUE(worker_chip_orch_comm_valid_notify_op(WorkerChipOrchNotifyOp::Set)); + EXPECT_FALSE(worker_chip_orch_comm_as_region_notify_op(static_cast(2), notify)); + + RegionWaitCmp cmp{}; + ASSERT_TRUE(worker_chip_orch_comm_as_region_wait_cmp(WorkerChipOrchWaitCmp::GE, cmp)); + EXPECT_EQ( + worker_chip_orch_comm_compare_counter(5, 5, WorkerChipOrchWaitCmp::GE), region_compare_counter(5, 5, cmp) + ); + EXPECT_FALSE(worker_chip_orch_comm_as_region_wait_cmp(static_cast(6), cmp)); +} + TEST(WorkerChipOrchCommTest, NotifyOpAndWaitCmpValidationRejectUnknownValues) { EXPECT_TRUE(worker_chip_orch_comm_valid_notify_op(WorkerChipOrchNotifyOp::Set)); EXPECT_TRUE(worker_chip_orch_comm_valid_notify_op(WorkerChipOrchNotifyOp::Add)); @@ -224,24 +256,4 @@ TEST(WorkerChipOrchCommTest, DescriptorAndSignalResultAreFixedSizePodTypes) { static_assert(std::is_trivially_copyable::value, "test result must be fixed-size"); } -TEST(WorkerChipOrchCommTest, LifecycleCreateWireStructsHaveFixedLayout) { - static_assert(std::is_standard_layout::value, "create request must be POD-like"); - static_assert( - std::is_trivially_copyable::value, "create request must be fixed-size" - ); - static_assert(std::is_standard_layout::value, "create reply must be POD-like"); - static_assert(std::is_trivially_copyable::value, "create reply must be fixed-size"); - - EXPECT_EQ(offsetof(WorkerChipRegionCreateRequest, magic_version), 0u); - EXPECT_EQ(offsetof(WorkerChipRegionCreateRequest, payload_bytes), sizeof(uint64_t) * 2); - EXPECT_EQ(sizeof(WorkerChipRegionCreateRequest), 32u); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, desc), 0u); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, access_profile), sizeof(uint64_t) * 6); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, device_id), sizeof(uint64_t) * 6 + sizeof(uint32_t) * 2); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, backing_shm), sizeof(uint64_t) * 6 + sizeof(uint32_t) * 2 + 4u); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, mapping_bytes), 96u); - EXPECT_EQ(offsetof(WorkerChipRegionCreateReply, shareable_handle), 104u); - EXPECT_EQ(sizeof(WorkerChipRegionCreateReply), 112u); -} - } // namespace diff --git a/tests/ut/cpp/common/test_worker_chip_orch_endpoint.cpp b/tests/ut/cpp/common/test_worker_chip_orch_endpoint.cpp index ee334a0e2e..1ff05d3dc1 100644 --- a/tests/ut/cpp/common/test_worker_chip_orch_endpoint.cpp +++ b/tests/ut/cpp/common/test_worker_chip_orch_endpoint.cpp @@ -49,6 +49,9 @@ TEST(WorkerChipOrchEndpointTest, DecodesDescriptorScalarsAndCounterRange) { ASSERT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::NONE) << endpoint.error().message; EXPECT_EQ(endpoint.descriptor().counter_base, desc.counter_base); EXPECT_EQ(endpoint.descriptor().counter_bytes, desc.counter_bytes); + EXPECT_EQ(endpoint.view().payload_span().base, desc.payload_base); + EXPECT_EQ(endpoint.view().counter_span().base, desc.counter_base); + EXPECT_FALSE(endpoint.view().failed()); uint64_t counter_addr = 0; ASSERT_TRUE(endpoint.counter_addr(8, counter_addr)) << endpoint.error().message; @@ -190,6 +193,11 @@ TEST(WorkerChipOrchEndpointTest, SignalWaitTimeoutCarriesStructuredMetadata) { EXPECT_EQ(endpoint.error().counter_operand, 1); EXPECT_EQ(endpoint.error().observed_counter, 0); EXPECT_EQ(observed, 0); + EXPECT_FALSE(endpoint.view().failed()); + + WorkerChipOrchPayloadView view{}; + EXPECT_FALSE(endpoint.payload_read(0, sizeof(uint32_t), view)); + EXPECT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::SIGNAL_TIMEOUT); } TEST(WorkerChipOrchEndpointTest, SignalWaitDoesNotTreatGreaterObservedValueAsProtocolError) { @@ -214,7 +222,131 @@ TEST(WorkerChipOrchEndpointTest, RejectsBadDescriptorScalars) { EXPECT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::BAD_DESCRIPTOR); EXPECT_EQ(endpoint.error().op, WorkerChipEndpointOp::INIT); - EXPECT_STRNE(endpoint.error().message, ""); + EXPECT_STREQ(endpoint.error().message, "invalid compatibility descriptor"); +} + +TEST(WorkerChipOrchEndpointTest, RejectsInvalidIndependentSpansWithFixedMessage) { + RegionStorage storage{}; + WorkerChipOrchRegionDesc desc = make_desc(&storage); + desc.payload_bytes = 0; + + WorkerChipOrchEndpoint endpoint(desc); + + EXPECT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::BAD_DESCRIPTOR); + EXPECT_STREQ(endpoint.error().message, "invalid independent local spans"); +} + +TEST(WorkerChipOrchEndpointTest, UnknownNotifyAndWaitStaySignalProtocolWithoutEnteringView) { + RegionStorage storage{}; + WorkerChipOrchEndpoint endpoint(make_desc(&storage)); + uint64_t counter_addr = 0; + ASSERT_TRUE(endpoint.counter_addr(0, counter_addr)) << endpoint.error().message; + + EXPECT_FALSE(endpoint.signal_notify(counter_addr, 1, static_cast(99))); + EXPECT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL); + EXPECT_FALSE(endpoint.view().failed()); + + WorkerChipOrchEndpoint second(make_desc(&storage)); + WorkerChipOrchSignalTestResult result{}; + EXPECT_FALSE(second.signal_test(counter_addr, 1, static_cast(99), result)); + EXPECT_EQ(second.error().kind, WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL); + EXPECT_FALSE(second.view().failed()); +} + +struct FailingCopyOps : LocalMemoryOps { + bool copy_bytes(void *, const void *, size_t) const { return false; } +}; + +struct FailingInvalidateOps : LocalMemoryOps { + bool invalidate(const void *, size_t) const { return false; } +}; + +struct FailingLoadOps : LocalMemoryOps { + bool load_i32(uint64_t, int32_t &) const { return false; } +}; + +TEST(WorkerChipOrchEndpointTest, IssuedPayloadFailuresMapToLocalOperation) { + RegionStorage storage{}; + WorkerChipOrchEndpointImpl reader(make_desc(&storage)); + WorkerChipOrchPayloadView view{}; + EXPECT_FALSE(reader.payload_read(0, 4, view)); + EXPECT_EQ(reader.error().kind, WorkerChipEndpointErrorKind::LOCAL_OPERATION); + + WorkerChipOrchEndpointImpl writer(make_desc(&storage)); + const uint32_t marker = 1; + EXPECT_FALSE(writer.payload_write(0, &marker, sizeof(marker))); + EXPECT_EQ(writer.error().kind, WorkerChipEndpointErrorKind::LOCAL_OPERATION); +} + +TEST(WorkerChipOrchEndpointTest, IssuedCounterFailuresMapToLocalOperation) { + RegionStorage storage{}; + uint64_t counter_addr = 0; + { + WorkerChipOrchEndpoint probe(make_desc(&storage)); + ASSERT_TRUE(probe.counter_addr(0, counter_addr)) << probe.error().message; + } + + WorkerChipOrchEndpointImpl notify_ep(make_desc(&storage)); + EXPECT_FALSE(notify_ep.signal_notify(counter_addr, 1, WorkerChipOrchNotifyOp::Add)); + EXPECT_EQ(notify_ep.error().kind, WorkerChipEndpointErrorKind::LOCAL_OPERATION); + + WorkerChipOrchEndpointImpl test_ep(make_desc(&storage)); + WorkerChipOrchSignalTestResult result{}; + EXPECT_FALSE(test_ep.signal_test(counter_addr, 1, WorkerChipOrchWaitCmp::EQ, result)); + EXPECT_EQ(test_ep.error().kind, WorkerChipEndpointErrorKind::LOCAL_OPERATION); + + WorkerChipOrchEndpointImpl wait_ep(make_desc(&storage)); + int32_t observed = 0; + EXPECT_FALSE(wait_ep.signal_wait(counter_addr, 1, WorkerChipOrchWaitCmp::EQ, 1, observed)); + EXPECT_EQ(wait_ep.error().kind, WorkerChipEndpointErrorKind::LOCAL_OPERATION); +} + +TEST(WorkerChipOrchEndpointTest, EnumConversionCoversSetAddAndAllComparisons) { + RegionNotifyOp notify{}; + ASSERT_TRUE(worker_chip_notify_op_to_region(WorkerChipOrchNotifyOp::Set, notify)); + EXPECT_EQ(notify, RegionNotifyOp::Set); + ASSERT_TRUE(worker_chip_notify_op_to_region(WorkerChipOrchNotifyOp::Add, notify)); + EXPECT_EQ(notify, RegionNotifyOp::Add); + EXPECT_FALSE(worker_chip_notify_op_to_region(static_cast(3), notify)); + + RegionWaitCmp cmp{}; + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::EQ, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::EQ); + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::NE, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::NE); + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::GT, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::GT); + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::GE, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::GE); + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::LT, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::LT); + ASSERT_TRUE(worker_chip_wait_cmp_to_region(WorkerChipOrchWaitCmp::LE, cmp)); + EXPECT_EQ(cmp, RegionWaitCmp::LE); + EXPECT_FALSE(worker_chip_wait_cmp_to_region(static_cast(9), cmp)); + EXPECT_EQ( + worker_chip_map_view_error_kind(RegionViewErrorKind::ISSUED_FAILURE), + WorkerChipEndpointErrorKind::LOCAL_OPERATION + ); + EXPECT_EQ( + worker_chip_map_view_error_kind(RegionViewErrorKind::INVALID_ENUM), WorkerChipEndpointErrorKind::SIGNAL_PROTOCOL + ); +} + +TEST(WorkerChipOrchEndpointTest, RejectsAbiMajor2AndOldContiguousMagic) { + RegionStorage storage{}; + WorkerChipOrchRegionDesc desc = make_desc(&storage); + desc.magic_version = worker_chip_orch_comm_pack_magic_version(WORKER_CHIP_ORCH_COMM_MAGIC, 2, 0); + std::array scalars{ + desc.magic_version, desc.region_id, desc.payload_base, + desc.payload_bytes, desc.counter_base, desc.counter_bytes, + }; + + WorkerChipOrchEndpoint endpoint(scalars.data(), scalars.size()); + + EXPECT_EQ(endpoint.error().kind, WorkerChipEndpointErrorKind::BAD_DESCRIPTOR); + EXPECT_EQ(endpoint.error().op, WorkerChipEndpointOp::INIT); + EXPECT_TRUE(endpoint.view().failed()); + EXPECT_STREQ(endpoint.error().message, "invalid compatibility descriptor"); } } // namespace diff --git a/tests/ut/py/test_worker/test_cancellation_journal.py b/tests/ut/py/test_worker/test_cancellation_journal.py index 16c8cea462..c5302c0eec 100644 --- a/tests/ut/py/test_worker/test_cancellation_journal.py +++ b/tests/ut/py/test_worker/test_cancellation_journal.py @@ -355,7 +355,7 @@ def fail_then_ok(): @pytest.mark.parametrize( ("method_name", "kind"), [ - ("_cleanup_worker_chip_regions", "region"), + ("_sweep_region_instances", "region"), ("_release_all_live_domains", "domain"), ("_flush_pending_remote_frees", "remote"), ("_release_all_buffers", "buffer"), diff --git a/tests/ut/py/test_worker/test_comm_provider.py b/tests/ut/py/test_worker/test_comm_provider.py new file mode 100644 index 0000000000..ca5096e22e --- /dev/null +++ b/tests/ut/py/test_worker/test_comm_provider.py @@ -0,0 +1,1638 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 +"""Unit tests for neutral comm-provider values and typed errors.""" + +from __future__ import annotations + +import ast +import dataclasses +import logging +import sys +from multiprocessing.shared_memory import SharedMemory +from pathlib import Path + +import pytest +import simpler.comm_provider as comm_provider_module +from simpler.buffer import BackendKind +from simpler.comm_provider import ( + POSIX_SHM_TOKEN_MAX_BYTES, + DeviceAllocationTarget, + HostAllocationTarget, + ImportCapability, + PosixShmImport, + ProviderCleanupFailure, + ProviderPartResourceState, + ProviderRegionResourceState, + ProviderRegionStore, + ProviderRegionStoreState, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationContext, + RegionAllocationError, + RegionAllocationResult, + RegionAllocationSpec, + RegionCleanupCause, + RegionControlError, + RegionControlErrorKind, + RegionEnvironmentKind, + RegionExportDescriptor, + RegionOperationKind, + RegionPartAllocation, + RegionPartAllocationSpec, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, + SimPosixShmAllocation, + VmmAllocation, + VmmShareableHandleImport, + validate_independent_local_views, +) + +_COMM_PROVIDER_PATH = Path(comm_provider_module.__file__).resolve() +_UINT64_MAX = (1 << 64) - 1 +_INT32_MAX = (1 << 31) - 1 + + +def _payload_spec(logical_bytes: int = 64, backing=BackendKind.VMM_WINDOW) -> RegionPartAllocationSpec: + return RegionPartAllocationSpec(planned_backing_kind=backing, logical_bytes=logical_bytes) + + +def _counter_spec(logical_bytes: int = 8, backing=BackendKind.VMM_WINDOW) -> RegionPartAllocationSpec: + return RegionPartAllocationSpec(planned_backing_kind=backing, logical_bytes=logical_bytes) + + +def _posix_export(logical_bytes: int, shm_name: str, *, mapping_bytes: int | None = None) -> RegionPartExportDescriptor: + return RegionPartExportDescriptor( + planned_backing_kind=BackendKind.VMM_WINDOW, + logical_bytes=logical_bytes, + mapping_bytes=logical_bytes if mapping_bytes is None else mapping_bytes, + import_capability=PosixShmImport(shm_name=shm_name), + ) + + +def _vmm_export(logical_bytes: int, handle: int = 7) -> RegionPartExportDescriptor: + return RegionPartExportDescriptor( + planned_backing_kind=BackendKind.VMM_WINDOW, + logical_bytes=logical_bytes, + mapping_bytes=logical_bytes, + import_capability=VmmShareableHandleImport(device_id=0, shareable_handle=handle), + ) + + +def test_comm_provider_module_does_not_import_worker_chip_or_w5a(): + tree = ast.parse(_COMM_PROVIDER_PATH.read_text(encoding="utf-8")) + imported: list[str] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imported.append(node.module or "") + forbidden_substrings = ( + "worker_chip", + "worker.py", + "mailbox", + "w5a", + "orchestrator", + "comm_region", + ) + for name in imported: + lowered = name.replace("_", "").lower() + assert "workerchip" not in lowered + assert not any(token in name for token in forbidden_substrings) + assert name not in {"simpler.worker", "worker"} + + +@pytest.mark.parametrize( + ("enum_cls", "name", "value"), + [ + (RegionPartKind, "INVALID", 0), + (RegionPartKind, "PAYLOAD", 1), + (RegionPartKind, "COUNTER", 2), + (RegionControlErrorKind, "NONE", 0), + (RegionControlErrorKind, "BAD_MAGIC_VERSION", 1), + (RegionControlErrorKind, "BAD_MESSAGE_SIZE", 2), + (RegionControlErrorKind, "INVALID_ENUM_VALUE", 3), + (RegionControlErrorKind, "RESERVED_NONZERO", 4), + (RegionControlErrorKind, "INVALID_FIELD_VALUE", 5), + (RegionControlErrorKind, "STORE_LIFECYCLE", 6), + (RegionControlErrorKind, "INTERNAL_INVARIANT", 7), + (RegionControlErrorKind, "BACKEND_FAILURE", 8), + (RegionOperationKind, "NONE", 0), + (RegionOperationKind, "MATERIALIZE", 1), + (RegionOperationKind, "ZERO_BYTES", 2), + (RegionOperationKind, "DESCRIBE", 3), + (RegionOperationKind, "LOCAL_VIEW", 4), + (RegionOperationKind, "RELEASE", 5), + (RegionCleanupCause, "NONE", 0), + (RegionCleanupCause, "BACKEND_ERROR", 1), + (RegionCleanupCause, "INTERRUPTED", 2), + (RegionCleanupCause, "BACKEND_STATE_MISMATCH", 3), + (ProviderReleaseStatus, "RELEASED", 1), + (ProviderReleaseStatus, "ALREADY_GONE", 2), + (ProviderReleaseStatus, "UNKNOWN_RESOURCE", 3), + (ProviderReleaseStatus, "CLEANUP_INCOMPLETE", 4), + ], +) +def test_typed_code_values_match_design_tables(enum_cls, name, value): + member = enum_cls[name] + assert int(member) == value + assert enum_cls(value) is member + + +@pytest.mark.parametrize( + "enum_cls", + [ + RegionPartKind, + RegionControlErrorKind, + RegionOperationKind, + RegionCleanupCause, + ProviderReleaseStatus, + BackendKind, + ], +) +def test_unknown_nonzero_enum_values_are_rejected(enum_cls): + with pytest.raises(ValueError): + enum_cls(99) + + +def test_state_enum_names_are_exactly_the_accepted_set(): + assert [state.name for state in ProviderRegionStoreState] == [ + "OPEN", + "CLOSING", + "CLOSE_FAILED", + "CLOSED", + ] + assert [state.name for state in ProviderRegionResourceState] == [ + "CREATING", + "ACTIVE", + "CLEANUP_PENDING", + ] + assert [state.name for state in ProviderPartResourceState] == [ + "SHELL", + "MATERIALIZING", + "READY", + "CLEANUP_PENDING", + "RELEASED", + ] + + +def test_allocation_spec_accepts_independent_part_sizes_and_backing_kinds(): + spec = RegionAllocationSpec( + payload=_payload_spec(96, BackendKind.VMM_WINDOW), + counter=_counter_spec(12, BackendKind.POSIX_SHM), + ) + assert spec.part(RegionPartKind.PAYLOAD).logical_bytes == 96 + assert spec.part(RegionPartKind.COUNTER).logical_bytes == 12 + assert spec.payload.planned_backing_kind is BackendKind.VMM_WINDOW + assert spec.counter.planned_backing_kind is BackendKind.POSIX_SHM + + +@pytest.mark.parametrize("logical_bytes", [0, -1, _UINT64_MAX + 1, True]) +def test_part_spec_rejects_non_positive_or_overflow_sizes(logical_bytes): + with pytest.raises((ValueError, TypeError)): + RegionPartAllocationSpec(planned_backing_kind=BackendKind.VMM_WINDOW, logical_bytes=logical_bytes) + + +@pytest.mark.parametrize("logical_bytes", [1, 2, 3, 5, 6, 7, 9]) +def test_counter_spec_requires_multiple_of_four(logical_bytes): + payload = _payload_spec() + counter = RegionPartAllocationSpec(planned_backing_kind=BackendKind.VMM_WINDOW, logical_bytes=logical_bytes) + with pytest.raises(ValueError, match="multiple of 4"): + RegionAllocationSpec(payload=payload, counter=counter) + + +def test_allocation_context_preserves_typed_targets(): + sim = RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=DeviceAllocationTarget(device_id=3), + ) + reserved = RegionAllocationContext( + environment_kind=RegionEnvironmentKind.ONBOARD, + target=HostAllocationTarget(), + ) + assert sim.target.device_id == 3 + assert isinstance(reserved.target, HostAllocationTarget) + with pytest.raises(ValueError): + DeviceAllocationTarget(device_id=_INT32_MAX + 1) + + +def test_export_descriptor_keeps_capability_and_omits_local_addresses(): + descriptor = RegionExportDescriptor( + payload=_posix_export(64, "/pto_payload_a"), + counter=_vmm_export(8, handle=11), + ) + assert dataclasses.fields(RegionPartExportDescriptor) + field_names = {field.name for field in dataclasses.fields(RegionPartExportDescriptor)} + assert "local_base" not in field_names + assert "local_addr" not in field_names + assert isinstance(descriptor.payload.import_capability, PosixShmImport) + assert isinstance(descriptor.counter.import_capability, VmmShareableHandleImport) + result = RegionAllocationResult(provider_resource_id=1, export_descriptor=descriptor) + assert result.provider_resource_id == 1 + with pytest.raises(ValueError): + RegionAllocationResult(provider_resource_id=0, export_descriptor=descriptor) + + +@pytest.mark.parametrize("shm_name", ["", "a" * (POSIX_SHM_TOKEN_MAX_BYTES + 1), "npu\u4e00", "name\x00x"]) +def test_posix_shm_token_rejects_empty_overlong_non_ascii_and_nul(shm_name): + with pytest.raises((ValueError, TypeError)): + PosixShmImport(shm_name=shm_name) + + +def test_posix_shm_token_accepts_bounded_ascii_and_mapping_may_exceed_logical_bytes(): + token = "p" * POSIX_SHM_TOKEN_MAX_BYTES + export = _posix_export(8, token, mapping_bytes=64) + assert export.mapping_bytes == 64 + with pytest.raises(ValueError): + _posix_export(16, "ok", mapping_bytes=15) + + +def test_local_view_accepts_base_zero_and_rejects_uint64_span_overflow(): + view = RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=0, logical_bytes=8) + assert view.local_base == 0 + with pytest.raises(ValueError, match="overflowed uint64"): + RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=_UINT64_MAX, logical_bytes=1) + with pytest.raises(ValueError, match="64-byte aligned"): + RegionPartLocalView(part=RegionPartKind.COUNTER, local_base=32, logical_bytes=8) + + +def test_independent_span_rules_allow_noncontiguous_nonoverlapping_bases(): + payload = RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=0x1000, logical_bytes=100) + counter = RegionPartLocalView(part=RegionPartKind.COUNTER, local_base=0x2000, logical_bytes=8) + validate_independent_local_views(payload, counter) + + +def test_independent_span_rules_reject_overlap_and_swapped_parts(): + payload = RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=0, logical_bytes=128) + counter = RegionPartLocalView(part=RegionPartKind.COUNTER, local_base=64, logical_bytes=8) + with pytest.raises(ValueError, match="must not overlap"): + validate_independent_local_views(payload, counter) + with pytest.raises(ValueError, match="part PAYLOAD"): + validate_independent_local_views(counter, payload) + + +def test_adjacent_spans_do_not_overlap(): + payload = RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=0, logical_bytes=64) + counter = RegionPartLocalView(part=RegionPartKind.COUNTER, local_base=64, logical_bytes=8) + validate_independent_local_views(payload, counter) + + +def test_release_result_allows_at_most_one_failure_per_part(): + payload_failure = ProviderCleanupFailure( + part=RegionPartKind.PAYLOAD, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_ERROR, + ) + counter_failure = ProviderCleanupFailure( + part=RegionPartKind.COUNTER, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.INTERRUPTED, + ) + result = ProviderReleaseResult( + provider_resource_id=9, + status=ProviderReleaseStatus.CLEANUP_INCOMPLETE, + failures=(payload_failure, counter_failure), + ) + assert len(result.failures) == 2 + with pytest.raises(ValueError, match="at most one failure per part"): + ProviderReleaseResult( + provider_resource_id=9, + status=ProviderReleaseStatus.CLEANUP_INCOMPLETE, + failures=(payload_failure, payload_failure), + ) + with pytest.raises(ValueError): + ProviderReleaseResult( + provider_resource_id=9, + status=ProviderReleaseStatus.RELEASED, + failures=(payload_failure,), + ) + + +def test_allocation_error_carries_debt_bit_and_no_cleanup_cause(): + error = RegionAllocationError( + provisional_resource_id=4, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.COUNTER, + failed_operation=RegionOperationKind.ZERO_BYTES, + cleanup_debt_remaining=True, + ) + assert error.provisional_resource_id == 4 + assert error.cleanup_debt_remaining is True + assert error.failed_part is RegionPartKind.COUNTER + assert error.failed_operation is RegionOperationKind.ZERO_BYTES + assert not hasattr(error, "typed_cause") + assert not hasattr(error, "cleanup_cause") + with pytest.raises(ValueError): + RegionAllocationError( + provisional_resource_id=0, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=False, + ) + with pytest.raises(ValueError): + RegionAllocationError( + provisional_resource_id=1, + control_kind=RegionControlErrorKind.BAD_MESSAGE_SIZE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=False, + ) + + +def test_control_error_rejects_none_kind_and_unknown_values(): + error = RegionControlError(RegionControlErrorKind.STORE_LIFECYCLE, "store closed") + assert error.kind is RegionControlErrorKind.STORE_LIFECYCLE + assert error.failed_part is RegionPartKind.INVALID + with pytest.raises(ValueError): + RegionControlError(RegionControlErrorKind.NONE) + with pytest.raises(ValueError): + RegionControlError(99) + + +class FakeRegionPartAllocation: + """Deterministic RegionPartAllocation shell for store tests.""" + + def __init__( + self, + part: RegionPartKind, + spec: RegionPartAllocationSpec, + *, + world: FakeShellWorld, + local_base: int, + shm_name: str, + mapping_bytes: int | None = None, + ) -> None: + self.part = part + self.spec = spec + self.world = world + self.calls: list[str] = [] + self.local_ledger: list[object] = [] + self.materialized = False + self.release_count = 0 + self.zero_calls: list[tuple[int, int]] = [] + self.fail_materialize: BaseException | type[BaseException] | None = None + self.fail_zero: BaseException | type[BaseException] | None = None + self.fail_mapping_bytes: BaseException | type[BaseException] | None = None + self.fail_import_capability: BaseException | type[BaseException] | None = None + self.fail_local_base: BaseException | type[BaseException] | None = None + self.raise_on_release: BaseException | type[BaseException] | None = None + self.release_step_failures: list[ProviderCleanupFailure] = [] + self._local_base = local_base + self._mapping_bytes = spec.logical_bytes if mapping_bytes is None else mapping_bytes + self._import_capability: ImportCapability = PosixShmImport(shm_name=shm_name) + + def _raise_configured(self, spec: BaseException | type[BaseException] | None) -> None: + if spec is None: + return + if isinstance(spec, BaseException): + raise spec + raise spec() + + def materialize(self) -> None: + self.calls.append("materialize") + self.world.record(self.part, "materialize") + if self.release_count != 0: + self.world.materialize_called_release = True + self.world.constructed_at_first_materialize.setdefault(self.part, tuple(self.world.constructed_parts)) + self._raise_configured(self.fail_materialize) + self.materialized = True + self.world.side_effects.append((self.part, "materialize")) + + def mapping_bytes(self) -> int: + self.calls.append("mapping_bytes") + self.world.record(self.part, "mapping_bytes") + if not self.materialized: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "mapping_bytes requires a materialized shell", + failed_part=self.part, + failed_operation=RegionOperationKind.DESCRIBE, + ) + self._raise_configured(self.fail_mapping_bytes) + return self._mapping_bytes + + def import_capability(self) -> ImportCapability: + self.calls.append("import_capability") + self.world.record(self.part, "import_capability") + if not self.materialized: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "import_capability requires a materialized shell", + failed_part=self.part, + failed_operation=RegionOperationKind.DESCRIBE, + ) + self._raise_configured(self.fail_import_capability) + return self._import_capability + + def local_base(self) -> int: + self.calls.append("local_base") + self.world.record(self.part, "local_base") + if not self.materialized: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "local_base requires a materialized shell", + failed_part=self.part, + failed_operation=RegionOperationKind.LOCAL_VIEW, + ) + self._raise_configured(self.fail_local_base) + return self._local_base + + def zero_bytes(self, offset: int, nbytes: int) -> None: + self.calls.append("zero_bytes") + self.world.record(self.part, "zero_bytes") + self.zero_calls.append((offset, nbytes)) + if not self.materialized: + raise RegionControlError( + RegionControlErrorKind.INTERNAL_INVARIANT, + "zero_bytes requires a materialized shell", + failed_part=self.part, + failed_operation=RegionOperationKind.ZERO_BYTES, + ) + self._raise_configured(self.fail_zero) + + def release_once(self) -> ProviderCleanupFailure | None: + self.calls.append("release_once") + self.world.record(self.part, "release_once") + self.release_count += 1 + if self.release_count > 1: + self.world.duplicate_releases.append(self.part) + self._raise_configured(self.raise_on_release) + for failure in self.release_step_failures: + self.local_ledger.append(failure) + if self.release_step_failures: + return self.release_step_failures[0] + return None + + +class FakeShellWorld: + def __init__(self) -> None: + self.calls: list[tuple[RegionPartKind, str]] = [] + self.constructed_parts: list[RegionPartKind] = [] + self.side_effects: list[tuple[RegionPartKind, str]] = [] + self.constructed_at_first_materialize: dict[RegionPartKind, tuple[RegionPartKind, ...]] = {} + self.duplicate_releases: list[RegionPartKind] = [] + self.materialize_called_release = False + + def record(self, part: RegionPartKind, name: str) -> None: + self.calls.append((part, name)) + + +class FakeShellFactory: + def __init__(self) -> None: + self.world = FakeShellWorld() + self.payloads: list[FakeRegionPartAllocation] = [] + self.counters: list[FakeRegionPartAllocation] = [] + self.fail_construct: dict[RegionPartKind, BaseException | type[BaseException]] = {} + self._seq = 0 + + def __call__( + self, + context: RegionAllocationContext, + part: RegionPartKind, + spec: RegionPartAllocationSpec, + ) -> RegionPartAllocation: + del context + self.world.record(part, "construct") + if part in self.fail_construct: + spec_or_exc = self.fail_construct[part] + if isinstance(spec_or_exc, BaseException): + raise spec_or_exc + raise spec_or_exc() + self._seq += 1 + local_base = 0x1000 if part is RegionPartKind.PAYLOAD else 0x2000 + shell = FakeRegionPartAllocation( + part, + spec, + world=self.world, + local_base=local_base, + shm_name=f"/{part.name[0].lower()}{self._seq}", + ) + self.world.constructed_parts.append(part) + if part is RegionPartKind.PAYLOAD: + self.payloads.append(shell) + else: + self.counters.append(shell) + return shell + + +def _sim_context() -> RegionAllocationContext: + return RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=DeviceAllocationTarget(device_id=0), + ) + + +def _allocation_spec() -> RegionAllocationSpec: + return RegionAllocationSpec(payload=_payload_spec(), counter=_counter_spec()) + + +def _open_store(factory: FakeShellFactory | None = None) -> tuple[ProviderRegionStore, FakeShellFactory]: + factory = FakeShellFactory() if factory is None else factory + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + return store, factory + + +def _mutating_shell_factory(factory: FakeShellFactory, mutate): + def _factory(context, kind, spec): + shell = FakeShellFactory.__call__(factory, context, kind, spec) + mutate(shell, kind) + return shell + + return _factory + + +def _backend_failure(part: RegionPartKind) -> ProviderCleanupFailure: + return ProviderCleanupFailure( + part=part, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_ERROR, + ) + + +def test_empty_sweep_closes_an_open_store_and_is_idempotent(): + store, _factory = _open_store() + assert store.state is ProviderRegionStoreState.OPEN + assert store.sweep() == () + assert store.state is ProviderRegionStoreState.CLOSED + assert store.sweep() == () + assert store.state is ProviderRegionStoreState.CLOSED + with pytest.raises(RegionControlError) as exc_info: + store.allocate_and_export(_allocation_spec()) + assert exc_info.value.kind is RegionControlErrorKind.STORE_LIFECYCLE + assert not hasattr(store, "retry_cleanup") + + +def test_allocate_and_export_installs_both_shells_before_materialize(): + store, factory = _open_store() + result = store.allocate_and_export(_allocation_spec()) + assert result.provider_resource_id == 1 + assert factory.world.constructed_at_first_materialize[RegionPartKind.PAYLOAD] == ( + RegionPartKind.PAYLOAD, + RegionPartKind.COUNTER, + ) + assert factory.world.calls[:4] == [ + (RegionPartKind.PAYLOAD, "construct"), + (RegionPartKind.COUNTER, "construct"), + (RegionPartKind.PAYLOAD, "materialize"), + (RegionPartKind.COUNTER, "materialize"), + ] + assert factory.payloads[0].zero_calls == [] + assert factory.counters[0].zero_calls == [(0, 8)] + assert factory.world.materialize_called_release is False + payload = store.local_view(1, RegionPartKind.PAYLOAD) + counter = store.local_view(1, RegionPartKind.COUNTER) + validate_independent_local_views(payload, counter) + descriptor = store.describe(1) + assert isinstance(descriptor.payload.import_capability, PosixShmImport) + assert isinstance(descriptor.counter.import_capability, PosixShmImport) + assert descriptor.payload.import_capability.shm_name != descriptor.counter.import_capability.shm_name + + +def test_ids_are_monotonic_nonzero_and_never_reused_after_burned_create_failure(): + store, factory = _open_store() + first = store.allocate_and_export(_allocation_spec()) + assert first.provider_resource_id == 1 + factory.fail_construct[RegionPartKind.COUNTER] = RuntimeError("construct") + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.provisional_resource_id == 2 + assert error.cleanup_debt_remaining is False + assert error.control_kind is RegionControlErrorKind.INTERNAL_INVARIANT + assert error.failed_part is RegionPartKind.COUNTER + assert factory.payloads[-1].release_count == 1 + assert factory.payloads[-1].materialized is False + assert store.state is ProviderRegionStoreState.OPEN + factory.fail_construct.clear() + third = store.allocate_and_export(_allocation_spec()) + assert third.provider_resource_id == 3 + released = store.release(1) + assert released.status is ProviderReleaseStatus.RELEASED + gone = store.release(1) + assert gone.status is ProviderReleaseStatus.ALREADY_GONE + unknown = store.release(99) + assert unknown.status is ProviderReleaseStatus.UNKNOWN_RESOURCE + with pytest.raises(RegionControlError) as id_exc: + store.release(0) + assert id_exc.value.kind is RegionControlErrorKind.INVALID_FIELD_VALUE + + +def test_allocate_after_successful_release_uses_a_new_id_and_cleans_once(): + store, factory = _open_store() + first = store.allocate_and_export(_allocation_spec()) + assert store.release(first.provider_resource_id).status is ProviderReleaseStatus.RELEASED + assert store.state is ProviderRegionStoreState.OPEN + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + second = store.allocate_and_export(_allocation_spec()) + assert second.provider_resource_id == 2 + assert store.release(second.provider_resource_id).status is ProviderReleaseStatus.RELEASED + assert factory.payloads[1].release_count == 1 + assert factory.counters[1].release_count == 1 + assert factory.world.duplicate_releases == [] + assert store.state is ProviderRegionStoreState.OPEN + + +def test_release_attempts_both_parts_and_enters_close_failed_on_first_cleanup_debt(): + store, factory = _open_store() + store.allocate_and_export(_allocation_spec()) + factory.payloads[0].release_step_failures = [_backend_failure(RegionPartKind.PAYLOAD)] + result = store.release(1) + assert result.status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert result.failures[0].part is RegionPartKind.PAYLOAD + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert store.state is ProviderRegionStoreState.CLOSE_FAILED + with pytest.raises(RegionControlError) as exc_info: + store.describe(1) + assert exc_info.value.kind is RegionControlErrorKind.STORE_LIFECYCLE + with pytest.raises(RegionControlError): + store.local_view(1, RegionPartKind.PAYLOAD) + with pytest.raises(RegionControlError): + store.release(1) + with pytest.raises(RegionControlError): + store.allocate_and_export(_allocation_spec()) + + +def test_sweep_first_cleans_untouched_resources_and_does_not_reissue_failed_cleanup(): + store, factory = _open_store() + store.allocate_and_export(_allocation_spec()) + store.allocate_and_export(_allocation_spec()) + factory.payloads[0].release_step_failures = [_backend_failure(RegionPartKind.PAYLOAD)] + first = store.release(1) + assert first.status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert factory.payloads[1].release_count == 0 + results = store.sweep() + assert [item.provider_resource_id for item in results] == [1, 2] + assert results[0].status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert results[1].status is ProviderReleaseStatus.RELEASED + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert factory.payloads[1].release_count == 1 + assert factory.counters[1].release_count == 1 + assert store.state is ProviderRegionStoreState.CLOSE_FAILED + again = store.sweep() + assert [item.provider_resource_id for item in again] == [1] + assert again[0].status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert factory.payloads[0].release_count == 1 + assert factory.world.duplicate_releases == [] + + +def test_counter_materialize_failure_cleans_both_installed_shells_and_stays_open(): + factory = FakeShellFactory() + + def _mutate(shell, kind): + if kind is RegionPartKind.COUNTER: + shell.fail_materialize = RuntimeError("counter materialize") + + store = ProviderRegionStore(_sim_context(), _shell_factory=_mutating_shell_factory(factory, _mutate)) + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.control_kind is RegionControlErrorKind.BACKEND_FAILURE + assert error.failed_part is RegionPartKind.COUNTER + assert error.failed_operation is RegionOperationKind.MATERIALIZE + assert error.cleanup_debt_remaining is False + assert factory.payloads[0].materialized is True + assert factory.counters[0].materialized is False + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert store.state is ProviderRegionStoreState.OPEN + + +@pytest.mark.parametrize( + ("attr", "exc", "part", "operation"), + [ + ("fail_zero", RuntimeError("zero"), RegionPartKind.COUNTER, RegionOperationKind.ZERO_BYTES), + ("fail_mapping_bytes", RuntimeError("map"), RegionPartKind.PAYLOAD, RegionOperationKind.DESCRIBE), + ("fail_mapping_bytes", RuntimeError("map"), RegionPartKind.COUNTER, RegionOperationKind.DESCRIBE), + ("fail_import_capability", RuntimeError("import"), RegionPartKind.PAYLOAD, RegionOperationKind.DESCRIBE), + ("fail_import_capability", RuntimeError("import"), RegionPartKind.COUNTER, RegionOperationKind.DESCRIBE), + ("fail_local_base", RuntimeError("base"), RegionPartKind.PAYLOAD, RegionOperationKind.LOCAL_VIEW), + ("fail_local_base", RuntimeError("base"), RegionPartKind.COUNTER, RegionOperationKind.LOCAL_VIEW), + ("fail_materialize", KeyboardInterrupt(), RegionPartKind.PAYLOAD, RegionOperationKind.MATERIALIZE), + ], +) +def test_create_path_failures_are_classified_and_cleanup_runs_once(attr, exc, part, operation): + factory = FakeShellFactory() + + def _mutate(shell, kind): + if kind is part: + setattr(shell, attr, exc) + + store = ProviderRegionStore(_sim_context(), _shell_factory=_mutating_shell_factory(factory, _mutate)) + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.failed_part is part + assert error.failed_operation is operation + assert error.control_kind is RegionControlErrorKind.BACKEND_FAILURE + assert error.cleanup_debt_remaining is False + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert store.state is ProviderRegionStoreState.OPEN + + +def test_create_cleanup_failure_retains_record_and_close_failed(): + factory = FakeShellFactory() + + def _mutate(shell, kind): + if kind is RegionPartKind.COUNTER: + shell.fail_materialize = RuntimeError("counter materialize") + if kind is RegionPartKind.PAYLOAD: + shell.release_step_failures = [_backend_failure(RegionPartKind.PAYLOAD)] + + store = ProviderRegionStore(_sim_context(), _shell_factory=_mutating_shell_factory(factory, _mutate)) + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.cleanup_debt_remaining is True + assert error.failed_part is RegionPartKind.COUNTER + assert store.state is ProviderRegionStoreState.CLOSE_FAILED + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + snapshot = store.sweep() + assert snapshot[0].status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert factory.payloads[0].release_count == 1 + + +def test_fake_release_keeps_later_failures_local_and_summarizes_the_first(): + store, factory = _open_store() + store.allocate_and_export(_allocation_spec()) + first = _backend_failure(RegionPartKind.PAYLOAD) + later = ProviderCleanupFailure( + part=RegionPartKind.PAYLOAD, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.INTERRUPTED, + ) + factory.payloads[0].release_step_failures = [first, later] + result = store.release(1) + assert result.failures == (first,) + assert factory.payloads[0].local_ledger == [first, later] + + +def test_overlapping_local_views_are_an_internal_invariant_and_are_cleaned(): + factory = FakeShellFactory() + + def _mutate(shell, _kind): + shell._local_base = 0 + + store = ProviderRegionStore(_sim_context(), _shell_factory=_mutating_shell_factory(factory, _mutate)) + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.control_kind is RegionControlErrorKind.INTERNAL_INVARIANT + assert error.failed_operation is RegionOperationKind.LOCAL_VIEW + assert error.cleanup_debt_remaining is False + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + + +def test_closed_dispatcher_is_an_internal_invariant_and_burns_the_id(): + store = ProviderRegionStore( + RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=HostAllocationTarget(), + ) + ) + with pytest.raises(RegionAllocationError) as exc_info: + store.allocate_and_export(_allocation_spec()) + error = exc_info.value + assert error.provisional_resource_id == 1 + assert error.control_kind is RegionControlErrorKind.INTERNAL_INVARIANT + assert error.cleanup_debt_remaining is False + assert store.state is ProviderRegionStoreState.OPEN + with pytest.raises(RegionAllocationError) as second: + store.allocate_and_export(_allocation_spec()) + assert second.value.provisional_resource_id == 2 + + +def test_id_exhaustion_fails_before_a_side_effect_and_does_not_wrap(): + store, factory = _open_store() + store._next_provider_resource_id = _UINT64_MAX + result = store.allocate_and_export(_allocation_spec()) + assert result.provider_resource_id == _UINT64_MAX + with pytest.raises(RegionControlError) as exc_info: + store.allocate_and_export(_allocation_spec()) + assert exc_info.value.kind is RegionControlErrorKind.INTERNAL_INVARIANT + assert factory.world.constructed_parts == [RegionPartKind.PAYLOAD, RegionPartKind.COUNTER] + gone = store.release(1) + assert gone.status is ProviderReleaseStatus.ALREADY_GONE + + +def test_release_interrupt_on_payload_still_attempts_counter(): + store, factory = _open_store() + store.allocate_and_export(_allocation_spec()) + factory.payloads[0].raise_on_release = KeyboardInterrupt() + result = store.release(1) + assert result.status is ProviderReleaseStatus.CLEANUP_INCOMPLETE + assert result.failures[0].typed_cause is RegionCleanupCause.INTERRUPTED + assert factory.counters[0].release_count == 1 + assert store.state is ProviderRegionStoreState.CLOSE_FAILED + + +def test_describe_and_local_view_reject_gone_and_unknown_ids(): + store, _factory = _open_store() + store.allocate_and_export(_allocation_spec()) + store.release(1) + with pytest.raises(RegionControlError) as gone: + store.describe(1) + assert gone.value.kind is RegionControlErrorKind.INVALID_FIELD_VALUE + with pytest.raises(RegionControlError) as unknown: + store.local_view(8, RegionPartKind.PAYLOAD) + assert unknown.value.kind is RegionControlErrorKind.INVALID_FIELD_VALUE + with pytest.raises(RegionControlError) as zero: + store.describe(0) + assert zero.value.kind is RegionControlErrorKind.INVALID_FIELD_VALUE + + +def test_unmaterialized_fake_facts_fail_and_materialize_does_not_release(): + factory = FakeShellFactory() + spec = _allocation_spec() + payload = factory(_sim_context(), RegionPartKind.PAYLOAD, spec.payload) + with pytest.raises(RegionControlError) as exc_info: + payload.mapping_bytes() + assert exc_info.value.kind is RegionControlErrorKind.INTERNAL_INVARIANT + payload.materialize() + assert payload.release_count == 0 + assert payload.mapping_bytes() == spec.payload.logical_bytes + + +def test_successful_sweep_releases_active_resources_in_id_order(): + store, factory = _open_store() + store.allocate_and_export(_allocation_spec()) + store.allocate_and_export(_allocation_spec()) + results = store.sweep() + assert [item.status for item in results] == [ + ProviderReleaseStatus.RELEASED, + ProviderReleaseStatus.RELEASED, + ] + assert [item.provider_resource_id for item in results] == [1, 2] + assert store.state is ProviderRegionStoreState.CLOSED + assert factory.payloads[0].release_count == 1 + assert factory.payloads[1].release_count == 1 + with pytest.raises(RegionControlError): + store.describe(1) + + +def test_sim_posix_store_creates_two_distinct_named_objects_and_zeros_only_counter(): + store = ProviderRegionStore(_sim_context()) + result = store.allocate_and_export(_allocation_spec()) + try: + from simpler.comm_provider import _posix_shm_create_name + + payload_name = result.export_descriptor.payload.import_capability.shm_name + counter_name = result.export_descriptor.counter.import_capability.shm_name + assert payload_name != counter_name + payload_shm = SharedMemory(name=_posix_shm_create_name(payload_name)) + counter_shm = SharedMemory(name=_posix_shm_create_name(counter_name)) + try: + assert payload_shm.buf is not None and payload_shm.buf.nbytes >= 64 + assert counter_shm.buf is not None and counter_shm.buf.nbytes >= 8 + assert bytes(counter_shm.buf[:8]) == b"\x00" * 8 + payload_view = store.local_view(result.provider_resource_id, RegionPartKind.PAYLOAD) + counter_view = store.local_view(result.provider_resource_id, RegionPartKind.COUNTER) + validate_independent_local_views(payload_view, counter_view) + finally: + payload_shm.close() + counter_shm.close() + finally: + released = store.release(result.provider_resource_id) + assert released.status is ProviderReleaseStatus.RELEASED + + +def test_sim_posix_collision_does_not_open_or_unlink_the_existing_object(): + from simpler.comm_provider import _generate_posix_shm_token, _posix_shm_create_name + + token = _generate_posix_shm_token() + existing = SharedMemory(name=_posix_shm_create_name(token), create=True, size=8) + try: + assert existing.buf is not None + existing.buf[:8] = b"KEEPKEEP" + shell = SimPosixShmAllocation( + _sim_context(), + RegionPartKind.PAYLOAD, + _payload_spec(8), + candidate_name=token, + ) + with pytest.raises(RegionControlError) as exc_info: + shell.materialize() + assert exc_info.value.kind is RegionControlErrorKind.BACKEND_FAILURE + assert shell.release_once() is None + still_there = SharedMemory(name=_posix_shm_create_name(token)) + try: + assert still_there.buf is not None + assert bytes(still_there.buf[:8]) == b"KEEPKEEP" + finally: + still_there.close() + finally: + existing.close() + existing.unlink() + + +def test_sim_posix_overlong_token_is_rejected_before_creation(): + with pytest.raises(ValueError): + SimPosixShmAllocation( + _sim_context(), + RegionPartKind.PAYLOAD, + _payload_spec(), + candidate_name="n" * (POSIX_SHM_TOKEN_MAX_BYTES + 1), + ) + + +def test_sim_posix_interrupted_create_unlinks_owned_name_once(): + from simpler.comm_provider import _generate_posix_shm_token, _posix_shm_create_name + + token = _generate_posix_shm_token() + created: dict[str, SharedMemory] = {} + + class _CreateThenInterrupt: + def __init__(self, name, *, create, size, **kwargs): + created["shm"] = SharedMemory(name=name, create=create, size=size) + raise KeyboardInterrupt() + + shell = SimPosixShmAllocation( + _sim_context(), + RegionPartKind.COUNTER, + _counter_spec(), + candidate_name=token, + shm_cls=_CreateThenInterrupt, + ) + with pytest.raises(KeyboardInterrupt): + shell.materialize() + failure = shell.release_once() + assert failure is None + assert shell.release_once() is None + with pytest.raises(FileNotFoundError): + SharedMemory(name=_posix_shm_create_name(token)) + if "shm" in created: + created["shm"].close() + + +def test_sim_posix_close_failure_still_unlinks_and_keeps_later_detail(): + from simpler.comm_provider import _posix_shm_create_name + + shell = SimPosixShmAllocation(_sim_context(), RegionPartKind.PAYLOAD, _payload_spec(8)) + shell.materialize() + token = shell.candidate_name + + def _boom_close(): + raise OSError("close failed") + + assert shell._shm is not None + shell._shm.close = _boom_close # type: ignore[method-assign] + failure = shell.release_once() + assert failure is not None + assert failure.typed_cause is RegionCleanupCause.BACKEND_ERROR + assert [step for step, _exc in shell.local_cleanup_details] == ["close"] + assert shell.release_once() is failure + with pytest.raises(FileNotFoundError): + SharedMemory(name=_posix_shm_create_name(token)) + + +def test_sim_posix_unlink_failure_is_summarized_after_successful_close(): + shell = SimPosixShmAllocation(_sim_context(), RegionPartKind.COUNTER, _counter_spec()) + shell.materialize() + + def _boom_unlink(): + raise OSError("unlink failed") + + assert shell._shm is not None + shell._shm.unlink = _boom_unlink # type: ignore[method-assign] + failure = shell.release_once() + assert failure is not None + assert [step for step, _exc in shell.local_cleanup_details] == ["unlink"] + shell._name_ownership_known = False + try: + if shell._shm is not None: + SharedMemory.unlink(shell._shm) + except FileNotFoundError: + pass + + +@pytest.mark.parametrize("platform", ["linux", "darwin"]) +def test_posix_token_has_no_slash_on_linux_and_macos(monkeypatch, platform): + from simpler.comm_provider import _generate_posix_shm_token, _posix_shm_create_name + + monkeypatch.setattr(sys, "platform", platform) + token = _generate_posix_shm_token() + assert token.startswith("smp_") + assert "/" not in token + assert _posix_shm_create_name(token) == token + shell = SimPosixShmAllocation(_sim_context(), RegionPartKind.PAYLOAD, _payload_spec(8, BackendKind.POSIX_SHM)) + try: + shell.materialize() + capability = shell.import_capability() + assert "/" not in shell.candidate_name + assert capability.shm_name == shell.candidate_name + assert "/" not in capability.shm_name + consumer = SharedMemory(name=_posix_shm_create_name(capability.shm_name)) + try: + assert consumer.buf is not None + finally: + consumer.close() + finally: + assert shell.release_once() is None + + +def test_posix_token_does_not_use_shared_memory_prepend_attr(): + assert "_prepend_leading_slash" not in _COMM_PROVIDER_PATH.read_text() + + +def test_posix_unlink_and_tracker_use_the_same_single_leading_slash(monkeypatch): + from multiprocessing import resource_tracker + + from simpler.comm_provider import _unlink_posix_shm_token + + unlinks: list[str] = [] + unregisters: list[tuple[str, str]] = [] + + class _FakePosix: + @staticmethod + def shm_unlink(name: str) -> None: + unlinks.append(name) + + monkeypatch.setitem(sys.modules, "_posixshmem", _FakePosix) + monkeypatch.setattr( + resource_tracker, + "unregister", + lambda name, kind: unregisters.append((name, kind)), + ) + _unlink_posix_shm_token("smp_deadbeefcafebabe01234567") + assert unlinks == ["/smp_deadbeefcafebabe01234567"] + assert unregisters == [("/smp_deadbeefcafebabe01234567", "shared_memory")] + + +def test_posix_repeated_unlink_treats_enoent_as_done(monkeypatch): + from multiprocessing import resource_tracker + + from simpler.comm_provider import _unlink_posix_shm_token + + class _FakePosix: + calls = 0 + + @staticmethod + def shm_unlink(name: str) -> None: + _FakePosix.calls += 1 + if _FakePosix.calls > 1: + raise FileNotFoundError(name) + + monkeypatch.setitem(sys.modules, "_posixshmem", _FakePosix) + monkeypatch.setattr(resource_tracker, "unregister", lambda *_args: None) + _unlink_posix_shm_token("smp_deadbeefcafebabe01234567") + _unlink_posix_shm_token("smp_deadbeefcafebabe01234567") + assert _FakePosix.calls == 2 + + +def test_posix_unregister_failure_warns_without_changing_result(monkeypatch, caplog): + from multiprocessing import resource_tracker + + from simpler.comm_provider import _unlink_posix_shm_token + + class _FakePosix: + @staticmethod + def shm_unlink(name: str) -> None: + del name + + monkeypatch.setitem(sys.modules, "_posixshmem", _FakePosix) + monkeypatch.setattr( + resource_tracker, + "unregister", + lambda *_args: (_ for _ in ()).throw(RuntimeError("tracker boom")), + ) + caplog.set_level(logging.WARNING, logger="simpler") + _unlink_posix_shm_token("smp_deadbeefcafebabe01234567") + assert "smp_deadbeefcafebabe01234567" in caplog.text + assert "tracker boom" in caplog.text + + +def test_sim_posix_first_cleanup_failure_is_close_when_both_steps_fail(): + from simpler.comm_provider import _posix_shm_create_name + + shell = SimPosixShmAllocation(_sim_context(), RegionPartKind.PAYLOAD, _payload_spec(8)) + shell.materialize() + token = shell.candidate_name + assert shell._shm is not None + real_unlink = shell._shm.unlink + + def _boom_close(): + raise OSError("close failed") + + def _boom_unlink(): + raise OSError("unlink failed") + + shell._shm.close = _boom_close # type: ignore[method-assign] + shell._shm.unlink = _boom_unlink # type: ignore[method-assign] + failure = shell.release_once() + assert failure is not None + assert [step for step, _exc in shell.local_cleanup_details] == ["close", "unlink"] + try: + real_unlink() + except FileNotFoundError: + pass + with pytest.raises(FileNotFoundError): + SharedMemory(name=_posix_shm_create_name(token)) + + +def _onboard_context(device_id: int = 0) -> RegionAllocationContext: + return RegionAllocationContext( + environment_kind=RegionEnvironmentKind.ONBOARD, + target=DeviceAllocationTarget(device_id=device_id), + ) + + +@pytest.fixture +def fake_vmm(): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_inspect, + _region_vmm_test_live_handles, + _region_vmm_test_reset_hooks, + _region_vmm_test_use_fake_driver, + ) + + _region_vmm_test_use_fake_driver() + _region_vmm_test_reset_hooks() + before = set(_region_vmm_test_live_handles()) + yield + _region_vmm_test_reset_hooks() + leftover = [handle for handle in _region_vmm_test_live_handles() if handle not in before] + unclean = [handle for handle in leftover if not _region_vmm_inspect(handle).release_once_done] + assert unclean == [], f"fake VMM leaked live handles: {unclean}" + + +def test_region_vmm_begin_has_no_physical_side_effect(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_issued_ops, + _region_vmm_test_live_handles, + ) + + before = set(_region_vmm_test_live_handles()) + handle = _region_vmm_begin(3) + inspect = _region_vmm_inspect(handle) + try: + assert handle not in before + assert inspect.present is True + assert inspect.device_id == 3 + assert inspect.physical_allocated is False + assert inspect.va_reserved is False + assert inspect.mapped is False + assert inspect.exported is False + assert inspect.mapping_bytes == 0 + assert _region_vmm_test_issued_ops() == [] + assert handle in _region_vmm_test_live_handles() + finally: + _region_vmm_release(handle) + assert _region_vmm_inspect(handle).present is False + + +@pytest.mark.parametrize( + ("stage", "when", "allocated", "reserved", "mapped", "exported"), + [ + ("physical_alloc", "before", False, False, False, False), + ("physical_alloc", "after", True, False, False, False), + ("va_reserve", "before", True, False, False, False), + ("va_reserve", "after", True, True, False, False), + ("map", "before", True, True, False, False), + ("map", "after", True, True, True, False), + ("set_access", "before", True, True, True, False), + ("set_access", "after", True, True, True, False), + ("export", "before", True, True, True, False), + ("export", "after", True, True, True, True), + ], +) +def test_region_vmm_allocate_failure_retains_record_without_rollback( + fake_vmm, stage, when, allocated, reserved, mapped, exported +): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(0) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage(stage, when) + with pytest.raises(RuntimeError, match="injected region VMM"): + _region_vmm_allocate_export(handle, 8) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + try: + assert inspect.present is True + assert inspect.physical_allocated is allocated + assert inspect.va_reserved is reserved + assert inspect.mapped is mapped + assert inspect.exported is exported + if when == "before": + assert stage not in issued + else: + assert stage in issued + finally: + _region_vmm_test_reset_hooks() + _region_vmm_release(handle) + assert _region_vmm_inspect(handle).present is False + + +def test_region_vmm_failed_unmap_skips_va_and_physical(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(0) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("unmap", "before") + with pytest.raises(RuntimeError, match="before unmap"): + _region_vmm_release(handle) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + assert inspect.present is True + assert inspect.mapped is True + assert inspect.va_reserved is True + assert inspect.physical_allocated is True + assert inspect.unmap_attempted is True + assert inspect.unmap_complete is False + assert inspect.release_once_done is True + assert "unmap" not in issued + assert "va_release" not in issued + assert "physical_free" not in issued + _region_vmm_test_reset_hooks() + with pytest.raises(RuntimeError, match="before unmap"): + _region_vmm_release(handle) + assert _region_vmm_test_issued_ops() == [] + assert _region_vmm_inspect(handle).mapped is True + + +def test_region_vmm_successful_unmap_attempts_va_and_physical_independently(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(0) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("va_release", "before") + with pytest.raises(RuntimeError, match="before va_release"): + _region_vmm_release(handle) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + assert inspect.present is True + assert inspect.mapped is False + assert inspect.unmap_complete is True + assert inspect.va_reserved is True + assert inspect.physical_allocated is False + assert inspect.physical_free_complete is True + assert inspect.first_cleanup_failure + assert inspect.local_cleanup_details == [inspect.first_cleanup_failure] + assert issued.count("unmap") == 1 + assert "va_release" not in issued + assert issued.count("physical_free") == 1 + _region_vmm_test_reset_hooks() + with pytest.raises(RuntimeError, match="before va_release"): + _region_vmm_release(handle) + assert _region_vmm_test_issued_ops() == [] + + +def test_region_vmm_already_gone_unmap_still_releases_va_and_physical(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + _region_vmm_test_set_unmap_already_gone, + ) + + handle = _region_vmm_begin(0) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_set_unmap_already_gone() + _region_vmm_release(handle) + assert _region_vmm_inspect(handle).present is False + assert _region_vmm_test_issued_ops() == ["bind_device", "unmap", "va_release", "physical_free"] + + +def test_region_vmm_two_handles_are_independent(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + ) + + first = _region_vmm_allocate_export(_region_vmm_begin(1), 8) + second = _region_vmm_allocate_export(_region_vmm_begin(2), 64) + try: + assert first.registry_handle != second.registry_handle + assert first.shareable_handle != second.shareable_handle + assert first.device_addr != second.device_addr + assert first.mapping_bytes == 64 + assert second.mapping_bytes == 64 + assert _region_vmm_inspect(first.registry_handle).device_id == 1 + assert _region_vmm_inspect(second.registry_handle).device_id == 2 + finally: + _region_vmm_release(first.registry_handle) + _region_vmm_release(second.registry_handle) + + +def test_region_vmm_two_handles_use_independent_mapping_sizes(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_release, + ) + + first = _region_vmm_allocate_export(_region_vmm_begin(1), 64) + second = _region_vmm_allocate_export(_region_vmm_begin(2), 128) + try: + assert first.registry_handle != second.registry_handle + assert first.shareable_handle != second.shareable_handle + assert first.mapping_bytes == 64 + assert second.mapping_bytes == 128 + finally: + _region_vmm_release(first.registry_handle) + _region_vmm_release(second.registry_handle) + + +def test_region_vmm_later_independent_errors_stay_local(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(0) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("va_release", "after") + with pytest.raises(RuntimeError, match="after va_release"): + _region_vmm_release(handle) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + assert inspect.present is True + assert inspect.unmap_complete is True + assert inspect.va_release_complete is True + assert inspect.physical_free_complete is True + assert inspect.physical_allocated is False + assert inspect.first_cleanup_failure + assert inspect.local_cleanup_details == [inspect.first_cleanup_failure] + assert "after va_release" in inspect.first_cleanup_failure + assert issued.count("unmap") == 1 + assert issued.count("va_release") == 1 + assert issued.count("physical_free") == 1 + _region_vmm_test_reset_hooks() + with pytest.raises(RuntimeError, match="after va_release"): + _region_vmm_release(handle) + assert _region_vmm_test_issued_ops() == [] + + +def test_region_vmm_physical_free_failure_retains_record_without_reissue(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(0) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("physical_free", "before") + with pytest.raises(RuntimeError, match="before physical_free"): + _region_vmm_release(handle) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + assert inspect.present is True + assert inspect.unmap_complete is True + assert inspect.va_reserved is False + assert inspect.physical_allocated is True + assert inspect.physical_free_attempted is True + assert inspect.physical_free_complete is False + assert "physical_free" not in issued + assert issued.count("unmap") == 1 + assert issued.count("va_release") == 1 + _region_vmm_test_reset_hooks() + with pytest.raises(RuntimeError, match="before physical_free"): + _region_vmm_release(handle) + assert _region_vmm_test_issued_ops() == [] + assert _region_vmm_inspect(handle).physical_allocated is True + + +def test_region_vmm_release_bind_failure_keeps_ledger_and_does_not_retry(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + ) + + handle = _region_vmm_begin(3) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("bind_device", "before") + with pytest.raises(RuntimeError, match="bind_device failed:") as first: + _region_vmm_release(handle) + inspect = _region_vmm_inspect(handle) + issued = _region_vmm_test_issued_ops() + assert inspect.present is True + assert inspect.mapped is True + assert inspect.va_reserved is True + assert inspect.physical_allocated is True + assert inspect.unmap_attempted is False + assert inspect.va_release_attempted is False + assert inspect.physical_free_attempted is False + assert inspect.release_once_done is True + assert inspect.first_cleanup_failure == str(first.value) + assert inspect.first_cleanup_failure.startswith("bind_device failed:") + assert inspect.local_cleanup_details == [inspect.first_cleanup_failure] + assert issued == ["bind_device"] + _region_vmm_test_reset_hooks() + with pytest.raises(RuntimeError, match="bind_device failed:") as second: + _region_vmm_release(handle) + assert str(second.value) == str(first.value) + assert _region_vmm_test_issued_ops() == [] + later = _region_vmm_inspect(handle) + assert later.mapped is True + assert later.va_reserved is True + assert later.physical_allocated is True + assert later.first_cleanup_failure == inspect.first_cleanup_failure + + +def test_region_vmm_zero_binds_record_device_before_zero(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + _region_vmm_zero_bytes, + ) + + handle = _region_vmm_begin(4) + _region_vmm_allocate_export(handle, 8) + other = _region_vmm_allocate_export(_region_vmm_begin(9), 8) + _region_vmm_test_reset_hooks() + _region_vmm_zero_bytes(handle, 0, 8) + assert _region_vmm_test_issued_ops() == ["bind_device", "zero_bytes"] + assert _region_vmm_inspect(handle).device_id == 4 + assert _region_vmm_inspect(other.registry_handle).device_id == 9 + _region_vmm_release(handle) + _region_vmm_release(other.registry_handle) + + +def test_region_vmm_zero_bind_failure_does_not_issue_zero(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_allocate_export, + _region_vmm_begin, + _region_vmm_inspect, + _region_vmm_release, + _region_vmm_test_fail_stage, + _region_vmm_test_issued_ops, + _region_vmm_test_reset_hooks, + _region_vmm_zero_bytes, + ) + + handle = _region_vmm_begin(4) + _region_vmm_allocate_export(handle, 8) + _region_vmm_test_reset_hooks() + _region_vmm_test_fail_stage("bind_device", "before") + with pytest.raises(RuntimeError, match="injected region VMM before bind_device"): + _region_vmm_zero_bytes(handle, 0, 8) + assert _region_vmm_test_issued_ops() == ["bind_device"] + inspect = _region_vmm_inspect(handle) + assert inspect.present is True + assert inspect.mapped is True + _region_vmm_test_reset_hooks() + _region_vmm_release(handle) + + +def test_vmm_allocation_stores_handle_before_allocate_and_zeros_only_logical_bytes(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_inspect, + _region_vmm_test_fail_stage, + _region_vmm_test_fake_bytes, + _region_vmm_test_reset_hooks, + ) + + shell = VmmAllocation(_onboard_context(4), RegionPartKind.COUNTER, _counter_spec(8)) + assert shell.registry_handle is None + _region_vmm_test_fail_stage("export", "before") + with pytest.raises(RegionControlError) as exc_info: + shell.materialize() + assert exc_info.value.kind is RegionControlErrorKind.BACKEND_FAILURE + handle = shell.registry_handle + assert handle is not None + inspect = _region_vmm_inspect(handle) + assert inspect.present is True + assert inspect.mapped is True + assert inspect.exported is False + _region_vmm_test_reset_hooks() + assert shell.release_once() is None + assert _region_vmm_inspect(handle).present is False + assert shell.release_once() is None + + ready = VmmAllocation(_onboard_context(4), RegionPartKind.COUNTER, _counter_spec(8)) + ready.materialize() + try: + ready.zero_bytes(0, 8) + backing = bytes(_region_vmm_test_fake_bytes(ready.registry_handle)) + assert backing[:8] == b"\x00" * 8 + assert backing[8:] == b"\xff" * (len(backing) - 8) + capability = ready.import_capability() + assert capability.device_id == 4 + assert ready.mapping_bytes() == 64 + assert ready.local_base() == int(_region_vmm_inspect(ready.registry_handle).device_addr) + finally: + assert ready.release_once() is None + + +def test_onboard_store_creates_two_independent_vmm_handles_and_zeros_only_counter(fake_vmm): + from _task_interface import ( # pyright: ignore[reportMissingImports] + _region_vmm_inspect, + _region_vmm_test_fake_bytes, + ) + + store = ProviderRegionStore(_onboard_context(5)) + result = store.allocate_and_export(_allocation_spec()) + resource = store._resources[result.provider_resource_id] + payload_shell = resource.parts[RegionPartKind.PAYLOAD].allocation + counter_shell = resource.parts[RegionPartKind.COUNTER].allocation + assert isinstance(payload_shell, VmmAllocation) + assert isinstance(counter_shell, VmmAllocation) + payload_handle = payload_shell.registry_handle + counter_handle = counter_shell.registry_handle + try: + payload = result.export_descriptor.payload.import_capability + counter = result.export_descriptor.counter.import_capability + assert isinstance(payload, VmmShareableHandleImport) + assert isinstance(counter, VmmShareableHandleImport) + assert payload.shareable_handle != counter.shareable_handle + assert payload.device_id == 5 + assert payload_handle != counter_handle + payload_view = store.local_view(result.provider_resource_id, RegionPartKind.PAYLOAD) + counter_view = store.local_view(result.provider_resource_id, RegionPartKind.COUNTER) + validate_independent_local_views(payload_view, counter_view) + payload_bytes = bytes(_region_vmm_test_fake_bytes(payload_handle)) + counter_bytes = bytes(_region_vmm_test_fake_bytes(counter_handle)) + assert payload_bytes[:64] == b"\xff" * 64 + assert counter_bytes[:8] == b"\x00" * 8 + assert counter_bytes[8:] == b"\xff" * (len(counter_bytes) - 8) + finally: + released = store.release(result.provider_resource_id) + assert released.status is ProviderReleaseStatus.RELEASED + assert _region_vmm_inspect(payload_handle).present is False + assert _region_vmm_inspect(counter_handle).present is False + + +def test_closed_dispatcher_routes_onboard_vmm_window_to_vmm_allocation(): + payload = comm_provider_module._closed_part_dispatcher( + _onboard_context(), + RegionPartKind.PAYLOAD, + _payload_spec(), + ) + sim = comm_provider_module._closed_part_dispatcher( + _sim_context(), + RegionPartKind.PAYLOAD, + _payload_spec(), + ) + assert isinstance(payload, VmmAllocation) + assert isinstance(sim, SimPosixShmAllocation) + assert payload.registry_handle is None diff --git a/tests/ut/py/test_worker/test_comm_provider_control.py b/tests/ut/py/test_worker/test_comm_provider_control.py new file mode 100644 index 0000000000..66bb8c3714 --- /dev/null +++ b/tests/ut/py/test_worker/test_comm_provider_control.py @@ -0,0 +1,845 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 +"""Unit tests for the private provider-region control wire.""" + +from __future__ import annotations + +import logging +import struct +from multiprocessing.shared_memory import SharedMemory + +import pytest +from simpler.buffer import BackendKind +from simpler.comm_provider import ( + PosixShmImport, + ProviderCleanupFailure, + ProviderRegionStore, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationError, + RegionAllocationResult, + RegionCleanupCause, + RegionControlError, + RegionControlErrorKind, + RegionExportDescriptor, + RegionOperationKind, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, + VmmShareableHandleImport, +) +from simpler.comm_provider_control import ( + ALLOCATE_COUNTER_EXPORT_OFFSET, + ALLOCATE_COUNTER_VIEW_OFFSET, + ALLOCATE_PAYLOAD_EXPORT_OFFSET, + ALLOCATE_PAYLOAD_VIEW_OFFSET, + ALLOCATE_REPLY_BYTES, + ALLOCATE_REQUEST_BYTES, + COMMIT_TAG_OFFSET, + PROVIDER_REGION_CTRL_MAGIC_VERSION, + RELEASE_COUNTER_FAILURE_OFFSET, + RELEASE_PAYLOAD_FAILURE_OFFSET, + RELEASE_REPLY_BYTES, + RELEASE_REQUEST_BYTES, + AllocateReplyTag, + ImportCapabilityWireKind, + ProviderAllocateClient, + ProviderReleaseClient, + RegionControlProtocolError, + ReleaseReplyTag, + _discard_control_shm, + decode_allocate_reply, + decode_allocate_request, + decode_release_reply, + decode_release_request, + encode_allocate_allocation_error_reply, + encode_allocate_request, + encode_allocate_request_error_reply, + encode_allocate_success_reply, + encode_release_error_reply, + encode_release_request, + encode_release_result_reply, + handle_ctrl_region_allocate, + handle_ctrl_region_release, + peek_allocate_reply_resource_id, +) + +from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _allocation_spec, _sim_context + +_OLD_MAGIC = 0x4C334C3200020000 + + +def _posix_part(logical_bytes: int, name: str, mapping_bytes: int | None = None) -> RegionPartExportDescriptor: + return RegionPartExportDescriptor( + planned_backing_kind=BackendKind.VMM_WINDOW, + logical_bytes=logical_bytes, + mapping_bytes=logical_bytes if mapping_bytes is None else mapping_bytes, + import_capability=PosixShmImport(shm_name=name), + ) + + +def _vmm_part(logical_bytes: int, handle: int = 9, mapping_bytes: int | None = None) -> RegionPartExportDescriptor: + return RegionPartExportDescriptor( + planned_backing_kind=BackendKind.VMM_WINDOW, + logical_bytes=logical_bytes, + mapping_bytes=logical_bytes if mapping_bytes is None else mapping_bytes, + import_capability=VmmShareableHandleImport(device_id=2, shareable_handle=handle), + ) + + +def _success_result() -> tuple[RegionAllocationResult, RegionPartLocalView, RegionPartLocalView]: + result = RegionAllocationResult( + provider_resource_id=11, + export_descriptor=RegionExportDescriptor(payload=_posix_part(64, "/pto_payload_a"), counter=_vmm_part(8, 21)), + ) + payload = RegionPartLocalView(part=RegionPartKind.PAYLOAD, local_base=0x1000, logical_bytes=64) + counter = RegionPartLocalView(part=RegionPartKind.COUNTER, local_base=0x2000, logical_bytes=8) + return result, payload, counter + + +def test_allocate_request_is_exactly_48_bytes_at_normative_offsets(): + spec = _allocation_spec() + buf = bytearray(ALLOCATE_REQUEST_BYTES) + encode_allocate_request(buf, spec) + magic, size, payload_kind, counter_kind, reserved20, payload_bytes, counter_bytes, reserved40 = struct.unpack_from( + " None: + self.store = store + self.allocate_calls = 0 + self.release_calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id + self.allocate_calls += 1 + req = SharedMemory(name=request_shm_name) + reply = SharedMemory(name=reply_shm_name) + assert req.buf is not None + assert reply.buf is not None + req_view = memoryview(req.buf)[:ALLOCATE_REQUEST_BYTES] + reply_view = memoryview(reply.buf)[:ALLOCATE_REPLY_BYTES] + try: + handle_ctrl_region_allocate(req_view, reply_view, self.store) + finally: + req_view.release() + reply_view.release() + req.close() + reply.close() + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id + self.release_calls += 1 + req = SharedMemory(name=request_shm_name) + reply = SharedMemory(name=reply_shm_name) + assert req.buf is not None + assert reply.buf is not None + req_view = memoryview(req.buf)[:RELEASE_REQUEST_BYTES] + reply_view = memoryview(reply.buf)[:RELEASE_REPLY_BYTES] + try: + handle_ctrl_region_release(req_view, reply_view, self.store) + finally: + req_view.release() + reply_view.release() + req.close() + reply.close() + + +def test_handler_roundtrip_calls_store_once_and_clients_are_one_shot(): + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + mailbox = _LoopbackMailbox(store) + allocated, payload, counter = ProviderAllocateClient(mailbox, 3).allocate(_allocation_spec()) + assert mailbox.allocate_calls == 1 + assert allocated.provider_resource_id == 1 + assert payload.logical_bytes == 64 + assert counter.logical_bytes == 8 + client = ProviderReleaseClient(mailbox, 3) + released = client.release(allocated.provider_resource_id) + assert released.status is ProviderReleaseStatus.RELEASED + assert mailbox.release_calls == 1 + again = client.release(allocated.provider_resource_id) + assert again.status is ProviderReleaseStatus.RELEASED + assert mailbox.release_calls == 1 + + +def test_handler_request_error_does_not_call_store(): + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + req = bytearray(ALLOCATE_REQUEST_BYTES) + encode_allocate_request(req, _allocation_spec()) + struct.pack_into(" None: + self.calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + raise AssertionError("allocate must not be called") + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + self.calls += 1 + raise RuntimeError("mailbox down") + + mailbox = _BoomMailbox() + client = ProviderReleaseClient(mailbox, 1) + with pytest.raises(RegionControlProtocolError): + client.release(4) + with pytest.raises(RegionControlProtocolError): + client.release(4) + assert mailbox.calls == 1 + + +@pytest.mark.parametrize("failure_timing", ["before_reply", "after_touching_reply"]) +def test_allocate_client_transport_failure_does_not_call_store_or_release(failure_timing): + class _BoomMailbox: + def __init__(self) -> None: + self.allocate_calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name + self.allocate_calls += 1 + if failure_timing == "after_touching_reply": + reply = SharedMemory(name=reply_shm_name) + reply.close() + raise RuntimeError("mailbox down") + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("allocate transport failure must not issue release") + + mailbox = _BoomMailbox() + client = ProviderAllocateClient(mailbox, 1) + with pytest.raises(RuntimeError, match="mailbox down"): + client.allocate(_allocation_spec()) + assert mailbox.allocate_calls == 1 + assert client.dispatch_started is True + assert client.committed_resource_id == 0 + + +def test_allocate_client_missing_commit_is_a_decode_failure(): + class _SilentMailbox: + def __init__(self) -> None: + self.allocate_calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + self.allocate_calls += 1 + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("uncommitted allocate must not issue release") + + mailbox = _SilentMailbox() + client = ProviderAllocateClient(mailbox, 1) + with pytest.raises(RegionControlProtocolError, match="commit"): + client.allocate(_allocation_spec()) + assert mailbox.allocate_calls == 1 + assert client.dispatch_started is True + assert client.committed_resource_id == 0 + + +def test_allocate_client_old_version_reply_is_a_decode_failure(): + class _OldVersionMailbox: + def __init__(self) -> None: + self.allocate_calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name + self.allocate_calls += 1 + reply = SharedMemory(name=reply_shm_name) + try: + assert reply.buf is not None + struct.pack_into(" None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("old-version allocate must not issue release") + + mailbox = _OldVersionMailbox() + client = ProviderAllocateClient(mailbox, 1) + with pytest.raises(RegionControlError) as exc_info: + client.allocate(_allocation_spec()) + assert exc_info.value.kind is RegionControlErrorKind.BAD_MAGIC_VERSION + assert mailbox.allocate_calls == 1 + assert client.dispatch_started is True + assert client.committed_resource_id == 0 + + +def test_allocate_client_uncommitted_reply_does_not_call_store_again(): + class _EmptyReplyMailbox: + def __init__(self) -> None: + self.calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + self.calls += 1 + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + raise AssertionError("release must not be called") + + mailbox = _EmptyReplyMailbox() + client = ProviderAllocateClient(mailbox, 1) + with pytest.raises(RegionControlProtocolError, match="commit"): + client.allocate(_allocation_spec()) + with pytest.raises(RegionControlProtocolError, match="commit"): + client.allocate(_allocation_spec()) + assert mailbox.calls == 2 + assert client.dispatch_started is True + assert client.committed_resource_id == 0 + + +class _CommitThenRaiseMailbox(_LoopbackMailbox): + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + super().control_region_allocate(worker_id, request_shm_name, reply_shm_name) + raise RuntimeError("mailbox down after commit") + + +class _MalformedSuccessMailbox: + def __init__(self) -> None: + self.allocate_calls = 0 + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name + self.allocate_calls += 1 + reply = SharedMemory(name=reply_shm_name) + try: + assert reply.buf is not None + struct.pack_into(" None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("malformed SUCCESS must not issue release") + + +def test_allocate_client_records_committed_id_when_mailbox_raises_after_handler(): + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + mailbox = _CommitThenRaiseMailbox(store) + client = ProviderAllocateClient(mailbox, 3) + with pytest.raises(RuntimeError, match="mailbox down after commit"): + client.allocate(_allocation_spec()) + assert client.dispatch_started is True + assert client.committed_resource_id == 1 + + +def test_allocate_client_malformed_success_does_not_commit_an_id(): + mailbox = _MalformedSuccessMailbox() + client = ProviderAllocateClient(mailbox, 1) + with pytest.raises(RegionControlError): + client.allocate(_allocation_spec()) + assert client.dispatch_started is True + assert client.committed_resource_id == 0 + + +def test_allocate_client_encode_failure_does_not_start_dispatch(monkeypatch): + from simpler import comm_provider_control as control + + def _boom(*_args, **_kwargs): + raise RuntimeError("encode failed") + + class _UnusedMailbox: + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("encode failure must not dispatch") + + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + del worker_id, request_shm_name, reply_shm_name + raise AssertionError("encode failure must not issue release") + + monkeypatch.setattr(control, "encode_allocate_request", _boom) + client = ProviderAllocateClient(_UnusedMailbox(), 1) + with pytest.raises(RuntimeError, match="encode failed"): + client.allocate(_allocation_spec()) + assert client.dispatch_started is False + assert client.committed_resource_id == 0 + + +def test_handler_success_does_not_release_the_active_resource(): + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + req = bytearray(ALLOCATE_REQUEST_BYTES) + encode_allocate_request(req, _allocation_spec()) + reply = bytearray(ALLOCATE_REPLY_BYTES) + handle_ctrl_region_allocate(memoryview(req), memoryview(reply), store) + assert struct.unpack_from(" None: + self.name = name + self.ops: list[str] = [] + self._close_error = close_error + self._unlink_error = unlink_error + + def close(self) -> None: + self.ops.append("close") + if self._close_error is not None: + raise self._close_error + + def unlink(self) -> None: + self.ops.append("unlink") + if self._unlink_error is not None: + raise self._unlink_error + + +def test_discard_control_shm_unlinks_after_close_failure(caplog): + shm = _RecordingShm("pto_close_fail", close_error=BufferError("still mapped")) + with caplog.at_level(logging.WARNING, logger="simpler"): + _discard_control_shm(shm) + assert shm.ops == ["close", "unlink"] + assert "pto_close_fail" in caplog.text + assert "close" in caplog.text + assert "still mapped" in caplog.text + + +def test_discard_control_shm_logs_unlink_failure_once(caplog): + shm = _RecordingShm("pto_unlink_fail", unlink_error=OSError("busy")) + with caplog.at_level(logging.WARNING, logger="simpler"): + _discard_control_shm(shm) + assert shm.ops == ["close", "unlink"] + assert "pto_unlink_fail" in caplog.text + assert "unlink" in caplog.text + assert "busy" in caplog.text + + +def test_discard_control_shm_treats_missing_unlink_as_done(): + shm = _RecordingShm("pto_missing", unlink_error=FileNotFoundError("gone")) + _discard_control_shm(shm) + assert shm.ops == ["close", "unlink"] + + +def test_allocate_and_release_survive_control_shm_cleanup_failure(monkeypatch, caplog): + from simpler import comm_provider_control as control + + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + mailbox = _LoopbackMailbox(store) + original = control._discard_control_shm + ops: list[str] = [] + + def _failing_discard(shm): + wrapper = _RecordingShm(shm.name, close_error=BufferError("still mapped"), unlink_error=OSError("busy")) + original(wrapper) + ops.extend(wrapper.ops) + original(shm) + + monkeypatch.setattr(control, "_discard_control_shm", _failing_discard) + with caplog.at_level(logging.WARNING, logger="simpler"): + allocated, _payload, _counter = ProviderAllocateClient(mailbox, 3).allocate(_allocation_spec()) + released = ProviderReleaseClient(mailbox, 3).release(allocated.provider_resource_id) + assert allocated.provider_resource_id == 1 + assert released.status is ProviderReleaseStatus.RELEASED + assert ops == ["close", "unlink", "close", "unlink", "close", "unlink", "close", "unlink"] + assert caplog.text.count("operation=close") == 4 + assert caplog.text.count("operation=unlink") == 4 diff --git a/tests/ut/py/test_worker/test_comm_region.py b/tests/ut/py/test_worker/test_comm_region.py index 509cf6dfc6..cd84eecdd0 100644 --- a/tests/ut/py/test_worker/test_comm_region.py +++ b/tests/ut/py/test_worker/test_comm_region.py @@ -6,13 +6,17 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 """Unit tests for the private region materializer.""" import dataclasses +import struct +from multiprocessing.shared_memory import SharedMemory from typing import Any, Optional, Union import pytest from simpler import comm_endpoints as ce +from simpler.comm_provider import RegionAllocationSpec, RegionPartAllocationSpec, RegionPartKind from simpler.comm_region import ( CounterPart, HostVmmCopyAccess, @@ -24,10 +28,12 @@ RefusalReason, RegionCounter, RegionInstance, + RegionInstanceRegistry, RegionInstanceState, RegionPartSpan, SignalTestResult, WaitCmp, + project_region_allocation_spec, validate_single_owner_region_shape, ) from simpler.orchestrator import _callback_frame_for, _callback_run @@ -92,6 +98,10 @@ def _materialize_default_region(worker: Worker): ) +def _tracked(worker: Worker) -> tuple[RegionInstance, ...]: + return tuple(worker._region_instance_registry._instances.values()) + + def _manual_two_member_context( worker: Worker, provider: ce.EndpointRecord, @@ -414,139 +424,170 @@ def test_shape_validation_rejects_duplicate_attachment_members(): ) -class _FakeCounter: - def __init__(self, calls: list[tuple]) -> None: - self._calls = calls - - def notify(self, value: int, op=NotifyOp.Set) -> None: - self._calls.append(("notify", value, op)) - - def test(self, cmp_value: int, cmp: WaitCmp): - self._calls.append(("test", cmp_value, cmp)) - return True - - def wait(self, cmp_value: int, cmp: WaitCmp, timeout: int): - self._calls.append(("wait", cmp_value, cmp, timeout)) - - -class _FakeRegion: - def __init__(self, calls: list[tuple], *, fail_mapping_close: bool = False) -> None: +class _FakeLease: + def __init__(self, calls: list[tuple], name: str, handle: int, *, fail_close: bool = False) -> None: self._calls = calls - self._fail_mapping_close = fail_mapping_close - self._worker_id = 1 - self.region_id = 42 - self._expired = False - self._released = False - self._chip_release_committed = False - - @property - def expired(self) -> bool: - return self._expired - - def payload_write(self, offset: int, host_buffer, nbytes=None) -> None: - self._calls.append(("payload_write", offset, host_buffer, nbytes)) - - def payload_read(self, offset: int, host_buffer, nbytes=None) -> None: - self._calls.append(("payload_read", offset, host_buffer, nbytes)) - - def counter(self, offset: int) -> _FakeCounter: - self._calls.append(("counter", offset)) - return _FakeCounter(self._calls) - - def _close_worker_host_mapping(self) -> None: - self._calls.append(("mapping_close", self._expired)) - if self._fail_mapping_close: + self._name = name + self.handle = handle + self.closed = False + self._fail_close = fail_close + + def close(self) -> None: + if self.closed: + return + self.closed = True + self._calls.append(("mapping_close", self._name)) + if self._fail_close: raise RuntimeError("mapping close failed") - def free(self) -> None: - self._calls.append(("free",)) - self._released = True - - def _expire(self) -> None: - self._calls.append(("expire",)) - self._expired = True - - -class _FakeMapping: - def __init__(self) -> None: - self.handle = 33 - self.payload_offset = 0 - self.payload_bytes = 64 - self.counter_offset = 64 - self.counter_bytes = 128 - class _FakeNativeWorker: def __init__(self, calls: list[tuple], *, fail_release: bool = False) -> None: self._calls = calls self._fail_release = fail_release - def control_worker_chip_region_release(self, worker_id: int, region_id: int) -> None: - self._calls.append(("release", worker_id, region_id)) - if self._fail_release: - raise RuntimeError("release failed") + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + from simpler.comm_provider import ProviderReleaseResult, ProviderReleaseStatus + from simpler.comm_provider_control import decode_release_request, encode_release_result_reply + + req_shm = SharedMemory(name=request_shm_name) + reply_shm = SharedMemory(name=reply_shm_name) + assert req_shm.buf is not None + assert reply_shm.buf is not None + req_buf = req_shm.buf + reply_buf = reply_shm.buf + try: + resource_id = decode_release_request(req_buf) + self._calls.append(("release", worker_id, resource_id)) + encode_release_result_reply( + reply_buf, + ProviderReleaseResult(provider_resource_id=int(resource_id), status=ProviderReleaseStatus.RELEASED), + ) + if self._fail_release: + raise RuntimeError("release failed") + finally: + del req_buf + del reply_buf + req_shm.close() + reply_shm.close() + + +def _committed_success(spec, *, resource_id: int = 42, shm_names=None, payload_base: int = 0, counter_base: int = 64): + from simpler.comm_provider import ( + PosixShmImport, + RegionAllocationResult, + RegionExportDescriptor, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, + ) + + names = shm_names or (f"/pto_payload_{resource_id}", f"/pto_counter_{resource_id}") + result = RegionAllocationResult( + provider_resource_id=int(resource_id), + export_descriptor=RegionExportDescriptor( + payload=RegionPartExportDescriptor( + spec.payload.planned_backing_kind, + int(spec.payload.logical_bytes), + int(spec.payload.logical_bytes), + PosixShmImport(names[0]), + ), + counter=RegionPartExportDescriptor( + spec.counter.planned_backing_kind, + int(spec.counter.logical_bytes), + int(spec.counter.logical_bytes), + PosixShmImport(names[1]), + ), + ), + ) + payload_view = RegionPartLocalView(RegionPartKind.PAYLOAD, int(payload_base), int(spec.payload.logical_bytes)) + counter_view = RegionPartLocalView(RegionPartKind.COUNTER, int(counter_base), int(spec.counter.logical_bytes)) + return result, payload_view, counter_view @pytest.fixture def region_worker(monkeypatch): - def build(*, fail_mapping_close: bool = False, fail_release: bool = False, device_ids=(8, 9)): + def build( + *, + fail_mapping_close: bool = False, + fail_release: bool = False, + fail_first_import: bool = False, + fail_second_import: bool = False, + allocate_error: Optional[BaseException] = None, + mutate_success=None, + device_ids=(8, 9), + ): worker = _l3(device_ids=device_ids) + worker._config = {**worker._config, "platform": "a2a3sim", "device_ids": list(device_ids)} calls: list[tuple] = [] - fake_region = _FakeRegion(calls, fail_mapping_close=fail_mapping_close) + leases: list[_FakeLease] = [] worker._worker = _FakeNativeWorker(calls, fail_release=fail_release) - - def create_region(worker_id: int, payload_bytes: int, counter_bytes: int): - calls.append(("create", worker_id, payload_bytes, counter_bytes)) - worker._live_worker_chip_regions.append(fake_region) - resources = worker._building_run_resources - if resources is not None: - resources.worker_chip_regions.append(fake_region) - resources.requires_ordered_cleanup = True - return fake_region - - monkeypatch.setattr(worker, "_create_worker_chip_region", create_region) - return worker, calls, fake_region + monkeypatch.setattr(worker, "_consume_worker_host_mapped_cleanup_error", lambda _api: None) + + def fake_allocate(self, spec): + calls.append(("allocate", int(spec.payload.logical_bytes), int(spec.counter.logical_bytes))) + if allocate_error is not None: + raise allocate_error + self.dispatch_started = True + self.committed_resource_id = 42 + result, payload_view, counter_view = _committed_success(spec) + if mutate_success is not None: + result, payload_view, counter_view = mutate_success(result, payload_view, counter_view) + return result, payload_view, counter_view + + def fake_import(_worker_id, _resource_id, export): + name = "payload" if not leases else "counter" + lease = _FakeLease(calls, name, handle=100 + len(leases), fail_close=fail_mapping_close) + calls.append(("import", name, int(export.logical_bytes))) + if fail_first_import and name == "payload": + raise RuntimeError("first import failed") + if fail_second_import and name == "counter": + raise RuntimeError("second import failed") + leases.append(lease) + return lease + + monkeypatch.setattr("simpler.comm_provider_control.ProviderAllocateClient.allocate", fake_allocate) + monkeypatch.setattr(worker, "_import_region_part_lease", fake_import) + return worker, calls, leases return build def test_worker_materializes_region_instance_and_closes_single_region(region_worker): - worker, calls, _fake_region = region_worker() + worker, calls, leases = region_worker() instance = _materialize_default_region(worker) assert instance.state is RegionInstanceState.LIVE assert instance.worker_id == 1 - with worker._control_reservation("test_region_instance"): - instance.payload_write(4, "src", nbytes=8) - instance.payload_read(12, "dst") - instance.counter(0).notify(3) - instance.counter(8).wait(3, WaitCmp.GE, timeout=10) - assert instance.counter(16).test(7, WaitCmp.EQ) is True + assert instance._payload_mapping is leases[0] + assert instance._counter_mapping is leases[1] + assert instance._payload_mapping is not None + assert instance._counter_mapping is not None + assert instance._payload_part is not None + assert instance._counter_part is not None + assert instance._payload_mapping is not instance._counter_mapping + assert instance._payload_mapping.handle != instance._counter_mapping.handle + assert instance._payload_part.span == RegionPartSpan(offset=0, nbytes=64) + assert instance._counter_part.span == RegionPartSpan(offset=0, nbytes=128) + with worker._control_reservation("test_region_instance"): instance.close() instance.close() assert instance.state is RegionInstanceState.CLOSED - assert worker._live_worker_chip_regions == [] + assert _tracked(worker) == () assert calls == [ - ("create", 1, 64, 128), - ("payload_write", 4, "src", 8), - ("payload_read", 12, "dst", None), - ("counter", 0), - ("notify", 3, NotifyOp.Set), - ("counter", 8), - ("wait", 3, WaitCmp.GE, 10), - ("counter", 16), - ("test", 7, WaitCmp.EQ), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] + assert all(lease.closed for lease in leases) def test_region_instance_constructs_separate_w4_part_slots_and_routes_access(region_worker, monkeypatch): - worker, _calls, fake_region = region_worker() - fake_region._worker_host_mapping = _FakeMapping() + worker, _calls, _leases = region_worker() instance = _materialize_default_region(worker) routed: list[tuple[Any, ...]] = [] @@ -573,14 +614,14 @@ def counter(self, offset): assert routed == [ ("payload_write", RegionPartSpan(offset=0, nbytes=64), 4, "src", 8), ("payload_read", RegionPartSpan(offset=0, nbytes=64), 12, "dst", None), - ("counter", RegionPartSpan(offset=64, nbytes=128), 16), + ("counter", RegionPartSpan(offset=0, nbytes=128), 16), ] -def test_region_instance_rejects_access_construction_from_non_host_vmm_copy_attachment(): - ctx = _accepted_context() - shape = validate_single_owner_region_shape(ctx) - consumer = shape.consumer.identity +def test_region_instance_rejects_access_construction_from_non_host_vmm_copy_attachment(region_worker): + worker, calls, _leases = region_worker() + ctx = _accepted_context(worker) + consumer = validate_single_owner_region_shape(ctx).consumer.identity payload_by_member = _attachments(ctx.plan.payload) payload_by_member[consumer] = dataclasses.replace( payload_by_member[consumer], @@ -589,42 +630,45 @@ def test_region_instance_rejects_access_construction_from_non_host_vmm_copy_atta ) bad_payload = dataclasses.replace(ctx.plan.payload, attachments=tuple(payload_by_member.values())) bad_ctx = dataclasses.replace(ctx, plan=dataclasses.replace(ctx.plan, payload=bad_payload)) - fake_region = _FakeRegion([]) - fake_region._worker_host_mapping = _FakeMapping() - instance = RegionInstance.planned(bad_ctx, shape) with pytest.raises(MaterializationRefusal) as excinfo: - instance._adopt_worker_chip_region(fake_region) + from simpler.comm_region import materialize_region_instance + + materialize_region_instance(bad_ctx) assert excinfo.value.reason is RefusalReason.UNSUPPORTED_ATTACHMENT + assert calls == [] + assert worker._region_instance_registry._instances == {} -def test_live_region_instance_rollback_reuses_single_region_cleanup(region_worker): - worker, calls, _fake_region = region_worker() +def test_live_region_instance_close_reuses_one_shot_cleanup(region_worker): + worker, calls, _leases = region_worker() instance = _materialize_default_region(worker) with worker._control_reservation("test_region_instance"): - instance.rollback() - instance.rollback() + instance.close() + instance.close() - assert instance.state is RegionInstanceState.ROLLED_BACK - assert worker._live_worker_chip_regions == [] + assert instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] with worker._control_reservation("test_region_instance"): instance.close() - assert instance.state is RegionInstanceState.ROLLED_BACK + assert instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () with pytest.raises(MaterializationError, match="not live"): instance.payload_write(0, "src") def test_region_instance_close_failure_marks_failed_and_poisons_worker(region_worker): - worker, calls, _fake_region = region_worker(fail_mapping_close=True) + worker, calls, _leases = region_worker(fail_mapping_close=True) instance = _materialize_default_region(worker) with worker._control_reservation("test_region_instance"): @@ -632,7 +676,7 @@ def test_region_instance_close_failure_marks_failed_and_poisons_worker(region_wo instance.close() assert instance.state is RegionInstanceState.CLOSE_FAILED - assert worker._live_worker_chip_regions == [] + assert _tracked(worker) == () with pytest.raises(RuntimeError, match="no further work is admitted"): worker._require_no_ordered_cleanup_failure("test") with pytest.raises(RuntimeError, match="not live"): @@ -640,50 +684,40 @@ def test_region_instance_close_failure_marks_failed_and_poisons_worker(region_wo with pytest.raises(RuntimeError, match="mapping close failed"): instance.close() assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] -def test_region_instance_rollback_failure_replays_cached_error_for_close(region_worker): - worker, calls, _fake_region = region_worker(fail_mapping_close=True) +def test_region_instance_close_failure_replays_cached_error(region_worker): + worker, calls, _leases = region_worker(fail_mapping_close=True) instance = _materialize_default_region(worker) with worker._control_reservation("test_region_instance"): with pytest.raises(RuntimeError, match="mapping close failed") as first_excinfo: - instance.rollback() + instance.close() with pytest.raises(RuntimeError, match="mapping close failed") as second_excinfo: - instance.rollback() - with pytest.raises(RuntimeError, match="mapping close failed") as close_excinfo: instance.close() - assert instance.state is RegionInstanceState.ROLLBACK_FAILED + assert instance.state is RegionInstanceState.CLOSE_FAILED + assert _tracked(worker) == () assert second_excinfo.value is first_excinfo.value - assert close_excinfo.value is first_excinfo.value assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] -def test_close_worker_chip_region_preserves_release_error_cause(region_worker): - worker, _calls, fake_region = region_worker(fail_mapping_close=True, fail_release=True) - - with pytest.raises(RuntimeError, match="mapping close failed") as excinfo: - worker._close_worker_chip_region(fake_region, poison_on_error=True) - - assert isinstance(excinfo.value.__cause__, RuntimeError) - assert str(excinfo.value.__cause__) == "release failed" - - def test_callback_region_close_before_submit_retired_from_run_cleanup(region_worker): - worker, calls, _fake_region = region_worker() + worker, calls, _leases = region_worker() resources = _RunResources() worker._building_run_resources = resources try: @@ -694,21 +728,21 @@ def test_callback_region_close_before_submit_retired_from_run_cleanup(region_wor worker._building_run_resources = None assert instance.state is RegionInstanceState.CLOSED - assert worker._live_worker_chip_regions == [] - assert resources.worker_chip_regions == [] - - worker._cleanup_worker_chip_regions(resources) + assert _tracked(worker) == () + worker._region_instance_registry.cleanup_run(resources) + assert _tracked(worker) == () assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] def test_callback_region_run_cleanup_then_later_close_is_idempotent(region_worker): - worker, calls, _fake_region = region_worker() + worker, calls, _leases = region_worker() resources = _RunResources() worker._building_run_resources = resources try: @@ -722,58 +756,62 @@ def test_callback_region_run_cleanup_then_later_close_is_idempotent(region_worke finally: worker._building_run_resources = None - worker._cleanup_worker_chip_regions(resources) - - assert instance.state is RegionInstanceState.LIVE - assert worker._live_worker_chip_regions == [] - assert resources.worker_chip_regions == [] - - with worker._control_reservation("test_region_instance"): - instance.close() + worker._region_instance_registry.cleanup_run(resources) assert instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () + instance.close() + assert instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] -def test_retire_worker_chip_region_tracking_does_not_retry_base_exception(): - class FailingTrackingList(list): - def __init__(self, values): - super().__init__(values) - self.assignments = 0 +def test_callback_region_data_plane_after_submit_uses_mapped_view(region_worker, monkeypatch): + worker, _calls, _leases = region_worker() + routed: list[str] = [] - def __setitem__(self, key, value): - self.assignments += 1 - raise KeyboardInterrupt("tracking interrupted") + def payload_write(self, offset, host_buffer, nbytes=None): + routed.append("payload_write") - worker = _l3(device_ids=[8, 9]) - region = object() - failing_tracking = FailingTrackingList([region]) - succeeding_tracking = [region] - resources = _RunResources(worker_chip_regions=failing_tracking) - worker._live_worker_chip_regions = succeeding_tracking + def payload_read(self, offset, host_buffer, nbytes=None): + routed.append("payload_read") - with pytest.raises(KeyboardInterrupt, match="tracking interrupted"): - worker._retire_worker_chip_region_tracking(region, resources) + def counter(self, offset): + routed.append("counter") + return "counter" - assert failing_tracking.assignments == 1 - assert succeeding_tracking == [] + monkeypatch.setattr(PayloadPart, "write", payload_write) + monkeypatch.setattr(PayloadPart, "read", payload_read) + monkeypatch.setattr(CounterPart, "counter", counter) + resources = _RunResources() + worker._building_run_resources = resources + try: + with _callback_run(19, worker): + instance = _materialize_default_region(worker) + frame = _callback_frame_for(worker) + assert frame is not None + frame.has_submitted_task = True + instance.payload_write(0, "src") + instance.payload_read(0, "dst") + assert instance.counter(0) == "counter" + with pytest.raises(RuntimeError, match="cannot follow a task submission"): + instance.close() + finally: + worker._building_run_resources = None -def test_materialize_create_failure_propagates_without_live_tracking(monkeypatch): - worker = _l3(device_ids=[8, 9]) - calls: list[tuple] = [] + assert routed == ["payload_write", "payload_read", "counter"] - def create_region(worker_id: int, payload_bytes: int, counter_bytes: int): - calls.append(("create", worker_id, payload_bytes, counter_bytes)) - raise RuntimeError("create failed") - monkeypatch.setattr(worker, "_create_worker_chip_region", create_region) +def test_materialize_tracks_before_allocate_and_create_failure_is_not_live(region_worker): + worker, calls, _leases = region_worker(allocate_error=RuntimeError("create failed")) with pytest.raises(RuntimeError, match="create failed"): worker._materialize_region_instance( @@ -782,12 +820,359 @@ def create_region(worker_id: int, payload_bytes: int, counter_bytes: int): ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=128), ) - assert worker._live_worker_chip_regions == [] - assert calls == [("create", 1, 64, 128)] + assert calls == [("allocate", 64, 128)] + assert _tracked(worker) == () + worker._require_no_ordered_cleanup_failure("test") + + +def test_materialize_tracks_before_allocate(region_worker, monkeypatch): + worker, calls, _leases = region_worker() + seen: list[int] = [] + from simpler.comm_provider_control import ProviderAllocateClient + + installed = ProviderAllocateClient.allocate + + def observe(self, spec): + seen.append(len(worker._region_instance_registry._instances)) + return installed(self, spec) + + monkeypatch.setattr(ProviderAllocateClient, "allocate", observe) + instance = _materialize_default_region(worker) + assert seen == [1] + assert instance.state is RegionInstanceState.LIVE + assert calls[:1] == [("allocate", 64, 128)] + + +def test_first_import_failure_releases_once_and_poisons(region_worker): + worker, calls, leases = region_worker(fail_first_import=True) + with pytest.raises(RuntimeError, match="first import failed"): + _materialize_default_region(worker) + + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("test") + assert calls == [ + ("allocate", 64, 128), + ("import", "payload", 64), + ("release", 1, 42), + ] + assert leases == [] + + +def test_second_import_failure_closes_first_lease_and_poisons(region_worker): + worker, calls, leases = region_worker(fail_second_import=True) + with pytest.raises(RuntimeError, match="second import failed"): + _materialize_default_region(worker) + + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("test") + assert calls == [ + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("release", 1, 42), + ] + assert leases[0].closed is True + assert len(leases) == 1 + + +def test_invalid_success_reply_releases_once_and_never_imports(region_worker): + from simpler.comm_provider import RegionPartKind, RegionPartLocalView + + def overlap(result, payload_view, _counter_view): + return result, payload_view, RegionPartLocalView(RegionPartKind.COUNTER, 0, 128) + + worker, calls, leases = region_worker(mutate_success=overlap) + with pytest.raises(RuntimeError, match="must not overlap"): + _materialize_default_region(worker) + + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("test") + assert calls == [ + ("allocate", 64, 128), + ("release", 1, 42), + ] + assert leases == [] + + +def test_invalid_success_compensation_debt_releases_once_and_survives_until_sweep(region_worker): + from simpler.comm_provider import RegionPartKind, RegionPartLocalView + + def overlap(result, payload_view, _counter_view): + return result, payload_view, RegionPartLocalView(RegionPartKind.COUNTER, 0, 128) + + worker, calls, leases = region_worker(fail_release=True, mutate_success=overlap) + with pytest.raises(RuntimeError, match="must not overlap"): + _materialize_default_region(worker) + + tracked = _tracked(worker) + assert len(tracked) == 1 + instance = tracked[0] + assert instance._state is RegionInstanceState.CLOSE_FAILED + assert instance._release_client is not None + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("test") + assert calls == [ + ("allocate", 64, 128), + ("release", 1, 42), + ] + assert leases == [] + worker._region_instance_registry.cleanup_run(None) + assert _tracked(worker) == (instance,) + assert calls.count(("release", 1, 42)) == 1 + worker._region_instance_registry.sweep() + assert _tracked(worker) == () + assert calls.count(("release", 1, 42)) == 1 + + +def test_successful_close_then_materialize_and_close_again(region_worker): + worker, calls, _leases = region_worker() + first = _materialize_default_region(worker) + with worker._control_reservation("test_region_instance"): + first.close() + assert first.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () + + second = _materialize_default_region(worker) + assert second.state is RegionInstanceState.LIVE + assert second is not first + with worker._control_reservation("test_region_instance"): + second.close() + assert second.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () + assert calls.count(("allocate", 64, 128)) == 2 + assert sum(1 for item in calls if item[0] == "import") == 4 + assert sum(1 for item in calls if item[0] == "mapping_close") == 4 + assert [item for item in calls if item[0] == "release"] == [ + ("release", 1, 42), + ("release", 1, 42), + ] + + +def test_committed_local_base_zero_is_not_treated_as_absent(region_worker): + worker, _calls, _leases = region_worker() + instance = _materialize_default_region(worker) + assert instance._payload_part is not None + assert instance._counter_part is not None + assert instance._payload_part.span.offset == 0 + assert instance._counter_part.span.offset == 0 + assert instance.state is RegionInstanceState.LIVE + + +def test_allocation_error_without_debt_retires_closed_without_release(region_worker): + from simpler.comm_provider import RegionAllocationError, RegionControlErrorKind, RegionOperationKind, RegionPartKind + + error = RegionAllocationError( + provisional_resource_id=9, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=False, + message="create failed cleanly", + ) + worker, calls, _leases = region_worker(allocate_error=error) + with pytest.raises(RegionAllocationError): + _materialize_default_region(worker) + + assert _tracked(worker) == () + worker._require_no_ordered_cleanup_failure("test") + assert calls == [("allocate", 64, 128)] + + +def test_allocation_error_with_debt_is_close_failed_without_release_client(region_worker): + from simpler.comm_provider import RegionAllocationError, RegionControlErrorKind, RegionOperationKind, RegionPartKind + + error = RegionAllocationError( + provisional_resource_id=9, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=True, + message="create left debt", + ) + worker, calls, _leases = region_worker(allocate_error=error) + with pytest.raises(RuntimeError, match="cleanup debt"): + _materialize_default_region(worker) + + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + assert tracked[0]._release_client is None + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("test") + assert calls == [("allocate", 64, 128)] + + +def test_registry_run_membership_cleanup_closes_only_that_run(region_worker): + worker, calls, _leases = region_worker() + first = _RunResources() + second = _RunResources() + worker._building_run_resources = first + try: + first_instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + worker._building_run_resources = second + try: + second_instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + + assert {first_instance, second_instance} == set(_tracked(worker)) + worker._region_instance_registry.cleanup_run(first) + assert first_instance.state is RegionInstanceState.CLOSED + assert second_instance.state is RegionInstanceState.LIVE + assert _tracked(worker) == (second_instance,) + assert [item for item in calls if item[0] == "release"] == [("release", 1, 42)] + + worker._region_instance_registry.cleanup_run(second) + assert second_instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () + assert [item for item in calls if item[0] == "release"] == [("release", 1, 42), ("release", 1, 42)] + + +def test_registry_run_cleanup_closes_live_instances(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + + assert instance.state is RegionInstanceState.LIVE + assert _tracked(worker) == (instance,) + worker._region_instance_registry.cleanup_run(resources) + assert instance.state is RegionInstanceState.CLOSED + assert _tracked(worker) == () + assert calls == [ + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), + ("release", 1, 42), + ] + + +def test_allocation_error_with_debt_survives_run_cleanup_until_sweep_without_release(region_worker): + from simpler.comm_provider import RegionAllocationError, RegionControlErrorKind, RegionOperationKind, RegionPartKind + + error = RegionAllocationError( + provisional_resource_id=9, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=True, + message="create left debt", + ) + worker, calls, _leases = region_worker(allocate_error=error) + resources = _RunResources() + worker._building_run_resources = resources + try: + with pytest.raises(RuntimeError, match="cleanup debt"): + _materialize_default_region(worker) + finally: + worker._building_run_resources = None + + tracked = _tracked(worker) + assert len(tracked) == 1 + instance = tracked[0] + assert instance._state is RegionInstanceState.CLOSE_FAILED + assert instance._release_client is None + worker._region_instance_registry.cleanup_run(resources) + assert _tracked(worker) == (instance,) + assert calls == [("allocate", 64, 128)] + worker._region_instance_registry.sweep() + assert _tracked(worker) == () + assert calls == [("allocate", 64, 128)] + + +def test_provider_release_failure_keeps_diagnostic_until_sweep(region_worker): + from simpler.comm_provider_control import RegionControlProtocolError + + worker, calls, _leases = region_worker(fail_release=True) + instance = _materialize_default_region(worker) + with worker._control_reservation("test_region_instance"): + with pytest.raises(RegionControlProtocolError, match="release"): + instance.close() + + assert instance.state is RegionInstanceState.CLOSE_FAILED + assert _tracked(worker) == (instance,) + worker._region_instance_registry.cleanup_run(None) + assert _tracked(worker) == (instance,) + release_calls = [item for item in calls if item[0] == "release"] + assert release_calls == [("release", 1, 42)] + worker._region_instance_registry.sweep() + assert _tracked(worker) == () + assert [item for item in calls if item[0] == "release"] == release_calls + + +def test_unpublished_close_failed_survives_cleanup_run_until_sweep(region_worker): + worker, calls, _leases = region_worker(fail_second_import=True) + resources = _RunResources() + worker._building_run_resources = resources + try: + with pytest.raises(RuntimeError, match="second import failed"): + _materialize_default_region(worker) + finally: + worker._building_run_resources = None + + tracked = _tracked(worker) + assert len(tracked) == 1 + instance = tracked[0] + worker._region_instance_registry.cleanup_run(resources) + assert _tracked(worker) == (instance,) + assert calls.count(("release", 1, 42)) == 1 + worker._region_instance_registry.sweep() + assert _tracked(worker) == () + assert calls.count(("release", 1, 42)) == 1 + + +def test_data_plane_does_not_consult_registry(region_worker, monkeypatch): + worker, _calls, _leases = region_worker() + instance = _materialize_default_region(worker) + registry = worker._region_instance_registry + + def boom(*_args, **_kwargs): + raise AssertionError("data-plane must not consult the region instance registry") + + for name in ( + "track", + "close", + "cleanup_run", + "sweep", + "record_data_plane_failure", + "_iter_run", + "_settle", + "_retire", + ): + monkeypatch.setattr(registry, name, boom) + monkeypatch.setattr(PayloadPart, "write", lambda *_args, **_kwargs: None) + monkeypatch.setattr(PayloadPart, "read", lambda *_args, **_kwargs: None) + monkeypatch.setattr(CounterPart, "counter", lambda *_args, **_kwargs: "counter") + + with worker._control_reservation("test_region_instance"): + instance.payload_write(0, "src") + instance.payload_read(0, "dst") + assert instance.counter(0) == "counter" + + assert not hasattr(RegionInstanceRegistry, "get") + assert not hasattr(RegionInstanceRegistry, "find") + assert not hasattr(RegionInstanceRegistry, "__getitem__") def test_live_region_instance_access_requires_control_context(region_worker): - worker, calls, _fake_region = region_worker() + worker, calls, _leases = region_worker() instance = _materialize_default_region(worker) with pytest.raises(RuntimeError, match="active orchestration/control context"): @@ -801,14 +1186,164 @@ def test_live_region_instance_access_requires_control_context(region_worker): instance.close() assert calls == [ - ("create", 1, 64, 128), - ("mapping_close", False), + ("allocate", 64, 128), + ("import", "payload", 64), + ("import", "counter", 128), + ("mapping_close", "payload"), + ("mapping_close", "counter"), ("release", 1, 42), - ("free",), - ("expire",), ] +def test_local_view_returns_completed_materialization_only(region_worker): + worker, _calls, _leases = region_worker() + instance = _materialize_default_region(worker) + assert instance.provider_resource_id == 42 + payload = instance.local_view(RegionPartKind.PAYLOAD) + counter = instance.local_view(RegionPartKind.COUNTER) + assert payload is instance._payload_local_view + assert counter is instance._counter_local_view + assert payload is not None and counter is not None + instance._state = RegionInstanceState.CLOSED + assert instance.local_view(RegionPartKind.PAYLOAD) is None + assert instance.local_view(RegionPartKind.COUNTER) is None + + +def test_live_resource_inventory_reports_registry_instance_count(region_worker): + worker, _calls, _leases = region_worker() + assert "region instance" not in worker._describe_live_resources() + _materialize_default_region(worker) + assert "1 region instance(s)" in worker._describe_live_resources() + + +def test_record_data_plane_failure_matches_single_region(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + error = RuntimeError("issued local operation failed") + worker._region_instance_registry.record_data_plane_failure(resources, instance.provider_resource_id, error) + assert instance.data_plane_error is error + assert instance._close_attempted is False + assert instance.state is RegionInstanceState.LIVE + assert [item for item in calls if item[0] == "release"] == [] + + +def test_record_data_plane_failure_matches_only_target_among_multiple(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + first = _materialize_default_region(worker) + second = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + second._provider_resource_id = 43 + error = RuntimeError("issued local operation failed") + worker._region_instance_registry.record_data_plane_failure(resources, first.provider_resource_id, error) + assert first.data_plane_error is error + assert second.data_plane_error is None + assert first._close_attempted is False + assert second._close_attempted is False + assert [item for item in calls if item[0] == "release"] == [] + + +def test_record_data_plane_failure_unknown_resource_is_routing_error(region_worker): + worker, _calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + with pytest.raises(MaterializationError, match="no region instance for resource 99"): + worker._region_instance_registry.record_data_plane_failure(resources, 99, RuntimeError("unused")) + assert instance.data_plane_error is None + assert instance._close_attempted is False + + +def test_record_data_plane_failure_duplicate_match_is_invariant_error(region_worker): + worker, _calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + first = _materialize_default_region(worker) + second = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + second._provider_resource_id = first.provider_resource_id + with pytest.raises(MaterializationError, match="duplicate region instances for resource 42"): + worker._region_instance_registry.record_data_plane_failure( + resources, first.provider_resource_id, RuntimeError("unused") + ) + assert first.data_plane_error is None + assert second.data_plane_error is None + + +def _endpoint_error(resource_id: int) -> RuntimeError: + return RuntimeError( + f"L3-L2 endpoint error op=payload_write kind=5 region={int(resource_id)} msg=issued local operation failed" + ) + + +def test_endpoint_unknown_resource_poisons_worker(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + assert worker._poison_worker_chip_region_from_endpoint_error(_endpoint_error(99), resources) is True + assert instance.data_plane_error is None + assert instance._close_attempted is False + assert [item for item in calls if item[0] == "release"] == [] + _assert_poisoned(worker) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) + + +def test_endpoint_duplicate_match_poisons_worker(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + first = _materialize_default_region(worker) + second = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + second._provider_resource_id = first.provider_resource_id + assert worker._poison_worker_chip_region_from_endpoint_error(_endpoint_error(first.provider_resource_id), resources) + assert first.data_plane_error is None + assert second.data_plane_error is None + assert [item for item in calls if item[0] == "release"] == [] + _assert_poisoned(worker) + + +def test_data_plane_poison_still_completes_registry_cleanup_and_release(region_worker): + worker, calls, _leases = region_worker() + resources = _RunResources() + worker._building_run_resources = resources + try: + instance = _materialize_default_region(worker) + finally: + worker._building_run_resources = None + error = _endpoint_error(instance.provider_resource_id) + assert worker._poison_worker_chip_region_from_endpoint_error(error, resources) is True + assert instance.data_plane_error is error + assert instance._close_attempted is False + assert instance.state is RegionInstanceState.LIVE + worker._region_instance_registry.cleanup_run(resources) + assert instance.state is RegionInstanceState.CLOSED + assert calls.count(("release", 1, 42)) == 1 + _assert_poisoned(worker, cause=error) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) + + def test_shape_validation_rejects_foreign_registry_context(): owner = _accepted_context(_l3(device_ids=[0, 1])) foreign_worker = _l3(device_ids=[0, 1]) @@ -835,3 +1370,363 @@ def test_shape_validation_rejects_stale_registry_epoch(): validate_single_owner_region_shape(ctx) assert excinfo.value.reason is RefusalReason.REGISTRY_MISMATCH + + +def test_projection_copies_admitted_part_backing_and_layout_bytes(): + ctx = _accepted_context() + shape = validate_single_owner_region_shape(ctx) + spec = project_region_allocation_spec(ctx.plan, ctx.layout) + + assert shape.worker_id == 1 + assert spec.payload.planned_backing_kind is ctx.plan.payload.backend_kind + assert spec.counter.planned_backing_kind is ctx.plan.counter.backend_kind + assert spec.payload.logical_bytes == ctx.layout.payload_bytes + assert spec.counter.logical_bytes == ctx.layout.counter_bytes + assert spec.payload.planned_backing_kind is ce.BackendKind.VMM_WINDOW + assert spec.counter.planned_backing_kind is ce.BackendKind.VMM_WINDOW + assert not hasattr(spec, "provider") + assert not hasattr(spec.payload, "device_id") + assert not hasattr(spec.payload, "worker_id") + + +def test_projection_copies_plan_backing_without_selecting_capability(): + ctx = _accepted_context() + posix_plan = dataclasses.replace( + ctx.plan, + payload=dataclasses.replace(ctx.plan.payload, backend_kind=ce.BackendKind.POSIX_SHM), + counter=dataclasses.replace(ctx.plan.counter, backend_kind=ce.BackendKind.POSIX_SHM), + ) + + spec = project_region_allocation_spec(posix_plan, ctx.layout) + + assert spec.payload.planned_backing_kind is ce.BackendKind.POSIX_SHM + assert spec.counter.planned_backing_kind is ce.BackendKind.POSIX_SHM + + +def test_projection_refuses_zero_layout_before_store_invocation(): + ctx = _accepted_context() + zero_layout = ce.RegionLayoutSpec(payload_bytes=0, counter_bytes=128) + + with pytest.raises(ValueError, match="logical_bytes must be positive"): + project_region_allocation_spec(ctx.plan, zero_layout) + + +def test_projection_requires_a_real_backend_plan(): + ctx = _accepted_context() + unsupported = ce.UnsupportedRegionPlan( + reason=ce.BackendUnsupportedReason.NO_DEFAULT_PROVIDER, + message="no default SingleOwner provider is available", + ) + + with pytest.raises(TypeError, match="requires a BackendPlan"): + project_region_allocation_spec(unsupported, ctx.layout) + with pytest.raises(TypeError, match="requires a RegionLayoutSpec"): + project_region_allocation_spec(ctx.plan, object()) + + +def test_compatibility_create_projects_admitted_w2_plan_not_byte_counts(monkeypatch): + worker = _l3(device_ids=[8, 9]) + captured: dict[str, Any] = {} + + def fake_allocate(_self, spec): + captured["spec"] = spec + raise RuntimeError("allocate must not run in this projection probe") + + monkeypatch.setattr("simpler.comm_provider_control.ProviderAllocateClient.allocate", fake_allocate) + + spec = worker._project_admitted_worker_chip_region_spec(1, 64, 128) + + assert spec.payload.planned_backing_kind is ce.BackendKind.VMM_WINDOW + assert spec.counter.planned_backing_kind is ce.BackendKind.VMM_WINDOW + assert spec.payload.logical_bytes == 64 + assert spec.counter.logical_bytes == 128 + assert "spec" not in captured + + +def test_compatibility_create_does_not_allocate_when_admission_refuses(monkeypatch): + worker = _l3(device_ids=[8, 9]) + calls: list[str] = [] + monkeypatch.setattr(worker, "_consume_worker_host_mapped_cleanup_error", lambda _api: None) + + def refuse(_ctx): + calls.append("admit") + raise MaterializationRefusal(RefusalReason.UNSUPPORTED_BACKEND_KIND, "refused") + + def allocate(_self, _spec): + calls.append("allocate") + raise AssertionError("unadmitted plans must not reach the provider store") + + monkeypatch.setattr("simpler.worker.validate_single_owner_region_shape", refuse) + monkeypatch.setattr("simpler.comm_provider_control.ProviderAllocateClient.allocate", allocate) + + with pytest.raises(MaterializationRefusal) as excinfo: + worker._create_worker_chip_region(1, 64, 128) + + assert excinfo.value.reason is RefusalReason.UNSUPPORTED_BACKEND_KIND + assert calls == ["admit"] + + +def test_compatibility_create_uses_projection_result_not_a_fabricated_spec(monkeypatch): + worker = _l3(device_ids=[8, 9]) + captured: dict[str, Any] = {} + monkeypatch.setattr(worker, "_consume_worker_host_mapped_cleanup_error", lambda _api: None) + + def fake_project(plan, layout): + captured["plan"] = plan + captured["layout"] = layout + return RegionAllocationSpec( + payload=RegionPartAllocationSpec(plan.payload.backend_kind, 96), + counter=RegionPartAllocationSpec(plan.counter.backend_kind, 8), + ) + + def fake_allocate(_self, spec): + captured["spec"] = spec + raise RuntimeError("stop after capturing the projected spec") + + monkeypatch.setattr("simpler.comm_region.project_region_allocation_spec", fake_project) + monkeypatch.setattr("simpler.comm_provider_control.ProviderAllocateClient.allocate", fake_allocate) + + with pytest.raises(RuntimeError, match="stop after capturing the projected spec"): + worker._create_worker_chip_region(1, 64, 128) + + plan = captured["plan"] + spec = captured["spec"] + assert isinstance(plan, ce.BackendPlan) + assert captured["layout"] == ce.RegionLayoutSpec(payload_bytes=64, counter_bytes=128) + assert spec.payload.logical_bytes == 96 + assert spec.counter.logical_bytes == 8 + assert spec.payload.planned_backing_kind is plan.payload.backend_kind + assert spec.counter.planned_backing_kind is plan.counter.backend_kind + registry = worker._get_endpoint_registry() + provider = registry.record_for(plan.topology_plan.provider_endpoint) + consumer = next(registry.record_for(member) for member in plan.ordered_members if member != provider.identity) + assert provider.path == "L3/L2[1]" + assert provider.deployment is ce.DEVICE_AICPU + assert consumer.path == "L3" + assert consumer.deployment is ce.HOST_CPU + assert set(plan.ordered_members) == {provider.identity, consumer.identity} + + +class _StoreControlMailbox: + def __init__( + self, + store, + *, + fail_before_allocate: Optional[BaseException] = None, + fail_after_allocate: Optional[BaseException] = None, + reply_mode: str = "handler", + ) -> None: + self.store = store + self.allocate_calls = 0 + self.release_calls = 0 + self.released_ids: list[int] = [] + self._fail_before_allocate = fail_before_allocate + self._fail_after_allocate = fail_after_allocate + self._reply_mode = reply_mode + + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + from simpler.comm_provider_control import ( + ALLOCATE_REPLY_BYTES, + ALLOCATE_REQUEST_BYTES, + COMMIT_TAG_OFFSET, + AllocateReplyTag, + handle_ctrl_region_allocate, + ) + + del worker_id + self.allocate_calls += 1 + if self._fail_before_allocate is not None: + raise self._fail_before_allocate + req = SharedMemory(name=request_shm_name) + reply = SharedMemory(name=reply_shm_name) + assert req.buf is not None + assert reply.buf is not None + req_view = memoryview(req.buf)[:ALLOCATE_REQUEST_BYTES] + reply_view = memoryview(reply.buf)[:ALLOCATE_REPLY_BYTES] + try: + if self._reply_mode == "handler": + handle_ctrl_region_allocate(req_view, reply_view, self.store) + elif self._reply_mode == "malformed": + struct.pack_into(" None: + from simpler.comm_provider_control import ( + RELEASE_REPLY_BYTES, + RELEASE_REQUEST_BYTES, + decode_release_request, + handle_ctrl_region_release, + ) + + del worker_id + self.release_calls += 1 + req = SharedMemory(name=request_shm_name) + reply = SharedMemory(name=reply_shm_name) + assert req.buf is not None + assert reply.buf is not None + req_view = memoryview(req.buf)[:RELEASE_REQUEST_BYTES] + reply_view = memoryview(reply.buf)[:RELEASE_REPLY_BYTES] + try: + self.released_ids.append(decode_release_request(req_view)) + handle_ctrl_region_release(req_view, reply_view, self.store) + finally: + req_view.release() + reply_view.release() + req.close() + reply.close() + + +def _live_control_worker(mailbox, monkeypatch, device_ids=(8, 9)): + worker = _l3(device_ids=device_ids) + worker._config = {**worker._config, "platform": "a2a3sim", "device_ids": list(device_ids)} + worker._worker = mailbox + monkeypatch.setattr(worker, "_consume_worker_host_mapped_cleanup_error", lambda _api: None) + return worker + + +def _assert_poisoned(worker, *, cause: Optional[BaseException] = None) -> RuntimeError: + with pytest.raises(RuntimeError, match="no further work is admitted") as poison_info: + worker._require_no_ordered_cleanup_failure("test") + if cause is not None: + chain: list[BaseException] = [] + current: Optional[BaseException] = poison_info.value + while current is not None and current not in chain: + chain.append(current) + current = current.__cause__ + assert any(item is cause or item == cause for item in chain) + return poison_info.value + + +def test_handler_commit_then_mailbox_error_releases_once_and_poisons(monkeypatch): + from simpler.comm_provider import ProviderRegionStore + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + mailbox = _StoreControlMailbox(store, fail_after_allocate=RuntimeError("mailbox down after commit")) + worker = _live_control_worker(mailbox, monkeypatch) + with pytest.raises(RuntimeError, match="mailbox down after commit") as exc_info: + _materialize_default_region(worker) + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._provider_resource_id == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + assert mailbox.release_calls == 1 + assert mailbox.released_ids == [1] + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + _assert_poisoned(worker, cause=exc_info.value) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) + + +def test_dispatch_empty_reply_does_not_release_and_poisons(monkeypatch): + from simpler.comm_provider import ProviderRegionStore + from simpler.comm_provider_control import RegionControlProtocolError + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + store = ProviderRegionStore(_sim_context(), _shell_factory=FakeShellFactory()) + mailbox = _StoreControlMailbox(store, reply_mode="empty") + worker = _live_control_worker(mailbox, monkeypatch) + with pytest.raises(RegionControlProtocolError): + _materialize_default_region(worker) + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._provider_resource_id == 0 + assert mailbox.release_calls == 0 + _assert_poisoned(worker) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) + + +def test_dispatch_malformed_reply_does_not_release_and_poisons(monkeypatch): + from simpler.comm_provider import ProviderRegionStore, RegionControlError + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + store = ProviderRegionStore(_sim_context(), _shell_factory=FakeShellFactory()) + mailbox = _StoreControlMailbox(store, reply_mode="malformed") + worker = _live_control_worker(mailbox, monkeypatch) + with pytest.raises(RegionControlError): + _materialize_default_region(worker) + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._provider_resource_id == 0 + assert mailbox.release_calls == 0 + _assert_poisoned(worker) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) + + +def test_request_encode_failure_does_not_dispatch_or_poison(monkeypatch): + from simpler import comm_provider_control as control + from simpler.comm_provider import ProviderRegionStore + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + store = ProviderRegionStore(_sim_context(), _shell_factory=FakeShellFactory()) + mailbox = _StoreControlMailbox(store) + worker = _live_control_worker(mailbox, monkeypatch) + + def _boom(*_args, **_kwargs): + raise RuntimeError("encode failed") + + monkeypatch.setattr(control, "encode_allocate_request", _boom) + with pytest.raises(RuntimeError, match="encode failed"): + _materialize_default_region(worker) + assert mailbox.allocate_calls == 0 + assert mailbox.release_calls == 0 + assert _tracked(worker) == () + worker._require_no_ordered_cleanup_failure("test") + + +def test_mailbox_error_before_handler_poisons_and_keeps_transport_type(monkeypatch): + from simpler.comm_provider import ProviderRegionStore + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + store = ProviderRegionStore(_sim_context(), _shell_factory=FakeShellFactory()) + mailbox = _StoreControlMailbox(store, fail_before_allocate=RuntimeError("mailbox down before handler")) + worker = _live_control_worker(mailbox, monkeypatch) + with pytest.raises(RuntimeError, match="mailbox down before handler") as exc_info: + _materialize_default_region(worker) + assert mailbox.allocate_calls == 1 + assert mailbox.release_calls == 0 + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._provider_resource_id == 0 + _assert_poisoned(worker, cause=exc_info.value) + with pytest.raises(RuntimeError, match="no further work is admitted") as admission_info: + _materialize_default_region(worker) + assert admission_info.value.__cause__ is worker._ordered_cleanup_error + + +def test_store_lifecycle_allocate_is_terminal_ambiguity(monkeypatch): + from simpler.comm_provider import ProviderRegionStore, RegionControlError, RegionControlErrorKind + + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + factory = FakeShellFactory() + store = ProviderRegionStore(_sim_context(), _shell_factory=factory) + store.sweep() + mailbox = _StoreControlMailbox(store) + worker = _live_control_worker(mailbox, monkeypatch) + with pytest.raises(RegionControlError) as exc_info: + _materialize_default_region(worker) + assert exc_info.value.kind is RegionControlErrorKind.STORE_LIFECYCLE + assert mailbox.release_calls == 0 + assert factory.world.calls == [] + tracked = _tracked(worker) + assert len(tracked) == 1 + assert tracked[0]._state is RegionInstanceState.CLOSE_FAILED + _assert_poisoned(worker, cause=exc_info.value) + with pytest.raises(RuntimeError, match="no further work is admitted"): + _materialize_default_region(worker) diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 70fd1ab8bf..c5a5889102 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -6,6 +6,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415, E402 """Unit tests for Worker (Python L3 wrapper over _Worker). Tests use SubWorker (fork/shm) as the only worker type — no NPU device required. @@ -66,6 +67,33 @@ _pack_py_callable_payload, ) + +def _native_control_region_release(recorder, worker_id, request_shm_name, reply_shm_name, *, fail=None): + from simpler.comm_provider import ProviderReleaseResult, ProviderReleaseStatus + from simpler.comm_provider_control import decode_release_request, encode_release_result_reply + + req_shm = SharedMemory(name=request_shm_name) + reply_shm = SharedMemory(name=reply_shm_name) + assert req_shm.buf is not None + assert reply_shm.buf is not None + req_buf = req_shm.buf + reply_buf = reply_shm.buf + try: + resource_id = decode_release_request(req_buf) + recorder.append((worker_id, resource_id)) + encode_release_result_reply( + reply_buf, + ProviderReleaseResult(provider_resource_id=int(resource_id), status=ProviderReleaseStatus.RELEASED), + ) + if fail is not None: + raise fail + finally: + del req_buf + del reply_buf + req_shm.close() + reply_shm.close() + + from ._harness import chip_callable, fake_chip_l3, requires_sim_binaries # --------------------------------------------------------------------------- @@ -186,6 +214,276 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No assert published_frame_counts == [2] +def _dummy_l2_domain(domain_id: int): + return worker_mod._L2GlobalDomain( + domain_id=domain_id, + generation=1, + domain_rank=0, + rank_count=1, + descriptor=cast(Any, object()), + local_window_base=0, + mapping_size=8, + requested_window_size=8, + ) + + +class _RecordingImportRegistry: + def __init__(self) -> None: + self.close_calls = 0 + self.error: Optional[BaseException] = None + + def close(self) -> None: + self.close_calls += 1 + if self.error is not None: + raise self.error + + +class _RecordingDomainImpl: + def __init__(self) -> None: + self.released: list[int] = [] + self.errors: dict[int, BaseException] = {} + + def comm_global_domain_release(self, domain_id: int) -> None: + self.released.append(int(domain_id)) + error = self.errors.get(int(domain_id)) + if error is not None: + raise error + + +class _TeardownPartShell: + def __init__(self, part, spec) -> None: + self.part = part + self.spec = spec + self.release_count = 0 + self.release_step_failures = [] + self._local_base = 0x1000 if part.name == "PAYLOAD" else 0x2000 + + def materialize(self) -> None: + return None + + def mapping_bytes(self) -> int: + return self.spec.logical_bytes + + def import_capability(self): + from simpler.comm_provider import PosixShmImport + + return PosixShmImport(shm_name=f"smp_{self.part.name.lower()}{id(self):016x}"[:32]) + + def local_base(self) -> int: + return self._local_base + + def zero_bytes(self, offset: int, nbytes: int) -> None: + del offset, nbytes + + def release_once(self): + self.release_count += 1 + if self.release_step_failures: + return self.release_step_failures[0] + return None + + +class _TeardownShellFactory: + def __init__(self) -> None: + self.payloads: list[_TeardownPartShell] = [] + self.counters: list[_TeardownPartShell] = [] + + def __call__(self, context, part, spec): + del context + shell = _TeardownPartShell(part, spec) + if part.name == "PAYLOAD": + self.payloads.append(shell) + else: + self.counters.append(shell) + return shell + + +def test_sweep_l2_global_domains_attempts_every_domain_and_lists_failures_by_id(): + impl = _RecordingDomainImpl() + impl.errors[7] = RuntimeError("domain 7 failed") + impl.errors[3] = RuntimeError("domain 3 failed") + store = worker_mod._L2GlobalDomainStore( + domains={ + 7: _dummy_l2_domain(7), + 3: _dummy_l2_domain(3), + 5: _dummy_l2_domain(5), + } + ) + with pytest.raises(RuntimeError, match="domain cleanup failed") as exc_info: + worker_mod._sweep_l2_global_domains(cast(Any, SimpleNamespace(_impl=impl)), store) + assert impl.released == [3, 5, 7] + assert store.domains == {} + message = str(exc_info.value) + assert message.index("domain 3:") < message.index("domain 7:") + assert "domain 3: domain 3 failed" in message + assert "domain 7: domain 7 failed" in message + assert "domain 5" not in message + assert exc_info.value.__cause__ is impl.errors[3] + + +def test_teardown_chip_process_resources_continues_and_aggregates_in_order(): + from simpler.buffer import BackendKind + from simpler.comm_provider import ( + DeviceAllocationTarget, + ProviderCleanupFailure, + ProviderRegionStore, + RegionAllocationContext, + RegionAllocationSpec, + RegionCleanupCause, + RegionEnvironmentKind, + RegionOperationKind, + RegionPartAllocationSpec, + RegionPartKind, + ) + + registry = _RecordingImportRegistry() + registry.error = RuntimeError("import close failed") + impl = _RecordingDomainImpl() + impl.errors[4] = RuntimeError("domain 4 failed") + domains = worker_mod._L2GlobalDomainStore(domains={4: _dummy_l2_domain(4)}) + factory = _TeardownShellFactory() + store = ProviderRegionStore( + RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=DeviceAllocationTarget(device_id=0), + ), + _shell_factory=factory, + ) + result = store.allocate_and_export( + RegionAllocationSpec( + payload=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=64), + counter=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=8), + ) + ) + factory.payloads[0].release_step_failures = [ + ProviderCleanupFailure( + part=RegionPartKind.PAYLOAD, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_ERROR, + ) + ] + + with pytest.raises(RuntimeError, match="chip process resource teardown failed") as exc_info: + worker_mod._teardown_chip_process_resources( + cast(Any, registry), + cast(Any, SimpleNamespace(_impl=impl)), + domains, + store, + ) + assert registry.close_calls == 1 + assert impl.released == [4] + assert factory.payloads[0].release_count == 1 + message = str(exc_info.value) + assert message.index("import close failed") < message.index("domain cleanup failed") + assert message.index("domain cleanup failed") < message.index(f"provider resource {result.provider_resource_id}") + assert "PAYLOAD" in message + assert "RELEASE" in message + assert "BACKEND_ERROR" in message + assert exc_info.value.__cause__ is registry.error + + +def test_teardown_chip_process_resources_lists_every_retained_resource(): + from simpler.buffer import BackendKind + from simpler.comm_provider import ( + DeviceAllocationTarget, + ProviderCleanupFailure, + ProviderRegionStore, + RegionAllocationContext, + RegionAllocationSpec, + RegionCleanupCause, + RegionEnvironmentKind, + RegionOperationKind, + RegionPartAllocationSpec, + RegionPartKind, + ) + + factory = _TeardownShellFactory() + store = ProviderRegionStore( + RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=DeviceAllocationTarget(device_id=0), + ), + _shell_factory=factory, + ) + spec = RegionAllocationSpec( + payload=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=64), + counter=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=8), + ) + first = store.allocate_and_export(spec) + second = store.allocate_and_export(spec) + factory.payloads[0].release_step_failures = [ + ProviderCleanupFailure( + part=RegionPartKind.PAYLOAD, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.BACKEND_ERROR, + ) + ] + factory.counters[1].release_step_failures = [ + ProviderCleanupFailure( + part=RegionPartKind.COUNTER, + backend_operation=RegionOperationKind.RELEASE, + typed_cause=RegionCleanupCause.INTERRUPTED, + ) + ] + with pytest.raises(RuntimeError, match="chip process resource teardown failed") as exc_info: + worker_mod._teardown_chip_process_resources( + cast(Any, _RecordingImportRegistry()), + cast(Any, SimpleNamespace(_impl=_RecordingDomainImpl())), + worker_mod._L2GlobalDomainStore(), + store, + ) + message = str(exc_info.value) + assert f"provider resource {first.provider_resource_id}" in message + assert f"provider resource {second.provider_resource_id}" in message + assert "PAYLOAD RELEASE BACKEND_ERROR" in message + assert "COUNTER RELEASE INTERRUPTED" in message + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert factory.payloads[1].release_count == 1 + assert factory.counters[1].release_count == 1 + + +def test_teardown_chip_process_resources_ignores_released_and_keeps_each_step_once(): + from simpler.buffer import BackendKind + from simpler.comm_provider import ( + DeviceAllocationTarget, + ProviderRegionStore, + ProviderRegionStoreState, + RegionAllocationContext, + RegionAllocationSpec, + RegionEnvironmentKind, + RegionPartAllocationSpec, + ) + + registry = _RecordingImportRegistry() + impl = _RecordingDomainImpl() + domains = worker_mod._L2GlobalDomainStore() + factory = _TeardownShellFactory() + store = ProviderRegionStore( + RegionAllocationContext( + environment_kind=RegionEnvironmentKind.SIM, + target=DeviceAllocationTarget(device_id=0), + ), + _shell_factory=factory, + ) + store.allocate_and_export( + RegionAllocationSpec( + payload=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=64), + counter=RegionPartAllocationSpec(planned_backing_kind=BackendKind.POSIX_SHM, logical_bytes=8), + ) + ) + worker_mod._teardown_chip_process_resources( + cast(Any, registry), + cast(Any, SimpleNamespace(_impl=impl)), + domains, + store, + ) + assert registry.close_calls == 1 + assert impl.released == [] + assert factory.payloads[0].release_count == 1 + assert factory.counters[0].release_count == 1 + assert store.state is ProviderRegionStoreState.CLOSED + + @pytest.mark.parametrize( ("platform", "runtime", "depth", "expected"), [ @@ -3987,34 +4285,30 @@ def __init__(self): def _release_slot_ref(self): self.releases += 1 - class Region: - def __init__(self, region_id): - self.region_id = region_id - self._worker_id = 0 - self.mapping_closed = False - self._expired = False - self._released = False - self._chip_release_committed = False - - @property - def expired(self): - return self._expired - - def _close_worker_host_mapping(self): - self.mapping_closed = True - - def free(self): - self._released = True + class FakeInstance: + def __init__(self, resource_id): + self.provider_resource_id = resource_id + self._close_attempted = False + self._state = worker_mod.RegionInstanceState.LIVE + self._cleanup_error = None + self._ever_live = True + self._provider_release_committed = False + self._data_plane_error = None + + def _retains_cleanup_only_reachability(self): + return False - def _expire(self): - self._expired = True + def _close_owned(self, *, poison_on_error): + self._close_attempted = True + self._state = worker_mod.RegionInstanceState.CLOSED + self._provider_release_committed = True class NativeWorker: def __init__(self): self.released_regions = [] - def control_worker_chip_region_release(self, worker_id, region_id): - self.released_regions.append((worker_id, region_id)) + def control_region_release(self, worker_id, request_shm_name, reply_shm_name): + _native_control_region_release(self.released_regions, worker_id, request_shm_name, reply_shm_name) class NativeOrchestrator: def __init__(self): @@ -4041,20 +4335,19 @@ def domain(name, allocation_id): first = worker_mod._RunResources() second = worker_mod._RunResources() first_ref, second_ref = SlotRef(), SlotRef() - first_region, second_region = Region(11), Region(22) + first_instance, second_instance = FakeInstance(11), FakeInstance(22) first_live, second_live = domain("first-live", 1), domain("second-live", 2) first_pending, second_pending = domain("first-pending", 3), domain("second-pending", 4) first.remote_slot_refs.append(cast(Any, first_ref)) second.remote_slot_refs.append(cast(Any, second_ref)) - first.worker_chip_regions.append(first_region) - second.worker_chip_regions.append(second_region) first.worker_chip_orch_comm_host_buffers[0x1000] = 64 second.worker_chip_orch_comm_host_buffers[0x2000] = 128 first.live_domains[first_live.name] = first_live second.live_domains[second_live.name] = second_live first.pending_release_domains.append(first_pending) second.pending_release_domains.append(second_pending) - worker._live_worker_chip_regions.extend([first_region, second_region]) + worker._region_instance_registry.track(cast(Any, first_instance), first) + worker._region_instance_registry.track(cast(Any, second_instance), second) worker._live_domains.update({first_live.name: first_live, second_live.name: second_live}) released_domains = [] @@ -4073,10 +4366,11 @@ def release_domain_now(handle): assert first_ref.releases == 1 assert second_ref.releases == 0 - assert native_worker.released_regions == [(0, 11)] - assert first_region.mapping_closed and first_region.expired - assert not second_region.mapping_closed and not second_region.expired - assert worker._live_worker_chip_regions == [second_region] + assert native_worker.released_regions == [] + assert first_instance._state is worker_mod.RegionInstanceState.CLOSED + assert second_instance._state is worker_mod.RegionInstanceState.LIVE + assert first_instance not in worker._region_instance_registry._instances.values() + assert second_instance in worker._region_instance_registry._instances.values() assert first.worker_chip_orch_comm_host_buffers == {} assert second.worker_chip_orch_comm_host_buffers == {0x2000: 128} assert released_domains == [first_pending, first_live] @@ -4090,34 +4384,36 @@ def test_region_mapping_failure_still_releases_child_and_poisons_successor(self) mapping_error = KeyboardInterrupt("mapping close") release_error = SystemExit("child release") - class Region: - region_id = 11 - _worker_id = 0 - + class FakeInstance: def __init__(self): - self._expired = False - self._released = False - self._chip_release_committed = False + self._close_attempted = False + self._state = worker_mod.RegionInstanceState.LIVE + self._cleanup_error = None + self._ever_live = True + self._provider_release_committed = False - @property - def expired(self): - return self._expired + def _retains_cleanup_only_reachability(self): + return False - def _close_worker_host_mapping(self): + def _close_owned(self, *, poison_on_error): + self._close_attempted = True + native_worker.released_regions.append((0, 11)) + self._provider_release_committed = True + self._state = worker_mod.RegionInstanceState.CLOSE_FAILED + if poison_on_error: + worker._record_unreclaimable( + "close_worker_chip_region: region 11 on worker 0 could not be " + "fully reclaimed; no further work is admitted", + mapping_error, + ) raise mapping_error - def free(self): - self._released = True - - def _expire(self): - self._expired = True - class NativeWorker: def __init__(self): self.released_regions = [] - def control_worker_chip_region_release(self, worker_id, region_id): - self.released_regions.append((worker_id, region_id)) + def control_region_release(self, worker_id, request_shm_name, reply_shm_name): + _native_control_region_release(self.released_regions, worker_id, request_shm_name, reply_shm_name) raise release_error class NativeOrchestrator: @@ -4130,83 +4426,20 @@ def _release_run(self, run_id): worker._orch = cast(Any, NativeOrchestrator()) resources = worker_mod._RunResources() resources.requires_ordered_cleanup = True - region = Region() - resources.worker_chip_regions.append(region) - worker._live_worker_chip_regions.append(region) + instance = FakeInstance() + worker._region_instance_registry.track(cast(Any, instance), resources) handle = RunHandle(worker, 1, (), resources) worker._accepted_run_handles.add(handle) assert worker._finalize_run_handle(handle, 1, None) is mapping_error assert native_worker.released_regions == [(0, 11)] - assert region.expired - assert resources.worker_chip_regions == [] - assert worker._live_worker_chip_regions == [] + assert instance._state is worker_mod.RegionInstanceState.CLOSE_FAILED + assert worker._region_instance_registry._instances == {} assert worker._ordered_cleanup_error is not None with pytest.raises(RuntimeError, match="no further work is admitted"): worker._require_no_ordered_cleanup_failure("submit") - def test_region_tracking_interrupt_still_retires_every_region(self): - interrupt = KeyboardInterrupt("region tracking removal") - - class InterruptingList(list): - def __init__(self, values): - super().__init__(values) - self.interrupted = False - - def __setitem__(self, key, value): - super().__setitem__(key, value) - if isinstance(key, slice) and not self.interrupted: - self.interrupted = True - raise interrupt - - class Region: - def __init__(self, region_id): - self.region_id = region_id - self._worker_id = 0 - self.mapping_closes = 0 - self.expires = 0 - self._released = False - self._chip_release_committed = False - - @property - def expired(self): - return self.expires > 0 - - def _close_worker_host_mapping(self): - self.mapping_closes += 1 - - def free(self): - self._released = True - - def _expire(self): - self.expires += 1 - - class NativeWorker: - def __init__(self): - self.released_regions = [] - - def control_worker_chip_region_release(self, worker_id, region_id): - self.released_regions.append((worker_id, region_id)) - - first, second = Region(11), Region(22) - resources = worker_mod._RunResources() - resources.worker_chip_regions = cast(Any, InterruptingList([first, second])) - worker = Worker(level=3, num_sub_workers=0) - native_worker = NativeWorker() - worker._worker = cast(Any, native_worker) - worker._live_worker_chip_regions.extend([first, second]) - - with pytest.raises(KeyboardInterrupt) as caught: - worker._cleanup_worker_chip_regions(resources) - - assert caught.value is interrupt - assert native_worker.released_regions == [(0, 11), (0, 22)] - assert (first.mapping_closes, first.expires) == (1, 1) - assert (second.mapping_closes, second.expires) == (1, 1) - assert resources.worker_chip_regions == [] - assert worker._live_worker_chip_regions == [] - def test_domain_released_after_its_run_retired_is_freed_inline(self): """A late release has no fence left to defer behind, so it frees now. @@ -6869,13 +7102,13 @@ def test_an_interrupt_through_a_region_create_refuses_further_work(self): child releases its own region before reporting one. """ worker = self._worker() - worker._config = {"platform": "a2a3sim"} + worker._config = {**worker._config, "platform": "a2a3sim", "device_ids": [0]} worker._validate_worker_chip_id = cast(Any, lambda _wid: None) def _interrupted(*_args): raise KeyboardInterrupt - worker._worker = cast(Any, SimpleNamespace(control_worker_chip_region_create=_interrupted)) + worker._worker = cast(Any, SimpleNamespace(control_region_allocate=_interrupted)) with pytest.raises(KeyboardInterrupt): worker._create_worker_chip_region(0, 4096, 64) assert worker._ordered_cleanup_error is not None @@ -6884,14 +7117,30 @@ def _interrupted(*_args): def test_an_ordinary_region_create_failure_does_not_refuse_further_work(self): worker = self._worker() - worker._config = {"platform": "a2a3sim"} + worker._config = {**worker._config, "platform": "a2a3sim", "device_ids": [0]} worker._validate_worker_chip_id = cast(Any, lambda _wid: None) def _failed(*_args): - raise RuntimeError("chip refused the region") + from simpler.comm_provider import ( + RegionAllocationError, + RegionControlErrorKind, + RegionOperationKind, + RegionPartKind, + ) + + raise RegionAllocationError( + provisional_resource_id=7, + control_kind=RegionControlErrorKind.BACKEND_FAILURE, + failed_part=RegionPartKind.PAYLOAD, + failed_operation=RegionOperationKind.MATERIALIZE, + cleanup_debt_remaining=False, + message="chip refused the region", + ) + + worker._worker = cast(Any, SimpleNamespace(control_region_allocate=_failed)) + from simpler.comm_provider import RegionAllocationError - worker._worker = cast(Any, SimpleNamespace(control_worker_chip_region_create=_failed)) - with pytest.raises(RuntimeError, match="chip refused the region"): + with pytest.raises(RegionAllocationError, match="chip refused the region"): worker._create_worker_chip_region(0, 4096, 64) assert worker._ordered_cleanup_error is None, "an ordinary failure must not shut the worker" diff --git a/tests/ut/py/test_worker/test_provider_region_onboard.py b/tests/ut/py/test_worker/test_provider_region_onboard.py new file mode 100644 index 0000000000..825134dac7 --- /dev/null +++ b/tests/ut/py/test_worker/test_provider_region_onboard.py @@ -0,0 +1,140 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 +"""Onboard proof that one provider region owns two independent VMM allocations.""" + +from __future__ import annotations + +import ctypes + +import pytest +from simpler import comm_provider_control, comm_region +from simpler import worker as worker_module +from simpler.task_interface import DataType +from simpler.worker import Worker +from simpler.worker_chip_orch_comm import NotifyOp, WaitCmp + +from simpler_setup.runtime_builder import RuntimeBuilder + +_RUNTIME = "tensormap_and_ringbuffer" +_PAYLOAD_BYTES = 16 +_COUNTER_BYTES = 128 +_HOST_MARKER = 0x11 + + +def _require_runtime(platform: str) -> None: + try: + RuntimeBuilder(platform=platform).get_binaries(_RUNTIME) + except (FileNotFoundError, RuntimeError) as e: + pytest.skip(f"{platform} runtime binaries unavailable: {e}") + + +def _install_lifecycle_spies(monkeypatch) -> list[tuple]: + events: list[tuple] = [] + original_import = worker_module._worker_host_mapped_region_import_onboard + original_close = comm_region._worker_host_mapped_region_close + original_release = comm_provider_control.ProviderReleaseClient.release + + def spy_import(device_id, shareable_handle, mapping_bytes, owner_token): + handle = original_import(device_id, shareable_handle, mapping_bytes, owner_token) + events.append(("import", int(handle), int(shareable_handle), int(mapping_bytes))) + return handle + + def spy_close(handle): + events.append(("close", int(handle))) + return original_close(handle) + + def spy_release(self, provider_resource_id): + events.append(("release", int(provider_resource_id))) + return original_release(self, provider_resource_id) + + monkeypatch.setattr(worker_module, "_worker_host_mapped_region_import_onboard", spy_import) + monkeypatch.setattr(comm_region, "_worker_host_mapped_region_close", spy_close) + monkeypatch.setattr(comm_provider_control.ProviderReleaseClient, "release", spy_release) + return events + + +def _assert_independent_handles(imports: list[tuple[int, int, int]]) -> None: + assert len(imports) == 2 + first_handle, first_shareable, first_bytes = imports[0] + second_handle, second_shareable, second_bytes = imports[1] + assert first_handle != second_handle + assert first_shareable != second_shareable + # Import order is PAYLOAD then COUNTER. mapping_bytes covers each logical + # span plus VMM granularity and need not equal the requested sizes. + assert first_bytes >= _PAYLOAD_BYTES + assert second_bytes >= _COUNTER_BYTES + + +def _assert_run_teardown(events: list[tuple], resource_id: int) -> None: + assert [item[0] for item in events] == ["import", "import", "close", "close", "release"] + _assert_independent_handles([item[1:] for item in events[:2]]) + assert {events[2][1], events[3][1]} == {events[0][1], events[1][1]} + assert events[4][1] == resource_id + + +@pytest.mark.requires_hardware +@pytest.mark.platforms(["a2a3", "a5"]) +@pytest.mark.device_count(1) +@pytest.mark.runtime("tensormap_and_ringbuffer") +def test_onboard_provider_region_host_two_lifecycles(st_platform, st_device_ids, monkeypatch): + _require_runtime(st_platform) + events = _install_lifecycle_spies(monkeypatch) + worker = Worker( + level=3, + device_ids=[int(st_device_ids[0])], + num_sub_workers=0, + platform=st_platform, + runtime=_RUNTIME, + ) + worker.init() + resource_ids: list[int] = [] + try: + + def orch(orch_handle, _args, _cfg): + host = orch_handle.alloc([_PAYLOAD_BYTES], DataType.UINT8) + buf = (ctypes.c_uint8 * _PAYLOAD_BYTES).from_address(int(host.base)) + for i in range(_PAYLOAD_BYTES): + buf[i] = _HOST_MARKER + region = orch_handle.create_worker_chip_region( + worker_id=0, payload_bytes=_PAYLOAD_BYTES, counter_bytes=_COUNTER_BYTES + ) + resource_ids.append(int(region.descriptor_scalars()[1])) + assert int(region._instance._payload_mapping) != int(region._instance._counter_mapping) + assert int(region._instance._payload_local_view.logical_bytes) == _PAYLOAD_BYTES + assert int(region._instance._counter_local_view.logical_bytes) == _COUNTER_BYTES + region.payload_write(0, host, nbytes=_PAYLOAD_BYTES) + region.counter(0).notify(1, NotifyOp.Set) + assert region.counter(0).test(1, WaitCmp.EQ).matched + for i in range(_PAYLOAD_BYTES): + buf[i] = 0 + region.payload_read(0, host, nbytes=_PAYLOAD_BYTES) + assert bytes(buf) == bytes([_HOST_MARKER] * _PAYLOAD_BYTES) + assert [item[0] for item in events] == ["import", "import"] + region.free() + assert region._instance._state is comm_region.RegionInstanceState.LIVE + assert [item[0] for item in events] == ["import", "import"] + + worker.run(orch) + first = events[:] + _assert_run_teardown(first, resource_ids[0]) + assert worker._region_instance_registry._instances == {} + assert worker._region_instance_registry._instances == {} + + events.clear() + worker.run(orch) + second = events[:] + _assert_run_teardown(second, resource_ids[1]) + assert {first[0][2], first[1][2]}.isdisjoint({second[0][2], second[1][2]}) + assert resource_ids[0] != resource_ids[1] + assert worker._region_instance_registry._instances == {} + assert worker._region_instance_registry._instances == {} + finally: + worker.close() + assert worker._region_instance_registry._instances == {} diff --git a/tests/ut/py/test_worker/test_worker_chip_message_queue.py b/tests/ut/py/test_worker/test_worker_chip_message_queue.py index f0f6bca6b7..abc3862309 100644 --- a/tests/ut/py/test_worker/test_worker_chip_message_queue.py +++ b/tests/ut/py/test_worker/test_worker_chip_message_queue.py @@ -16,9 +16,25 @@ from typing import Optional import pytest -from simpler import comm_region, worker_chip_orch_comm +from simpler import comm_region from simpler import worker as worker_module from simpler.buffer import AccessMode, BackendKind, CanonicalIdentity, mint_owner_instance_id, wrap_fork_inherited +from simpler.comm_provider import ( + PosixShmImport, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationResult, + RegionExportDescriptor, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, +) +from simpler.comm_provider_control import ( + decode_allocate_request, + decode_release_request, + encode_allocate_success_reply, + encode_release_result_reply, +) from simpler.orchestrator import Orchestrator from simpler.task_interface import DataType, get_element_size from simpler.worker import _IDLE, _OFF_STATE, Worker, _buffer_field_addr, _mailbox_store_i32 @@ -37,13 +53,42 @@ WaitCmp, WorkerChipOrchRegion, WorkerChipOrchRegionDesc, - WorkerChipRegionAccessProfile, - WorkerHostRegionMapping, ) _DESC_BID = itertools.count(1) +class _CompatInstance: + def __init__(self, payload_handle: int, counter_handle: int, payload_bytes: int, counter_bytes: int) -> None: + self.worker_id = 0 + self._data_plane_error = None + self._payload_part = comm_region.PayloadPart( + comm_region.RegionPartSpan(offset=0, nbytes=int(payload_bytes)), + comm_region.HostVmmCopyAccess(payload_handle), + ) + self._counter_part = comm_region.CounterPart( + comm_region.RegionPartSpan(offset=0, nbytes=int(counter_bytes)), + comm_region.HostVmmCopyAccess(counter_handle), + ) + + @property + def state(self): + return comm_region.RegionInstanceState.LIVE + + @property + def data_plane_error(self): + return self._data_plane_error + + def payload_write(self, offset, host_buffer, nbytes=None): + self._payload_part.write(offset, host_buffer, nbytes) + + def payload_read(self, offset, host_buffer, nbytes=None): + self._payload_part.read(offset, host_buffer, nbytes) + + def counter(self, offset): + return self._counter_part.counter(offset) + + def _fake_alloc_handle(orch, nbytes): """A FORK_SHM Buffer over a bare _FakeCOrch alloc — mirrors Orchestrator.alloc for the low-level tests that drive WorkerChipQueue with a fake C orch directly.""" @@ -72,7 +117,7 @@ class _FakeCWorker: def __init__(self): self.next_region_id = 1 - def control_worker_chip_region_create(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: req_shm = SharedMemory(name=request_shm_name) reply_shm = SharedMemory(name=reply_shm_name) req_buf = req_shm.buf @@ -80,28 +125,33 @@ def control_worker_chip_region_create(self, worker_id: int, request_shm_name: st assert req_buf is not None assert reply_buf is not None try: - req = worker_chip_orch_comm._REGION_CREATE_REQUEST.unpack_from(req_buf, 0) - payload_bytes = int(req[2]) - counter_bytes = int(req[3]) - counter_offset = ((payload_bytes + 63) // 64) * 64 + spec = decode_allocate_request(req_buf) + payload_bytes = int(spec.payload.logical_bytes) + counter_bytes = int(spec.counter.logical_bytes) region_id = self.next_region_id self.next_region_id += 1 - backing_name = f"queue-direct-{region_id}".encode() - worker_chip_orch_comm._REGION_CREATE_REPLY.pack_into( + result = RegionAllocationResult( + provider_resource_id=region_id, + export_descriptor=RegionExportDescriptor( + payload=RegionPartExportDescriptor( + BackendKind.VMM_WINDOW, + payload_bytes, + payload_bytes, + PosixShmImport(f"queue-direct-{region_id}-p"), + ), + counter=RegionPartExportDescriptor( + BackendKind.VMM_WINDOW, + counter_bytes, + counter_bytes, + PosixShmImport(f"queue-direct-{region_id}-c"), + ), + ), + ) + encode_allocate_success_reply( reply_buf, - 0, - 0x4C334C3200020000, - region_id, - 0x1000_0000, - payload_bytes, - 0x1000_0000 + counter_offset, - counter_bytes, - int(WorkerChipRegionAccessProfile.SIM_POSIX_SHM), - 0, - 0, - backing_name + b"\x00" * (worker_chip_orch_comm._CTRL_SHM_TOKEN_BYTES - len(backing_name)), - counter_offset + counter_bytes, - 0, + result, + RegionPartLocalView(RegionPartKind.PAYLOAD, 0x1000_0000, payload_bytes), + RegionPartLocalView(RegionPartKind.COUNTER, 0x1000_1000, counter_bytes), ) finally: del req_buf @@ -109,8 +159,24 @@ def control_worker_chip_region_create(self, worker_id: int, request_shm_name: st req_shm.close() reply_shm.close() - def control_worker_chip_region_release(self, worker_id: int, region_id: int) -> None: - pass + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + req_shm = SharedMemory(name=request_shm_name) + reply_shm = SharedMemory(name=reply_shm_name) + req_buf = req_shm.buf + reply_buf = reply_shm.buf + assert req_buf is not None + assert reply_buf is not None + try: + resource_id = decode_release_request(req_buf) + encode_release_result_reply( + reply_buf, + ProviderReleaseResult(provider_resource_id=int(resource_id), status=ProviderReleaseStatus.RELEASED), + ) + finally: + del req_buf + del reply_buf + req_shm.close() + reply_shm.close() class _FakeCOrch: @@ -142,16 +208,19 @@ def __init__(self): self.fail_next_cmd: Optional[str] = None self.original_helpers: list[tuple[object, str, object]] = [] - def import_region(self, _token: str, mapping_bytes: int, _owner_token: str) -> int: + def import_region(self, token: str, mapping_bytes: int, _owner_token: str) -> int: + if str(token).endswith("-c"): + self.counter_mapping_offset = 0 + self.counter_base = 0x1000_1000 + return 2 self.payload = bytearray(int(mapping_bytes)) self.counters = {} - self.counter_mapping_offset = int(mapping_bytes) - WORKER_CHIP_QUEUE_COUNTER_BYTES - self.counter_base = self.payload_base + self.counter_mapping_offset + self.counter_mapping_offset = 0 self.requests.append( ( _FakeRequest( cmd="alloc_region", - payload_bytes=self.counter_mapping_offset, + payload_bytes=int(mapping_bytes), counter_bytes=WORKER_CHIP_QUEUE_COUNTER_BYTES, ), 0.0, @@ -250,53 +319,22 @@ def _make_orchestrator() -> tuple[Orchestrator, Worker, SharedMemory, _FakeClien fake_client = _FakeClient() fake_client.original_helpers = [ (worker_module, "_worker_host_mapped_region_import_sim", worker_module._worker_host_mapped_region_import_sim), - ( - worker_chip_orch_comm, - "_worker_host_mapped_payload_write", - worker_chip_orch_comm._worker_host_mapped_payload_write, - ), - ( - worker_chip_orch_comm, - "_worker_host_mapped_payload_read", - worker_chip_orch_comm._worker_host_mapped_payload_read, - ), - ( - worker_chip_orch_comm, - "_worker_host_mapped_counter_notify", - worker_chip_orch_comm._worker_host_mapped_counter_notify, - ), - ( - worker_chip_orch_comm, - "_worker_host_mapped_counter_test", - worker_chip_orch_comm._worker_host_mapped_counter_test, - ), - ( - worker_chip_orch_comm, - "_worker_host_mapped_counter_wait", - worker_chip_orch_comm._worker_host_mapped_counter_wait, - ), - ( - worker_chip_orch_comm, - "_worker_host_mapped_region_close", - worker_chip_orch_comm._worker_host_mapped_region_close, - ), (comm_region, "_host_vmm_copy_to", comm_region._host_vmm_copy_to), (comm_region, "_host_vmm_copy_from", comm_region._host_vmm_copy_from), (comm_region, "_region_counter_notify", comm_region._region_counter_notify), (comm_region, "_region_counter_test", comm_region._region_counter_test), (comm_region, "_region_counter_wait", comm_region._region_counter_wait), + (comm_region, "_worker_host_mapped_region_close", comm_region._worker_host_mapped_region_close), ] worker._lifecycle = worker_module._Lifecycle.READY worker._worker = _FakeCWorker() worker._chip_shms = [shm] worker._worker_chip_test_fake_client = fake_client + reservation = worker._control_reservation("test_worker_chip_queue") + reservation.__enter__() + worker._worker_chip_test_reservation = reservation worker_module._worker_host_mapped_region_import_sim = fake_client.import_region - worker_chip_orch_comm._worker_host_mapped_payload_write = fake_client.payload_write - worker_chip_orch_comm._worker_host_mapped_payload_read = fake_client.payload_read - worker_chip_orch_comm._worker_host_mapped_counter_notify = fake_client.counter_notify - worker_chip_orch_comm._worker_host_mapped_counter_test = fake_client.counter_test - worker_chip_orch_comm._worker_host_mapped_counter_wait = fake_client.counter_wait - worker_chip_orch_comm._worker_host_mapped_region_close = lambda _handle: None + comm_region._worker_host_mapped_region_close = lambda _handle: None comm_region._host_vmm_copy_to = fake_client.payload_write comm_region._host_vmm_copy_from = fake_client.payload_read comm_region._region_counter_notify = fake_client.counter_notify @@ -306,6 +344,10 @@ def _make_orchestrator() -> tuple[Orchestrator, Worker, SharedMemory, _FakeClien def _close(worker: Worker, shm: SharedMemory) -> None: + reservation = getattr(worker, "_worker_chip_test_reservation", None) + if reservation is not None: + reservation.__exit__(None, None, None) + worker._worker_chip_test_reservation = None worker._close_worker_chip_orch_comm() fake_client = getattr(worker, "_worker_chip_test_fake_client", None) if fake_client is not None: @@ -460,8 +502,7 @@ def fail_alloc(_shape, _dtype, _identity): with pytest.raises(RuntimeError, match="injected alloc failure"): orch.create_worker_chip_queue(worker_id=0, depth=4, input_arena_bytes=128, output_arena_bytes=128) - assert len(worker._live_worker_chip_regions) == 1 - assert worker._live_worker_chip_regions[0]._released is True + assert len(worker._region_instance_registry._instances) == 1 finally: orch._o.alloc = original_alloc_ref _close(worker, shm) @@ -549,7 +590,7 @@ def test_worker_host_mapped_queue_ordinary_input_uses_direct_payload_write(monke orch = _FakeCOrch() layout = make_worker_chip_queue_layout(4, 128, 128) desc = WorkerChipOrchRegionDesc( - magic_version=0x4C334C3200020000, + magic_version=0x4C334C3200030000, region_id=1, payload_base=0x1000_0000, payload_bytes=layout.payload_bytes, @@ -558,19 +599,8 @@ def test_worker_host_mapped_queue_ordinary_input_uses_direct_payload_write(monke ) region = WorkerChipOrchRegion( object(), - 0, + _CompatInstance(44, 45, desc.payload_bytes, desc.counter_bytes), desc, - WorkerHostRegionMapping( - worker_id=0, - region_id=1, - access_profile=WorkerChipRegionAccessProfile.SIM_POSIX_SHM, - total_bytes=desc.counter_base - desc.payload_base + desc.counter_bytes, - payload_offset=0, - payload_bytes=desc.payload_bytes, - counter_offset=desc.counter_base - desc.payload_base, - counter_bytes=desc.counter_bytes, - handle=44, - ), ) queue = WorkerChipQueue( orch, @@ -602,7 +632,7 @@ def counter_notify(_handle: int, offset: int, value: int, _op: int) -> None: assert (layout.input_arena_offset, b"ordinary") in payload_writes assert len(orch._buffers) == alloc_count - assert counters[region._worker_host_mapping.counter_offset + layout.input_desc_tail_offset] == 1 + assert counters[layout.input_desc_tail_offset] == 1 def test_direct_mapped_ordinary_host_bytearray_does_not_allocate_queue_buffer(): @@ -985,7 +1015,7 @@ def test_expired_queue_rejects_later_operations_without_abort_flag(): orch, worker, shm, fake_client = _make_orchestrator() try: queue = orch.create_worker_chip_queue(worker_id=0, depth=4, input_arena_bytes=128, output_arena_bytes=128) - queue.region._expire() + queue.region._instance._state = comm_region.RegionInstanceState.CLOSED fake_client.requests.clear() with pytest.raises(RuntimeError, match="expired"): diff --git a/tests/ut/py/test_worker/test_worker_chip_orch_comm.py b/tests/ut/py/test_worker/test_worker_chip_orch_comm.py index 977f8f52de..046ff98c9a 100644 --- a/tests/ut/py/test_worker/test_worker_chip_orch_comm.py +++ b/tests/ut/py/test_worker/test_worker_chip_orch_comm.py @@ -6,26 +6,58 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +# ruff: noqa: PLC0415 +import contextlib import ctypes import gc import importlib import os +import struct import threading import time from concurrent.futures import ThreadPoolExecutor from multiprocessing.shared_memory import SharedMemory from pathlib import Path -from typing import Any, Optional, cast +from typing import Any, Optional, Union, cast +from unittest.mock import MagicMock import pytest +from simpler import comm_endpoints as ce from simpler import comm_region, worker_chip_orch_comm from simpler import worker as worker_module -from simpler.buffer import mint_owner_instance_id, wrap_fork_inherited +from simpler.buffer import BackendKind, mint_owner_instance_id, wrap_fork_inherited +from simpler.comm_provider import ( + PosixShmImport, + ProviderRegionStore, + ProviderReleaseResult, + ProviderReleaseStatus, + RegionAllocationResult, + RegionControlError, + RegionControlErrorKind, + RegionExportDescriptor, + RegionPartExportDescriptor, + RegionPartKind, + RegionPartLocalView, + VmmShareableHandleImport, +) +from simpler.comm_provider_control import ( + ALLOCATE_REPLY_BYTES, + ALLOCATE_REQUEST_BYTES, + AllocateReplyTag, + decode_allocate_reply, + decode_allocate_request, + decode_release_request, + encode_allocate_request, + encode_allocate_success_reply, + encode_release_result_reply, +) +from simpler.orchestrator import Orchestrator from simpler.task_interface import DataType from simpler.worker import ( _IDLE, _OFF_STATE, + MAILBOX_SIZE, Worker, _buffer_field_addr, _mailbox_store_i32, @@ -40,6 +72,9 @@ _task_interface_ext = cast(Any, importlib.import_module("_task_interface")) _TASK_INTERFACE_CPP = Path(__file__).resolve().parents[4] / "python" / "bindings" / "task_interface.cpp" +_WORKER_PY = Path(__file__).resolve().parents[4] / "python" / "simpler" / "worker.py" +_ACCESS_ONBOARD_VMM = 1 +_ACCESS_SIM_POSIX_SHM = 2 def test_worker_host_mapped_backend_types_are_removed_from_native_implementation(): @@ -52,34 +87,86 @@ def test_worker_host_mapped_backend_types_are_removed_from_native_implementation "WorkerHostMappedRegionRegistry", "WorkerHostMappedRegionHandle", "class WorkerHostMappedRegion", + "ChipChildOnboardRegionExport", + "region_vmm_legacy_create", + "_l3_child_onboard_region_create", + "_l3_child_onboard_region_close", ): assert legacy_type not in source +def test_combined_host_worker_chip_region_helpers_are_removed(): + source = _WORKER_PY.read_text(encoding="utf-8") + for legacy_name in ( + "_HostWorkerChipRegion", + "_HostWorkerChipRegionStore", + "_release_host_worker_chip_region", + "_create_sim_worker_chip_region", + "_create_onboard_worker_chip_region", + "_sweep_host_" + "worker_chip" + "_regions", + "_close_worker_host_mapping", + ): + assert legacy_name not in source + + +def test_old_create_reply_codec_and_combined_mapping_remain_deleted(): + worker_source = _WORKER_PY.read_text(encoding="utf-8") + region_source = Path(comm_region.__file__).read_text(encoding="utf-8") + binding_source = _TASK_INTERFACE_CPP.read_text(encoding="utf-8") + for source in (worker_source, region_source, binding_source): + for legacy_name in ( + "validate_region_create_reply", + "decode_region_create_reply", + "WorkerHostRegionMapping", + "worker_chip_orch_region_access.h", + "ChipChildOnboardRegionExport", + "region_vmm_legacy_create", + ): + assert legacy_name not in source + access_header = ( + Path(__file__).resolve().parents[4] + / "src" + / "common" + / "platform" + / "include" + / "host" + / "worker_chip_orch_region_access.h" + ) + assert not access_header.exists() + + class _FakeDirectCWorker: def __init__( self, *, payload_base: int = 0xDEAD_0000, - access_profile: int = int(worker_chip_orch_comm.WorkerChipRegionAccessProfile.SIM_POSIX_SHM), + counter_base: int = 0xDEAD_1000, + access_profile: int = _ACCESS_SIM_POSIX_SHM, device_id: int = 0, shareable_handle: int = 0xABCDEF, - magic_version: int = 0x4C334C3200020000, + reply_magic: Optional[int] = None, region_id: Optional[int] = None, mapping_bytes: Optional[int] = None, + corrupt_access_profile: bool = False, ): self.create_calls: list[tuple[int, str, str]] = [] self.release_calls: list[tuple[int, int]] = [] self.next_region_id = 1 self.payload_base = int(payload_base) + self.counter_base = int(counter_base) self.access_profile = int(access_profile) self.device_id = int(device_id) self.shareable_handle = int(shareable_handle) - self.magic_version = int(magic_version) + self.reply_magic = reply_magic self.region_id = region_id self.mapping_bytes = mapping_bytes + self.corrupt_access_profile = bool(corrupt_access_profile) + self.allocate_specs: list[Any] = [] + + def close(self) -> None: + return None - def control_worker_chip_region_create(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + def control_region_allocate(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: self.create_calls.append((int(worker_id), str(request_shm_name), str(reply_shm_name))) req_shm = SharedMemory(name=request_shm_name) reply_shm = SharedMemory(name=reply_shm_name) @@ -88,45 +175,74 @@ def control_worker_chip_region_create(self, worker_id: int, request_shm_name: st assert req_buf is not None assert reply_buf is not None try: - req = worker_chip_orch_comm._REGION_CREATE_REQUEST.unpack_from(req_buf, 0) - payload_bytes = int(req[2]) - counter_bytes = int(req[3]) - counter_offset = ((payload_bytes + 63) // 64) * 64 + spec = decode_allocate_request(req_buf) + self.allocate_specs.append(spec) region_id = int(self.region_id) if self.region_id is not None else self.next_region_id if self.region_id is None: self.next_region_id += 1 - backing_name = f"sim-direct-{region_id}".encode() - if self.access_profile == int(worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM): - backing_name = b"" - shareable_handle = ( - self.shareable_handle - if self.access_profile == int(worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM) - else 0 + sim = self.access_profile == _ACCESS_SIM_POSIX_SHM + if self.corrupt_access_profile: + sim = not sim + payload_bytes = int(spec.payload.logical_bytes) + counter_bytes = int(spec.counter.logical_bytes) + payload_mapping = payload_bytes if self.mapping_bytes is None else int(self.mapping_bytes) + counter_mapping = counter_bytes if self.mapping_bytes is None else int(self.mapping_bytes) + if sim: + payload_cap: Union[PosixShmImport, VmmShareableHandleImport] = PosixShmImport( + f"sim-direct-{region_id}-p" + ) + counter_cap: Union[PosixShmImport, VmmShareableHandleImport] = PosixShmImport( + f"sim-direct-{region_id}-c" + ) + else: + payload_cap = VmmShareableHandleImport(self.device_id, self.shareable_handle) + counter_cap = VmmShareableHandleImport(self.device_id, self.shareable_handle + 1) + result = RegionAllocationResult( + provider_resource_id=max(region_id, 1), + export_descriptor=RegionExportDescriptor( + payload=RegionPartExportDescriptor( + BackendKind.VMM_WINDOW, payload_bytes, max(payload_mapping, payload_bytes), payload_cap + ), + counter=RegionPartExportDescriptor( + BackendKind.VMM_WINDOW, counter_bytes, max(counter_mapping, counter_bytes), counter_cap + ), + ), ) - worker_chip_orch_comm._REGION_CREATE_REPLY.pack_into( + encode_allocate_success_reply( reply_buf, - 0, - self.magic_version, - region_id, - self.payload_base, - payload_bytes, - self.payload_base + counter_offset, - counter_bytes, - self.access_profile, - 0, - self.device_id, - backing_name + b"\x00" * (worker_chip_orch_comm._CTRL_SHM_TOKEN_BYTES - len(backing_name)), - counter_offset + counter_bytes if self.mapping_bytes is None else int(self.mapping_bytes), - shareable_handle, + result, + RegionPartLocalView(RegionPartKind.PAYLOAD, self.payload_base, payload_bytes), + RegionPartLocalView(RegionPartKind.COUNTER, self.counter_base, counter_bytes), ) + if region_id == 0: + struct.pack_into(" None: - self.release_calls.append((int(worker_id), int(region_id))) + def control_region_release(self, worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + req_shm = SharedMemory(name=request_shm_name) + reply_shm = SharedMemory(name=reply_shm_name) + req_buf = req_shm.buf + reply_buf = reply_shm.buf + assert req_buf is not None + assert reply_buf is not None + try: + resource_id = decode_release_request(req_buf) + self.release_calls.append((int(worker_id), int(resource_id))) + encode_release_result_reply( + reply_buf, + ProviderReleaseResult(provider_resource_id=int(resource_id), status=ProviderReleaseStatus.RELEASED), + ) + finally: + del req_buf + del reply_buf + req_shm.close() + reply_shm.close() class _EndpointFailingOrch: @@ -156,32 +272,6 @@ def _release_run(self, run_id: int) -> None: assert run_id == 1 -class _NamedOnboardRegionExport: - def __init__(self) -> None: - self.device_addr = 0xD00D_0000 - self.mapping_bytes = 65536 - self.shareable_handle = 0xABCDEF - self.registry_handle = 77 - - -class _FakeChipWorkerForRegionCreate: - """Collaborator double for the in-process L3-L2 region-create handlers. - - Distinct from ``_harness.FakeChipWorker``, which stands in for the class the - forked chip child instantiates: this one is never bound as - ``worker.ChipWorker``, so it implements only the device-memory surface a - region create touches. - """ - - device_id = 2 - - def __init__(self) -> None: - self.copy_calls: list[tuple[int, int]] = [] - - def copy_to(self, dst: int, _src: int, size: int) -> None: - self.copy_calls.append((int(dst), int(size))) - - def _write_ctrl_shm_name(buf: memoryview, offset: int, name: str) -> None: encoded = name.encode("utf-8") assert len(encoded) < worker_module._CTRL_SHM_NAME_BYTES @@ -191,7 +281,7 @@ def _write_ctrl_shm_name(buf: memoryview, offset: int, name: str) -> None: def _make_started_sim_worker() -> tuple[Worker, SharedMemory, _FakeDirectCWorker]: worker = Worker(level=3, device_ids=[0], platform="a2a3sim", runtime="tensormap_and_ringbuffer") - shm = SharedMemory(create=True, size=4096) + shm = SharedMemory(create=True, size=MAILBOX_SIZE) assert shm.buf is not None _mailbox_store_i32(_buffer_field_addr(shm.buf, _OFF_STATE), _IDLE) fake_c_worker = _FakeDirectCWorker() @@ -203,11 +293,11 @@ def _make_started_sim_worker() -> tuple[Worker, SharedMemory, _FakeDirectCWorker def _make_started_onboard_worker(platform: str = "a2a3") -> tuple[Worker, SharedMemory, _FakeDirectCWorker]: worker = Worker(level=3, device_ids=[2], platform=platform, runtime="tensormap_and_ringbuffer") - shm = SharedMemory(create=True, size=4096) + shm = SharedMemory(create=True, size=MAILBOX_SIZE) assert shm.buf is not None _mailbox_store_i32(_buffer_field_addr(shm.buf, _OFF_STATE), _IDLE) fake_c_worker = _FakeDirectCWorker( - access_profile=int(worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM), + access_profile=_ACCESS_ONBOARD_VMM, device_id=2, ) worker._lifecycle = worker_module._Lifecycle.READY @@ -248,23 +338,24 @@ def test_sim_direct_region_uses_lifecycle_control_and_worker_host_metadata(monke region = worker._create_worker_chip_region(0, 64, 128) payload = wrap_fork_inherited(0x1234_0000, 16, mint_owner_instance_id(), 1, "L3") - region.payload_write(0, payload, nbytes=8) - region.payload_read(8, payload, nbytes=8) - result = region.counter(64).test(7, WaitCmp.EQ) - region.counter(64).notify(3, NotifyOp.Set) + with worker._control_reservation("test_sim_direct_region"): + region.payload_write(0, payload, nbytes=8) + region.payload_read(8, payload, nbytes=8) + result = region.counter(64).test(7, WaitCmp.EQ) + region.counter(64).notify(3, NotifyOp.Set) assert len(fake_c_worker.create_calls) == 1 - assert region.descriptor_scalars() == [0x4C334C3200020000, 1, 0xDEAD_0000, 64, 0xDEAD_0040, 128] + assert region.descriptor_scalars() == [0x4C334C3200030000, 1, 0xDEAD_0000, 64, 0xDEAD_1000, 128] assert 99 not in region.descriptor_scalars() - worker_host_mapping = region._worker_host_mapping - assert worker_host_mapping is not None - assert worker_host_mapping.handle != region.descriptor.payload_base - assert worker_host_mapping.counter_offset == 64 - assert calls[0] == ("import", "sim-direct-1", 192, worker._owner_id) - assert calls[1][0:3] == ("write", 99, 0) - assert calls[2][0:3] == ("read", 99, 8) - assert calls[3] == ("test", 99, 128, 7, int(WaitCmp.EQ)) - assert calls[4] == ("notify", 99, 128, 3, int(NotifyOp.Set)) + assert region._instance._payload_mapping == 99 + assert region._instance._counter_mapping == 99 + assert region._instance._payload_mapping != region.descriptor.payload_base + assert calls[0] == ("import", "sim-direct-1-p", 64, worker._owner_id) + assert calls[1] == ("import", "sim-direct-1-c", 128, worker._owner_id) + assert calls[2][0:3] == ("write", 99, 0) + assert calls[3][0:3] == ("read", 99, 8) + assert calls[4] == ("test", 99, 64, 7, int(WaitCmp.EQ)) + assert calls[5] == ("notify", 99, 64, 3, int(NotifyOp.Set)) assert result == SignalTestResult(matched=True, observed=7) finally: worker._close_worker_chip_orch_comm() @@ -291,17 +382,40 @@ def test_onboard_direct_region_imports_vmm_shareable_handle_and_uses_worker_host ) region = worker._create_worker_chip_region(0, 64, 128) - region.counter(64).notify(9, NotifyOp.Set) + with worker._control_reservation("test_onboard_direct_region"): + region.counter(64).notify(9, NotifyOp.Set) assert len(fake_c_worker.create_calls) == 1 - assert region.descriptor_scalars() == [0x4C334C3200020000, 1, 0xDEAD_0000, 64, 0xDEAD_0040, 128] + assert region.descriptor_scalars() == [0x4C334C3200030000, 1, 0xDEAD_0000, 64, 0xDEAD_1000, 128] assert 123 not in region.descriptor_scalars() - worker_host_mapping = region._worker_host_mapping - assert worker_host_mapping is not None - assert worker_host_mapping.access_profile == worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM - assert worker_host_mapping.counter_offset == 64 - assert calls[0] == ("import_onboard", 2, 0xABCDEF, 192, worker._owner_id) - assert calls[1] == ("notify", 123, 128, 9, int(NotifyOp.Set)) + assert region._instance._payload_mapping == 123 + assert calls[0] == ("import_onboard", 2, 0xABCDEF, 64, worker._owner_id) + assert calls[1] == ("import_onboard", 2, 0xABCDEF + 1, 128, worker._owner_id) + assert calls[2] == ("notify", 123, 64, 9, int(NotifyOp.Set)) + finally: + worker._close_worker_chip_orch_comm() + shm.close() + shm.unlink() + + +def test_sim_direct_create_sends_projected_vmm_window_spec(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + try: + monkeypatch.setattr( + worker_module, + "_worker_host_mapped_region_import_sim", + lambda _token, _mapping_bytes, _owner_token: 99, + ) + + region = worker._create_worker_chip_region(0, 64, 128) + + assert len(fake_c_worker.allocate_specs) == 1 + spec = fake_c_worker.allocate_specs[0] + assert spec.payload.planned_backing_kind is BackendKind.VMM_WINDOW + assert spec.counter.planned_backing_kind is BackendKind.VMM_WINDOW + assert spec.payload.logical_bytes == 64 + assert spec.counter.logical_bytes == 128 + assert region.region_id == 1 finally: worker._close_worker_chip_orch_comm() shm.close() @@ -322,7 +436,9 @@ def test_sim_direct_create_import_failure_rolls_back_l2_host_region(monkeypatch) assert fake_c_worker.create_calls assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [] + unpublished = tuple(worker._region_instance_registry._instances.values()) + assert len(unpublished) == 1 + assert unpublished[0]._state is comm_region.RegionInstanceState.CLOSE_FAILED finally: worker._close_worker_chip_orch_comm() shm.close() @@ -331,13 +447,15 @@ def test_sim_direct_create_import_failure_rolls_back_l2_host_region(monkeypatch) def test_direct_create_decode_failure_rolls_back_l2_host_region(): worker, shm, fake_c_worker = _make_started_sim_worker() - fake_c_worker.access_profile = 99 + fake_c_worker.access_profile = _ACCESS_ONBOARD_VMM try: - with pytest.raises(ValueError, match="99 is not a valid"): + with pytest.raises(RuntimeError, match="committed import capability does not match"): worker._create_worker_chip_region(0, 64, 128) assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [] + unpublished = tuple(worker._region_instance_registry._instances.values()) + assert len(unpublished) == 1 + assert unpublished[0]._state is comm_region.RegionInstanceState.CLOSE_FAILED finally: worker._close_worker_chip_orch_comm() shm.close() @@ -345,34 +463,38 @@ def test_direct_create_decode_failure_rolls_back_l2_host_region(): @pytest.mark.parametrize( - ("reply_updates", "match"), + ("reply_updates", "exc_type", "match", "released_id"), [ - ({"magic_version": 0xBAD}, "magic_version is invalid"), - ({"region_id": 0}, "region_id must be nonzero"), + ({"reply_magic": 0xBAD}, RegionControlError, "unsupported provider control version", None), + ({"reply_magic": 0x4C334C3200020000}, RegionControlError, "unsupported provider control version", None), + ({"region_id": 0}, RegionControlError, "SUCCESS requires a resource id", None), ( - {"access_profile": int(worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM)}, - "access_profile must be sim_posix_shm", + {"access_profile": _ACCESS_ONBOARD_VMM}, + RuntimeError, + "committed import capability does not match", + 1, ), ], ) -def test_direct_create_validation_failure_rolls_back_l2_host_region(reply_updates, match): +def test_direct_create_validation_failure_rolls_back_l2_host_region(reply_updates, exc_type, match, released_id): worker, shm, fake_c_worker = _make_started_sim_worker() for name, value in reply_updates.items(): setattr(fake_c_worker, name, value) - expected_region_id = int(reply_updates.get("region_id", 1)) try: - with pytest.raises(RuntimeError, match=match): + with pytest.raises(exc_type, match=match): worker._create_worker_chip_region(0, 64, 128) - assert fake_c_worker.release_calls == ([(0, expected_region_id)] if expected_region_id else []) - assert worker._live_worker_chip_regions == [] + assert fake_c_worker.release_calls == ([] if released_id is None else [(0, released_id)]) + unpublished = tuple(worker._region_instance_registry._instances.values()) + assert len(unpublished) == 1 + assert unpublished[0]._state is comm_region.RegionInstanceState.CLOSE_FAILED finally: worker._close_worker_chip_orch_comm() shm.close() shm.unlink() -def test_onboard_direct_mapping_bytes_too_small_rolls_back_l2_host_region(monkeypatch): +def test_onboard_direct_mapping_bytes_cover_each_independent_part(monkeypatch): worker, shm, fake_c_worker = _make_started_onboard_worker() fake_c_worker.mapping_bytes = 191 calls: list[tuple] = [] @@ -383,12 +505,14 @@ def test_onboard_direct_mapping_bytes_too_small_rolls_back_l2_host_region(monkey lambda *args: calls.append(args) or 123, ) - with pytest.raises(RuntimeError, match="onboard_vmm reply mapping_bytes is smaller"): - worker._create_worker_chip_region(0, 64, 128) + region = worker._create_worker_chip_region(0, 64, 128) - assert calls == [] - assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [] + assert calls == [ + (2, 0xABCDEF, 191, worker._owner_id), + (2, 0xABCDEF + 1, 191, worker._owner_id), + ] + assert region._instance._payload_mapping == 123 + assert region._instance._counter_mapping == 123 finally: worker._close_worker_chip_orch_comm() shm.close() @@ -408,69 +532,63 @@ def test_onboard_direct_mapping_allows_granularity_aligned_mapping(monkeypatch): region = worker._create_worker_chip_region(0, 64, 128) - assert calls == [(2, 0xABCDEF, 65536, worker._owner_id)] - assert region._worker_host_mapping is not None - assert region._worker_host_mapping.total_bytes == 192 + assert calls == [ + (2, 0xABCDEF, 65536, worker._owner_id), + (2, 0xABCDEF + 1, 65536, worker._owner_id), + ] + assert region._instance._payload_mapping == 123 + assert region._instance._counter_mapping == 123 finally: worker._close_worker_chip_orch_comm() shm.close() shm.unlink() -@pytest.mark.parametrize("interrupted_publication", ["worker", "run"]) -def test_direct_region_create_rolls_back_partially_published_region(monkeypatch, interrupted_publication): - class _AppendThenInterrupt(list): - def append(self, item) -> None: - super().append(item) - raise KeyboardInterrupt(f"interrupted {interrupted_publication} publication") - +def test_direct_region_create_rolls_back_when_compat_wrapper_fails(monkeypatch): worker, shm, fake_c_worker = _make_started_sim_worker() resources = worker_module._RunResources() worker._building_run_resources = resources close_calls: list[int] = [] monkeypatch.setattr(worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 55) monkeypatch.setattr( - worker_chip_orch_comm, + comm_region, "_worker_host_mapped_region_close", lambda handle: close_calls.append(int(handle)), ) - if interrupted_publication == "worker": - worker._live_worker_chip_regions = _AppendThenInterrupt() - else: - resources.worker_chip_regions = _AppendThenInterrupt() + monkeypatch.setattr( + worker_module, + "worker_chip_orch_region_desc_from_local_views", + lambda *_args, **_kwargs: (_ for _ in ()).throw(KeyboardInterrupt("interrupted wrapper publication")), + ) try: - with pytest.raises(KeyboardInterrupt, match=f"interrupted {interrupted_publication} publication"): + with pytest.raises(KeyboardInterrupt, match="interrupted wrapper publication"): worker._create_worker_chip_region(0, 64, 128) - assert worker._live_worker_chip_regions == [] - assert resources.worker_chip_regions == [] + assert worker._region_instance_registry._instances == {} assert resources.requires_ordered_cleanup is False - assert close_calls == [55] + assert close_calls == [55, 55] assert fake_c_worker.release_calls == [(0, 1)] finally: worker._building_run_resources = None - worker._live_worker_chip_regions.clear() - resources.worker_chip_regions.clear() worker._close_worker_chip_orch_comm() shm.close() shm.unlink() def test_direct_region_create_mapping_rollback_failure_poisons_worker(monkeypatch): - class _AppendThenInterrupt(list): - def append(self, item) -> None: - super().append(item) - raise KeyboardInterrupt("interrupted publication") - worker, shm, fake_c_worker = _make_started_sim_worker() - worker._live_worker_chip_regions = _AppendThenInterrupt() monkeypatch.setattr(worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 55) monkeypatch.setattr( - worker_chip_orch_comm, + comm_region, "_worker_host_mapped_region_close", lambda _handle: (_ for _ in ()).throw(RuntimeError("mapping close failed")), ) + monkeypatch.setattr( + worker_module, + "worker_chip_orch_region_desc_from_local_views", + lambda *_args, **_kwargs: (_ for _ in ()).throw(KeyboardInterrupt("interrupted wrapper publication")), + ) try: with pytest.raises(RuntimeError, match="rollback could not close the L3 Host mapping") as excinfo: @@ -478,11 +596,10 @@ def append(self, item) -> None: assert isinstance(excinfo.value.__cause__, RuntimeError) assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [] + assert worker._region_instance_registry._instances == {} with pytest.raises(RuntimeError, match="no further work is admitted"): worker._require_no_ordered_cleanup_failure("submit") finally: - worker._live_worker_chip_regions.clear() worker._close_worker_chip_orch_comm() shm.close() shm.unlink() @@ -555,7 +672,8 @@ def interrupt_ack(_owner_token: str, _observed: str) -> None: assert caught.value is ack_interrupt assert fake_c_worker.release_calls == [(0, 1)] - assert worker._ordered_cleanup_error is worker._worker_host_mapped_cleanup_error + assert worker._ordered_cleanup_error is not None + assert worker._worker_host_mapped_cleanup_error is not None finally: worker._close_worker_chip_orch_comm() shm.close() @@ -792,44 +910,14 @@ def acknowledge_cleanup_error(owner_token: str, observed: str) -> None: def test_close_replay_does_not_double_release_worker_chip_region_after_mapping_close_failure(monkeypatch): - class _FakeCloseWorker(_FakeDirectCWorker): - def __init__(self) -> None: - super().__init__() - self.close_calls = 0 - - def close(self) -> None: - self.close_calls += 1 - - worker = Worker(level=3, device_ids=[0], platform="a2a3sim", runtime="tensormap_and_ringbuffer") - fake_c_worker = _FakeCloseWorker() - worker._lifecycle = worker_module._Lifecycle.READY - worker._worker = fake_c_worker + worker, shm, fake_c_worker = _make_started_sim_worker() worker._init_owner_thread = threading.current_thread() - mapping = worker_chip_orch_comm.WorkerHostRegionMapping( - worker_id=0, - region_id=1, - access_profile=worker_chip_orch_comm.WorkerChipRegionAccessProfile.SIM_POSIX_SHM, - total_bytes=192, - payload_offset=0, - payload_bytes=64, - counter_offset=64, - counter_bytes=128, - handle=77, - ) - region = worker_chip_orch_comm.WorkerChipOrchRegion( - worker, - 0, - worker_chip_orch_comm.WorkerChipOrchRegionDesc( - magic_version=0x4C334C3200020000, - region_id=1, - payload_base=0xDEAD_0000, - payload_bytes=64, - counter_base=0xDEAD_0040, - counter_bytes=128, - ), - mapping, + handles = iter((77, 78)) + monkeypatch.setattr( + worker_module, + "_worker_host_mapped_region_import_sim", + lambda _token, _size, _owner_token: next(handles), ) - worker._live_worker_chip_regions.append(region) close_calls: list[int] = [] fail_next_mapping_close = True @@ -840,20 +928,26 @@ def close_mapping(handle: int) -> None: fail_next_mapping_close = False raise RuntimeError("mapping close failed") - monkeypatch.setattr(worker_chip_orch_comm, "_worker_host_mapped_region_close", close_mapping) + monkeypatch.setattr(comm_region, "_worker_host_mapped_region_close", close_mapping) - with pytest.raises(RuntimeError, match="mapping close failed"): - worker.close() + try: + region = worker._create_worker_chip_region(0, 64, 128) + with pytest.raises(RuntimeError, match="mapping close failed"): + worker.close() - assert close_calls == [77] - assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [region] + assert close_calls == [77, 78] + assert fake_c_worker.release_calls == [(0, 1)] + assert worker._region_instance_registry._instances == {} + assert region.expired - worker.close() + worker.close() - assert close_calls == [77, 77] - assert fake_c_worker.release_calls == [(0, 1)] - assert worker._live_worker_chip_regions == [] + assert close_calls == [77, 78] + assert fake_c_worker.release_calls == [(0, 1)] + finally: + with contextlib.suppress(FileNotFoundError): + shm.close() + shm.unlink() def test_concurrent_close_publishes_joiner_cleanup_error_to_every_caller(monkeypatch): @@ -1072,50 +1166,28 @@ def test_native_mapping_cleanup_errors_are_keyed_by_owner_token(): assert _task_interface_ext._worker_host_mapped_region_take_cleanup_error(owner_token) == "" -def test_onboard_region_create_handler_uses_named_export_fields(monkeypatch): - req_shm = SharedMemory(create=True, size=worker_chip_orch_comm._REGION_CREATE_REQUEST_BYTES) - reply_shm = SharedMemory(create=True, size=worker_chip_orch_comm._REGION_CREATE_REPLY_BYTES) +def test_region_allocate_handler_commits_store_success_into_reply_shm(): + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _allocation_spec, _sim_context + + req_shm = SharedMemory(create=True, size=ALLOCATE_REQUEST_BYTES) + reply_shm = SharedMemory(create=True, size=ALLOCATE_REPLY_BYTES) req_buf = cast(memoryview, req_shm.buf) reply_buf = cast(memoryview, reply_shm.buf) ctrl_storage = bytearray(worker_module._OFF_ARGS + 2 * worker_module._CTRL_SHM_NAME_BYTES) ctrl_buf = memoryview(ctrl_storage) - cw = _FakeChipWorkerForRegionCreate() - typed_cw = cast(Any, cw) - close_calls: list[int] = [] try: - worker_chip_orch_comm.WorkerChipRegionCreateRequest( - magic_version=0x4C334C3200020000, - request_bytes=worker_chip_orch_comm._REGION_CREATE_REQUEST_BYTES, - payload_bytes=64, - counter_bytes=128, - ).encode_into(req_buf) + encode_allocate_request(req_buf, _allocation_spec()) _write_ctrl_shm_name(ctrl_buf, worker_module._OFF_ARGS, req_shm.name) _write_ctrl_shm_name(ctrl_buf, worker_module._OFF_ARGS + worker_module._CTRL_SHM_NAME_BYTES, reply_shm.name) - monkeypatch.setattr( - worker_module, - "_l3_child_onboard_region_create", - lambda _nbytes: _NamedOnboardRegionExport(), - ) - monkeypatch.setattr( - worker_module, - "_l3_child_onboard_region_close", - lambda handle: close_calls.append(int(handle)), - ) - - store = worker_module._HostWorkerChipRegionStore() - worker_module._handle_ctrl_worker_chip_region_create(typed_cw, ctrl_buf, "a2a3", store) - - reply = worker_chip_orch_comm.decode_region_create_reply(reply_buf) - assert reply.desc.scalars() == [0x4C334C3200020000, 1, 0xD00D_0000, 64, 0xD00D_0040, 128] - assert reply.access_profile == worker_chip_orch_comm.WorkerChipRegionAccessProfile.ONBOARD_VMM - assert reply.device_id == 2 - assert reply.mapping_bytes == 65536 - assert reply.shareable_handle == 0xABCDEF - assert cw.copy_calls == [(0xD00D_0040, 128)] - - region = store.regions.pop(1) - worker_module._release_host_worker_chip_region(region) - assert close_calls == [77] + store = ProviderRegionStore(_sim_context(), _shell_factory=FakeShellFactory()) + worker_module._handle_ctrl_region_allocate(ctrl_buf, store) + tag, result, payload_view, counter_view, error = decode_allocate_reply(reply_buf) + assert tag is AllocateReplyTag.SUCCESS + assert error is None + assert result is not None and result.provider_resource_id == 1 + assert payload_view is not None and counter_view is not None + released = store.release(result.provider_resource_id) + assert released.status is ProviderReleaseStatus.RELEASED finally: ctrl_buf.release() del req_buf @@ -1126,40 +1198,31 @@ def test_onboard_region_create_handler_uses_named_export_fields(monkeypatch): reply_shm.unlink() -def test_region_create_handler_rejects_abi_mismatch(): - req_shm = SharedMemory(create=True, size=worker_chip_orch_comm._REGION_CREATE_REQUEST_BYTES) - reply_shm = SharedMemory(create=True, size=worker_chip_orch_comm._REGION_CREATE_REPLY_BYTES) +def test_region_allocate_handler_writes_request_error_for_bad_magic(): + from tests.ut.py.test_worker.test_comm_provider import FakeShellFactory, _sim_context + + req_shm = SharedMemory(create=True, size=ALLOCATE_REQUEST_BYTES) + reply_shm = SharedMemory(create=True, size=ALLOCATE_REPLY_BYTES) req_buf = cast(memoryview, req_shm.buf) + reply_buf = cast(memoryview, reply_shm.buf) ctrl_storage = bytearray(worker_module._OFF_ARGS + 2 * worker_module._CTRL_SHM_NAME_BYTES) ctrl_buf = memoryview(ctrl_storage) - store = worker_module._HostWorkerChipRegionStore() try: + req_buf[:ALLOCATE_REQUEST_BYTES] = b"\x00" * ALLOCATE_REQUEST_BYTES + struct.pack_into(" None: - events.append(("release", int(region_id))) - original_release(worker_id, region_id) + def release(worker_id: int, request_shm_name: str, reply_shm_name: str) -> None: + req_shm = SharedMemory(name=request_shm_name) + req_buf = req_shm.buf + assert req_buf is not None + try: + events.append(("release", int(decode_release_request(req_buf)))) + finally: + del req_buf + req_shm.close() + original_release(worker_id, request_shm_name, reply_shm_name) try: - fake_c_worker.control_worker_chip_region_release = release + fake_c_worker.control_region_release = release monkeypatch.setattr( worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 77 ) monkeypatch.setattr( - worker_chip_orch_comm, + comm_region, "_worker_host_mapped_region_close", lambda handle: events.append(("close", int(handle))), ) region = worker._create_worker_chip_region(0, 64, 128) region.free() - worker._cleanup_worker_chip_regions() + worker._sweep_region_instances() - assert events == [("close", 77), ("release", 1)] + assert events == [("close", 77), ("close", 77), ("release", 1)] + with pytest.raises(RuntimeError, match="expired"): + region.descriptor_scalars() + assert region._instance._close_attempted + finally: + worker._close_worker_chip_orch_comm() + shm.close() + shm.unlink() + + +def test_compat_wrapper_expires_after_run_cleanup(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + resources = worker_module._RunResources() + worker._building_run_resources = resources + try: + monkeypatch.setattr( + worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 55 + ) + region = worker._create_worker_chip_region(0, 64, 128) + assert region.expired is False + worker._building_run_resources = None + worker._region_instance_registry.cleanup_run(resources) + assert region.expired is True with pytest.raises(RuntimeError, match="expired"): region.descriptor_scalars() + assert fake_c_worker.release_calls == [(0, 1)] + assert worker._region_instance_registry._instances == {} + finally: + worker._building_run_resources = None + worker._close_worker_chip_orch_comm() + shm.close() + shm.unlink() + + +def test_compat_free_is_logical_and_keeps_region_instance_live(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + try: + monkeypatch.setattr( + worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 55 + ) + region = worker._create_worker_chip_region(0, 64, 128) + instance = region._instance + region.free() + region.free() + + with pytest.raises(RuntimeError, match="has been released"): + region.descriptor_scalars() + assert instance._state is comm_region.RegionInstanceState.LIVE + assert fake_c_worker.release_calls == [] + assert worker._region_instance_registry._instances[id(instance)] is instance + finally: + worker._close_worker_chip_orch_comm() + shm.close() + shm.unlink() + + +def test_compat_descriptor_uses_independent_local_views_and_bumped_magic(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + fake_c_worker.payload_base = 0x1000 + fake_c_worker.counter_base = 0x9000 + try: + monkeypatch.setattr( + worker_module, "_worker_host_mapped_region_import_sim", lambda _token, _size, _owner_token: 55 + ) + region = worker._create_worker_chip_region(0, 64, 128) + scalars = region.descriptor_scalars() + + assert len(scalars) == 6 + assert scalars[0] == worker_chip_orch_comm._REGION_MAGIC_VERSION + assert scalars[0] != 0x4C334C3200020000 + assert scalars[1] == 1 + assert scalars[2] == 0x1000 + assert scalars[3] == 64 + assert scalars[4] == 0x9000 + assert scalars[5] == 128 + assert region._instance is not None + assert region._instance._payload_part is not region._instance._counter_part + finally: + worker._close_worker_chip_orch_comm() + shm.close() + shm.unlink() + + +def test_public_create_worker_chip_region_uses_admitted_w2_plan_and_two_imports(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + imports: list[tuple[str, int]] = [] + captured: dict[str, Any] = {} + original_ctx = worker._admitted_worker_chip_region_context + + def spy_ctx(worker_id, payload_bytes, counter_bytes): + ctx = original_ctx(worker_id, payload_bytes, counter_bytes) + captured["ctx"] = ctx + return ctx + + monkeypatch.setattr(worker, "_admitted_worker_chip_region_context", spy_ctx) + monkeypatch.setattr( + worker_module, + "_worker_host_mapped_region_import_sim", + lambda token, mapping_bytes, owner_token: imports.append((str(token), int(mapping_bytes))) + or (100 + len(imports)), + ) + try: + orch = Orchestrator(MagicMock(), worker) + region = orch.create_worker_chip_region(worker_id=0, payload_bytes=64, counter_bytes=128) + ctx = captured["ctx"] + plan = ctx.plan + assert isinstance(plan, ce.BackendPlan) + assert isinstance(plan.topology_plan, ce.SingleOwnerPlan) + registry = worker._get_endpoint_registry() + provider = registry.record_for(plan.topology_plan.provider_endpoint) + consumer = next(registry.record_for(member) for member in plan.ordered_members if member != provider.identity) + assert provider.path == "L3/L2[0]" + assert provider.deployment is ce.DEVICE_AICPU + assert consumer.path == "L3" + assert consumer.deployment is ce.HOST_CPU + assert set(plan.ordered_members) == {provider.identity, consumer.identity} + assert [item[0] for item in imports] == ["sim-direct-1-p", "sim-direct-1-c"] + assert [item[1] for item in imports] == [64, 128] + assert region._instance._payload_part is not region._instance._counter_part + scalars = region.descriptor_scalars() + assert len(scalars) == 6 + assert scalars[0] == worker_chip_orch_comm._REGION_MAGIC_VERSION + assert scalars[0] != 0x4C334C3200020000 + region.free() + assert fake_c_worker.release_calls == [] + assert region._instance._state is comm_region.RegionInstanceState.LIVE + worker._sweep_region_instances() + assert fake_c_worker.release_calls == [(0, 1)] + assert worker._region_instance_registry._instances == {} finally: worker._close_worker_chip_orch_comm() shm.close() @@ -1699,3 +1897,77 @@ def orch(orch_handle, _args, _cfg): worker.run(orch) finally: worker.close() + + +@pytest.mark.parametrize("platform", ["a2a3sim", "a5sim"]) +def test_sim_public_api_two_shm_objects_and_two_consecutive_lifecycles(platform, monkeypatch): + try: + RuntimeBuilder(platform=platform).get_binaries("tensormap_and_ringbuffer") + except FileNotFoundError as e: + pytest.skip(f"{platform} runtime binaries unavailable: {e}") + + imports: list[tuple[str, int]] = [] + original_import = worker_module._worker_host_mapped_region_import_sim + + def spy_import(token, mapping_bytes, owner_token): + imports.append((str(token), int(mapping_bytes))) + return original_import(token, mapping_bytes, owner_token) + + monkeypatch.setattr(worker_module, "_worker_host_mapped_region_import_sim", spy_import) + + worker = Worker( + level=3, + device_ids=[0], + platform=platform, + runtime="tensormap_and_ringbuffer", + num_sub_workers=0, + ) + worker.init() + leftover_names: list[str] = [] + try: + + def orch(orch_handle, _args, _cfg): + host = orch_handle.alloc([16], DataType.UINT8) + buf_t = ctypes.c_uint8 * 16 + buf = buf_t.from_address(int(host.base)) + for i in range(16): + buf[i] = (i + 7) & 0xFF + region = orch_handle.create_worker_chip_region(worker_id=0, payload_bytes=16, counter_bytes=128) + region.payload_write(0, host) + for i in range(16): + buf[i] = 0 + region.payload_read(0, host) + assert bytes(buf) == bytes((i + 7) & 0xFF for i in range(16)) + region.counter(0).notify(3, NotifyOp.Set) + assert region.counter(0).test(3, WaitCmp.EQ).matched + region.free() + + worker.run(orch) + first_pair = imports[:2] + leftover_names.extend(name for name, _size in first_pair) + assert len(first_pair) == 2 + assert first_pair[0][0] != first_pair[1][0] + assert {first_pair[0][1], first_pair[1][1]} == {16, 128} + assert worker._region_instance_registry._instances == {} + assert worker._region_instance_registry._instances == {} + for name in leftover_names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + worker.run(orch) + second_pair = imports[2:] + leftover_names.extend(name for name, _size in second_pair) + assert len(second_pair) == 2 + assert second_pair[0][0] != second_pair[1][0] + assert {name for name, _size in first_pair}.isdisjoint({name for name, _size in second_pair}) + assert worker._region_instance_registry._instances == {} + assert worker._region_instance_registry._instances == {} + for name in leftover_names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + worker.close() + assert worker._region_instance_registry._instances == {} + for name in leftover_names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name)