primebeaker 0.4.0__tar.gz

primebeaker-0.4.0/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include examples *.md *.toml *.jsonl

primebeaker-0.4.0/PKG-INFO

@@ -0,0 +1,457 @@
+Metadata-Version: 2.4
+Name: primebeaker
+Version: 0.4.0
+Summary: Self-contained Prime-RL environments, TOML tooling, and Beaker launchers.
+Requires-Python: ~=3.12.0
+Description-Content-Type: text/markdown
+Requires-Dist: fire>=0.7.1
+Requires-Dist: pydantic>=2.7
+Requires-Dist: literegistry-tool-client<0.2,>=0.1.0
+Provides-Extra: evaluation
+Requires-Dist: jtc[evals,harness]==0.3.0; extra == "evaluation"
+Provides-Extra: watcher
+Requires-Dist: fastapi; extra == "watcher"
+Requires-Dist: jtc==0.3.0; extra == "watcher"
+Requires-Dist: PyYAML>=6.0; extra == "watcher"
+Requires-Dist: uvicorn; extra == "watcher"
+Requires-Dist: wandb; extra == "watcher"
+Provides-Extra: runtime
+Requires-Dist: aiohttp; extra == "runtime"
+Requires-Dist: datasets; extra == "runtime"
+Requires-Dist: httpx; extra == "runtime"
+Requires-Dist: literegistry[base_deployment,podman_beaker]==1.0.48; extra == "runtime"
+Requires-Dist: PyYAML>=6.0; extra == "runtime"
+Requires-Dist: uvicorn; extra == "runtime"
+Requires-Dist: verifiers[harbor]==0.2.1; extra == "runtime"
+Provides-Extra: test
+Requires-Dist: pytest>=8; extra == "test"
+Provides-Extra: full
+Requires-Dist: primebeaker[evaluation,runtime,watcher]; extra == "full"
+
+# PrimeBeaker
+
+PrimeBeaker is a standalone package for running Prime-RL SFT and RL on Beaker.
+It includes:
+
+- lossless SFT/RL TOML loading and editing;
+- single-node and multi-node Beaker preview/submission;
+- heterogeneous RL placement (mixed trainer/inference node plus dedicated inference nodes);
+- distributed multi-node SFT through torchrun/FSDP;
+- all 16 runnable Verifiers environments, using LiteRegistry tool clients;
+- package-local prompt templates, tool-call wire parsing, and tool schemas;
+- an immutable Prime-RL image catalog and reproducible Dockerfiles.
+- safe whole-topology checkpoint discovery and resume;
+- a thin CLI adapter over LiteRegistry's native Beaker service deployment.
+- managed Python evaluation scheduling, readiness, signal forwarding, and cleanup.
+- managed iterative rejection-sampling scheduling and resume, with JTC owning
+  only the judge/filter/compile loop.
+- typed standard, search-agent, and rubric-audit YAML lifecycle descriptions,
+  plus registry-backed scheduling for every packaged JTC workflow—including
+  all verifier, TMAX/Podman, inference, and rubric-generation variants.
+
+## Install
+
+Tool clients are provided exclusively by the standalone
+`literegistry-tool-client` package. Import them directly from
+`literegistry_tool_client`; PrimeBeaker does not re-export that API.
+Normal installation resolves the dependency from PyPI.
+
+For TOML and launch tooling:
+
+```bash
+pip install -e .
+```
+
+For environments, the LiteRegistry gateway, and multi-node RL:
+
+```bash
+pip install -e '.[runtime]'
+```
+
+For JTC evaluation scheduling as well:
+
+```bash
+pip install -e '.[runtime,evaluation]'
+primebeaker evaluation preview --config examples/configs/eval/python_only_smoke.yaml
+primebeaker evaluation submit --config examples/configs/eval/python_only_smoke.yaml
+```
+
+For the training/evaluation provenance watcher:
+
+```bash
+pip install -e '.[watcher]'
+primebeaker watcher status --database=watcher.sqlite3
+primebeaker watcher serve --database=watcher.sqlite3 --port=8790
+```
+
+The watcher implementation and lifecycle live in PrimeBeaker. It reads JTC
+training configs and evaluation outcomes but JTC never imports or launches it.
+
+PrimeBeaker owns these YAMLs and the complete service/Beaker lifecycle. JTC is
+the one-way application dependency: it supplies the workload contracts and
+Python workers in the evaluation image. Evaluation image construction lives in
+JTC at `docker/Dockerfile.eval` and installs released PyPI wheels only.
+
+The runtime extra installs LiteRegistry's Python launchers, but service stacks
+also require their separate container images in Beaker. Complete the runtime
+installation by building the official Dockerfiles from the exact installed
+LiteRegistry companion versions and importing the images into your workspace:
+
+```bash
+beaker account whoami
+docker version
+primebeaker services images install --workspace=ai2/oe-agents
+```
+
+The command downloads the matching source distributions, builds both the base
+services and Podman/mirror stacks, uploads every image, and prints the immutable
+Beaker IDs as `launcher_args`. Use those values for `services launch` or
+`services podman launch`. To install only the images required by the Podman
+terminal verifier, pass `--stack=podman`. A local LiteRegistry checkout can be
+used without downloading sources via `--source-root=/path/to/literegistry`.
+If the base deployment will run local search, also pass
+`--build-local-search --jtc-build-context=/path/to/jtc`; that image needs JTC's
+Lucene build assets and is therefore not part of the default build.
+
+This is an explicit post-install step because Python package installation must
+not silently mutate Docker or a remote Beaker workspace. It requires a running
+Docker daemon, an authenticated `beaker` CLI, network access to the configured
+Python and container indexes, and write access to the selected workspace.
+
+PrimeBeaker targets Python 3.12, matching the cataloged Prime-RL runtime. The
+`runtime` extra pins `verifiers[harbor]==0.2.1`, whose released PyPI wheel
+contains the client and environment APIs used by PrimeBeaker. No Git checkout is
+required for the runtime extra.
+
+The authenticated `beaker` executable, Docker, and a Prime-RL GPU image are
+external requirements. See [`src/primebeaker/images/README.md`](src/primebeaker/images/README.md) for immutable
+training-image locations and exact rebuild/publish instructions.
+
+## Tiny runnable examples
+
+[`examples/`](examples/README.md) contains two-step SFT and RL setups. Each
+uses four complete training records and two complete validation records, preserving
+the existing data format without bundling a corpus. Both are ready for the
+Fire-based `preview` and `submit` commands documented there.
+
+## Environments
+
+```bash
+python -m primebeaker.environments
+python -m primebeaker.environments tool-label
+```
+
+TOMLs use normal module paths such as:
+
+```text
+primebeaker.environments.jtc_tool_label_env:load_environment
+primebeaker.environments.jtc_reward_model_env:load_environment
+primebeaker.environments.jtc_search_agent_env:load_environment
+primebeaker.environments.podman_terminal_env:load_environment
+```
+
+The Podman terminal solver expects each dataset row to contain a prompt and an
+initial image under `original_image`, `image`, or `container_image`. It starts a
+fresh container directly from that image; saved trace fields are neither required
+nor read. The model solves the task in that live container using exactly two
+model-facing tools, `bash` and `submit`, and must finish with:
+
+```text
+echo TERMINAL_COMPLETE
+```
+
+Only after that marker succeeds, the environment uses the same Podman client to
+clear any stale or model-written reward file, execute the configured unit-test
+command against the model's final filesystem, capture its structured
+stdout/stderr/exit status, and read `/logs/verifier/reward.txt`. Its numeric
+value is clamped to `[0, 1]`; a missing or invalid file contributes exactly zero.
+Four independent rubric weights are configured directly in the Prime-RL TOML:
+
+```toml
+[orchestrator.train.source.legacy.args]
+reward_file_path = "/logs/verifier/reward.txt"
+test_command = "bash /tests/test.sh"
+test_timeout = 600
+reward_file_weight = 1.0
+termination_reward_weight = 0.1
+podman_failure_penalty_weight = 1.0
+fake_tool_penalty_weight = 0.1
+```
+
+The raw rubric values are respectively the captured unit-test score, `1` for a
+successful standalone completion command, `-1` for a Podman setup or execution
+failure, and `-1` per attempted tool name that the environment does not provide. A rollout that merely fails to terminate gets no
+termination credit rather than a separate missing-marker penalty. Failure phase, type,
+message, and an OOM indicator are retained in rollout state for diagnosis.
+
+The launcher rewrites copied legacy environment paths to the `primebeaker`
+namespace in a temporary runtime TOML; it never edits the source TOML.
+
+## Single-node launch
+
+```bash
+primebeaker sft preview --toml /weka/path/sft.toml
+
+primebeaker rl preview --toml /weka/path/rl.toml \
+  --env TERMINAL_SERVER_URL=http://service/terminal
+
+primebeaker rl submit --toml /weka/path/rl.toml \
+  --image beaker://IMMUTABLE_PRIMEBEAKER_IMAGE \
+  --workspace ai2/oe-agents \
+  --cluster ai2/holmes
+```
+
+## Multi-node RL
+
+The same command automatically selects the multi-node backend when the TOML
+contains `deployment.type = "multi_node"`:
+
+```bash
+primebeaker rl preview \
+  --toml /weka/path/multinode-rl.toml \
+  --image beaker://IMMUTABLE_PRIMEBEAKER_IMAGE \
+  --registry redis://registry-host:6379 \
+  --required-service terminal=8 \
+  --required-service localsearch:corpus=8
+```
+
+The generated Beaker task uses replicas, leader selection, synchronized start,
+host networking, and failure/preemption propagation. Replica 0 runs the
+trainer, orchestrator, global vLLM router, and local tool gateway. It may also
+run an inference GPU slice. Later replicas run the dedicated inference slices.
+All replicas share rendezvous state through the mounted Weka dataset and use
+one W&B run ID.
+
+The retained heterogeneous layout matches the extracted production launcher:
+exactly one trainer node and one Prime-RL inference replica, with any number of
+configured inference nodes. The TOML's `inference.parallel.tp`, `.dp`, and
+`api_server_count` must describe every exposed inference rank.
+
+Tool workers and Redis remain external services. `--required-service NAME=N`
+adds an exact startup barrier without bundling service deployment into the
+training package.
+
+## Joint service + RL runs
+
+The standalone `services`, `rl`, and `sft` commands remain available. For
+an RL run that owns its service stack, put both launch descriptions in one
+`primebeaker.run/v1` YAML and submit them together:
+
+```bash
+primebeaker run preview --config /weka/path/run.yaml
+primebeaker run launch --config /weka/path/run.yaml
+```
+
+The top-level `service` object is passed to LiteRegistry's native
+`BaseDeploymentConfig`. The `rl` object accepts the same launch settings as
+the standalone RL command: `toml`, `image`, `workspace`, `clusters`,
+`priority`, `min_runtime_hours`, `mount_path`, `dataset`,
+`working_dir`, `scratch_dir`, `home_dir`, `wandb_secret`, `hf_secret`,
+`wandb_run_id`, `wandb_entity`, `environment`, `secrets`,
+`setup_command`, `budget`, `description`,
+`rendezvous_timeout_seconds`, `registry`, `gateway_port`,
+`gateway_workers`, and `required_services`. YAML uses objects for
+`environment`, `secrets`, and `required_services`, and a list for
+`clusters`. A relative `toml` path is resolved relative to the run YAML.
+
+When `rl.registry` is omitted, it is derived from `service.registry` or
+`service.head_registry`; for example, `sqlite:///weka/run.sqlite3` becomes
+`head+sqlite:///weka/run.sqlite3`. This allows both experiments to be
+submitted immediately. Training retains its existing registry/service startup
+barriers.
+
+By default, `lifecycle.cleanup_service_on_training_exit` is true. The joint
+launcher injects the service experiment ID and a `BEAKER_TOKEN` secret into
+trainer replica 0, which stops the owned service experiment during trainer
+shutdown. Set `lifecycle.beaker_token_secret` when the Beaker workspace secret
+has a different name. If RL submission fails, the local launcher rolls the
+service experiment back immediately. See the
+[WebTerminal search-agent example](examples/search-agent-webterminal/README.md)
+for complete joint and independent launch forms.
+
+The independent SFT and RL commands share these placement/runtime flags:
+`image`, `workspace`, `cluster`, `priority`, `min_runtime_hours`,
+`mount_path`, `dataset`, `working_dir`, `scratch_dir`, `home_dir`,
+`wandb_secret`, `hf_secret`, `wandb_run_id`, `env`, `secret`,
+`setup_command`, `budget`, and `description`. Multi-node launches also
+accept `rendezvous_timeout_seconds`; RL additionally accepts `registry`,
+`wandb_entity`, `gateway_port`, `gateway_workers`, and
+`required_service`.
+
+## LiteRegistry services
+
+PrimeBeaker does not maintain a second service-stack implementation. The
+`services` command delegates directly to LiteRegistry's
+`literegistry-base-deployment` package, which owns gateway, Redis, Python,
+terminal, web search/fetch, cache, local-search, and vLLM deployment:
+
+```bash
+primebeaker services preview \
+  --head-registry=/weka/gfaria/registries/example \
+  --service-cluster=ai2/jupiter \
+  --python-replicas=2 \
+  --terminal-replicas=2 \
+  --web-search-replicas=2
+
+primebeaker services launch \
+  --head-registry=/weka/gfaria/registries/example \
+  --service-cluster=ai2/jupiter
+
+REGISTRY=redis://YOUR_REDIS_HOST:6379 primebeaker services podman preview \
+  --registry=redis://YOUR_REDIS_HOST:6379 \
+  --service-cluster=ai2/jupiter \
+  --podman-replicas=4 \
+  --docker-mirror-replicas=2
+
+REGISTRY=redis://YOUR_REDIS_HOST:6379 primebeaker services podman yaml preview \
+  --config=examples/podman-terminal/services.yaml
+```
+
+All service flags are passed into LiteRegistry's native
+`BaseDeploymentConfig`; PrimeBeaker adds no stack rendering or coordination
+logic. The `podman` subgroup similarly delegates to the native
+`literegistry-podman-beaker` package. Its YAML form is `services podman yaml`:
+it uses the same strict `primebeaker.services/v1` schema as `services yaml` but
+passes the `services` object to the Podman launcher. The same operations remain
+available through those two upstream executables.
+
+Both native launchers use `literegistry.coop.ports` for collision-safe
+dynamic ports and child supervision, plus `literegistry.coop.endpoints` for
+healthy endpoint publication and shutdown cleanup. The base launcher can publish through a `head+file://`, `head+sqlite://`, or
+`head+redis://` registry URI. The native Podman launcher instead requires a
+direct, persistent `redis://` or `rediss://` endpoint shared with training.
+Pass the applicable URI to `primebeaker rl preview|submit|resume --registry=...`.
+
+## Managed Python evaluations
+
+PrimeBeaker owns the infrastructure half of an evaluation lifecycle: it starts
+LiteRegistry's native service deployment, waits for its published Redis and
+gateway endpoints, checks the required service roster, runs one Python module
+inside a separate Beaker evaluation task, forwards termination signals, and
+stops the owned service experiment in a `finally`/process-exit cleanup path.
+
+The evaluation task command is always:
+
+```text
+python3 -m primebeaker.evaluation_worker
+```
+
+The application module and its arguments are validated data in a
+`PythonEvaluationRequest`; arbitrary commands and shell launchers are not
+accepted. JTC uses this API with its standard, search, audit, or generic
+`jtc.eval.workflow_worker` entry point, so JTC retains benchmark/workflow logic
+while PrimeBeaker owns scheduling and lifetime. The generic worker resolves a
+stable name from JTC's packaged workflow registry; it does not execute an
+arbitrary module or shell command. The image must contain released
+`primebeaker` and `jtc`
+wheels. The supported image build installs only exact versions from public
+PyPI; see
+[`src/primebeaker/images/README.md`](src/primebeaker/images/README.md).
+
+The generic Fire CLI accepts either an inline JSON object or a JSON file:
+
+```bash
+primebeaker evaluation preview --request-json=/weka/path/evaluation-request.json
+primebeaker evaluation submit --request-json=/weka/path/evaluation-request.json
+```
+
+Most users should use `primebeaker evaluation preview|submit --config=...` with
+a PrimeBeaker-owned YAML. A top-level registered JTC workflow name selects the
+generic worker, while `arguments` contains only that workflow's Python keyword
+arguments and `services`/`required_services` describe the managed topology.
+
+## Resume a multi-node RL experiment
+
+Resume always creates a fresh complete topology. It first requires one step
+to have every trainer rank shard, trainer `.metadata`, and the orchestrator's
+`progress.pt`. It refuses to submit while any source job is active, preserves
+the W&B run ID, moves stale broadcast handshakes aside, and writes an attempt
+manifest under `<output_dir>/primebeaker_resume/`.
+
+With no `--resume-step`, the latest complete common checkpoint is selected:
+
+```bash
+primebeaker rl resume \
+  --from-experiment=01SOURCE \
+  --dry-run
+
+primebeaker rl resume \
+  --from-experiment=01SOURCE \
+  --resume-step=220 \
+  --registry=head+file:///weka/gfaria/registries/example
+```
+
+The source Beaker experiment supplies the image, workspace, clusters, Weka
+mount, output directory, gateway settings, and training TOML. Any corresponding
+resume flag is an explicit override. A source without an external registry
+must be given `--registry`; `primebeaker services launch` can create that stack.
+
+## Multi-node SFT
+
+PrimeBeaker also translates Prime-RL's SLURM-oriented multi-node SFT topology
+to Beaker:
+
+```toml
+[deployment]
+type = "multi_node"
+num_nodes = 4
+gpus_per_node = 8
+```
+
+```bash
+primebeaker sft preview \
+  --toml /weka/path/multinode-sft.toml \
+  --image beaker://IMMUTABLE_PRIMEBEAKER_IMAGE
+```
+
+Each Beaker replica starts one torchrun node. Replica 0 publishes the rendezvous
+address and materializes the validated Prime-RL trainer config; all nodes join
+the same distributed job.
+
+## Output and clients
+
+Preview is non-mutating. `--write-spec` writes under
+`<scratch-dir>/primebeaker/beaker_experiments/`; submit writes the same spec and
+invokes `beaker experiment create`.
+
+`literegistry_tool_client` is the canonical API for terminal, Python, Podman,
+search, fetch, judge, reward-model, web-terminal, and submit clients. Service
+URLs and credentials remain explicit rather than being tied to JTC
+infrastructure.
+
+## Publishing to PyPI
+
+The `Publish to PyPI` GitHub Actions workflow builds and validates a source
+distribution and wheel with Python 3.12, then publishes them when a GitHub
+release is published. Running the workflow manually only builds and validates
+artifacts, so it can be used to check packaging before a release.
+
+First release (the project does not exist on PyPI yet):
+
+1. Create the `pypi` environment in the repository's GitHub Settings →
+   Environments.
+2. Sign in to PyPI and open [account publishing](https://pypi.org/manage/account/publishing/).
+   Add a pending GitHub publisher with these exact values:
+
+   | Field | Value |
+   | --- | --- |
+   | PyPI project name | `primebeaker` |
+   | Owner | `goncalorafaria` |
+   | Repository | `primebeaker` |
+   | Workflow filename | `publish.yml` |
+   | Environment | `pypi` |
+
+3. Commit and push the workflow and package source, then publish a GitHub
+   release tagged `v0.4.0` at that commit (matching the current package version).
+
+The first successful workflow upload creates the PyPI project automatically
+and converts the pending publisher into a permanent Trusted Publisher. No
+manual initial upload, PyPI API token, or GitHub secret is required. See
+[PyPI's first-project instructions](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/).
+
+For each release, update `[project].version` in `pyproject.toml`, commit and push
+it together with the workflow, then publish a GitHub release tagged `v<version>`
+(or `<version>`) at that commit. The workflow rejects a tag that does not match
+the package version. Choose a version that has not already been uploaded to
+PyPI; existing releases cannot be overwritten. Draft releases and tag pushes
+alone do not publish a package.