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

2024-09-23 23:10:03 -07:00 · 2024-09-23 23:10:03 -07:00 · b8747e8a7c
commit b8747e8a7c
parent 3185fb0cca
1 changed files with 9 additions and 1 deletions
--- 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)