LiteLLM integration is broken

[tcapelle]

Script to reproduce:

# pip install weave litellm
import weave
from litellm import completion

weave.init("litellm_bug")

class Model(weave.Model):
    system_prompt: str = "You are a helpful assistant."

    @weave.op()
    def predict(self, prompt: str) -> str:
        return completion(
            messages=[{"role": "system", "content": self.system_prompt}, {"role": "user", "content": prompt}],
            model="gpt-4o")

model = Model()
output = model.predict("Hello, world!")
print(output)

Stack Trace:

Traceback (most recent call last):
  File "/Users/tcapelle/work/evalForge/litellm_bug.py", line 23, in <module>
    output = model.predict("Hello, world!")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/weave/trace/op.py", line 575, in wrapper
    res, _ = _do_call(
             ^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/weave/trace/op.py", line 393, in _do_call
    execute_result = _execute_call(
                     ^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/weave/trace/op.py", line 309, in _execute_call
    handle_exception(e)
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/weave/trace/op.py", line 307, in _execute_call
    res = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/work/evalForge/litellm_bug.py", line 18, in predict
    return completion(
           ^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/litellm/utils.py", line 1013, in wrapper
    raise e
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/litellm/utils.py", line 797, in wrapper
    logging_obj, kwargs = function_setup(
                          ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/litellm/utils.py", line 612, in function_setup
    raise e
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/litellm/utils.py", line 581, in function_setup
    logging_obj = LiteLLMLogging(
                  ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/site-packages/litellm/litellm_core_utils/litellm_logging.py", line 236, in __init__
    self.messages = copy.deepcopy(messages)
                    ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 196, in _deepcopy_list
    append(deepcopy(a, memo))
           ^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 221, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 162, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 259, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 221, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 162, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 259, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 221, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 162, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 259, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 221, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 162, in deepcopy
    y = _reconstruct(x, memo, *rv)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 259, in _reconstruct
    state = deepcopy(state, memo)
            ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 136, in deepcopy
    y = copier(x, memo)
        ^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 221, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
                             ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tcapelle/miniforge3/envs/weave/lib/python3.12/copy.py", line 151, in deepcopy
    rv = reductor(4)
         ^^^^^^^^^^^
TypeError: cannot pickle '_thread.RLock' object

Removing the self.system_prompt usage from weave.Model predict seems to fix it?!

Claude response

I see the issue. The error occurs because LiteLLM is trying to do a deep copy of the messages, and there seems to be a threading lock object getting caught up in the attribute reference. This is likely due to how Weave handles class attributes internally.

[kdaniel21]

I am experiencing the same issue. An even simpler reproduction would be this (without a dependency on LiteLLM):

import weave
import copy
import typing

# weave.init(...)

class Model(weave.Model):
    system_prompt: typing.List[typing.Any] = []

    @weave.op()
    def predict(self) -> str:
        copy.deepcopy(self.system_prompt)
        return
    
model = Model()
model.predict()

This bug also seems to have been introduced in 0.51.14 (possibly #2434?).

Removing the self.system_prompt usage from weave.Model predict seems to fix it?!

Disclaimer: I am not really familiar with the codebase, so feel free to correct me. But as far as I remember and understood, all properties are wrapped in WeaveList, WeaveDict etc. data types so that they can be tracked. I'd think that when executing an operation, this wrapped objects are persisted (as metadata for the model), and this happens in a separate thread, which is where the lock gets added, and something goes wrong.
This would also explain why removing the weave.op decorator solves the problem, as the wrapped object probably would not get persisted, and there would be no lock.

This also seems to be related to this issue (for some strange reason I cannot reference it properly)

I wrote a small helper that seems to get rid of the issue by recursively converting the wrapped object to its "Python equivalent":

def unbox_deep(boxed_value):
    if isinstance(boxed_value, WeaveList):
        return [unbox_deep(element) for element in boxed_value]
    if isinstance(boxed_value, WeaveDict):
        return { key: unbox_deep(boxed_value[key]) for key in boxed_value.keys()}
    
    return unbox(boxed_value)

Then, simply unboxing the value before copying should work just as expected:

copy.deepcopy(unbox_deep(self.system_prompt)))

[andrewtruong]

Should be resolved with this set of PRs:
#2868