Cleanup resources at the end of each job (#198)

openreview · Dec 18, 2024 · 643a230 · 643a230
1 parent d004554
commit 643a230
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/expertise/service/expertise.py b/expertise/service/expertise.py
@@ -3,6 +3,8 @@
 import time
 import os
 import json
+import torch
+import gc
 from csv import reader
 import openreview
 from openreview import OpenReviewException
@@ -219,6 +221,10 @@ async def worker_process(self, job, token):
             self.update_status(config, JobStatus.RUN_EXPERTISE)
             execute_expertise(config=config.to_json())
             self.update_status(config, JobStatus.COMPLETED)
+
+            # Explicitly cleanup resources
+            torch.cuda.empty_cache()
+            gc.collect()
         except Exception as e:
             self.update_status(config, JobStatus.ERROR, str(e))
             # Re raise exception so that it appears in the queue