AstraFlow

https://infini-ai-lab.github.io/astraflow/docs/
Admin
AstraFlow is an asynchronous RL training system for large reasoning and agentic models on...
Tokens:26,194
Snippets:315
Trust Score:9.7
Update:3 months ago
Show doc for...
Tokens:
Raw
### POST /notify_version Response Examples

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Example responses for successful updates, version skips, and failed pulls.

```python
# Successful update
{
    "ok": True,                     # envelope — handler ran without raising
    "result": {
        "ok": True,                 # manager — update succeeded
        "model_id": "model0",
        "version": 5,
        "pulled": True,
        "pull_result": {
            "mode": "full",         # or "delta"
            "shm_path": "/dev/shm/astraflow_weights/.../model0/model.safetensors",
        },
        "timing": {
            "pull_s": 1.05,
            "pause_s": 0.12,
            "load_s": 2.81,
            "resume_s": 0.08,
        },
    },
}

# Fast skip (version already loaded)
{
    "ok": True,
    "result": {
        "ok": True,
        "model_id": "model0",
        "pulled": False,
        "reason": "version=5 <= local=5",
    },
}

# Pull/load failed — envelope still ok, inner ok=False
{
    "ok": True,
    "result": {
        "ok": False,
        "model_id": "model0",
        "reason": "TCP pull timed out",
    },
}
```

--------------------------------

### Install AstraFlow Core and Dev Tools

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Installs core dependencies and development tools using uv.

```bash
uv pip install -e ".[dev]"
```

--------------------------------

### Verify Installation

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation-amd.html

Checks the versions of installed packages and confirms GPU availability and Flash Attention functionality.

```python
python -c "
import astraflow, torch, transformers
print(f'astraflow:    {astraflow.version.__version__}')
print(f'torch:        {torch.__version__}, hip={torch.version.hip}, GPUs: {torch.cuda.device_count()}')
print(f'transformers: {transformers.__version__}')
"

python -c "
import flash_attn, sglang
from flash_attn import flash_attn_varlen_func
print(f'flash-attn: {flash_attn.__version__} (varlen OK)')
print(f'sglang:     {sglang.__version__}')
"
```

--------------------------------

### Install Retrieval Server Dependencies

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/search.html

Sets up the conda environment and installs requirements for the local FAISS retrieval server.

```bash
conda create -n rag-retriever python=3.10 -y
conda activate rag-retriever
pip install -r astraEnv/ASearcher/requirements-rag-server.txt
```

--------------------------------

### Install Megatron training backend dependencies

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Installs Transformer Engine and apex from source to support Megatron-LM features. Requires nvcc to be available on the system PATH.

```bash
# nvcc must be on PATH (same CUDA toolchain as the flash-attn build above)
export CUDA_HOME=/usr/local/cuda-13.0
export PATH="$CUDA_HOME/bin:$PATH"
export NVTE_FRAMEWORK=pytorch

# Transformer Engine (required for the Megatron backend with TP/SP).
# The prebuilt transformer-engine wheels link libcublas.so.12 and do NOT load
# on a CUDA 13 install (ImportError: libcublas.so.12). Build TE from source
# against your CUDA 13 toolkit instead; nvidia-mathdx provides the build-time
# cuBLASDx / cuDNN frontend headers.
uv pip install nvidia-mathdx==25.6.0
uv pip install -v --no-build-isolation \
  "git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.13"

# apex (optional — Megatron falls back to Torch Norm / torch Adam if absent).
# APEX_CPP_EXT/APEX_CUDA_EXT select the fused kernels; FORCE_CUDA=1 builds them
# without a visible GPU.
git clone --depth 1 https://github.com/NVIDIA/apex.git /tmp/apex
cd /tmp/apex
FORCE_CUDA=1 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 \
  uv pip install -v --no-build-isolation .
cd - && rm -rf /tmp/apex
```

--------------------------------

### Install SGLang Inference Backend

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Installs the SGLang inference backend using project extras to ensure correct dependency overrides.

```bash
uv pip install -e ".[sglang]"
```

--------------------------------

### POST /eval_start

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Starts the evaluation process.

```APIDOC
## POST /eval_start

### Description
Initiates the evaluation window.

### Method
POST

### Endpoint
/eval_start
```

--------------------------------

### Install uv Package Manager

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Installs the uv package manager, which is required for handling project dependencies.

```bash
pip install -U "uv>=0.10"
```

--------------------------------

### Install Flash Attention

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Builds FlashAttention-2 from source, requiring CUDA 13 and sufficient temporary storage space.

```bash
# nvcc must be on PATH and match torch's CUDA (13.0 for torch 2.11+cu130)
export CUDA_HOME=/usr/local/cuda-13.0
export PATH="$CUDA_HOME/bin:$PATH"

# nvcc writes GBs of intermediate files to $TMPDIR. Point it at local scratch
# with plenty of space — NOT a small/NFS-quota'd home, or the build fails with
# "nvFatbin error: empty input" or "Disk quota exceeded" from truncated temps.
export TMPDIR=/tmp/fa-build && mkdir -p "$TMPDIR"

uv pip install "flash-attn==2.8.3" --no-build-isolation
```

--------------------------------

### Launch Qwen3-1.7B Training Run

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/quickstart.html

Executes the all-in-one script to start the AstraFlow orchestrator, RaaS server, and trainer processes.

```bash
bash examples/math/qwen3-1.7b-m2po-2gpus-full/scripts/run_qwen3-1.7b-m2po-2gpus-full.sh
```

--------------------------------

### Launch Retrieval Server

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/search.html

Starts the FAISS retrieval server on specified GPUs. Ensure the server is running before initiating training.

```bash
cd astraEnv/ASearcher
conda activate rag-retriever
export RAG_SERVER_ADDR_DIR=./tmp-log/rag_server_addrs
export PORT=7000
export USE_FAISS_GPU=1   # set 0 to disable GPU FAISS
bash scripts/launch_rag_server.sh 6,7
```

--------------------------------

### Action Protocol Examples

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/textcraft-recursive.html

Demonstrates the XML/JSON action surface used by the model to interact with the environment.

```xml
<action type="get_info">{"items": ["stick", "oak_planks"]}</action>
<action type="view_inventory">{}</action>
<action type="craft">{"ingredients": {"oak_log": 1}, "target": ["oak_planks", 4]}</action>
<action type="spawn">{"subtasks": [
  {"targets": {"oak_planks": 16}, "max_steps": 8},
  {"targets": {"stick": 8}, "max_steps": 5}
]}</action>
<action type="finish">{"message": "crafted 4 wooden_pickaxe"}</action>
```

--------------------------------

### Register Trainer with POST /ready

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Payload for signaling trainer readiness to AstraFlow. Omit model_id for single-model setups and recovered_version for fresh starts.

```python
{
    "train_batch_size": 512,
    "model_id": "model0",              # omit for single-model
    "sender_endpoint": "10.0.0.1:18861",
    "recovered_version": 0,            # omit for fresh start
}
```

```python
{"ok": True}
```

--------------------------------

### GET /get_capabilities

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/weight-manager.html

Retrieves supported strategies and delta readiness.

```APIDOC
## GET /get_capabilities

### Description
Returns the supported strategies and the current delta readiness status.

### Method
GET

### Endpoint
/get_capabilities
```

--------------------------------

### GET /batch

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Retrieves a data batch for the trainer.

```APIDOC
## GET /batch

### Description
Retrieves a data batch. This is called every step.

### Method
GET

### Endpoint
/batch
```

--------------------------------

### Launch Multi-Agent Math Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/multi-agent.html

Executes the training scripts for the Qwen3-8B actor-verifier multi-agent setup using either DAPO or DeepScaleR datasets.

```bash
# DAPO data (dynamic sampling on)
bash examples/math-multi-agent/qwen3-8b-actor-verifier-m2po-delta-dapo-data/scripts/run_qwen3-8b-actor-verifier-m2po-delta-dapo-data.sh

# DeepScaleR data (dynamic sampling on)
bash examples/math-multi-agent/qwen3-8b-actor-verifier-m2po-delta-deepscaler-data/scripts/run_qwen3-8b-actor-verifier-m2po-delta-deepscaler-data.sh
```

--------------------------------

### Implement AstraFlow Custom Trainer Loop

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

This example demonstrates the full lifecycle of a custom trainer, including initializing the sender agent, managing the training loop, and handling weight synchronization.

