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

Author: iabhi4

litellm/litellm_core_utils/core_helpers.py

@@ -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

litellm/main.py

@@ -116,6 +116,7 @@
 
 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,

tests/test_litellm/litellm_core_utils/test_core_helpers.py

@@ -9,7 +9,11 @@
     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