cortex-training 0.0.3__py3-none-any.whl

cortex_training/client.py

@@ -0,0 +1,2866 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HTTP client for the Cortex Training REST API.
+
+Single-purpose, production-shaped client:
+
+- One auth path: Programmatic Access Token (PAT) — see ``CortexTrainingClient.from_pat``.
+- One CreateJob shape: a list of typed :class:`SubJobConfig` (each carries either
+  a :class:`TrainingConfig` or an :class:`InferenceConfig`). The client-side
+  validators mirror the server's own required-field checks so an invalid job
+  fails locally instead of after a round trip.
+- One wire format: ``{job_id?, sub_job_configs:[SubJobConfig]}`` with per-sub-job
+  ``training_config`` / ``inference_config`` blocks. See
+  ``docs/reference/rest-api.md`` for the full schema.
+
+The server keeps the config blocks open (``additionalProperties``), so unknown
+keys pass through unchanged. The strict, typed shape lives here in the client;
+use the ``extra`` fields for anything the typed layer does not model yet.
+"""
+
+from __future__ import annotations
+
+import base64
+import functools
+import gzip
+import hashlib
+import inspect
+import json
+import logging
+import os
+import re
+import shutil
+import tempfile
+import threading
+import time
+from collections.abc import Iterator
+from collections.abc import Mapping
+from dataclasses import dataclass
+from dataclasses import field
+from enum import Enum
+from pathlib import Path
+from pathlib import PurePosixPath
+from typing import Any
+
+import requests
+from requests.adapters import HTTPAdapter
+from tenacity import Retrying
+from tenacity import retry_if_exception
+from tenacity import stop_after_attempt
+from tenacity import wait_exponential_jitter
+from urllib3.exceptions import NewConnectionError
+
+from cortex_training import wire
+from cortex_training.telemetry import CachedSessionTokenProvider
+from cortex_training.telemetry import OtlpMetricEmitter
+
+logger = logging.getLogger(__name__)
+
+_STDOUT_CHUNK_RE = re.compile(
+    r"^console\.(?P<timestamp>\d{8}-\d{6})\.(?P<suffix>[^./\\]+)\.gz$"
+)
+_GPU_METRICS_CHUNK_RE = re.compile(
+    r"^gpu\.(?P<timestamp>\d{8}-\d{6})\.(?P<suffix>[^./\\]+)\.gz$"
+)
+_STREAM_COPY_BUFFER_BYTES = 1024 * 1024
+
+# Env var that unlocks create-job debug options. These are an internal-only
+# capability and are deliberately not documented for external use: the client
+# refuses to forward a create-job request carrying a `debug` block unless this
+# is set to a truthy value. The server gates them independently.
+DEBUG_OPTIONS_ENV = "CORTEX_TRAINING_ENABLE_DEBUG_OPTIONS"
+DISABLE_TELEMETRY_ENV = "CORTEX_TRAINING_DISABLE_TELEMETRY"
+ENABLE_SUCCESS_TELEMETRY_ENV = "CORTEX_TRAINING_ENABLE_SUCCESS_TELEMETRY"
+
+
+def _env_flag_enabled(name: str) -> bool:
+    return os.environ.get(name, "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+
+
+def _debug_options_enabled() -> bool:
+    """True when create-job debug options are explicitly enabled via env var."""
+    return _env_flag_enabled(DEBUG_OPTIONS_ENV)
+
+
+def _telemetry_disabled() -> bool:
+    """True when OTLP client metrics are explicitly disabled via env var."""
+    return _env_flag_enabled(DISABLE_TELEMETRY_ENV)
+
+
+def _success_telemetry_enabled() -> bool:
+    """True when successful automatic operation outcomes should be emitted."""
+    return _env_flag_enabled(ENABLE_SUCCESS_TELEMETRY_ENV)
+
+
+# HTTP statuses worth retrying: the request was well-formed, so the same call
+# may succeed once transient load/infra conditions clear.
+#   429 Too Many Requests   - rate limited; back off and retry
+#   500 Internal Server Err - transient server-side failure
+#   502 Bad Gateway         - upstream/proxy hiccup in front of the backend
+#   503 Service Unavailable - server temporarily overloaded or restarting
+#   504 Gateway Timeout     - upstream didn't respond in time
+#   409 Conflict            - the target zone is restarting
+#   404 Not Found           - a just-created job may not be visible yet, so a
+#                             valid id can 404 briefly and succeed on retry.
+# Note the trade-off on 404: a *mistyped* job id is also retried (up to
+# `max_retries`, with backoff) before the error surfaces. Every other 4xx is
+# excluded because it signals a client/config error that won't fix itself.
+_TRANSIENT_STATUSES = {429, 500, 502, 503, 504, 404, 409}
+_CHUNK_GROUP_RESTART_REQUIRED = "chunk_group_restart_required"
+_CHUNK_GROUP_ERROR_CODES = {
+    _CHUNK_GROUP_RESTART_REQUIRED,
+    "chunk_group_conflict",
+    "chunk_group_missing_chunks",
+}
+
+
+class ChunkGroupError(requests.exceptions.HTTPError):
+    """Base class for structured chunk-group failures."""
+
+    def __init__(self, message: str, *, response: requests.Response, detail: dict):
+        super().__init__(
+            message,
+            response=response,
+            request=getattr(response, "request", None),
+        )
+        self.detail = detail
+
+
+class ChunkGroupRestartError(ChunkGroupError):
+    """The server lost an incomplete group and the bounded replay also failed."""
+
+
+class ChunkGroupConflictError(ChunkGroupError):
+    """The same chunk-group identity was reused inconsistently."""
+
+
+def _is_transient(exc: BaseException) -> bool:
+    if isinstance(exc, (requests.exceptions.ConnectionError, requests.exceptions.Timeout)):
+        return True
+    if isinstance(exc, requests.exceptions.HTTPError):
+        resp = exc.response
+        return resp is not None and resp.status_code in _TRANSIENT_STATUSES
+    return False
+
+
+def _iter_error_dicts(value: Any, *, depth: int = 0) -> Iterator[dict]:
+    """Yield dictionaries from direct or service-wrapped JSON error bodies."""
+    if depth > 8:
+        return
+    if isinstance(value, dict):
+        yield value
+        for child in value.values():
+            yield from _iter_error_dicts(child, depth=depth + 1)
+        return
+    if isinstance(value, list):
+        for child in value:
+            yield from _iter_error_dicts(child, depth=depth + 1)
+        return
+    if not isinstance(value, str):
+        return
+
+    decoder = json.JSONDecoder()
+    for index, char in enumerate(value):
+        if char not in "[{":
+            continue
+        try:
+            decoded, _ = decoder.raw_decode(value[index:])
+        except (TypeError, ValueError):
+            continue
+        yield from _iter_error_dicts(decoded, depth=depth + 1)
+
+
+def _chunk_group_error_detail(response: requests.Response | None) -> dict | None:
+    if response is None:
+        return None
+    try:
+        body = response.json()
+    except Exception:
+        return None
+    for candidate in _iter_error_dicts(body):
+        if candidate.get("code") in _CHUNK_GROUP_ERROR_CODES:
+            return candidate
+        message = str(candidate.get("message") or "")
+        if candidate.get("chunk_group_id") and message == "request chunk group is missing chunks":
+            return {
+                **candidate,
+                "code": "chunk_group_missing_chunks",
+            }
+    return None
+
+
+_METRIC_IDENTITY_FIELDS = (
+    "job_id",
+    "sub_job_id",
+    "sub_job_type",
+    "source_job_id",
+    "source_sub_job_id",
+    "target_sub_job_id",
+    "target_sub_job_ids",
+    "request_id",
+    "checkpoint_id",
+)
+_REQUEST_ID_RESULT_OPERATIONS = {
+    "forward_backward",
+    "forward",
+    "generate",
+    "generate_stream",
+    "step",
+    "save",
+    "load",
+    "weight_sync",
+    "bootstrap_router_replay",
+    "router_replay_discard",
+    "reset_prefix_cache",
+    "cancel_request",
+}
+_METRIC_ERROR_MESSAGE_LIMIT = 512
+_SECRET_IN_ERROR_RE = re.compile(
+    r"(?i)(authorization[\"']?\s*[:=]\s*"
+    r"(?:bearer\s+|snowflake\s+token\s*=\s*)?|"
+    r"bearer\s+|snowflake\s+token\s*=\s*|"
+    r"(?:password|secret|credential|pat|api[_-]?key|private[_-]?key|"
+    r"access[_-]?token|refresh[_-]?token|session[_-]?token|token)"
+    r"[\"']?\s*[:=]\s*)"
+    r"(?:\"[^\"]*\"|'[^']*'|\S+)"
+)
+
+
+def _safe_metric_error_message(exc: BaseException) -> str:
+    """Return a bounded, single-line error message with obvious secrets redacted."""
+    message = " ".join(str(exc).split())
+    message = _SECRET_IN_ERROR_RE.sub(r"\1<redacted>", message)
+    return message[:_METRIC_ERROR_MESSAGE_LIMIT]
+
+
+def _metric_error_attributes(exc: BaseException) -> dict[str, Any]:
+    attributes: dict[str, Any] = {"error.type": type(exc).__name__}
+    response = getattr(exc, "response", None)
+    if response is None:
+        return attributes
+    status_code = getattr(response, "status_code", None)
+    if isinstance(status_code, int):
+        attributes["error.http_status"] = status_code
+    headers = getattr(response, "headers", {}) or {}
+    request_id = headers.get("x-snowflake-request-id")
+    if isinstance(request_id, str) and request_id:
+        attributes["snowflake.request_id"] = request_id
+    try:
+        body = response.json()
+    except Exception:
+        return attributes
+    for candidate in _iter_error_dicts(body):
+        code = candidate.get("code") or candidate.get("error_code")
+        if code is not None:
+            attributes["error.code"] = str(code)
+            break
+    return attributes
+
+
+def _metric_identity(
+    operation: str,
+    arguments: Mapping[str, Any],
+    result: Any = None,
+) -> dict[str, Any]:
+    attributes = {}
+    for field in _METRIC_IDENTITY_FIELDS:
+        value = arguments.get(field)
+        if value is not None:
+            attributes[field] = (
+                list(value) if isinstance(value, (list, tuple)) else value
+            )
+    if operation == "create_job" and "job_id" not in attributes:
+        if isinstance(result, str):
+            attributes["job_id"] = result
+        elif isinstance(result, dict) and result.get("job_id") is not None:
+            attributes["job_id"] = result["job_id"]
+    if operation in _REQUEST_ID_RESULT_OPERATIONS and "request_id" not in attributes:
+        if isinstance(result, str):
+            attributes["request_id"] = result
+        elif isinstance(result, dict) and result.get("request_id") is not None:
+            attributes["request_id"] = result["request_id"]
+    return attributes
+
+
+def _track_operation(operation: str):
+    """Record one outcome metric around a public client operation."""
+
+    def decorate(method):
+        signature = inspect.signature(method)
+
+        @functools.wraps(method)
+        def wrapped(self, *args, **kwargs):
+            if getattr(self, "_metric_emitter", None) is None:
+                return method(self, *args, **kwargs)
+            state = getattr(self, "_operation_metric_state", None)
+            if state is None:
+                return method(self, *args, **kwargs)
+            active = getattr(state, "active", set())
+            if active:
+                return method(self, *args, **kwargs)
+            try:
+                bound = signature.bind(self, *args, **kwargs)
+                bound.apply_defaults()
+            except TypeError:
+                return method(self, *args, **kwargs)
+            except Exception:
+                logger.debug("client operation metric setup failed", exc_info=True)
+                return method(self, *args, **kwargs)
+            active = set()
+            active.add(operation)
+            state.active = active
+            state.attempt_count = 0
+            state.request_count = 0
+            started = time.monotonic()
+            try:
+                result = method(self, *args, **kwargs)
+            except Exception as exc:
+                self._emit_operation_outcome(
+                    operation,
+                    started,
+                    bound.arguments,
+                    error=exc,
+                )
+                raise
+            else:
+                self._emit_operation_outcome(
+                    operation,
+                    started,
+                    bound.arguments,
+                    result=result,
+                )
+                return result
+            finally:
+                state.active = set()
+                state.attempt_count = 0
+                state.request_count = 0
+
+        return wrapped
+
+    return decorate
+
+
+def _is_chunk_post_transient(exc: BaseException) -> bool:
+    if isinstance(exc, requests.exceptions.HTTPError):
+        if _chunk_group_error_detail(exc.response) is not None:
+            return False
+    return _is_transient(exc)
+
+
+def _is_connect_error(exc: BaseException) -> bool:
+    """Only failures proving the request never reached the server (safe for create_job)."""
+    if isinstance(exc, requests.exceptions.ConnectTimeout):
+        return True
+    if isinstance(exc, requests.exceptions.ConnectionError):
+        cause = exc.args[0] if exc.args else None
+        reason = getattr(cause, "reason", None)
+        return isinstance(cause, NewConnectionError) or isinstance(reason, NewConnectionError)
+    return False
+
+
+# ─── Canonical types ─────────────────────────────────────────────────────
+
+
+class JobType(str, Enum):
+    """Sub-job types supported by Cortex Training.
+
+    Matches the ``job_type`` enum in the REST schema; see
+    ``docs/reference/rest-api.md`` section 8.1.
+    """
+
+    TRAINING = "training"
+    SAMPLING = "sampling"
+    LOG_PROBABILITY = "log_probability"
+
+
+class Hardware(str, Enum):
+    """GPU hardware for create job and capacity.
+
+    Wire values match the REST ``hardware`` field; see
+    ``docs/reference/rest-api.md`` sections 5.1 and 5.4. Omitting it defaults to
+    H200 on the server. Unknown values are rejected client-side.
+    """
+
+    H200 = "H200"
+    B200 = "B200"
+    B300 = "B300"
+
+
+def _hardware_wire(hardware: Hardware | str) -> str:
+    """Return the wire spelling for ``hardware``."""
+    if isinstance(hardware, Hardware):
+        return hardware.value
+    try:
+        return Hardware(hardware).value
+    except ValueError as exc:
+        allowed = ", ".join(member.value for member in Hardware)
+        raise ValueError(f"hardware must be one of {allowed}, got {hardware!r}") from exc
+
+
+def _effective_primerl_config(extra: dict) -> dict:
+    prime_rl = extra.get("prime_rl")
+    if isinstance(prime_rl, dict):
+        # The training runtime expands the nested prime_rl block over the training
+        # config, so validate the same effective values here.
+        return {**extra, **prime_rl}
+    return extra
+
+
+def _validate_primerl_lm_head_config(extra: dict, *, location: str) -> None:
+    cfg = _effective_primerl_config(extra)
+    fused_cross_entropy = cfg.get("fused_cross_entropy", "liger")
+    if fused_cross_entropy is False:
+        return
+
+    token_chunk_size = cfg.get("fused_lm_head_token_chunk_size")
+    token_chunked_lm_head = isinstance(token_chunk_size, int) and not isinstance(token_chunk_size, bool)
+    if cfg.get("fp32_lm_head", False) or token_chunked_lm_head:
+        raise ValueError(
+            f"{location} cannot combine fused_cross_entropy with "
+            "fp32_lm_head or fused_lm_head_token_chunk_size. "
+            "Set fused_cross_entropy=false to use fp32/chunked LM head, "
+            "or remove fp32/chunked LM-head knobs to use fused CE."
+        )
+
+
+@dataclass
+class TrainingConfig:
+    """Training hyperparameters for a training sub-job.
+
+    Required fields mirror the server-side validator: ``max_seq_len > 0``,
+    ``train_batch_size > 0``, ``n_gpus > 0``, and a non-empty ``optimizer``.
+
+    ``extra`` carries any additional fields the training worker consumes. The
+    server keeps this block open (additionalProperties), so unknown keys flow
+    through unchanged. See ``docs/reference/rest-api.md`` section 8.2.
+    """
+
+    optimizer: dict
+    max_seq_len: int
+    train_batch_size: int
+    n_gpus: int
+    gradient_clipping: float | None = None
+    multiplex_job_id: str | None = None
+    # When False, resuming from a checkpoint loads weights only and starts the
+    # optimizer fresh. Required to change DP size (the DP-sharded optimizer
+    # cannot be resized); None leaves the server default (True).
+    load_optimizer_states: bool | None = None
+    extra: dict = field(default_factory=dict)
+
+    def validate(self) -> None:
+        if not isinstance(self.optimizer, dict) or not self.optimizer:
+            raise ValueError("training.optimizer is required and must be a non-empty dict")
+        if self.max_seq_len <= 0:
+            raise ValueError("training.max_seq_len must be > 0")
+        if self.train_batch_size <= 0:
+            raise ValueError("training.train_batch_size must be > 0")
+        if self.n_gpus <= 0:
+            raise ValueError("training.n_gpus must be > 0")
+        _validate_primerl_lm_head_config(self.extra, location="training.extra")
+        prime_rl = self.extra.get("prime_rl")
+        if isinstance(prime_rl, dict):
+            _validate_primerl_lm_head_config(prime_rl, location="training.extra.prime_rl")
+
+    def to_wire(self) -> dict:
+        out: dict = {
+            "optimizer": self.optimizer,
+            "max_seq_len": self.max_seq_len,
+            "train_batch_size": self.train_batch_size,
+            "n_gpus": self.n_gpus,
+        }
+        if self.gradient_clipping is not None:
+            out["gradient_clipping"] = self.gradient_clipping
+        if self.multiplex_job_id is not None:
+            out["multiplex_job_id"] = self.multiplex_job_id
+        if self.load_optimizer_states is not None:
+            out["load_optimizer_states"] = self.load_optimizer_states
+        for k, v in self.extra.items():
+            out.setdefault(k, v)
+        return out
+
+
+@dataclass
+class InferenceConfig:
+    """Sampling/log-probability config for an inference sub-job.
+
+    Required fields mirror the server-side validator: ``max_seq_len > 0`` and
+    ``n_gpus > 0``. ``extra`` carries vLLM-style passthrough keys (e.g.
+    ``gpu_memory_utilization``). See ``docs/reference/rest-api.md`` section 8.3.
+    """
+
+    max_seq_len: int
+    n_gpus: int
+    multiplex_job_id: str | None = None
+    extra: dict = field(default_factory=dict)
+
+    def validate(self) -> None:
+        if self.max_seq_len <= 0:
+            raise ValueError("sampling.max_seq_len must be > 0")
+        if self.n_gpus <= 0:
+            raise ValueError("sampling.n_gpus must be > 0")
+
+    def to_wire(self) -> dict:
+        out: dict = {"max_seq_len": self.max_seq_len, "n_gpus": self.n_gpus}
+        if self.multiplex_job_id is not None:
+            out["multiplex_job_id"] = self.multiplex_job_id
+        for k, v in self.extra.items():
+            out.setdefault(k, v)
+        return out
+
+
+@dataclass
+class SubJobConfig:
+    """One sub-job within a CreateJob request.
+
+    Matches the ``SubJobConfig`` schema in ``docs/reference/rest-api.md``
+    section 8.1.
+
+    Exactly one of ``training`` or ``sampling`` must be set, matching the
+    sub-job's ``job_type`` (see :meth:`validate`).
+    """
+
+    job_type: JobType
+    model_name: str
+    training: TrainingConfig | None = None
+    sampling: InferenceConfig | None = None
+    global_batch_size: int | None = None
+    dtype: str | None = None
+    seed: int | None = None
+    model_post_init: list[str] | None = None
+    source_checkpoint_info: dict | None = None
+
+    @classmethod
+    def training_job(
+        cls,
+        model_name: str,
+        *,
+        optimizer: dict,
+        max_seq_len: int,
+        train_batch_size: int,
+        n_gpus: int,
+        gradient_clipping: float | None = None,
+        multiplex_job_id: str | None = None,
+        load_optimizer_states: bool | None = None,
+        extra_training: dict | None = None,
+        global_batch_size: int | None = None,
+        dtype: str | None = None,
+        seed: int | None = None,
+        model_post_init: list[str] | None = None,
+        source_checkpoint_info: dict | None = None,
+    ) -> "SubJobConfig":
+        """Build a training :class:`SubJobConfig`. All training fields are explicit.
+
+        Args:
+            model_name: Model identifier (e.g., "Qwen/Qwen3-1.7B").
+            optimizer: Optimizer config dict (e.g., {"name": "AdamW", "lr": 0.0001}).
+            max_seq_len: Maximum sequence length for training.
+            train_batch_size: Per-device batch size.
+            n_gpus: Number of GPUs (DP size).
+            gradient_clipping: Optional gradient clipping threshold.
+            multiplex_job_id: Optional multiplex job identifier.
+            load_optimizer_states: Whether to load optimizer states from
+                checkpoints. Set to ``False`` when loading across different DP
+                sizes (``n_gpus``) since optimizer states are DP-sharded and
+                cannot be resized. When ``None``, uses server default (True).
+                **This is a creation-time setting and cannot be changed later.**
+            extra_training: Additional training config keys (passthrough).
+            global_batch_size: Global batch size across all GPUs.
+            dtype: Model dtype (e.g., "bfloat16").
+            seed: Random seed.
+            model_post_init: List of post-initialization operations.
+            source_checkpoint_info: Optional checkpoint to resume from at
+                creation time (use ``{"checkpoint_id": "...", "source_job_id": "..."}`).
+        """
+        return cls(
+            job_type=JobType.TRAINING,
+            model_name=model_name,
+            training=TrainingConfig(
+                optimizer=optimizer,
+                max_seq_len=max_seq_len,
+                train_batch_size=train_batch_size,
+                n_gpus=n_gpus,
+                gradient_clipping=gradient_clipping,
+                multiplex_job_id=multiplex_job_id,
+                load_optimizer_states=load_optimizer_states,
+                extra=dict(extra_training) if extra_training else {},
+            ),
+            global_batch_size=global_batch_size,
+            dtype=dtype,
+            seed=seed,
+            model_post_init=model_post_init,
+            source_checkpoint_info=source_checkpoint_info,
+        )
+
+    @classmethod
+    def sampling_job(
+        cls,
+        model_name: str,
+        *,
+        max_seq_len: int,
+        n_gpus: int,
+        multiplex_job_id: str | None = None,
+        extra_sampling: dict | None = None,
+        job_type: JobType = JobType.SAMPLING,
+        global_batch_size: int | None = None,
+        dtype: str | None = None,
+        seed: int | None = None,
+        model_post_init: list[str] | None = None,
+        source_checkpoint_info: dict | None = None,
+    ) -> "SubJobConfig":
+        """Build a sampling/log-probability :class:`SubJobConfig`."""
+        if job_type not in (JobType.SAMPLING, JobType.LOG_PROBABILITY):
+            raise ValueError(f"sampling_job() only accepts SAMPLING or LOG_PROBABILITY, got {job_type!r}")
+        return cls(
+            job_type=job_type,
+            model_name=model_name,
+            sampling=InferenceConfig(
+                max_seq_len=max_seq_len,
+                n_gpus=n_gpus,
+                multiplex_job_id=multiplex_job_id,
+                extra=dict(extra_sampling) if extra_sampling else {},
+            ),
+            global_batch_size=global_batch_size,
+            dtype=dtype,
+            seed=seed,
+            model_post_init=model_post_init,
+            source_checkpoint_info=source_checkpoint_info,
+        )
+
+    def validate(self) -> None:
+        # Mirrors the server-side CreateJob validation.
+        if not self.model_name:
+            raise ValueError("sub_job.model_name is required")
+        if self.training is not None and self.sampling is not None:
+            raise ValueError("sub_job.training and sub_job.sampling are mutually exclusive")
+        if self.job_type == JobType.TRAINING:
+            if self.training is None:
+                raise ValueError("training sub-job requires a `training` block")
+            self.training.validate()
+        elif self.job_type in (JobType.SAMPLING, JobType.LOG_PROBABILITY):
+            if self.sampling is None:
+                raise ValueError(f"{self.job_type.value} sub-job requires a `sampling` block")
+            self.sampling.validate()
+        else:  # pragma: no cover - JobType enum is closed
+            raise ValueError(f"unknown job_type: {self.job_type!r}")
+
+    def to_wire(self) -> dict:
+        wire: dict = {
+            "job_type": self.job_type.value,
+            "model_name": self.model_name,
+        }
+        if self.global_batch_size is not None:
+            wire["global_batch_size"] = self.global_batch_size
+        if self.dtype is not None:
+            wire["dtype"] = self.dtype
+        if self.seed is not None:
+            wire["seed"] = self.seed
+        if self.model_post_init is not None:
+            wire["model_post_init"] = list(self.model_post_init)
+        if self.training is not None:
+            wire["training_config"] = self.training.to_wire()
+        if self.sampling is not None:
+            wire["inference_config"] = self.sampling.to_wire()
+        if self.source_checkpoint_info is not None:
+            wire["source_checkpoint_info"] = self.source_checkpoint_info
+        return wire
+
+
+# ─── Forward-backward payload helpers ────────────────────────────────────
+
+
+def _load_torch():
+    import torch
+
+    return torch
+
+
+def _tensor_dtype(torch, dtype_name: str | None):
+    if dtype_name is None:
+        return torch.long
+    aliases = {
+        "bool": torch.bool,
+        "boolean": torch.bool,
+        "float": torch.float32,
+        "float32": torch.float32,
+        "float64": torch.float64,
+        "double": torch.float64,
+        "bfloat16": torch.bfloat16,
+        "bf16": torch.bfloat16,
+        "float16": torch.float16,
+        "fp16": torch.float16,
+        "int": torch.int32,
+        "int32": torch.int32,
+        "int64": torch.int64,
+        "long": torch.long,
+    }
+    try:
+        return aliases[dtype_name]
+    except KeyError as exc:
+        raise ValueError(f"unsupported tensor dtype: {dtype_name}") from exc
+
+
+def _json_tensor(value: Any, *, default_dtype: str | None = "long"):
+    torch = _load_torch()
+    if isinstance(value, dict):
+        if "data" not in value:
+            raise ValueError("tensor object must contain a data field")
+        dtype_name = value.get("dtype", default_dtype)
+        value = value["data"]
+    else:
+        dtype_name = default_dtype
+    return torch.tensor(value, dtype=_tensor_dtype(torch, dtype_name))
+
+
+def _read_fwd_bwd_args(payload: dict[str, Any]) -> tuple:
+    args = payload.get("args", [])
+    if args is None:
+        return ()
+    if not isinstance(args, list):
+        raise ValueError("fwd-bwd payload args must be a JSON list")
+    return tuple(args)
+
+
+def _tensorize_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    out = {}
+    for key, value in kwargs.items():
+        if isinstance(value, dict) and "data" in value:
+            out[key] = _json_tensor(value)
+        elif isinstance(value, list):
+            out[key] = _json_tensor(value)
+        else:
+            out[key] = value
+    return out
+
+
+def _require_bool(value: Any, name: str) -> bool:
+    if not isinstance(value, bool):
+        raise ValueError(f"{name} must be a boolean")
+    return value
+
+
+def _positive_int(value: Any, name: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+    return value
+
+
+def _repeat_texts(texts: Any, batch_size: Any) -> list[str]:
+    if isinstance(texts, str):
+        text_list = [texts]
+    elif isinstance(texts, list) and texts and all(isinstance(x, str) for x in texts):
+        text_list = texts
+    else:
+        raise ValueError("fwd-bwd payload texts must be a non-empty string list")
+
+    if batch_size is None:
+        return text_list
+    batch = _positive_int(batch_size, "fwd-bwd payload batch_size")
+    return [text_list[i % len(text_list)] for i in range(batch)]
+
+
+def _load_tokenizer(tokenizer_spec: Any):
+    try:
+        from transformers import AutoTokenizer
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "No module named 'transformers'",
+            name="transformers",
+        ) from exc
+
+    if isinstance(tokenizer_spec, str):
+        model_name = tokenizer_spec
+        kwargs: dict[str, Any] = {}
+    elif isinstance(tokenizer_spec, dict):
+        model_name = tokenizer_spec.get("model_name") or tokenizer_spec.get("name") or tokenizer_spec.get("path")
+        if not isinstance(model_name, str) or not model_name:
+            raise ValueError("tokenizer.model_name is required")
+        kwargs = {}
+        for key in ("trust_remote_code", "use_fast", "revision"):
+            if key in tokenizer_spec:
+                kwargs[key] = tokenizer_spec[key]
+    else:
+        raise ValueError("fwd-bwd payload tokenizer must be a string or object")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
+    if getattr(tokenizer, "pad_token_id", None) is None:
+        eos_token = getattr(tokenizer, "eos_token", None)
+        if eos_token is not None:
+            tokenizer.pad_token = eos_token
+    return tokenizer
+
+
+def _tokenize_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    tokenizer_spec = payload.get("tokenizer", payload.get("model_name"))
+    if tokenizer_spec is None:
+        raise ValueError("fwd-bwd payload requires tokenizer or input_ids")
+
+    tokenizer = _load_tokenizer(tokenizer_spec)
+    texts = _repeat_texts(payload.get("texts"), payload.get("batch_size"))
+    padding = payload.get("padding", "max_length")
+    truncation = payload.get("truncation", True)
+    add_special_tokens = payload.get("add_special_tokens", True)
+    max_length = payload.get("max_length")
+    if padding == "max_length":
+        max_length = _positive_int(max_length, "fwd-bwd payload max_length")
+    elif max_length is not None:
+        max_length = _positive_int(max_length, "fwd-bwd payload max_length")
+    _require_bool(truncation, "fwd-bwd payload truncation")
+    _require_bool(add_special_tokens, "fwd-bwd payload add_special_tokens")
+
+    encode_kwargs = {
+        "return_tensors": "pt",
+        "padding": padding,
+        "truncation": truncation,
+        "add_special_tokens": add_special_tokens,
+    }
+    if max_length is not None:
+        encode_kwargs["max_length"] = max_length
+    encoded = tokenizer(texts, **encode_kwargs)
+    if "input_ids" not in encoded:
+        raise ValueError("tokenizer output did not include input_ids")
+    return encoded
+
+
+def _ensure_rank2(tensor, name: str) -> None:
+    if len(tensor.shape) != 2:
+        raise ValueError(f"{name} must have shape [batch, seq_len]")
+
+
+def _build_position_ids(input_ids, position_spec: Any):
+    torch = _load_torch()
+    if position_spec is None or position_spec is False:
+        return None
+    if position_spec is True or position_spec == "arange":
+        batch_size, seq_len = input_ids.shape
+        return torch.arange(seq_len, dtype=torch.long).unsqueeze(0).expand(batch_size, -1).contiguous()
+    if isinstance(position_spec, (dict, list)):
+        return _json_tensor(position_spec)
+    raise ValueError("fwd-bwd payload position_ids must be true, false, arange, or tensor data")
+
+
+def _label_config(payload: dict[str, Any]) -> tuple[Any, int, bool]:
+    labels = payload.get("labels", payload.get("label_strategy", "next_token"))
+    if isinstance(labels, dict) and "data" in labels:
+        return labels, int(payload.get("ignore_index", -100)), True
+    if isinstance(labels, dict):
+        strategy = labels.get("strategy", "next_token")
+        ignore_index = int(labels.get("ignore_index", payload.get("ignore_index", -100)))
+        mask_padding = labels.get("mask_padding", payload.get("mask_padding", True))
+        return strategy, ignore_index, _require_bool(mask_padding, "labels.mask_padding")
+    return (
+        labels,
+        int(payload.get("ignore_index", -100)),
+        _require_bool(
+            payload.get("mask_padding", True),
+            "fwd-bwd payload mask_padding",
+        ),
+    )
+
+
+def _build_labels(input_ids, attention_mask, payload: dict[str, Any]):
+    torch = _load_torch()
+    labels, ignore_index, mask_padding = _label_config(payload)
+    if labels is None or labels == "none":
+        return None
+    if isinstance(labels, (dict, list)):
+        return _json_tensor(labels)
+    if labels in ("next_token", "shifted_input_ids"):
+        out = torch.roll(input_ids, shifts=-1, dims=1)
+        out[:, -1] = ignore_index
+        if mask_padding and attention_mask is not None:
+            target_mask = torch.roll(attention_mask, shifts=-1, dims=1)
+            target_mask[:, -1] = 0
+            out = out.masked_fill(target_mask == 0, ignore_index)
+        return out
+    if labels in ("input_ids", "self"):
+        out = input_ids.clone()
+        if mask_padding and attention_mask is not None:
+            out = out.masked_fill(attention_mask == 0, ignore_index)
+        return out
+    raise ValueError("fwd-bwd labels strategy must be next_token, input_ids, none, or tensor data")
+
+
+def build_forward_backward_kwargs(payload: dict[str, Any]) -> dict[str, Any]:
+    """Build tensor kwargs for a forward-backward training batch.
+
+    ``payload`` is intentionally JSON-shaped so CLI examples can stay readable:
+    callers may provide either tokenizable ``texts`` plus a ``tokenizer`` or
+    direct tensor data under ``input_ids`` / ``labels`` / ``position_ids``.
+    """
+    if not isinstance(payload, dict):
+        raise ValueError("fwd-bwd payload must be a JSON object")
+
+    if "kwargs" in payload:
+        kwargs = payload["kwargs"]
+        if not isinstance(kwargs, dict):
+            raise ValueError("fwd-bwd payload kwargs must be an object")
+        return _tensorize_kwargs(kwargs)
+
+    torch = _load_torch()
+    encoded = None
+    if "input_ids" in payload:
+        input_ids = _json_tensor(payload["input_ids"])
+        attention_mask = (
+            _json_tensor(payload["attention_mask"])
+            if isinstance(payload.get("attention_mask"), (dict, list))
+            else None
+        )
+    else:
+        encoded = _tokenize_payload(payload)
+        input_ids = encoded["input_ids"]
+        attention_mask = encoded.get("attention_mask")
+
+    _ensure_rank2(input_ids, "input_ids")
+    if attention_mask is not None:
+        _ensure_rank2(attention_mask, "attention_mask")
+        if attention_mask.shape != input_ids.shape:
+            raise ValueError("attention_mask must have the same shape as input_ids")
+
+    kwargs: dict[str, Any] = {"input_ids": input_ids.contiguous()}
+
+    include_attention = payload.get("include_attention_mask", False)
+    if "attention_mask" in payload and isinstance(payload["attention_mask"], bool):
+        include_attention = payload["attention_mask"]
+    include_attention = _require_bool(
+        include_attention,
+        "fwd-bwd payload include_attention_mask",
+    )
+    if include_attention:
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+        kwargs["attention_mask"] = attention_mask.contiguous()
+
+    position_spec = payload.get("position_ids")
+    if position_spec is None and payload.get("include_position_ids", False):
+        position_spec = "arange"
+    position_ids = _build_position_ids(input_ids, position_spec)
+    if position_ids is not None:
+        _ensure_rank2(position_ids, "position_ids")
+        if position_ids.shape != input_ids.shape:
+            raise ValueError("position_ids must have the same shape as input_ids")
+        kwargs["position_ids"] = position_ids.contiguous()
+
+    labels = _build_labels(input_ids, attention_mask, payload)
+    if labels is not None:
+        _ensure_rank2(labels, "labels")
+        if labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids")
+        kwargs["labels"] = labels.contiguous()
+
+    return kwargs
+
+
+def serialize_forward_backward_args(args: tuple = (), kwargs: dict[str, Any] | None = None) -> bytes:
+    """Serialize ``{"args": args, "kwargs": kwargs}`` for Cortex Training fwd-bwd."""
+    if not isinstance(args, tuple):
+        raise ValueError("forward-backward args must be a tuple")
+    if kwargs is None:
+        kwargs = {}
+    if not isinstance(kwargs, dict):
+        raise ValueError("forward-backward kwargs must be a dict")
+
+    return wire.dumps(
+        {"args": args, "kwargs": kwargs},
+        metadata={"response_options": {"format": "dssst1", "delivery": "chunked"}},
+    )
+
+
+def build_forward_backward_payload(spec: dict[str, Any]) -> bytes:
+    """Build a serialized fwd-bwd byte payload from readable JSON."""
+    if not isinstance(spec, dict):
+        raise ValueError("fwd-bwd JSON must be an object")
+    payload = spec.get("payload", spec)
+    if not isinstance(payload, dict):
+        raise ValueError("fwd-bwd payload must be an object")
+    args = _read_fwd_bwd_args(payload)
+    kwargs = build_forward_backward_kwargs(payload)
+    return serialize_forward_backward_args(args, kwargs)
+
+
+def _sql_string_literal(value: str) -> str:
+    """Return ``value`` as a single-quoted Snowflake SQL string."""
+    return "'" + value.replace("'", "''") + "'"
+
+
+def _validate_artifact_relative_path(path: str) -> str:
+    """Reject paths that could escape or ambiguously address the run stage."""
+    candidate = PurePosixPath(path)
+    if (
+        not path
+        or path.startswith("/")
+        or "\\" in path
+        or any(part in {"", ".", ".."} for part in path.split("/"))
+        or candidate.name in {"", ".", ".."}
+    ):
+        raise ValueError(f"unsafe experiment artifact path: {path!r}")
+    return path
+
+
+# ─── HTTP client ─────────────────────────────────────────────────────────
+
+
+class CortexTrainingClient:
+    """HTTP client for the Cortex Training REST API (``cortex-training``).
+
+    Construct with :meth:`from_pat`::
+
+        client = CortexTrainingClient.from_pat(
+            host="ACCOUNT.snowflakecomputing.com",
+            pat="...",
+            database="CORTEX_TRAINING_DB",
+            schema="PUBLIC",
+        )
+
+    Against a local or otherwise compatible server, construct the client
+    directly with an explicit ``base_url``; this skips PAT auth::
+
+        client = CortexTrainingClient(
+            base_url="http://localhost:8084",
+            database="MY_DB",
+            schema="PUBLIC",
+        )
+
+    Then build one or more sub-jobs and submit::
+
+        sub = SubJobConfig.training_job(
+            model_name="Qwen/Qwen3-1.7B",
+            optimizer={"name": "AdamW", "lr": 1e-5},
+            max_seq_len=2048,
+            train_batch_size=1,
+            n_gpus=1,
+        )
+        job_id = client.create_job(sub_jobs=[sub])
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        database: str,
+        schema: str,
+        endpoint: str = "cortex-training",
+        poll_interval: float = 0.5,
+        poll_timeout: float = 1800.0,
+        poll_backoff_multiplier: float = 1.25,
+        poll_max_interval: float = 6.0,
+        pool_maxsize: int = 1024,
+        max_retries: int = 10,
+    ):
+        if poll_interval <= 0:
+            raise ValueError("poll_interval must be > 0")
+        if poll_timeout <= 0:
+            raise ValueError("poll_timeout must be > 0")
+        if poll_backoff_multiplier < 1.0:
+            raise ValueError("poll_backoff_multiplier must be >= 1.0")
+        if poll_max_interval < poll_interval:
+            raise ValueError("poll_max_interval must be >= poll_interval")
+        if pool_maxsize <= 0:
+            raise ValueError("pool_maxsize must be > 0")
+        if max_retries < 0:
+            raise ValueError("max_retries must be >= 0")
+        self.base_url = base_url.rstrip("/")
+        self.database = database
+        self.schema = schema
+        self.endpoint = endpoint
+        self.poll_interval = poll_interval
+        self.poll_timeout = poll_timeout
+        self.poll_backoff_multiplier = poll_backoff_multiplier
+        self.poll_max_interval = poll_max_interval
+        self.max_retries = max_retries
+        self._fwd_bwd_send_count = 0
+        self._fwd_bwd_request_debug: dict[str, dict[str, Any]] = {}
+        self._generate_request_ids: set[str] = set()
+        # Per-job cache of the sampling sub-job's ``inference_config.max_seq_len``
+        # (the value the backend launches vLLM with as ``max_model_len``). Used
+        # to validate generate prompt lengths client-side. A cached ``None``
+        # means the job was fetched but exposes no sampling max_seq_len (skip
+        # validation); a missing key means it has not been resolved yet (e.g. a
+        # transient get_job failure), so a later call will retry.
+        self._sampling_max_seq_len: dict[str, int | None] = {}
+        self._artifact_connection_config: dict[str, str] | None = None
+        self._metric_emitter: OtlpMetricEmitter | None = None
+        self._operation_metric_state = threading.local()
+        self._session = requests.Session()
+        adapter = HTTPAdapter(pool_connections=pool_maxsize, pool_maxsize=pool_maxsize)
+        self._session.mount("https://", adapter)
+        self._session.mount("http://", adapter)
+
+    @classmethod
+    def from_pat(
+        cls,
+        host: str,
+        pat: str,
+        database: str,
+        schema: str,
+        endpoint: str = "cortex-training",
+        verify_ssl: bool = True,
+        *,
+        telemetry_timeout: float = 3.0,
+        **kwargs: Any,
+    ) -> "CortexTrainingClient":
+        """Authenticate with a Programmatic Access Token (PAT).
+
+        Sends ``Authorization: Bearer <pat>`` plus the
+        ``X-Snowflake-Authorization-Token-Type: PROGRAMMATIC_ACCESS_TOKEN`` header
+        on every request. No login round-trip.
+        """
+        client = cls(
+            base_url=f"https://{host}",
+            database=database,
+            schema=schema,
+            endpoint=endpoint,
+            **kwargs,
+        )
+        client._session.headers["Authorization"] = f"Bearer {pat}"
+        client._session.headers["X-Snowflake-Authorization-Token-Type"] = "PROGRAMMATIC_ACCESS_TOKEN"
+        client._session.verify = verify_ssl
+        client._artifact_connection_config = {"host": host, "pat": pat}
+        if not _telemetry_disabled():
+            token_provider = CachedSessionTokenProvider(
+                client.base_url,
+                pat,
+                verify_ssl=verify_ssl,
+                timeout=telemetry_timeout,
+            )
+            client._metric_emitter = OtlpMetricEmitter(
+                client.base_url,
+                token_provider,
+                verify_ssl=verify_ssl,
+                timeout=telemetry_timeout,
+            )
+        return client
+
+    @property
+    def _prefix(self) -> str:
+        return f"{self.base_url}/api/v2/databases/{self.database}/schemas/{self.schema}/{self.endpoint}"
+
+    def emit_metric(
+        self,
+        operation: str,
+        value: Any = 1,
+        *,
+        attributes: dict[str, Any] | None = None,
+    ) -> None:
+        """Emit a best-effort OTLP log record used as a client metric.
+
+        PAT clients lazily exchange the PAT for a cached session token on the
+        first call. Local/mock clients, or clients constructed with
+        ``CORTEX_TRAINING_DISABLE_TELEMETRY`` set, treat this method as a no-op.
+        All authentication, discovery, and export errors are ignored.
+        """
+        try:
+            if self._metric_emitter is not None:
+                self._metric_emitter.emit(operation, value, attributes=attributes)
+        except Exception:
+            return
+
+    def close(self) -> None:
+        """Flush telemetry briefly and close this client's HTTP sessions."""
+        try:
+            if self._metric_emitter is not None:
+                self._metric_emitter.close()
+        except Exception:
+            logger.debug("client telemetry close failed", exc_info=True)
+        finally:
+            self._session.close()
+
+    def __enter__(self) -> "CortexTrainingClient":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.close()
+
+    def _emit_operation_outcome(
+        self,
+        operation: str,
+        started: float,
+        arguments: Mapping[str, Any],
+        *,
+        result: Any = None,
+        error: BaseException | None = None,
+    ) -> None:
+        """Emit one structured outcome without changing the operation result."""
+        if self._metric_emitter is None:
+            return
+        if error is None and not _success_telemetry_enabled():
+            return
+        try:
+            attempt_count = int(
+                getattr(self._operation_metric_state, "attempt_count", 0)
+            )
+            request_count = int(
+                getattr(self._operation_metric_state, "request_count", 0)
+            )
+            value: dict[str, Any] = {
+                "attempt_count": attempt_count,
+                "duration_ms": round((time.monotonic() - started) * 1000, 3),
+                "request_count": request_count,
+                "retry_count": max(attempt_count - request_count, 0),
+                "success": error is None,
+            }
+            attributes = _metric_identity(operation, arguments, result)
+            if error is not None:
+                value["error_message"] = _safe_metric_error_message(error)
+                attributes.update(_metric_error_attributes(error))
+            self.emit_metric(operation, value, attributes=attributes)
+        except Exception:
+            logger.debug("client operation metric construction failed", exc_info=True)
+
+    @staticmethod
+    def _debug_json_summary(value: Any, *, limit: int = 2000) -> str:
+        try:
+            text = json.dumps(value, sort_keys=True, default=str)
+        except TypeError:
+            text = repr(value)
+        if len(text) > limit:
+            return f"{text[:limit]}...<truncated {len(text) - limit} chars>"
+        return text
+
+    @classmethod
+    def _debug_response_summary(cls, resp: requests.Response, *, limit: int = 2000) -> str:
+        try:
+            return cls._debug_json_summary(resp.json(), limit=limit)
+        except ValueError:
+            text = getattr(resp, "text", "")
+        except Exception as exc:
+            text = f"<failed to read response body: {type(exc).__name__}: {exc}>"
+        text = str(text).replace("\n", "\\n")
+        if len(text) > limit:
+            return f"{text[:limit]}...<truncated {len(text) - limit} chars>"
+        return text
+
+    @staticmethod
+    def _debug_context_label(debug_context: dict[str, Any]) -> str:
+        parts = ["[cortex-training fwd-bwd]"]
+        for key in ("phase", "call", "job_id", "request_id", "payload_bytes"):
+            value = debug_context.get(key)
+            if value is not None:
+                parts.append(f"{key}={value}")
+        return " ".join(parts)
+
+    def _send(self, method: str, url: str, *, retry_on=_is_transient, **kwargs) -> requests.Response:
+        fn = getattr(self._session, method.lower())
+        debug_context = kwargs.pop("debug_context", None)
+        debug_label = self._debug_context_label(debug_context) if debug_context is not None else None
+        attempt_no = 0
+        max_attempts = 1 + self.max_retries
+        try:
+            state = self._operation_metric_state
+            if getattr(state, "active", set()):
+                state.request_count = getattr(state, "request_count", 0) + 1
+        except Exception:
+            logger.debug("client operation request counting failed", exc_info=True)
+
+        def attempt() -> requests.Response:
+            nonlocal attempt_no
+            attempt_no += 1
+            try:
+                state = self._operation_metric_state
+                if getattr(state, "active", set()):
+                    state.attempt_count = getattr(state, "attempt_count", 0) + 1
+            except Exception:
+                logger.debug("client operation attempt counting failed", exc_info=True)
+            if debug_label is not None:
+                logger.debug(
+                    "%s sending %s %s attempt=%d/%d",
+                    debug_label,
+                    method.upper(),
+                    url,
+                    attempt_no,
+                    max_attempts,
+                )
+            try:
+                resp = fn(url, **kwargs)
+            except Exception as exc:
+                if debug_label is not None:
+                    logger.debug(
+                        "%s request exception %s %s attempt=%d/%d: %s: %s",
+                        debug_label,
+                        method.upper(),
+                        url,
+                        attempt_no,
+                        max_attempts,
+                        type(exc).__name__,
+                        exc,
+                    )
+                raise
+            status_code = getattr(resp, "status_code", None)
+            headers = getattr(resp, "headers", {}) or {}
+            sf_request_id = headers.get("x-snowflake-request-id")
+            if not isinstance(sf_request_id, str) or not sf_request_id:
+                sf_request_id = None
+            if sf_request_id:
+                req = getattr(resp, "request", None)
+                logger.debug(
+                    "snowflake request_id=%s  %s %s  status=%d",
+                    sf_request_id,
+                    getattr(req, "method", method.upper()),
+                    getattr(req, "path_url", url),
+                    status_code,
+                )
+            if debug_label is not None:
+                try:
+                    status_int = int(status_code)
+                except (TypeError, ValueError):
+                    status_int = None
+                outcome = (
+                    "successful response" if status_int is not None and 200 <= status_int < 400 else "failed response"
+                )
+                snowflake = f" snowflake_request_id={sf_request_id}" if sf_request_id else ""
+                logger.debug(
+                    "%s %s %s %s attempt=%d/%d status=%s%s body=%s",
+                    debug_label,
+                    outcome,
+                    method.upper(),
+                    url,
+                    attempt_no,
+                    max_attempts,
+                    status_code,
+                    snowflake,
+                    self._debug_response_summary(resp),
+                )
+            resp.raise_for_status()
+            return resp
+
+        retryer = Retrying(
+            retry=retry_if_exception(retry_on),
+            stop=stop_after_attempt(1 + self.max_retries),
+            wait=wait_exponential_jitter(initial=0.5, max=10.0),
+            reraise=True,
+        )
+        return retryer(attempt)
+
+    # ─── Job management ──────────────────────────────────────────────────
+
+    @_track_operation("create_job")
+    def create_job(
+        self,
+        sub_jobs: list[SubJobConfig],
+        job_id: str | None = None,
+        experiment_name: str | None = None,
+        hardware: Hardware | str | None = None,
+        idle_timeout_seconds: int | None = None,
+        pending_timeout_seconds: int | None = None,
+    ) -> str:
+        """Create a job from a list of sub-jobs and return its server job_id.
+
+        Each :class:`SubJobConfig` is validated client-side before the request
+        is sent (see :meth:`SubJobConfig.validate`). A job supports zero or one
+        ``training`` sub-job and any number of ``sampling`` /
+        ``log_probability`` sub-jobs. ``job_id`` is optional;
+        when omitted the server generates one. ``experiment_name`` is optional;
+        when omitted the server auto-creates an experiment for the job.
+        ``hardware`` is optional (:class:`Hardware`.H200 / .B200 / .B300); when
+        set, all sub-jobs in the job run on that GPU type. Omitted defaults to
+        H200.
+        ``idle_timeout_seconds`` bounds how long the job may sit idle before the
+        server reclaims it: ``0`` disables reclamation, any other value must be
+        between 300 and 604,800. ``pending_timeout_seconds`` bounds how long the
+        job may wait for capacity before failing with ``pending_timeout``: between
+        300 and 604,800, with no disable value. Omit either to use the server
+        default, currently 30 minutes for idle and 24 hours for pending.
+        """
+        if not sub_jobs:
+            raise ValueError("create_job requires a non-empty sub_jobs list")
+        if idle_timeout_seconds is not None:
+            if isinstance(idle_timeout_seconds, bool) or not isinstance(
+                idle_timeout_seconds, int
+            ):
+                raise ValueError("idle_timeout_seconds must be an integer")
+            if idle_timeout_seconds != 0 and not (
+                300 <= idle_timeout_seconds <= 604_800
+            ):
+                raise ValueError(
+                    "idle_timeout_seconds must be 0 or between 300 and 604800"
+                )
+        if pending_timeout_seconds is not None:
+            if isinstance(pending_timeout_seconds, bool) or not isinstance(
+                pending_timeout_seconds, int
+            ):
+                raise ValueError("pending_timeout_seconds must be an integer")
+            if not (300 <= pending_timeout_seconds <= 604_800):
+                raise ValueError(
+                    "pending_timeout_seconds must be between 300 and 604800"
+                )
+        for sj in sub_jobs:
+            sj.validate()
+        body: dict = {"sub_job_configs": [sj.to_wire() for sj in sub_jobs]}
+        if job_id is not None:
+            body["job_id"] = job_id
+        if experiment_name is not None:
+            body["experiment_name"] = experiment_name
+        if hardware is not None:
+            body["hardware"] = _hardware_wire(hardware)
+        if idle_timeout_seconds is not None:
+            body["idle_timeout_seconds"] = idle_timeout_seconds
+        if pending_timeout_seconds is not None:
+            body["pending_timeout_seconds"] = pending_timeout_seconds
+        return self.create_job_from_body(body)["job_id"]
+
+    @_track_operation("create_job")
+    def create_job_from_body(self, body: dict) -> dict:
+        """Create a job from a raw REST CreateJob request body.
+
+        This is useful for tooling that already has the REST JSON payload,
+        while :meth:`create_job` remains the typed path for Python callers.
+
+        A job supports zero or one ``training`` sub-job and any number of
+        ``sampling`` / ``log_probability`` sub-jobs; a second training sub-job
+        raises :class:`ValueError` before the request is sent.
+        """
+        if not isinstance(body, dict):
+            raise ValueError("create_job_from_body requires a JSON object")
+        sub_job_configs = body.get("sub_job_configs")
+        if not isinstance(sub_job_configs, list) or not sub_job_configs:
+            raise ValueError("create_job_from_body requires a non-empty sub_job_configs list")
+        training_sub_jobs = sum(
+            1
+            for cfg in sub_job_configs
+            if isinstance(cfg, dict)
+            and str(cfg.get("job_type") or "").strip().lower() == JobType.TRAINING.value
+        )
+        if training_sub_jobs > 1:
+            raise ValueError("at most one training sub-job is supported per job")
+        if body.get("debug") and not _debug_options_enabled():
+            raise ValueError(
+                "create-job debug options are an internal-only capability; set "
+                f"{DEBUG_OPTIONS_ENV}=1 to send a request with a `debug` block"
+            )
+        resp = self._send("POST", self._prefix, json=body, retry_on=_is_connect_error)
+        return resp.json()
+
+    @staticmethod
+    def _normalize_job_status(status) -> str:
+        """Normalize full enum names (``JOB_STATE_RUNNING``) to short form (``running``).
+
+        Coerces to ``str`` first: a backend running ahead of the API layer can
+        return a JobState enum value the API layer doesn't yet know (e.g. a
+        newly added state), which is then rendered as a raw integer rather than
+        a name. Tolerating it keeps :meth:`wait_for_job` polling — an unknown
+        status is neither ``running`` nor terminal, so it stays in-progress —
+        instead of crashing on ``int.lower()``.
+        """
+        return str(status or "").lower().removeprefix("job_state_")
+
+    def _sleep_with_backoff(self, delay: float, deadline: float) -> float:
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return delay
+        time.sleep(min(delay, remaining))
+        return min(delay * self.poll_backoff_multiplier, self.poll_max_interval)
+
+    @_track_operation("wait_for_job")
+    def wait_for_job(self, job_id: str) -> dict:
+        """Poll until the job reaches ``running``. Raises on terminal states or timeout."""
+        deadline = time.monotonic() + self.poll_timeout
+        delay = self.poll_interval
+        while time.monotonic() < deadline:
+            job = self.get_job(job_id)
+            status = self._normalize_job_status(job.get("status", ""))
+            if status == "running":
+                return job
+            if status in ("failed", "done", "cancelled", "canceled"):
+                reason = job.get("reason", "")
+                raise RuntimeError(f"Job {job_id} reached terminal state '{status}': {reason}")
+            delay = self._sleep_with_backoff(delay, deadline)
+        raise TimeoutError(f"Job {job_id} did not become running within {self.poll_timeout}s")
+
+    @_track_operation("get_job")
+    def get_job(self, job_id: str) -> dict:
+        resp = self._send("GET", f"{self._prefix}/{job_id}")
+        return resp.json()
+
+    @_track_operation("list_jobs")
+    def list_jobs(self, status: str | None = None) -> list:
+        params = {}
+        if status is not None:
+            params["status"] = status
+        resp = self._send("GET", self._prefix, params=params)
+        return resp.json().get("jobs", [])
+
+    @_track_operation("cancel_job")
+    def cancel_job(self, job_id: str) -> None:
+        # The REST API uses colon-action syntax: /{jobId}:cancel
+        self._send("POST", f"{self._prefix}/{job_id}:cancel")
+
+    @_track_operation("get_capacity")
+    def get_capacity(self, hardware: Hardware | str | None = None) -> dict:
+        """Return the calling account's GPU ceiling and current usage.
+
+        Backed by the account-scoped endpoint ``/cortex-training/capacity``
+        (not under ``/{job_id}``). The account is resolved server-side from the
+        caller's session — never from a request field — so a caller can only
+        ever read its own account's capacity. Optional ``hardware``
+        (:class:`Hardware`.H200 / .B200 / .B300) is sent as a query parameter and
+        scopes the numbers to that GPU type; omitted defaults to H200.
+
+        The returned dict always carries all six fields. The server emits
+        proto3 JSON, which omits zero/false fields (an account holding nothing
+        under a zero ceiling is literally ``{}``), so we fill in the documented
+        defaults:
+
+        - ``has_reservation`` (bool): whether ``max_total_gpus`` is a
+          *guaranteed* commitment rather than a best-effort per-account cap.
+          When ``False`` the account draws on the shared pool, so
+          ``available_gpus`` is an upper bound.
+        - ``max_total_gpus`` (int): the account's ceiling — its commitment when
+          ``has_reservation``, else its per-account cap. ``-1`` means no
+          ceiling, ``0`` a real quota of zero, ``> 0`` the limit. This is the
+          canonical ceiling; prefer it over ``reserved_gpus``.
+        - ``reserved_gpus`` (int, **deprecated**): mirrors ``max_total_gpus``
+          when ``has_reservation`` is ``True``, else 0.
+        - ``in_use_gpus`` (int): GPUs the account holds — jobs in ``placing``,
+          ``initializing``, or ``running``. Queued work is *not* counted here;
+          it is in ``pending_gpus``.
+        - ``pending_gpus`` (int): GPUs requested by jobs still waiting for
+          capacity (``pending``). It claims quota, so it reduces
+          ``available_gpus``.
+        - ``available_gpus`` (int): what a new job could start with right now —
+          ceiling headroom capped by what is schedulable. ``0`` does not mean
+          blocked: a submit within the ceiling is still accepted and queued.
+
+        See ``docs/reference/rest-api.md`` section 5.4 for the authoritative
+        field list.
+        """
+        params = {}
+        if hardware is not None:
+            params["hardware"] = _hardware_wire(hardware)
+        resp = self._send(
+            "GET",
+            f"{self._prefix}/capacity",
+            **({"params": params} if params else {}),
+        )
+        body = resp.json()
+        return {
+            "has_reservation": bool(body.get("has_reservation", False)),
+            "max_total_gpus": int(body.get("max_total_gpus", 0)),
+            "reserved_gpus": int(body.get("reserved_gpus", 0)),
+            "in_use_gpus": int(body.get("in_use_gpus", 0)),
+            "pending_gpus": int(body.get("pending_gpus", 0)),
+            "available_gpus": int(body.get("available_gpus", 0)),
+        }
+
+    @_track_operation("get_experiment_run")
+    def get_experiment_run(self, job_id: str) -> dict[str, str]:
+        """Return ``{experiment_name, experiment_run_name}`` for ``job_id``.
+
+        Backed by the endpoint ``/cortex-training/{job_id}/experiment-run``.
+        """
+        resp = self._send("GET", f"{self._prefix}/{job_id}/experiment-run")
+        return resp.json()
+
+    def _query_sql_row(self, statement: str) -> list[Any]:
+        """Execute a synchronous SQL statement and return its first row."""
+        resp = self._send(
+            "POST",
+            f"{self.base_url}/api/v2/statements",
+            json={"statement": statement, "database": self.database, "schema": self.schema},
+        )
+        rows = resp.json().get("data")
+        if not isinstance(rows, list) or not rows or not isinstance(rows[0], list) or not rows[0]:
+            raise ValueError(f"SQL query returned no rows: {statement}")
+        return rows[0]
+
+    def _experiment_run_uri(self, job_id: str) -> tuple[str, str]:
+        """Return the experiment run's ``snow://`` URI and run name."""
+        run = self.get_experiment_run(job_id)
+        try:
+            experiment_name = run["experiment_name"]
+            run_name = run["experiment_run_name"]
+        except KeyError as exc:
+            raise ValueError(f"experiment-run response missing field: {exc.args[0]}") from exc
+        if not isinstance(experiment_name, str) or not experiment_name:
+            raise ValueError("experiment-run response has invalid experiment_name")
+        if not isinstance(run_name, str) or not run_name:
+            raise ValueError("experiment-run response has invalid experiment_run_name")
+        return (
+            f"snow://experiment/{experiment_name}/versions/{run_name}/",
+            run_name,
+        )
+
+    def _open_experiment_artifact_connection(self) -> Any:
+        """Open a connector session for Snowflake experiment artifact LIST/GET."""
+        config = self._artifact_connection_config
+        if config is None:
+            raise RuntimeError(
+                "experiment artifact download requires a PAT-authenticated client"
+            )
+        user, account, role = self._query_sql_row(
+            "SELECT CURRENT_USER(), CURRENT_ACCOUNT_NAME(), CURRENT_ROLE()"
+        )
+        if not isinstance(user, str) or not user:
+            raise ValueError("SQL identity response missing current user")
+        if not isinstance(account, str) or not account:
+            raise ValueError("SQL identity response missing current account")
+
+        import snowflake.connector
+
+        kwargs: dict[str, Any] = {
+            "host": config["host"],
+            "account": account,
+            "user": user,
+            "authenticator": "PROGRAMMATIC_ACCESS_TOKEN",
+            "token": config["pat"],
+            "database": self.database,
+            "schema": self.schema,
+        }
+        if isinstance(role, str) and role:
+            kwargs["role"] = role
+        return snowflake.connector.connect(**kwargs)
+
+    @staticmethod
+    def _list_experiment_artifacts(
+        connection: Any,
+        run_uri: str,
+        run_name: str,
+        artifact_root: str,
+    ) -> list[str]:
+        """List safe run-relative file paths below ``artifact_root``."""
+        root = artifact_root.strip("/") + "/"
+        _validate_artifact_relative_path(root.removesuffix("/"))
+        cursor = connection.cursor()
+        try:
+            rows = cursor.execute(
+                f"LIST {_sql_string_literal(run_uri + root)}"
+            ).fetchall()
+        finally:
+            cursor.close()
+
+        run_prefix = f"/versions/{run_name}/"
+        paths: list[str] = []
+        for row in rows:
+            name = row[0] if row else None
+            if not isinstance(name, str) or not name.startswith(run_prefix):
+                raise ValueError(
+                    f"experiment artifact LIST returned an unexpected path: {name!r}"
+                )
+            relative_path = _validate_artifact_relative_path(
+                name.removeprefix(run_prefix)
+            )
+            if not relative_path.startswith(root):
+                raise ValueError(
+                    "experiment artifact LIST returned a file outside the "
+                    f"requested root {root!r}: {relative_path!r}"
+                )
+            paths.append(relative_path)
+        return sorted(set(paths))
+
+    @staticmethod
+    def _download_experiment_artifact(
+        connection: Any,
+        run_uri: str,
+        relative_path: str,
+        target_dir: Path,
+    ) -> Path:
+        """Download one exact experiment artifact and return its local path."""
+        relative_path = _validate_artifact_relative_path(relative_path)
+        target_dir.mkdir(parents=True, exist_ok=True)
+        target_uri = target_dir.resolve().as_uri() + "/"
+        cursor = connection.cursor()
+        try:
+            cursor.execute(
+                "GET "
+                f"{_sql_string_literal(run_uri + relative_path)} "
+                f"{_sql_string_literal(target_uri)}"
+            ).fetchall()
+        finally:
+            cursor.close()
+        downloaded = target_dir / PurePosixPath(relative_path).name
+        if not downloaded.is_file():
+            raise RuntimeError(
+                f"Snowflake GET did not create expected artifact {downloaded}"
+            )
+        return downloaded
+
+    @_track_operation("fetch_execution_logs")
+    def fetch_execution_logs(self, job_id: str) -> list[dict[str, str]]:
+        """Download every log file under the job's experiment run stage.
+
+        Returns a list of
+        ``{"sub_job_id", "filename", "artifact_uri", "content"}``
+        entries, one per object found under a ``_logs/<sub_job_id>/`` subtree
+        in the run's stage. Filenames such as ``execution.jsonl`` and
+        ``server.log`` are both included; bodies are decoded as UTF-8.
+
+        ``artifact_uri`` is the logical ``snow://`` URI. Files are retrieved
+        through Snowflake's supported experiment artifact ``LIST``/``GET``
+        interface; no object-store credentials are exposed to the caller.
+        """
+        run_uri, run_name = self._experiment_run_uri(job_id)
+        connection = self._open_experiment_artifact_connection()
+        try:
+            paths: list[str] = []
+            for artifact_root in ("_logs", "checkpoints/_logs"):
+                paths.extend(
+                    self._list_experiment_artifacts(
+                        connection, run_uri, run_name, artifact_root
+                    )
+                )
+
+            results: list[dict[str, str]] = []
+            with tempfile.TemporaryDirectory(prefix="cortex-training-execution-logs-") as temp:
+                temp_root = Path(temp)
+                for relative_path in sorted(set(paths)):
+                    if relative_path.startswith("_logs/"):
+                        after_logs = relative_path.removeprefix("_logs/")
+                    elif relative_path.startswith("checkpoints/_logs/"):
+                        after_logs = relative_path.removeprefix("checkpoints/_logs/")
+                    else:
+                        continue
+                    sub_job_id, separator, filename = after_logs.partition("/")
+                    if (
+                        not separator
+                        or not filename
+                        or not (
+                            sub_job_id == job_id
+                            or sub_job_id.startswith(f"{job_id}:")
+                        )
+                        or sub_job_id in {".", ".."}
+                        or "\\" in sub_job_id
+                    ):
+                        continue
+                    downloaded = self._download_experiment_artifact(
+                        connection,
+                        run_uri,
+                        relative_path,
+                        temp_root,
+                    )
+                    artifact_uri = run_uri + relative_path
+                    try:
+                        content = downloaded.read_text(encoding="utf-8")
+                    finally:
+                        downloaded.unlink(missing_ok=True)
+                    results.append(
+                        {
+                            "sub_job_id": sub_job_id,
+                            "filename": filename,
+                            "artifact_uri": artifact_uri,
+                            "content": content,
+                        }
+                    )
+            return results
+        finally:
+            connection.close()
+
+    def _download_gzip_artifacts(
+        self,
+        job_id: str,
+        output_dir: str | os.PathLike[str],
+        *,
+        artifact_name: str,
+        chunk_pattern: re.Pattern[str],
+        destination_name: str,
+        temporary_prefix: str,
+    ) -> list[dict[str, Any]]:
+        """Reconstruct staged gzip chunks into one atomic file per sub-job."""
+        run_uri, run_name = self._experiment_run_uri(job_id)
+        connection = self._open_experiment_artifact_connection()
+        root = Path(output_dir).expanduser()
+        downloads: dict[str, dict[str, Any]] = {}
+        artifact_roots = (f"_{artifact_name}", f"checkpoints/_{artifact_name}")
+        try:
+            paths_by_root = {
+                artifact_root: self._list_experiment_artifacts(
+                    connection, run_uri, run_name, artifact_root
+                )
+                for artifact_root in artifact_roots
+            }
+            chunks_by_sub_job: dict[str, dict[str, Any]] = {}
+            for artifact_root, paths in paths_by_root.items():
+                prefix = artifact_root + "/"
+                for relative_path in paths:
+                    after_root = relative_path.removeprefix(prefix)
+                    sub_job_id, separator, filename = after_root.partition("/")
+                    if (
+                        not separator
+                        or sub_job_id in {"", ".", ".."}
+                        or not (
+                            sub_job_id == job_id
+                            or sub_job_id.startswith(f"{job_id}:")
+                        )
+                        or "/" in filename
+                        or "\\" in sub_job_id
+                        or chunk_pattern.fullmatch(filename) is None
+                    ):
+                        continue
+                    chunks = chunks_by_sub_job.get(sub_job_id)
+                    if chunks is None:
+                        chunks = {"artifact_root": artifact_root, "paths": []}
+                        chunks_by_sub_job[sub_job_id] = chunks
+                    elif chunks["artifact_root"] != artifact_root:
+                        raise ValueError(
+                            f"{artifact_name} chunks for {sub_job_id!r} appear "
+                            "under multiple artifact roots"
+                        )
+                    chunks["paths"].append(relative_path)
+
+            with tempfile.TemporaryDirectory(
+                prefix=f"cortex-training-{artifact_name}-artifacts-"
+            ) as temp:
+                artifact_temp = Path(temp)
+                for sub_job_id in sorted(chunks_by_sub_job):
+                    chunk_paths = sorted(chunks_by_sub_job[sub_job_id]["paths"])
+                    destination = root / sub_job_id / destination_name
+                    destination.parent.mkdir(parents=True, exist_ok=True)
+                    output = tempfile.NamedTemporaryFile(
+                        mode="wb",
+                        dir=str(destination.parent),
+                        prefix=temporary_prefix,
+                        suffix=".tmp",
+                        delete=False,
+                    )
+                    state = {
+                        "destination": destination,
+                        "temporary": Path(output.name),
+                        "output": output,
+                        "chunk_count": 0,
+                        "first_path": chunk_paths[0],
+                        "last_path": chunk_paths[-1],
+                    }
+                    downloads[sub_job_id] = state
+                    for relative_path in chunk_paths:
+                        downloaded = self._download_experiment_artifact(
+                            connection,
+                            run_uri,
+                            relative_path,
+                            artifact_temp,
+                        )
+                        try:
+                            with downloaded.open("rb") as compressed:
+                                with gzip.GzipFile(fileobj=compressed, mode="rb") as source:
+                                    shutil.copyfileobj(
+                                        source,
+                                        output,
+                                        length=_STREAM_COPY_BUFFER_BYTES,
+                                    )
+                        finally:
+                            downloaded.unlink(missing_ok=True)
+                        state["chunk_count"] += 1
+
+            results: list[dict[str, Any]] = []
+            for sub_job_id in sorted(downloads):
+                state = downloads[sub_job_id]
+                output = state["output"]
+                output.flush()
+                os.fsync(output.fileno())
+                output.close()
+                os.replace(state["temporary"], state["destination"])
+                results.append(
+                    {
+                        "sub_job_id": sub_job_id,
+                        "filename": state["destination"].name,
+                        "saved_path": str(state["destination"]),
+                        "chunk_count": state["chunk_count"],
+                        "first_artifact_uri": run_uri + state["first_path"],
+                        "last_artifact_uri": run_uri + state["last_path"],
+                    }
+                )
+            return results
+        except BaseException:
+            for state in downloads.values():
+                output = state["output"]
+                if not output.closed:
+                    output.close()
+                state["temporary"].unlink(missing_ok=True)
+            raise
+        finally:
+            connection.close()
+
+    def download_stdout_logs(
+        self, job_id: str, output_dir: str | os.PathLike[str]
+    ) -> list[dict[str, Any]]:
+        """Reconstruct sealed console artifacts into one file per sub-job."""
+        return self._download_gzip_artifacts(
+            job_id,
+            output_dir,
+            artifact_name="stdout",
+            chunk_pattern=_STDOUT_CHUNK_RE,
+            destination_name="stdout.log",
+            temporary_prefix=".stdout-",
+        )
+
+    def download_metrics(
+        self, job_id: str, output_dir: str | os.PathLike[str]
+    ) -> list[dict[str, Any]]:
+        """Reconstruct staged GPU metric chunks into one JSONL file per sub-job."""
+        return self._download_gzip_artifacts(
+            job_id,
+            output_dir,
+            artifact_name="metrics",
+            chunk_pattern=_GPU_METRICS_CHUNK_RE,
+            destination_name="gpu.jsonl",
+            temporary_prefix=".gpu-metrics-",
+        )
+
+    # ─── Data-plane async operations ─────────────────────────────────────
+
+    # Maximum payload size for forward-backward requests.
+    # Payloads beyond this limit will fail fast with a message.
+    _MAX_FWD_BWD_BYTES = 60 * 1024 * 1024  # 60 MB
+
+    # Maximum payload size for generate / generate_stream request bodies.
+    # The encoded JSON body must stay under this cap; oversized requests fail
+    # fast client-side with a clear message.
+    _MAX_GENERATE_BYTES = 60 * 1024 * 1024  # 60 MB
+
+    def _check_generate_payload(self, op: str, body: bytes) -> None:
+        if len(body) > self._MAX_GENERATE_BYTES:
+            size_mb = len(body) / (1024 * 1024)
+            raise ValueError(
+                f"{op} payload is {size_mb:.1f} MB, which exceeds the "
+                "maximum allowed size of "
+                f"{self._MAX_GENERATE_BYTES / (1024 * 1024):.0f} MB. "
+                "Reduce the number of prompts or shorten each prompt to "
+                "bring the serialized request under the limit."
+            )
+
+    def _resolve_sampling_max_seq_len(self, job_id: str) -> int | None:
+        """Return the sampling sub-job's ``max_seq_len``, or ``None`` if unknown.
+
+        The value is read from the job's sampling ``inference_config`` (the same
+        number the backend launches vLLM with as ``max_model_len``) and cached
+        per ``job_id``. A successful lookup — including a definitive "no sampling
+        sub-job" miss — is cached; a transient ``get_job`` failure is not, so a
+        later call retries rather than disabling validation permanently.
+        """
+        if job_id in self._sampling_max_seq_len:
+            return self._sampling_max_seq_len[job_id]
+        try:
+            job = self.get_job(job_id)
+        except Exception:
+            # Don't let a transient lookup failure break generate, and don't
+            # cache the miss — the real request will surface any hard error.
+            return None
+        max_seq_len: int | None = None
+        for sub in job.get("sub_jobs", []) or []:
+            if not isinstance(sub, dict):
+                continue
+            cfg = sub.get("inference_config")
+            if not isinstance(cfg, dict):
+                continue
+            raw = cfg.get("max_seq_len")
+            # Proto numbers can arrive as floats (e.g. 2048.0); coerce to int.
+            if isinstance(raw, (int, float)) and raw > 0:
+                max_seq_len = int(raw)
+                break
+        self._sampling_max_seq_len[job_id] = max_seq_len
+        return max_seq_len
+
+    def _check_prompt_lengths(self, op: str, job_id: str, prompts: list[str | list[int]]) -> None:
+        """Fail fast when a pre-tokenized prompt cannot fit the sampling window.
+
+        Mirrors vLLM's decoder-prompt check (``vllm/v1/engine/input_processor``):
+        a generation prompt must be strictly shorter than ``max_model_len`` so
+        at least one output token fits, i.e. it is rejected when
+        ``len(prompt_token_ids) >= max_seq_len``. Only pre-tokenized prompts
+        (``list[int]``) are validated here, since their token count is exact;
+        string prompts would require replicating the server tokenizer and are
+        left to fail server-side.
+        """
+        if not any(isinstance(p, list) for p in prompts):
+            return
+        max_seq_len = self._resolve_sampling_max_seq_len(job_id)
+        if max_seq_len is None:
+            return
+        for i, prompt in enumerate(prompts):
+            if not isinstance(prompt, list):
+                continue
+            n = len(prompt)
+            if n >= max_seq_len:
+                raise ValueError(
+                    f"{op} prompt at index {i} has {n} tokens, which does not "
+                    f"fit the sampling job's max_seq_len of {max_seq_len}: a "
+                    "prompt must be shorter than max_seq_len to leave room for "
+                    "at least one generated token (vLLM rejects "
+                    "len(prompt) >= max_model_len). Shorten the prompt to at "
+                    f"most {max_seq_len - 1} tokens, or recreate the job with a "
+                    "larger max_seq_len."
+                )
+
+    def _post_octet_request_chunks(
+        self,
+        *,
+        job_id: str,
+        path_suffix: str,
+        operation: str,
+        frame: bytes,
+        max_bytes: int,
+        force_chunk: bool = False,
+        allow_group_restart: bool = False,
+        debug_context: dict[str, Any] | None = None,
+    ) -> dict:
+        chunks = wire.encode_byte_chunks(
+            frame,
+            kind="request",
+            operation=operation,
+            max_bytes=max_bytes,
+            force_chunk=force_chunk,
+        )
+        group_metadata = wire.read_byte_chunk_metadata(chunks[0])
+        chunk_group_id = (
+            str(group_metadata["chunk_group_id"])
+            if group_metadata is not None and group_metadata.get("chunk_group_id")
+            else None
+        )
+        group_restarts = 0
+        idx = 0
+        final_body: dict | None = None
+        while idx < len(chunks):
+            chunk = chunks[idx]
+            try:
+                resp = self._send(
+                    "POST",
+                    f"{self._prefix}/{job_id}/{path_suffix}",
+                    data=chunk,
+                    headers={"Content-Type": "application/octet-stream"},
+                    retry_on=(_is_chunk_post_transient if allow_group_restart else _is_transient),
+                    debug_context=debug_context,
+                )
+            except requests.exceptions.HTTPError as exc:
+                detail = _chunk_group_error_detail(exc.response)
+                if detail is None:
+                    raise
+                code = detail.get("code")
+                if allow_group_restart and code == _CHUNK_GROUP_RESTART_REQUIRED and group_restarts < 1:
+                    group_restarts += 1
+                    idx = 0
+                    final_body = None
+                    logger.warning(
+                        "%s chunk group %s lost server state; replaying once from chunk 0",
+                        operation,
+                        chunk_group_id,
+                    )
+                    continue
+                error_type = (
+                    ChunkGroupRestartError if code == _CHUNK_GROUP_RESTART_REQUIRED else ChunkGroupConflictError
+                )
+                raise error_type(
+                    str(detail.get("message") or code),
+                    response=exc.response,
+                    detail=detail,
+                ) from exc
+
+            body = resp.json()
+            if idx < len(chunks) - 1:
+                if isinstance(body, dict) and body.get("request_id"):
+                    raise RuntimeError(f"{operation} chunk {idx} unexpectedly returned request_id")
+            else:
+                final_body = body
+            idx += 1
+        if final_body is None:
+            raise RuntimeError(f"{operation} produced no request body")
+        return final_body
+
+    @_track_operation("forward_backward")
+    def forward_backward(self, job_id: str, data: bytes) -> str:
+        """Submit a forward+backward pass. Returns request_id."""
+        self._fwd_bwd_send_count += 1
+        call = self._fwd_bwd_send_count
+        debug_context = {
+            "phase": "submit",
+            "call": call,
+            "job_id": job_id,
+            "payload_bytes": len(data),
+        }
+        debug_label = self._debug_context_label(debug_context)
+        logger.debug("%s payload total bytes: %d", debug_label, len(data))
+        body = self._post_octet_request_chunks(
+            job_id=job_id,
+            path_suffix="forward-backward",
+            operation="fwd-bwd",
+            frame=data,
+            max_bytes=self._MAX_FWD_BWD_BYTES,
+            force_chunk=True,
+            allow_group_restart=True,
+            debug_context=debug_context,
+        )
+        request_id = body["request_id"]
+        request_debug_context = {
+            "phase": "poll",
+            "call": call,
+            "job_id": job_id,
+            "request_id": request_id,
+            "payload_bytes": len(data),
+        }
+        self._fwd_bwd_request_debug[request_id] = request_debug_context
+        logger.debug(
+            "%s submitted request_id=%s body=%s",
+            debug_label,
+            request_id,
+            self._debug_json_summary(body),
+        )
+        return request_id
+
+    @_track_operation("generate")
+    def generate(
+        self,
+        job_id: str,
+        prompts: list[str | list[int]],
+        sampling_params: dict | list[dict | None] | None = None,
+        routing_key: str | list[str | None] | None = None,
+        strict: bool | None = None,
+    ) -> str:
+        """Submit a generate request to a sampling sub-job. Returns request_id.
+
+        The typed fields are serialized as a DSSST1 safetensors frame and sent
+        as ``application/octet-stream``. Oversized frames are byte-chunked
+        across multiple POSTs; the final POST returns the pollable request id.
+        Poll with :meth:`poll_request`.
+
+        Prompt-length validation is asymmetric by input type:
+
+        * **Pre-tokenized prompts** (``list[int]``) are checked client-side
+          against the sampling sub-job's ``max_seq_len`` and fail fast with a
+          ``ValueError`` when ``len(prompt) >= max_seq_len`` — the token count
+          is exact, so this never rejects a prompt vLLM would accept.
+        * **String prompts** are *not* validated here. Counting their tokens
+          would require replicating the server tokenizer (including special
+          tokens), so an over-long string instead fails server-side in vLLM
+          with a ``maximum model length`` error rather than client-side.
+        """
+        payload: dict = {"prompts": prompts}
+        if sampling_params is not None:
+            payload["sampling_params"] = sampling_params
+        if routing_key is not None:
+            payload["routing_key"] = routing_key
+        if strict is not None:
+            payload["strict"] = strict
+        self._check_prompt_lengths("generate", job_id, prompts)
+        body = wire.dumps(
+            payload,
+            metadata={"response_options": {"format": "dssst1", "delivery": "chunked"}},
+        )
+        response = self._post_octet_request_chunks(
+            job_id=job_id,
+            path_suffix="generate",
+            operation="generate",
+            frame=body,
+            max_bytes=self._MAX_GENERATE_BYTES,
+        )
+        request_id = response["request_id"]
+        self._generate_request_ids.add(request_id)
+        return request_id
+
+    @_track_operation("generate_stream")
+    def generate_stream(
+        self,
+        job_id: str,
+        prompts: list[str | list[int]],
+        sampling_params: dict | list[dict | None] | None = None,
+        routing_key: str | list[str | None] | None = None,
+        strict: bool | None = None,
+    ) -> dict:
+        """Start a streaming generate request. Returns the response body.
+
+        Progress is read with :meth:`get_request_status` using the returned
+        ``request_id``. Cancellation uses :meth:`cancel_request`.
+
+        Prompt-length validation matches :meth:`generate`: pre-tokenized
+        (``list[int]``) prompts are checked client-side against the sampling
+        sub-job's ``max_seq_len`` and fail fast, while over-long string
+        prompts are left to fail server-side in vLLM (see :meth:`generate`).
+        """
+        payload: dict = {"prompts": prompts}
+        if sampling_params is not None:
+            payload["sampling_params"] = sampling_params
+        if routing_key is not None:
+            payload["routing_key"] = routing_key
+        if strict is not None:
+            payload["strict"] = strict
+        self._check_prompt_lengths("generate_stream", job_id, prompts)
+        body = json.dumps(payload).encode("utf-8")
+        self._check_generate_payload("generate_stream", body)
+        resp = self._send(
+            "POST",
+            f"{self._prefix}/{job_id}/generate-stream",
+            data=body,
+            headers={"Content-Type": "application/octet-stream"},
+        )
+        return resp.json()
+
+    @_track_operation("step")
+    def step(self, job_id: str, learning_rate: float | None = None) -> str:
+        """Submit an optimizer step. Returns request_id."""
+        body: dict = {}
+        if learning_rate is not None:
+            body["learning_rate"] = learning_rate
+        resp = self._send("POST", f"{self._prefix}/{job_id}/step", json=body)
+        return resp.json()["request_id"]
+
+    @_track_operation("save")
+    def save(
+        self,
+        job_id: str,
+        checkpoint_id: str | None = None,
+        checkpoint_type: str | None = None,
+    ) -> str:
+        """Submit a checkpoint save. Returns request_id."""
+        body: dict = {}
+        if checkpoint_id is not None:
+            body["checkpoint_id"] = checkpoint_id
+        if checkpoint_type is not None:
+            normalized_type = checkpoint_type.lower()
+            if normalized_type not in ("resumable", "weights-only"):
+                raise ValueError("checkpoint_type must be 'resumable' or 'weights-only'")
+            body["checkpoint_type"] = normalized_type
+        resp = self._send("POST", f"{self._prefix}/{job_id}/save", json=body)
+        return resp.json()["request_id"]
+
+    @_track_operation("load")
+    def load(
+        self,
+        job_id: str,
+        checkpoint_id: str,
+        source_job_id: str | None = None,
+        *,
+        target_sub_job_id: str | None = None,
+    ) -> str:
+        """Submit a checkpoint load into an existing job. Returns request_id.
+
+        Args:
+            job_id: The job to load the checkpoint into.
+            checkpoint_id: Checkpoint identifier (from a prior save).
+            source_job_id: Optional. Load from another job's checkpoint store.
+            target_sub_job_id: Optional. Route the load to a specific training
+                sub-job. Format: ``"{job_id}:training:{index}"``. Omit to use
+                the session's default training sub-job.
+
+        Returns:
+            request_id for polling via ``poll_request()``.
+
+        **When to use target_sub_job_id:**
+
+        A job has at most one training sub-job, so omitting this routes to that
+        sub-job. Pass it when you want explicit routing rather than relying on
+        the server's default resolution.
+
+        **Discovering sub-job IDs:**
+
+        .. code-block:: python
+
+            job = client.get_job(job_id)
+            for sj in job["sub_jobs"]:
+                if sj["job_type"] == "training":
+                    sub_job_id = sj["sub_job_id"]
+                    n_gpus = sj["training_config"]["n_gpus"]
+                    print(f"Training sub-job {sub_job_id}: {n_gpus} GPUs")
+
+        **DP size compatibility:**
+
+        If the target sub-job has a different DP size (``n_gpus``) than the
+        checkpoint was saved from, you **must** have created the job with
+        ``load_optimizer_states=False`` in the ``SubJobConfig``. The optimizer
+        states are DP-sharded and cannot be resized. This setting is configured
+        at job creation time and cannot be changed at runtime.
+
+        **Server validation:**
+
+        The server validates that ``target_sub_job_id`` belongs to this session
+        (returns 400 if not) and is a training sub-job (returns 501 for
+        non-training sub-jobs like sampling).
+        """
+        body: dict = {"checkpoint_id": checkpoint_id}
+        if source_job_id is not None:
+            body["source_job_id"] = source_job_id
+        if target_sub_job_id is not None:
+            body["target_sub_job_id"] = target_sub_job_id
+        resp = self._send("POST", f"{self._prefix}/{job_id}/load", json=body)
+        return resp.json()["request_id"]
+
+    def _operation(
+        self,
+        job_id: str,
+        operation_type: str,
+        *,
+        payload: dict | bytes | bytearray | memoryview | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Submit a generic data-plane operation to the routed sub-job."""
+        body: dict = {"operation_type": operation_type}
+        if sub_job_id is not None:
+            body["sub_job_id"] = sub_job_id
+        if sub_job_type is not None:
+            body["sub_job_type"] = sub_job_type
+        if isinstance(payload, (bytes, bytearray, memoryview)):
+            payload = {
+                "payload_b64": base64.b64encode(bytes(payload)).decode("ascii"),
+                "content_type": "application/octet-stream",
+            }
+        if payload is not None:
+            body["payload"] = payload
+        print(f"***** POST operation {self._prefix}/{job_id}/operation {body=}")
+        resp = self._send("POST", f"{self._prefix}/{job_id}/operation", json=body)
+        return resp.json()
+
+    @_track_operation("forward")
+    def forward(
+        self,
+        job_id: str,
+        payload: dict | bytes | bytearray | memoryview | None = None,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Submit a generic forward operation. Returns the response body."""
+        if isinstance(payload, (bytes, bytearray, memoryview)) and len(payload) > self._MAX_FWD_BWD_BYTES:
+            size_mb = len(payload) / (1024 * 1024)
+            raise ValueError(
+                f"forward payload is {size_mb:.1f} MB, which exceeds the "
+                "maximum allowed size of "
+                f"{self._MAX_FWD_BWD_BYTES / (1024 * 1024):.0f} MB. "
+                "Reduce the batch size or sequence length in your training "
+                "configuration to bring the serialized input batch under "
+                "the limit."
+            )
+        return self._operation(
+            job_id,
+            "forward",
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    def fwd(
+        self,
+        job_id: str,
+        payload: dict | bytes | bytearray | memoryview | None = None,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Compatibility alias for :meth:`forward`."""
+        return self.forward(
+            job_id,
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    def fwd_no_grad(
+        self,
+        job_id: str,
+        payload: dict | bytes | bytearray | memoryview | None = None,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Compatibility alias for :meth:`forward`."""
+        return self.forward(
+            job_id,
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    @_track_operation("weight_sync")
+    def weight_sync(
+        self,
+        job_id: str,
+        source_sub_job_id: str,
+        target_sub_job_ids: list[str],
+        *,
+        weight_format: str | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> str:
+        """Sync model weights from a training sub-job into one or more sampling
+        sub-jobs (used in RL training loops). Returns request_id.
+
+        Multi-sub-job sessions require the operation envelope to include a
+        routing hint; by default the operation is routed through the source
+        training sub-job.
+
+        Args:
+            job_id: The job owning both sub-jobs.
+            source_sub_job_id: Training sub-job to read weights from, e.g.
+                ``f"{job_id}:training:0"``.
+            target_sub_job_ids: Sampling sub-jobs to update, e.g.
+                ``[f"{job_id}:sampling:0"]``.
+            weight_format: ``"vllm"`` (server default) or ``"hf"`` for full
+                weights, or ``"lora"`` to broadcast only the adapter tensors.
+
+        Example::
+
+            request_id = client.weight_sync(
+                job_id,
+                source_sub_job_id=f"{job_id}:training:0",
+                target_sub_job_ids=[f"{job_id}:sampling:0"],
+                weight_format="lora",
+            )
+            client.poll_request(job_id, request_id)
+        """
+        body = {
+            "source_sub_job_id": source_sub_job_id,  # training
+            "target_sub_job_ids": list(target_sub_job_ids),  # sampling
+        }
+        if weight_format is not None:
+            # "lora" broadcasts only the trained adapter tensors.
+            body["weight_format"] = weight_format
+        return self._operation(
+            job_id,
+            "weight-sync",
+            payload=body,
+            sub_job_id=sub_job_id or source_sub_job_id,  # training
+            sub_job_type=(sub_job_type if sub_job_type is not None else "training"),
+        )["request_id"]
+
+    @_track_operation("bootstrap_router_replay")
+    def bootstrap_router_replay(
+        self,
+        job_id: str,
+        source_sub_job_id: str,
+        target_sub_job_id: str,
+        *,
+        max_cache_bytes: int | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Bootstrap router replay from one sub-job to another."""
+        payload: dict = {
+            "source_sub_job_id": source_sub_job_id,
+            "target_sub_job_id": target_sub_job_id,
+        }
+        if max_cache_bytes is not None:
+            payload["max_cache_bytes"] = max_cache_bytes
+        return self._operation(
+            job_id,
+            "bootstrap-router-replay",
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    @_track_operation("router_replay_discard")
+    def router_replay_discard(
+        self,
+        job_id: str,
+        sample_ids: list[str] | None = None,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+        extra: dict | None = None,
+    ) -> dict:
+        """Discard router replay samples on a routed sub-job."""
+        payload = dict(extra or {})
+        payload.setdefault("sample_ids", list(sample_ids or []))
+        return self._operation(
+            job_id,
+            "router-replay-discard",
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    @_track_operation("reset_prefix_cache")
+    def reset_prefix_cache(
+        self,
+        job_id: str,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = "sampling",
+        drain: bool = True,
+        timeout_s: float = 60.0,
+        retry_interval_s: float = 0.1,
+        extra: dict | None = None,
+    ) -> dict:
+        """Reset the sampling prefix cache for a Cortex Training job.
+
+        Operation-specific fields are carried inside the opaque operation
+        payload. ``extra`` allows new reset options to flow through without a
+        cortex-training release.
+        """
+        payload = dict(extra or {})
+        payload.setdefault("drain", drain)
+        payload.setdefault("timeout_s", timeout_s)
+        payload.setdefault("retry_interval_s", retry_interval_s)
+        return self._operation(
+            job_id,
+            "reset-prefix-cache",
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    # ─── Async request polling ────────────────────────────────────────────
+
+    @staticmethod
+    def _normalize_request_status(status: str) -> str:
+        """Normalize full enum names (``REQUEST_STATE_DONE``) to short form (``done``)."""
+        return status.lower().removeprefix("request_state_")
+
+    @staticmethod
+    def _decode_result_payload(result: dict) -> dict | None:
+        if not isinstance(result, dict):
+            return None
+        if result.get("wire_format") != wire.WIRE_FORMAT_VERSION:
+            return None
+        if result.get("encoding") != "base64":
+            raise RuntimeError("DSSST1 result payload must use base64 encoding")
+        raw = result.get("payload_b64")
+        if not isinstance(raw, str) or not raw:
+            raise RuntimeError("DSSST1 result payload missing payload_b64")
+        payload = base64.b64decode(raw)
+        decoded = wire.loads(payload)
+        if not isinstance(decoded, dict):
+            raise RuntimeError("DSSST1 result payload did not decode to a dict")
+        return decoded
+
+    @staticmethod
+    def _decode_result_chunk_event(event: dict) -> bytes | None:
+        if not isinstance(event, dict) or event.get("type") != "result_chunk":
+            return None
+        raw = event.get("payload_b64")
+        if not isinstance(raw, str) or not raw:
+            raise RuntimeError("result_chunk event missing payload_b64")
+        payload = base64.b64decode(raw)
+        expected = event.get("payload_sha256")
+        if expected is not None and hashlib.sha256(payload).hexdigest() != expected:
+            raise RuntimeError("result_chunk payload_sha256 mismatch")
+        return payload
+
+    @staticmethod
+    def _restore_generate_result_lists(value: Any, torch_module: Any | None = None) -> Any:
+        torch_module = _load_torch() if torch_module is None else torch_module
+        if torch_module.is_tensor(value):
+            return value.cpu().tolist()
+        if isinstance(value, dict):
+            return {
+                key: CortexTrainingClient._restore_generate_result_lists(item, torch_module)
+                for key, item in value.items()
+            }
+        if isinstance(value, list):
+            return [CortexTrainingClient._restore_generate_result_lists(item, torch_module) for item in value]
+        if isinstance(value, tuple):
+            return tuple(CortexTrainingClient._restore_generate_result_lists(item, torch_module) for item in value)
+        return value
+
+    def _normalize_generate_result_if_needed(self, request_id: str, result: dict) -> dict:
+        if request_id not in self._generate_request_ids:
+            return result
+        self._generate_request_ids.discard(request_id)
+        if isinstance(result, dict) and "results" in result:
+            result = dict(result)
+            result["results"] = self._restore_generate_result_lists(result["results"])
+        return result
+
+    @staticmethod
+    def _decode_stream_result_event(event: Any) -> Any:
+        if not isinstance(event, dict) or event.get("type") != "result":
+            return event
+        result = event.get("result")
+        if not isinstance(result, dict):
+            return event
+        decoded = CortexTrainingClient._decode_result_payload(result)
+        if decoded is None:
+            return event
+        event = dict(event)
+        event["result"] = CortexTrainingClient._restore_generate_result_lists(decoded)
+        return event
+
+    @staticmethod
+    def _decode_stream_result_events(status: dict) -> dict:
+        events = status.get("events")
+        if not isinstance(events, list):
+            return status
+        decoded = [CortexTrainingClient._decode_stream_result_event(event) for event in events]
+        if decoded == events:
+            return status
+        status = dict(status)
+        status["events"] = decoded
+        return status
+
+    @_track_operation("poll_request")
+    def poll_request(self, job_id: str, request_id: str) -> dict:
+        """Poll until the async request completes. Returns the result dict."""
+        deadline = time.monotonic() + self.poll_timeout
+        delay = self.poll_interval
+        debug_context = self._fwd_bwd_request_debug.get(request_id)
+        debug_label = self._debug_context_label(debug_context) if debug_context is not None else None
+        result_chunks: list[bytes] = []
+        cursor: str | None = None
+        while time.monotonic() < deadline:
+            status = self.get_request_status(job_id, request_id, cursor=cursor)
+            state = self._normalize_request_status(status.get("status", ""))
+            events = status.get("events") or []
+            received_chunk = False
+            for event in events:
+                chunk = self._decode_result_chunk_event(event)
+                if chunk is not None:
+                    result_chunks.append(chunk)
+                    received_chunk = True
+            next_cursor = status.get("next_cursor")
+            if isinstance(next_cursor, str) and next_cursor:
+                cursor = next_cursor
+                # Drain available result chunks without backing off.
+                continue
+            if state in ("completed", "done", "succeeded"):
+                if result_chunks:
+                    result = wire.decode_result_chunks(result_chunks)
+                else:
+                    result = status.get("result") or {}
+                    decoded = self._decode_result_payload(result)
+                    if decoded is not None:
+                        result = decoded
+                result = self._normalize_generate_result_if_needed(request_id, result)
+                if debug_label is not None:
+                    logger.debug(
+                        "%s completed state=%s result=%s",
+                        debug_label,
+                        state,
+                        self._debug_json_summary(result),
+                    )
+                    self._fwd_bwd_request_debug.pop(request_id, None)
+                return result
+            if state in ("failed", "cancelled", "canceled"):
+                error = status.get("error", "")
+                if debug_label is not None:
+                    logger.debug(
+                        "%s failed state=%s error=%s body=%s",
+                        debug_label,
+                        state,
+                        error,
+                        self._debug_json_summary(status),
+                    )
+                    self._fwd_bwd_request_debug.pop(request_id, None)
+                self._generate_request_ids.discard(request_id)
+                raise RuntimeError(f"Request {request_id} ended with state '{state}': {error}")
+            if received_chunk:
+                # Defensive: if the server returns a chunk without next_cursor
+                # before terminal status, poll again immediately rather than
+                # sleeping.
+                continue
+            delay = self._sleep_with_backoff(delay, deadline)
+        if debug_label is not None:
+            logger.debug("%s timed out after %ss", debug_label, self.poll_timeout)
+            self._fwd_bwd_request_debug.pop(request_id, None)
+        self._generate_request_ids.discard(request_id)
+        raise TimeoutError(f"Request {request_id} did not complete within {self.poll_timeout}s")
+
+    @_track_operation("get_request_status")
+    def get_request_status(
+        self,
+        job_id: str,
+        request_id: str,
+        *,
+        max_events: int | None = None,
+        cursor: str | None = None,
+    ) -> dict:
+        params = {}
+        if max_events is not None:
+            params["max_events"] = max_events
+        if cursor:
+            params["cursor"] = cursor
+        debug_context = self._fwd_bwd_request_debug.get(request_id)
+        resp = self._send(
+            "GET",
+            f"{self._prefix}/{job_id}/requests/{request_id}",
+            params=params or None,
+            debug_context=debug_context,
+        )
+        return self._decode_stream_result_events(resp.json())
+
+    @_track_operation("cancel_request")
+    def cancel_request(
+        self,
+        job_id: str,
+        request_id: str,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        return self._operation(
+            job_id,
+            "cancel-request",
+            payload={"request_id": request_id},
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    # ─── Logs (read-only) ────────────────────────────────────────────────
+
+    def tail_logs(
+        self,
+        job_id: str,
+        *,
+        cursor: str | None = None,
+        max_lines: int | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Fetch one page of a sub-job's logs from ``cursor`` (read-only).
+
+        The target is the sub-job's zone-manager (Ray head) pod, whose stdout
+        already includes worker output (Ray ``log_to_driver``); identify it via
+        ``sub_job_id`` (or ``sub_job_type`` when unambiguous). Returns
+        ``{"entries": [...], "next_cursor": str, "eof": bool}``. Re-call with
+        ``next_cursor`` to continue; an empty ``cursor`` reads from the start.
+        """
+        payload: dict = {}
+        if cursor is not None:
+            payload["cursor"] = cursor
+        if max_lines is not None:
+            payload["max_lines"] = max_lines
+        return self._operation(
+            job_id,
+            "tail-logs",
+            payload=payload,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    def _stream_pages(
+        self,
+        fetch,
+        entries_key: str,
+        *,
+        follow: bool,
+        cursor: str | None = None,
+        poll_interval: float | None = None,
+    ) -> Iterator[dict]:
+        """Drive a cursor-paged tail: call ``fetch(cursor)`` for a page, yield
+        its entries (under ``entries_key``), advance the cursor, and poll with
+        backoff. With ``follow`` False, stop once a page reports ``eof``.
+        """
+        base_delay = poll_interval or self.poll_interval
+        delay = base_delay
+        while True:
+            page = fetch(cursor)
+            cursor = page.get("next_cursor", cursor)
+            entries = page.get(entries_key) or []
+            for entry in entries:
+                yield entry
+            if entries:
+                delay = base_delay  # made progress; poll again promptly
+                continue
+            if page.get("eof", False) and not follow:
+                return
+            time.sleep(delay)
+            delay = min(delay * self.poll_backoff_multiplier, self.poll_max_interval)
+
+    def stream_logs(
+        self,
+        job_id: str,
+        *,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+        cursor: str | None = None,
+        max_lines: int | None = None,
+        poll_interval: float | None = None,
+        follow: bool = True,
+    ) -> Iterator[dict]:
+        """Yield a sub-job's log entries, polling with backoff.
+
+        Drains the log then, with ``follow=True`` (the default), keeps polling
+        for new entries until the caller stops iterating — suitable for a live
+        tail. With ``follow=False`` it stops once the log reports ``eof``.
+        """
+        return self._stream_pages(
+            lambda cur: self.tail_logs(
+                job_id,
+                cursor=cur,
+                max_lines=max_lines,
+                sub_job_id=sub_job_id,
+                sub_job_type=sub_job_type,
+            ),
+            "entries",
+            follow=follow,
+            cursor=cursor,
+            poll_interval=poll_interval,
+        )
+
+    # ─── Zone scheduling events (read-only) ───────────────────────────────
+
+    def tail_events(
+        self,
+        job_id: str,
+        *,
+        cursor: str | None = None,
+        max_events: int | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+    ) -> dict:
+        """Fetch one page of zone scheduling events for the session.
+
+        Returns ``{"events": [...], "next_cursor": str, "eof": bool}``. Served
+        by the zone dispatcher itself, not routed to a Zone Manager.
+
+        Note: the server does not currently accept this operation; see
+        ``docs/reference/rest-api.md`` section 7.9.
+        """
+        payload: dict = {}
+        if cursor is not None:
+            payload["cursor"] = cursor
+        if max_events is not None:
+            payload["max_events"] = max_events
+        return self._operation(
+            job_id,
+            "zmd-events",
+            payload=payload or None,
+            sub_job_id=sub_job_id,
+            sub_job_type=sub_job_type,
+        )
+
+    def stream_events(
+        self,
+        job_id: str,
+        *,
+        cursor: str | None = None,
+        max_events: int | None = None,
+        sub_job_id: str | None = None,
+        sub_job_type: str | None = None,
+        poll_interval: float | None = None,
+        follow: bool = True,
+    ) -> Iterator[dict]:
+        """Yield zone scheduling events for the session, polling with backoff."""
+        return self._stream_pages(
+            lambda cur: self.tail_events(
+                job_id,
+                cursor=cur,
+                max_events=max_events,
+                sub_job_id=sub_job_id,
+                sub_job_type=sub_job_type,
+            ),
+            "events",
+            follow=follow,
+            cursor=cursor,
+            poll_interval=poll_interval,
+        )
+
+    # ─── Checkpoints ─────────────────────────────────────────────────────
+
+    @_track_operation("list_checkpoints")
+    def list_checkpoints(self, job_id: str) -> list:
+        resp = self._send("GET", f"{self._prefix}/{job_id}/checkpoints")
+        return resp.json().get("checkpoints", [])
+
+    @_track_operation("export_checkpoint")
+    def export_checkpoint(self, job_id: str, checkpoint_id: str) -> dict:
+        # The REST API uses colon-action syntax: /{jobId}/checkpoints/{cpId}:export
+        resp = self._send(
+            "POST",
+            f"{self._prefix}/{job_id}/checkpoints/{checkpoint_id}:export",
+        )
+        return resp.json()
+
+    @_track_operation("delete_checkpoint")
+    def delete_checkpoint(self, job_id: str, checkpoint_id: str) -> None:
+        """Delete a checkpoint. Returns None on success."""
+        self._send("DELETE", f"{self._prefix}/{job_id}/checkpoints/{checkpoint_id}")