diff --git a/vllm/worker/model_runner_base.py b/vllm/worker/model_runner_base.py index 975b88c0e79a..86883cf15244 100644 --- 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)