freesolo-flash-dev 1.2.12__tar.gz → 1.2.13__tar.gz

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.12
+Version: 1.2.13
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/adapters/lora_rank.py

@@ -7,11 +7,11 @@ import json
 import math
 import re
 from collections.abc import Callable, Mapping
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from decimal import Decimal, InvalidOperation
 from typing import TYPE_CHECKING, Any
 
-from flash.core.catalog import serving_context_cap, serving_lora_rank_cap
+from flash.core.catalog import lora_expert_count, serving_context_cap, serving_lora_rank_cap
 
 
 class ServingPreflightError(ValueError):
@@ -279,6 +279,117 @@ def rank_from_adapter_config(config: Mapping[str, Any], *, source: str) -> int:
     )
 
 
+_LORA_A_INFIX = ".lora_A."
+_LORA_B_INFIX = ".lora_B."
+
+
+@dataclass(frozen=True)
+class DeclaredLoraRanks:
+    """The default rank, per-module overrides, and parameter-target module paths in a config."""
+
+    default: int | None = None
+    by_module: Mapping[str, int] = field(default_factory=dict)
+    stacked_rank_modules: tuple[str, ...] = ()
+    stacked_rank_multiplier: int | None = None
+
+    def __bool__(self) -> bool:
+        return self.default is not None or bool(self.by_module)
+
+
+def declared_lora_ranks(config: Mapping[str, Any]) -> DeclaredLoraRanks:
+    """Read the per-module LoRA rank structure out of ``adapter_config.json``."""
+    if not isinstance(config, Mapping):
+        return DeclaredLoraRanks()
+    try:
+        default = _positive_int(config["r"], source="adapter_config.json", field="r")
+    except (KeyError, ValueError):
+        default = None
+
+    by_module: dict[str, int] = {}
+    pattern = config.get("rank_pattern")
+    if isinstance(pattern, Mapping):
+        for module, value in pattern.items():
+            try:
+                by_module[str(module)] = _positive_int(
+                    value, source="adapter_config.json", field="rank_pattern"
+                )
+            except ValueError:
+                continue
+
+    stacked_rank_modules: list[str] = []
+    targets = config.get("target_parameters")
+    if isinstance(targets, list):
+        for target in targets:
+            if not isinstance(target, str):
+                continue
+            module, separator, _ = target.rpartition(".")
+            if separator and module:
+                stacked_rank_modules.append(module)
+
+    base_model = str(config.get("base_model_name_or_path") or "").strip()
+    return DeclaredLoraRanks(
+        default=default,
+        by_module=by_module,
+        stacked_rank_modules=tuple(stacked_rank_modules),
+        stacked_rank_multiplier=lora_expert_count(base_model),
+    )
+
+
+def _rank_for_module(module_path: str, declared: DeclaredLoraRanks) -> int | None:
+    """Resolve a module's rank with PEFT's ordered, anchored ``rank_pattern`` matching."""
+    for module, rank in declared.by_module.items():
+        if not module:
+            continue
+        try:
+            # entries are regexes in PEFT, deliberately not escaped here either
+            if re.match(rf"(.*\.)?({module})$", module_path):
+                return rank
+        except re.error:
+            continue
+    return declared.default
+
+
+def _module_uses_target_parameters(module_path: str, declared: DeclaredLoraRanks) -> bool:
+    """Whether this serialized module came from one of PEFT's targeted parameters."""
+    return any(
+        module_path == module or module_path.endswith(f".{module}")
+        for module in declared.stacked_rank_modules
+    )
+
+
+def lora_tensor_rank_disagrees(key: str, shape: Any, declared: DeclaredLoraRanks) -> bool:
+    """Return whether a 2-D LoRA weight provably contradicts its configured rank.
+
+    Ordinary module weights carry exactly ``r`` on the LoRA axis. A 3-D parameter targeted through
+    PEFT's ``target_parameters`` serializes under its parent module and stacks every model expert on
+    that axis, so only that module may carry exactly ``r * num_experts``. The distinction must be per
+    module because one config may contain ordinary ``target_modules`` and fused parameters together.
+    """
+    if not declared:
+        return False
+    is_a = _LORA_A_INFIX in key
+    is_b = _LORA_B_INFIX in key
+    if is_a == is_b:  # neither, or a key somehow claiming both
+        return False
+    if not isinstance(shape, (list, tuple)) or len(shape) != 2:
+        return False
+    dims = [d for d in shape if isinstance(d, int) and not isinstance(d, bool) and d > 0]
+    if len(dims) != 2:
+        return False
+
+    infix = _LORA_A_INFIX if is_a else _LORA_B_INFIX
+    module_path = key.partition(infix)[0]
+    rank = _rank_for_module(module_path, declared)
+    if rank is None:
+        return False
+    axis = dims[0] if is_a else dims[1]
+    if _module_uses_target_parameters(module_path, declared):
+        multiplier = declared.stacked_rank_multiplier
+        expected = rank * multiplier if multiplier is not None else rank
+        return axis != expected
+    return axis != rank
+
+
 def alpha_from_adapter_config(config: Mapping[str, Any], *, source: str) -> int:
     """Return the maximum LoRA alpha across default and per-module metadata."""
     return _max_adapter_value(

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/core/catalog.py

@@ -148,6 +148,9 @@ class ModelInfo:
     mamba_block_size: int = 0
     # grouped (in_features, out_features, count) for peft all-linear on the full loaded model.
     lora_target_shapes: tuple[tuple[int, int, int], ...] = ()
+    # fused routed experts stacked on the lora rank axis for target_parameters. zero means no
+    # catalog-backed stacked allowance; ordinary and unknown target parameters carry exactly r.
+    lora_expert_count: int = 0
 
     @property
     def is_moe(self) -> bool:
@@ -174,6 +177,7 @@ class ModelInfo:
             "linear_value_head_dim",
             "linear_conv_kernel_dim",
             "lora_target_shapes",
+            "lora_expert_count",
         ):
             data.pop(key, None)
         serving = data["serving"]
