Refine statemachine

Author: HubertZhang

checkpoint_engine/worker.py

@@ -70,29 +70,35 @@ def update_weights_from_ipc(
         socket.send_string(msg)
         socket.recv()  # wait for ack
         raise
+    # State machine:
+    # + receive tensor_metadata -> update_weights
+    # + receive Exception -> raise and stop
+    # + receive None first time -> release resources
+    # + receive None second time -> call post_hook and stop
     try:
         released = False
         while True:
             payload: list[FlattenedTensorMetadata] | Exception | None = socket.recv_pyobj()
-            if payload is None:  # done signal
-                # TODO: refine stm logic, wrap all messages instead of None and Exception
-                device_manager.device_module.synchronize()
-                if not released:
-                    released = True
-                    buffer = None
-                    del ipc_handle
-
-                    gc.collect()
-                    device_manager.device_module.ipc_collect()
-                    device_manager.device_module.empty_cache()
-                    device_manager.device_module.synchronize()
-                    socket.send(b"")
-                    continue
+            if released:
+                assert payload is None, "Should not receive any payload after released"
                 if post_hook is not None:
                     post_hook()
                 device_manager.device_module.synchronize()
                 socket.send(b"")
                 break
+            if payload is None:  # done signal
+                # TODO: wrap all messages to an object instead of None and Exception
+                device_manager.device_module.synchronize()
+                released = True
+                buffer = None
+                del ipc_handle
+
+                gc.collect()
+                device_manager.device_module.ipc_collect()
+                device_manager.device_module.empty_cache()
+                device_manager.device_module.synchronize()
+                socket.send(b"")
+                continue
             if isinstance(payload, list):  # still updating weights
                 try:
                     run(_extract_weights(payload, buffer))