ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

studio: scope cancel-cleanup to in-flight tmp dirs; walk back tool_call_id (#5488)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-05-18 00:01:48 -07:00 • committed by GitHub
commit d79fd92798
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 473 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

29
studio/backend/core/training/training.py
View file

@ -19,6 +19,7 @@ import math
import multiprocessing as mp
import os
import queue
import re
import shutil
import threading
import time
@ -40,11 +41,18 @@ from utils.paths import outputs_root
logger = get_logger(__name__)
_HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$")
def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
"""Remove ``checkpoint-<int>`` subdirs after a cancelled run.
Only paths whose realpath is under outputs_root are touched."""
"""Remove only HF Trainer ``tmp-checkpoint-<step>/`` partials after a cancel.
Completed ``checkpoint-<int>/`` dirs and any non-numeric-suffix tmp dir
are user-owned and survive. Symlinked output_dir / children are skipped
so containment cannot be bypassed.
"""
out = Path(output_dir)
if not out.exists():
if not out.exists() or not out.is_dir() or out.is_symlink():
return
try:
out_real = out.resolve()
@ -54,7 +62,6 @@ def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
try:
out_real.relative_to(out_root_real)
except ValueError:
# Refuse to delete anything outside the configured outputs root.
logger.warning(
"Skipping checkpoint cleanup - %s is not under outputs_root %s",
out_real,
@ -62,14 +69,10 @@ def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
)
return
removed = 0
for entry in out.iterdir() if out.is_dir() else []:
if not entry.is_dir():
for entry in out.iterdir():
if not entry.is_dir() or entry.is_symlink():
continue
name = entry.name
if not name.startswith("checkpoint-"):
continue
tail = name[len("checkpoint-") :]
if not tail.isdigit():
if not _HF_TMP_CHECKPOINT_RE.match(entry.name):
continue
try:
shutil.rmtree(entry, ignore_errors = False)
@ -77,7 +80,7 @@ def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None:
except OSError as exc:
logger.warning("Could not remove %s: %s", entry, exc)
logger.info(
"Cancelled-run cleanup removed %d checkpoint dir(s) under %s",
"Cancelled-run cleanup removed %d in-flight tmp-checkpoint dir(s) under %s",
removed,
out,
)
@ -378,8 +381,6 @@ class TrainingBackend:
if self._pump_thread is not None and self._pump_thread.is_alive():
self._pump_thread.join(timeout = 8.0)
# Drop checkpoint-* dirs on explicit cancel only; stop-and-save
# keeps its artifacts.
if cancelled and output_dir:
try:
_cleanup_cancelled_checkpoints(output_dir)

93
studio/backend/models/inference.py
View file

@ -393,15 +393,12 @@ ContentPart = Annotated[
class ChatMessage(BaseModel):
"""
A single message in the conversation.
"""Single message in a chat conversation.
``content`` may be a plain string (text-only) or a list of
content parts for multimodal messages (OpenAI vision format).
Assistant messages that only contain tool calls may set ``content``
to ``None`` with ``tool_calls`` populated. ``role="tool"`` messages
carry the result of a client-executed tool call and require
``tool_call_id`` per the OpenAI spec.
``content`` is a string or a list of multimodal content parts. Assistant
messages with only ``tool_calls`` populated may set ``content=None``.
Missing ``tool_call_id`` on ``role="tool"`` is resolved at the
``ChatCompletionRequest`` layer by walking back to the preceding assistant.
"""
role: Literal["system", "user", "assistant", "tool"] = Field(
@ -433,17 +430,11 @@ class ChatMessage(BaseModel):
raise ValueError('"name" is only valid on role="tool" messages.')
if self.role == "tool":
if not self.tool_call_id:
# Frontend's second-round POST drops the streamed id;
# synthesise one so the request round-trips.
import secrets as _secrets
self.tool_call_id = f"call_{_secrets.token_hex(8)}"
# tool_call_id resolution happens at ChatCompletionRequest scope.
if not self.content:
raise ValueError('role="tool" messages require non-empty "content".')
elif self.role == "assistant":
# Tolerate the post-Stop empty-assistant sentinel by
# collapsing content="" to None.
# Post-Stop sentinel: collapse content="" / [] to None.
if (self.content == "" or self.content == []) and not self.tool_calls:
self.content = None
else: # "user" | "system"
@ -631,6 +622,76 @@ class ChatCompletionRequest(BaseModel):
),
)
@model_validator(mode = "after")
def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest":
"""Fill missing tool_call_id by walking back to the preceding assistant.
OpenAI / Anthropic passthrough require the result id to match the
assistant's tool_calls[].id. Prefer function.name match, else first
unconsumed tool_call; synth random id only if no candidate exists.
Crossing a user turn breaks the lookup.
"""
# Pre-mark explicit ids first so a sibling missing-id result does not
# steal one already claimed by name.
consumed: set[tuple[int, int]] = set()
def _mark_consumed(start_idx: int, tool_call_id: str) -> None:
for asst_idx in range(start_idx - 1, -1, -1):
prev = self.messages[asst_idx]
if prev.role == "user":
break
if prev.role != "assistant" or not prev.tool_calls:
continue
for tc_idx, tc in enumerate(prev.tool_calls):
if isinstance(tc, dict) and tc.get("id") == tool_call_id:
consumed.add((asst_idx, tc_idx))
return
for tool_idx, msg in enumerate(self.messages):
if msg.role == "tool" and msg.tool_call_id:
_mark_consumed(tool_idx, msg.tool_call_id)
for tool_idx, msg in enumerate(self.messages):
if msg.role != "tool" or msg.tool_call_id:
continue
picked: str | None = None
for asst_idx in range(tool_idx - 1, -1, -1):
prev = self.messages[asst_idx]
if prev.role != "assistant" or not prev.tool_calls:
if prev.role == "user":
break
continue
name_match = None
fallback = None
for tc_idx, tc in enumerate(prev.tool_calls):
if (asst_idx, tc_idx) in consumed:
continue
if not isinstance(tc, dict):
continue
tc_id = tc.get("id")
if not tc_id:
continue
function = tc.get("function")
function_name = (
function.get("name") if isinstance(function, dict) else None
)
if msg.name and function_name == msg.name:
name_match = (tc_id, asst_idx, tc_idx)
break
if fallback is None:
fallback = (tc_id, asst_idx, tc_idx)
chosen = name_match or fallback
if chosen is not None:
picked, a, t = chosen
consumed.add((a, t))
break
if picked is None:
import secrets as _secrets
picked = f"call_{_secrets.token_hex(8)}"
msg.tool_call_id = picked
return self
# ── OpenAI shell-tool container management ─────────────────────

180
studio/backend/tests/test_cleanup_cancelled_checkpoints.py Normal file
View file

@ -0,0 +1,180 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Tests for core/training/training.py:_cleanup_cancelled_checkpoints."""
import os
import sys
from pathlib import Path
import pytest
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
@pytest.fixture
def outputs_setup(tmp_path, monkeypatch):
"""Point outputs_root() at a temp dir so cleanup is allowed to run on it.
The training module binds ``outputs_root`` at import time
(``from utils.paths import outputs_root``), so we have to patch
the symbol on the importer module, not on storage_roots.
"""
from core.training import training as training_mod
monkeypatch.setattr(training_mod, "outputs_root", lambda: tmp_path)
return tmp_path
def _mk_dir(parent: Path, name: str) -> Path:
p = parent / name
p.mkdir()
(p / "marker.txt").write_text(name)
return p
def test_completed_checkpoints_are_preserved(outputs_setup):
"""The big regression: prior to this fix, every completed
checkpoint-N/ was rmtree'd on Cancel, destroying resume points."""
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-1"
out.mkdir()
ckpts = [_mk_dir(out, f"checkpoint-{n}") for n in (200, 400, 600)]
tmp = _mk_dir(out, "tmp-checkpoint-800")
_cleanup_cancelled_checkpoints(out)
for c in ckpts:
assert c.exists(), f"completed {c.name} was destroyed"
assert (c / "marker.txt").exists()
assert not tmp.exists(), "in-flight tmp-checkpoint-800 should be removed"
def test_in_flight_tmp_checkpoints_removed(outputs_setup):
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-2"
out.mkdir()
_mk_dir(out, "tmp-checkpoint-100")
_mk_dir(out, "tmp-checkpoint-200")
_mk_dir(out, "checkpoint-50") # completed, kept
_cleanup_cancelled_checkpoints(out)
assert not (out / "tmp-checkpoint-100").exists()
assert not (out / "tmp-checkpoint-200").exists()
assert (out / "checkpoint-50").exists()
def test_non_checkpoint_dirs_left_alone(outputs_setup):
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-3"
out.mkdir()
_mk_dir(out, "logs")
_mk_dir(out, "tensorboard")
_mk_dir(out, "checkpoint-final") # non-int suffix, kept
_mk_dir(out, "checkpoint-best")
_mk_dir(out, "tmp-checkpoint-99")
_cleanup_cancelled_checkpoints(out)
for n in ("logs", "tensorboard", "checkpoint-final", "checkpoint-best"):
assert (out / n).exists(), f"{n} should be preserved"
assert not (out / "tmp-checkpoint-99").exists()
def test_output_dir_outside_outputs_root_is_refused(tmp_path, monkeypatch):
"""Containment check: even if a bug passed an output_dir outside
outputs_root, the cleanup must refuse to touch it."""
from core.training import training as training_mod
from core.training.training import _cleanup_cancelled_checkpoints
inside = tmp_path / "inside"
inside.mkdir()
monkeypatch.setattr(training_mod, "outputs_root", lambda: inside)
outside = tmp_path / "outside"
outside.mkdir()
_mk_dir(outside, "tmp-checkpoint-1")
_cleanup_cancelled_checkpoints(outside)
assert (
outside / "tmp-checkpoint-1"
).exists(), "must not rmtree under a path outside outputs_root"
def test_symlinked_output_dir_skipped(outputs_setup):
"""A symlinked output_dir is skipped so the realpath check can't be
leveraged to delete content via a symlink trick."""
from core.training.training import _cleanup_cancelled_checkpoints
real = outputs_setup / "real-run"
real.mkdir()
_mk_dir(real, "tmp-checkpoint-1")
link = outputs_setup / "link-run"
try:
link.symlink_to(real, target_is_directory = True)
except (OSError, NotImplementedError):
pytest.skip("symlinks not supported on this filesystem / platform")
_cleanup_cancelled_checkpoints(link)
assert (real / "tmp-checkpoint-1").exists(), "symlinked output_dir must be skipped"
def test_missing_output_dir_is_noop(outputs_setup):
from core.training.training import _cleanup_cancelled_checkpoints
_cleanup_cancelled_checkpoints(outputs_setup / "does-not-exist")
# Should not raise; nothing to assert beyond non-failure.
def test_symlinked_child_skipped(outputs_setup):
"""A symlinked tmp-checkpoint-* child must not be deleted, so the
realpath bypass cannot redirect rmtree to arbitrary content."""
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-symchild"
out.mkdir()
target = outputs_setup / "external"
target.mkdir()
(target / "important.txt").write_text("keep me")
link = out / "tmp-checkpoint-99"
try:
link.symlink_to(target, target_is_directory = True)
except (OSError, NotImplementedError):
pytest.skip("symlinks not supported on this filesystem / platform")
_cleanup_cancelled_checkpoints(out)
assert (
target / "important.txt"
).exists(), "symlink target outside outputs_root must not be rmtree'd"
def test_non_numeric_tmp_checkpoint_suffix_preserved(outputs_setup):
"""HF Trainer's partials are tmp-checkpoint-<step>. A user-named
tmp-checkpoint-final / tmp-checkpoint-backup / tmp-checkpoint-notes
must NOT be deleted by the cancel cleanup."""
from core.training.training import _cleanup_cancelled_checkpoints
out = outputs_setup / "run-non-numeric"
out.mkdir()
numeric = _mk_dir(out, "tmp-checkpoint-100")
user_final = _mk_dir(out, "tmp-checkpoint-final")
user_backup = _mk_dir(out, "tmp-checkpoint-backup")
user_notes = _mk_dir(out, "tmp-checkpoint-user-notes")
_cleanup_cancelled_checkpoints(out)
assert not numeric.exists(), "in-flight tmp-checkpoint-100 should be removed"
assert user_final.exists(), "user dir tmp-checkpoint-final must be preserved"
assert user_backup.exists(), "user dir tmp-checkpoint-backup must be preserved"
assert user_notes.exists(), "user dir tmp-checkpoint-user-notes must be preserved"

190
studio/backend/tests/test_inference_model_validation.py
View file

@ -34,3 +34,193 @@ def test_nonblank_chat_template_override_is_preserved_verbatim():
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"

18
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -125,21 +125,23 @@ class TestChatMessageToolRoles:
)
assert msg.content is None
def test_tool_role_missing_tool_call_id_synthesised(self):
# Frontend drops the id on second-round POST; validator synthesises one.
def test_tool_role_missing_tool_call_id_left_for_request_validator(self):
# Per-message: missing tool_call_id is now allowed at this layer.
# ChatCompletionRequest's walkback fills it in from the prior
# assistant tool_calls; see test_inference_model_validation.py for
# the resolution coverage.
msg = ChatMessage(role = "tool", content = '{"temperature": 72}')
assert msg.tool_call_id is not None
assert msg.tool_call_id.startswith("call_")
assert len(msg.tool_call_id) >= len("call_") + 8
assert msg.tool_call_id is None
assert msg.content == '{"temperature": 72}'
def test_tool_role_empty_tool_call_id_synthesised(self):
def test_tool_role_empty_tool_call_id_left_for_request_validator(self):
msg = ChatMessage(
role = "tool",
tool_call_id = "",
content = '{"temperature": 72}',
)
assert msg.tool_call_id is not None
assert msg.tool_call_id.startswith("call_")
# Empty-string is treated the same as missing by the walkback.
assert msg.tool_call_id in (None, "")
# ── Role-aware content requirements ────────────────────────────