freesolo-flash-dev 1.2.36__tar.gz → 1.2.37__tar.gz

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.36
+Version: 1.2.37
 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.36 → freesolo_flash_dev-1.2.37}/flash/cli/commands/__init__.py

@@ -399,6 +399,7 @@ def cmd_train(args) -> int:
         # equally not a verification -- so treat anything but an explicit True as unverified.
         affordability_verified = status.pop("affordability_verified", None) is True
         cost = "and cost" if affordability_verified else "but NOT cost"
+        budget_validated = prompt_budget_validation_suffix(status)
         environment = (
             "it did NOT import or run your environment.py. packaged input/output fields and the "
             "statically readable training contract were tokenized together; tokens, retention, "
@@ -411,8 +412,8 @@ def cmd_train(args) -> int:
         )
         print(
             "dry-run validated: config/schema, model+algorithm compatibility, lora rank, "
-            f"runtime-secret presence, warm-start source, serving context cap, {cost}. "
-            f"{environment}",
+            f"runtime-secret presence, warm-start source, serving context cap{budget_validated}, "
+            f"{cost}. {environment}",
             file=sys.stderr,
         )
         if not affordability_verified:
@@ -430,8 +431,10 @@ def cmd_train(args) -> int:
         else:
             print(json.dumps(status, indent=2))
         _print_unpacked_batch_warning(status, spec)  # after the payload, so stdout stays parseable
+        print_status_prompt_budget_warning(status)
         _print_reasoning_loss_warning(status)
         return 0
