Solved an issue where lighteval vllm would hang indefinitely in multi node settings

Author: ncassereau

src/lighteval/models/vllm/vllm_model.py

@@ -93,6 +93,7 @@ class VLLMModelConfig:
     )
     pairwise_tokenization: bool = False  # whether to tokenize the context and continuation separately or together.
     generation_parameters: GenerationParameters = None  # sampling parameters to use for generation
+    enforce_eager: bool = False  # whether or not to disable cuda graphs with vllm
 
     subfolder: Optional[str] = None
 
@@ -136,13 +137,19 @@ def tokenizer(self):
         return self._tokenizer
 
     def cleanup(self):
-        destroy_model_parallel()
+        if ray is not None:
+            ray.get(ray.remote(destroy_model_parallel).remote())
+        else:
+            destroy_model_parallel()
         if self.model is not None:
             del self.model.llm_engine.model_executor.driver_worker
         self.model = None
         gc.collect()
         ray.shutdown()
-        destroy_distributed_environment()
+        if ray is not None:
+            ray.get(ray.remote(destroy_distributed_environment).remote())
+        else:
+            destroy_distributed_environment()
         torch.cuda.empty_cache()
 
     @property
@@ -182,6 +189,7 @@ def _create_auto_model(self, config: VLLMModelConfig, env_config: EnvConfig) ->
             "max_model_len": self._max_length,
             "swap_space": 4,
             "seed": 1234,
+            "enforce_eager": config.enforce_eager,
         }
         if int(config.data_parallel_size) > 1:
             self.model_args["worker_use_ray"] = True