diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index e4abb64b8b..d2c2316d45 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -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-`` subdirs after a cancelled run. - Only paths whose realpath is under outputs_root are touched.""" + """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. + + Completed ``checkpoint-/`` 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) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 19b5c3dcd3..d03a8dff44 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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 ───────────────────── diff --git a/studio/backend/tests/test_cleanup_cancelled_checkpoints.py b/studio/backend/tests/test_cleanup_cancelled_checkpoints.py new file mode 100644 index 0000000000..0d09f027cf --- /dev/null +++ b/studio/backend/tests/test_cleanup_cancelled_checkpoints.py @@ -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-. 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" diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py index 219affade3..ebd9c6c722 100644 --- a/studio/backend/tests/test_inference_model_validation.py +++ b/studio/backend/tests/test_inference_model_validation.py @@ -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" diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index e87faa3b32..638cbc12c8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -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 ────────────────────────────