[MISC] Skip dumping inputs when unpicklable (vllm-project#8744)

EmbeddedLLM · Sep 24, 2024 · b8747e8 · b8747e8
1 parent 3185fb0
commit b8747e8
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/vllm/worker/model_runner_base.py b/vllm/worker/model_runner_base.py
@@ -137,7 +137,15 @@ def _wrapper(*args, **kwargs):
                                                       for t in kv_caches
                                                       if is_tensor(t)]
 
-                    pickle.dump(dumped_inputs, filep)
+                    try:
+                        pickle.dump(dumped_inputs, filep)
+                    except Exception as pickle_err:
+                        logger.warning(
+                            "Failed to pickle inputs of failed execution: %s",
+                            str(pickle_err))
+                        raise type(err)(f"Error in model execution: "
+                                        f"{str(err)}") from err
+
                     logger.info(
                         "Completed writing input of failed execution to %s.",
                         filename)