### 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 ```