fix: avoid deepcopy crash with non-pickleables in Gemini/Vertex

This commit is contained in:
iabhi4 2025-09-10 12:59:03 -07:00
parent 48d3aad68f
commit 384ad7e99c
3 changed files with 69 additions and 10 deletions

View File

@ -228,9 +228,11 @@ def safe_deep_copy(data):
"""
Safe Deep Copy
The LiteLLM Request has some object that can-not be pickled / deep copied
Use this function to safely deep copy the LiteLLM Request
The LiteLLM request may contain objects that cannot be pickled/deep-copied
(e.g., tracing spans, locks, clients).
This helper deep-copies each top-level key independently; on failure keeps
original ref
"""
import copy
@ -255,9 +257,22 @@ def safe_deep_copy(data):
"litellm_parent_otel_span"
)
data["litellm_metadata"]["litellm_parent_otel_span"] = "placeholder"
new_data = copy.deepcopy(data)
# Step 2: re-add the litellm_parent_otel_span after doing a deep copy
# Step 2: Per-key deepcopy with fallback
if isinstance(data, dict):
new_data = {}
for k, v in data.items():
try:
new_data[k] = copy.deepcopy(v)
except Exception:
new_data[k] = v
else:
try:
new_data = copy.deepcopy(data)
except Exception:
new_data = data
# Step 3: re-add the litellm_parent_otel_span after doing a deep copy
if isinstance(data, dict) and litellm_parent_otel_span is not None:
if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]:
data["metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span
@ -268,4 +283,4 @@ def safe_deep_copy(data):
data["litellm_metadata"][
"litellm_parent_otel_span"
] = litellm_parent_otel_span
return new_data
return new_data

View File

@ -116,6 +116,7 @@ from litellm.utils import (
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
from .litellm_core_utils.core_helpers import safe_deep_copy
from .litellm_core_utils.fallback_utils import (
async_completion_with_fallbacks,
completion_with_fallbacks,
@ -2772,8 +2773,7 @@ def completion( # type: ignore # noqa: PLR0915
)
api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE")
new_params = deepcopy(optional_params)
new_params = safe_deep_copy(optional_params or {})
response = vertex_chat_completion.completion( # type: ignore
model=model,
messages=messages,
@ -2817,7 +2817,7 @@ def completion( # type: ignore # noqa: PLR0915
api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE")
new_params = deepcopy(optional_params)
new_params = safe_deep_copy(optional_params or {})
if vertex_partner_models_chat_completion.is_vertex_partner_model(model):
model_response = vertex_partner_models_chat_completion.completion(
model=model,

View File

@ -9,7 +9,11 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide
from litellm.litellm_core_utils.core_helpers import (
get_litellm_metadata_from_kwargs,
safe_divide,
safe_deep_copy
)
def test_get_litellm_metadata_from_kwargs():
@ -127,3 +131,43 @@ def test_safe_divide_weight_scenario():
expected_zero = [0, 0, 0]
assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}"
def test_safe_deep_copy_with_non_pickleables_and_span():
"""
Verify safe_deep_copy:
- does not crash when non-pickleables are present,
- preserves structure/keys,
- deep-copies JSON-y payloads (e.g., messages),
- keeps non-pickleables by reference,
- redacts OTEL span in the copy and restores it in the original.
"""
import threading
rlock = threading.RLock()
data = {
"metadata": {"litellm_parent_otel_span": rlock, "x": 1},
"messages": [{"role": "user", "content": "hi"}],
"optional_params": {"handle": rlock},
"ok": True,
}
copied = safe_deep_copy(data)
# Structure preserved
assert set(copied.keys()) == set(data.keys())
# Messages are deep-copied (new object, same content)
assert copied["messages"] is not data["messages"]
assert copied["messages"][0] == data["messages"][0]
# Non-pickleable subtree kept by reference (no crash)
assert copied["optional_params"] is data["optional_params"]
assert copied["optional_params"]["handle"] is rlock
# OTEL span: redacted in the copy, restored in original
assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder"
assert data["metadata"]["litellm_parent_otel_span"] is rlock
# Other simple fields unchanged
assert copied["ok"] is True
assert copied["metadata"]["x"] == 1