Unload vllm model

rl/llm/engines/local.py

@@ -1,4 +1,5 @@
 import datetime
+import gc
 import math
 import os
 import socket
@@ -377,7 +378,19 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        import torch
+        import torch.distributed
+        from vllm.model_executor.parallel_utils.parallel_state import (
+            destroy_model_parallel,
+        )
+
+        LOGGER.info("Unloading VLLM model from GPU memory...")
+        destroy_model_parallel()
         del self.vllm
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.distributed.destroy_process_group()
+        LOGGER.info("VLLM model unloaded.")
 
     def generate(self, prompt: InferenceInput) -> InferenceOutput:
         return self.batch_generate([prompt])[0]

rl/utils/io.py

@@ -161,6 +161,12 @@ def write_csv(
 
 
 def download(url: str, dest: str | Path) -> None:
+    """Download a file from a URL to a destination path, with a progress bar.
+
+    Args:
+        url: The URL to download from.
+        dest: The destination path.
+    """
     dest = Path(dest)
     response = requests.get(url, stream=True)
     total_size = int(response.headers.get("content-length", 0))