Skip to content

fix(playground): rewind through the agent, not the global engine - #111

Open
kevin9327 wants to merge 1 commit into
cactus-compute:mainfrom
kevin9327:fix/playground-reset
Open

fix(playground): rewind through the agent, not the global engine#111
kevin9327 wants to merge 1 commit into
cactus-compute:mainfrom
kevin9327:fix/playground-reset

Conversation

@kevin9327

Copy link
Copy Markdown
Contributor

Problem

The playground rewinds the conversation by calling the process-global engine directly:

def complete(self, tools_json, query):
    from .. import Needle, _lib
    ...
        else:
            _lib().needle_reset()

def reset(self):
    from .. import _lib
    ...
        _lib().needle_reset()

Needle.reset does not do that, and for good reason:

def reset(self):
    self._bind()
    if self._worker is not None:
        self._worker.reset()
    else:
        _lib(self._generation).needle_reset()

Two things go wrong when the playground is holding a tuned model, which is exactly what it does
after /load-model or after a finetune finishes and calls engine.load_weights(out):

  • A tuned agent runs in a FineTuneWorker subprocess. _lib().needle_reset() resets the in-process
    engine, which that agent is not using, so the worker keeps the whole prior conversation. Every
    follow-up query on the same toolset, and every press of Reset, silently carries the earlier turns.
  • _lib() defaults to generation 2. A .cact whose tag makes it generation 3 loads a v3 engine for
    the agent, while these two calls load and reset a second, v2 engine that nothing is talking to.

Engine.reset also drops the agent right after, so on a tuned model each reset abandons a live
worker subprocess and waits for __del__ to reap it.

Fix

Ask the agent, which already knows whether it is a worker and which generation it is:

  • complete calls self.agent.reset().
  • reset calls self.agent.close(), since it discards the agent immediately afterwards and a worker
    subprocess should be closed rather than left to the garbage collector. The next complete builds a
    new agent, which re-inits the engine, so the conversation is fresh either way.

The _lib import is no longer needed in either method.

Tests

tests/test_playground.py, new. It stubs Needle and makes needle._lib raise, so reaching the
global engine is a test failure rather than something you have to notice:

  • test_a_second_query_on_the_same_tools_rewinds_that_agent — the second query rewinds the agent it
    is about to use.
  • test_new_tools_build_a_fresh_agent_without_rewinding — a changed toolset builds a new agent and
    does not rewind the old one.
  • test_reset_closes_the_agent_it_drops/reset closes the agent before dropping it.
  • test_a_tuned_agent_is_rewound_through_itself — the tuned path goes through the agent too.

Three of the four fail before this change with the playground reached the process-global engine.

How I tested

Windows 11, Python 3.12.10. pytest -q -m "not slow" gives 68 passed, 9 skipped, up from 64 passed
by the four new tests; nothing else moves. The 9 skips are the six engine tests and the three worker
tests that want a C compiler.

test_lora.py, test_render.py, test_run.py, test_build.py and test_finetune.py are excluded:
they need flax, whose orbax-checkpoint dependency will not install on Windows because its paths
exceed MAX_PATH. Nothing here touches them, and CI runs them on Ubuntu.

I did not drive the browser playground end to end, because building a tuned .cact needs that same
flax install. The reasoning above is from the code and is pinned by the tests.

Signed-off-by: kevin9327 <kevin9327@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant