fix: force ps to quit when error occur during updating

[specture724]

resolve #38

send an Exception to PS when update_weights_from_ipc execute run failed. _update_per_bucket needs to ensure all processes not to get an Exception before continuing next updating iteration

[Copilot]

Pull Request Overview

This PR implements error handling and recovery mechanisms for distributed weight updates in the checkpoint engine. The changes add support for detecting and handling remote errors during weight updates, ensuring proper cleanup even when errors occur.

Key changes:

• Added error handling in worker processes to send exceptions back through ZMQ sockets
• Modified the parameter server to detect remote errors via non-empty responses and coordinate error recovery using all-reduce
• Moved process group cleanup and cache clearing into a finally block to ensure proper resource cleanup on both success and failure paths

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
checkpoint_engine/worker.py	Wrapped weight update execution in try-except to send exceptions back via ZMQ socket
checkpoint_engine/ps.py	Added error detection via response checking, coordinated error handling with all-reduce, and moved cleanup to finally block
tests/test_update.py	Added comprehensive test infrastructure for error scenarios with new test cases and subprocess-based execution

Comments suppressed due to low confidence (2)

tests/test_update.py:62

• Variable weights is not used.

        weights = weights  # Do some fake processing

tests/test_update.py:62

• This assignment assigns a variable to itself.

        weights = weights  # Do some fake processing

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[weixiao-huang]

LGTM

[SongXiaoXi]

LGTM

[blahgeek]

lgtm