fix(server): resolve seven server defects found on 96dab686

src/twinkle/loss/base.py

@@ -12,6 +12,14 @@ class Loss:
     def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput:
         ...
 
+    def validate_inputs(self, inputs, **kwargs) -> None:
+        """Reject a batch before any rank starts its forward pass.
+
+        Loss implementations whose alignment rules can fail per sample override
+        this hook. The default accepts every batch.
+        """
+        return None
+
     def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float:
         if len(indices) == len(inputs):
             return 1.0

src/twinkle/loss/grpo.py

@@ -152,10 +152,10 @@ def _pad_and_align_to_batch(
           * Response-only form (e.g. ``old_logps`` from vLLM): length equals
             the number of trainable positions in ``mask[i]`` and is scattered
             directly onto those positions.
-          * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward,
-            right-padded to ``mask.shape[1]``): length ``>= mask.shape[1]``;
-            we slice to ``seq_len`` and index by ``mask[i]`` to extract the
-            trainable positions, then scatter.
+          * Full-sequence form (e.g. ``ref_logps`` from a ref-model forward):
+            every trainable mask position must exist in the per-sample row.
+            Rows need not include the batch's right-padding; we index by
+            ``mask[i]`` to extract the trainable positions, then scatter.
 
         Any other length is a real bug and triggers a hard assert — never
         silently truncate, since that misaligns IS ratios to the wrong tokens.
@@ -196,25 +196,96 @@ def _pad_and_align_to_batch(
             pos = mask[i].nonzero(as_tuple=True)[0]
             n_pos = len(pos)
             n_sample = sample.numel()
+            last_pos = int(pos[-1].item()) if n_pos else -1
+            if not self._alignment_supported(n_pos, n_sample, last_pos):
+                self._raise_alignment_mismatch(i, n_pos, n_sample, seq_len)
 
             if n_sample == 1:
                 result[i, pos] = sample.item()
             elif n_sample == n_pos:
                 # Response-only form (e.g. old_logps from vLLM).
                 result[i, pos] = sample
-            elif n_pos == 0 or (n_sample > 0 and pos[-1].item() < n_sample):
+            else:
                 # Variable-length full-sequence form. The processor right-pads the
                 # batch, but per-sample RL fields from Tinker remain unpadded. They
                 # are valid when every selected mask position exists in this row.
                 result[i, pos] = sample[pos]
-            else:
-                raise AssertionError(f'data/mask length mismatch at sample {i}: '
-                                     f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} '
-                                     '(expected n_sample == n_pos for response-only form, or all masked positions '
-                                     'to exist in the per-sample full-sequence form)')
 
         return result
 
+    @staticmethod
+    def _alignment_supported(n_pos: int, n_sample: int, last_pos: int) -> bool:
+        """Whether one ragged row fits the response-only or full-sequence form."""
+        return (n_sample == 1 or n_sample == n_pos or n_pos == 0 or (n_sample > 0 and last_pos < n_sample))
+
+    @staticmethod
+    def _raise_alignment_mismatch(sample_index: int, n_pos: int, n_sample: int, seq_len: int):
+        raise AssertionError(f'data/mask length mismatch at sample {sample_index}: '
+                             f'n_pos={n_pos}, n_sample={n_sample}, seq_len={seq_len} '
+                             '(expected n_sample == n_pos for response-only form, or all masked positions '
+                             'to exist in the per-sample full-sequence form)')
+
+    @staticmethod
+    def _alignment_rows(data, batch_size: int):
+        """Return one row per sample, or None when the shape is not attributable."""
+        import numpy as np
+        import torch
+
+        if torch.is_tensor(data):
+            if data.dim() == 1 and data.shape[0] == batch_size:
+                return [data[index].reshape(1) for index in range(batch_size)]
+            if data.dim() >= 2 and data.shape[0] == batch_size:
+                return list(data)
+            return None
+        if isinstance(data, np.ndarray):
+            return GRPOLoss._alignment_rows(torch.as_tensor(data), batch_size)
+        if isinstance(data, (list, tuple)) and len(data) == batch_size:
+            return list(data)
+        return None
+
+    def validate_inputs(self, inputs, *, old_logps=None, ref_logps=None, advantages=None, **kwargs) -> None:
+        """Reject ragged per-sample lengths before any rank starts its forward.
+
+        ``_pad_and_align_to_batch`` already rejects these shapes, but it runs inside
+        the loss on each rank independently. Raising there lets one rank return
+        while the others enter the DDP backward collective, which deadlocks the
+        replica. Running the same rule on the full, unsliced batch first makes
+        every rank fail together before any collective is entered.
+        """
+        import torch
+
+        if old_logps is None and ref_logps is None and advantages is None:
+            return
+        samples = inputs if isinstance(inputs, list) else [inputs]
+        if not samples:
+            return
+        for data in (old_logps, ref_logps, advantages):
+            if data is None:
+                continue
+            rows = self._alignment_rows(data, len(samples))
+            if rows is None:
+                # The shape does not map one row per sample; the in-loss check
+                # remains the only applicable validation for this batch.
+                continue
+            for index, (sample, row) in enumerate(zip(samples, rows)):
+                if not isinstance(sample, dict):
+                    continue
+                labels = sample.get('labels')
+                if labels is None:
+                    continue
+                if not torch.is_tensor(labels):
+                    labels = torch.as_tensor(labels)
+                if labels.dim() == 1:
+                    labels = labels.unsqueeze(0)
+                mask = self._resolve_loss_mask(sample, labels)[0]
+                pos = mask.nonzero(as_tuple=True)[0]
+                row_tensor = row if torch.is_tensor(row) else torch.as_tensor(row)
+                n_sample = row_tensor.numel()
+                n_pos = len(pos)
+                last_pos = int(pos[-1].item()) if n_pos else -1
+                if not self._alignment_supported(n_pos, n_sample, last_pos):
+                    self._raise_alignment_mismatch(index, n_pos, n_sample, labels.shape[-1])
+
     def _resolve_loss_mask(self, inputs: Dict, labels: 'torch.Tensor') -> 'torch.Tensor':
         """Positions this loss may score: trainable *and* log-prob-bearing.
 

src/twinkle/server/checkpoint/checkpoint_manager.py

@@ -456,12 +456,21 @@ def delete(self, model_id: str, checkpoint_id: str) -> bool:
         """Remove a checkpoint and refresh only its kind's last-checkpoint pointer.
 
         Sampler links own versions under .versions; legacy directories remain
