Shut down preemption sync manager if enabled.

PiperOrigin-RevId: 469497215
khcs · Aug 23, 2022 · a73a6a8 · a73a6a8
1 parent e8ec454
commit a73a6a8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax/_src/distributed.py b/jax/_src/distributed.py
@@ -87,6 +87,8 @@ def shutdown(self):
     if self.service:
       self.service.shutdown()
       self.service = None
+    if self.preemption_sync_manager:
+      self.preemption_sync_manager = None
 
   def initialize_preemption_sync_manager(self):
     if self.preemption_sync_manager is not None: