diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 990197969..4185031c0 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -383,30 +383,33 @@ inline void bind_worker(nb::module_ &m) { .def( "add_next_level_worker", - [](Worker &self, uint64_t mailbox_ptr) { - self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr)); + [](Worker &self, uint64_t mailbox_ptr, int child_pid) { + self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), child_pid); }, - nb::arg("mailbox_ptr"), + nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, "Add a NEXT_LEVEL sub-worker. `mailbox_ptr` is the address of a " "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is " - "Python-managed (fork + _chip_process_loop)." + "Python-managed (fork + _chip_process_loop). `child_pid` is that " + "forked child, used to detect an exit before mailbox completion." ) .def( "add_next_level_worker_at", - [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr) { - self.add_next_level_worker(worker_id, reinterpret_cast(mailbox_ptr)); + [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid) { + self.add_next_level_worker(worker_id, reinterpret_cast(mailbox_ptr), child_pid); }, - nb::arg("worker_id"), nb::arg("mailbox_ptr"), "Add a NEXT_LEVEL sub-worker with an explicit worker id." + nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, + "Add a NEXT_LEVEL sub-worker with an explicit worker id." ) .def( "add_sub_worker", - [](Worker &self, uint64_t mailbox_ptr) { - self.add_worker(WorkerType::SUB, reinterpret_cast(mailbox_ptr)); + [](Worker &self, uint64_t mailbox_ptr, int child_pid) { + self.add_worker(WorkerType::SUB, reinterpret_cast(mailbox_ptr), child_pid); }, - nb::arg("mailbox_ptr"), + nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, "Add a SUB sub-worker. `mailbox_ptr` is the address of a " "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is " - "Python-managed (fork + _sub_worker_loop)." + "Python-managed (fork + _sub_worker_loop). `child_pid` is that " + "forked child, used to detect an exit before mailbox completion." ) .def( "add_remote_l3_socket", diff --git a/python/simpler/worker.py b/python/simpler/worker.py index cfd755011..3be93c2ef 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -925,6 +925,17 @@ def _mailbox_addr(shm: SharedMemory) -> int: return ctypes.addressof(ctypes.c_char.from_buffer(buf)) +def _require_matching_pids(shms: list[SharedMemory], pids: list[int], kind: str) -> None: + """Guard the shm/pid pairing the C++ liveness check depends on. + + Registering a mailbox against the wrong child pid would make the endpoint + watch an unrelated process, so a length mismatch is rejected instead of + being silently truncated by ``zip``. + """ + if len(shms) != len(pids): + raise RuntimeError(f"{kind} worker shm/pid count mismatch: {len(shms)} mailboxes, {len(pids)} pids") + + def _buffer_field_addr(buf, offset: int) -> int: """Absolute address of a field inside a shared-memory buffer. @@ -4276,20 +4287,25 @@ def _setup(inner=inner_worker): dw = self._worker assert dw is not None - # Register chip workers as NEXT_LEVEL (L3) + # Register chip workers as NEXT_LEVEL (L3). The child pid lets the C++ + # endpoint fail a dispatch whose child died instead of spinning on a + # mailbox that can no longer be completed. if device_ids: - for shm in self._chip_shms: - dw.add_next_level_worker(_mailbox_addr(shm)) + _require_matching_pids(self._chip_shms, self._chip_pids, "chip") + for shm, pid in zip(self._chip_shms, self._chip_pids): + dw.add_next_level_worker(_mailbox_addr(shm), pid) # Register Worker children as NEXT_LEVEL (L4+) if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): raise RuntimeError("explicit NEXT_LEVEL worker ids require a rebuilt _task_interface module") + _require_matching_pids(self._next_level_shms, self._next_level_pids, "next_level") for idx, shm in enumerate(self._next_level_shms): worker_id = self._next_level_worker_ids[idx] - dw.add_next_level_worker_at(worker_id, _mailbox_addr(shm)) + dw.add_next_level_worker_at(worker_id, _mailbox_addr(shm), self._next_level_pids[idx]) - for shm in self._sub_shms: - dw.add_sub_worker(_mailbox_addr(shm)) + _require_matching_pids(self._sub_shms, self._sub_pids, "sub") + for shm, pid in zip(self._sub_shms, self._sub_pids): + dw.add_sub_worker(_mailbox_addr(shm), pid) # Start Scheduler + WorkerThreads (C++ threads start here, after fork) dw.init() diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index f5f52b01e..457276583 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -68,15 +68,15 @@ Worker::~Worker() { if (initialized_) close(); } -void Worker::add_worker(WorkerType type, void *mailbox) { +void Worker::add_worker(WorkerType type, void *mailbox, int child_pid) { if (initialized_) throw std::runtime_error("Worker: add_worker after init"); - if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox); - else manager_.add_sub(mailbox); + if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox, child_pid); + else manager_.add_sub(mailbox, child_pid); } -void Worker::add_next_level_worker(int32_t worker_id, void *mailbox) { +void Worker::add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid) { if (initialized_) throw std::runtime_error("Worker: add_next_level_worker after init"); - manager_.add_next_level_at(worker_id, mailbox); + manager_.add_next_level_at(worker_id, mailbox, child_pid); } void Worker::add_remote_l3_socket( diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index 1ced949bd..f8e2e82b0 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -74,8 +74,10 @@ class Worker { // MAILBOX_SIZE-byte MAP_SHARED region; the real worker (a `ChipWorker` // for NEXT_LEVEL, a Python callable for SUB) lives in the forked // child and consumes the mailbox via the Python child loop. - void add_worker(WorkerType type, void *mailbox); - void add_next_level_worker(int32_t worker_id, void *mailbox); + // `child_pid` is the forked child servicing `mailbox`, or -1 when the + // caller owns no waitable child. + void add_worker(WorkerType type, void *mailbox, int child_pid = -1); + void add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid = -1); // Register a REMOTE_L3 endpoint only after its session runner completed // prestart and reported HELLO READY on the command lane. diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index ecef136ae..f17c6b956 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -54,6 +55,26 @@ std::string format_digest(const uint8_t *digest) { return out; } +// Task-dispatch spin iterations between child liveness samples. Paired with +// the loop's 50 us sleep this is a ~10 ms sampling period. +constexpr int kChildLivenessPollInterval = 200; + +// Wall-clock period between child liveness samples on spin loops that do not +// sleep between iterations. +constexpr std::chrono::milliseconds kChildLivenessPollPeriod{10}; + +std::string child_status_message(int child_pid, int status) { + std::string msg = "child process pid=" + std::to_string(child_pid) + " exited before mailbox completion"; + if (WIFEXITED(status)) { + msg += " (exit_status=" + std::to_string(WEXITSTATUS(status)) + ")"; + } else if (WIFSIGNALED(status)) { + msg += " (signal=" + std::to_string(WTERMSIG(status)) + ")"; + } else { + msg += " (status=" + std::to_string(status) + ")"; + } + return msg; +} + } // namespace namespace { @@ -126,12 +147,41 @@ void WorkerEndpoint::control_l3_l2_region_release(uint64_t) { // LocalMailboxEndpoint — mailbox helpers // ============================================================================= -LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox) : - mailbox_(mailbox) { +LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid) : + mailbox_(mailbox), + child_pid_(child_pid) { if (mailbox == nullptr) throw std::invalid_argument("LocalMailboxEndpoint: null mailbox"); caps_.worker_id = worker_id; } +std::string LocalMailboxEndpoint::check_child_death() { + if (child_dead_) return child_death_reason_; + if (child_pid_ <= 0) return {}; + + int status = 0; + pid_t r = 0; + do { + r = waitpid(static_cast(child_pid_), &status, WNOHANG); + } while (r < 0 && errno == EINTR); + + if (r == 0) return {}; + + if (r < 0 && errno != ECHILD) { + // Any other waitpid() failure says nothing about the child, so the + // caller keeps polling rather than tearing down a live worker. + return {}; + } + + child_dead_ = true; + if (r < 0) { + child_death_reason_ = "child process pid=" + std::to_string(child_pid_) + + " is no longer waitable (reaped elsewhere) before mailbox completion"; + } else { + child_death_reason_ = child_status_message(child_pid_, status); + } + return child_death_reason_; +} + MailboxState LocalMailboxEndpoint::read_mailbox_state() const { volatile int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_STATE); int32_t v; @@ -318,8 +368,21 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // Signal child process. write_mailbox_state(MailboxState::TASK_READY); - // Spin-poll until child signals TASK_DONE. + // Spin-poll until child signals TASK_DONE. A child that dies without + // publishing TASK_DONE would otherwise leave this loop spinning forever, + // so its liveness is sampled every kChildLivenessPollInterval iterations + // (~10 ms at the 50 us sleep below), amortizing the waitpid() syscall. + int poll_count = 0; while (read_mailbox_state() != MailboxState::TASK_DONE) { + if (++poll_count >= kChildLivenessPollInterval) { + poll_count = 0; + std::string death = check_child_death(); + if (!death.empty()) { + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run: " + death; + return completion; + } + } std::this_thread::sleep_for(std::chrono::microseconds(50)); } @@ -348,13 +411,13 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // WorkerManager // ============================================================================= -void WorkerManager::add_next_level(void *mailbox) { - add_next_level_at(static_cast(next_level_entries_.size()), mailbox); +void WorkerManager::add_next_level(void *mailbox, int child_pid) { + add_next_level_at(static_cast(next_level_entries_.size()), mailbox, child_pid); } -void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox) { +void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox, int child_pid) { if (worker_id < 0) throw std::invalid_argument("WorkerManager::add_next_level_at: negative worker_id"); - next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox}); + next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox, child_pid}); } void WorkerManager::add_next_level_endpoint(std::unique_ptr endpoint) { @@ -362,7 +425,7 @@ void WorkerManager::add_next_level_endpoint(std::unique_ptr endp next_level_endpoint_entries_.push_back(std::move(endpoint)); } -void WorkerManager::add_sub(void *mailbox) { sub_entries_.push_back(mailbox); } +void WorkerManager::add_sub(void *mailbox, int child_pid) { sub_entries_.push_back(LocalSubEntry{mailbox, child_pid}); } void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { if (ring == nullptr) throw std::invalid_argument("WorkerManager::start: null ring"); @@ -391,16 +454,18 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { auto make_next_level_threads = [&]() { for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); - auto endpoint = std::make_unique(entry.worker_id, entry.mailbox); + auto endpoint = std::make_unique(entry.worker_id, entry.mailbox, entry.child_pid); wt->start(ring, on_complete, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } }; - auto make_sub_threads = [&](const std::vector &entries, + auto make_sub_threads = [&](const std::vector &entries, std::vector> &threads) { for (size_t i = 0; i < entries.size(); ++i) { auto wt = std::make_unique(); - auto endpoint = std::make_unique(static_cast(i), entries[i]); + auto endpoint = std::make_unique( + static_cast(i), entries[i].mailbox, entries[i].child_pid + ); wt->start(ring, on_complete, std::move(endpoint)); threads.push_back(std::move(wt)); } @@ -500,11 +565,24 @@ void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeo std::chrono::steady_clock::now() + std::chrono::duration_cast(std::chrono::duration(timeout_s)); } + auto next_liveness_check = std::chrono::steady_clock::now() + kChildLivenessPollPeriod; while (read_mailbox_state() != MailboxState::CONTROL_DONE) { - if (std::chrono::steady_clock::now() >= deadline) { + auto now = std::chrono::steady_clock::now(); + if (now >= deadline) { mailbox_control_timed_out_ = true; throw std::runtime_error(std::string(op_name) + " timed out waiting for CONTROL_DONE"); } + if (now >= next_liveness_check) { + next_liveness_check = now + kChildLivenessPollPeriod; + std::string death = check_child_death(); + if (!death.empty()) { + // The mailbox is poisoned rather than reset to IDLE: with the + // child gone no later command can complete, so admitting one + // would restore the hang this check exists to break. + mailbox_control_timed_out_ = true; + throw std::runtime_error(std::string(op_name) + ": " + death); + } + } } int32_t err = 0; std::memcpy(&err, mbox() + MAILBOX_OFF_ERROR, sizeof(int32_t)); diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 49567f079..c8a73f639 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -249,7 +249,11 @@ class WorkerEndpoint { class LocalMailboxEndpoint : public WorkerEndpoint { public: - LocalMailboxEndpoint(int32_t worker_id, void *mailbox); + // `child_pid` is the forked child servicing this mailbox, or -1 when the + // caller does not own a waitable child. A valid pid enables liveness + // checks that turn a child that dies before publishing TASK_DONE / + // CONTROL_DONE into a reported failure instead of an endless spin-poll. + LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid = -1); const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override; @@ -302,11 +306,22 @@ class LocalMailboxEndpoint : public WorkerEndpoint { void *mailbox_{nullptr}; std::mutex mailbox_mu_; bool mailbox_control_timed_out_{false}; + int child_pid_{-1}; + // Set once the child has been reaped or observed unwaitable; the exit + // status is only available from the reaping waitpid(), so it is retained + // here to describe every later operation on this dead mailbox. + bool child_dead_{false}; + std::string child_death_reason_; char *mbox() const { return static_cast(mailbox_); } MailboxState read_mailbox_state() const; void write_mailbox_state(MailboxState s); void run_control_command(const char *op_name, double timeout_s = -1.0); + + // Returns a description of the child's death, or an empty string while it + // is still running. Never throws: callers decide whether a dead child is + // reported as a completion outcome or an exception. + std::string check_child_death(); }; // ============================================================================= @@ -459,10 +474,12 @@ class WorkerManager { // Register a worker. `mailbox` is a MAILBOX_SIZE-byte MAP_SHARED // region; the real worker (a `ChipWorker` for NEXT_LEVEL, a Python // callable for SUB) lives in the forked child. - void add_next_level(void *mailbox); - void add_next_level_at(int32_t worker_id, void *mailbox); + // `child_pid` is the forked child servicing `mailbox`; pass -1 when the + // caller owns no waitable child. + void add_next_level(void *mailbox, int child_pid = -1); + void add_next_level_at(int32_t worker_id, void *mailbox, int child_pid = -1); void add_next_level_endpoint(std::unique_ptr endpoint); - void add_sub(void *mailbox); + void add_sub(void *mailbox, int child_pid = -1); void start(Ring *ring, const OnCompleteFn &on_complete); void stop(); @@ -542,9 +559,14 @@ class WorkerManager { struct LocalNextLevelEntry { int32_t worker_id{-1}; void *mailbox{nullptr}; + int child_pid{-1}; + }; + struct LocalSubEntry { + void *mailbox{nullptr}; + int child_pid{-1}; }; std::vector next_level_entries_; - std::vector sub_entries_; + std::vector sub_entries_; std::vector> next_level_endpoint_entries_; std::vector> next_level_threads_; diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index fd0cf6ed7..395d0137c 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -10,12 +10,15 @@ */ #include +#include +#include #include #include #include #include #include +#include #include #include #include @@ -359,6 +362,51 @@ TEST(WorkerManagerTest, StartRejectsDuplicateNextLevelWorkerId) { EXPECT_TRUE(threw); } +// A child that dies without publishing CONTROL_DONE must be reported, not +// waited on forever. The mailbox stays at CONTROL_REQUEST exactly as it would +// if the real `_chip_process_loop` had crashed mid-command. Run in a worker +// thread with a bounded join so a regression fails the test instead of +// hanging the suite. +TEST(WorkerManagerTest, ControlCommandFailsWhenChildExitsBeforeCompletion) { + void *mailbox = + mmap(nullptr, MAILBOX_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, /*fd=*/-1, /*offset=*/0); + ASSERT_NE(mailbox, MAP_FAILED); + std::memset(mailbox, 0, MAILBOX_SIZE); + + pid_t child = fork(); + ASSERT_GE(child, 0); + if (child == 0) { + _exit(3); + } + + LocalMailboxEndpoint endpoint(/*worker_id=*/0, mailbox, static_cast(child)); + + std::promise result; + auto done = result.get_future(); + std::thread caller([&] { + try { + endpoint.control_malloc(64); + result.set_value(""); + } catch (const std::runtime_error &e) { + result.set_value(e.what()); + } + }); + + ASSERT_EQ(done.wait_for(std::chrono::seconds(10)), std::future_status::ready) + << "control_malloc did not observe the dead child; it is spinning on CONTROL_DONE"; + std::string message = done.get(); + caller.join(); + + EXPECT_NE(message.find("child process pid=" + std::to_string(child)), std::string::npos) << message; + EXPECT_NE(message.find("exit_status=3"), std::string::npos) << message; + + // The endpoint is poisoned once the child is gone: a later command reports + // rather than resuming the spin. + EXPECT_THROW(endpoint.control_free(0), std::runtime_error); + + ASSERT_EQ(munmap(mailbox, MAILBOX_SIZE), 0); +} + TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { Ring allocator; allocator.init(/*heap_bytes=*/0);