```python
from astraflow.train_worker.trainer.astraflow_client import AstraFlowClient

# --- Startup ---

# 1. Start the weight sender agent (reusable library).
#    This launches a subprocess that serves /get_buffer_info,
#    /register_sglang_instance, and /request_transfer.
sender = start_sender_agent(model)

# 2. Connect to AstraFlow and signal readiness.
client = AstraFlowClient(
    service_url="http://astraflow-host:8000",
)
client.initialize()          # wait for AstraFlow service to be ready
client.signal_ready(
    train_batch_size=512,
    sender_endpoint=sender.endpoint,   # e.g. "10.0.0.1:18861"
)

# --- Training Loop ---

version = 0
while training:
    # Pull batch (blocks until data is available)
    batch, buffer_stats = client.get_batch(version=version)

    # Your training logic
    loss = train_step(batch)

    # Copy updated weights to shared-memory buffer
    copy_weights_to_buffer()
    sender.notify_buffer_ready(version + 1)

    # Notify AstraFlow — it broadcasts to all RaaS instances
    version += 1
    should_eval = (version % eval_freq == 0)
    eval_results = client.notify_version(
        version=version,
        run_eval=should_eval,
    )

    if eval_results:
        log_eval(eval_results)

# --- Shutdown ---
client.drain_pending_notifications()
client.shutdown_service()
```

--------------------------------

### Setup AgentBench Environment

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/agentbench.html

Commands to create a dedicated conda environment and pull the necessary Docker images for AgentBench tasks.

```bash
cd astraEnv/AgentBench
conda create -n agent-bench python=3.9
conda activate agent-bench
pip install -r requirements.txt
conda install podman

# pull the image for the environment you want to train on
docker pull longinyu/agentbench-alfworld   # for ALFWorld
docker pull longinyu/agentbench-webshop    # for WebShop
```

--------------------------------

### GET /get_buffer_info

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Retrieves buffer information from the trainer.

```APIDOC
## GET /get_buffer_info

### Description
Retrieves buffer information. This is called once per RaaS.

### Method
GET

### Endpoint
/get_buffer_info
```

--------------------------------

### Verify AstraFlow installation

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Prints version information and CUDA availability for the core AstraFlow environment.

```python
python -c "
import astraflow, torch, transformers
print(f'astraflow:    {astraflow.version.__version__}')
print(f'torch:        {torch.__version__}, CUDA: {torch.cuda.is_available()}, GPUs: {torch.cuda.device_count()}')
print(f'transformers: {transformers.__version__}')
"
```

--------------------------------

### GET /get_version

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/weight-manager.html

Retrieves the current weight version.

```APIDOC
## GET /get_version

### Description
Returns the current weight version from the sender agent.

### Method
GET

### Endpoint
/get_version
```

--------------------------------

### GET /status

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Checks the readiness status of the RaaS service.

```APIDOC
## GET /status

### Description
Returns the readiness status of the RaaS service.

### Method
GET

### Endpoint
/status
```

--------------------------------

### GET /status

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/raas.html

Performs a readiness check and acts as a heartbeat for the orchestrator.

```APIDOC
## GET /status

### Description
Returns the current readiness state of the RaaS instance. Used as a heartbeat by the orchestrator.

### Method
GET

### Endpoint
/status

### Response
#### Success Response (200)
- **status** (string) - The current state: "ready", "idle", "starting", or "error"
```

--------------------------------

### POST /shutdown Request and Response

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Example request and response for the /shutdown endpoint.

```python
{}
```

```python
{
    "ok": True,
    "result": "shutting down",
}
```

--------------------------------

### Verify Flash Attention and SGLang

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Checks the installed versions of the Flash Attention and SGLang libraries.

```python
python -c "
import flash_attn, sglang
print(f'flash-attn: {flash_attn.__version__}')
print(f'sglang:     {sglang.__version__}')
"
```

--------------------------------

### Run Astraflow Docker Containers

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Commands to launch Astraflow containers for different training backends. Ensure the NVIDIA Container Toolkit is installed to support GPU access.

```bash
# FSDP backend (default) — covers most recipes
docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1

# Megatron-LM backend — adds Transformer Engine + apex (Step 5 above, pre-built in).
# Use this for `backend: megatron` (TP/PP/EP, MoE, large models).
docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1.megatron
```

--------------------------------

### Verify Offline Dataset Loading

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/math-offline.html

Example log output confirming that the AstraFlow service is successfully loading datasets from the local offline directory.

```text
Auto-derived offline_dir for dataset 'deepscaler': data-data/math/deepscaler
Loading DeepScaleR dataset from offline path: data-data/math/deepscaler
Auto-derived offline_dir for dataset 'aime24': data-data/math/aime24
... (same for aime25, amc, minerva, math500)
```

--------------------------------

### GET /availability

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/raas.html

Returns capacity information for load-balanced routing.

```APIDOC
## GET /availability

### Description
Returns capacity metrics for load-balanced routing.

### Method
GET

### Endpoint
/availability

### Response
#### Success Response (200)
- **available** (integer) - Available capacity
- **inflight** (integer) - Number of requests currently in flight
```

--------------------------------

### Verify Megatron backend installation

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Checks for the successful import of Transformer Engine and apex modules.

