fix(playground): rewind through the agent, not the global engine - #111
Open
kevin9327 wants to merge 1 commit into
Open
fix(playground): rewind through the agent, not the global engine#111kevin9327 wants to merge 1 commit into
kevin9327 wants to merge 1 commit into
Conversation
Signed-off-by: kevin9327 <kevin9327@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The playground rewinds the conversation by calling the process-global engine directly:
Needle.resetdoes not do that, and for good reason:Two things go wrong when the playground is holding a tuned model, which is exactly what it does
after
/load-modelor after a finetune finishes and callsengine.load_weights(out):FineTuneWorkersubprocess._lib().needle_reset()resets the in-processengine, 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.cactwhose tag makes it generation 3 loads a v3 engine forthe agent, while these two calls load and reset a second, v2 engine that nothing is talking to.
Engine.resetalso drops the agent right after, so on a tuned model each reset abandons a liveworker 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:
completecallsself.agent.reset().resetcallsself.agent.close(), since it discards the agent immediately afterwards and a workersubprocess should be closed rather than left to the garbage collector. The next
completebuilds anew agent, which re-inits the engine, so the conversation is fresh either way.
The
_libimport is no longer needed in either method.Tests
tests/test_playground.py, new. It stubsNeedleand makesneedle._libraise, so reaching theglobal 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 itis about to use.
test_new_tools_build_a_fresh_agent_without_rewinding— a changed toolset builds a new agent anddoes not rewind the old one.
test_reset_closes_the_agent_it_drops—/resetcloses 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 passedby 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.pyandtest_finetune.pyare excluded:they need
flax, whoseorbax-checkpointdependency will not install on Windows because its pathsexceed
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
.cactneeds that sameflaxinstall. The reasoning above is from the code and is pinned by the tests.