+    local_budget = warn_before_paid_submit(client, spec)
     status = client.create_run(
         payload,
         runtime_secrets=runtime_secrets,
@@ -440,6 +443,7 @@ def cmd_train(args) -> int:
     run_id = status["run_id"]
     _print_unpacked_batch_warning(status, spec)  # a real submit overrides batch_size the same way
     _print_reasoning_loss_warning(status)  # and trains on whatever reasoning the template kept
+    print_warmstart_context_supplement(local_budget, status)
     logger.info(
         "submitted run %s: model=%s algorithm=%s gpu=%s",
         run_id,
@@ -963,9 +967,13 @@ from flash.cli.commands.deploy import (  # noqa: E402,F401
     cmd_undeploy,
 )
 
-# re-exported at the bottom rather than imported at the top: `train_cost` resolves names back
-# through this package, so a top import would be circular. `cmd_train` below calls these, and the
-# estimate tests import `_warn_if_wandb_requested_without_key` from this package by name.
+# imported at the bottom because the train helpers resolve patched command-package names lazily.
+from flash.cli.commands.prompt_budget import (  # noqa: E402
+    print_status_prompt_budget_warning,
+    print_warmstart_context_supplement,
+    prompt_budget_validation_suffix,
+    warn_before_paid_submit,
+)
 from flash.cli.commands.train_cost import (  # noqa: E402,F401
     _client_train_schema,
     _cmd_train_cost,

freesolo_flash_dev-1.2.37/flash/cli/commands/prompt_budget.py

@@ -0,0 +1,146 @@
+"""prompt-budget warnings and warm-start context lookup for training commands."""
+
+from __future__ import annotations
+
+import sys
+from typing import TypeGuard, cast
+
+from flash import __version__
+from flash.cli.ui import render
+from flash.engine.plan.prompt_budget import PromptBudget, rl_prompt_budget
+
+
+def _positive_int(value: object) -> TypeGuard[int]:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+def _budget_from_status(status: object) -> PromptBudget | None:
+    budget = status.get("prompt_budget") if isinstance(status, dict) else None
+    if not isinstance(budget, dict):
+        return None
+    if budget.get("algorithm") not in {"grpo", "opd", "rl"}:
+        return None
+    if budget.get("context_source") not in {"authored", "recipe_default"}:
+        return None
+    if budget.get("prompt_budget_is_upper_bound") is not True:
+        return None
+    engine_len = budget.get("engine_len")
+    max_completion = budget.get("max_completion")
+    prompt_budget = budget.get("prompt_budget")
+    if not _positive_int(engine_len):
+        return None
+    if not _positive_int(max_completion):
+        return None
+    if not _positive_int(prompt_budget):
+        return None
+    if prompt_budget != engine_len - max_completion:
+        return None
+    source_context = budget.get("warm_start_context")
+    if source_context is not None and not _positive_int(source_context):
+        return None
+    return cast("PromptBudget", budget)
+
+
+def prompt_budget_validation_suffix(status: object) -> str:
+    """qualify a dry-run summary only when the server returned a budget."""
+    return ", prompt budget (upper bound)" if _budget_from_status(status) is not None else ""
+
+
+def warmstart_source_context(client, spec) -> int | None:
+    """read a warm-start source's authored context through the cli's authenticated client."""
+    from flash.core.catalog import samples_on_policy
+    from flash.schema import parse_checkpoint_ref
+
+    ref = getattr(spec.train, "init_from_adapter", "")
+    if not ref or not samples_on_policy(spec.algorithm):
+        return None
+    parsed = parse_checkpoint_ref(ref)
+    if parsed is None:
+        return None
+    try:
+        source = client.get_run(parsed[0])
+    # reporting is best-effort and must never block a valid submission
+    except Exception:
+        return None
+    train = (source or {}).get("spec", {}).get("train", {}) if isinstance(source, dict) else {}
+    context = train.get("max_context_tokens") if isinstance(train, dict) else None
+    if isinstance(context, bool) or not isinstance(context, int) or context <= 0:
+        return None
+    return context
+
+
+def _warmstart_context_sentence(context: object) -> str | None:
+    if isinstance(context, bool) or not isinstance(context, int) or context <= 0:
+        return None
+    return (
+        f"The warm-start source was configured with max_context_tokens={context}; "
+        "it is NOT inherited."
+    )
+
+
+def prompt_budget_warning(
+    budget: object,
+    *,
+    derived_locally: bool = False,
+    include_warm_start_context: bool = True,
+) -> str | None:
+    """render the defaulted-budget warning for one prompt-budget descriptor."""
+    if not isinstance(budget, dict) or budget.get("context_source") != "recipe_default":
+        return None
+    try:
+        prompt_budget = int(budget["prompt_budget"])
+        max_completion = int(budget["max_completion"])
+    except (KeyError, TypeError, ValueError):
+        return None
+    algorithm = str(budget.get("algorithm") or "").upper() or "RL"
+    message = (
+        f"train.max_context_tokens is unset, so {algorithm} derives a prompt budget of at most "
+        f"{prompt_budget} tokens from the recipe default (engine {budget.get('engine_len')} minus "
+        f"max_completion_tokens {max_completion}), not from the model's context. The worker clamps "
+        "the engine length to the model architecture first, so its budget can be smaller than "
+        "this. Prompts over the budget are DROPPED, not truncated, so the run trains on a "
+        "shorter, biased subset."
+    )
+    if derived_locally:
+        message += (
+            f" (Derived locally by this CLI {__version__}; if the control plane is on a "
+            "different release its recipe default may differ.)"
+        )
+    if include_warm_start_context:
+        source = _warmstart_context_sentence(budget.get("warm_start_context"))
+        if source:
+            message += f" {source}"
+    return message + " Set train.max_context_tokens explicitly to choose the budget."
+
+
+def _print_warning(message: str | None) -> None:
+    if message:
+        print(render.warn(message) if render.styled() else f"warning: {message}", file=sys.stderr)
+
+
+def print_status_prompt_budget_warning(status: object) -> None:
+    """print the server-derived warning carried by a submit or dry-run response."""
+    _print_warning(prompt_budget_warning(_budget_from_status(status)))
+
+
+def warn_before_paid_submit(client, spec) -> PromptBudget | None:
+    """print the locally derived warning before create_run can start provisioning."""
+    budget = rl_prompt_budget(spec, warm_start_context=warmstart_source_context(client, spec))
+    _print_warning(prompt_budget_warning(budget, derived_locally=True))
+    return budget
+
+
+def print_warmstart_context_supplement(local_budget: PromptBudget | None, status: object) -> None:
+    """print only context the server resolved after the generic pre-submit warning."""
+    served_budget = _budget_from_status(status)
+    if served_budget is None or (
+        local_budget is not None and local_budget.get("warm_start_context")
+    ):
+        return
+    _print_warning(_warmstart_context_sentence(served_budget.get("warm_start_context")))
+
+
+def warn_cost_prompt_budget(spec) -> None:
+    """print the local prompt-budget warning for the offline ``--cost`` path."""
+    budget = rl_prompt_budget(spec)
+    _print_warning(prompt_budget_warning(budget, derived_locally=True))

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/cli/commands/train_cost.py

@@ -91,6 +91,9 @@ def _cmd_train_cost_offline(spec) -> int:
         print(render.cost_panel(estimate))
     else:
         print(estimate.breakdown())
+    from flash.cli.commands.prompt_budget import warn_cost_prompt_budget
+
+    warn_cost_prompt_budget(spec)
     return 0
 
 

freesolo_flash_dev-1.2.37/flash/engine/plan/prompt_budget.py

@@ -0,0 +1,63 @@
+"""derived grpo/opd prompt-budget data reported before paid training."""
+
+from __future__ import annotations
+
+from typing import Literal, NotRequired, TypedDict
+
+from flash.engine.plan.vram import (
+    grpo_completion_len,
+    grpo_rollout_seq_len,
+    opd_completion_len,
+    opd_rollout_seq_len,
+)
+
+
+class PromptBudget(TypedDict):
+    """serialized prompt-budget descriptor carried by prepared and persisted runs."""
+
+    algorithm: Literal["grpo", "opd", "rl"]
+    engine_len: int
+    max_completion: int
+    prompt_budget: int
+    context_source: Literal["authored", "recipe_default"]
+    prompt_budget_is_upper_bound: Literal[True]
+    warm_start_context: NotRequired[int]
+
+
+def rl_prompt_budget(spec, *, warm_start_context: int | None = None) -> PromptBudget | None:
+    """derive the pre-clamp prompt budget used by a grpo/opd worker.
+
+    the worker clamps ``engine_len`` to the model architecture before subtracting the completion
+    allowance, so this control-plane value is an upper bound. sft returns ``None`` because it
+    truncates over-long rows and reports that behavior through its workload profile instead.
+    """
+    from typing import cast
+
+    from flash.core.catalog import samples_on_policy
+
+    raw_algorithm = getattr(spec, "algorithm", "")
+    if not samples_on_policy(raw_algorithm):
+        return None
+    algorithm = cast("Literal['grpo', 'opd', 'rl']", raw_algorithm)
+    train = spec.train
+    authored = int(getattr(train, "max_context_tokens", 0) or 0)
+    max_tokens = getattr(train, "max_completion_tokens", None)
+    thinking = bool(getattr(spec, "thinking", False))
+    if algorithm == "opd":
+        max_completion = opd_completion_len(max_tokens, thinking)
+        engine_len = opd_rollout_seq_len(authored, max_tokens, thinking)
+    else:
+        max_completion = grpo_completion_len(max_tokens, thinking)
+        engine_len = grpo_rollout_seq_len(authored, max_tokens, thinking)
+    budget: PromptBudget = {
+        "algorithm": algorithm,
+        "engine_len": int(engine_len),
+        "max_completion": int(max_completion),
+        "prompt_budget": int(engine_len) - int(max_completion),
+        "context_source": "authored" if authored else "recipe_default",
+        "prompt_budget_is_upper_bound": True,
+    }
+    source_context = int(warm_start_context or 0)
+    if source_context > 0:
+        budget["warm_start_context"] = source_context
+    return budget

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/engine/worker/opd_train_runner.py

@@ -157,8 +157,12 @@ def _prepare_prompts(
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     thinking_prefill = _opd_train._thinking_prefill_text(tokenizer)
-    requested_len = request.knobs.max_length or (
-        _opd_train.RECIPE.opd.max_prompt_len + request.knobs.max_completion
+    from flash.engine.plan.vram import opd_rollout_seq_len
+
+    requested_len = opd_rollout_seq_len(
+        request.knobs.max_length,
+        request.knobs.max_completion,
+        bool(_opd_train._w.THINKING),
     )
     # clamp to the architecture BEFORE deriving the prompt budget, so every downstream length agrees.
     # clamping only the engine would admit prompts sized against the unclamped budget and then fail

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/engine/worker/train/rl/inputs.py

@@ -208,15 +208,15 @@ def _grpo_is_multimodal(train, message_prompts):
 
 
 def _resolve_sequence_lengths(model_id, model_revision, train_spec, rl, gcfg, tok, multi_turn):
-    # same resolver the spec-parse prompt-budget guard uses, so a run rejected at submit is exactly
-    # the run this would have failed here after the gpu was already paid for.
-    from flash.engine.plan.vram import grpo_completion_len
+    from flash.engine.plan.vram import grpo_completion_len, grpo_rollout_seq_len
 
-    max_completion = grpo_completion_len(gcfg.get("max_tokens"), bool(_w.THINKING))
+    thinking = bool(_w.THINKING)
+    max_tokens = gcfg.get("max_tokens")
+    max_completion = grpo_completion_len(max_tokens, thinking)
     train_ctx = (
         train_spec.max_context_tokens if (train_spec and train_spec.max_context_tokens) else 0
     )
-    requested_len = int(train_ctx or max(1024, rl.max_prompt_len + max_completion))
+    requested_len = grpo_rollout_seq_len(train_ctx, max_tokens, thinking)
     # clamp to the architecture BEFORE deriving the prompt budget, so every downstream length agrees.
     # clamping only the engine would admit prompts up to the unclamped budget and then fail them at
     # rollout, and would let the token budget pack more than the one-sequence memory floor intends.

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/runner/__init__.py

@@ -34,6 +34,7 @@ from flash.core.spec import (  # noqa: F401
     GpuSpec,
     JobSpec,
 )
+from flash.engine.plan.prompt_budget import PromptBudget
 from flash.providers._lifecycle.poll import _attempt_int as _attempt_int
 from flash.teacher.retry_contract import (
     OPD_RETRY_CONTRACT_STATUS_KEY,
@@ -190,6 +191,10 @@ class RunStatus:
     gpu_status: dict | None = None
     workload_profile_input_digest: str | None = None
     workload_profile: dict | None = None
+    # submit-time derived grpo/opd prompt budget from flash.engine.plan.prompt_budget. workers drop
+    # over-budget prompts rather than truncating them, so record the value before gpu allocation.
+    # none for sft, which reports truncation through workload_profile, and for older records.
+    prompt_budget: PromptBudget | None = None
     effective_preparation: dict | None = None
 
     def to_dict(self) -> dict:
@@ -424,6 +429,8 @@ class PreparedJob:
     worker_spec: JobSpec
     estimated_cost_usd: float
     adapter_identity: dict | None = None
+    # derived grpo/opd prompt budget, resolved where the warm-start source is authorized and read.
+    prompt_budget: PromptBudget | None = None
 
 
 _BILLING_FIELDS = frozenset({"billing_state", "billing_error", "billing_charge"})

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/runner/preparation.py

@@ -143,7 +143,7 @@ def _prepare_init_from_adapter(
     owner_org_id: str = "",
     owner_key_id: int | None = None,
     token: str | None = None,
-) -> tuple[JobSpec, JobSpec, dict | None]:
+) -> tuple[JobSpec, JobSpec, dict | None, int | None]:
     """prepare public and worker specs with source-authoritative adapter metadata.
 
     Failures here are genuinely about the warm-start source, so they are tagged
@@ -168,11 +168,11 @@ def _prepare_init_from_adapter_inner(
     owner_org_id: str = "",
     owner_key_id: int | None = None,
     token: str | None = None,
-) -> tuple[JobSpec, JobSpec, dict | None]:
+) -> tuple[JobSpec, JobSpec, dict | None, int | None]:
     _runner()._require_supported_adapter_continuation(spec)
     ref = spec.train.init_from_adapter
     if not ref:
-        return spec, spec, None
+        return spec, spec, None, None
     from flash.adapters.lora_rank import (
         adapter_artifact_identity,
         load_hf_adapter_config,
@@ -265,7 +265,8 @@ def _prepare_init_from_adapter_inner(
         worker_spec,
         train=replace(worker_spec.train, lora_rank=metadata.rank, lora_alpha=metadata.alpha),
     )
-    return public_spec, worker_spec, identity
+    source_context = int(getattr(src_spec.train, "max_context_tokens", 0) or 0) or None
+    return public_spec, worker_spec, identity, source_context
 
 
 def _mark_warmstart_source(worker_spec: JobSpec, child_run_id: str) -> None:

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/flash/runner/submit.py

@@ -127,11 +127,13 @@ def prepare_job(
     owner_org_id = _runner()._context_org_id(billing_context) or _runner()._context_org_id(
         platform_context
     )
-    public_spec, worker_spec, adapter_identity = _runner()._prepare_init_from_adapter(
-        spec,
-        owner_org_id=owner_org_id,
-        owner_key_id=owner_key_id,
-        token=os.environ.get("HF_TOKEN"),
+    public_spec, worker_spec, adapter_identity, warm_start_context = (
+        _runner()._prepare_init_from_adapter(
+            spec,
+            owner_org_id=owner_org_id,
+            owner_key_id=owner_key_id,
+            token=os.environ.get("HF_TOKEN"),
+        )
     )
     # these read-only gates belong to preparation: every submit path passes here exactly once, and
     # callers receive the pinned worker spec before quoting, affordability, persistence, or allocation.
@@ -149,11 +151,20 @@ def prepare_job(
     from flash.cost.spec import estimate_for_spec
 
     estimated_cost_usd = float(estimate_for_spec(worker_spec).total_usd)
+    # derive the rl prompt budget from the same resolved spec the quote is built from, so the
+    # reported budget describes the run that was actually priced and submitted.
+    from flash.engine.plan.prompt_budget import rl_prompt_budget
+
+    prompt_budget = rl_prompt_budget(
+        worker_spec,
+        warm_start_context=warm_start_context,
+    )
     return _runner().PreparedJob(
         public_spec=public_spec,
         worker_spec=worker_spec,
         estimated_cost_usd=estimated_cost_usd,
         adapter_identity=adapter_identity,
+        prompt_budget=prompt_budget,
     )
 
 
@@ -267,6 +278,7 @@ def submit_job(
         platform_context=platform_context,
         workload_profile_input_digest=worker_spec.workload_profile_input_digest or None,
         workload_profile=worker_spec.workload_profile or None,
+        prompt_budget=prepared.prompt_budget,
         effective_preparation={
             "worker_spec": worker_spec.to_internal_dict(),
             "workload_profile": worker_spec.workload_profile or None,

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.36"
+version = "1.2.37"
 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.36"
+version = "1.2.37"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.2.36 → freesolo_flash_dev-1.2.37}/tests/test_client_server_integration.py

@@ -237,6 +237,11 @@ def test_dry_run_create_previews_on_the_server(make_client) -> None:
     # Persisted like any run — retrievable and listed — but no GPU was ever allocated.
     assert client.get_run(run_id)["state"] == "dry_run"
     assert run_id in [r["run_id"] for r in client.list_runs()]
+    # the derived prompt budget travels with the preview and remains visible through runs status.
+    budget = created["prompt_budget"]
+    assert budget["context_source"] == "recipe_default"
+    assert budget["prompt_budget"] == budget["engine_len"] - budget["max_completion"]
+    assert client.get_run(run_id)["prompt_budget"] == budget
 
 
 def test_logs_offset_paging_roundtrip(make_client) -> None: