Skip to content

Release source-rebinding closure after successful trajectory tokenization #1010

Closed

Description

@bradhilton

Successful trajectory tokenization retains completed source graphs until cyclic GC, even after the caller drops both the input and returned result. This predates #977: it reproduces on 6f0ca471 and maintained main 333b8bdd.

_serialization._rebind_history_sources creates a recursive local visit function. Its closure contains both a reference to visit itself and the canonical exchange list/identity map. Dropping the returned tokenized object therefore leaves a library-created cycle that still owns the exchanges and request data.

The public CPU reproduction distinguishes intentional caller ownership from retention:

  • History-only tokenization releases inputs immediately after caller/result deletion.
  • Trajectory, trace and group tokenization retain the inputs until gc.collect().
  • A finite burst of 32 successful trajectory calls retains all 32 otherwise-unreferenced inputs with GC disabled; collection releases all 32.
  • The live retained visitor closure contains its own function object and the matching canonical exchange. The test harness returns only weak references after the call.

This is deferred reclamation, not evidence of a permanent leak with GC enabled. It can increase memory pressure between collections. It is separate from the sticky-exception retention correction in #977 and is not a qualification gate.

Minimal repository-local reproduction (using an existing public test fixture):

import gc
import sys
import weakref

sys.path.insert(0, 'tests/unit/trajectories')
from test_tokenize import _chat_exchange
from art.trajectories import Trajectory, TrajectoryExchanges

class Options(dict):
    pass

def run():
    value = Options(key=['stable'])
    reference = weakref.ref(value)
    exchange = _chat_exchange([1], [2])
    exchange.request['metadata'] = {'shared': value}
    trajectory = Trajectory(
        exchanges=TrajectoryExchanges(chat_completions=[exchange])
    )
    result = trajectory.tokenize()
    assert result.tokens == [1, 2]
    return reference

gc.collect()
gc.disable()
try:
    reference = run()
    assert reference() is not None  # No caller/result owns it now.
    gc.collect()
    assert reference() is None
finally:
    gc.enable()

Proposed fix: clear the recursive visitor cell in a finally block after source rebinding, preserving traversal, canonical source identity, exception identity and legitimate returned-result ownership. Regression coverage should exercise success, caller-held results, repeated calls, rebinding errors and callbacks.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions