fix(abstractions): double deallocation in inference with continuous b…

…ackends
PKU-Alignment · Dec 21, 2024 · 26b552c · 26b552c
1 parent d2f3147
commit 26b552c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/abstractions/backends.py b/src/abstractions/backends.py
@@ -378,6 +378,7 @@ def vllm_free_gpu_memory():
 
             try:
                 sgl.set_default_backend(sgl.RuntimeEndpoint(f"http://localhost:{port}"))
+                backend_key = None
                 connected = True
                 backend = None
                 print("Connected to backend.", flush=True)
@@ -671,7 +672,9 @@ def sglang_free_gpu_memory():
             with open(f"{root}/output/backend_history.json", "r") as f:
                 backend_history = json.load(f)
 
-            backend_history.pop(backend_key)
+            if backend_key:
+                backend_history.pop(backend_key)
+
             with open(f"{root}/output/backend_history.json", "w") as f:
                 json.dump(backend_history, f)