From b8747e8a7c318ab774862f94ccbdbba5b7d9dd4a Mon Sep 17 00:00:00 2001
From: Cody Yu <hao.yu.cody@gmail.com>
Date: Mon, 23 Sep 2024 23:10:03 -0700
Subject: [PATCH] [MISC] Skip dumping inputs when unpicklable (#8744)

---
 vllm/worker/model_runner_base.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/vllm/worker/model_runner_base.py b/vllm/worker/model_runner_base.py
index 975b88c0..86883cf1 100644
--- a/vllm/worker/model_runner_base.py
+++ b/vllm/worker/model_runner_base.py
@@ -137,7 +137,15 @@ def dump_input_when_exception(exclude_args: Optional[List[int]] = None,
                                                       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)