Skip to content

Fix: improved unsafe pickle strategy#1545

Merged
jmartin-tech merged 9 commits intoNVIDIA:mainfrom
jmartin-tech:fix/improved_unsafe_pickle_strategy
Jan 14, 2026
Merged

Fix: improved unsafe pickle strategy#1545
jmartin-tech merged 9 commits intoNVIDIA:mainfrom
jmartin-tech:fix/improved_unsafe_pickle_strategy

Conversation

@jmartin-tech
Copy link
Collaborator

Inspired by #1543, and address the required review comments in that PR.

Fixes #1471

This revises custom pickle code for all class implementing Configurable to suppress volatile attributes that cannot be pickled without changes to the original instance.

This pattern use the following aspects:

  • _unsafe_attributes is a List[str] of attributes defined in the class as needing to be suppressed
  • serialization
    • copies self.dict in the original instance
    • clears the unsafe entries
    • clears extra_dependency_names entries
  • deserialization
    • calls _load_deps to restore any modules references
    • calls _load_unsafe to restore any registered volatile attributes

The above patten avoids modification of the original instance ensuring multi-processing will not impact state, and removes the need for _clear_client by simplifying the registration of volatile attributes.

Testing showed that promptinject probes were holding a local module reference that could also impact pickle support and has a been refactored to allow all plugins to pass testing.

Verification

List the steps needed to make sure this thing works

  • new unit test for pickle support
  • Execute parallel_attempt and parallel_request runs using a compatible generator
  • Execute parallel_attempt and parallel_request runs using a compatible generator and translation using Google Cloud Translate

motlaharsh0909-lgtm and others added 5 commits December 26, 2025 17:49
GoogleTranslator stores a live Google Cloud translation client on the
instance, which cannot be pickled when parallel_attempts > 1 and causes
PicklingError during multiprocessing startup.

This change prevents the client from being pickled and allows it to be
recreated per worker process, following the same pattern used by
RivaTranslator.# On branch fix-googletranslator-pickling

Signed-off-by: Harsh Motla <[email protected]>
This better reflects the expected outcome of the method

Signed-off-by: Jeffrey Martin <[email protected]>
@jmartin-tech jmartin-tech requested a review from leondz January 9, 2026 14:46
Moves fileformats.FileIsExecutable to the new `_load_deps` deferred loading pattern

Signed-off-by: Jeffrey Martin <[email protected]>
Signed-off-by: Jeffrey Martin <[email protected]>
All "plugins" are of class `Configurable`, since serialization is also
coupled to the naming conventions defined in `_load_deps()`, consolidate
the method in the common class.

Signed-off-by: Jeffrey Martin <[email protected]>
Signed-off-by: Jeffrey Martin <[email protected]>
@jmartin-tech jmartin-tech merged commit 6442217 into NVIDIA:main Jan 14, 2026
15 checks passed
@github-actions github-actions bot locked and limited conversation to collaborators Jan 14, 2026
@jmartin-tech jmartin-tech deleted the fix/improved_unsafe_pickle_strategy branch January 19, 2026 22:46
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Race condition accessing generators during multiprocessing

2 participants