* studio: scope cancel-cleanup to in-flight tmp dirs; walk back tool_call_id Two follow-ups to #5375's training and chat hardening. _cleanup_cancelled_checkpoints used to rmtree every checkpoint-N directory on Cancel. That is the opposite of what the user expects. A user cancelling an 8h run with save_steps=2000 loses every completed checkpoint they could have resumed from. The 67 MB residue the audit memo flagged is the HF Trainer atomic-rename partial (tmp-checkpoint-N), not the completed ones. The cleanup now targets only tmp-checkpoint subdirs; completed checkpoint-N directories are user-owned and stay. Symlinked output_dir and symlinked children are skipped so the realpath containment cannot be levered into deleting arbitrary content via a symlink trick. ChatMessage._validate_role_shape stamped a random secrets.token_hex id on tool messages with no tool_call_id. That id is uncorrelated with the prior assistant tool_calls id, so strict passthrough backends (OpenAI, Anthropic) reject the request as orphaned and llama.cpp treats the tool result as "no preceding call" and hallucinates. The synthesis moves up to ChatCompletionRequest, where the whole conversation is visible: for each tool message missing an id we walk back to the most recent assistant turn with tool_calls (stopping at user turns), prefer a function.name match, otherwise take the first unconsumed tool_call. Synthesis is the fallback when no candidate assistant turn exists, preserving the prior round-trip guarantee for orphaned tool messages. Tests: - test_cleanup_cancelled_checkpoints.py (new): pins that completed checkpoint subdirs survive, tmp-checkpoint partials are removed, non-int suffixes (checkpoint-final, checkpoint-best) are left alone, output_dir outside outputs_root is refused, symlinked output_dir and symlinked child are both skipped, missing dir is a no-op. - test_inference_model_validation.py: 6 new walkback cases covering name-match preference, first-unconsumed fallback, explicit-id passthrough, multi-tool-result pairing, synth-on-no-parent, and no-cross-user-turn invariant. - test_openai_tool_passthrough.py: the two ChatMessage-level synth-on-missing tests are rewritten to assert that the per- message validator now leaves tool_call_id untouched; resolution coverage lives in the request-level tests above. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: explicit tool_call_id reserve, numeric tmp-checkpoint suffix only Reviewer follow-ups to the training-cleanup + tool_call_id walkback PR. tool_call_id walkback: a mixed assistant turn with [call_a, call_b] followed by a tool result that carried tool_call_id="call_a" and a sibling tool result with no id resolved to ['call_a', 'call_a'] because the explicit id never reserved call_a in the consumed set. Added a pre-pass over the message list that walks back from every role="tool" message carrying an explicit id and marks the matching (asst_idx, tc_idx) consumed, then the missing-id walkback runs against that pre-populated set. The second result now resolves to call_b. While here, also harden the function-shape check: if a provider ships a malformed tool_call where `function` is a string rather than a dict, the old `(tc.get("function") or {}).get("name")` raised AttributeError on the string's .get; now isinstance-gated so the walkback falls through to the fallback id without raising. Cancel cleanup: `tmp-checkpoint-*` is too broad. HF Trainer's in-flight partials are always `tmp-checkpoint-<integer-step>`, so constrain the cleanup regex to `^tmp-checkpoint-\d+$`. A user folder named `tmp-checkpoint-final`, `tmp-checkpoint-backup`, or `tmp-checkpoint-user-notes` is now preserved. ChatMessage docstring still pointed at the pre-PR contract that required `tool_call_id` on every role="tool" message. Updated to say missing ids are accepted at message scope and resolved at ChatCompletionRequest scope. Inline comment above the cancel-cleanup call now describes the actual behaviour (in-flight tmp partials, completed checkpoints preserved). Test: - python -m pytest studio/backend/tests/test_inference_model_validation.py studio/backend/tests/test_cleanup_cancelled_checkpoints.py studio/backend/tests/test_openai_tool_passthrough.py -q -> 76 passed (was 67 before this commit; +2 walkback regression tests, +1 numeric-suffix preservation test) * studio: trim verbose comments in cleanup + tool_call_id walkback Move the HF tmp-checkpoint regex to module scope as a named constant. Drop the multi-paragraph docstring on _cleanup_cancelled_checkpoints and the inline call-site rationale; the function name + the test class already cover the why. Compress _resolve_missing_tool_call_ids docstring from a six-line explanation to two. Same logic, fewer in-flow tutorials. 76 tests in cleanup + inference-model-validation + tool-passthrough pass. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
226 lines
6.8 KiB
Python
226 lines
6.8 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
import os
|
|
import sys
|
|
|
|
_backend = os.path.join(os.path.dirname(__file__), "..")
|
|
sys.path.insert(0, _backend)
|
|
|
|
from models.inference import LoadRequest
|
|
|
|
|
|
def _base_load_request(**overrides):
|
|
data = {
|
|
"model_path": "unsloth/test-model-GGUF",
|
|
"hf_token": None,
|
|
"max_seq_length": 4096,
|
|
"load_in_4bit": True,
|
|
"is_lora": False,
|
|
"gguf_variant": "Q4_K_M",
|
|
}
|
|
data.update(overrides)
|
|
return LoadRequest.model_validate(data)
|
|
|
|
|
|
def test_blank_chat_template_override_normalizes_to_none():
|
|
req = _base_load_request(chat_template_override = " \n\t")
|
|
|
|
assert req.chat_template_override is None
|
|
|
|
|
|
def test_nonblank_chat_template_override_is_preserved_verbatim():
|
|
template = " {{ messages }} "
|
|
req = _base_load_request(chat_template_override = template)
|
|
|
|
assert req.chat_template_override == template
|
|
|
|
|
|
# ---------- ChatCompletionRequest tool_call_id walkback ----------
|
|
|
|
from models.inference import ChatCompletionRequest
|
|
|
|
|
|
def _req(messages, **overrides):
|
|
payload = {"model": "x", "messages": messages, **overrides}
|
|
return ChatCompletionRequest.model_validate(payload)
|
|
|
|
|
|
def test_tool_message_inherits_id_from_prior_assistant_tool_call():
|
|
req = _req(
|
|
[
|
|
{"role": "user", "content": "what is 2+2"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_real123",
|
|
"type": "function",
|
|
"function": {"name": "calc", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "name": "calc", "content": "4"}, # no tool_call_id
|
|
]
|
|
)
|
|
assert req.messages[-1].tool_call_id == "call_real123"
|
|
|
|
|
|
def test_tool_message_with_explicit_id_unchanged():
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_a",
|
|
"type": "function",
|
|
"function": {"name": "search", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_user_supplied", "content": "ok"},
|
|
]
|
|
)
|
|
assert req.messages[-1].tool_call_id == "call_user_supplied"
|
|
|
|
|
|
def test_walkback_prefers_function_name_match():
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_x",
|
|
"type": "function",
|
|
"function": {"name": "search", "arguments": "{}"},
|
|
},
|
|
{
|
|
"id": "call_y",
|
|
"type": "function",
|
|
"function": {"name": "calc", "arguments": "{}"},
|
|
},
|
|
],
|
|
},
|
|
{"role": "tool", "name": "calc", "content": "4"},
|
|
]
|
|
)
|
|
assert req.messages[-1].tool_call_id == "call_y"
|
|
|
|
|
|
def test_walkback_takes_first_unconsumed_when_no_name():
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_a",
|
|
"type": "function",
|
|
"function": {"name": "calc", "arguments": "{}"},
|
|
},
|
|
{
|
|
"id": "call_b",
|
|
"type": "function",
|
|
"function": {"name": "search", "arguments": "{}"},
|
|
},
|
|
],
|
|
},
|
|
{"role": "tool", "content": "first result"},
|
|
{"role": "tool", "content": "second result"},
|
|
]
|
|
)
|
|
assert req.messages[-2].tool_call_id == "call_a"
|
|
assert req.messages[-1].tool_call_id == "call_b"
|
|
|
|
|
|
def test_walkback_falls_back_to_synth_when_no_assistant_turn():
|
|
req = _req(
|
|
[
|
|
{"role": "user", "content": "hi"},
|
|
{"role": "tool", "content": "orphan"},
|
|
]
|
|
)
|
|
tcid = req.messages[-1].tool_call_id
|
|
assert tcid is not None and tcid.startswith("call_") and len(tcid) > 5
|
|
|
|
|
|
def test_walkback_does_not_cross_user_turn():
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "old_call",
|
|
"type": "function",
|
|
"function": {"name": "calc", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "old_call", "content": "4"},
|
|
{"role": "user", "content": "next turn"},
|
|
{"role": "tool", "content": "no parent in this turn"},
|
|
]
|
|
)
|
|
last = req.messages[-1].tool_call_id
|
|
# The walkback must NOT pick old_call because a user turn intervenes;
|
|
# falls back to synth.
|
|
assert last is not None
|
|
assert last != "old_call"
|
|
assert last.startswith("call_")
|
|
|
|
|
|
def test_walkback_skips_explicitly_consumed_tool_call_id():
|
|
"""Sibling tool result with an explicit id must reserve its assistant
|
|
slot so a follow-up missing-id result picks the OTHER tool call."""
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_a",
|
|
"type": "function",
|
|
"function": {"name": "calc", "arguments": "{}"},
|
|
},
|
|
{
|
|
"id": "call_b",
|
|
"type": "function",
|
|
"function": {"name": "search", "arguments": "{}"},
|
|
},
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_a", "content": "4"},
|
|
{"role": "tool", "content": "second result"},
|
|
]
|
|
)
|
|
assert [m.tool_call_id for m in req.messages if m.role == "tool"] == [
|
|
"call_a",
|
|
"call_b",
|
|
]
|
|
|
|
|
|
def test_walkback_handles_malformed_function_string():
|
|
"""A tool_call with ``function`` as a string (provider quirk) must not
|
|
raise; resolution falls back to fallback id selection."""
|
|
req = _req(
|
|
[
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{"id": "call_a", "type": "function", "function": "calc"},
|
|
],
|
|
},
|
|
{"role": "tool", "name": "calc", "content": "4"},
|
|
]
|
|
)
|
|
assert req.messages[-1].tool_call_id == "call_a"
|