feat: client-side embedding training & multi-turn rollout support#247
feat: client-side embedding training & multi-turn rollout support#247Yunnglin wants to merge 1 commit into
Conversation
Expose the core-library embedding training (InfoNCE) and trainable multi-turn rollout capabilities through twinkle_client, and clean up the client_generator technical debt. - Remove client_tools/client_generator.py; convert affected client modules (dataloader/dataset/processor/model/sampler) to hand-maintained sources by stripping their AUTO-GENERATED headers. - Extract bridge-token stitching into the shared pure function twinkle_agentic/rollout/bridge.py::extend_with_bridge; MultiTurnRollout now calls it instead of duplicating the logic. - Add ClientMultiTurnRollout (src/twinkle_client/rollout/) driving HTTP sampling via vLLMSampler, reusing ToolManager and extend_with_bridge. - Enhance MockSampler with opt-in multi-turn knobs (new_input_feature, configurable stop_reason / tool_call_text / tool_call_turns) for local CPU-only multi-turn e2e; backward compatible by default. - Add embedding e2e scaffolding + local mock protocol-link validation and GPU-gated numeric/SP-CP tests; add multi-turn property/unit tests, local mock e2e, and GPU-gated alignment / Gateway-logprobs tests. - Document client embedding call sequence and the Tinker Gateway indirect multi-turn approach with its capability boundaries. Local CPU suite: 98 passed, GPU-gated cases skip cleanly.
There was a problem hiding this comment.
Code Review
This pull request introduces client-side multi-turn agentic rollout orchestration (ClientMultiTurnRollout) and adds documentation and tests for embedding training and multi-turn tool calling. Feedback focuses on ensuring proper JSON serialization of PyTorch tensors and template outputs in the HTTP client path, as well as removing an unused variable in the mock sampler.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if isinstance(obj, np.ndarray): | ||
| return obj.tolist() |
There was a problem hiding this comment.
The _to_plain function recursively converts numpy arrays to lists, but it does not handle PyTorch Tensor objects. In multimodal scenarios, mm_token_type_ids is padded and concatenated as a torch.Tensor (lines 139-147), which remains a Tensor after _to_plain is called. When the client attempts to serialize this to JSON to send over HTTP, it will raise a TypeError. We should update _to_plain to also convert PyTorch tensors to lists if they are present, using a safe string-based type check to avoid importing torch in environments where it is not installed.
| if isinstance(obj, np.ndarray): | |
| return obj.tolist() | |
| if isinstance(obj, np.ndarray): | |
| return obj.tolist() | |
| if type(obj).__name__ == 'Tensor' and hasattr(obj, 'tolist'): | |
| return obj.tolist() |
| from twinkle.data_format import Trajectory | ||
| from twinkle.data_format.sampling import SamplingParams | ||
| from twinkle.template.base import Template | ||
| from twinkle_agentic.rollout.bridge import extend_with_bridge |
There was a problem hiding this comment.
Import _to_plain from twinkle_agentic.rollout.bridge so that we can convert the initial pif returned by self.template.encode into a plain dictionary with serializable types before sending it over HTTP.
| from twinkle_agentic.rollout.bridge import extend_with_bridge | |
| from twinkle_agentic.rollout.bridge import extend_with_bridge, _to_plain |
| pif = self.template.encode(traj, add_generation_prompt=True) | ||
| pif.setdefault('messages', list(traj.get('messages', []))) |
There was a problem hiding this comment.
The initial pif returned by self.template.encode is not run through _to_plain before being used. If the template returns an InputFeature object containing numpy arrays or PyTorch tensors (which is standard for real templates), calling .setdefault might fail if it's not a dict, and sending it over HTTP on line 139 will cause a JSON serialization error. We should convert the initial pif to a plain dictionary using _to_plain just like the core-library MultiTurnRollout does.
| pif = self.template.encode(traj, add_generation_prompt=True) | |
| pif.setdefault('messages', list(traj.get('messages', []))) | |
| pif = self.template.encode(traj, add_generation_prompt=True) | |
| pif = _to_plain(pif) | |
| pif.setdefault('messages', list(traj.get('messages', []))) |
| # Recognised ctor / per-call knobs that tune multi-turn control flow. Kept in a | ||
| # set so the ctor can log only *genuinely* unknown kwargs (a real backend | ||
| # signature drift) while still accepting these opt-in multi-turn knobs. | ||
| _MULTI_TURN_KNOBS = ('stop_reason', 'tool_call_text', 'tool_call_turns') |
There was a problem hiding this comment.
The _MULTI_TURN_KNOBS tuple is defined here but never used anywhere in the codebase. Furthermore, the comment mentions it is 'Kept in a set', but it is defined as a tuple. Since the constructor explicitly defines stop_reason, tool_call_text, and tool_call_turns as keyword-only arguments, Python's native argument binding handles them automatically, making this variable completely redundant. We should remove this unused variable and its comment.
Summary
Exposes the core-library embedding training (InfoNCE) and trainable multi-turn rollout capabilities through
twinkle_client, and cleans up theclient_generatortechnical debt.tinkerclient stays untouched — only documented for its Gateway-based indirect (generation-only) multi-turn path.Technical debt cleanup
client_tools/client_generator.py; convert affected client modules (dataloader/dataset/processor/model/sampler) to hand-maintained sources by stripping theirAUTO-GENERATEDheaders (byte-for-byte otherwise unchanged).Shared bridge function
twinkle_agentic/rollout/bridge.py::extend_with_bridge.MultiTurnRolloutnow calls it instead of maintaining its own copy, so the core-lib and client paths cannot drift.Client multi-turn rollout
ClientMultiTurnRollout(src/twinkle_client/rollout/) driving HTTP sampling viavLLMSampler.sample(), reusingToolManagerand `extend_with_bridTechnical debt cleanup
client_tools/client_generator.py; convert affected client modules (dataloader/dataset/processor/model/sampler) to hand-maintained sources be. Backward compatible by default.Verification
Shared bridge function
twinkle_agentic/rollout/bridge.py::extend_with_bridge.MultiTurnRolloutnow calls it instead of maintaining its own copy, so the core-lib and cU-g- Extract bridge-token stmeClient multi-turn rollout
ClientMultiTurnRollout(src/twinkle_client/rollout/) driving HTTP sampling viavLLMSampler.sample(), reusingToolManagerand `extend_with_bridTechnical debt cleanup
client_totib- AddClientMultiTurnRollou-### Technical debt cleanup
client_tools/client_generator.py; convert affected client modules (dataloader/dataset/processor/model/samplega- Removeclient_tools/clGPVerification
Shared bridge function
'
Notes
tests/server/start_e2e_server.py) and cannot be validated locally.