@@ -500,6 +504,7 @@ MODELS: dict[str, ModelInfo] = {
             (4608, 2048, 1),
             (4608, 4608, 1),
         ),
+        lora_expert_count=256,
         vocab_size=248_320,
         algos=ALGORITHMS,
         min_vram_gb=141,
@@ -569,6 +574,14 @@ def serving_lora_rank_cap(model: str | ModelInfo | None) -> int | None:
     return int(info.serving.max_lora_rank)
 
 
+def lora_expert_count(model: str | ModelInfo | None) -> int | None:
+    """Return the exact fused-expert stack size for LoRA target parameters, when cataloged."""
+    info = _model_info_for_serving(model)
+    if info is None or info.lora_expert_count <= 0:
+        return None
+    return int(info.lora_expert_count)
+
+
 def serving_context_cap(model: str | ModelInfo | None) -> int | None:
     """Return the model's served ``max_model_len``, or None without a local serving entry.
 

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/engine/worker/train/rl/child/multiturn.py

@@ -62,11 +62,23 @@ def post_json(url: str, path: str, payload: dict) -> dict:
         with urllib.request.urlopen(request) as response:
             body = response.read()
     except urllib.error.HTTPError as error:
+        # decode inside the guard, raise outside it: `suppress(Exception)` around the raise as well
+        # swallowed the detail-bearing RuntimeError it was meant to let through, so the bridge's own
+        # error text ("can't start new thread") could never reach the user and every failure read as
+        # a bare "returned HTTP 400".
+        detail = None
         with contextlib.suppress(Exception):
-            detail = json.loads(error.read().decode("utf-8"))["error"]
-            raise RuntimeError(f"flash multi-turn bridge rejected {path}: {detail}") from error
+            detail = str(json.loads(error.read().decode("utf-8"))["error"])
+        # 5xx is the bridge failing to serve a valid request (resource exhaustion); 4xx is this
+        # request being rejected. naming which one it was stops a server fault from being read as a
+        # malformed payload.
+        fault = "could not serve" if error.code >= 500 else "rejected"
+        if detail:
+            raise RuntimeError(
+                f"flash multi-turn bridge {fault} {path} (HTTP {error.code}): {detail}"
+            ) from error
         raise RuntimeError(
-            f"flash multi-turn bridge returned HTTP {error.code} for {path}"
+            f"flash multi-turn bridge {fault} {path}: HTTP {error.code} with no error detail"
         ) from error
     except (OSError, http.client.HTTPException) as error:
         raise RuntimeError(

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/engine/worker/train/rl/multi_turn.py

@@ -39,6 +39,70 @@ _SINGLE_TURN_SCORE_BATCH_SIZE = 64
 _SINGLE_TURN_SCORE_FLUSH_WAIT_S = 0.1
 _SINGLE_TURN_SCORE_SHUTDOWN_WAIT_S = 5.0
 
+# a body that could not be decoded at all: bad Content-Length, bad utf-8, or bad json.
+_UNDECODABLE_BODY_ERRORS = (ValueError, TypeError, UnicodeDecodeError)
+
+
+class _BadRequest(Exception):
+    """A request THIS BRIDGE is deliberately rejecting because the request itself is wrong.
+
+    Raised by the validating code rather than inferred from exception type at the handler, which
+    cannot distinguish the bridge rejecting a bad index from a user env raising IndexError inside
+    its own scoring -- and reporting the latter as a client error is the bug the 503 split fixes.
+    """
+
+
+class _BadIndex(_BadRequest, IndexError):
+    """An index outside the dataset. Also an IndexError: that is what it genuinely is."""
+
+
+class _BadSession(_BadRequest, KeyError):
+    """An unknown or duplicate session id. Also a KeyError, for the same reason.
+
+    ``KeyError.__str__`` reprs its argument, which would double-quote the message; this reports the
+    message as written so the client sees the same text every other rejection produces.
+    """
+
+    def __str__(self) -> str:
+        return self.args[0] if self.args else ""
+
+
+def _request_field(payload: dict, key: str):
+    """Read one required payload field, or reject the request.
+
+    A decoded body that is not an object at all (``[]``, ``"x"``, ``3``) is a malformed request, so
+    it is rejected here rather than raising TypeError out of the subscript and reading as a fault.
+    """
+    if not isinstance(payload, dict):
+        raise _BadRequest(f"request body must be a json object, got {type(payload).__name__}")
+    try:
+        return payload[key]
+    except KeyError as exc:
+        raise _BadRequest(f"missing required field {key!r}") from exc
+
+
+def request_int(payload: dict, key: str) -> int:
+    """Read one required integer payload field, or reject the request."""
+    value = _request_field(payload, key)
+    try:
+        return int(value)
+    # OverflowError: json decodes 1e309 to inf, and int(inf) overflows rather than ValueError-ing.
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise _BadRequest(f"field {key!r} is not an integer: {value!r}") from exc
+
+
+def request_session_id(payload: dict) -> str:
+    """Read the required session id, or reject the request."""
+    return str(_request_field(payload, "session_id"))
+
+
+# ONLY the bridge's own deliberate rejections are client errors. classifying by exception TYPE
+# instead cannot work: a user env raising IndexError/KeyError deep inside its own scoring would be
+# reported as a malformed request -- the same "blame the caller for this side's failure" bug this
+# module's 503 split exists to fix, one layer down. so the rejecting code raises _BadRequest itself
+# and anything else reaching the handler is this side failing to serve a well-formed request.
+_BAD_REQUEST_ERRORS = (_BadRequest,)
+
 # size the listen backlog for the full prompts_per_step * group_size connection burst.
 # socketserver's default of 5 resets overflowed clients, and bridge_post intentionally does not retry.
 # a fixed backlog only moves the cliff; linux may clamp this value to somaxconn.
@@ -113,10 +177,10 @@ class MultiTurnBridge:
         self._scorer.close(_MULTI_TURN_SCORE_SHUTDOWN_WAIT_S)
 
     def _session(self, payload: dict) -> dict:
-        session_id = str(payload["session_id"])
+        session_id = request_session_id(payload)
         session = self._sessions.get(session_id)
         if session is None:
-            raise KeyError(f"unknown multi-turn session {session_id}")
+            raise _BadSession(f"unknown multi-turn session {session_id}")
         # every touch renews the lease, so only a session nobody is driving can ever age out. all
         # callers already hold the lock.
         session["touched_at"] = time.monotonic()
@@ -137,12 +201,12 @@ class MultiTurnBridge:
         return stale
 
     def start(self, payload: dict) -> dict:
-        index = int(payload["index"])
+        index = request_int(payload, "index")
         if index < 0 or index >= len(self._examples):
-            raise IndexError(
+            raise _BadIndex(
                 f"multi-turn example index {index} is outside [0, {len(self._examples)})"
             )
-        session_id = str(payload["session_id"])
+        session_id = request_session_id(payload)
         example = self._examples[index]
         with self._lock:
             # swept here rather than on a timer thread: a session is only ever abandoned by an
@@ -150,7 +214,7 @@ class MultiTurnBridge:
             # exactly here. no extra thread, and nothing to shut down.
             reaped = self._reap_abandoned_sessions()
             if session_id in self._sessions:
-                raise KeyError(f"duplicate multi-turn session {session_id}")
+                raise _BadSession(f"duplicate multi-turn session {session_id}")
             state = self._env.new_rollout_state(example)
             # new_rollout_state calls start_episode again after dataset preparation. adopt the dataset's
             # prompt so randomized envs do not generate for episode a and score episode b; keep the
@@ -216,7 +280,7 @@ class MultiTurnBridge:
             return score_rollouts(self._env, requests)
 
     def score(self, payload: dict) -> dict:
-        turn_count = int(payload["turn_count"])
+        turn_count = request_int(payload, "turn_count")
         with self._lock:
             session = self._session(payload)
             state = session["state"]
@@ -282,7 +346,7 @@ class MultiTurnBridge:
 
     def close(self, payload: dict) -> dict:
         with self._lock:
-            self._sessions.pop(str(payload["session_id"]), None)
+            self._sessions.pop(request_session_id(payload), None)
         return {"closed": True}
 
     def open_sessions(self) -> int:
@@ -371,9 +435,9 @@ def start_reward_server(
     )
 
     def _score_route(payload: dict) -> dict:
-        index = int(payload["index"])
+        index = request_int(payload, "index")
         if index < 0 or index >= example_count:
-            raise IndexError(f"reward example index {index} is outside [0, {example_count})")
+            raise _BadIndex(f"reward example index {index} is outside [0, {example_count})")
         solution_str = payload.get("solution_str", "")
         if score_batcher is not None:
             return {"score": float(score_batcher.submit((index, solution_str)))}
@@ -397,14 +461,37 @@ def start_reward_server(
             try:
                 n = int(self.headers.get("Content-Length", "0"))
                 payload = json.loads(self.rfile.read(n).decode("utf-8"))
-                result = route(payload)
-            except Exception as exc:
-                print(f"[rl-verl] reward server request failed: {exc}", flush=True)
-                body = json.dumps({"error": str(exc)}).encode()
-                self.send_response(400)
+            except _UNDECODABLE_BODY_ERRORS as exc:
+                status = 400
+                detail = f"{type(exc).__name__}: {exc}"
             else:
+                try:
+                    result = route(payload)
+                except _BAD_REQUEST_ERRORS as exc:
+                    # the request itself is wrong: a malformed field, an index outside the dataset,
+                    # or an unknown/duplicate session. no retry or capacity change makes it succeed.
+                    # reported without the internal class name -- the message already says what is
+                    # wrong, and "_BadIndex:" only names a private type the caller cannot act on.
+                    status = 400
+                    detail = str(exc)
+                except Exception as exc:
+                    # the request was fine and this side could not serve it -- thread exhaustion
+                    # ("can't start new thread") under rollout concurrency is the case that drove
+                    # this split. 400 blamed the caller's payload and sent every reader to score
+                    # code that cannot produce one, so a resource fault reads as a 5xx here.
+                    status = 503
+                    detail = f"{type(exc).__name__}: {exc}"
+                else:
+                    status = 200
+            if status == 200:
                 body = json.dumps(result).encode()
-                self.send_response(200)
+            else:
+                print(
+                    f"[rl-verl] reward server request failed ({status}) for {self.path}: {detail}",
+                    flush=True,
+                )
+                body = json.dumps({"error": detail}).encode()
+            self.send_response(status)
             self.send_header("Content-Type", "application/json")
             self.send_header("Content-Length", str(len(body)))
             self.end_headers()

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/engine/worker/train/rl/shims.py

@@ -618,6 +618,21 @@ def render_reward_module(url_env: str = "FLASH_VERL_REWARD_URL") -> str:
         "        with urllib.request.urlopen(req) as r:\n"
         "            payload = json.loads(r.read().decode())\n"
         "            return float(payload['score'])\n"
+        # HTTPError is a URLError subclass, so it must be handled FIRST or the bridge's own error
+        # text -- the reason the request failed, e.g. "can't start new thread" -- is dropped for a
+        # bare "HTTP Error 503: Service Unavailable" that names no cause.
+        "    except urllib.error.HTTPError as exc:\n"
+        "        detail = ''\n"
+        "        try:\n"
+        "            detail = str(json.loads(exc.read().decode())['error'])\n"
+        "        except Exception:\n"
+        "            detail = ''\n"
+        "        fault = 'could not serve' if exc.code >= 500 else 'rejected'\n"
+        "        if detail:\n"
+        "            raise RuntimeError('flash reward bridge %s the request (HTTP %s): %s'\n"
+        "                               % (fault, exc.code, detail)) from exc\n"
+        "        raise RuntimeError('flash reward bridge %s the request: HTTP %s with no error detail'\n"
+        "                           % (fault, exc.code)) from exc\n"
         "    except urllib.error.URLError as exc:\n"
         "        raise RuntimeError('flash reward bridge request failed: %s' % exc) from exc\n"
         "    except Exception as exc:\n"

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/serve/export.py

@@ -27,6 +27,11 @@ from flash.adapters.artifacts import (
     has_loadable_adapter_weights,
     is_adapter_weight_filename,
 )
+from flash.adapters.lora_rank import (
+    DeclaredLoraRanks,
+    declared_lora_ranks,
+    lora_tensor_rank_disagrees,
+)
 from flash.serve.deploy import ServingError
 
 logger = get_logger(__name__)
@@ -395,6 +400,9 @@ def _normalize_adapter_key_namespace(adapter_dir: Path) -> str:
         raise ValueError("no adapter_model weight file to normalize")
     if paths[0].name.endswith(".safetensors"):
         scans = [_scan_safetensors(path) for path in paths]
+        # before any rewrite: the shapes are what they are, and an adapter whose tensors contradict
+        # its own config is unloadable regardless of which key namespace it ends up in.
+        _verify_export_tensor_ranks(adapter_dir, scans)
     elif any(_bin_may_carry_infixed_keys(path) for path in paths):
         # one shard mentioning the infix commits every shard to a full scan. liveness and
         # post-rename collisions are properties of the union of the shards, so a shard that never
@@ -445,6 +453,9 @@ def _normalize_export_adapter_keys(adapter_dir: Path) -> str:
     """
     try:
         return _normalize_adapter_key_namespace(adapter_dir)
+    except _AdapterRankMismatch:
+        # a self-contradicting adapter, not an unreadable key namespace. reported as itself.
+        raise
     except Exception as exc:
         raise ServingError(
             f"could not normalize exported adapter keys ({exc}); refusing to export weights "
@@ -452,6 +463,70 @@ def _normalize_export_adapter_keys(adapter_dir: Path) -> str:
         ) from exc
 
 
+class _AdapterRankMismatch(ValueError):
+    """The adapter's tensors contradict its own config, so it is unloadable as published.
+
+    Distinct from the errors `_normalize_export_adapter_keys` wraps: that wrapper reports a key
+    namespace it could not normalize, which is a different defect with a different remedy, and
+    letting this one be reported as that would send the user looking in the wrong place. A
+    ValueError so it reaches the export route as a 404 beside the sibling "no loadable LoRA
+    adapter" refusal, rather than as a 502 that blames the Hub for a bad artifact.
+    """
+
+
+def _declared_export_ranks(adapter_dir: Path) -> DeclaredLoraRanks:
+    """The rank structure ``adapter_config.json`` declares, or an empty one we cannot judge from.
+
+    Empty rather than an error: a config carrying no readable rank is the pre-existing shape of
+    every fixture and of adapters PEFT wrote without one, and export has never required it. The
+    check below exists to catch a config and its tensors DISAGREEING, which needs both halves.
+    """
+    try:
+        config = json.loads((adapter_dir / "adapter_config.json").read_text(encoding="utf-8"))
+    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
+        return DeclaredLoraRanks()
+    return declared_lora_ranks(config) if isinstance(config, dict) else DeclaredLoraRanks()
+
+
+def _verify_export_tensor_ranks(adapter_dir: Path, scans: list[_WeightScan]) -> None:
+    """Refuse an export whose LoRA tensors do not carry the rank its config declares.
+
+    A serving engine sizes its LoRA slots from ``adapter_config.json`` and then binds the tensors,
+    so the two disagreeing is not a cosmetic metadata bug: it is an adapter that cannot be loaded.
+    That is not hypothetical -- a 35B MoE SFT run published `r: 32` beside expert tensors export
+    checked only for presence and key-namespace readability, so it succeeded and printed a Hub URL
+    while deploying the identical artifact failed. Export exiting 0 is what made a broken adapter
+    look like a finished deliverable.
+
+    `lora_tensor_rank_disagrees` resolves the rank PER MODULE the way PEFT does, so a
+    `rank_pattern` override is judged against its own rank rather than a summary of the config.
+
+    Only safetensors headers are read, because they carry shapes as metadata and are already parsed
+    here. A `.bin` would have to be loaded through torch, which the control plane does not install,
+    and paying that cost is not warranted for a representation the trainer no longer writes.
+    """
+    declared = _declared_export_ranks(adapter_dir)
+    if not declared:
+        return
+    mismatched = [
+        f"{key} has shape {descriptor.get('shape')}"
+        for scan in scans
+        for key, descriptor in (scan.header or {}).items()
+        if key != "__metadata__"
+        and isinstance(descriptor, dict)
+        and lora_tensor_rank_disagrees(key, descriptor.get("shape"), declared)
+    ]
+    if not mismatched:
+        return
+    shown = ", ".join(sorted(mismatched)[:3])
+    raise _AdapterRankMismatch(
+        f"adapter_config.json declares r={declared.default} but {len(mismatched)} LoRA tensor(s) "
+        f"do not carry the rank configured for their module (e.g. {shown}); a serving engine sizes "
+        "its LoRA slots from the config, so this adapter cannot be loaded and exporting it would "
+        "publish a broken artifact"
+    )
+
+
 def _rewrite_adapter_config_base_model(
     adapter_dir: Path, base_model: str, base_model_revision: str = ""
 ) -> bool:

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/serve/preflight.py

