### 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
{"items": ["stick", "oak_planks"]}
{}
{"ingredients": {"oak_log": 1}, "target": ["oak_planks", 4]}
{"subtasks": [
{"targets": {"oak_planks": 16}, "max_steps": 8},
{"targets": {"stick": 8}, "max_steps": 5}
]}
{"message": "crafted 4 wooden_pickaxe"}
```
--------------------------------
### 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
```