fix: avoid deepcopy crash with non-pickleables in Gemini/Vertex
This commit is contained in:
parent
48d3aad68f
commit
384ad7e99c
@ -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
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user