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

@@ -208,13 +208,83 @@ def _pad_and_align_to_batch(
                 # 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)')
+                self._raise_alignment_mismatch(i, n_pos, n_sample, seq_len)
 
         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

@@ -22,13 +22,35 @@
 
 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()
 
+_TERMINAL_STATUSES = frozenset({'completed', 'failed', 'cancelled'})
+
+
+async def _resolve_sampler_envelope(gateway: GatewayServer, sampler_data: dict[str, Any]) -> 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.
+    """
+    envelope = types.TaskEnvelope.model_validate(sampler_data)
+    while envelope.status not in _TERMINAL_STATUSES:
+        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
+
 
 def _register_openai_routes(app: FastAPI, self_fn: Callable[[], GatewayServer]) -> None:
     """Register OpenAI-compatible routes on the gateway FastAPI app."""
@@ -105,8 +127,37 @@ async def chat_completions(
                 )
 
             sampler_data = json.loads(response.body)
+            envelope = await _resolve_sampler_envelope(self, sampler_data)
+            if envelope is None:
+                return JSONResponse(
+                    status_code=404,
+                    content=make_error(
+                        message=f'Sampling task {sampler_data.get("request_id")} was not found',
+                        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 +278,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/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/transformers_model.py

@@ -103,6 +103,26 @@ def forward_only(self, *, inputs: InputFeature | list[InputFeature] | Trajectory
         output = super().forward_only(inputs=inputs, **kwargs)
         return to_cpu_safe_output(output)
 
+    @remote_function(dispatch='all', collect='none')
+    def preflight_loss_inputs(self,
+                              *,
+                              inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory],
+                              adapter_name: str = None,
+                              **kwargs):
+        """Validate the full, unsliced batch on every rank before any forward.
+
+        The loss's own alignment check runs after a rank has already forwarded
+        its slice. If one rank raises there while another proceeds to backward,
+        the DDP collective deadlocks the replica (BUG-05). Running the same
+        check on the full batch first makes every rank fail together while no
+        rank holds an in-flight forward or collective.
+        """
+        adapter_name = adapter_name or self._get_default_group()
+        group = self.optimizer_group.get(adapter_name)
+        loss_instance = getattr(group, 'loss_instance', None)
+        if loss_instance is not None:
+            loss_instance.validate_inputs(inputs, **kwargs)
+
     @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward_backward(self, *, inputs: InputFeature | list[InputFeature] | Trajectory | list[Trajectory], **kwargs):
         """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O)."""

src/twinkle/server/model/twinkle_handlers.py

@@ -137,11 +137,15 @@ def first_element(data):
                 for key in inputs:
                     if isinstance(inputs[key], list) and isinstance(first_element(inputs[key]), (int, float)):
                         inputs[key] = torch.tensor(inputs[key])
+            kwargs = backend_kwargs(body)
+            resolved_adapter_name = self.resolve_model_adapter_name(adapter_name)
+            preflight = getattr(self.model, 'preflight_loss_inputs', None)
+            if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')):
+                # Full, unsliced batch on every rank: a loss alignment error must
+                # fail all ranks before any of them enters its backward collective.
+                await self.call_backend(preflight, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs)
             ret = await self.call_backend(
-                self.model.forward_backward,
-                inputs=all_inputs,
-                adapter_name=self.resolve_model_adapter_name(adapter_name),
-                **backend_kwargs(body))
+                self.model.forward_backward, inputs=all_inputs, adapter_name=resolved_adapter_name, **kwargs)
             return {'result': ret}
 
         return await run_submit(
@@ -241,11 +245,13 @@ async def forward_backward_from_data_plane(
         async def _call(self, body, adapter_name, token):
             raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane)
             kwargs = merge_forward_kwargs(backend_kwargs(body), field_kwargs)
+            inputs = to_backend_inputs(raw_inputs)
+            resolved_adapter_name = self.resolve_model_adapter_name(adapter_name)
+            preflight = getattr(self.model, 'preflight_loss_inputs', None)
+            if preflight is not None and any(key in kwargs for key in ('old_logps', 'ref_logps', 'advantages')):
+                await self.call_backend(preflight, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs)
             ret = await self.call_backend(
-                self.model.forward_backward,
-                inputs=to_backend_inputs(raw_inputs),
-                adapter_name=self.resolve_model_adapter_name(adapter_name),
-                **kwargs)
+                self.model.forward_backward, inputs=inputs, adapter_name=resolved_adapter_name, **kwargs)
             return {'result': ret}
 
         return await run_submit(
@@ -582,7 +588,14 @@ async def _task():
                 session_id=session_id,
             )
             try:
-                self.register_resource(adapter_name, token, session_id)
+                # Re-check at creation time: the admission check above runs
+                # before the async state write, so two concurrent creates can
+                # both pass it. The claim below is one synchronous step, which
+                # makes exactly one of them win.
+                if self.is_full_mode:
+                    self.claim_full_mode_resource(adapter_name, token, session_id)
+                else:
+                    self.register_resource(adapter_name, token, session_id)
                 if self.is_full_mode:
                     # No PEFT adapter to add; the default optimizer group is used.
                     self.set_resource_state(adapter_name, 'grad_ready', False)

src/twinkle/server/task_queue/worker.py

@@ -35,6 +35,22 @@
 logger = get_logger()
 
 
+def _exception_summary(exc: BaseException) -> str:
+    """Build a single-line summary from the deepest available cause.
+
+    A RayTaskError's own ``str`` starts with the Ray call header (actor id,
+    pid, repr); the actionable error (e.g. the loss assertion) is in ``cause``.
+    Preferring the cause keeps the summary useful without enlarging the wire
+    field, while the full traceback still travels separately.
+    """
+    cause = getattr(exc, 'cause', None)
+    if not isinstance(cause, BaseException):
+        cause = exc.__cause__
+    if isinstance(cause, BaseException):
+        return f'{type(cause).__name__}: {cause}'
+    return f'{type(exc).__name__}: {exc}'
+
+
 def _reason_code_for_server_error(exc: TwinkleServerError) -> str:
     """Map typed execution failures to protocol-independent reason codes."""
     mappings: tuple[tuple[type[TwinkleServerError], str], ...] = (
@@ -481,7 +497,7 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None:
             exec_time = time.monotonic() - exec_start
             # error is a single-line summary; the full traceback goes only to the
             # traceback field, never into `error`.
-            error = f'{type(exc).__name__}: {exc}'
+            error = _exception_summary(exc)
             logger.error(f'[ComputeWorker] Task {task.request_id} FAILED after {exec_time:.2f}s, '
                          f'type={task_type}:\n{traceback.format_exc(limit=3)}')
             await self._store_task_failed(

src/twinkle_client/_future.py

@@ -99,6 +99,7 @@ def _unwrap(env: TaskEnvelope, model_cls) -> Any:
             request_id=env.request_id,
             error_code=p.error_code,
             details=p.details,
+            traceback=p.traceback,
         )
     if env.status == 'cancelled':
         p = env.error

src/twinkle_client/exceptions.py

@@ -80,13 +80,15 @@ def __init__(
         request_id: str,
         error_code: int | None = None,
         details: list[dict[str, Any]] | None = None,
+        traceback: str | None = None,
     ) -> None:
         super().__init__(error)
         self.error = error
         self.category = category
         self.request_id = request_id
         self.error_code = error_code
         self.details = details
+        self.traceback = traceback
 
 
 class TaskCancelledError(Exception):