-        supported. Missing checkpoints or traversal attempts return False.
+        supported. Missing checkpoints, traversal attempts, and paths that are
+        not a real checkpoint all return False.
         """
         # Basic safety check to prevent directory traversal
         if '..' in checkpoint_id:
             return False
 
+        # Delete only a checkpoint that the manager itself can resolve. This
+        # rejects kind roots (``weights``), the run root (``.``), aliases whose
+        # metadata names another checkpoint (``sampler_weights/latest``), and
+        # any directory that has no checkpoint metadata.
+        checkpoint = self.get(model_id, checkpoint_id)
+        if checkpoint is None or checkpoint.checkpoint_id != checkpoint_id:
+            return False
+
         ckpt_dir = self.get_ckpt_dir(model_id, checkpoint_id)
 
         if not ckpt_dir.exists():

src/twinkle/server/gateway/openai_handlers.py

@@ -8,7 +8,9 @@
 """
 from __future__ import annotations
 
+import asyncio
 import json
+import os
 import uuid
 from collections.abc import Callable
 from fastapi import Depends, FastAPI, Request
@@ -22,13 +24,68 @@
 
 import httpx
 
+import twinkle.protocol.types as types
+from twinkle.server.lifecycle.envelope import envelope_from_record
 from twinkle.server.utils import get_template_for_model
 from twinkle.utils.logger import get_logger
 from . import routes
 from .openai_bridge import make_error, translate_chat_request, translate_response, translate_stream_chunk
+from .use_cases import poll_future
 
 logger = get_logger()
 
+_OPENAI_SAMPLE_TIMEOUT_ENV = 'TWINKLE_OPENAI_SAMPLE_TIMEOUT'
+_DEFAULT_OPENAI_SAMPLE_TIMEOUT = 1800.0
+
+
+class SamplerWaitTimeoutError(TimeoutError):
+    """The sampler task did not reach a terminal state inside the gateway's budget."""
+
+
+def _sample_wait_timeout() -> float:
+    return float(os.environ.get(_OPENAI_SAMPLE_TIMEOUT_ENV, str(_DEFAULT_OPENAI_SAMPLE_TIMEOUT)))
+
+
+async def _resolve_sampler_envelope(gateway: GatewayServer, sampler_data: dict[str, Any],
+                                    request: Request) -> types.TaskEnvelope | None:
+    """Resolve the sampler's TaskEnvelope to a terminal state.
+
+    ``/twinkle/sample`` always answers with a TaskEnvelope; a slow sample is
+    still ``queued``/``running`` when the proxy returns. Polling the shared
+    state backend is the only way to obtain its result without changing the
+    upstream wire contract. Returns ``None`` when the future record was lost.
+
+    ``TWINKLE_OPENAI_SAMPLE_TIMEOUT`` bounds the entire terminal-state wait,
+    including state reads. Client disconnection is checked between long polls.
+    """
+    envelope = types.TaskEnvelope.model_validate(sampler_data)
+    if envelope.status in types.TERMINAL_STATUSES:
+        return envelope
+
+    async def wait_for_terminal(envelope: types.TaskEnvelope) -> types.TaskEnvelope | None:
+        while envelope.status not in types.TERMINAL_STATUSES:
+            if await request.is_disconnected():
+                raise SamplerWaitTimeoutError('Client disconnected while waiting for the sampling task.')
+            # ``poll_future`` waits out its whole window before reporting a missing
+            # record, so check once here: a lost record is a terminal answer, not a
+            # reason to keep the coroutine alive.
+            if await gateway.state.get_future(envelope.request_id) is None:
+                return None
+            outcome = await poll_future(gateway.state, envelope.request_id)
+            if outcome.record is None:
+                return None
+            envelope = envelope_from_record(envelope.request_id, outcome.record)
+        return envelope
+
+    timeout = _sample_wait_timeout()
+    try:
+        return await asyncio.wait_for(wait_for_terminal(envelope), timeout=timeout)
+    except SamplerWaitTimeoutError:
+        raise
+    except asyncio.TimeoutError:
+        raise SamplerWaitTimeoutError(
+            f'Sampling task {envelope.request_id} did not finish within {timeout:g}s.') from None
+
 
 def _register_openai_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None:
     """Register OpenAI-compatible routes on the gateway FastAPI app."""
@@ -105,8 +162,43 @@ async def chat_completions(
                 )
 
             sampler_data = json.loads(response.body)
+            try:
+                envelope = await _resolve_sampler_envelope(self, sampler_data, request)
+            except SamplerWaitTimeoutError as e:
+                return JSONResponse(
+                    status_code=504,
+                    content=make_error(message=str(e), error_type='timeout'),
+                )
+            if envelope is None:
+                return JSONResponse(
+                    status_code=500,
+                    content=make_error(
+                        message=f'Sampling task {sampler_data.get("request_id")} was lost by the server',
+                        error_type='server_error',
+                    ),
+                )
+            if envelope.status != 'completed':
+                error = envelope.error
+                message = error.error if error is not None else f'Sampling task {envelope.status}'
+                error_type = ('invalid_request_error'
+                              if error is not None and error.category.value == 'user' else 'server_error')
+                return JSONResponse(
+                    status_code=error.error_code if error is not None else 500,
+                    content=make_error(message=message, error_type=error_type),
+                )
+
+            sampler_result = envelope.result or {}
+            if not sampler_result.get('samples'):
+                return JSONResponse(
+                    status_code=502,
+                    content=make_error(
+                        message='Sampler completed without returning any samples',
+                        error_type='server_error',
+                    ),
+                )
+
             request_id = f'chatcmpl-{uuid.uuid4().hex[:24]}'
-            openai_response = translate_response(sampler_data, model, request_id)
+            openai_response = translate_response(sampler_result, model, request_id)
             resp_headers = {}
             replica_id = response.headers.get('x-twinkle-replica-id')
             if replica_id:
@@ -227,9 +319,11 @@ async def _ensure_template(
         return
 
     template_cls = get_template_for_model(base_model)
+    # The sampler's strict schema accepts only template_cls / adapter_name /
+    # init_kwargs. It already knows its own model_id and injects it when the
+    # template constructor needs one, so sending model_id here is rejected 422.
     set_template_body = json.dumps({
         'template_cls': template_cls,
-        'model_id': base_model,
         'adapter_name': '',
     }).encode()
 

src/twinkle/server/gateway/use_cases.py

@@ -18,6 +18,8 @@
 from twinkle.server.checkpoint import create_checkpoint_manager, create_training_run_manager
 from twinkle.server.lifecycle.poll_config import long_poll_window, retrieve_poll_interval
 
+# Deliberately local: this layer works on canonical records and must not import the
+# protocol models (see test_gateway_services_do_not_import_protocol_models).
 _TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'})
 
 

src/twinkle/server/lifecycle/submit.py

@@ -198,7 +198,11 @@ async def _task():
     if seq_id is not None:
         session_id = get_session_id_from_request(request) or request.state.request_id
         dedup_key = f'seq::{session_id}::{adapter_name or "-"}::{seq_id}'
-        ttl = int(self.task_queue_config.effective_execution_timeout) + 60
+        # The claim must outlive the whole legal lifetime of the original task:
+        # it may wait in the queue and then execute. A TTL that covers only the
+        # execution bound lets a retry re-enqueue while the original is still
+        # queued, i.e. applies one gradient update twice.
+        ttl = int(self.task_queue_config.queue_timeout + self.task_queue_config.effective_execution_timeout) + 60
         prior_request_id = await self.state.claim_seq(dedup_key, request_id, ttl)
         if prior_request_id is not None:
             return await self._peek_terminal(prior_request_id, fallback_status='pending')

src/twinkle/server/model/app.py

@@ -269,6 +269,16 @@ def assert_full_mode_available(self, adapter_name: str | None = None) -> None:
         if holder is not None:
             raise FullModeBusyError(holder)
 
+    def claim_full_mode_resource(self, adapter_name: str, token: str, session_id: str) -> None:
+        """Atomically take the exclusive full-mode resource for one tenant.
+
+        The check and the registration are synchronous neighbours: no coroutine
+        can interleave between them, so two concurrent creates cannot both pass
+        the availability check before either is recorded.
+        """
+        self.assert_full_mode_available(adapter_name)
+        self.register_resource(adapter_name, token, session_id)
+
     async def _on_adapter_expired(self, adapter_name: str) -> None:
         self.fail_pending_tasks_for_model(adapter_name, reason='Adapter expired')
         await self._cleanup_adapter(adapter_name)

src/twinkle/server/model/backends/common.py

@@ -10,7 +10,8 @@
 from tinker import types
 from typing import Any, List
 
-from twinkle import DeviceMesh
+from twinkle import DeviceMesh, remote_function
+from twinkle.server.model.tinker_datum import datum_to_input_feature, extract_rl_features_for_loss
 from twinkle.server.validation.tinker_loss import validate_tinker_loss
 from twinkle.template import Template
 
@@ -173,6 +174,35 @@ def _tinker_setup_loss(self, loss_fn: str, inputs, adapter_name: str, kwargs: di
                 grpo_beta = kwargs.pop('beta', 0.0)
                 self.set_loss('GRPOLoss', adapter_name=adapter_name, epsilon=epsilon, beta=grpo_beta)
 
+    @remote_function(dispatch='all', collect='none')
+    def preflight_loss_inputs(self, *, inputs, adapter_name: str = None, loss_fn: str = None, **kwargs) -> None:
+        """Validate the full, unsliced batch on every rank before any forward.
+
+        Each backend's ``forward_backward`` is dispatched with ``slice_dp``, so a
+        rank that rejects its slice inside the loss has already forwarded while
+        its peers enter the backward collective, which deadlocks the replica
+        (BUG-05). Running the loss's own rule on the full batch first, on every
+        rank, makes them fail together before any collective is entered.
+
+        ``loss_fn`` is the tinker route's loss selector: tinker batches carry
+        their RL features inside the datums, so the loss has to be set up and
+        the datums converted before the same rule applies. Twinkle-native
+        batches already pass ``old_logps``/``ref_logps``/``advantages`` as
+        keyword arguments and leave ``loss_fn`` unset.
+        """
+        adapter_name = adapter_name or self._get_default_group()
+        if loss_fn is not None:
+            loss_values = extract_rl_features_for_loss(inputs)
+            loss_values.pop('ref_outputs', None)
+            self._tinker_setup_loss(loss_fn, inputs, adapter_name, kwargs)
+            inputs = datum_to_input_feature(inputs, self.get_template(adapter_name))
+            loss_values.update(kwargs)
+        else:
+            loss_values = kwargs
+        loss = getattr(self.optimizer_group.get(adapter_name), 'loss_instance', None)
+        if loss is not None:
+            loss.validate_inputs(inputs, **loss_values)
+
     def _ensure_dpo_metric(self, adapter_name: str, beta: float):
         """Add DPOMetric for the adapter if not already present.
 

src/twinkle/server/model/tinker_handlers.py

@@ -184,6 +184,14 @@ async def _do_forward_backward():
                 datum_list = body.forward_backward_input.data
                 loss_fn = body.forward_backward_input.loss_fn
                 loss_fn_config = body.forward_backward_input.loss_fn_config or {}
+                preflight = getattr(self.model, 'preflight_loss_inputs', None)
+                if preflight is not None and loss_fn == 'importance_sampling' and not validate_tinker_loss(
+                        loss_fn, datum_list, data_world_size=self.data_world_size):
+                    # GRPO datums carry ragged per-sample targets; the loss check inside
+                    # the sliced backend would deadlock the peer ranks (BUG-05), so the
+                    # full batch is validated on every rank first.
+                    await self.call_backend(
+                        preflight, inputs=datum_list, adapter_name=model_adapter, loss_fn=loss_fn, **loss_fn_config)
                 output, loss = await self.call_backend(
                     self.model.tinker_forward_backward,
                     inputs=datum_list,