[example] add Fireworks Eval Only

examples/eval/fireworks/README.md

@@ -0,0 +1,56 @@
+# Eval-only generation with Fireworks
+
+Run SkyRL evaluation rollouts against a hosted [Fireworks](https://fireworks.ai) endpoint instead
+of a self-hosted vLLM server — no GPUs, no engine startup, no weight loading. The stock
+`SkyRLGymGenerator` and eval loop are unchanged; only the inference client is swapped.
+
+## Why
+
+Standing up vLLM is the expensive part of iterating on everything *around* the model: environment
+logic, reward functions, chat templates, dataset formatting, eval metrics. For that kind of quick
+prototyping you don't need your own server at all — you need an endpoint that behaves like one.
+This example points the eval-only entrypoint at Fireworks so you can validate a full
+generate → env step → reward → metrics loop from a laptop, then switch to the real vLLM-backed
+setup for training with no changes to your environment or data.
+
+Fireworks was chosen specifically (rather than any OpenAI-compatible API) because SkyRL's generator is
+token-in/token-out: it sends prompts as token ids and expects the generated token ids back.
+Fireworks' `/completions` accepts a pre-tokenized integer-array `prompt` and, with
+`return_token_ids=true`, returns the generated `token_ids` — so the tokens SkyRL records are
+exactly what the served model consumed and produced, with no re-tokenization drift.
+
+This is generation/eval-only: there is no weight sync to a hosted endpoint, so training still
+requires the vllm-backed entrypoints.
+
+## Quickstart
+
+```bash
+# 1. Prepare the GSM8K dataset
+uv run examples/train/gsm8k/gsm8k_dataset.py --output_dir $HOME/data/gsm8k
+
+# 2. Set your Fireworks API key
+export FIREWORKS_AI_API_KEY=<your_key_here>
+
+# 3. Run eval (defaults: gpt-oss-20b on Fireworks serverless)
+bash examples/eval/fireworks/run_eval_fireworks.sh
+```
+
+## Configuration conventions
+
+No custom config fields — the example reuses existing knobs:
+
+| Knob | Meaning here |
+|---|---|
+| `trainer.policy.model.path` | The served model's HF id (e.g. `openai/gpt-oss-20b`). Consumed **only** for the tokenizer in eval-only mode — it must be the served model's tokenizer, since prompts are sent as raw token ids (a mismatch silently degrades generations). |
+| `generator.inference_engine.served_model_name` | The Fireworks model id (e.g. `accounts/fireworks/models/gpt-oss-20b`), routed as the request `model`. |
+| `FIREWORKS_AI_API_KEY` (env) | API key, sent as `Authorization: Bearer`. |
+| `FIREWORKS_BASE_URL` (env, optional) | Server root override (no `/v1` suffix) for self-hosted OpenAI-compatible endpoints. |
+
+## Files
+
+- `fireworks_client.py` — `FireworksInferenceClient`, an `InferenceEngineInterface` implementation
+  over the `fireworks-ai` SDK. Converts the vLLM-shaped sampling params the stock eval path emits
+  to the subset Fireworks accepts, and no-ops the control plane (weight sync raises).
+- `main_eval_fireworks.py` — `FireworksEvalOnlyEntrypoint`, an `EvalOnlyEntrypoint` subclass that
+  overrides `get_inference_client()` to build the client above.
+- `run_eval_fireworks.sh` — GSM8K launcher (installs the `fireworks` uv extra).

examples/eval/fireworks/fireworks_client.py

@@ -0,0 +1,247 @@
+"""External Fireworks inference client (generation/eval only).
+
+Fireworks' OpenAI-compatible ``/completions`` accepts a pre-tokenized integer-array ``prompt``
+and, with ``return_token_ids=true``, returns the generated integer ``token_ids``. This gives
+token-in/token-out against an external endpoint with no re-tokenization drift, so the stock
+``SkyRLGymGenerator`` works unchanged.
+
+Built on the Fireworks v1 SDK (``fireworks-ai``, installed via the ``fireworks`` uv extra):
+``prompt`` accepts ``Iterable[Iterable[int]]``, ``return_token_ids`` is a first-class request
+param, and the response ``Choice`` declares ``token_ids``/``prompt_token_ids``. The SDK carries
+auth, retries with backoff, timeouts, and connection pooling. This module does not import vllm
+and has no control plane: wake/sleep/etc. are no-ops, weight sync raises.
+
+Sampling params arrive vLLM-shaped: the stock eval path emits them via
+``get_sampling_params_for_backend(backend="vllm", ...)``, and :func:`_to_fireworks_sampling_params`
+converts that dict to the subset Fireworks accepts at request time.
+
+Part of the standalone Fireworks eval example (``examples/eval/fireworks``); constructed by
+``FireworksEvalOnlyEntrypoint`` in ``main_eval_fireworks.py``.
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+from fireworks import AsyncFireworks
+from loguru import logger
+
+from skyrl.backends.skyrl_train.inference_servers.base import (
+    InferenceEngineInput,
+    InferenceEngineInterface,
+    InferenceEngineOutput,
+)
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+_GEN_EVAL_ONLY = (
+    "Fireworks is a generation/eval-only backend (external hosted endpoint with no weight sync). "
+    "Use the vllm-backed entrypoints for training."
+)
+
+# Server root of the Fireworks data plane. The SDK appends `/v1/completions` to an overridden
+# base_url, so this must NOT end in `/v1` (the constructor guards user-supplied values).
+DEFAULT_FIREWORKS_BASE_URL = "https://api.fireworks.ai/inference"
+
+# vLLM-only sampling keys the Fireworks /completions API rejects. `skip_special_tokens` and
+# `include_stop_str_in_output` are unnecessary: the client decodes responses locally from token
+# ids, and Fireworks includes a matched stop string's tokens in `token_ids` (only its `text`
+# field, which the client ignores, excludes it). `min_tokens` has no counterpart.
+_FIREWORKS_UNSUPPORTED_KEYS = ("min_tokens", "skip_special_tokens", "include_stop_str_in_output")
+
+
+def _to_fireworks_sampling_params(sampling_params: Dict[str, Any]) -> Dict[str, Any]:
+    """Convert a vLLM-shaped sampling-params dict to the subset Fireworks accepts.
+
+    Runs per request inside ``generate()``, and the vLLM emitter always includes the
+    vLLM-only keys, so those are dropped silently (see ``_FIREWORKS_UNSUPPORTED_KEYS`` for why
+    each is safe). Out-of-range ``top_k`` values are dropped with a warning (Fireworks accepts
+    ``0..100``); the ``top_k=-1`` / ``min_p<=0`` disable-sentinels and ``None`` values are
+    dropped silently (absence disables them; ``None`` would serialize as ``null`` via
+    ``extra_body``).
+    """
+    params = dict(sampling_params)
+    for key in _FIREWORKS_UNSUPPORTED_KEYS:
+        params.pop(key, None)
+    top_k = params.get("top_k")
+    if top_k is not None and not (isinstance(top_k, int) and 0 <= top_k <= 100):
+        if top_k != -1:  # -1 is the disable sentinel; absence disables top_k on Fireworks
+            logger.warning(f"Dropping sampling param `top_k={top_k}`: Fireworks accepts 0..100.")
+        del params["top_k"]
+    if params.get("min_p") is not None and params["min_p"] <= 0:
+        del params["min_p"]
+    return {key: value for key, value in params.items() if value is not None}
+
+
+class FireworksInferenceClient(InferenceEngineInterface):
+    def __init__(
+        self,
+        model_name: str,
+        tokenizer: "PreTrainedTokenizerBase",
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        max_retries: int = 3,
+        request_timeout: float = 600.0,
+    ):
+        """Args:
+        model_name: Fireworks model id used as the request ``model`` (e.g.
+            ``accounts/fireworks/models/gpt-oss-20b``).
+        tokenizer: The policy tokenizer; must be the served model's tokenizer since prompts
+            are sent as raw token ids.
+        base_url: Server root without ``/v1`` (defaults to the Fireworks data plane).
+        api_key: API key sent as ``Authorization: Bearer``. Always passed explicitly so the
+            SDK never falls back to the ``FIREWORKS_API_KEY`` env var; ``"EMPTY"`` placeholder
+            keeps keyless self-hosted endpoints constructible.
+        max_retries: SDK retry budget (backoff on 408/409/429/5xx and ``x-should-retry``).
+        request_timeout: Per-request timeout in seconds. Overrides the SDK's 60s default,
+            which is too short for long generations.
+        """
+        base_url = (base_url or DEFAULT_FIREWORKS_BASE_URL).rstrip("/")
+        if base_url.endswith("/v1"):
+            raise ValueError(
+                "base_url is the server root; the fireworks SDK appends /v1/completions. "
+                "Use e.g. https://api.fireworks.ai/inference, not .../inference/v1."
+            )
+        self._base_url = base_url
+        self._model_name = model_name
+        self._tokenizer = tokenizer
+        self._client = AsyncFireworks(
+            base_url=self._base_url,
+            api_key=api_key or "EMPTY",
+            max_retries=max_retries,
+            timeout=request_timeout,
+        )
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    def get_endpoint_url(self) -> str:
+        return self._base_url
+
+    async def generate(
+        self,
+        input_batch: InferenceEngineInput,
+        model: Optional[str] = None,
+    ) -> InferenceEngineOutput:
+        prompt_token_ids = input_batch.get("prompt_token_ids")
+        if prompt_token_ids is None:
+            raise ValueError("FireworksInferenceClient only accepts `prompt_token_ids`, not `prompts`.")
+        if input_batch.get("mm_features"):
+            raise NotImplementedError("FireworksInferenceClient does not support multi-modal features.")
+
+        sampling_params = _to_fireworks_sampling_params(input_batch.get("sampling_params", {}))
+        if sampling_params.get("n", 1) > 1:
+            raise ValueError("n > 1 is not supported. Use `config.generator.n_samples_per_prompt` instead.")
+        want_logprobs = sampling_params.get("logprobs") is not None
+
+        # model/prompt/return_token_ids are typed SDK params; the sampling dict rides extra_body
+        # (merged into the JSON request body) so additional_kwargs passthrough keeps working.
+        completion = await self._client.completions.create(
+            model=model or self._model_name,
+            prompt=prompt_token_ids,
+            return_token_ids=True,
+            extra_body=sampling_params,
+        )
+        if not completion.choices:
+            raise RuntimeError(f"Fireworks returned no choices: {completion!r}")
+        choices = sorted(completion.choices, key=lambda choice: choice.index)
+
+        response_ids: List[List[int]] = []
+        responses: List[str] = []
+        stop_reasons: List[str] = []
+        response_logprobs: List[Optional[List[float]]] = []
+        for choice in choices:
+            token_ids = choice.token_ids
+            # Re-encoding `choice.text` locally would silently reintroduce the retokenization
+            # drift this backend exists to avoid, so a missing field is a hard error.
+            assert token_ids is not None, (
+                f"Fireworks response missing `token_ids` for choice {choice.index} despite " "return_token_ids=true."
+            )
+            response_ids.append(list(token_ids))
+            # Decode locally to guarantee the InferenceEngineOutput invariant:
+            # tokenizer.decode(response_ids[i], skip_special_tokens=True) == responses[i].
+            responses.append(self._tokenizer.decode(token_ids, skip_special_tokens=True))
+            stop_reasons.append(choice.finish_reason or "stop")
+            if want_logprobs:
+                logprobs = self._extract_logprobs(choice)
+                # Silently emitting None here would surface far downstream as a confusing
+                # length-validation failure on GeneratorOutput["rollout_logprobs"].
+                if logprobs is None:
+                    raise RuntimeError(
+                        f"Sampling params requested logprobs but Fireworks returned none (or an "
+                        f"unrecognized shape) for choice {choice.index}. Set "
+                        f"generator.eval_sampling_params.logprobs=null (and "
+                        f"generator.sampling_params.logprobs=null) if logprobs are not needed."
+                    )
+                response_logprobs.append(logprobs)
+
+        return InferenceEngineOutput(
+            responses=responses,
+            response_ids=response_ids,
+            stop_reasons=stop_reasons,
+            response_logprobs=response_logprobs if want_logprobs else None,
+            prompt_logprobs=None,
+            rollout_expert_indices=None,
+        )
+
+    @staticmethod
+    def _extract_logprobs(choice: Any) -> Optional[List[float]]:
+        """Extract per-token logprobs from either Fireworks response shape.
+
+        ``choice.logprobs`` is a union of the classic completions shape (``LogProbs``, carries
+        ``token_logprobs``; what the live endpoint returns for integer ``logprobs``) and the
+        OpenAI chat-style shape (``NewLogProbs``, carries ``content`` items with ``.logprob``).
+        Null entries map to 0.0 rather than being dropped so the result stays aligned 1:1 with
+        the generated token ids (downstream validation asserts equal lengths).
+        """
+        logprobs = choice.logprobs
+        if logprobs is None:
+            return None
+        token_logprobs = getattr(logprobs, "token_logprobs", None)
+        if token_logprobs is not None:
+            return [logprob if logprob is not None else 0.0 for logprob in token_logprobs]
+        content = getattr(logprobs, "content", None)
+        if content is not None:
+            return [item.logprob if item.logprob is not None else 0.0 for item in content]
+        return None
+
+    async def completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client.completions.create(**request_payload.get("json", {}))
+        return response.model_dump()
+
+    async def chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client.chat.completions.create(**request_payload.get("json", {}))
+        return response.model_dump()
+
+    async def render_chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        raise NotImplementedError("render_chat_completion is not supported for the Fireworks backend.")
+
+    async def wake_up(self, *args: Any, **kwargs: Any):
+        return {}
+
+    async def sleep(self, *args: Any, **kwargs: Any):
+        return {}
+
+    async def reset_prefix_cache(self, reset_running_requests: bool = False):
+        return {}
+
+    async def pause_generation(self) -> None:
+        return
+
+    async def resume_generation(self) -> None:
+        return
+
+    async def finish_session(self, session_id: str) -> None:
+        return
+
+    async def teardown(self):
+        await self._client.close()
+
+    async def get_world_size(self) -> Tuple[int, int]:
+        raise NotImplementedError(_GEN_EVAL_ONLY)
+
+    async def init_weight_update_communicator(self, init_info):
+        raise NotImplementedError(_GEN_EVAL_ONLY)
+
+    async def update_named_weights(self, request):
+        raise NotImplementedError(_GEN_EVAL_ONLY)

examples/eval/fireworks/main_eval_fireworks.py

@@ -0,0 +1,83 @@
+"""Eval-only entrypoint that generates against the external Fireworks endpoint.
+
+Mirrors ``skyrl.train.entrypoints.main_generate`` but overrides ``get_inference_client`` to build
+:class:`~examples.eval.fireworks.fireworks_client.FireworksInferenceClient` — no local inference
+engines and no vLLM. Token-in/token-out is preserved (prompts are sent as raw token ids and
+Fireworks returns the generated ``token_ids``), so the stock ``SkyRLGymGenerator`` works unchanged.
+
+Configuration conventions (no custom config fields):
+  - ``trainer.policy.model.path`` is the served model's HF id (e.g. ``openai/gpt-oss-20b``). In
+    eval-only mode it is consumed solely for the tokenizer and dataset tokenization, and it MUST
+    be the served model's tokenizer: token ids are consumed raw by the server, so a mismatch
+    degrades generations silently instead of erroring.
+  - ``generator.inference_engine.served_model_name`` is the Fireworks model id (e.g.
+    ``accounts/fireworks/models/gpt-oss-20b``). ``resolve_policy_model_name`` routes it as the
+    request ``model``.
+  - ``FIREWORKS_AI_API_KEY`` (env) is the API key; ``FIREWORKS_BASE_URL`` (env, optional)
+    overrides the server root for self-hosted OpenAI-compatible endpoints (no ``/v1`` suffix).
+
+Requires the ``fireworks`` uv extra; see ``run_eval_fireworks.sh``.
+"""
+
+import asyncio
+import os
+import sys
+
+import ray
+from loguru import logger
+
+from skyrl.backends.skyrl_train.inference_servers.base import InferenceEngineInterface
+from skyrl.train.config import SkyRLTrainConfig
+from skyrl.train.entrypoints.main_generate import EvalOnlyEntrypoint
+from skyrl.train.utils.utils import initialize_ray, validate_generator_cfg
+
+
+class FireworksEvalOnlyEntrypoint(EvalOnlyEntrypoint):
+    def get_inference_client(self) -> InferenceEngineInterface:
+        try:
+            from examples.eval.fireworks.fireworks_client import FireworksInferenceClient
+        except ImportError as e:
+            raise ImportError(
+                "The Fireworks eval example requires the fireworks-ai SDK. Install the `fireworks` "
+                "extra, e.g. `uv run --isolated --extra fireworks ...`."
+            ) from e
+
+        ie_cfg = self.cfg.generator.inference_engine
+        assert ie_cfg.served_model_name, (
+            "Set generator.inference_engine.served_model_name to the Fireworks model id "
+            "(e.g. accounts/fireworks/models/gpt-oss-20b)"
+        )
+        api_key = os.environ.get("FIREWORKS_AI_API_KEY")
+        assert api_key, "Please set a fireworks ai api key to `FIREWORKS_AI_API_KEY` env-var."
+
+        # Surface the tokenizer/model pairing at startup: model.path must be the served model's
+        # tokenizer since prompts are sent as raw token ids (see module docstring).
+        logger.info(
+            f"Fireworks eval: served_model_name={ie_cfg.served_model_name}, "
+            f"tokenizer (trainer.policy.model.path)={self.cfg.trainer.policy.model.path}"
+        )
+        return FireworksInferenceClient(
+            model_name=ie_cfg.served_model_name,
+            tokenizer=self.tokenizer,
+            base_url=os.environ.get("FIREWORKS_BASE_URL"),
+            api_key=api_key,
+        )
+
+
+@ray.remote(num_cpus=1)
+def eval_entrypoint(cfg: SkyRLTrainConfig) -> dict:
+    exp = FireworksEvalOnlyEntrypoint(cfg)
+    inference_engine_client = exp.get_inference_client()
+    return asyncio.run(exp.run(inference_engine_client))
+
+
+def main() -> None:
+    cfg = SkyRLTrainConfig.from_cli_overrides(sys.argv[1:])
+    validate_generator_cfg(cfg)
+    initialize_ray(cfg)
+    metrics = ray.get(eval_entrypoint.remote(cfg))
+    logger.info(f"Metrics from eval only run: {metrics}")
+
+
+if __name__ == "__main__":
+    main()