```python
python -c "
import transformer_engine.pytorch  # noqa: F401
print('transformer-engine: OK')
try:
    from apex.normalization import FusedLayerNorm  # noqa: F401
    print('apex: OK')
except ImportError:
    print('apex: not installed (Torch Norm fallback)')
"
```

--------------------------------

### Run Offline Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/math-offline.html

Execute the training script to start the AstraFlow service using the pre-downloaded local datasets.

```bash
bash examples/math/offline/qwen3-8b-m2po-full-offline/scripts/run_qwen3-8b-m2po-full-offline.sh
```

--------------------------------

### GET /get_buffer_info

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Retrieves buffer information from the trainer's weight sender.

```APIDOC
## GET /get_buffer_info

### Description
Endpoint served by the weight sender agent to provide information about the current weight buffer.

### Method
GET

### Endpoint
/get_buffer_info
```

--------------------------------

### GET /status Response

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Heartbeat endpoint returning current readiness status as raw JSON.

```json
{
  "status": "ready",
  "message": "optional details"
}
```

--------------------------------

### Retrieve Training Batch with GET /batch

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Request format for pulling training data and the corresponding response structure. The response is a direct payload without an envelope.

```http
GET /batch?model_id=model0&version=42
```

```python
{
    "batch": {
        "input_ids": Tensor,         # [batch, seq_len]
        "rewards": Tensor,           # [batch, seq_len]
        "logprobs": Tensor,          # [batch, seq_len]
        "loss_mask": Tensor,         # [batch, seq_len]
        # ... other fields depending on workflow
    },
    "buffer_stats": {
        "buffer/size": 1024,
        "buffer/staleness_mean": 1.2,
        # ... other buffer/filter metrics for wandb
    },
}
```

--------------------------------

### Launch Qwen3-1.7B Math Reasoning Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/math.html

Executes the training process for the 1.7B model using either delta or full weight transfer on 2 GPUs.

```bash
# delta weight transfer
bash examples/math/qwen3-1.7b-m2po-2gpus-delta/scripts/run_qwen3-1.7b-m2po-2gpus-delta.sh

# full weight transfer
bash examples/math/qwen3-1.7b-m2po-2gpus-full/scripts/run_qwen3-1.7b-m2po-2gpus-full.sh
```

--------------------------------

### Run Qwen2.5-7B-Instruct Search Recipes

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/search.html

Executes the training script for ASearcher agents using either delta or full weight transfer modes.

```bash
# delta weight transfer
bash examples/search/qwen2.5-7b-instruct-m2po-delta/scripts/run_qwen2.5-7b-instruct-m2po-delta.sh

# full weight transfer
bash examples/search/qwen2.5-7b-instruct-m2po-full/scripts/run_qwen2.5-7b-instruct-m2po-full.sh
```

--------------------------------

### Launch Qwen3-8B Math Reasoning Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/math.html

Executes the full-scale training process for the 8B model on 8 GPUs.

```bash
bash examples/math/qwen3-8b-m2po-full/scripts/run_qwen3-8b-m2po-full.sh
```

--------------------------------

### Launch WebShop Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/agentbench.html

Scripts to launch the WebShop environment server, AstraFlow service, RaaS inference server, and trainer using full or delta weight transfer.

```bash
# full weight transfer
bash examples/webshop/qwen2.5-7b-instruct-m2po-full/scripts/run_qwen2.5-7b-instruct-m2po-full.sh

# delta weight transfer
bash examples/webshop/qwen2.5-7b-instruct-m2po-delta/scripts/run_qwen2.5-7b-instruct-m2po-delta.sh
```

--------------------------------

### Download LiveCodeBench v5 Dataset

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/code.html

Manual download required for LiveCodeBench v5 evaluation data before launching training recipes.

```bash
huggingface-cli download inclusionAI/AReaL-boba-2-RL-Code \
  --repo-type dataset \
  --local-dir ./data-data/AReaL-boba-2-RL-Code
```

--------------------------------

### Download Math Datasets

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/math-offline.html

Execute this script from the repository root to populate the local dataset directory.

```bash
python examples/math/offline/download_math_datasets.py --root data-data/math
```

--------------------------------

### GET /batch

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Retrieves a data batch from the AstraFlow service.

```APIDOC
## GET /batch

### Description
Fetches a batch of data from AstraFlow for the training process.

### Method
GET

### Endpoint
/batch
```

--------------------------------

### GET /get_buffer_info

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/weight-manager.html

Retrieves buffer size and tensor metadata.

```APIDOC
## GET /get_buffer_info

### Description
Returns information regarding the buffer size and associated tensor metadata.

### Method
GET

### Endpoint
/get_buffer_info
```

--------------------------------

### GET /status

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Retrieves the current health status of the RaaS service.

```APIDOC
## GET /status

### Description
Performs a heartbeat check on the RaaS service.

### Method
GET

### Endpoint
/status
```

--------------------------------

### Download Qwen Model

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/textcraft-recursive.html

Downloads the required Qwen3-4B-Instruct-2507 model from Hugging Face.

```bash
huggingface-cli download Qwen/Qwen3-4B-Instruct-2507
```

--------------------------------

### GET /availability

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Checks the current capacity or availability of the RaaS service.

```APIDOC
## GET /availability

### Description
Returns the current capacity or availability status of the RaaS service.

### Method
GET

### Endpoint
/availability
```

--------------------------------

### Launch Single-Agent Baseline Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/multi-agent.html

Executes the training script for the single-agent Qwen3-8B baseline using the DAPO dataset.

```bash
bash examples/math-multi-agent/qwen3-8b-single-agent-m2po-delta-dapo-data/scripts/run_qwen3-8b-single-agent-m2po-delta-dapo-data.sh
```

--------------------------------

### POST /submit

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/raas.html

Submits a prompt for rollout generation.

```APIDOC
## POST /submit

### Description
Submits a single prompt for rollout generation.

### Method
POST

### Endpoint
/submit

### Response
#### Success Response (200)
- **task_id** (string) - The unique identifier for the submitted task
```

--------------------------------

### POST /notify_version Request Payload

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Example of a pickled Python dictionary for the /notify_version endpoint.

```python
{
    "model_id": "model0",
    "version": 5,
    "sender_endpoint": "10.0.0.1:19861",
}
```

--------------------------------

### Launch Multi-Agent Codegen + Verifier Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/code.html

Launch the all-in-one script for the two-agent recipe on the first node.

```bash
bash examples/code-multi-agent/qwen3-8b-codegen-verifier-m2po-full-2node/scripts/run_qwen3-8b-codegen-verifier-m2po-full-2node.sh
```

--------------------------------

### POST /pull Request Payload

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Example of a pickled Python dictionary for the /pull endpoint.

```python
{
    "max_items": 64,
    "timeout": 0.5,
}
```

--------------------------------

### GET /availability Response

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Capacity reporting endpoint used by the orchestrator for load balancing.

```json
{
  "available": 12,
  "inflight": 4,
  "max_concurrency": 16
}
```

--------------------------------

### POST /submit

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Enqueues a single rollout for execution.

```APIDOC
## POST /submit

### Description
Enqueues one rollout for processing.

### Method
POST

### Endpoint
/submit
```

--------------------------------

### Launch Single-Agent Qwen3-8B Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/code.html

Launch scripts for single-agent code generation using either delta or full weight transfer modes.

```bash
# delta weight transfer
bash examples/code/qwen3-8b-m2po-delta/scripts/run_qwen3-8b-m2po-delta.sh

# full weight transfer
bash examples/code/qwen3-8b-m2po-full/scripts/run_qwen3-8b-m2po-full.sh
```

--------------------------------

### POST /ready

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Registers the trainer with AstraFlow.

```APIDOC
## POST /ready

### Description
Registers the trainer with the AstraFlow system. This is called once.

### Method
POST

### Endpoint
/ready
```

--------------------------------

### Download Corpus and Build Index

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/search.html

Downloads the Wikipedia 2018 knowledge base and builds the FAISS index. Note that index building can take several hours.

```bash
cd astraEnv/ASearcher
conda activate rag-retriever
export WIKI2018_WORK_DIR=data/wiki2018
mkdir -p "$WIKI2018_WORK_DIR"
huggingface-cli download inclusionAI/ASearcher-Local-Knowledge \
  --repo-type dataset --local-dir "$WIKI2018_WORK_DIR" --local-dir-use-symlinks False
bash scripts/build_index.sh
```

--------------------------------

### Offload Model Parameters

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/weight-manager.html

The primary integration point for the Trainer to initiate weight offloading.

```python
wm.offload(model.named_parameters(), version, rank, world_size)
```

--------------------------------

### GET /batch

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Pulls a training batch from the AstraFlow buffer. This request blocks until sufficient data is available.

```APIDOC
## GET /batch

### Description
Pull a training batch. Blocks until sufficient data is available in the buffer.

### Method
GET

### Endpoint
/batch

### Parameters
#### Query Parameters
- **model_id** (str) - Optional - Which model’s data to pull (multi-model)
- **version** (int) - Optional - Trainer’s current version (used for multi-model barrier sync)

### Response
#### Success Response (200)
- **batch** (dict) - The training batch data
- **buffer_stats** (dict) - Metrics regarding the buffer state
```

--------------------------------

### WeightManager Offload Workflow

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/megatron-weight-sync.html

Visual representation of the offloading process from optimizer step to RaaS pull.

```text
optimizer.step()
   │
   ▼
WeightManager.offload(export_hf_named_params(...), version, ...)
   │   DP-head ranks write each (hf_name, tensor) into the INACTIVE
   │   half of the HF double buffer, in fixed order; non-heads barrier.
   ▼
notify_buffer_ready  ──►  sender swaps active/inactive
   │
   ▼                       sender._compute_delta() diffs HF-inactive vs
   │                       HF-active  → indices in HF space  ✓
   ▼
RaaS pulls full or delta (unchanged from FSDP)
```

--------------------------------

### Eval Start and End Requests and Responses

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Marks the boundaries of an evaluation window. Requests are empty pickled dictionaries.

```python
{}
```

```python
{"ok": True, "result": {}}  # ack only; payload is implementation-defined
```

--------------------------------

### Enable Delta Weight Transfer

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/delta-weight-transfer.html

Configure the trainer environment or RaaS YAML settings to enable delta weight transfer strategies.

```bash
export WEIGHT_TRANSFER_STRATEGIES="full,delta"
```

```yaml
raas:
  delta_full_sync_interval: 10   # full resync every 10 steps
```

--------------------------------

### POST /pull Response Envelope

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Example response structure for the /pull endpoint containing task results or failure states.

```python
{
    "ok": True,
    "result": [
        # Successful task — `result` holds whatever arun_episode returned
        {
            "task_id": 42,
            "result": {
                "input_ids": [...],       # list[int] or torch.Tensor
                "output_ids": [...],
                "output_versions": [5, 5, 5, 5],  # per-token weight version
                "rewards": [0.0, 0.0, 0.0, 1.0],
                # ... workflow-defined fields
            },
        },
        # Rejected sample — workflow returned None
        {"task_id": 43, "result": None},
        # Per-task failure — workflow raised; error rides inside `result`
        {
            "task_id": 44,
            "result": {"ok": False, "error": "RuntimeError('timeout')"},
        },
    ],
}
```

--------------------------------

### POST /submit

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Submits a sample to the RaaS service.

```APIDOC
## POST /submit

### Description
Submits a sample for processing.

### Method
POST

### Endpoint
/submit
```

--------------------------------

### Run Smoke Test

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation-amd.html

Executes a math recipe with a limited step count to verify the training pipeline.

```bash
SMOKE_STEPS=2 bash examples/math/qwen3-8b-m2po-full/scripts/run_qwen3-8b-m2po-full_amd.sh
```

--------------------------------

### Get Buffer Information Response

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Returns the model tensor layout and buffer length. Called once when a RaaS instance connects.

```json
{
  "single_buffer_length": 3489660928,
  "tensors_meta": [
    ["model.layers.0.self_attn.q_proj.weight", [[2048, 2048], "bfloat16"]],
    ["model.layers.0.self_attn.k_proj.weight", [[512, 2048], "bfloat16"]]
  ]
}
```

--------------------------------

### Launch ALFWorld Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/agentbench.html

Scripts to launch the ALFWorld environment server, AstraFlow service, RaaS inference server, and trainer using full or delta weight transfer.

```bash
# full weight transfer
bash examples/alfworld/qwen2.5-7b-instruct-m2po-full/scripts/run_qwen2.5-7b-instruct-m2po-full.sh

# delta weight transfer
bash examples/alfworld/qwen2.5-7b-instruct-m2po-delta/scripts/run_qwen2.5-7b-instruct-m2po-delta.sh
```

--------------------------------

### Visualize WeightManager Dependency Structure

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/weight-manager.html

Illustrates the unidirectional dependency flow where both Trainer and RaaS import from the WeightManager component.

```text
train_worker/ ──imports──► weight_manager/
raas/         ──imports──► weight_manager/transfer/

train_worker/ ✗ raas/      (no cross-dependency)
raas/ ✗ train_worker/      (no cross-dependency)
```

--------------------------------

### Configure Trainer Weight Transfer Ports

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/multi-agent-weight-transfer.html

Set unique HTTP ports for each trainer agent to facilitate weight transfer via environment variables.

```bash
# TCP weight-transfer ports (one per trainer)
export WEIGHT_TRANSFER_HTTP_PORT_MODEL0=19861
export WEIGHT_TRANSFER_HTTP_PORT_MODEL1=19862

# Trainer model0
WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" \
  torchrun ... --trainer trainer_model0

# Trainer model1
WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL1}" \
  torchrun ... --trainer trainer_model1
```

--------------------------------

### Prepare TextCraft Dataset

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/textcraft-recursive.html

Synthesizes the required training and validation tasks locally from the bundled recipe database.

```python
python -c "from astraflow.dataflow.dataset.textcraft import download_dataset; download_dataset()"
```

--------------------------------

### GET /get_buffer_info

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Queries the trainer's model tensor layout to allow RaaS to allocate its receive buffer. This is called once when a RaaS instance first connects.

```APIDOC
## GET /get_buffer_info

### Description
Queries the trainer's model tensor layout to allocate its receive buffer. Called once when a RaaS instance first connects.

### Method
GET

### Endpoint
/get_buffer_info

### Response
#### Success Response (200)
- **single_buffer_length** (integer) - The length of the buffer.
- **tensors_meta** (array) - List of tensor metadata including names and shapes.

#### Response Example
{
  "single_buffer_length": 3489660928,
  "tensors_meta": [
    ["model.layers.0.self_attn.q_proj.weight", [[2048, 2048], "bfloat16"]],
    ["model.layers.0.self_attn.k_proj.weight", [[512, 2048], "bfloat16"]]
  ]
}
```

--------------------------------

### Create Conda Environment

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation.html

Initializes a new Conda environment with Python 3.12.

```bash
conda create -n astraflow python=3.12 -y
conda activate astraflow
```

--------------------------------

### Implement RaaS Server Skeleton

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

A FastAPI-based skeleton for a RaaS manager, including methods for workflow registration, task submission, and model weight synchronization.

```python
# my_raas/server.py
import asyncio
import pickle
from typing import Any
from fastapi import FastAPI, Request
from fastapi.responses import Response, JSONResponse

class Manager:
    def __init__(self):
        self._status = "starting"
        self._workflows: dict[str, Any] = {}
        self._inflight: dict[int, asyncio.Task] = {}
        self._done: dict[int, Any] = {}
        self._next_id = 0
        self._weight_versions: dict[str, int] = {}
        self._weight_locks: dict[str, asyncio.Lock] = {}
        # ... your inference engine handles go here

    async def bootstrap(self):
        # launch inference engines, load base model, etc.
        self._status = "ready"

    def get_status(self) -> dict:
        return {"status": self._status}

    async def get_availability(self) -> dict:
        return {
            "available": max(0, MAX_CONCURRENCY - len(self._inflight)),
            "inflight": len(self._inflight),
        }

    def register_workflow(self, *, workflow_id, workflow_cls,
                          reward_fn=None, gconfig_overrides=None, **kw):
        cls = resolve_workflow(workflow_cls)
        rfn = resolve_reward(reward_fn) if reward_fn else None
        self._workflows[workflow_id] = cls(reward_fn=rfn, **kw)
        return {}

    async def submit(self, *, data, workflow_id="default") -> int:
        wf = self._workflows[workflow_id]
        tid = self._next_id; self._next_id += 1
        task = asyncio.create_task(wf.arun_episode(self.engine, data))
        task.add_done_callback(lambda t, i=tid: self._done.update(
            {i: (t.result() if not t.exception() else
                 {"ok": False, "error": repr(t.exception())})))
        self._inflight[tid] = task
        return tid

    async def pull_completed(self, *, max_items=256, timeout=0.0) -> list:
        if not self._done and timeout > 0:
            await asyncio.sleep(timeout)  # simple implementation
        out = []
        for tid in list(self._done)[:max_items]:
            out.append({"task_id": tid, "result": self._done.pop(tid)})
            self._inflight.pop(tid, None)
        return out

    async def notify_version(self, *, model_id, version, sender_endpoint):
        if version <= self._weight_versions.get(model_id, 0):
            return {"ok": True, "pulled": False}
        lock = self._weight_locks.setdefault(model_id, asyncio.Lock())
        async with lock:
            loop = asyncio.get_event_loop()
            shm_path = await loop.run_in_executor(
                None, self.weight_receiver.pull,
                sender_endpoint, model_id, version)
            await loop.run_in_executor(None, self.engine.pause_generation)
            await loop.run_in_executor(
                None, self.engine.load_weights_from_path, shm_path)
            await loop.run_in_executor(None, self.engine.continue_generation)
            self._weight_versions[model_id] = version
            self.engine.set_version(version)
        return {"ok": True, "model_id": model_id, "version": version}

    # ... eval_start / eval_end / eval_submit / eval_pull /
    # reset_training_engine / destroy ...

def _ok(result): return Response(
    pickle.dumps({"ok": True, "result": result}),
    media_type="application/octet-stream")
def _err(exc): return Response(
    pickle.dumps({"ok": False, "error": repr(exc)}),
    media_type="application/octet-stream", status_code=500)

def build_app(mgr: Manager) -> FastAPI:
    app = FastAPI()

    @app.get("/status")
    async def status(): return mgr.get_status()

    @app.get("/availability")
    async def avail(): return await mgr.get_availability()

    @app.post("/register_workflow")
    async def reg_wf(r: Request):
        try: p = pickle.loads(await r.body())
        except Exception as e: return _err(e)
        try: return _ok(mgr.register_workflow(**p))
        except Exception as e: return _err(e)

    # ... /submit, /pull, /notify_version, /shutdown ...
    # To support eval, also wire: /reset_training_engine,
    # /eval_start, /eval_end, /eval_submit, /eval_pull.
    return app
```

