From a73a6a8de082c58667ab732d12762b8c5c9f1dbd Mon Sep 17 00:00:00 2001
From: jax authors <no-reply@google.com>
Date: Tue, 23 Aug 2022 10:37:34 -0700
Subject: [PATCH] Shut down preemption sync manager if enabled.

PiperOrigin-RevId: 469497215
---
 jax/_src/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/jax/_src/distributed.py b/jax/_src/distributed.py
index 8d6d04216abb..a11c8088e75f 100644
--- 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: