freesolo-flash-dev 0.2.25__py3-none-any.whl

flash/envs/adapter/rubric.py

@@ -0,0 +1,222 @@
+"""Rubric/reward/introspection helpers for the verifiers adapter.
+
+Leaf helpers split out of ``flash.envs.adapter``: rubric flattening, judge discovery,
+reward-func invocation, eval-metric guarding, and env-shape introspection. None reference
+the rest of the adapter package (no import cycle); the package ``__init__`` re-exports them.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+
+# The judge-related kwarg names a reward func may declare, sourced from a JudgeRubric.
+# Single source of truth for both ``_judge_kwargs`` and ``_AVAILABLE_REWARD_KWARGS``.
+_JUDGE_KWARG_NAMES = ("judge", "judge_client", "judge_model", "judge_prompt")
+
+# The kwargs this adapter can supply to a reward func. The non-judge keys are exactly the
+# ones built into the ``available`` dict in VerifiersEnvironment._reward_available; the judge
+# keys come from ``_judge_kwargs``. Deriving the frozenset from these shared names avoids the
+# manual "keep in sync" coupling (adding a kwarg below without updating the set would
+# re-trigger the false "requires unavailable arg" failure).
+_BASE_REWARD_KWARG_NAMES = (
+    "completion",
+    "prompt",
+    "answer",
+    "info",
+    "state",
+    "parser",
+    "task",
+)
+_AVAILABLE_REWARD_KWARGS = frozenset(_BASE_REWARD_KWARG_NAMES + _JUDGE_KWARG_NAMES)
+
+
+def _reward_requires_unavailable_args(func) -> str | None:
+    """Name of a required arg this adapter cannot supply, or None.
+
+    Group/batch reward funcs declare plural required params (``completions``,
+    ``prompts``, ``answers``, ...). The worker scores one completion at a time and has no
+    batch, so such a func would be called without its required argument and silently score
+    0.0 — train/eval on an all-zero signal. Detect it so the caller can fail fast."""
+    try:
+        params = inspect.signature(func).parameters.values()
+    except (TypeError, ValueError):
+        return None  # builtins/uninspectable: _invoke_reward passes everything
+    for p in params:
+        if p.kind in (p.VAR_KEYWORD, p.VAR_POSITIONAL):
+            continue
+        if p.default is inspect.Parameter.empty and p.name not in _AVAILABLE_REWARD_KWARGS:
+            return p.name
+    return None
+
+
+def _run_async(coro):
+    """Run an awaitable to completion from sync code, even inside a running loop."""
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(coro)
+    # Already inside a loop (rare for the worker): run in a fresh loop on a thread.
+    import concurrent.futures
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
+        return ex.submit(lambda: asyncio.run(coro)).result()
+
+
+def _call_dataset_getter(obj, method_name: str, *, seed: int, n: int = -1):
+    """Call a verifiers dataset getter, binding (n, seed) when it declares them.
+
+    verifiers exposes get_dataset/get_eval_dataset as get_X(n=-1, seed=0); some published envs
+    declare them WITHOUT defaults, so a no-arg call raised TypeError, swallowed into an empty
+    dataset (a paid run over no data). Bind ``n`` (default -1 = all rows — the adapter does its
+    own fixed subset selection; callers pass a positive cap to avoid materializing a huge split)
+    and the seed when the signature declares them; a genuine failure propagates (fail loudly)
+    instead of silently emptying the split."""
+    fn = getattr(obj, method_name, None)
+    if not callable(fn):
+        return None
+    try:
+        param_names = set(inspect.signature(fn).parameters)
+    except (TypeError, ValueError):
+        param_names = set()
+    kwargs = {}
+    if "n" in param_names:
+        kwargs["n"] = n
+    if "seed" in param_names:
+        kwargs["seed"] = seed
+    return fn(**kwargs)
+
+
+def _rows_to_list(ds) -> list[dict]:
+    if ds is None:
+        return []
+    try:
+        return [dict(r) for r in ds]
+    except Exception:
+        return list(ds)
+
+
+def _flatten_rubric(rubric) -> list[tuple]:
+    """Collect ``(func, weight)`` pairs from a rubric, recursing into ``RubricGroup``.
+
+    verifiers composes rubrics (e.g. a ``RubricGroup`` wrapping a ``MathRubric`` plus a
+    ``MultiTurnMonitorRubric``); the real reward funcs live on the *nested* rubrics while the
+    group's own ``funcs`` is empty. Flattening finds them all.
+    """
+    funcs = list(getattr(rubric, "funcs", None) or getattr(rubric, "reward_funcs", None) or [])
+    weights = list(
+        getattr(rubric, "weights", None) or getattr(rubric, "reward_weights", None) or []
+    )
+    if len(weights) < len(funcs):
+        weights += [1.0] * (len(funcs) - len(weights))
+    pairs = list(zip(funcs, weights, strict=False))
+    for sub in getattr(rubric, "rubrics", None) or []:
+        pairs.extend(_flatten_rubric(sub))
+    return pairs
+
+
+def _find_judge_rubric(rubric):
+    """Return the first ``JudgeRubric`` in a rubric tree (or None), for judge-arg injection."""
+    if rubric is None:
+        return None
+    try:
+        import verifiers as vf
+
+        judge_cls = getattr(vf, "JudgeRubric", None)
+    except ImportError:
+        judge_cls = None
+    if judge_cls is not None and isinstance(rubric, judge_cls):
+        return rubric
+    # Duck-type fallback: anything exposing a `judge` method + a judge_client attr.
+    if callable(getattr(rubric, "judge", None)) and hasattr(rubric, "judge_client"):
+        return rubric
+    for sub in getattr(rubric, "rubrics", None) or []:
+        found = _find_judge_rubric(sub)
+        if found is not None:
+            return found
+    return None
+
+
+def _judge_kwargs(judge_rubric) -> dict:
+    """The judge-related kwargs a reward func may declare, sourced from a JudgeRubric."""
+    if judge_rubric is None:
+        return {}
+    return {name: getattr(judge_rubric, name, None) for name in _JUDGE_KWARG_NAMES}
+
+
+def _invoke_reward(func, available: dict) -> float:
+    """Call a verifiers reward func passing only the kwargs it declares; await if async.
+
+    Exceptions PROPAGATE. ``scores_breakdown`` invokes this for *weighted* reward funcs, so an
+    exception here is a real (weighted) reward func genuinely failing (e.g. a JudgeRubric judge
+    raising on an API/rate-limit error, or a parse error on row data). Swallowing it as 0.0
+    would silently train/score on an all-zero signal and waste a paid run, so we fail loudly
+    instead. (Unweighted monitor funcs are skipped entirely by ``scores_breakdown``.)
+    """
+    try:
+        params = inspect.signature(func).parameters
+        if any(p.kind == p.VAR_KEYWORD for p in params.values()):
+            kwargs = dict(available)
+        else:
+            kwargs = {k: v for k, v in available.items() if k in params}
+    except (TypeError, ValueError):
+        kwargs = dict(available)
+    result = func(**kwargs)
+    if inspect.isawaitable(result):
+        result = _run_async(result)
+    return float(result or 0.0)
+
+
+def _substring_answer_score(completion: str, example: dict) -> float:
+    """Fallback reward for an env with no rubric: 1.0 iff the example ``answer`` is a
+    non-empty substring of the completion, else 0.0."""
+    answer = str(example.get("answer") or "")
+    return 1.0 if answer and answer in (completion or "") else 0.0
+
+
+def _unique_key(name: str, existing) -> str:
+    """``name``, or the first ``name_1``/``name_2``/… not already a key of ``existing``.
+
+    Probe for an unused exact key so two rubric funcs that share a name both survive — a
+    prefix/length heuristic can recompute a suffix that collides with an already-recorded
+    key (e.g. ``score`` vs ``score_detail``) and silently overwrite a scorer.
+    """
+    if name not in existing:
+        return name
+    i = 1
+    while f"{name}_{i}" in existing:
+        i += 1
+    return f"{name}_{i}"
+
+
+def _is_multi_turn(vf_env) -> bool:
+    """True for a tool/multi-turn verifiers env (NOT a plain SingleTurnEnv)."""
+    try:
+        import verifiers as vf
+    except ImportError:
+        return False
+    tool = getattr(vf, "ToolEnv", None)
+    multi = getattr(vf, "MultiTurnEnv", None)
+    single = getattr(vf, "SingleTurnEnv", None)
+    if tool is not None and isinstance(vf_env, tool):
+        return True
+    if multi is not None and isinstance(vf_env, multi):
+        # SingleTurnEnv subclasses MultiTurnEnv in verifiers; exempt it.
+        return single is None or not isinstance(vf_env, single)
+    return False
+
+
+def _is_tool_env(vf_env) -> bool:
+    """True for a verifiers ``ToolEnv`` or any subclass (Stateful/Sandbox/Python).
+
+    Tool envs expose Python tool callables; the worker hands those to TRL's
+    ``GRPOTrainer(tools=...)`` so TRL drives the tool-call loop natively (it owns generation,
+    tool execution, and assistant-only token masking). A *pure* ``MultiTurnEnv`` (env turns are
+    arbitrary content, e.g. a simulated user) is multi-turn but NOT a tool env, and takes the
+    ``rollout_func`` path instead."""
+    try:
+        import verifiers as vf
+    except ImportError:
+        return False
+    tool = getattr(vf, "ToolEnv", None)
+    return tool is not None and isinstance(vf_env, tool)

flash/envs/base.py

@@ -0,0 +1,52 @@
+"""Small, serializable environment interface for SFT/RL jobs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class Environment(Protocol):
+    id: str
+
+    def dataset(self) -> list[dict]:
+        """Return the training rows (the only split used; eval is on the serving side)."""
+
+    def prompt_messages(self, example: dict) -> list[dict]:
+        """Chat messages fed to the model for one example."""
+
+    def sft_completion(self, example: dict) -> list[dict]:
+        """Gold completion messages appended after the prompt for one SFT example — a multi-turn
+        trajectory or a single assistant turn."""
+
+    def reward(self, completion: str, example: dict, state: dict | None = None) -> float:
+        """Scalar RL reward for a completion."""
+
+    def grade(self, completion: str, example: dict, state: dict | None = None) -> bool:
+        """Boolean correctness scorer the reward can build on."""
+
+
+@dataclass
+class BaseEnvironment:
+    id: str
+
+    def dataset(self) -> list[dict]:
+        raise NotImplementedError
+
+    def prompt_messages(self, example: dict) -> list[dict]:
+        return [{"role": "user", "content": str(example.get("input") or "")}]
+
+    def sft_completion(self, example: dict) -> list[dict]:
+        # Single-turn default: one target assistant turn from the record's scalar ``output``.
+        # FreesoloEnvironment overrides this to support multi-turn target trajectories via the
+        # freesolo-sdk (``Environment.sft_completion`` -> ``datasets.target_messages``).
+        return [{"role": "assistant", "content": str(example.get("output") or "")}]
+
+    def reward(self, completion: str, example: dict, state: dict | None = None) -> float:
+        return 1.0 if self.grade(completion, example, state) else 0.0
+
+    def grade(self, completion: str, example: dict, state: dict | None = None) -> bool:
+        gold = str(example.get("output") or "").strip()
+        # A missing/empty output must NOT grade every completion correct (`"" in x` is
+        # always True) — treat it as unscorable -> incorrect.
+        return bool(gold) and gold in (completion or "")

flash/envs/registry.py

@@ -0,0 +1,62 @@
+"""Environment registry used by specs, worker, CLI, and server.
+
+Every managed run names a Freesolo SDK environment by Hub slug.
+The canonical generated environment entrypoint is ``environment.py:load_environment``.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from .._fileio import read_json_or_empty, secure_json_write
+from .base import Environment
+
+# Manifest of local Freesolo environment ids (written by `flash env install`).
+INSTALLED_MANIFEST = Path(
+    os.environ.get("FLASH_ENVS_MANIFEST", str(Path.home() / ".flash" / "envs.json"))
+)
+
+
+def load_installed_manifest() -> dict:
+    return read_json_or_empty(INSTALLED_MANIFEST)
+
+
+def list_installed_environments() -> list[str]:
+    """Freesolo environment ids recorded via `flash env install`."""
+    return sorted(load_installed_manifest())
+
+
+def record_installed_env(env_id: str, package: str, extras: dict | None = None) -> None:
+    manifest = load_installed_manifest()
+    manifest[env_id] = {"package": package, **(extras or {})}
+    # The manifest can hold a credentialed --extra-index-url, so write it with private perms.
+    secure_json_write(INSTALLED_MANIFEST, manifest)
+
+
+def worker_pip_for_env(env_id: str) -> list[str]:
+    """Pip deps the GPU worker needs to run a Freesolo environment."""
+    return ["freesolo"]
+
+
+def load_environment(
+    env_id: str, params: dict | None = None, resolved_sha: str | None = None
+) -> Environment:
+    """Load a Freesolo SDK environment and wrap it in Flash's protocol.
+
+    ``resolved_sha`` is the optional resolve-once hint (the control-plane-pinned commit sha for the
+    env's GitHub ref). None/"" preserves today's behavior — the adapter resolves the ref itself.
+    """
+    params = params or {}
+    from .adapter import load_freesolo_environment
+
+    if not env_id:
+        raise ValueError(
+            "no environment specified: set [environment] id to the id returned by "
+            "`flash env push --name <name>` (for example 'your-name/your-env')"
+        )
+    # User [environment.params] are freeform and forwarded verbatim to the SDK loader. The
+    # control-plane resolve-once pin is passed out-of-band as a POSITIONAL-ONLY argument, so a user
+    # param of ANY name (even "pinned_sha"/"resolved_sha") lands in **params and reaches the SDK
+    # unchanged — it can never bind to or disable the pin. None/"" keeps today's behavior.
+    return load_freesolo_environment(env_id, resolved_sha or None, **params)

flash/mcp/__init__.py

@@ -0,0 +1 @@
+"""MCP integration package."""

flash/mcp/server.py

@@ -0,0 +1,85 @@
+"""Minimal stdio MCP-style bridge for coding agents.
+
+This intentionally avoids a hard dependency on a specific MCP SDK while exposing
+the stable JSON tools that agents need. Requests are newline-delimited JSON:
+{"tool": "list_models", "args": {...}}.
+
+Run-lifecycle tools call the managed Flash control plane with the same stored
+credentials as the CLI (`flash login`); dry-run validation stays local.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Callable
+
+from flash.catalog import public_model_rows
+from flash.client import client_from_config
+from flash.client.runtime_secrets import runtime_secrets_from_local_env
+from flash.client.specs import spec_payload
+from flash.schema import spec_from_dict
+
+
+def list_models(args: dict) -> dict:
+    return {"models": public_model_rows()}
+
+
+def create_train_run(args: dict) -> dict:
+    spec = spec_from_dict(args, run_id=args.get("run_id"))
+    if args.get("dry_run"):
+        # Fully local: validate without credentials, a server, or a GPU.
+        return {"run_id": spec.run_id, "state": "dry_run", "spec": spec.to_dict()}
+    return client_from_config().create_run(
+        spec_payload(spec),
+        runtime_secrets=runtime_secrets_from_local_env(keys=spec.environment.secrets),
+    )
+
+
+def get_run_status(args: dict) -> dict:
+    return client_from_config().get_run(args["run_id"])
+
+
+def get_run_logs(args: dict) -> dict:
+    page = client_from_config().get_logs(args["run_id"], offset=int(args.get("offset", 0)))
+    return {"run_id": args["run_id"], **{k: page[k] for k in ("logs", "offset", "state")}}
+
+
+def deploy_adapter_tool(args: dict) -> dict:
+    return client_from_config().deploy(
+        args["run_id"],
+        dry_run=bool(args.get("dry_run", False)),
+    )
+
+
+TOOLS: dict[str, Callable[[dict], dict]] = {
+    "list_models": list_models,
+    "create_training_run": create_train_run,
+    "get_run_status": get_run_status,
+    "get_run_logs": get_run_logs,
+    "deploy_adapter": deploy_adapter_tool,
+}
+
+
+def handle(payload: dict) -> dict:
+    tool = payload.get("tool")
+    if tool not in TOOLS:
+        raise ValueError(f"unknown tool {tool!r}; choose one of {sorted(TOOLS)}")
+    return TOOLS[tool](payload.get("args") or {})
+
+
+def main() -> int:
+    for line in sys.stdin:
+        try:
+            line = line.strip()
+            if not line:
+                continue
+            response = {"ok": True, "result": handle(json.loads(line))}
+        except Exception as exc:
+            response = {"ok": False, "error": str(exc)}
+        print(json.dumps(response), flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

flash/providers/__init__.py

@@ -0,0 +1,59 @@
+"""Pluggable GPU substrates.
+
+The training worker (``flash.engine.worker``) reads a JobSpec from the environment, pulls code
+from the HF dataset repo, and streams artifacts/heartbeats/metrics back to it. The provider
+owns pricing, provisioning, polling, cancellation, and teardown.
+
+  runpod      serverless Flash endpoints (always on)
+  lambda      Lambda Cloud GPU instances (instance-based complement; iff LAMBDA_API_KEY set)
+  hyperstack  Hyperstack GPU VMs (instance-based complement; iff HYPERSTACK_API_KEY set)
+
+This module is the registry: ``get_provider(name)`` / ``PROVIDER_NAMES``.
+``allocator.allocate`` iterates the active provider list below; ``available_providers`` narrows
+it to the ones configured on THIS control plane (Lambda/Hyperstack are opt-in via their operator
+keys, so a box without them silently behaves exactly as the RunPod-only setup).
+"""
+
+from __future__ import annotations
+
+from functools import cache
+
+from flash.providers.base import Provider
+
+# Active provider order is also the tie-break preference (RunPod wins price ties, then Lambda).
+PROVIDER_NAMES: tuple[str, ...] = ("runpod", "lambda", "hyperstack")
+
+
+def get_provider(name: str) -> Provider:
+    """The ``Provider`` singleton for a registered name (raises on unknown)."""
+    # Normalize BEFORE the cache so "RunPod"/"runpod"/" runpod " share one cache entry.
+    return _get_provider((name or "").strip().lower())
+
+
+@cache
+def _get_provider(key: str) -> Provider:
+    if key == "runpod":
+        from flash.providers.runpod import PROVIDER
+
+        return PROVIDER
+    if key == "lambda":
+        from flash.providers.lambdalabs import PROVIDER
+
+        return PROVIDER
+    if key == "hyperstack":
+        from flash.providers.hyperstack import PROVIDER
+
+        return PROVIDER
+    raise KeyError(f"unknown provider {key!r} (known: {', '.join(PROVIDER_NAMES)})")
+
+
+def available_providers() -> tuple[str, ...]:
+    """Provider NAMES usable from this control plane right now: a provider is available when it
+    ``is_configured()`` (creds present). RunPod is always on; Lambda/Hyperstack join only when
+    their operator keys are present."""
+    return tuple(n for n in PROVIDER_NAMES if get_provider(n).is_configured())
+
+
+def configured_providers() -> list[Provider]:
+    """The ``Provider`` objects available right now (see ``available_providers``)."""
+    return [get_provider(n) for n in available_providers()]

flash/providers/_auth.py

@@ -0,0 +1,24 @@
+"""Shared operator-credential helpers for the GPU providers.
+
+Every provider authenticates the same way: a single API key read ONLY from an
+environment variable on the control-plane host (never config files, never shipped to
+workers). The per-provider ``auth.py`` modules wrap these with their own env-var name
+and error message.
+"""
+
+from __future__ import annotations
+
+import os
+
+
+def load_provider_key(env_var: str) -> str | None:
+    """Provider API key from ``env_var`` (operator configuration), or None."""
+    return os.environ.get(env_var) or None
+
+
+def ensure_provider_auth(env_var: str, missing_message: str) -> str:
+    """Return the provider key from ``env_var``; raise ``missing_message`` if unset."""
+    key = load_provider_key(env_var)
+    if not key:
+        raise RuntimeError(missing_message)
+    return key