--------------------------------

### Run AstraFlow Container

Source: https://infini-ai-lab.github.io/astraflow/docs/get-started/installation-amd.html

Executes the AstraFlow container with necessary device access and resource limits for AMD GPU training.

```bash
docker run --rm \
  --device=/dev/kfd --device=/dev/dri --group-add video \
  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
  --ipc=host --network=host --shm-size=512g --ulimit nofile=65536:65536 \
  -v /home:/home \
  -e HF_TOKEN -e WANDB_API_KEY \
  -it astraflow:rocm
```

--------------------------------

### Version Barrier Coordination Flow

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/multi-agent-weight-transfer.html

Visual representation of the asynchronous weight loading and barrier synchronization sequence.

```text
model0: notify_version(v=5) ──► (async fire-and-forget load) ──► barrier (wait)
                                                                       │
model1: notify_version(v=5) ──► (async fire-and-forget load) ──► barrier (leader)
                                                                       │
                                                                       ▼
                                                       (eval? run reset + sync load per model + eval)
                                                                       │
                                                                       ▼
                                                       notify_all() → both return
```

--------------------------------

### POST /register_sglang_instance

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Registers an SGLang instance with the trainer.

```APIDOC
## POST /register_sglang_instance

### Description
Registers an SGLang instance. This is called once per RaaS.

### Method
POST

### Endpoint
/register_sglang_instance
```

--------------------------------

### POST /submit Request and Response

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-raas.html

Enqueues a rollout task. The handler should return immediately after creating an asynchronous task.

```python
{
    "data": {"prompt": "What is 2+2?", "answer": "4"},
    "workflow_id": "rlvr-math",
}
```

```python
{
    "ok": True,
    "result": {"task_id": 42},
}
```

--------------------------------

### Run Recursive Agent Training

Source: https://infini-ai-lab.github.io/astraflow/docs/recipes/textcraft-recursive.html

Executes the training script for the recursive agent using the provided shell configuration.

```bash
bash examples/textcraft-recursive-agent/qwen3-4b-recursive/scripts/run_qwen3-4b-recursive.sh
```

--------------------------------

### POST /ready

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Registers the custom trainer with the AstraFlow service.

```APIDOC
## POST /ready

### Description
Registers the custom trainer instance with the AstraFlow service to begin the training lifecycle.

### Method
POST

### Endpoint
/ready
```

--------------------------------

### POST /ready

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Signals that the trainer is ready to receive data. AstraFlow begins data acquisition only after both RaaS and the trainer have signaled readiness.

```APIDOC
## POST /ready

### Description
Signal that the trainer is ready to receive data. AstraFlow starts data acquisition only after both RaaS and trainer have signalled readiness.

### Method
POST

### Endpoint
/ready

### Request Body
- **train_batch_size** (int) - Required - Number of examples per training batch
- **model_id** (str) - Optional - Model identifier for multi-model setups
- **sender_endpoint** (str) - Required - "host:port" of the trainer’s weight sender HTTP server
- **recovered_version** (int) - Optional - Weight version to resume from after checkpoint recovery

### Response
#### Success Response (200)
- **ok** (bool) - Success status
```

--------------------------------

### POST /register_sglang_instance

Source: https://infini-ai-lab.github.io/astraflow/docs/architecture/custom-trainer.html

Registers an SGLang instance with the weight sender.

```APIDOC
## POST /register_sglang_instance

### Description
Registers a new SGLang instance to the trainer's weight sender agent.

### Method
POST

### Endpoint
/register_sglang_instance
```
Tokens:
Raw