@@ -18,6 +18,17 @@ from flash.engine.plan.vram import opd_completion_len
 
 SERVING_PROMPT_TOKEN_ALLOWANCE = 256
 
+# ceiling on what one deployment smoke may generate, independent of whatever the run trains at.
+# the smoke asks a fixed trivial question under a wall clock that must ALSO cover cold-starting the
+# base model and loading the adapter, so a budget inherited from training spends that deadline on
+# tokens nobody reads: an sft run raised to an 8192 context asked a thinking 27B for 8192 tokens and
+# timed out the deployment. that also coupled the knobs the wrong way, since raising
+# max_context_tokens to avoid training truncation made the run harder to deploy.
+# sized at the largest thinking budget any algorithm resolves to by default, so no
+# default-configured run is smoked at less than it is today, and a generation that still hits this
+# has failed to answer a trivial prompt rather than run out of room.
+SMOKE_COMPLETION_TOKEN_CEILING = 2048
+
 
 class ExternalSchemaReference(ValueError):
     def __init__(self, ref: str) -> None:
@@ -68,21 +79,15 @@ def validate_structured_output_patterns(constraint: dict) -> None:
             raise ValueError(f"{field} is invalid: {exc}") from exc
 
 
-def resolve_smoke_completion_tokens(spec: JobSpec) -> int:
-    """Resolve the deployment smoke's completion budget for one generation.
-
-    This sizes a single smoke request, not the run's training context. sft is the reason the two
-    cannot share a number: an sft run has no completion budget of its own, because the worker trains
-    one packed prompt+completion block bounded by ``max_context_tokens``.
-    """
+def _run_completion_budget(spec: JobSpec) -> int:
+    """The completion budget the run itself trains at, before the smoke's own ceiling applies."""
     explicit = spec.train.max_completion_tokens
     positive_explicit = int(explicit) if explicit is not None and int(explicit) > 0 else None
     if spec.algorithm == "opd":
         return opd_completion_len(positive_explicit, spec.thinking)
     if spec.algorithm == "sft":
-        # SFT ignores max_completion_tokens, so mirror the worker: use explicit max_context_tokens with
-        # the recipe fallback. do not subtract the smoke prompt; slight over-allocation is safe, while
-        # under-allocation can reject a valid checkpoint. serving clamps the final request.
+        # SFT ignores max_completion_tokens, so mirror the worker: use explicit max_context_tokens
+        # with the recipe fallback.
         context = spec.train.max_context_tokens
         if context is not None and int(context) > 0:
             return int(context)
@@ -93,6 +98,28 @@ def resolve_smoke_completion_tokens(spec: JobSpec) -> int:
     return int(recipe.max_completion_len_thinking if spec.thinking else recipe.max_completion_len)
 
 
+def resolve_smoke_completion_tokens(spec: JobSpec, *, constrained: bool = False) -> int:
+    """Resolve the deployment smoke's completion budget for one generation.
+
+    This sizes a single smoke request, not the run's training context, and the two are deliberately
+    decoupled: the smoke's question is fixed and trivial, so what a run trains at says nothing about
+    how many tokens answering it needs. Taking the run's number verbatim let a long training context
+    spend the smoke's whole wall-clock budget generating, which is why the ceiling exists.
+
+    A configured grammar is exempt. Deployment registers it as the adapter's serving default, so the
+    smoke generates under it too, and the shortest string it admits can exceed the ceiling -- a long
+    ``choice``, a fixed-repetition ``regex``, a schema with a large ``minLength``. Capping there
+    would truncate the only legal answer and reject an adapter that serves correctly. Deciding that
+    from the constraint needs the tokenizer and a minimum-length analysis of three grammar dialects,
+    so the run's own budget stands in for it: the run trains within that number, so it bounds the
+    grammar it was trained to satisfy.
+    """
+    budget = _run_completion_budget(spec)
+    if constrained:
+        return budget
+    return min(budget, SMOKE_COMPLETION_TOKEN_CEILING)
+
+
 def resolve_effective_completion_tokens(spec: JobSpec) -> int | None:
     """Resolve the completion budget the serving context guard must fit, if the run has one.
 

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/flash/server/routes/serving_smoke.py

@@ -332,10 +332,16 @@ def _run_deployment_smoke(
     train = getattr(spec, "train", None)
     constraint = parse_structured_outputs(getattr(train, "structured_outputs", ""))
     max_tokens = 256
-    # thinking spends tokens before content, so use the run's own completion budget regardless of
-    # grammar; 256 tokens can truncate the think block and reject a healthy deployment.
+    # thinking spends tokens before content regardless of grammar, so 256 can truncate the think
+    # block and reject a healthy deployment. the resolver widens it to what the run reasons within,
+    # bounded by a smoke-specific ceiling -- this budget is spent from the same wall clock that must
+    # also cold-start the base model and load the adapter, so it cannot track the training context.
+    # a grammar lifts that ceiling: it is the adapter's serving default, so the smoke generates under
+    # it, and the shortest string it admits may be longer than the ceiling allows.
     if spec.thinking:
-        max_tokens = max(256, resolve_smoke_completion_tokens(spec))
+        max_tokens = max(
+            256, resolve_smoke_completion_tokens(spec, constrained=constraint is not None)
+        )
         serving_capacity = serving_completion_token_capacity(
             spec, prompt_allowance=SERVING_PROMPT_TOKEN_ALLOWANCE
         )

{freesolo_flash_dev-1.2.12 → freesolo_flash_dev-1.2.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.12"
+version = "1.2.13"
 description ="Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -137,7 +137,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "1.2.12"
+version = "1.2.13"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]