Fix VLLM subclass instance attribute access inside trace context

[Butanium]

Generated with claude code to fix #583, unsure how good this is as it uses pickle to serialize some of the objects

Summary

• Fix VLLM subclass instance attribute access inside trace context
• Preserve user-defined attributes during Envoy serialization/deserialization
• Add test coverage for subclass attribute access patterns

Problem

When subclassing VLLM, instance attributes (including self.* access inside properties) work outside the trace context but fail inside it:

class MyVLLM(VLLM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._my_attr = "hello"
    
    @property
    def my_attr(self):
        return self._my_attr  # accesses self.*

model = MyVLLM("gpt2", dispatch=True)

# Works outside trace
print(model.my_attr)  # "hello"

with model.trace("Hello world"):
    print(model.my_attr)  # FAILS: AttributeError: GPT2LMHeadModel has no attribute _my_attr

Root Cause

When VLLM traces code, the Mediator is serialized and sent to the VLLM worker. The original Envoy.__getstate__ only preserved internal attributes (_alias, _children, etc.), losing user-defined attributes like _my_attr. On deserialization, these attributes were missing.

Proposed Solution

Modified Envoy.__getstate__ and __setstate__ to:

1. Identify user-defined attributes (those not in _INTERNAL_ATTRS)
2. Check if each attribute is pickleable (skip unpickleable ones like vllm_entrypoint)
3. Serialize and restore these user-defined attributes

Also fixed a typo: self.fake_inputs → self._fake_inputs (missing underscore)

Test plan

• New tests in TestSubclassAttributes:
  • test_instance_attr_in_trace - Basic instance attribute access
  • test_property_with_self_access - Properties that access self.*
  • test_init_attrs - Attributes set in __init__
  • test_dynamic_attrs - Dynamically added attributes
• Existing VLLM tests pass
• LanguageModel still works correctly

🤖 Generated with Claude Code

[Butanium]

I'm unsure if this is optimal, as using pickle is not super stable

[Copilot]

Pull request overview

This PR fixes a bug where VLLM subclasses could not access instance attributes inside trace contexts. When tracing, the Mediator is serialized and sent to VLLM workers, but the original Envoy.__getstate__ only preserved internal attributes, losing user-defined attributes during deserialization.

• Modified Envoy.__getstate__ to identify and preserve pickleable user-defined attributes
• Modified Envoy.__setstate__ to restore user-defined attributes with backward compatibility
• Fixed typo: self.fake_inputs → self._fake_inputs (lines 1140-1141)

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/nnsight/intervention/envoy.py	Added _INTERNAL_ATTRS frozenset to distinguish internal from user attributes; modified __getstate__ to serialize pickleable user attributes; modified __setstate__ to restore user attributes and fix typo
tests/test_vllm.py	Added TestSubclassAttributes class with 4 tests covering instance attributes, properties with self access, init attributes, and dynamic attributes in trace contexts

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Testing pickleability with pickle.dumps(value) fully serializes each value, which can be expensive for large objects. This happens during every __getstate__ call (i.e., every time the Envoy is pickled for sending to workers).

Consider using a lighter check, such as checking if the object has a __reduce__ or __reduce_ex__ method, or catching the exception during the actual serialization in the return statement instead of pre-testing. Alternatively, you could check the type against known unpickleable types (like vllm_entrypoint which is likely an LLM instance).

[Copilot]

The serialization.py inject function (lines 42-44) copies attributes from the root envoy that are not in the deserialized envoy's __dict__. With this change, user-defined attributes are now included in the serialized state and will be in _self.__dict__ after original_setstate(_self, state) is called. This means those attributes won't be overwritten by the inject function.

However, if a user attribute exists on the root envoy but has been modified or is different in the serialized state, the inject function will not overwrite it, potentially causing inconsistencies. Consider whether user attributes should be synchronized with the root envoy or if the serialized version should take precedence. The current behavior (serialized version takes precedence) seems reasonable, but this should be documented or tested.

[Copilot]

The exception handling catches AttributeError which can occur if an object has a __getstate__ that raises AttributeError. However, this is an uncommon case. More commonly, objects that can't be pickled will raise TypeError or pickle.PicklingError.

The current exception tuple (pickle.PicklingError, TypeError, AttributeError) is reasonable, but be aware that catching AttributeError broadly could potentially hide legitimate attribute errors during pickling. Consider logging skipped attributes in debug mode to help troubleshoot serialization issues.

[Butanium]

this is indeed ugly

[Butanium]

I think this is probably not good enough and will cause weird pickle error to people so i'm closing this