Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 73133cdb2d | |||
| abb23bdeef |
@@ -1,29 +0,0 @@
|
||||
name: Badge - ChartQA
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - ChartQA
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-chartqa.yml', label: 'chartqa', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -10,8 +10,6 @@ on:
|
||||
- Examples - Tinker
|
||||
- Examples - Azure
|
||||
- Examples - Claude Code
|
||||
- Examples - RAG
|
||||
- Examples - ChartQA
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
@@ -39,7 +37,5 @@ jobs:
|
||||
{ workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-claude-code.yml', label: 'examples-claude-code.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-rag.yml', label: 'examples-rag.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-chartqa.yml', label: 'examples-chartqa.stable', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
|
||||
@@ -7,8 +7,6 @@ on:
|
||||
- Examples - Spider
|
||||
- Examples - APO
|
||||
- Examples - Unsloth
|
||||
- Examples - RAG
|
||||
- Examples - Claude Code
|
||||
- GPU Test
|
||||
types: [completed]
|
||||
|
||||
@@ -34,8 +32,6 @@ jobs:
|
||||
{ workflow: 'examples-spider.yml', label: 'spider.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-apo.yml', label: 'apo.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-unsloth.yml', label: 'unsloth.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-claude-code.yml', label: 'claude-code.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-rag.yml', label: 'rag.latest', variants: ['latest'] },
|
||||
{ workflow: 'tests-full.yml', label: 'tests-full.latest', variants: ['latest'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
name: Badge - RAG
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - RAG
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-rag.yml', label: 'rag', variants: ['legacy', 'stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -24,6 +24,6 @@ jobs:
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-spider.yml', label: 'spider', variants: ['stable'] },
|
||||
{ workflow: 'examples-spider.yml', label: 'spider', variants: ['stable', 'legacy'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
|
||||
+29
-283
@@ -3,15 +3,12 @@ permissions:
|
||||
contents: read
|
||||
on:
|
||||
workflow_dispatch:
|
||||
schedule:
|
||||
# Every Monday and Thursday at 3 AM UTC+8
|
||||
- cron: '0 19 * * 0,3'
|
||||
|
||||
jobs:
|
||||
benchmark:
|
||||
name: ${{ matrix.workload.kind }} (${{ matrix.backend.id }}, ${{ matrix.workload.display }})
|
||||
runs-on: ${{ matrix.workload.runner }}
|
||||
timeout-minutes: ${{ matrix.workload.timeout }}
|
||||
name: Benchmark (${{ matrix.backend.id }}, ${{ matrix.scenario.display }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
@@ -20,15 +17,10 @@ jobs:
|
||||
compose_file: compose.prometheus-memory-store.yml
|
||||
- id: mongo
|
||||
compose_file: compose.prometheus-mongo-store.yml
|
||||
workload:
|
||||
- id: scenario-minimal-scale
|
||||
scenario:
|
||||
- id: minimal-production
|
||||
display: Minimal production scale
|
||||
kind: scenario
|
||||
store_workers: 4
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 45
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 4096
|
||||
@@ -36,14 +28,9 @@ jobs:
|
||||
--n-runners 32
|
||||
--max-rounds 6
|
||||
--sleep-seconds 0.5
|
||||
- id: scenario-medium-scale
|
||||
- id: medium-production
|
||||
display: Medium production scale
|
||||
kind: scenario
|
||||
store_workers: 16
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 45
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 10000
|
||||
@@ -51,75 +38,40 @@ jobs:
|
||||
--n-runners 100
|
||||
--max-rounds 10
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-midhigh-scale
|
||||
display: Mid-high production scale
|
||||
kind: scenario
|
||||
store_workers: 24
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 20000
|
||||
--batch-size 2048
|
||||
--n-runners 256
|
||||
--max-rounds 8
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-large-batch
|
||||
- id: large-batch
|
||||
display: Large batch waves
|
||||
kind: scenario
|
||||
store_workers: 64
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu-high
|
||||
timeout: 120
|
||||
store_workers: 32
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 50000
|
||||
--total-tasks 100000
|
||||
--batch-size 8192
|
||||
--n-runners 256
|
||||
--max-rounds 6
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-long-queues
|
||||
- id: long-queues
|
||||
display: Long rollout queues
|
||||
kind: scenario
|
||||
store_workers: 48
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 120
|
||||
store_workers: 32
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
--total-tasks 50000
|
||||
--total-tasks 100000
|
||||
--batch-size 1024
|
||||
--n-runners 256
|
||||
--remaining-tasks 4096
|
||||
--max-rounds 4
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-high-concurrency
|
||||
- id: high-concurrency
|
||||
display: High-throughput concurrent requests
|
||||
kind: scenario
|
||||
store_workers: 96
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu-high
|
||||
timeout: 120
|
||||
store_workers: 32
|
||||
args: >-
|
||||
--mode single
|
||||
--total-tasks 50000
|
||||
--total-tasks 100000
|
||||
--concurrency 2048
|
||||
--n-runners 256
|
||||
--max-rounds 2
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-heavy-traces
|
||||
- id: heavy-traces
|
||||
display: Heavy rollouts with deep traces
|
||||
kind: scenario
|
||||
store_workers: 64
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu-high
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
--total-tasks 10000
|
||||
@@ -128,65 +80,15 @@ jobs:
|
||||
--n-runners 512
|
||||
--max-rounds 20
|
||||
--sleep-seconds 1.0
|
||||
|
||||
- id: micro-worker
|
||||
display: Update worker
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: worker
|
||||
- id: micro-dequeue-empty
|
||||
display: Dequeue empty
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-empty
|
||||
- id: micro-rollout
|
||||
display: Rollout + span
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: rollout
|
||||
- id: micro-dequeue-update-attempt
|
||||
display: Dequeue + update attempt
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-update-attempt
|
||||
- id: micro-dequeue-only
|
||||
display: Dequeue only
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-only
|
||||
- id: micro-metrics
|
||||
display: Multi-metric fan-out
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 15
|
||||
cli: metrics
|
||||
env:
|
||||
PYTHONUNBUFFERED: "1"
|
||||
STORE_URL: http://localhost:4747
|
||||
STORE_API_URL: http://localhost:4747/v1/agl
|
||||
PROM_URL: http://localhost:9090
|
||||
GITHUB_ACTIONS_TIMEOUT_MINUTES: ${{ matrix.workload.timeout }}
|
||||
WORKLOAD_KIND: ${{ matrix.workload.kind }}
|
||||
WORKLOAD_ID: ${{ matrix.workload.id }}
|
||||
SCENARIO_ID: ${{ matrix.scenario.id }}
|
||||
BACKEND_ID: ${{ matrix.backend.id }}
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
ARTIFACT_DIR: artifacts/${{ matrix.scenario.id }}-${{ matrix.backend.id }}
|
||||
COMPOSE_FILE: ${{ matrix.backend.compose_file }}
|
||||
AGL_STORE_N_WORKERS: ${{ matrix.workload.store_workers }}
|
||||
ANALYSIS_FILE: ${{ format('analysis-{0}.log', matrix.workload.id) }}
|
||||
SUMMARY_FILE: ${{ format('summary-{0}.log', matrix.workload.id) }}
|
||||
PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
ARTIFACT_NAME: ${{ format('{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
AGL_STORE_N_WORKERS: ${{ matrix.scenario.store_workers }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
@@ -220,65 +122,37 @@ jobs:
|
||||
set -euo pipefail
|
||||
for attempt in {1..60}; do
|
||||
if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
|
||||
sleep 1
|
||||
curl -fsS "$STORE_API_URL/rollouts" # Warm up the scraper
|
||||
sleep 15 # Allow some time for the baseline metrics to be established
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "Store did not become ready in time" >&2
|
||||
# show logs for debugging
|
||||
cd docker && docker compose -f "$COMPOSE_FILE" logs app
|
||||
exit 1
|
||||
|
||||
- name: Prepare artifact directory
|
||||
run: mkdir -p "$ARTIFACT_DIR"
|
||||
|
||||
- name: Record workload start
|
||||
- name: Record benchmark start
|
||||
run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: (Scenario) Run ${{ matrix.workload.display }} workload
|
||||
if: ${{ matrix.workload.kind == 'scenario' }}
|
||||
- name: Run ${{ matrix.scenario.display }} workload
|
||||
run: |
|
||||
set -euo pipefail
|
||||
uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
|
||||
--store-url "$STORE_URL" \
|
||||
${{ matrix.workload.args }}
|
||||
${{ matrix.scenario.args }}
|
||||
|
||||
- name: (Micro) Run ${{ matrix.workload.display }}
|
||||
if: ${{ matrix.workload.kind == 'micro' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
uv run --locked --no-sync python -m tests.benchmark.micro_benchmark \
|
||||
--store-url "$STORE_URL" \
|
||||
--summary-file "$ARTIFACT_DIR/$SUMMARY_FILE" \
|
||||
"${{ matrix.workload.cli }}" | tee "$ARTIFACT_DIR/${{ matrix.workload.id }}.txt"
|
||||
|
||||
- name: Record workload end
|
||||
- name: Record benchmark end
|
||||
if: ${{ always() }}
|
||||
run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Show micro benchmark summary
|
||||
if: ${{ always() && matrix.workload.kind == 'micro' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
summary_file="$ARTIFACT_DIR/$SUMMARY_FILE"
|
||||
if [ -f "$summary_file" ]; then
|
||||
echo "Micro benchmark summary ($WORKLOAD_ID/$BACKEND_ID):"
|
||||
cat "$summary_file"
|
||||
else
|
||||
echo "Summary file not found: $summary_file"
|
||||
fi
|
||||
|
||||
- name: Run workload analysis
|
||||
- name: Run benchmark analysis
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
|
||||
echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/$ANALYSIS_FILE"
|
||||
echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/analysis.txt"
|
||||
exit 1
|
||||
fi
|
||||
uv run --locked --no-sync python -m tests.benchmark.analysis \
|
||||
@@ -286,22 +160,7 @@ jobs:
|
||||
--store-url "$STORE_API_URL" \
|
||||
--start "$BENCHMARK_START" \
|
||||
--end "$BENCHMARK_END" \
|
||||
| tee "$ARTIFACT_DIR/$ANALYSIS_FILE"
|
||||
|
||||
- name: Collect docker logs
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
cd docker
|
||||
readarray -t services < <(docker compose -f "$COMPOSE_FILE" config --services)
|
||||
if [ "${#services[@]}" -eq 0 ]; then
|
||||
echo "No services defined in compose file."
|
||||
exit 0
|
||||
fi
|
||||
for service in "${services[@]}"; do
|
||||
docker compose -f "$COMPOSE_FILE" logs "$service" > "../$ARTIFACT_DIR/docker-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
|
||||
done
|
||||
| tee "$ARTIFACT_DIR/analysis.txt"
|
||||
|
||||
- name: Stop ${{ matrix.backend.id }} Prometheus stack
|
||||
if: ${{ always() }}
|
||||
@@ -316,126 +175,13 @@ jobs:
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
if [ -d docker/data/prometheus ]; then
|
||||
tar -C docker/data -czf "$ARTIFACT_DIR/${PROM_ARCHIVE_BASENAME}.tar.gz" prometheus
|
||||
tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-${SCENARIO_ID}-${BACKEND_ID}.tar.gz" prometheus
|
||||
fi
|
||||
|
||||
- name: Upload workload artifacts
|
||||
- name: Upload benchmark artifacts
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: ${{ env.ARTIFACT_NAME }}
|
||||
path: ${{ env.ARTIFACT_DIR }}
|
||||
if-no-files-found: error
|
||||
|
||||
collection-benchmarks:
|
||||
name: collection (${{ matrix.backend.id }}, ${{ matrix.workload.id }})
|
||||
runs-on: ${{ matrix.backend.runner }}
|
||||
timeout-minutes: 15
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
backend:
|
||||
- id: memory
|
||||
needs_mongo: false
|
||||
runner: ubuntu-latest
|
||||
- id: mongo
|
||||
needs_mongo: true
|
||||
runner: ubuntu-latest
|
||||
workload:
|
||||
- id: high-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 2048
|
||||
type: insert
|
||||
- id: medium-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 128
|
||||
type: insert
|
||||
- id: low-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 4
|
||||
type: insert
|
||||
- id: high-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 2048
|
||||
type: dequeue
|
||||
- id: medium-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 128
|
||||
type: dequeue
|
||||
- id: low-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 4
|
||||
type: dequeue
|
||||
env:
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.backend.id, matrix.workload.id) }}
|
||||
SUMMARY_FILE: ${{ format('artifacts/{0}-{1}/summary-{0}-{1}.jsonl', matrix.backend.id, matrix.workload.id) }}
|
||||
ARTIFACT_NAME: ${{ format('collections-{0}-{1}', matrix.backend.id, matrix.workload.id) }}
|
||||
MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --extra mongo --group core-stable --group dev
|
||||
|
||||
- name: Launch MongoDB
|
||||
if: ${{ matrix.backend.needs_mongo }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f compose.mongo.yml down -v || true
|
||||
docker compose -f compose.mongo.yml up -d --quiet-pull
|
||||
for attempt in {1..60}; do
|
||||
if docker compose -f compose.mongo.yml exec -T mongo mongosh --quiet --eval 'db.runCommand({ping:1})' >/dev/null 2>&1; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "MongoDB did not become ready in time" >&2
|
||||
docker compose -f compose.mongo.yml logs mongo
|
||||
exit 1
|
||||
|
||||
- name: Run collection benchmark
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
echo "Running collection benchmark (backend=${{ matrix.backend.id }}, workload=${{ matrix.workload.id }})"
|
||||
uv run --locked --no-sync python -m tests.benchmark.collection_benchmark \
|
||||
"${{ matrix.workload.type }}" \
|
||||
--backend "${{ matrix.backend.id }}" \
|
||||
--total-tasks "${{ matrix.workload.total_tasks }}" \
|
||||
--concurrency "${{ matrix.workload.concurrency }}" \
|
||||
--task-prefix "${{ matrix.backend.id }}-${{ matrix.workload.id }}" \
|
||||
--summary-file "$SUMMARY_FILE" \
|
||||
--mongo-uri "$MONGO_URI" \
|
||||
--mongo-database agentlightning_collection_bench
|
||||
|
||||
- name: Show collection benchmark summary
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ -f "$SUMMARY_FILE" ]; then
|
||||
echo "Collection benchmark summary (${{ matrix.backend.id }}):"
|
||||
cat "$SUMMARY_FILE"
|
||||
else
|
||||
echo "Summary file not found: $SUMMARY_FILE"
|
||||
fi
|
||||
|
||||
- name: Stop MongoDB
|
||||
if: ${{ always() && matrix.backend.needs_mongo }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f compose.mongo.yml down -v || true
|
||||
|
||||
- name: Upload collection artifacts
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: ${{ env.ARTIFACT_NAME }}
|
||||
name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
|
||||
path: ${{ env.ARTIFACT_DIR }}
|
||||
if-no-files-found: error
|
||||
|
||||
@@ -171,12 +171,12 @@ jobs:
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-stable
|
||||
--group dev --group experiment --group agents --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-${{ matrix.setup-script }}
|
||||
--group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
@@ -270,104 +270,6 @@ jobs:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Setup Docker environments
|
||||
run: ./scripts/mongodb_docker_run.sh
|
||||
shell: bash
|
||||
|
||||
- name: Training with MongoDB
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --mongo-uri mongodb://localhost:27017/?replicaSet=rs0
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_mongo
|
||||
|
||||
- name: Validate training with MongoDB
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_mongo.outputs.project_name }} ${{ steps.calc_x_train_mongo.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with LoRA
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --lora
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_lora
|
||||
if: matrix.setup-script != 'legacy'
|
||||
|
||||
- name: Validate training with LoRA
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_lora.outputs.project_name }} ${{ steps.calc_x_train_lora.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
if: matrix.setup-script != 'legacy'
|
||||
|
||||
- name: Training with trajectory level aggregation
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --trajectory-level
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_trajectory_level
|
||||
|
||||
- name: Validate training with trajectory level aggregation
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_trajectory_level.outputs.project_name }} ${{ steps.calc_x_train_trajectory_level.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with Weave
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --weave
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_weave
|
||||
|
||||
- name: Validate training with Weave
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_weave.outputs.project_name }} ${{ steps.calc_x_train_weave.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with external store
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
@@ -1,168 +0,0 @@
|
||||
name: Examples - ChartQA
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 6 AM UTC+8
|
||||
- cron: "0 22 * * *"
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-chartqa, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'ChartQA - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('ChartQA - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
chartqa:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-chartqa' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: ChartQA (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group image --group langchain --group vllm-0-10-2 --group torch-gpu-stable
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: dependencies-chartqa-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare ChartQA dataset
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/chartqa
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1fWRt9hehg8_uV7BDWSCwKTycM60JcmGN/view?usp=sharing" -O chartqa-data.zip
|
||||
unzip chartqa-data.zip
|
||||
rm chartqa-data.zip
|
||||
shell: bash
|
||||
|
||||
- name: ChartQA sanity check with GPT
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/chartqa
|
||||
uv run python debug_chartqa_agent.py
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: Run vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
uv run --no-sync vllm serve Qwen/Qwen2-VL-2B-Instruct \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--max-model-len 4096 \
|
||||
--allowed-local-media-path "$(pwd)/data" \
|
||||
--enable-prefix-caching \
|
||||
--port 8088 &
|
||||
|
||||
VLLM_READY=0
|
||||
for i in {1..100}; do
|
||||
if curl -sSf http://localhost:8088/v1/models > /dev/null 2>&1; then
|
||||
echo "vLLM server is ready!"
|
||||
VLLM_READY=1
|
||||
break
|
||||
fi
|
||||
echo "Waiting for vLLM server to be ready... (${i})"
|
||||
sleep 5
|
||||
done
|
||||
if [[ "$VLLM_READY" != "1" ]]; then
|
||||
echo "vLLM server failed to start!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: ChartQA sanity check with vLLM
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
uv run python debug_chartqa_agent.py
|
||||
shell: bash
|
||||
env:
|
||||
USE_LLM_PROXY: "1"
|
||||
OPENAI_API_BASE: http://localhost:8088/v1
|
||||
OPENAI_MODEL: Qwen/Qwen2-VL-2B-Instruct
|
||||
|
||||
- name: Stop vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
pkill -f vllm
|
||||
for i in {1..60}; do
|
||||
if ! pgrep -f vllm; then
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: ChartQA training
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_chartqa_agent.py ci
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: chartqa_train
|
||||
|
||||
- name: Validate ChartQA training
|
||||
run: |
|
||||
set -euo pipefail
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.chartqa_train.outputs.project_name }} ${{ steps.chartqa_train.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -1,179 +0,0 @@
|
||||
name: Examples - RAG
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 6 AM UTC+8
|
||||
- cron: '0 22 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-rag, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'RAG - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('RAG - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
rag:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-rag' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: RAG (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group rag --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group rag --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: dependencies-rag-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare RAG dataset
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/rag
|
||||
mkdir -p data
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" -O data/dataset_tiny.parquet
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" -O data/chunks_candidate_tiny.pkl
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" -O data/index_hnsw_faiss_n32e40_tiny.index
|
||||
|
||||
- name: Run WIKI Retriever MCP Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/rag
|
||||
uv run python wiki_retriever_mcp.py &
|
||||
for i in {1..20}; do
|
||||
sleep 5
|
||||
if nc -z localhost 8099; then
|
||||
echo "MCP server is up!"
|
||||
exit 0
|
||||
else
|
||||
echo "Waiting for MCP server to start..."
|
||||
fi
|
||||
done
|
||||
echo "MCP server failed to start within expected time."
|
||||
exit 1
|
||||
|
||||
- name: Run vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
vllm serve Qwen/Qwen2.5-1.5B-Instruct \
|
||||
--enable-auto-tool-choice \
|
||||
--tool-call-parser hermes \
|
||||
--port 8000 &
|
||||
|
||||
VLLM_READY=0
|
||||
for i in {1..100}; do
|
||||
if curl -sSf http://localhost:8000/v1/models > /dev/null 2>&1; then
|
||||
echo "vLLM server is ready!"
|
||||
VLLM_READY=1
|
||||
break
|
||||
fi
|
||||
echo "Waiting for vLLM server to be ready... (${i})"
|
||||
sleep 5
|
||||
done
|
||||
if [[ "$VLLM_READY" != "1" ]]; then
|
||||
echo "vLLM server failed to start!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Run RAG Sanity check
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/rag
|
||||
uv run python rag_agent.py
|
||||
shell: bash
|
||||
|
||||
- name: Stop vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
pkill -f vllm
|
||||
for i in {1..60}; do
|
||||
if ! pgrep -f vllm; then
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: RAG training
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/rag
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_rag.py fast
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: rag_train
|
||||
|
||||
- name: Validate RAG training
|
||||
run: |
|
||||
set -ex
|
||||
# Allow up to 5 rollouts to fail to produce rewards
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.rag_train.outputs.project_name }} ${{ steps.rag_train.outputs.run_name }} --reward-tolerance 5
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -33,7 +33,8 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
# legacy is omitted because langchain doesn't work with legacy vllm versions
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
@@ -57,13 +58,13 @@ jobs:
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group langchain --group torch-gpu-stable
|
||||
--group dev --group experiment --group agents --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable)
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script == 'stable'
|
||||
--group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
|
||||
@@ -68,23 +68,13 @@ jobs:
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
# TODO: Currently only test the client tracer implementation.
|
||||
- name: Tinker LLM sanity check (tracer text)
|
||||
- name: Tinker LLM sanity check
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python -m tests.test_tinker_llm tracer-text
|
||||
shell: bash
|
||||
env:
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker LLM sanity check (tracer tool)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python -m tests.test_tinker_llm tracer-tool
|
||||
# TODO: Currently only test the client tracer implementation.
|
||||
python -m tests.test_tinker_llm
|
||||
shell: bash
|
||||
env:
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
@@ -1,18 +0,0 @@
|
||||
# Pre-defined workflow with workflow_dispatch trigger,
|
||||
# convenient for testing and debugging.
|
||||
|
||||
name: Playground
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
playground:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
- name: Run script
|
||||
run: |
|
||||
echo "Hello, world!"
|
||||
@@ -27,43 +27,13 @@ jobs:
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-gpu' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Full Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})
|
||||
name: GPU Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
|
||||
|
||||
runs-on: ${{ matrix.mark.runs-on }}
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
mark:
|
||||
- id: store
|
||||
display-name: Store
|
||||
pytest-mark: 'store' # store tests should not require gpu
|
||||
runs-on: ubuntu-latest
|
||||
has-gpu: false
|
||||
# AgentOps needs to be separated because it injects tricky global state.
|
||||
- id: agentops
|
||||
display-name: AgentOps
|
||||
pytest-mark: 'agentops' # including agentops+litellm tests here
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
has-gpu: true
|
||||
# Similar for Weave.
|
||||
- id: weave
|
||||
display-name: Weave
|
||||
pytest-mark: 'weave'
|
||||
runs-on: ubuntu-latest # No GPU tests for Weave.
|
||||
has-gpu: false
|
||||
# Other tests that require GPU
|
||||
- id: gpu
|
||||
display-name: GPU required
|
||||
pytest-mark: '(gpu or llmproxy) and not agentops'
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
has-gpu: true
|
||||
# Other uncovered tests
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not weave and not gpu and not llmproxy'
|
||||
runs-on: ubuntu-latest
|
||||
has-gpu: false
|
||||
env:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
@@ -73,7 +43,6 @@ jobs:
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
if: matrix.mark.has-gpu
|
||||
run: nvidia-smi
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
@@ -82,32 +51,16 @@ jobs:
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.env.python-version }}
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
|
||||
- name: Sync dependencies (latest, gpu)
|
||||
if: matrix.env.setup-script == 'latest' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group torch-gpu-stable
|
||||
# Don't install vllm/pytorch on CPU counterparts
|
||||
- name: Sync dependencies (latest, cpu)
|
||||
if: matrix.env.setup-script == 'latest' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group core-stable
|
||||
- name: Sync dependencies (stable, gpu)
|
||||
if: matrix.env.setup-script == 'stable' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.env.setup-script }}
|
||||
- name: Sync dependencies (stable, cpu)
|
||||
if: matrix.env.setup-script == 'stable' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group core-stable
|
||||
# Don't install langchain for legacy dependency because it has conflicts with torch.
|
||||
- name: Sync dependencies (legacy, gpu)
|
||||
if: matrix.env.setup-script == 'legacy' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group torch-gpu-legacy
|
||||
- name: Sync dependencies (legacy, cpu)
|
||||
if: matrix.env.setup-script == 'legacy' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group core-legacy
|
||||
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
@@ -117,22 +70,62 @@ jobs:
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: dependencies-tests-full-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
|
||||
name: dependencies-tests-full-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Setup Docker environments
|
||||
run: ./scripts/mongodb_docker_run.sh
|
||||
run: |
|
||||
set -euo pipefail
|
||||
|
||||
cd docker
|
||||
|
||||
# Setup data directories
|
||||
./setup.sh
|
||||
|
||||
# Start Dockers
|
||||
docker compose -f compose.mongo.yml up -d
|
||||
|
||||
SERVICE_NAME=mongo
|
||||
TIMEOUT=60 # seconds
|
||||
SLEEP=2
|
||||
|
||||
cid="$(docker compose -f compose.mongo.yml ps -q "$SERVICE_NAME")"
|
||||
if [ -z "$cid" ]; then
|
||||
echo "Service $SERVICE_NAME is not running"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Waiting for $SERVICE_NAME to become healthy..."
|
||||
end=$((SECONDS + TIMEOUT))
|
||||
|
||||
while [ "$SECONDS" -lt "$end" ]; do
|
||||
status="$(docker inspect -f '{{.State.Health.Status}}' "$cid")"
|
||||
echo "Current status: $status"
|
||||
|
||||
if [ "$status" = "healthy" ]; then
|
||||
echo "$SERVICE_NAME is healthy ✅"
|
||||
exit 0
|
||||
elif [ "$status" = "unhealthy" ]; then
|
||||
echo "$SERVICE_NAME is unhealthy ❌"
|
||||
docker logs "$cid" || true
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep "$SLEEP"
|
||||
done
|
||||
|
||||
echo "Timed out waiting for $SERVICE_NAME to become healthy after ${TIMEOUT}s"
|
||||
docker logs "$cid" || true
|
||||
exit 1
|
||||
shell: bash
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
@@ -142,10 +135,9 @@ jobs:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
# mongo, openai, gpu, all enabled by default
|
||||
- name: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}${{ matrix.env.setup-script == 'legacy' && ' and not langchain' || '' }}"
|
||||
uv run pytest -v --durations=0 tests
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
@@ -186,15 +178,11 @@ jobs:
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group torch-gpu-stable
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script == 'stable'
|
||||
# Don't install langchain for legacy dependency because it has conflicts with torch.
|
||||
- name: Sync dependencies (legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-legacy
|
||||
if: matrix.setup-script == 'legacy'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
@@ -234,14 +222,6 @@ jobs:
|
||||
python write_traces.py agentops
|
||||
sleep 5
|
||||
|
||||
- name: Write Traces with Operations
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python write_traces.py operation
|
||||
sleep 5
|
||||
|
||||
- name: Write Traces via Otel Tracer with Client
|
||||
run: |
|
||||
set -euo pipefail
|
||||
@@ -340,24 +320,3 @@ jobs:
|
||||
echo "Waiting for llm_proxy.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: MultiMetrics backend example
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python write_metrics.py --duration 8 --prom-port 9105 --prom-host 0.0.0.0 2>&1 | tee metrics.log &
|
||||
pid=$!
|
||||
|
||||
for attempt in $(seq 1 20); do
|
||||
if curl -sSf http://localhost:9105/metrics | grep -q minimal_requests_total; then
|
||||
echo "Metrics endpoint responding"
|
||||
wait $pid
|
||||
cat metrics.log
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
echo "Metrics endpoint did not respond"
|
||||
exit 1
|
||||
|
||||
+14
-55
@@ -19,8 +19,7 @@ jobs:
|
||||
lint:
|
||||
strategy:
|
||||
matrix:
|
||||
setup: [fast, slow, next]
|
||||
fail-fast: false
|
||||
setup: [fast, slow]
|
||||
name: Lint - ${{ matrix.setup }}
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
@@ -33,14 +32,10 @@ jobs:
|
||||
- name: Sync dependencies (fast)
|
||||
run: uv sync --frozen --group dev --no-default-groups
|
||||
if: matrix.setup == 'fast'
|
||||
- name: Upgrade dependencies (next)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup == 'next'
|
||||
- name: Sync dependencies (slow)
|
||||
run: |
|
||||
uv sync --frozen \
|
||||
--extra apo \
|
||||
--extra weave \
|
||||
--extra verl \
|
||||
--extra mongo \
|
||||
--group dev \
|
||||
@@ -49,9 +44,8 @@ jobs:
|
||||
--group trl \
|
||||
--group tinker \
|
||||
--group agents \
|
||||
--group langchain \
|
||||
--no-default-groups
|
||||
if: matrix.setup != 'fast'
|
||||
if: matrix.setup == 'slow'
|
||||
# This pre-commit skips JavaScript on purpose.
|
||||
- name: Run pre-commit
|
||||
uses: pre-commit/action@v3.0.1
|
||||
@@ -66,7 +60,7 @@ jobs:
|
||||
if: matrix.setup == 'fast'
|
||||
- name: Run pyright (slow)
|
||||
run: uv run --locked --no-sync pyright -p pyrightconfig.json
|
||||
if: matrix.setup != 'fast'
|
||||
if: matrix.setup == 'slow'
|
||||
|
||||
lint-js:
|
||||
name: Lint - JavaScript
|
||||
@@ -77,8 +71,6 @@ jobs:
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Run ESLint
|
||||
@@ -111,10 +103,6 @@ jobs:
|
||||
- name: Set source commit for docs
|
||||
run: |
|
||||
echo "SOURCE_COMMIT=${{ github.sha }}" >> $GITHUB_ENV
|
||||
- name: Verify OpenAPI specification is up-to-date
|
||||
run: |
|
||||
uv run --locked --no-sync python scripts/export_openapi.py
|
||||
git diff --exit-code docs/assets/store-openapi.json
|
||||
- name: Build documentation
|
||||
run: uv run --locked --no-sync mkdocs build --strict
|
||||
- name: Upload docs artifact
|
||||
@@ -127,32 +115,7 @@ jobs:
|
||||
test:
|
||||
strategy:
|
||||
matrix:
|
||||
mark:
|
||||
# store has many tests and is a good isolated group.
|
||||
- id: store
|
||||
display-name: Store
|
||||
pytest-mark: 'store'
|
||||
# AgentOps needs to be separated because it injects tricky global state.
|
||||
- id: agentops
|
||||
display-name: AgentOps
|
||||
pytest-mark: 'agentops'
|
||||
# Similar for Weave.
|
||||
- id: weave
|
||||
display-name: Weave
|
||||
pytest-mark: 'weave'
|
||||
# litellm proxy tests are slow
|
||||
- id: llmproxy
|
||||
display-name: LLM proxy
|
||||
pytest-mark: 'llmproxy'
|
||||
# Robustness of utilities is important. There are many tests.
|
||||
- id: utils
|
||||
display-name: Utilities
|
||||
pytest-mark: 'utils'
|
||||
# unmarked tests: adapter, execution engine, etc.
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not weave and not llmproxy and not utils'
|
||||
env:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.11'
|
||||
@@ -163,7 +126,7 @@ jobs:
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
|
||||
name: Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})
|
||||
name: Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
@@ -171,16 +134,16 @@ jobs:
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.env.python-version }}
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-stable
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
|
||||
if: matrix.env.setup-script != 'latest'
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
@@ -190,15 +153,13 @@ jobs:
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: dependencies-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
|
||||
name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
@@ -206,12 +167,12 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests -m "not mongo and not openai and not gpu and (${{ matrix.mark.pytest-mark }})"
|
||||
uv run pytest -v --durations=0 tests -m "not mongo"
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
|
||||
test-js:
|
||||
name: Test (JavaScript)
|
||||
name: Test - JavaScript
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
@@ -221,8 +182,6 @@ jobs:
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
+1
-3
@@ -1,10 +1,8 @@
|
||||
# Agentlightning specific files
|
||||
verl_old
|
||||
meta-llama/**
|
||||
**/debug/**/*.png
|
||||
**/debug/**/*.json
|
||||
debug/*.png
|
||||
requirements-freeze*.txt
|
||||
/playground
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
|
||||
@@ -3,14 +3,13 @@ repos:
|
||||
rev: v6.0.0
|
||||
hooks:
|
||||
- id: end-of-file-fixer
|
||||
exclude: (.*store-openapi\.json$)
|
||||
- id: trailing-whitespace
|
||||
- id: check-yaml
|
||||
exclude: ^mkdocs\.yml$
|
||||
- id: check-toml
|
||||
- id: check-added-large-files
|
||||
args: ["--maxkb=1024"]
|
||||
exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)|(.*store-openapi\.json$)
|
||||
exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)
|
||||
- id: check-shebang-scripts-are-executable
|
||||
- id: detect-private-key
|
||||
- repo: https://github.com/pycqa/isort
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
# Repository Guidelines
|
||||
|
||||
## Architecture Overview
|
||||
Agent Lightning runs through a continuous loop: runners and tracers emit spans, `LightningStore` (`agentlightning/store/`) keeps them synchronized, and algorithms in `agentlightning/algorithm/` consume those traces to improve behavior.
|
||||
|
||||
## Project Structure & Module Organization
|
||||
- `agentlightning/`: adapters, execution stack, training loop, tracer, reward logic, and the `agl` CLI.
|
||||
- `docs/` & `examples/`: narrative and procedural docs (assets in `docs/assets/`, navigation in `mkdocs.yml`) plus runnable workflows whose READMEs point to their companion how-to guides. `docs/how-to` covers task-focused instructions, while `docs/tutorials` explains concepts and subsystems.
|
||||
- `dashboard/`, `scripts/`, `tests/`: UI bundles, release/dataset/CI automation, and mirrored coverage of the runtime tree. Record download steps rather than committing binaries.
|
||||
|
||||
## Build, Test, and Development Commands
|
||||
- `uv sync --group dev` — provision tooling once per environment.
|
||||
- `uv run --no-sync pytest -v` — execute the full suite; add a path or `-k expr` to narrow the run.
|
||||
- `uv run --no-sync pyright` — enforce static typing parity with CI.
|
||||
- `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` and `uv run --no-sync mkdocs build --strict` — keep formatting tidy and documentation valid.
|
||||
Always commit the refreshed `uv.lock` when dependencies shift, and mention optional groups (VERL, APO, GPU) in PR notes.
|
||||
|
||||
## Common Issues & Fixes
|
||||
- When `uv run` errors with `Permission denied` under `~/.cache`, override both cache locations inline: ``UV_CACHE="$(pwd)/.cache_uv" XDG_CACHE_HOME="$(pwd)/.cache_xdg" uv run --no-sync <command>``.
|
||||
|
||||
## Coding Style & Naming Conventions
|
||||
- Target `requires-python >= 3.10`, four-space indentation, 120-character lines (though docstrings may run longer), and formatter-owned diffs (Black + isort, `black` profile). Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
|
||||
- Maintain exhaustive type hints (pyright enforces them) and prefer shared dataclasses or Pydantic models from `agentlightning.types`.
|
||||
- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points. Use mkdocs styles: `[][]` syntax for cross-references and single backticks for inline code blocks.
|
||||
- Writing logs is encouraged, especially for long functions with multiple steps and try-except blocks that catch all exceptions. Use `logging.getLogger(__name__)` to get loggers. Distinguish between DEBUG, INFO, WARNING, and ERROR logs.
|
||||
|
||||
## Testing Guidelines
|
||||
- Mirror runtime directories under `tests/` and match filenames for quick traceability.
|
||||
- Parametrize pytest cases and apply markers (`openai`, `gpu`, `agentops`, `mongo`, `llmproxy`) so optional suites can be skipped via selectors like `-m "not mongo"` yet still exercised in CI.
|
||||
- Lean on fixtures, favor real stores/spans/agents over mocks, and drive coverage across the majority of branches.
|
||||
- If an imported module is missing from the environment, check whether `uv sync` has been run with the right groups. Do not make stubs for external dependencies unless necessary.
|
||||
|
||||
## Example Contributions
|
||||
- Ship each example with a README that includes smoke-test instructions so maintainers can validate quickly. The README must contain an "Included Files" section summarizing every file and its role.
|
||||
- Keep runnable example modules self-contained with a module-level docstring describing CLI usage. Document important or educational classes/functions with targeted docstrings and inline comments where clarity matters.
|
||||
- Add a CI workflow per example named `examples-<name>.yml` in `.github/workflows/`. Register it in `badge-<name>.yml`, `badge-examples.yml`, and `badge-latest.yml` when applicable so badges stay accurate.
|
||||
|
||||
## Commit & Pull Request Guidelines
|
||||
- Branch from a fresh `main` using `feature/<slug>`, `fix/<slug>`, `docs/<slug>`, or `chore/<slug>`.
|
||||
- Write imperative, scoped commits, reference issues with `Fixes #123`, and rerun pre-commit plus the relevant pytest/doc builds before pushing.
|
||||
- Use PR descriptions to summarize intent, list verification commands, call out dependency or docs-navigation updates, and link new docs/examples via `mkdocs.yml` or `examples/README.md`. Include logs for dashboard changes.
|
||||
@@ -57,7 +57,6 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.
|
||||
|
||||
- [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
|
||||
- [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
|
||||
- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl) and their blog [*Stop Wrestling with Your Agent RL: How Youtu-Agent Achieved Stable, 128-GPU Scaling Without Breaking a Sweat*](https://spotted-coconut-df8.notion.site/Stop-Wrestling-with-Your-Agent-RL-How-Youtu-Agent-Achieved-Stable-128-GPU-Scaling-Without-Breaking-2ca5e8f089ba80539a98c582b65e0233).
|
||||
|
||||
## ⚡ Architecture
|
||||
|
||||
|
||||
@@ -12,56 +12,13 @@ from opentelemetry.sdk.trace import ReadableSpan
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agentlightning.emitter.reward import get_reward_value
|
||||
from agentlightning.semconv import AGL_OPERATION, AGL_REWARD, LightningSpanAttributes
|
||||
from agentlightning.types import Span, Triplet
|
||||
from agentlightning.utils.otel import filter_and_unflatten_attributes
|
||||
|
||||
from .base import TraceAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _attributes_get_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[str]:
|
||||
"""Get a string from the attributes, if present.
|
||||
If there are multiple matches, the first one is returned.
|
||||
"""
|
||||
for key in keys:
|
||||
if key in attributes:
|
||||
if isinstance(attributes[key], str):
|
||||
return attributes[key]
|
||||
else:
|
||||
logger.warning(f"Attribute {key} is found but is not a string: {attributes[key]}")
|
||||
return None
|
||||
|
||||
|
||||
def _attributes_get_ids_multiple(attributes: Dict[str, Any], keys: List[str]) -> Optional[List[int]]:
|
||||
"""Get a list of integers from the attributes, if present.
|
||||
If there are multiple matches, the first one is returned.
|
||||
"""
|
||||
for key in keys:
|
||||
if key in attributes:
|
||||
if (isinstance(attributes[key], list) or isinstance(attributes[key], tuple)) and all(
|
||||
isinstance(x, int) for x in attributes[key]
|
||||
):
|
||||
return list(attributes[key])
|
||||
else:
|
||||
logger.warning(f"Attribute {key} is found but is not a list of integers: {attributes[key]}")
|
||||
return None
|
||||
|
||||
|
||||
def _attributes_unflatten_multiple(
|
||||
attributes: Dict[str, Any], keys: List[str]
|
||||
) -> Union[Dict[str, Any], List[Any], None]:
|
||||
"""Unflatten the attributes, if present.
|
||||
If there are multiple matches, the first one is returned.
|
||||
"""
|
||||
for key in keys:
|
||||
result = filter_and_unflatten_attributes(attributes, key)
|
||||
if result:
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
class Transition(BaseModel):
|
||||
"""A single transition within a reinforcement learning trajectory.
|
||||
|
||||
@@ -174,7 +131,7 @@ class TraceTree:
|
||||
if not should_visit(node):
|
||||
return False
|
||||
agent_name = node.agent_name()
|
||||
vis_name = node.id[-8:] + " (" + node.span.name + ")"
|
||||
vis_name = node.id[:8] + " (" + node.span.name + ")"
|
||||
if agent_name is not None:
|
||||
vis_name += " [" + agent_name + "]"
|
||||
dot.node(node.id, vis_name) # type: ignore
|
||||
@@ -351,19 +308,6 @@ class TraceTree:
|
||||
if agent_name is not None:
|
||||
return agent_name
|
||||
|
||||
# Case 6: Weave
|
||||
is_agent_type = attributes.get("type") == "agent"
|
||||
if is_agent_type:
|
||||
agent_name = cast(Optional[str], attributes.get("agentlightning.operation.input.name"))
|
||||
if agent_name is not None:
|
||||
return agent_name
|
||||
|
||||
# Case 7: Weave + LangChain
|
||||
if self.span.name.startswith("langchain.Chain."):
|
||||
attributes_lc_name = cast(Optional[str], attributes.get("lc_name"))
|
||||
if attributes_lc_name is not None:
|
||||
return attributes_lc_name
|
||||
|
||||
def maybe_reward_dict(self) -> dict[str, Any]:
|
||||
"""Return a reward payload if the span encodes one.
|
||||
|
||||
@@ -383,17 +327,7 @@ class TraceTree:
|
||||
`True` when the span payload describes a reward, otherwise `False`.
|
||||
"""
|
||||
maybe_reward = self.maybe_reward_dict()
|
||||
if maybe_reward and maybe_reward.get("type") == "reward": # type: ignore
|
||||
return True
|
||||
|
||||
# Agent-lightning 0.3+
|
||||
if (
|
||||
self.span.name == AGL_OPERATION
|
||||
and self.span.attributes.get(LightningSpanAttributes.OPERATION_NAME.value) == AGL_REWARD
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
return maybe_reward and maybe_reward.get("type") == "reward" # type: ignore
|
||||
|
||||
def find_llm_calls(
|
||||
self,
|
||||
@@ -430,9 +364,7 @@ class TraceTree:
|
||||
is_llm_call = False
|
||||
if is_llm_call:
|
||||
# Check the response id
|
||||
response_id = _attributes_get_multiple(
|
||||
self.span.attributes, ["gen_ai.response.id", "agentlightning.operation.output.id"]
|
||||
)
|
||||
response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id") # type: ignore
|
||||
if response_id is None and within_llm_call is True:
|
||||
is_llm_call = False
|
||||
if (
|
||||
@@ -444,8 +376,7 @@ class TraceTree:
|
||||
|
||||
if is_llm_call:
|
||||
llm_calls.append((self, within_matching_subtree)) # type: ignore
|
||||
if existing_llm_call_response_ids is None:
|
||||
existing_llm_call_response_ids = set()
|
||||
existing_llm_call_response_ids = existing_llm_call_response_ids or set()
|
||||
if response_id is not None:
|
||||
existing_llm_call_response_ids.add(response_id)
|
||||
if within_llm_call is not None:
|
||||
@@ -559,12 +490,12 @@ class TraceTree:
|
||||
assign_to: List[Tuple[str, int]] = []
|
||||
for child in item.children:
|
||||
if child.id in llm_call_ids:
|
||||
assign_to.append((child.id, child.end_time)) # type: ignore
|
||||
assign_to.append(child.id) # type: ignore
|
||||
|
||||
agentops_output = child.maybe_reward_dict()
|
||||
agentops_output = item.maybe_reward_dict()
|
||||
if agentops_output and agentops_output.get("type") == "reward":
|
||||
for assign_to_id, assign_to_end_time in reversed(assign_to):
|
||||
if assign_to_end_time > child.start_time: # type: ignore
|
||||
if assign_to_end_time > item.start_time: # type: ignore
|
||||
# This reward happens before the end of the LLM call.
|
||||
continue
|
||||
if assign_to_id in rewards:
|
||||
@@ -574,129 +505,28 @@ class TraceTree:
|
||||
|
||||
return rewards
|
||||
|
||||
def extract_prompt_image_urls(self, prompt_raw_content: Any) -> List[str]:
|
||||
"""Extract image URLs from the span attributes, in order of appearance.
|
||||
|
||||
Args:
|
||||
prompt_raw_content: The raw content of the prompt, which can be in one of several formats:
|
||||
|
||||
- List[dict]: A list of message entries, each being a dict with at least a "content" key.
|
||||
- Dict[str, Any]: A dictionary, often with numeric string keys (e.g., `{"0": {...}, "1": {...}}`), where each value is a message entry.
|
||||
If the dict does not have numeric keys, it is treated as a single message entry.
|
||||
"""
|
||||
message_entries: List[Any] = []
|
||||
if isinstance(prompt_raw_content, list):
|
||||
message_entries = cast(List[Any], prompt_raw_content)
|
||||
elif isinstance(prompt_raw_content, dict):
|
||||
# Common when the attributes expand to {"0": {...}, "prompt_filter_results": ...}
|
||||
numeric_keys = [
|
||||
key
|
||||
for key in cast(Dict[str, Any], prompt_raw_content).keys()
|
||||
if isinstance(key, str) and key.isdigit() # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
]
|
||||
if numeric_keys:
|
||||
for key in sorted(numeric_keys, key=int):
|
||||
message_entries.append(prompt_raw_content[key])
|
||||
else:
|
||||
message_entries = [prompt_raw_content]
|
||||
else:
|
||||
return []
|
||||
|
||||
image_urls: List[str] = []
|
||||
for message in cast(List[Dict[str, Any]], message_entries):
|
||||
if (
|
||||
not isinstance(message, dict) # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
or "content" not in message
|
||||
):
|
||||
continue
|
||||
content = message["content"]
|
||||
if isinstance(content, str):
|
||||
try:
|
||||
content = json.loads(content) # This content should now be a list
|
||||
except json.JSONDecodeError:
|
||||
logger.debug(f"Failed to parse message content as JSON: {content}")
|
||||
continue
|
||||
if isinstance(content, list):
|
||||
for content_part in cast(List[Dict[str, Any]], content):
|
||||
if not isinstance(content_part, dict): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
continue
|
||||
if content_part.get("type") == "image_url":
|
||||
image_url_dict = cast(Dict[str, Any], content_part.get("image_url"))
|
||||
if not isinstance(image_url_dict, dict): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
continue
|
||||
if "url" in image_url_dict:
|
||||
image_urls.append(image_url_dict["url"])
|
||||
return image_urls
|
||||
|
||||
def span_to_triplet(self, span: Span, agent_name: str) -> Triplet:
|
||||
"""Convert a span to a triplet.
|
||||
|
||||
Subclass can override this method to add more fields to the triplet,
|
||||
such as chat messages and tool calls.
|
||||
"""
|
||||
prompt_token_ids = (
|
||||
_attributes_get_ids_multiple(
|
||||
span.attributes,
|
||||
[
|
||||
"prompt_token_ids",
|
||||
"agentlightning.operation.output.prompt_token_ids", # Weave tracer
|
||||
],
|
||||
)
|
||||
or []
|
||||
)
|
||||
response_token_ids = (
|
||||
_attributes_get_ids_multiple(
|
||||
span.attributes,
|
||||
[
|
||||
"response_token_ids",
|
||||
"agentlightning.operation.output.response_token_ids.0", # Weave tracer
|
||||
"agentlightning.operation.output.choices.0.token_ids", # Weave tracer with newer vLLM
|
||||
"agentlightning.operation.output.choices.0.provider_specific_fields.token_ids", # new vLLM + new OpenAI client SDK
|
||||
],
|
||||
)
|
||||
or []
|
||||
)
|
||||
prompt_token_ids = span.attributes.get("prompt_token_ids", []) # type: ignore
|
||||
response_token_ids = span.attributes.get("response_token_ids", []) # type: ignore
|
||||
response_id = span.attributes.get("gen_ai.response.id", None) # type: ignore
|
||||
|
||||
response_id = _attributes_get_multiple(
|
||||
span.attributes, ["gen_ai.response.id", "agentlightning.operation.output.id"]
|
||||
)
|
||||
request_metadata = _attributes_unflatten_multiple(
|
||||
span.attributes, ["gen_ai.request", "agentlightning.operation.input"]
|
||||
)
|
||||
response_metadata = _attributes_unflatten_multiple(
|
||||
span.attributes, ["gen_ai.response", "agentlightning.operation.output"]
|
||||
)
|
||||
# Special handling for Weave tracer: messages are handled separately
|
||||
if isinstance(request_metadata, dict):
|
||||
request_metadata.pop("messages", None)
|
||||
if isinstance(response_metadata, dict):
|
||||
response_metadata.pop("choices", None)
|
||||
response_metadata.pop("prompt_token_ids", None)
|
||||
response_metadata.pop("response_token_ids", None)
|
||||
|
||||
prompt_raw_content = _attributes_unflatten_multiple(
|
||||
span.attributes, ["gen_ai.prompt", "agentlightning.operation.input.messages"]
|
||||
)
|
||||
completion_raw_content = _attributes_unflatten_multiple(
|
||||
span.attributes, ["gen_ai.completion", "agentlightning.operation.output.choices"]
|
||||
)
|
||||
image_urls = self.extract_prompt_image_urls(prompt_raw_content)
|
||||
prompt_payload = {"token_ids": prompt_token_ids, "raw_content": prompt_raw_content, "image_urls": image_urls}
|
||||
response_payload = {"token_ids": response_token_ids, "raw_content": completion_raw_content}
|
||||
|
||||
# FIXME: logprob doesn't support Weave tracer yet.
|
||||
logprobs_content = span.attributes.get("logprobs.content", None) # type: ignore
|
||||
if isinstance(logprobs_content, str):
|
||||
logprobs_content = json.loads(logprobs_content)
|
||||
response_payload["logprobs"] = logprobs_content
|
||||
response: Dict[str, Any] = {"token_ids": response_token_ids, "logprobs": logprobs_content}
|
||||
else:
|
||||
response = {"token_ids": response_token_ids}
|
||||
|
||||
return Triplet(
|
||||
prompt=prompt_payload,
|
||||
response=response_payload,
|
||||
prompt={"token_ids": prompt_token_ids},
|
||||
response=response,
|
||||
reward=None,
|
||||
metadata=dict(
|
||||
request=request_metadata, response=response_metadata, response_id=response_id, agent_name=agent_name
|
||||
),
|
||||
metadata=dict(response_id=response_id, agent_name=agent_name),
|
||||
)
|
||||
|
||||
def to_trajectory(
|
||||
|
||||
@@ -7,44 +7,23 @@ APO with textual gradients that read rollout spans and outputs to modify the pro
|
||||
- rollout: same pattern as your example, but task is a dict (T_task)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Counter,
|
||||
Dict,
|
||||
Generic,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Set,
|
||||
Tuple,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
|
||||
|
||||
import poml
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from agentlightning.adapter.messages import TraceToMessages
|
||||
from agentlightning.algorithm.base import Algorithm
|
||||
from agentlightning.algorithm.utils import batch_iter_over_dataset, with_llm_proxy, with_store
|
||||
from agentlightning.algorithm.utils import batch_iter_over_dataset
|
||||
from agentlightning.reward import find_final_reward
|
||||
from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
@@ -381,10 +360,8 @@ class APO(Algorithm, Generic[T_task]):
|
||||
)
|
||||
return new_prompt
|
||||
|
||||
@with_store
|
||||
async def get_rollout_results(
|
||||
self,
|
||||
store: LightningStore,
|
||||
rollout: List[Rollout],
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
@@ -402,6 +379,7 @@ class APO(Algorithm, Generic[T_task]):
|
||||
List of rollout results formatted for APO processing.
|
||||
"""
|
||||
rollout_results: List[RolloutResultForAPO] = []
|
||||
store = self.get_store()
|
||||
adapter = self.get_adapter()
|
||||
for r in rollout:
|
||||
spans = await store.query_spans(r.rollout_id)
|
||||
@@ -798,12 +776,8 @@ class APO(Algorithm, Generic[T_task]):
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
@with_llm_proxy()
|
||||
@with_store
|
||||
async def run(
|
||||
self,
|
||||
store: LightningStore, # Injected by decorator - callers should not provide this parameter
|
||||
llm_proxy: Optional[LLMProxy], # Injected by decorator - callers should not provide this parameter
|
||||
train_dataset: Optional[Dataset[T_task]] = None,
|
||||
val_dataset: Optional[Dataset[T_task]] = None,
|
||||
) -> None:
|
||||
|
||||
@@ -5,16 +5,11 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, List, Literal, Optional
|
||||
from typing import Any, List, Literal, Optional
|
||||
|
||||
from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
|
||||
|
||||
from .base import Algorithm
|
||||
from .utils import with_llm_proxy, with_store
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -41,8 +36,6 @@ class Baseline(FastAlgorithm):
|
||||
finish, and logs every collected span and reward. It is primarily useful as
|
||||
a smoke test for the platform plumbing rather than a performant trainer.
|
||||
|
||||
The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started.
|
||||
|
||||
Args:
|
||||
n_epochs: Number of dataset passes to execute for both the train and val
|
||||
splits during developer experiments.
|
||||
@@ -187,12 +180,8 @@ class Baseline(FastAlgorithm):
|
||||
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
|
||||
@with_llm_proxy()
|
||||
@with_store
|
||||
async def run(
|
||||
self,
|
||||
store: LightningStore, # Injected by decorator - callers should not provide this parameter
|
||||
llm_proxy: Optional[LLMProxy], # Injected by decorator - callers should not provide this parameter
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> None:
|
||||
@@ -213,6 +202,8 @@ class Baseline(FastAlgorithm):
|
||||
logger.debug(f"Train indices: {train_indices}")
|
||||
logger.debug(f"Val indices: {val_indices}")
|
||||
|
||||
store = self.get_store()
|
||||
|
||||
# Currently we only supports a single resource update at the start.
|
||||
initial_resources = self.get_initial_resources()
|
||||
if initial_resources is not None:
|
||||
|
||||
@@ -1,42 +1,11 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
import random
|
||||
from collections.abc import Coroutine
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Concatenate,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
ParamSpec,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
overload,
|
||||
)
|
||||
from typing import Iterator, List, Sequence, TypeVar
|
||||
|
||||
from agentlightning.types import Dataset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
from .base import Algorithm
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
T_algo = TypeVar("T_algo", bound="Algorithm")
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
|
||||
@@ -72,106 +41,3 @@ def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterat
|
||||
if len(current_batch) == batch_size:
|
||||
yield [dataset[index] for index in current_batch]
|
||||
current_batch = []
|
||||
|
||||
|
||||
def with_store(
|
||||
func: Callable[Concatenate[T_algo, LightningStore, P], Coroutine[Any, Any, R]],
|
||||
) -> Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]]:
|
||||
"""Inject the algorithm's `LightningStore` into coroutine methods.
|
||||
|
||||
The decorator calls `Algorithm.get_store()` once per invocation and passes the
|
||||
resulting store as an explicit argument to the wrapped coroutine. Decorated
|
||||
methods therefore receive the resolved store even when invoked by helper
|
||||
utilities rather than directly by the algorithm.
|
||||
|
||||
Args:
|
||||
func: The coroutine that expects `(self, store, *args, **kwargs)`.
|
||||
|
||||
Returns:
|
||||
A coroutine wrapper that automatically retrieves the store and forwards it
|
||||
to `func`.
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: T_algo, *args: P.args, **kwargs: P.kwargs) -> R:
|
||||
store = self.get_store()
|
||||
return await func(self, store, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@overload
|
||||
def with_llm_proxy(
|
||||
required: Literal[False] = False,
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[Concatenate[T_algo, Optional[LLMProxy], P], Coroutine[Any, Any, R]]],
|
||||
Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
|
||||
]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def with_llm_proxy(
|
||||
required: Literal[True],
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[Concatenate[T_algo, LLMProxy, P], Coroutine[Any, Any, R]]],
|
||||
Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
|
||||
]: ...
|
||||
|
||||
|
||||
def with_llm_proxy(
|
||||
required: bool = False,
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[..., Coroutine[Any, Any, Any]]],
|
||||
Callable[..., Coroutine[Any, Any, Any]],
|
||||
]:
|
||||
"""Resolve and optionally lifecycle-manage the configured LLM proxy.
|
||||
|
||||
Args:
|
||||
required: When True, raises `ValueError` if the algorithm does not have an
|
||||
[`LLMProxy`][agentlightning.LLMProxy] set. When False, the wrapped coroutine receives
|
||||
`None` if no proxy is available.
|
||||
auto_start: When True, [`LLMProxy.start()`][agentlightning.LLMProxy.start] is invoked if the proxy is not
|
||||
already running before calling `func` and [`LLMProxy.stop()`][agentlightning.LLMProxy.stop] is
|
||||
called afterwards.
|
||||
|
||||
Returns:
|
||||
A decorator that injects the [`LLMProxy`][agentlightning.LLMProxy] (or `None`) as the first
|
||||
argument after `self` and manages automatic startup/shutdown when requested.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[..., Coroutine[Any, Any, Any]],
|
||||
) -> Callable[..., Coroutine[Any, Any, Any]]:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: Algorithm, *args: Any, **kwargs: Any) -> Any:
|
||||
llm_proxy = self.get_llm_proxy()
|
||||
|
||||
if required and llm_proxy is None:
|
||||
raise ValueError(
|
||||
"LLM proxy is required but not configured. Call set_llm_proxy() before using this method."
|
||||
)
|
||||
|
||||
auto_started = False
|
||||
if auto_start and llm_proxy is not None:
|
||||
if llm_proxy.is_running():
|
||||
logger.info("Proxy is already running, skipping start")
|
||||
else:
|
||||
logger.info("Starting proxy, managed by the algorithm")
|
||||
await llm_proxy.start()
|
||||
auto_started = True
|
||||
|
||||
try:
|
||||
# At type level, overloads guarantee that if `required=True`
|
||||
# then `func` expects a non-optional LLMProxy.
|
||||
return await func(self, llm_proxy, *args, **kwargs)
|
||||
finally:
|
||||
if auto_started and llm_proxy is not None:
|
||||
logger.info("Stopping proxy, managed by the algorithm")
|
||||
await llm_proxy.stop()
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Optional, Type
|
||||
from typing import Any, Optional
|
||||
|
||||
from hydra import compose, initialize
|
||||
from omegaconf import OmegaConf
|
||||
@@ -12,10 +10,6 @@ from agentlightning.client import AgentLightningClient
|
||||
from agentlightning.types import Dataset
|
||||
from agentlightning.verl.entrypoint import run_ppo # type: ignore
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.verl.daemon import AgentModeDaemon
|
||||
from agentlightning.verl.trainer import AgentLightningTrainer
|
||||
|
||||
|
||||
class VERL(Algorithm):
|
||||
"""VERL-powered algorithm that delegates training to the VERL PPO runner.
|
||||
@@ -29,28 +23,6 @@ class VERL(Algorithm):
|
||||
config: Dictionary mirroring the overrides passed to the VERL CLI. The
|
||||
overrides are merged with VERL's packaged defaults via Hydra before
|
||||
launching training.
|
||||
trainer_cls: Optional override for the trainer class. Experimental.
|
||||
daemon_cls: Optional override for the daemon class. Experimental.
|
||||
|
||||
!!! note "Trajectory aggregation (experimental)"
|
||||
|
||||
Trajectory-level aggregation merges an entire multi-turn rollout into a single,
|
||||
masked training sample so GPU time is spent once per trajectory rather than N times
|
||||
per turn. Enable it via:
|
||||
|
||||
```python
|
||||
config["agentlightning"]["trace_aggregator"] = {
|
||||
"level": "trajectory",
|
||||
"trajectory_max_prompt_length": ...,
|
||||
"trajectory_max_response_length": ...,
|
||||
}
|
||||
```
|
||||
|
||||
Keep conversations structured (message lists rather than manual string
|
||||
concatenation) so prefix matching can stitch traces, and toggle `debug=True` plus
|
||||
`unmatch_log_dir` when you need to inspect retokenization or chat-template
|
||||
mismatches. See [this blog post](https://agent-lightning.github.io/posts/trajectory_level_aggregation/)
|
||||
for more details.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
@@ -118,12 +90,7 @@ class VERL(Algorithm):
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: dict[str, Any],
|
||||
trainer_cls: Optional[Type[AgentLightningTrainer]] = None,
|
||||
daemon_cls: Optional[Type[AgentModeDaemon]] = None,
|
||||
):
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
super().__init__()
|
||||
|
||||
# Compose the base config exactly like your decorator:
|
||||
@@ -135,8 +102,6 @@ class VERL(Algorithm):
|
||||
# Allow adding new fields
|
||||
OmegaConf.set_struct(base_cfg, False)
|
||||
self.config = OmegaConf.merge(base_cfg, override_conf)
|
||||
self.trainer_cls = trainer_cls
|
||||
self.daemon_cls = daemon_cls
|
||||
|
||||
def run(
|
||||
self,
|
||||
@@ -154,11 +119,6 @@ class VERL(Algorithm):
|
||||
adapter have been garbage-collected when using the V1 execution
|
||||
mode.
|
||||
"""
|
||||
from agentlightning.verl.daemon import AgentModeDaemon
|
||||
from agentlightning.verl.trainer import AgentLightningTrainer
|
||||
|
||||
trainer_cls = self.trainer_cls or AgentLightningTrainer
|
||||
daemon_cls = self.daemon_cls or AgentModeDaemon
|
||||
try:
|
||||
store = self.get_store()
|
||||
except Exception:
|
||||
@@ -170,8 +130,6 @@ class VERL(Algorithm):
|
||||
store=None,
|
||||
llm_proxy=None,
|
||||
adapter=None,
|
||||
trainer_cls=trainer_cls,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
else:
|
||||
print("Store is set. Assuming v1 execution mode.")
|
||||
@@ -184,8 +142,6 @@ class VERL(Algorithm):
|
||||
store=store,
|
||||
llm_proxy=llm_proxy,
|
||||
adapter=adapter,
|
||||
trainer_cls=trainer_cls,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
|
||||
def get_client(self) -> AgentLightningClient:
|
||||
|
||||
@@ -12,7 +12,6 @@ from typing import Dict, Iterable, Tuple
|
||||
_SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
|
||||
"vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
|
||||
"store": ("agentlightning.cli.store", "Run a LightningStore server."),
|
||||
"prometheus": ("agentlightning.cli.prometheus", "Serve Prometheus metrics from the multiprocess registry."),
|
||||
"agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
|
||||
}
|
||||
|
||||
|
||||
@@ -1,115 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Serve Prometheus metrics from the Agent Lightning multiprocess registry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
from fastapi import FastAPI
|
||||
from prometheus_client import make_asgi_app # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
from agentlightning.logging import setup as setup_logging
|
||||
from agentlightning.utils.metrics import get_prometheus_registry
|
||||
from agentlightning.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def ensure_prometheus_dir() -> str:
|
||||
"""Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
|
||||
|
||||
directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
|
||||
if directory is None:
|
||||
raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
|
||||
|
||||
Path(directory).mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Serving Prometheus multiprocess metrics from %s", directory)
|
||||
return directory
|
||||
|
||||
|
||||
def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
|
||||
"""Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
|
||||
|
||||
Args:
|
||||
metrics_path: URL path to expose the Prometheus metrics endpoint on.
|
||||
|
||||
Returns:
|
||||
A FastAPI application ready to serve metrics.
|
||||
"""
|
||||
|
||||
if not metrics_path.startswith("/"):
|
||||
raise ValueError("metrics_path must start with '/'.")
|
||||
|
||||
normalized_path = metrics_path.rstrip("/")
|
||||
if normalized_path in ("", "/"):
|
||||
raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
|
||||
|
||||
app = FastAPI(title="Agent Lightning Prometheus exporter", docs_url=None, redoc_url=None)
|
||||
metrics_app = make_asgi_app(registry=get_prometheus_registry()) # pyright: ignore[reportUnknownVariableType]
|
||||
app.mount(normalized_path, metrics_app) # pyright: ignore[reportUnknownArgumentType]
|
||||
|
||||
@app.get("/health")
|
||||
async def healthcheck() -> dict[str, str]: # pyright: ignore[reportUnusedFunction]
|
||||
return {"status": "ok"}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
|
||||
parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
|
||||
parser.add_argument(
|
||||
"--metrics-path",
|
||||
default="/v1/prometheus",
|
||||
help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
||||
help="Configure the logging level for the metrics server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--access-log",
|
||||
action="store_true",
|
||||
help="Enable uvicorn access logs. Disabled by default to reduce noise.",
|
||||
)
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
setup_logging(args.log_level)
|
||||
ensure_prometheus_dir()
|
||||
|
||||
try:
|
||||
app = create_prometheus_app(args.metrics_path)
|
||||
except ValueError as exc:
|
||||
logger.error("Failed to configure prometheus app: %s", exc)
|
||||
return 1
|
||||
|
||||
launcher_args = PythonServerLauncherArgs(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level=getattr(logging, args.log_level),
|
||||
access_log=args.access_log,
|
||||
healthcheck_url="/health",
|
||||
)
|
||||
launcher = PythonServerLauncher(app, launcher_args)
|
||||
|
||||
try:
|
||||
asyncio.run(launcher.run_forever())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Received shutdown signal. Stopping Prometheus server.")
|
||||
except RuntimeError as exc:
|
||||
logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -7,18 +7,11 @@ from __future__ import annotations
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Iterable, List
|
||||
from typing import Iterable
|
||||
|
||||
from agentlightning import setup_logging
|
||||
from agentlightning.store.client_server import LightningStoreServer
|
||||
from agentlightning.store.memory import InMemoryLightningStore
|
||||
from agentlightning.utils.metrics import (
|
||||
ConsoleMetricsBackend,
|
||||
MetricsBackend,
|
||||
MultiMetricsBackend,
|
||||
PrometheusMetricsBackend,
|
||||
setup_multiprocess_prometheus,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -40,10 +33,9 @@ def main(argv: Iterable[str] | None = None) -> int:
|
||||
help="Configure the logging level for the store.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tracker",
|
||||
nargs="+",
|
||||
choices=["prometheus", "console"],
|
||||
help="Enable metrics tracking. Repeat for multiple trackers.",
|
||||
"--prometheus",
|
||||
action="store_true",
|
||||
help="Enable Prometheus metrics.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n-workers",
|
||||
@@ -71,36 +63,12 @@ def main(argv: Iterable[str] | None = None) -> int:
|
||||
|
||||
setup_logging(args.log_level)
|
||||
|
||||
trackers: List[MetricsBackend] = []
|
||||
if args.tracker:
|
||||
if "prometheus" in args.tracker:
|
||||
logger.info("Enabling Prometheus metrics tracking.")
|
||||
if args.n_workers > 1:
|
||||
# This has to be done before prometheus_client is imported
|
||||
setup_multiprocess_prometheus()
|
||||
logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
|
||||
trackers.append(PrometheusMetricsBackend())
|
||||
|
||||
if "console" in args.tracker:
|
||||
logger.info("Enabling console metrics tracking.")
|
||||
trackers.append(ConsoleMetricsBackend())
|
||||
|
||||
if len(trackers) == 0:
|
||||
tracker: MetricsBackend | None = None
|
||||
elif len(trackers) == 1:
|
||||
tracker = trackers[0]
|
||||
else:
|
||||
tracker = MultiMetricsBackend(trackers)
|
||||
|
||||
if args.backend == "memory":
|
||||
store = InMemoryLightningStore(
|
||||
thread_safe=True, # Using thread_safe store for server
|
||||
tracker=tracker,
|
||||
)
|
||||
store = InMemoryLightningStore(prometheus=args.prometheus)
|
||||
elif args.backend == "mongo":
|
||||
from agentlightning.store.mongo import MongoLightningStore
|
||||
|
||||
store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
|
||||
store = MongoLightningStore(client=args.mongo_uri, prometheus=args.prometheus)
|
||||
else:
|
||||
raise ValueError(f"Invalid backend: {args.backend}")
|
||||
|
||||
@@ -116,7 +84,7 @@ def main(argv: Iterable[str] | None = None) -> int:
|
||||
port=args.port,
|
||||
cors_allow_origins=args.cors_origins,
|
||||
launch_mode=launch_mode,
|
||||
tracker=tracker,
|
||||
prometheus=args.prometheus,
|
||||
n_workers=args.n_workers,
|
||||
)
|
||||
try:
|
||||
|
||||
@@ -1,17 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Convenient helpers for creating spans / traces.
|
||||
|
||||
All emitters operate in two modes, switchable via the `propagate` parameter.
|
||||
The emitters first [`SpanCreationRequest`][agentlightning.SpanCreationRequest] object, then:
|
||||
|
||||
1. When `propagate` is True, this creation request will be propagated to the active tracer
|
||||
and a [`Span`][agentlightning.Span] instance will be created (possibly deferred).
|
||||
2. When `propagate` is False, the creation request will be returned directly. Useful for cases
|
||||
when you don't have a tracer but you want to create a creation request for later use.
|
||||
"""
|
||||
|
||||
from .annotation import emit_annotation, operation
|
||||
from .annotation import emit_annotation
|
||||
from .exception import emit_exception
|
||||
from .message import emit_message, get_message_value
|
||||
from .object import emit_object, get_object_value
|
||||
@@ -27,7 +16,6 @@ from .reward import (
|
||||
|
||||
__all__ = [
|
||||
"reward",
|
||||
"operation",
|
||||
"emit_reward",
|
||||
"get_reward_value",
|
||||
"get_rewards_from_span",
|
||||
|
||||
@@ -1,38 +1,19 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Helpers for emitting annotation/operation spans."""
|
||||
"""Helpers for emitting annotation spans."""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
ContextManager,
|
||||
Dict,
|
||||
Optional,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
from typing import Any, Dict
|
||||
|
||||
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import SpanCoreFields, SpanRecordingContext, TraceStatus
|
||||
from agentlightning.utils.otel import check_attributes_sanity, flatten_attributes, sanitize_attributes
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
_FnType = TypeVar("_FnType", bound=Callable[..., Any])
|
||||
from agentlightning.semconv import AGL_ANNOTATION
|
||||
from agentlightning.utils.otel import flatten_attributes, get_tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanCoreFields:
|
||||
def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> ReadableSpan:
|
||||
"""Emit a new annotation span.
|
||||
|
||||
This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward].
|
||||
@@ -46,325 +27,22 @@ def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanC
|
||||
Args:
|
||||
annotation: Dictionary containing annotation key-value pairs.
|
||||
Representatives are rewards, tags, and metadata.
|
||||
propagate: Whether to propagate the span to tracers automatically.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
"""
|
||||
annotation_attributes = flatten_attributes(annotation, expand_leaf_lists=False)
|
||||
check_attributes_sanity(annotation_attributes)
|
||||
sanitized_attributes = sanitize_attributes(annotation_attributes)
|
||||
logger.debug("Emitting annotation span with keys %s", sanitized_attributes.keys())
|
||||
annotation_attributes = flatten_attributes(annotation)
|
||||
if any(not isinstance(v, (str, int, float, bool, bytes)) for v in annotation_attributes.values()):
|
||||
raise TypeError("All annotation attributes must be primitive types (str, int, float, bool, bytes)")
|
||||
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot emit annotation span.")
|
||||
else:
|
||||
tracer = DummyTracer()
|
||||
|
||||
return tracer.create_span(
|
||||
name=AGL_ANNOTATION,
|
||||
attributes=sanitized_attributes,
|
||||
status=TraceStatus(status_code="OK"),
|
||||
# TODO: this should use a tracer from current context rather than the singleton
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
span = tracer.start_span(
|
||||
AGL_ANNOTATION,
|
||||
attributes=annotation_attributes,
|
||||
)
|
||||
logger.debug("Emitting annotation span with keys %s", annotation_attributes)
|
||||
with span:
|
||||
pass
|
||||
if not isinstance(span, ReadableSpan):
|
||||
raise ValueError(f"Span is not a ReadableSpan: {span}")
|
||||
|
||||
|
||||
class OperationContext:
|
||||
"""Context manager and decorator for tracing operations.
|
||||
|
||||
This class manages a tracer-backed span for a logical unit of work. It can be
|
||||
used either:
|
||||
|
||||
* As a decorator, in which case inputs and outputs are inferred
|
||||
automatically from the wrapped function's signature.
|
||||
* As a context manager, in which case inputs and outputs can be recorded
|
||||
explicitly via [`set_input`][agentlightning.emitter.annotation.OperationContext.set_input]
|
||||
and [`set_output`][agentlightning.emitter.annotation.OperationContext.set_output].
|
||||
|
||||
Attributes:
|
||||
name: Human-readable span name.
|
||||
initial_attributes: Attributes applied when the span is created.
|
||||
tracer: Tracer implementation used to create spans.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
|
||||
"""Initialize a new operation context.
|
||||
|
||||
Args:
|
||||
name: Human-readable name of the span.
|
||||
attributes: Initial attributes attached to the span. Values are
|
||||
JSON-serialized where necessary.
|
||||
propagate: Whether the span should be sent to active exporters.
|
||||
"""
|
||||
self.name = name
|
||||
self.initial_attributes = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
self.propagate = propagate
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot trace operation spans.")
|
||||
self.tracer = tracer
|
||||
else:
|
||||
self.tracer = DummyTracer()
|
||||
self._ctx_manager: Optional[ContextManager[SpanRecordingContext]] = None
|
||||
self._recording_context: Optional[SpanRecordingContext] = None
|
||||
self._span: Optional[SpanCoreFields] = None
|
||||
|
||||
def __enter__(self) -> "OperationContext":
|
||||
"""Enter the context manager and start a new span.
|
||||
|
||||
Returns:
|
||||
The current :class:`OperationContext` instance with an active span.
|
||||
"""
|
||||
sanitized_attrs = sanitize_attributes(self.initial_attributes)
|
||||
self._ctx_manager = self.tracer.operation_context(self.name, attributes=sanitized_attrs)
|
||||
recording_context = self._ctx_manager.__enter__()
|
||||
self._recording_context = recording_context
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> None:
|
||||
"""Exit the context manager and finish the span."""
|
||||
if self._ctx_manager:
|
||||
self._ctx_manager.__exit__(exc_type, exc_val, exc_tb)
|
||||
if self._recording_context:
|
||||
self._span = self._recording_context.get_recorded_span()
|
||||
self._ctx_manager = None
|
||||
self._recording_context = None
|
||||
|
||||
def span(self) -> SpanCoreFields:
|
||||
"""Get the span that was created by this context manager."""
|
||||
if self._span is None:
|
||||
raise RuntimeError("Span is not ready yet.")
|
||||
return self._span
|
||||
|
||||
def set_input(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Record input arguments on the current span.
|
||||
|
||||
Positional arguments are stored under the `input.args.<index>` attributes,
|
||||
and keyword arguments are stored under `input.<name>` attributes.
|
||||
|
||||
This is intended for use inside a `with operation(...) as op` block.
|
||||
|
||||
Args:
|
||||
*args: Positional arguments to record.
|
||||
**kwargs: Keyword arguments to record.
|
||||
"""
|
||||
if not self._recording_context:
|
||||
raise RuntimeError("No recording context found. Cannot set input.")
|
||||
|
||||
prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
attributes: Dict[str, Any] = {}
|
||||
if args:
|
||||
for idx, value in enumerate(args):
|
||||
flattened = flatten_attributes({str(idx): value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{prefix}.args.{nested_key}"] = nested_value
|
||||
if kwargs:
|
||||
for key, value in kwargs.items():
|
||||
flattened = flatten_attributes({key: value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{prefix}.{nested_key}"] = nested_value
|
||||
if attributes:
|
||||
self._recording_context.record_attributes(sanitize_attributes(attributes))
|
||||
|
||||
def set_output(self, output: Any) -> None:
|
||||
"""Record the output value on the current span.
|
||||
|
||||
This is intended for use inside a `with operation(...) as op` block.
|
||||
|
||||
Args:
|
||||
output: The output value to record.
|
||||
"""
|
||||
if not self._recording_context:
|
||||
raise RuntimeError("No recording context found. Cannot set output.")
|
||||
|
||||
flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: output})
|
||||
self._recording_context.record_attributes(sanitize_attributes(flattened))
|
||||
|
||||
def __call__(self, fn: _FnType) -> _FnType:
|
||||
"""Wrap a callable so its execution is traced in a span.
|
||||
|
||||
When used as a decorator, a new span is created for each call to
|
||||
the wrapped function. The bound arguments are recorded as input
|
||||
attributes, the return value is recorded as an output attribute,
|
||||
and any exception is recorded and marks the span as an error.
|
||||
|
||||
Args:
|
||||
fn: The function or coroutine function to wrap.
|
||||
|
||||
Returns:
|
||||
The wrapped callable.
|
||||
"""
|
||||
function_name = fn.__name__
|
||||
|
||||
sig = inspect.signature(fn)
|
||||
|
||||
sanitized_init_attrs = sanitize_attributes(
|
||||
{LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
|
||||
)
|
||||
|
||||
def _record_auto_inputs(
|
||||
recording_ctx: SpanRecordingContext, args: Tuple[Any, ...], kwargs: Dict[str, Any]
|
||||
) -> None:
|
||||
"""Bind arguments to signature and log them on the span."""
|
||||
attributes: Dict[str, Any] = {}
|
||||
try:
|
||||
bound = sig.bind(*args, **kwargs)
|
||||
bound.apply_defaults()
|
||||
for name, value in bound.arguments.items():
|
||||
parameter = sig.parameters.get(name)
|
||||
if parameter and parameter.kind is inspect.Parameter.VAR_POSITIONAL:
|
||||
attr_prefix = f"{LightningSpanAttributes.OPERATION_INPUT.value}.{name}"
|
||||
for idx, item in enumerate(value):
|
||||
flattened = flatten_attributes({str(idx): item})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{attr_prefix}.{nested_key}"] = nested_value
|
||||
else:
|
||||
flattened = flatten_attributes({name: value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
|
||||
except Exception:
|
||||
if args:
|
||||
for idx, value in enumerate(args):
|
||||
flattened = flatten_attributes({str(idx): value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.args.{nested_key}"] = (
|
||||
nested_value
|
||||
)
|
||||
if kwargs:
|
||||
flattened = flatten_attributes({"kwargs": kwargs})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
|
||||
if attributes:
|
||||
recording_ctx.record_attributes(sanitize_attributes(attributes))
|
||||
|
||||
def _record_auto_outputs(recording_ctx: SpanRecordingContext, result: Any) -> None:
|
||||
"""Record the output value on the span."""
|
||||
flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: result})
|
||||
recording_ctx.record_attributes(sanitize_attributes(flattened))
|
||||
|
||||
if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
|
||||
|
||||
@functools.wraps(fn)
|
||||
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Async wrapper that traces the wrapped coroutine."""
|
||||
with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
|
||||
_record_auto_inputs(recording_ctx, args, kwargs)
|
||||
result = await fn(*args, **kwargs)
|
||||
_record_auto_outputs(recording_ctx, result)
|
||||
return result
|
||||
|
||||
return cast(_FnType, async_wrapper)
|
||||
|
||||
else:
|
||||
|
||||
@functools.wraps(fn)
|
||||
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Sync wrapper that traces the wrapped callable."""
|
||||
with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
|
||||
_record_auto_inputs(recording_ctx, args, kwargs)
|
||||
result = fn(*args, **kwargs)
|
||||
_record_auto_outputs(recording_ctx, result)
|
||||
return result
|
||||
|
||||
return cast(_FnType, sync_wrapper)
|
||||
|
||||
|
||||
@overload
|
||||
def operation(
|
||||
fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
|
||||
) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(
|
||||
*, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
|
||||
) -> OperationContext: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(fn: _FnType, *, name: Optional[str] = None, **additional_attributes: Any) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(*, name: Optional[str] = None, **additional_attributes: Any) -> OperationContext: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(fn: _FnType, **additional_attributes: Any) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(**additional_attributes: Any) -> OperationContext: ...
|
||||
|
||||
|
||||
def operation(
|
||||
fn: Optional[_FnType] = None,
|
||||
*,
|
||||
propagate: bool = True,
|
||||
name: Optional[str] = None,
|
||||
**additional_attributes: Any,
|
||||
) -> Union[_FnType, OperationContext]:
|
||||
"""Entry point for tracking operations.
|
||||
|
||||
This helper can be used either as a decorator or as a context manager.
|
||||
The span name is fixed to [`AGL_OPERATION`][agentlightning.semconv.AGL_OPERATION];
|
||||
custom span names are not supported. Any keyword arguments are recorded as span attributes.
|
||||
|
||||
Usage as a decorator:
|
||||
|
||||
```python
|
||||
@operation
|
||||
def func(...):
|
||||
...
|
||||
|
||||
@operation(category="compute")
|
||||
def func(...):
|
||||
...
|
||||
```
|
||||
|
||||
Usage as a context manager:
|
||||
|
||||
```python
|
||||
with operation(user_id=123) as op:
|
||||
op.set_input(data=data)
|
||||
# ... do work ...
|
||||
op.set_output(result)
|
||||
```
|
||||
|
||||
Args:
|
||||
fn: When used as `@operation`, this is the wrapped function.
|
||||
When used as `operation(**attrs)`, this should be omitted (or
|
||||
left as `None`) and only keyword attributes are provided.
|
||||
propagate: Whether spans should use the active span processor. When False,
|
||||
spans will stay local and not be exported.
|
||||
name: Optional alias that populates
|
||||
[`LightningSpanAttributes.OPERATION_NAME`][agentlightning.semconv.LightningSpanAttributes.OPERATION_NAME]
|
||||
when `additional_attributes` does not already define it.
|
||||
**additional_attributes: Additional span attributes to attach at
|
||||
creation time.
|
||||
|
||||
Returns:
|
||||
Either a wrapped callable (when used as a decorator) or an
|
||||
[`OperationContext`][agentlightning.emitter.annotation.OperationContext]
|
||||
(when used as a context manager factory).
|
||||
"""
|
||||
|
||||
if name is not None:
|
||||
if LightningSpanAttributes.OPERATION_NAME.value in additional_attributes:
|
||||
raise ValueError("Cannot specify both `name` and `additional_attributes.operation_name`.")
|
||||
additional_attributes[LightningSpanAttributes.OPERATION_NAME.value] = name
|
||||
|
||||
# Case 1: Used as @operation (bare decorator or with attributes)
|
||||
if callable(fn):
|
||||
# Create context with fixed name, then immediately wrap the function
|
||||
return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn)
|
||||
|
||||
# Case 2: Used as operation(...) / with operation(...)
|
||||
# Custom span names are intentionally not supported; use AGL_OPERATION.
|
||||
if fn is not None:
|
||||
raise ValueError("Custom span names are intentionally not supported when used as a context manager.")
|
||||
return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)
|
||||
return span
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import traceback
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from opentelemetry.semconv.attributes import exception_attributes
|
||||
|
||||
from agentlightning.semconv import AGL_EXCEPTION
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import TraceStatus
|
||||
from agentlightning.utils.otel import flatten_attributes, format_exception_attributes, sanitize_attributes
|
||||
from agentlightning.utils.otel import get_tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -32,23 +32,25 @@ def emit_exception(
|
||||
"""
|
||||
if not isinstance(exception, BaseException): # type: ignore
|
||||
raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
|
||||
span_attributes = format_exception_attributes(exception)
|
||||
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
|
||||
span_attributes = {
|
||||
exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
|
||||
exception_attributes.EXCEPTION_MESSAGE: str(exception),
|
||||
exception_attributes.EXCEPTION_ESCAPED: True,
|
||||
}
|
||||
if stacktrace.strip():
|
||||
span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
|
||||
|
||||
if attributes:
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
span_attributes.update(attributes)
|
||||
|
||||
logger.debug("Emitting exception span for %s", type(exception).__name__)
|
||||
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot emit exception span.")
|
||||
else:
|
||||
tracer = DummyTracer()
|
||||
tracer.create_span(
|
||||
span = tracer.start_span(
|
||||
AGL_EXCEPTION,
|
||||
attributes=span_attributes,
|
||||
# The exception span is successful by itself.
|
||||
status=TraceStatus(status_code="OK"),
|
||||
)
|
||||
logger.debug("Emitting exception span for %s", type(exception).__name__)
|
||||
with span:
|
||||
span.record_exception(exception)
|
||||
# We don't set the status of the span here. They have other semantics.
|
||||
|
||||
@@ -4,10 +4,8 @@ import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import Attributes, SpanLike
|
||||
from agentlightning.utils.otel import flatten_attributes, sanitize_attributes
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otel import get_tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -29,21 +27,17 @@ def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, prop
|
||||
if not isinstance(message, str): # type: ignore
|
||||
raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
|
||||
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot emit message span.")
|
||||
else:
|
||||
tracer = DummyTracer()
|
||||
span_attributes: Attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
span_attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
|
||||
if attributes:
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
logger.debug("Emitting message span with message: %s", message)
|
||||
tracer.create_span(
|
||||
span_attributes.update(attributes)
|
||||
span = tracer.start_span(
|
||||
AGL_MESSAGE,
|
||||
attributes=span_attributes,
|
||||
)
|
||||
logger.debug("Emitting message span with message: %s", message)
|
||||
with span:
|
||||
pass
|
||||
|
||||
|
||||
def get_message_value(span: SpanLike) -> Optional[str]:
|
||||
|
||||
@@ -6,15 +6,13 @@ import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import SpanCoreFields, SpanLike, TraceStatus
|
||||
from agentlightning.utils.otel import flatten_attributes, full_qualified_name, sanitize_attributes
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otel import full_qualified_name, get_tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> SpanCoreFields:
|
||||
def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
|
||||
"""Emit an object's serialized representation as an OpenTelemetry span.
|
||||
|
||||
Args:
|
||||
@@ -27,29 +25,20 @@ def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propag
|
||||
"""
|
||||
span_attributes = encode_object(object)
|
||||
if attributes:
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
|
||||
span_attributes.update(attributes)
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
span = tracer.start_span(
|
||||
AGL_OBJECT,
|
||||
attributes=span_attributes,
|
||||
)
|
||||
attr_length = 0
|
||||
if LightningSpanAttributes.OBJECT_JSON.value in span_attributes:
|
||||
attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_JSON.value])
|
||||
elif LightningSpanAttributes.OBJECT_LITERAL.value in span_attributes:
|
||||
attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_LITERAL.value])
|
||||
logger.debug("Emitting object span with payload size %d characters", attr_length)
|
||||
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot emit object span.")
|
||||
else:
|
||||
# Do not actually propagate to any store or tracer backend.
|
||||
tracer = DummyTracer()
|
||||
|
||||
return tracer.create_span(
|
||||
name=AGL_OBJECT,
|
||||
attributes=span_attributes,
|
||||
status=TraceStatus(status_code="OK"),
|
||||
)
|
||||
with span:
|
||||
pass
|
||||
|
||||
|
||||
def encode_object(object: Any) -> Dict[str, Any]:
|
||||
|
||||
@@ -20,10 +20,13 @@ from typing import (
|
||||
cast,
|
||||
)
|
||||
|
||||
import agentops
|
||||
from agentops.sdk.decorators import operation
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
|
||||
from agentlightning.types import SpanCoreFields, SpanLike
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otel import filter_and_unflatten_attributes
|
||||
|
||||
from .annotation import emit_annotation
|
||||
@@ -58,8 +61,6 @@ _FnType = TypeVar("_FnType", bound=Callable[..., Any])
|
||||
|
||||
def _agentops_initialized() -> bool:
|
||||
"""Return `True` when the AgentOps client has been configured."""
|
||||
import agentops
|
||||
|
||||
return agentops.get_client().initialized
|
||||
|
||||
|
||||
@@ -80,8 +81,6 @@ def reward(fn: _FnType) -> _FnType:
|
||||
Wrapped callable that preserves the original signature.
|
||||
"""
|
||||
|
||||
from agentops.sdk.decorators import operation
|
||||
|
||||
def wrap_result(result: Optional[float]) -> _RewardSpanData:
|
||||
"""Normalize the reward value into the span payload format."""
|
||||
if result is None:
|
||||
@@ -147,7 +146,7 @@ def emit_reward(
|
||||
primary_key: str | None = None,
|
||||
attributes: Dict[str, Any] | None = None,
|
||||
propagate: bool = True,
|
||||
) -> SpanCoreFields:
|
||||
) -> ReadableSpan:
|
||||
"""Emit a reward value as an OpenTelemetry span.
|
||||
|
||||
Examples:
|
||||
@@ -173,7 +172,11 @@ def emit_reward(
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
Returns:
|
||||
Span core fields capturing the recorded reward.
|
||||
Readable span capturing the recorded reward.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provided reward cannot be interpreted as a float or the
|
||||
resulting span is not a [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) instance.
|
||||
"""
|
||||
logger.debug(f"Emitting reward: {reward}")
|
||||
reward_dimensions: List[RewardDimension] = []
|
||||
|
||||
@@ -143,10 +143,6 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
|
||||
logger.debug("Algorithm bundle starting against endpoint %s", wrapper_store.endpoint)
|
||||
await algorithm(wrapper_store, stop_evt)
|
||||
logger.debug("Algorithm bundle completed successfully")
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Algorithm received CancelledError; signaling stop event")
|
||||
stop_evt.set()
|
||||
raise
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
|
||||
stop_evt.set()
|
||||
@@ -183,10 +179,6 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
|
||||
logger.debug("Runner %s executing with provided store", worker_id)
|
||||
await runner(client_store, worker_id, stop_evt)
|
||||
logger.debug("Runner %s completed successfully", worker_id)
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Runner %s received CancelledError; signaling stop event", worker_id)
|
||||
stop_evt.set()
|
||||
raise
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
|
||||
stop_evt.set()
|
||||
@@ -218,13 +210,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
|
||||
def _runner_sync(runner: RunnerBundle, worker_id: int, store: LightningStore, stop_evt: ExecutionEvent) -> None:
|
||||
# Runners are executed in child processes; each process owns its own
|
||||
# event loop to keep the asyncio scheduler isolated.
|
||||
try:
|
||||
asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Runner (asyncio) %s received KeyboardInterrupt; exiting gracefully", worker_id)
|
||||
except BaseException as exc:
|
||||
logger.exception("Runner (asyncio) %s crashed by %s; signaling stop event", worker_id, exc)
|
||||
raise
|
||||
asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
|
||||
|
||||
for i in range(self.n_runners):
|
||||
process = cast(
|
||||
@@ -248,13 +234,7 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
|
||||
"""Used when `main_process == "runner"`."""
|
||||
|
||||
def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
|
||||
try:
|
||||
asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Algorithm (asyncio.run) received KeyboardInterrupt; exiting gracefully")
|
||||
except BaseException as exc:
|
||||
logger.exception("Algorithm (asyncio.run) crashed by %s; signaling stop event", exc)
|
||||
raise
|
||||
asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
|
||||
|
||||
process = cast(
|
||||
multiprocessing.Process,
|
||||
|
||||
@@ -6,7 +6,6 @@ AGENTOPS_INSTALLED: bool = False
|
||||
AGENTOPS_LANGCHAIN_INSTALLED: bool = False
|
||||
LITELLM_INSTALLED: bool = False
|
||||
VLLM_INSTALLED: bool = False
|
||||
WEAVE_INSTALLED: bool = False
|
||||
|
||||
try:
|
||||
from . import agentops # type: ignore
|
||||
|
||||
@@ -1,500 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import warnings
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Callable, Dict, Iterator, List
|
||||
|
||||
import weave.trace.weave_init
|
||||
from pydantic import validate_call
|
||||
from weave.trace_server import trace_server_interface as tsi
|
||||
from weave.trace_server.ids import generate_id
|
||||
from weave.trace_server_bindings.client_interface import TraceServerClientInterface
|
||||
from weave.trace_server_bindings.models import ServerInfoRes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"instrument_weave",
|
||||
"uninstrument_weave",
|
||||
"InMemoryWeaveTraceServer",
|
||||
]
|
||||
|
||||
|
||||
class InMemoryWeaveTraceServer(TraceServerClientInterface):
|
||||
"""A minimal in-memory implementation of the TraceServerInterface.
|
||||
|
||||
It stores calls and objects in local dictionaries and returns valid Pydantic
|
||||
responses to satisfy the Weave client and FullTraceServerInterface protocol.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Minimal storage to allow basic querying in tests
|
||||
self.calls: Dict[str, tsi.CallSchema] = {}
|
||||
self.partial_calls: Dict[str, Dict[str, Any]] = {}
|
||||
self.objs: Dict[str, Any] = {}
|
||||
self.files: Dict[str, bytes] = {}
|
||||
self.feedback: List[tsi.FeedbackCreateReq] = []
|
||||
|
||||
self._call_threading_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, *args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
|
||||
return cls()
|
||||
|
||||
def server_info(self) -> ServerInfoRes:
|
||||
return ServerInfoRes(min_required_weave_python_version="0.52.22")
|
||||
|
||||
def ensure_project_exists(self, entity: str, project: str) -> tsi.EnsureProjectExistsRes:
|
||||
return tsi.EnsureProjectExistsRes(project_name=project)
|
||||
|
||||
# --- Call API ---
|
||||
|
||||
@validate_call
|
||||
def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
|
||||
# NOTE: It's not necessary that call_end must be called after call_start.
|
||||
request_content = req.start.model_dump(exclude_none=True)
|
||||
|
||||
# If id needs to be generated here, it's very likely we won't be able to find the call later.
|
||||
# This is just to make the type checker happy.
|
||||
call_id = request_content.get("id") or generate_id()
|
||||
trace_id = request_content.get("trace_id") or generate_id()
|
||||
request_content["id"] = call_id
|
||||
request_content["trace_id"] = trace_id
|
||||
|
||||
with self._call_threading_lock:
|
||||
if call_id in self.partial_calls:
|
||||
# call_end has already been called for this call.
|
||||
kwargs = {**request_content, **self.partial_calls[call_id]}
|
||||
self.calls[call_id] = tsi.CallSchema(**kwargs)
|
||||
del self.partial_calls[call_id]
|
||||
else:
|
||||
self.partial_calls[call_id] = request_content
|
||||
|
||||
return tsi.CallStartRes(id=call_id, trace_id=trace_id)
|
||||
|
||||
@validate_call
|
||||
def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
|
||||
request_content = req.end.model_dump(exclude_none=True)
|
||||
call_id = req.end.id
|
||||
|
||||
with self._call_threading_lock:
|
||||
if call_id in self.partial_calls:
|
||||
# End request always override the start request content.
|
||||
kwargs = {**self.partial_calls[call_id], **request_content}
|
||||
self.calls[call_id] = tsi.CallSchema(**kwargs)
|
||||
del self.partial_calls[call_id]
|
||||
else:
|
||||
self.partial_calls[call_id] = request_content
|
||||
return tsi.CallEndRes()
|
||||
|
||||
@validate_call
|
||||
def call_start_batch(self, req: tsi.CallCreateBatchReq) -> tsi.CallCreateBatchRes:
|
||||
for item in req.batch:
|
||||
if isinstance(item, tsi.CallStartReq):
|
||||
self.call_start(item)
|
||||
elif isinstance(item, tsi.CallEndReq):
|
||||
self.call_end(item)
|
||||
return tsi.CallCreateBatchRes(res=[])
|
||||
|
||||
@validate_call
|
||||
def call_read(self, req: tsi.CallReadReq) -> tsi.CallReadRes:
|
||||
call_data = self.calls.get(req.id)
|
||||
return tsi.CallReadRes(call=call_data)
|
||||
|
||||
@validate_call
|
||||
def calls_query(self, req: tsi.CallsQueryReq) -> tsi.CallsQueryRes:
|
||||
return tsi.CallsQueryRes(calls=list(self.calls_query_stream(req)))
|
||||
|
||||
@validate_call
|
||||
def calls_query_stream(self, req: tsi.CallsQueryReq) -> Iterator[tsi.CallSchema]:
|
||||
yield from self.calls.values()
|
||||
|
||||
@validate_call
|
||||
def calls_delete(self, req: tsi.CallsDeleteReq) -> tsi.CallsDeleteRes:
|
||||
num_deleted = 0
|
||||
for call_id in req.call_ids:
|
||||
if call_id in self.calls:
|
||||
del self.calls[call_id]
|
||||
num_deleted += 1
|
||||
return tsi.CallsDeleteRes(num_deleted=num_deleted)
|
||||
|
||||
@validate_call
|
||||
def call_update(self, req: tsi.CallUpdateReq) -> tsi.CallUpdateRes:
|
||||
return tsi.CallUpdateRes()
|
||||
|
||||
@validate_call
|
||||
def calls_query_stats(self, req: tsi.CallsQueryStatsReq) -> tsi.CallsQueryStatsRes:
|
||||
return tsi.CallsQueryStatsRes(count=len(self.calls))
|
||||
|
||||
# --- Cost API ---
|
||||
|
||||
@validate_call
|
||||
def cost_create(self, req: tsi.CostCreateReq) -> tsi.CostCreateRes:
|
||||
return tsi.CostCreateRes(ids=[(generate_id(), generate_id()) for _ in req.costs])
|
||||
|
||||
@validate_call
|
||||
def cost_query(self, req: tsi.CostQueryReq) -> tsi.CostQueryRes:
|
||||
return tsi.CostQueryRes(results=[])
|
||||
|
||||
@validate_call
|
||||
def cost_purge(self, req: tsi.CostPurgeReq) -> tsi.CostPurgeRes:
|
||||
return tsi.CostPurgeRes()
|
||||
|
||||
# --- Object API (Legacy V1) ---
|
||||
|
||||
@validate_call
|
||||
def obj_create(self, req: tsi.ObjCreateReq) -> tsi.ObjCreateRes:
|
||||
digest = generate_id()
|
||||
self.objs[digest] = req.obj
|
||||
return tsi.ObjCreateRes(digest=digest)
|
||||
|
||||
@validate_call
|
||||
def obj_read(self, req: tsi.ObjReadReq) -> tsi.ObjReadRes:
|
||||
return tsi.ObjReadRes(obj=self.objs.get(req.digest, {}))
|
||||
|
||||
@validate_call
|
||||
def objs_query(self, req: tsi.ObjQueryReq) -> tsi.ObjQueryRes:
|
||||
return tsi.ObjQueryRes(objs=[])
|
||||
|
||||
@validate_call
|
||||
def obj_delete(self, req: tsi.ObjDeleteReq) -> tsi.ObjDeleteRes:
|
||||
return tsi.ObjDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Table API ---
|
||||
|
||||
@validate_call
|
||||
def table_create(self, req: tsi.TableCreateReq) -> tsi.TableCreateRes:
|
||||
return tsi.TableCreateRes(digest=generate_id(), row_digests=[])
|
||||
|
||||
@validate_call
|
||||
def table_create_from_digests(self, req: tsi.TableCreateFromDigestsReq) -> tsi.TableCreateFromDigestsRes:
|
||||
return tsi.TableCreateFromDigestsRes(digest=generate_id())
|
||||
|
||||
@validate_call
|
||||
def table_update(self, req: tsi.TableUpdateReq) -> tsi.TableUpdateRes:
|
||||
return tsi.TableUpdateRes(digest=generate_id(), updated_row_digests=[])
|
||||
|
||||
@validate_call
|
||||
def table_query(self, req: tsi.TableQueryReq) -> tsi.TableQueryRes:
|
||||
return tsi.TableQueryRes(rows=[])
|
||||
|
||||
@validate_call
|
||||
def table_query_stream(self, req: tsi.TableQueryReq) -> Iterator[tsi.TableRowSchema]:
|
||||
yield from []
|
||||
|
||||
@validate_call
|
||||
def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
|
||||
return tsi.TableQueryStatsRes(count=0)
|
||||
|
||||
@validate_call
|
||||
def table_query_stats_batch(self, req: tsi.TableQueryStatsBatchReq) -> tsi.TableQueryStatsBatchRes:
|
||||
return tsi.TableQueryStatsBatchRes(tables=[])
|
||||
|
||||
# --- Ref API ---
|
||||
|
||||
@validate_call
|
||||
def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
|
||||
return tsi.RefsReadBatchRes(vals=[])
|
||||
|
||||
# --- File API ---
|
||||
|
||||
def file_create(self, req: tsi.FileCreateReq) -> tsi.FileCreateRes:
|
||||
self.files[req.name] = req.content
|
||||
return tsi.FileCreateRes(digest=generate_id())
|
||||
|
||||
def file_content_read(self, req: tsi.FileContentReadReq) -> tsi.FileContentReadRes:
|
||||
return tsi.FileContentReadRes(content=self.files.get(req.digest, b"dummy_content"))
|
||||
|
||||
def files_stats(self, req: tsi.FilesStatsReq) -> tsi.FilesStatsRes:
|
||||
total_size = sum(len(c) for c in self.files.values())
|
||||
return tsi.FilesStatsRes(total_size_bytes=total_size)
|
||||
|
||||
# --- Feedback API ---
|
||||
|
||||
@validate_call
|
||||
def feedback_create(self, req: tsi.FeedbackCreateReq) -> tsi.FeedbackCreateRes:
|
||||
req.id = req.id or generate_id()
|
||||
self.feedback.append(req)
|
||||
return tsi.FeedbackCreateRes(
|
||||
id=req.id,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
wb_user_id="dummy_user",
|
||||
payload=req.payload,
|
||||
)
|
||||
|
||||
def feedback_create_batch(self, req: tsi.FeedbackCreateBatchReq) -> tsi.FeedbackCreateBatchRes:
|
||||
results: List[tsi.FeedbackCreateRes] = []
|
||||
for item in req.batch:
|
||||
res = self.feedback_create(item)
|
||||
results.append(res)
|
||||
return tsi.FeedbackCreateBatchRes(res=results)
|
||||
|
||||
@validate_call
|
||||
def feedback_query(self, req: tsi.FeedbackQueryReq) -> tsi.FeedbackQueryRes:
|
||||
return tsi.FeedbackQueryRes(result=[])
|
||||
|
||||
@validate_call
|
||||
def feedback_purge(self, req: tsi.FeedbackPurgeReq) -> tsi.FeedbackPurgeRes:
|
||||
self.feedback.clear()
|
||||
return tsi.FeedbackPurgeRes()
|
||||
|
||||
@validate_call
|
||||
def feedback_replace(self, req: tsi.FeedbackReplaceReq) -> tsi.FeedbackReplaceRes:
|
||||
return tsi.FeedbackReplaceRes(
|
||||
id=req.id or generate_id(),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
wb_user_id="dummy",
|
||||
payload={},
|
||||
)
|
||||
|
||||
# --- Action API ---
|
||||
|
||||
@validate_call
|
||||
def actions_execute_batch(self, req: tsi.ActionsExecuteBatchReq) -> tsi.ActionsExecuteBatchRes:
|
||||
return tsi.ActionsExecuteBatchRes()
|
||||
|
||||
# --- Execute LLM API ---
|
||||
|
||||
@validate_call
|
||||
def completions_create(self, req: tsi.CompletionsCreateReq) -> tsi.CompletionsCreateRes:
|
||||
return tsi.CompletionsCreateRes(response={"choices": [{"text": "dummy completion"}]})
|
||||
|
||||
@validate_call
|
||||
def completions_create_stream(self, req: tsi.CompletionsCreateReq) -> Iterator[dict[str, Any]]:
|
||||
yield {"choices": [{"text": "dummy "}]}
|
||||
yield {"choices": [{"text": "stream"}]}
|
||||
|
||||
# --- Execute Image Generation API ---
|
||||
|
||||
@validate_call
|
||||
def image_create(self, req: tsi.ImageGenerationCreateReq) -> tsi.ImageGenerationCreateRes:
|
||||
return tsi.ImageGenerationCreateRes(response={})
|
||||
|
||||
# --- Project Statistics API ---
|
||||
|
||||
@validate_call
|
||||
def project_stats(self, req: tsi.ProjectStatsReq) -> tsi.ProjectStatsRes:
|
||||
return tsi.ProjectStatsRes(
|
||||
trace_storage_size_bytes=0,
|
||||
objects_storage_size_bytes=0,
|
||||
tables_storage_size_bytes=0,
|
||||
files_storage_size_bytes=0,
|
||||
)
|
||||
|
||||
# --- Thread API ---
|
||||
|
||||
@validate_call
|
||||
def threads_query_stream(self, req: tsi.ThreadsQueryReq) -> Iterator[tsi.ThreadSchema]:
|
||||
yield from []
|
||||
|
||||
# --- Evaluation API (V1) ---
|
||||
|
||||
@validate_call
|
||||
def evaluate_model(self, req: tsi.EvaluateModelReq) -> tsi.EvaluateModelRes:
|
||||
return tsi.EvaluateModelRes(call_id=generate_id())
|
||||
|
||||
@validate_call
|
||||
def evaluation_status(self, req: tsi.EvaluationStatusReq) -> tsi.EvaluationStatusRes:
|
||||
return tsi.EvaluationStatusRes(status=tsi.EvaluationStatusNotFound())
|
||||
|
||||
# --- OTEL API ---
|
||||
|
||||
def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
|
||||
return tsi.OtelExportRes()
|
||||
|
||||
# ==========================================
|
||||
# Object Interface (V2 APIs)
|
||||
# ==========================================
|
||||
|
||||
# --- Ops ---
|
||||
def op_create(self, req: tsi.OpCreateReq) -> tsi.OpCreateRes:
|
||||
return tsi.OpCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
|
||||
|
||||
def op_read(self, req: tsi.OpReadReq) -> tsi.OpReadRes:
|
||||
return tsi.OpReadRes(op=None) # type: ignore
|
||||
|
||||
def op_list(self, req: tsi.OpListReq) -> Iterator[tsi.OpReadRes]:
|
||||
yield from []
|
||||
|
||||
def op_delete(self, req: tsi.OpDeleteReq) -> tsi.OpDeleteRes:
|
||||
return tsi.OpDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Datasets ---
|
||||
def dataset_create(self, req: tsi.DatasetCreateReq) -> tsi.DatasetCreateRes:
|
||||
return tsi.DatasetCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
|
||||
|
||||
def dataset_read(self, req: tsi.DatasetReadReq) -> tsi.DatasetReadRes:
|
||||
return tsi.DatasetReadRes(dataset=None) # type: ignore
|
||||
|
||||
def dataset_list(self, req: tsi.DatasetListReq) -> Iterator[tsi.DatasetReadRes]:
|
||||
yield from []
|
||||
|
||||
def dataset_delete(self, req: tsi.DatasetDeleteReq) -> tsi.DatasetDeleteRes:
|
||||
return tsi.DatasetDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Scorers ---
|
||||
def scorer_create(self, req: tsi.ScorerCreateReq) -> tsi.ScorerCreateRes:
|
||||
return tsi.ScorerCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0, scorer=generate_id())
|
||||
|
||||
def scorer_read(self, req: tsi.ScorerReadReq) -> tsi.ScorerReadRes:
|
||||
return tsi.ScorerReadRes(scorer=None) # type: ignore
|
||||
|
||||
def scorer_list(self, req: tsi.ScorerListReq) -> Iterator[tsi.ScorerReadRes]:
|
||||
yield from []
|
||||
|
||||
def scorer_delete(self, req: tsi.ScorerDeleteReq) -> tsi.ScorerDeleteRes:
|
||||
return tsi.ScorerDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Evaluations (V2) ---
|
||||
def evaluation_create(self, req: tsi.EvaluationCreateReq) -> tsi.EvaluationCreateRes:
|
||||
return tsi.EvaluationCreateRes(
|
||||
digest=generate_id(), object_id=generate_id(), version_index=0, evaluation_ref=generate_id()
|
||||
)
|
||||
|
||||
def evaluation_read(self, req: tsi.EvaluationReadReq) -> tsi.EvaluationReadRes:
|
||||
return tsi.EvaluationReadRes(evaluation=None) # type: ignore
|
||||
|
||||
def evaluation_list(self, req: tsi.EvaluationListReq) -> Iterator[tsi.EvaluationReadRes]:
|
||||
yield from []
|
||||
|
||||
def evaluation_delete(self, req: tsi.EvaluationDeleteReq) -> tsi.EvaluationDeleteRes:
|
||||
return tsi.EvaluationDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Models ---
|
||||
def model_create(self, req: tsi.ModelCreateReq) -> tsi.ModelCreateRes:
|
||||
return tsi.ModelCreateRes(
|
||||
digest=generate_id(), object_id=generate_id(), version_index=0, model_ref=generate_id()
|
||||
)
|
||||
|
||||
def model_read(self, req: tsi.ModelReadReq) -> tsi.ModelReadRes:
|
||||
return tsi.ModelReadRes(model=None) # type: ignore
|
||||
|
||||
def model_list(self, req: tsi.ModelListReq) -> Iterator[tsi.ModelReadRes]:
|
||||
yield from []
|
||||
|
||||
def model_delete(self, req: tsi.ModelDeleteReq) -> tsi.ModelDeleteRes:
|
||||
return tsi.ModelDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Evaluation Runs ---
|
||||
def evaluation_run_create(self, req: tsi.EvaluationRunCreateReq) -> tsi.EvaluationRunCreateRes:
|
||||
return tsi.EvaluationRunCreateRes(evaluation_run_id=generate_id())
|
||||
|
||||
def evaluation_run_read(self, req: tsi.EvaluationRunReadReq) -> tsi.EvaluationRunReadRes:
|
||||
return tsi.EvaluationRunReadRes(evaluation_run=None) # type: ignore
|
||||
|
||||
def evaluation_run_list(self, req: tsi.EvaluationRunListReq) -> Iterator[tsi.EvaluationRunReadRes]:
|
||||
yield from []
|
||||
|
||||
def evaluation_run_delete(self, req: tsi.EvaluationRunDeleteReq) -> tsi.EvaluationRunDeleteRes:
|
||||
return tsi.EvaluationRunDeleteRes(num_deleted=0)
|
||||
|
||||
def evaluation_run_finish(self, req: tsi.EvaluationRunFinishReq) -> tsi.EvaluationRunFinishRes:
|
||||
return tsi.EvaluationRunFinishRes(success=True)
|
||||
|
||||
# --- Predictions ---
|
||||
def prediction_create(self, req: tsi.PredictionCreateReq) -> tsi.PredictionCreateRes:
|
||||
return tsi.PredictionCreateRes(prediction_id=generate_id())
|
||||
|
||||
def prediction_read(self, req: tsi.PredictionReadReq) -> tsi.PredictionReadRes:
|
||||
return tsi.PredictionReadRes(prediction=None) # type: ignore
|
||||
|
||||
def prediction_list(self, req: tsi.PredictionListReq) -> Iterator[tsi.PredictionReadRes]:
|
||||
yield from []
|
||||
|
||||
def prediction_delete(self, req: tsi.PredictionDeleteReq) -> tsi.PredictionDeleteRes:
|
||||
return tsi.PredictionDeleteRes(num_deleted=0)
|
||||
|
||||
def prediction_finish(self, req: tsi.PredictionFinishReq) -> tsi.PredictionFinishRes:
|
||||
return tsi.PredictionFinishRes(success=True)
|
||||
|
||||
# --- Scores ---
|
||||
def score_create(self, req: tsi.ScoreCreateReq) -> tsi.ScoreCreateRes:
|
||||
return tsi.ScoreCreateRes(score_id=generate_id())
|
||||
|
||||
def score_read(self, req: tsi.ScoreReadReq) -> tsi.ScoreReadRes:
|
||||
return tsi.ScoreReadRes(score=None) # type: ignore
|
||||
|
||||
def score_list(self, req: tsi.ScoreListReq) -> Iterator[tsi.ScoreReadRes]:
|
||||
yield from []
|
||||
|
||||
def score_delete(self, req: tsi.ScoreDeleteReq) -> tsi.ScoreDeleteRes:
|
||||
return tsi.ScoreDeleteRes(num_deleted=0)
|
||||
|
||||
|
||||
# Module-level storage for originals
|
||||
_original_init_weave_get_server: Callable[..., Any] | None = None
|
||||
_original_get_entity_project_from_project_name: Callable[..., Any] | None = None
|
||||
_original_get_username: Callable[..., Any] | None = None
|
||||
|
||||
|
||||
def init_weave_get_server_factory(server: InMemoryWeaveTraceServer) -> Callable[..., Any]:
|
||||
# Bypass the usage of Weave remote server
|
||||
def init_weave_get_server(*args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
|
||||
return server
|
||||
|
||||
return init_weave_get_server
|
||||
|
||||
|
||||
def get_entity_project_from_project_name_factory(entity_name: str) -> tuple[str, str]:
|
||||
# Bypass the usage of API
|
||||
try:
|
||||
assert _original_get_entity_project_from_project_name is not None
|
||||
if _original_get_entity_project_from_project_name is not get_entity_project_from_project_name_factory:
|
||||
return _original_get_entity_project_from_project_name(entity_name)
|
||||
else:
|
||||
warnings.warn("W&B integration might have been repeatedly/recursively instrumented.")
|
||||
return "agl", "weave"
|
||||
except weave.trace.weave_init.WeaveWandbAuthenticationException:
|
||||
# In case API is not available.
|
||||
return "agl", "weave"
|
||||
|
||||
|
||||
def get_username() -> str:
|
||||
# Bypass the usage of API
|
||||
try:
|
||||
assert _original_get_username is not None
|
||||
return _original_get_username()
|
||||
except RuntimeError:
|
||||
return "agl"
|
||||
except Exception as exc:
|
||||
warnings.warn(f"Unexpected error in get_username. Using default username. Error: {exc}")
|
||||
return "agl"
|
||||
|
||||
|
||||
def instrument_weave(server: InMemoryWeaveTraceServer):
|
||||
"""Patch the Weave/W&B integration to bypass actual network calls for testing."""
|
||||
|
||||
global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
|
||||
_original_init_weave_get_server = weave.trace.weave_init.init_weave_get_server
|
||||
_original_get_entity_project_from_project_name = weave.trace.weave_init.get_entity_project_from_project_name
|
||||
_original_get_username = weave.trace.weave_init.get_username
|
||||
weave.trace.weave_init.init_weave_get_server = init_weave_get_server_factory(server)
|
||||
weave.trace.weave_init.get_entity_project_from_project_name = get_entity_project_from_project_name_factory
|
||||
weave.trace.weave_init.get_username = get_username
|
||||
|
||||
|
||||
def uninstrument_weave():
|
||||
"""Restore the original Weave/W&B integration methods and HTTP requests."""
|
||||
global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
|
||||
|
||||
if _original_init_weave_get_server is not None:
|
||||
weave.trace.weave_init.init_weave_get_server = _original_init_weave_get_server
|
||||
_original_init_weave_get_server = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
|
||||
if _original_get_entity_project_from_project_name is not None:
|
||||
weave.trace.weave_init.get_entity_project_from_project_name = _original_get_entity_project_from_project_name
|
||||
_original_get_entity_project_from_project_name = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
|
||||
if _original_get_username is not None:
|
||||
weave.trace.weave_init.get_username = _original_get_username
|
||||
_original_get_username = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
@@ -198,7 +198,6 @@ class LitAgent(Generic[T]):
|
||||
* `float` representing the final reward.
|
||||
* `List[ReadableSpan]` with OpenTelemetry spans.
|
||||
* `List[Span]` with Agent Lightning spans.
|
||||
* `List[SpanCoreFields]` with Agent Lightning spans.
|
||||
"""
|
||||
raise NotImplementedError("Agents must implement the `rollout` method.")
|
||||
|
||||
|
||||
+78
-264
@@ -33,8 +33,8 @@ from opentelemetry.sdk.trace import ReadableSpan
|
||||
from agentlightning.litagent import LitAgent
|
||||
from agentlightning.reward import emit_reward, find_final_reward
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.tracer.agentops import AgentOpsTracer
|
||||
from agentlightning.tracer.base import Tracer
|
||||
from agentlightning.tracer.otel import OtelTracer
|
||||
from agentlightning.types import (
|
||||
AttemptedRollout,
|
||||
Hook,
|
||||
@@ -43,7 +43,6 @@ from agentlightning.types import (
|
||||
RolloutMode,
|
||||
RolloutRawResult,
|
||||
Span,
|
||||
SpanCoreFields,
|
||||
)
|
||||
from agentlightning.utils.system_snapshot import system_snapshot
|
||||
|
||||
@@ -74,9 +73,8 @@ class LitAgentRunner(Runner[T_task]):
|
||||
max_rollouts: Optional[int] = None,
|
||||
poll_interval: float = 5.0,
|
||||
heartbeat_interval: float = 10.0,
|
||||
interval_jitter: float = 0.5,
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
|
||||
heartbeat_include_gpu: bool = False,
|
||||
interval_jitter: float = 0.1,
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
|
||||
) -> None:
|
||||
"""Initialize the agent runner.
|
||||
|
||||
@@ -90,10 +88,7 @@ class LitAgentRunner(Runner[T_task]):
|
||||
poll_interval - interval_jitter and poll_interval + interval_jitter.
|
||||
This is to avoid the overload caused by the synchronization of the runners.
|
||||
heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
|
||||
"thread" is the default and recommended mode as it prevents blocking the event loop
|
||||
under load. Use "asyncio" for simpler deployments with low worker counts.
|
||||
heartbeat_include_gpu: Whether to include GPU stats in heartbeat snapshots.
|
||||
Querying GPU stats can be slow under load, so this is disabled by default.
|
||||
"asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
|
||||
"""
|
||||
super().__init__()
|
||||
self._tracer = tracer
|
||||
@@ -102,7 +97,6 @@ class LitAgentRunner(Runner[T_task]):
|
||||
self._heartbeat_interval = heartbeat_interval
|
||||
self._interval_jitter = interval_jitter
|
||||
self._heartbeat_launch_mode = heartbeat_launch_mode
|
||||
self._heartbeat_include_gpu = heartbeat_include_gpu
|
||||
self._random_state = random.Random()
|
||||
|
||||
# Set later
|
||||
@@ -282,84 +276,50 @@ class LitAgentRunner(Runner[T_task]):
|
||||
"""
|
||||
store = self.get_store()
|
||||
|
||||
trace_spans: list[Span] = []
|
||||
result_recognized: bool = False
|
||||
trace_spans: list[ReadableSpan] | list[Span] = []
|
||||
|
||||
# Case 0: result is None
|
||||
if raw_result is None:
|
||||
trace_spans = self._tracer.get_last_trace()
|
||||
result_recognized = True
|
||||
|
||||
# Case 1: result is a float (final reward)
|
||||
if isinstance(raw_result, (bool, int, float)):
|
||||
if isinstance(raw_result, (bool, int)):
|
||||
logger.warning(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Reward is not a number, got: {type(raw_result)}. "
|
||||
"Auto converting to float."
|
||||
)
|
||||
raw_result = float(raw_result)
|
||||
if isinstance(raw_result, float):
|
||||
# Preserve the existing spans before another span is emitted
|
||||
trace_spans = list(self._tracer.get_last_trace())
|
||||
# This will NOT emit another span to the tracer
|
||||
reward_span_core_fields = emit_reward(raw_result, propagate=False)
|
||||
reward_span = emit_reward(raw_result, propagate=False)
|
||||
# We add it to the store manually
|
||||
sequence_id = await store.get_next_span_sequence_id(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
reward_span = Span.from_core_fields(
|
||||
reward_span_core_fields,
|
||||
rollout_id=rollout.rollout_id,
|
||||
attempt_id=rollout.attempt.attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
await store.add_span(reward_span)
|
||||
result_recognized = True
|
||||
await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
|
||||
trace_spans.append(reward_span)
|
||||
|
||||
# Case 2-4: result is a list
|
||||
if isinstance(raw_result, list):
|
||||
# For rollout methods that return a list, we assume that the returned spans
|
||||
# are the complete span set from the whole rollout
|
||||
trace_spans = raw_result
|
||||
|
||||
# Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
|
||||
if isinstance(self._tracer, OtelTracer):
|
||||
|
||||
if not isinstance(
|
||||
self._tracer, AgentOpsTracer
|
||||
): # TODO: this should be replaced with general OpenTelemetry tracer in next version
|
||||
for span in raw_result:
|
||||
await store.add_otel_span(
|
||||
rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
|
||||
"The traces should have already been added to the store. "
|
||||
"Returning the traces from the rollout will result in duplicate spans."
|
||||
"No need to return anything from rollout."
|
||||
)
|
||||
for span in raw_result:
|
||||
added_span = await store.add_otel_span(
|
||||
rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
|
||||
)
|
||||
if added_span is not None:
|
||||
trace_spans.append(added_span)
|
||||
else:
|
||||
logger.error(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Failed to add OpenTelemetry span to the store: {span}"
|
||||
)
|
||||
result_recognized = True
|
||||
|
||||
# Case 3: result is a list of Span (agentlightning spans)
|
||||
elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
|
||||
# Add the spans directly to the store
|
||||
for span in raw_result:
|
||||
await store.add_span(cast(Span, span))
|
||||
trace_spans = [cast(Span, span) for span in raw_result]
|
||||
result_recognized = True
|
||||
|
||||
# Case 4: result is a list of SpanCoreFields (agentlightning spans)
|
||||
elif len(raw_result) > 0 and all(isinstance(t, SpanCoreFields) for t in raw_result):
|
||||
# Add the spans directly to the store too, but needs to get sequence id first
|
||||
sequence_ids = await store.get_many_span_sequence_ids(
|
||||
[(rollout.rollout_id, rollout.attempt.attempt_id) for _ in range(len(raw_result))]
|
||||
)
|
||||
trace_spans = [
|
||||
Span.from_core_fields(
|
||||
cast(SpanCoreFields, span_core_fields),
|
||||
rollout_id=rollout.rollout_id,
|
||||
attempt_id=rollout.attempt.attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
for span_core_fields, sequence_id in zip(raw_result, sequence_ids, strict=True)
|
||||
]
|
||||
await store.add_many_spans(trace_spans)
|
||||
result_recognized = True
|
||||
trace_spans = raw_result
|
||||
|
||||
# Left over cases for list
|
||||
elif len(raw_result) == 0:
|
||||
@@ -367,8 +327,7 @@ class LitAgentRunner(Runner[T_task]):
|
||||
f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
|
||||
"Please check your rollout implementation."
|
||||
)
|
||||
trace_spans = []
|
||||
result_recognized = True
|
||||
trace_spans = raw_result
|
||||
|
||||
else:
|
||||
types = [type(t).__name__ for t in raw_result][:10]
|
||||
@@ -377,55 +336,17 @@ class LitAgentRunner(Runner[T_task]):
|
||||
f"but got: {', '.join(types)}..."
|
||||
)
|
||||
|
||||
if not result_recognized:
|
||||
raise TypeError(
|
||||
f"Invalid raw result type. It's expected to be none, float, or a list of ReadableSpan or Span, "
|
||||
f"but got: {type(raw_result).__name__}..."
|
||||
)
|
||||
|
||||
return trace_spans
|
||||
|
||||
async def _emit_heartbeat(self, store: LightningStore) -> None:
|
||||
"""Send a heartbeat tick to the store.
|
||||
|
||||
Args:
|
||||
store: The lightning store to update.
|
||||
"""
|
||||
logger.debug(f"{self._log_prefix()} Preparing to emit heartbeat.")
|
||||
"""Send a heartbeat tick to the store."""
|
||||
worker_id = self.get_worker_id()
|
||||
|
||||
try:
|
||||
snapshot = await asyncio.wait_for(
|
||||
asyncio.to_thread(system_snapshot, self._heartbeat_include_gpu),
|
||||
timeout=self._heartbeat_interval,
|
||||
)
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat snapshot acquired.")
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"%s Heartbeat snapshot acquisition timed out after %.1fs, skipping.",
|
||||
self._log_prefix(),
|
||||
self._heartbeat_interval,
|
||||
)
|
||||
return
|
||||
await store.update_worker(worker_id, system_snapshot())
|
||||
except asyncio.CancelledError:
|
||||
# bypass the exception
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("%s Unable to acquire heartbeat snapshot.", self._log_prefix())
|
||||
return
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(store.update_worker(worker_id, snapshot), timeout=self._heartbeat_interval)
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat updated successfully.")
|
||||
except asyncio.CancelledError:
|
||||
# bypass the exception
|
||||
raise
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"%s update worker heartbeat timed out after %.1fs, skipping.",
|
||||
self._log_prefix(),
|
||||
self._heartbeat_interval,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("%s Unable to update worker heartbeat.", self._log_prefix())
|
||||
|
||||
@@ -440,161 +361,51 @@ class LitAgentRunner(Runner[T_task]):
|
||||
return None
|
||||
|
||||
if self._heartbeat_launch_mode == "asyncio":
|
||||
return self._start_heartbeat_asyncio_loop(store)
|
||||
if self._heartbeat_launch_mode == "thread":
|
||||
return self._start_heartbeat_thread_loop(store)
|
||||
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
def _start_heartbeat_asyncio_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
|
||||
"""Start a background heartbeat loop using asyncio.
|
||||
|
||||
Args:
|
||||
store: The lightning store to update.
|
||||
|
||||
Returns:
|
||||
An async stopper function that can be used to stop the heartbeat loop.
|
||||
"""
|
||||
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
async def heartbeat_loop() -> None:
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
# Run _emit_heartbeat in thread pool to avoid blocking the event loop.
|
||||
# Timeout at the interval - if it takes longer, the data is stale anyway.
|
||||
async def heartbeat_loop() -> None:
|
||||
while not stop_event.is_set():
|
||||
await self._emit_heartbeat(store)
|
||||
except Exception:
|
||||
logger.exception("%s Heartbeat failed.", self._log_prefix())
|
||||
with suppress(asyncio.TimeoutError):
|
||||
with suppress(asyncio.TimeoutError):
|
||||
interval = self._heartbeat_interval + self._random_state.uniform(
|
||||
-self._interval_jitter, self._interval_jitter
|
||||
)
|
||||
interval = max(interval, 0.01)
|
||||
await asyncio.wait_for(stop_event.wait(), timeout=interval)
|
||||
|
||||
task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
|
||||
|
||||
async def stop() -> None:
|
||||
stop_event.set()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
return stop
|
||||
|
||||
if self._heartbeat_launch_mode == "thread":
|
||||
stop_evt = threading.Event()
|
||||
|
||||
def thread_worker() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
while not stop_evt.is_set():
|
||||
loop.run_until_complete(self._emit_heartbeat(store))
|
||||
interval = self._heartbeat_interval + self._random_state.uniform(
|
||||
-self._interval_jitter, self._interval_jitter
|
||||
)
|
||||
interval = max(interval, 0.01)
|
||||
await asyncio.wait_for(stop_event.wait(), timeout=interval)
|
||||
stop_evt.wait(interval)
|
||||
|
||||
task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
|
||||
thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
|
||||
thread.start()
|
||||
|
||||
async def stop() -> None:
|
||||
stop_event.set()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
async def stop() -> None:
|
||||
stop_evt.set()
|
||||
await asyncio.to_thread(thread.join)
|
||||
|
||||
return stop
|
||||
return stop
|
||||
|
||||
def _start_heartbeat_thread_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
|
||||
"""Start a background heartbeat loop using threading.
|
||||
|
||||
It uses two threads: one to produce the snapshot and one to consume it,
|
||||
to avoid either of them blocking the event loop.
|
||||
|
||||
Args:
|
||||
store: The lightning store to update.
|
||||
|
||||
Returns:
|
||||
An async stopper function that can be used to stop the heartbeat loop.
|
||||
"""
|
||||
stop_evt = threading.Event()
|
||||
lock = threading.Lock()
|
||||
|
||||
latest_snapshot = None
|
||||
latest_ts = 0.0 # time.monotonic() when snapshot was captured
|
||||
|
||||
# Consider snapshot stale after ~1 interval plus jitter slack.
|
||||
stale_after = self._heartbeat_interval + self._interval_jitter + 1.0
|
||||
|
||||
worker_id = self.get_worker_id()
|
||||
|
||||
def producer() -> None:
|
||||
nonlocal latest_snapshot, latest_ts
|
||||
while not stop_evt.is_set():
|
||||
try:
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat producer: acquiring snapshot.")
|
||||
snap = system_snapshot(self._heartbeat_include_gpu) # sync
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat producer: snapshot acquired.")
|
||||
ts = time.monotonic()
|
||||
with lock:
|
||||
latest_snapshot = snap
|
||||
latest_ts = ts
|
||||
except Exception:
|
||||
logger.warning("%s Heartbeat producer: system_snapshot failed.", self._log_prefix(), exc_info=True)
|
||||
|
||||
interval = self._heartbeat_interval + self._random_state.uniform(
|
||||
-self._interval_jitter, self._interval_jitter
|
||||
)
|
||||
stop_evt.wait(max(interval, 0.01))
|
||||
|
||||
def consumer() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
last_warned_ts = None # Track which snapshot we've already warned about
|
||||
try:
|
||||
while not stop_evt.is_set():
|
||||
with lock:
|
||||
snap = latest_snapshot
|
||||
ts = latest_ts
|
||||
|
||||
wait_interval = max(
|
||||
self._heartbeat_interval
|
||||
+ self._random_state.uniform(-self._interval_jitter, self._interval_jitter),
|
||||
0.01,
|
||||
)
|
||||
|
||||
if snap is None:
|
||||
# probably just started
|
||||
logger.debug("%s Heartbeat consumer: no snapshot yet; skipping update.", self._log_prefix())
|
||||
stop_evt.wait(wait_interval)
|
||||
continue
|
||||
|
||||
age = time.monotonic() - ts
|
||||
if age > stale_after:
|
||||
# Only warn once per stale snapshot (check if we haven't warned about this timestamp yet)
|
||||
if last_warned_ts != ts:
|
||||
logger.warning(
|
||||
"%s Heartbeat consumer: snapshot stale (age=%.2fs > %.2fs); skipping update.",
|
||||
self._log_prefix(),
|
||||
age,
|
||||
stale_after,
|
||||
)
|
||||
last_warned_ts = ts
|
||||
stop_evt.wait(wait_interval)
|
||||
continue
|
||||
|
||||
try:
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat consumer: updating worker.")
|
||||
loop.run_until_complete(
|
||||
asyncio.wait_for(
|
||||
store.update_worker(worker_id, snap),
|
||||
timeout=self._heartbeat_interval,
|
||||
)
|
||||
)
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat consumer: worker updated.")
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"%s Heartbeat consumer: update timed out after %.1fs.",
|
||||
self._log_prefix(),
|
||||
self._heartbeat_interval,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("%s Heartbeat consumer: update failed.", self._log_prefix(), exc_info=True)
|
||||
|
||||
stop_evt.wait(wait_interval)
|
||||
finally:
|
||||
with suppress(Exception):
|
||||
loop.stop()
|
||||
with suppress(Exception):
|
||||
loop.close()
|
||||
|
||||
t_prod = threading.Thread(target=producer, name=f"{worker_id}-heartbeat-producer", daemon=True)
|
||||
t_cons = threading.Thread(target=consumer, name=f"{worker_id}-heartbeat-consumer", daemon=True)
|
||||
t_prod.start()
|
||||
t_cons.start()
|
||||
|
||||
async def stop() -> None:
|
||||
stop_evt.set()
|
||||
await asyncio.to_thread(t_prod.join)
|
||||
await asyncio.to_thread(t_cons.join)
|
||||
|
||||
return stop
|
||||
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
|
||||
|
||||
async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
|
||||
"""Sleep until the next poll interval, with optional event-based interruption.
|
||||
@@ -650,8 +461,6 @@ class LitAgentRunner(Runner[T_task]):
|
||||
logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
|
||||
return rollout_id
|
||||
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Resources fetched (id={resources_update.resources_id}).")
|
||||
|
||||
trace_spans: List[ReadableSpan] | List[Span] = []
|
||||
has_exception: bool = False
|
||||
|
||||
@@ -659,11 +468,9 @@ class LitAgentRunner(Runner[T_task]):
|
||||
await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
|
||||
|
||||
start_time = time.time()
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Prepared for trace context.")
|
||||
async with self._tracer.trace_context(
|
||||
name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
|
||||
):
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Entered trace context.")
|
||||
await self._trigger_hooks(
|
||||
hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
|
||||
)
|
||||
@@ -675,27 +482,21 @@ class LitAgentRunner(Runner[T_task]):
|
||||
rollout_method = (
|
||||
agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
|
||||
)
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Starting async rollout method.")
|
||||
result = await rollout_method(
|
||||
next_rollout.input, resources=resources_update.resources, rollout=next_rollout
|
||||
)
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Async rollout method completed.")
|
||||
else:
|
||||
rollout_method = (
|
||||
agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
|
||||
)
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Starting sync rollout method.")
|
||||
result = rollout_method(
|
||||
next_rollout.input, resources=resources_update.resources, rollout=next_rollout
|
||||
)
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Sync rollout method completed.")
|
||||
|
||||
await self._trigger_hooks(
|
||||
hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
|
||||
)
|
||||
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} Trace context exited.")
|
||||
|
||||
# Possible exceptions in post_process will be caught in the overall exception handler
|
||||
trace_spans = await self._post_process_rollout_result(next_rollout, result)
|
||||
last_reward = find_final_reward(trace_spans)
|
||||
@@ -765,7 +566,6 @@ class LitAgentRunner(Runner[T_task]):
|
||||
while not (event is not None and event.is_set()):
|
||||
logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
|
||||
next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
|
||||
logger.debug(f"{self._log_prefix()} Next rollout retrieved: {next_rollout}")
|
||||
if next_rollout is None:
|
||||
logger.debug(
|
||||
f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."
|
||||
@@ -777,6 +577,16 @@ class LitAgentRunner(Runner[T_task]):
|
||||
if next_rollout is None:
|
||||
return
|
||||
|
||||
try:
|
||||
# Claim the rollout but updating the current worker id
|
||||
await store.update_attempt(
|
||||
next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
|
||||
)
|
||||
except Exception:
|
||||
# This exception could happen if the rollout is dequeued and the other end died for some reason
|
||||
logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
|
||||
continue
|
||||
|
||||
# Execute the step
|
||||
await self._step_impl(next_rollout)
|
||||
|
||||
@@ -830,8 +640,12 @@ class LitAgentRunner(Runner[T_task]):
|
||||
else:
|
||||
resources_id = None
|
||||
|
||||
attempted_rollout = await self.get_store().start_rollout(
|
||||
input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
|
||||
attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
|
||||
# Register the attempt as running by the current worker
|
||||
await self.get_store().update_attempt(
|
||||
attempted_rollout.rollout_id,
|
||||
attempted_rollout.attempt.attempt_id,
|
||||
worker_id=self.get_worker_id(),
|
||||
)
|
||||
rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from agentlightning.client import AgentLightningClient
|
||||
from agentlightning.litagent import LitAgent
|
||||
from agentlightning.litagent.litagent import is_v0_1_rollout_api
|
||||
from agentlightning.tracer.base import Tracer
|
||||
from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet
|
||||
from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
|
||||
|
||||
from .base import Runner
|
||||
|
||||
@@ -99,7 +99,7 @@ class LegacyAgentRunner(Runner[Any]):
|
||||
trace: Any = None
|
||||
final_reward: Optional[float] = None
|
||||
triplets: Optional[List[Triplet]] = None
|
||||
trace_spans: Optional[List[SpanLike]] = None
|
||||
trace_spans: Optional[List[ReadableSpan]] = None
|
||||
|
||||
# Handle different types of results from the agent
|
||||
# Case 1: result is a float (final reward)
|
||||
@@ -108,14 +108,10 @@ class LegacyAgentRunner(Runner[Any]):
|
||||
# Case 2: result is a list of Triplets
|
||||
if isinstance(result, list) and all(isinstance(t, Triplet) for t in result):
|
||||
triplets = result # type: ignore
|
||||
# Case 3.1: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if isinstance(result, list) and all(isinstance(t, (ReadableSpan)) for t in result):
|
||||
# Case 3: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if isinstance(result, list) and all(isinstance(t, ReadableSpan) for t in result):
|
||||
trace_spans = result # type: ignore
|
||||
trace = [json.loads(readable_span.to_json()) for readable_span in trace_spans] # type: ignore
|
||||
# Case 3.2: result is a list of Span (Agent-lightning spans)
|
||||
if isinstance(result, list) and all(isinstance(t, Span) for t in result):
|
||||
trace_spans = result # type: ignore
|
||||
trace = [span.model_dump() for span in trace_spans] # type: ignore
|
||||
# Case 4: result is a list of dict (trace JSON)
|
||||
if isinstance(result, list) and all(isinstance(t, dict) for t in result):
|
||||
trace = result
|
||||
@@ -127,9 +123,10 @@ class LegacyAgentRunner(Runner[Any]):
|
||||
|
||||
# If the agent has tracing enabled, use the tracer's last trace if not already set
|
||||
if self.tracer and (trace is None or trace_spans is None):
|
||||
trace_spans = self.tracer.get_last_trace() # type: ignore
|
||||
if trace_spans:
|
||||
trace = [cast(Span, span).model_dump() for span in trace_spans]
|
||||
spans = self.tracer.get_last_trace()
|
||||
if spans:
|
||||
trace = [json.loads(readable_span.to_json()) for readable_span in spans]
|
||||
trace_spans = spans
|
||||
|
||||
# Always extract triplets from the trace using TracerTraceToTriplet
|
||||
if trace_spans:
|
||||
|
||||
@@ -29,14 +29,6 @@ AGL_EXCEPTION = "agentlightning.exception"
|
||||
Used by the exception emitter to record exception details.
|
||||
"""
|
||||
|
||||
AGL_OPERATION = "agentlightning.operation"
|
||||
"""Agent-lightning's standard span name for functions.
|
||||
Wrap function or code-blocks as operations.
|
||||
"""
|
||||
|
||||
AGL_REWARD = "agentlightning.reward"
|
||||
"""Agent-lightning's standard span name for reward operations."""
|
||||
|
||||
AGL_VIRTUAL = "agentlightning.virtual"
|
||||
"""Agent-lightning's standard span name for virtual operations.
|
||||
|
||||
@@ -56,9 +48,6 @@ class LightningResourceAttributes(Enum):
|
||||
SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
|
||||
"""Resource name for span sequence ID in Agent-lightning spans."""
|
||||
|
||||
TRACER_NAME = "agentlightning.tracer.name"
|
||||
"""Which tracer is used to create this span."""
|
||||
|
||||
|
||||
class LightningSpanAttributes(Enum):
|
||||
"""Attribute names that commonly appear in Agent-lightning spans.
|
||||
@@ -95,15 +84,6 @@ class LightningSpanAttributes(Enum):
|
||||
OBJECT_JSON = "agentlightning.object.json"
|
||||
"""Attribute name for object serialized value (JSON) in object spans."""
|
||||
|
||||
OPERATION_NAME = "agentlightning.operation.name"
|
||||
"""Attribute name for operation name in operation spans, normally the function name."""
|
||||
|
||||
OPERATION_INPUT = "agentlightning.operation.input"
|
||||
"""Attribute name for operation input in operation spans."""
|
||||
|
||||
OPERATION_OUTPUT = "agentlightning.operation.output"
|
||||
"""Attribute name for operation output in operation spans."""
|
||||
|
||||
|
||||
class RewardAttributes(Enum):
|
||||
"""Multi-dimensional reward attributes will look like:
|
||||
|
||||
@@ -10,12 +10,10 @@ from agentlightning.types import (
|
||||
Attempt,
|
||||
AttemptedRollout,
|
||||
AttemptStatus,
|
||||
EnqueueRolloutRequest,
|
||||
NamedResources,
|
||||
ResourcesUpdate,
|
||||
Rollout,
|
||||
RolloutConfig,
|
||||
RolloutMode,
|
||||
RolloutStatus,
|
||||
Span,
|
||||
TaskInput,
|
||||
@@ -158,11 +156,10 @@ class LightningStore:
|
||||
async def start_rollout(
|
||||
self,
|
||||
input: TaskInput,
|
||||
mode: RolloutMode | None = None,
|
||||
mode: Literal["train", "val", "test"] | None = None,
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
worker_id: str | None = None,
|
||||
) -> AttemptedRollout:
|
||||
"""Register a rollout and immediately create its first attempt.
|
||||
|
||||
@@ -185,7 +182,6 @@ class LightningStore:
|
||||
resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
|
||||
config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
|
||||
metadata: Free-form metadata persisted verbatim with the rollout.
|
||||
worker_id: Optional worker identifier to associate the new attempt with.
|
||||
|
||||
Returns:
|
||||
The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
|
||||
@@ -231,22 +227,6 @@ class LightningStore:
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
|
||||
"""Persist multiple rollouts in `queuing` state.
|
||||
|
||||
The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
|
||||
per request and preserves the input ordering. Subclasses can override to provide
|
||||
more efficient bulk enqueue semantics.
|
||||
|
||||
Args:
|
||||
rollouts: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
|
||||
parameters. Each entry requires `input` and can optionally include other fields.
|
||||
|
||||
Returns:
|
||||
Rollouts enqueued in the same order as `rollouts`.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
|
||||
"""Claim the oldest queued rollout and transition it to `preparing`.
|
||||
|
||||
@@ -263,9 +243,6 @@ class LightningStore:
|
||||
* Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
|
||||
(e.g., `last_dequeue_time`) when `worker_id` is provided.
|
||||
|
||||
Args:
|
||||
worker_id: Optional worker identifier to associate the claimed attempt with.
|
||||
|
||||
Returns:
|
||||
The next attempt to execute, or `None` when no eligible rollouts are queued.
|
||||
|
||||
@@ -274,30 +251,7 @@ class LightningStore:
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def dequeue_many_rollouts(
|
||||
self,
|
||||
*,
|
||||
limit: int = 1,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> Sequence[AttemptedRollout]:
|
||||
"""Claim up to `limit` queued rollouts without blocking.
|
||||
|
||||
The implementation can repeatedly invokes
|
||||
[`dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] until reaching
|
||||
the requested limit or the queue is empty. Subclasses can override it to fetch
|
||||
multiple rollouts atomically.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of rollouts to claim. Non-positive values return an empty list.
|
||||
worker_id: Optional worker identifier passed through to each dequeue call.
|
||||
|
||||
Returns:
|
||||
Attempted rollouts claimed in FIFO order. May contain fewer than `limit` entries
|
||||
when the queue is exhausted.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
|
||||
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
|
||||
"""Create a manual retry attempt for an existing rollout.
|
||||
|
||||
This is typically invoked by runners that wish to retry outside of the
|
||||
@@ -308,7 +262,6 @@ class LightningStore:
|
||||
|
||||
Args:
|
||||
rollout_id: Unique identifier of the rollout receiving a new attempt.
|
||||
worker_id: Optional worker identifier to associate the new attempt with.
|
||||
|
||||
Returns:
|
||||
The rollout paired with its newly-created attempt.
|
||||
@@ -766,8 +719,7 @@ class LightningStore:
|
||||
|
||||
When `attempt_id` is `"latest"` the update must target the attempt with the highest
|
||||
`sequence_id`; otherwise it must target the specific attempt. Implementations should
|
||||
propagate status changes to the rollout (for example
|
||||
via [`rollout_status_from_attempt()`][agentlightning.store.utils.rollout_status_from_attempt])
|
||||
propagate status changes to the rollout (for example via [`propagate_status()`][agentlightning.store.utils.propagate_status])
|
||||
once the latest attempt transitions to a terminal state.
|
||||
|
||||
Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],
|
||||
|
||||
@@ -47,7 +47,6 @@ from agentlightning.types import (
|
||||
Attempt,
|
||||
AttemptedRollout,
|
||||
AttemptStatus,
|
||||
EnqueueRolloutRequest,
|
||||
NamedResources,
|
||||
PaginatedResult,
|
||||
ResourcesUpdate,
|
||||
@@ -59,12 +58,10 @@ from agentlightning.types import (
|
||||
Worker,
|
||||
WorkerStatus,
|
||||
)
|
||||
from agentlightning.utils.metrics import MetricsBackend, get_prometheus_registry
|
||||
from agentlightning.utils.otlp import handle_otlp_export, spans_from_proto
|
||||
from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncher, PythonServerLauncherArgs
|
||||
|
||||
from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
|
||||
from .collection.base import resolve_error_type
|
||||
from .utils import LATENCY_BUCKETS
|
||||
|
||||
server_logger = logging.getLogger("agentlightning.store.server")
|
||||
@@ -84,26 +81,12 @@ class RolloutRequest(BaseModel):
|
||||
resources_id: Optional[str] = None
|
||||
config: Optional[RolloutConfig] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
worker_id: Optional[str] = None
|
||||
|
||||
|
||||
class DequeueRolloutRequest(BaseModel):
|
||||
worker_id: Optional[str] = None
|
||||
|
||||
|
||||
class StartAttemptRequest(BaseModel):
|
||||
worker_id: Optional[str] = None
|
||||
|
||||
|
||||
class EnqueueManyRolloutsRequest(BaseModel):
|
||||
rollouts: List[EnqueueRolloutRequest]
|
||||
|
||||
|
||||
class DequeueManyRolloutsRequest(BaseModel):
|
||||
limit: int = 1
|
||||
worker_id: Optional[str] = None
|
||||
|
||||
|
||||
class QueryRolloutsRequest(BaseModel):
|
||||
status_in: Optional[List[RolloutStatus]] = Field(FastAPIQuery(default=None))
|
||||
rollout_id_in: Optional[List[str]] = Field(FastAPIQuery(default=None))
|
||||
@@ -240,7 +223,7 @@ class LightningStoreServer(LightningStore):
|
||||
launcher_args: The arguments to use for the server launcher.
|
||||
It's not allowed to set `host`, `port`, `launch_mode` together with `launcher_args`.
|
||||
n_workers: The number of workers to run in the server. Only applicable for `mp` launch mode.
|
||||
tracker: The metrics tracker to use for the server.
|
||||
prometheus: Whether to enable Prometheus metrics.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -252,7 +235,7 @@ class LightningStoreServer(LightningStore):
|
||||
launch_mode: LaunchMode = "thread",
|
||||
launcher_args: PythonServerLauncherArgs | None = None,
|
||||
n_workers: int = 1,
|
||||
tracker: MetricsBackend | None = None,
|
||||
prometheus: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.store = store
|
||||
@@ -270,7 +253,6 @@ class LightningStoreServer(LightningStore):
|
||||
port=port,
|
||||
launch_mode=launch_mode,
|
||||
healthcheck_url=API_V1_AGL_PREFIX + "/health",
|
||||
n_workers=n_workers,
|
||||
)
|
||||
|
||||
store_capabilities = self.store.capabilities
|
||||
@@ -290,7 +272,7 @@ class LightningStoreServer(LightningStore):
|
||||
app=self.app,
|
||||
args=self.launcher_args,
|
||||
)
|
||||
self._tracker = tracker
|
||||
self._prometheus = prometheus
|
||||
|
||||
self._lock: threading.Lock = threading.Lock()
|
||||
self._cors_allow_origins = self._normalize_cors_origins(cors_allow_origins)
|
||||
@@ -335,6 +317,7 @@ class LightningStoreServer(LightningStore):
|
||||
return {
|
||||
"launcher_args": self.launcher_args,
|
||||
"server_launcher": self.server_launcher,
|
||||
"_prometheus": self._prometheus,
|
||||
"_owner_pid": self._owner_pid,
|
||||
}
|
||||
|
||||
@@ -352,12 +335,11 @@ class LightningStoreServer(LightningStore):
|
||||
self.store = None
|
||||
self.launcher_args = state["launcher_args"]
|
||||
self.server_launcher = state["server_launcher"]
|
||||
self._tracker = None
|
||||
self._prometheus = state["_prometheus"]
|
||||
self._owner_pid = state["_owner_pid"]
|
||||
self._cors_allow_origins = state.get("_cors_allow_origins")
|
||||
self._client = None
|
||||
self._lock = threading.Lock()
|
||||
self._prometheus_registry = None
|
||||
# Do NOT reconstruct app, _uvicorn_config, _uvicorn_server
|
||||
# to avoid transferring server state to subprocess
|
||||
|
||||
@@ -438,10 +420,9 @@ class LightningStoreServer(LightningStore):
|
||||
api = APIRouter(prefix=API_V1_PREFIX)
|
||||
|
||||
# The outermost-layer of monitoring
|
||||
if self._tracker is not None:
|
||||
self._setup_metrics(api=api, app=self.app)
|
||||
if self._prometheus:
|
||||
self._setup_prometheus(api=api, app=self.app)
|
||||
|
||||
# TODO: This should only be enabled in development mode.
|
||||
@self.app.middleware("http")
|
||||
async def _app_exception_handler( # pyright: ignore[reportUnusedFunction]
|
||||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||||
@@ -477,9 +458,7 @@ class LightningStoreServer(LightningStore):
|
||||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||||
):
|
||||
# If not API request, just pass through
|
||||
if not request.url.path.startswith(API_V1_AGL_PREFIX) and not request.url.path.startswith(
|
||||
API_V1_PREFIX + "/traces"
|
||||
):
|
||||
if not request.url.path.startswith(API_V1_AGL_PREFIX):
|
||||
return await call_next(request)
|
||||
|
||||
start = time.perf_counter()
|
||||
@@ -543,38 +522,22 @@ class LightningStoreServer(LightningStore):
|
||||
async def health(): # pyright: ignore[reportUnusedFunction]
|
||||
return {"status": "ok"}
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=List[Rollout])
|
||||
async def enqueue_rollouts( # pyright: ignore[reportUnusedFunction]
|
||||
request: EnqueueManyRolloutsRequest,
|
||||
) -> List[Rollout]:
|
||||
enqueue_requests = request.rollouts
|
||||
if not enqueue_requests:
|
||||
return []
|
||||
if len(enqueue_requests) == 1:
|
||||
single = enqueue_requests[0]
|
||||
rollout = await self.enqueue_rollout(
|
||||
input=single.input,
|
||||
mode=single.mode,
|
||||
resources_id=single.resources_id,
|
||||
config=single.config,
|
||||
metadata=single.metadata,
|
||||
)
|
||||
return [rollout]
|
||||
rollouts = await self.enqueue_many_rollouts(enqueue_requests)
|
||||
return list(rollouts)
|
||||
@api.post(API_AGL_PREFIX + "/queues/rollouts/enqueue", status_code=201, response_model=Rollout)
|
||||
async def enqueue_rollout(request: RolloutRequest): # pyright: ignore[reportUnusedFunction]
|
||||
return await self.enqueue_rollout(
|
||||
input=request.input,
|
||||
mode=request.mode,
|
||||
resources_id=request.resources_id,
|
||||
config=request.config,
|
||||
metadata=request.metadata,
|
||||
)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=List[AttemptedRollout])
|
||||
async def dequeue_rollouts( # pyright: ignore[reportUnusedFunction]
|
||||
request: DequeueManyRolloutsRequest | None = Body(None),
|
||||
) -> List[AttemptedRollout]:
|
||||
payload = request or DequeueManyRolloutsRequest()
|
||||
if payload.limit <= 0:
|
||||
return []
|
||||
if payload.limit == 1:
|
||||
single = await self.dequeue_rollout(worker_id=payload.worker_id)
|
||||
return [single] if single else []
|
||||
rollouts = await self.dequeue_many_rollouts(limit=payload.limit, worker_id=payload.worker_id)
|
||||
return list(rollouts)
|
||||
@api.post(API_AGL_PREFIX + "/queues/rollouts/dequeue", response_model=Optional[AttemptedRollout])
|
||||
async def dequeue_rollout( # pyright: ignore[reportUnusedFunction]
|
||||
request: DequeueRolloutRequest | None = Body(None),
|
||||
):
|
||||
worker_id = request.worker_id if request else None
|
||||
return await self.dequeue_rollout(worker_id=worker_id)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/rollouts", status_code=201, response_model=AttemptedRollout)
|
||||
async def start_rollout(request: RolloutRequest): # pyright: ignore[reportUnusedFunction]
|
||||
@@ -584,7 +547,6 @@ class LightningStoreServer(LightningStore):
|
||||
resources_id=request.resources_id,
|
||||
config=request.config,
|
||||
metadata=request.metadata,
|
||||
worker_id=request.worker_id,
|
||||
)
|
||||
|
||||
@api.get(API_AGL_PREFIX + "/rollouts", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
|
||||
@@ -603,24 +565,6 @@ class LightningStoreServer(LightningStore):
|
||||
)
|
||||
return _build_paginated_response(results, limit=params.limit, offset=params.offset)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/rollouts/search", response_model=PaginatedResult[Union[AttemptedRollout, Rollout]])
|
||||
async def search_rollouts(request: QueryRolloutsRequest): # pyright: ignore[reportUnusedFunction]
|
||||
_validate_paginated_request(request, Rollout)
|
||||
status_in = request.status_in if "status_in" in request.model_fields_set else None
|
||||
rollout_id_in = request.rollout_id_in if "rollout_id_in" in request.model_fields_set else None
|
||||
# Get all rollouts from the underlying store
|
||||
results = await self.query_rollouts(
|
||||
status_in=status_in,
|
||||
rollout_id_in=rollout_id_in,
|
||||
rollout_id_contains=request.rollout_id_contains,
|
||||
filter_logic=request.filter_logic,
|
||||
sort_by=request.sort_by,
|
||||
sort_order=request.sort_order,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return _build_paginated_response(results, limit=request.limit, offset=request.offset)
|
||||
|
||||
@api.get(API_AGL_PREFIX + "/rollouts/{rollout_id}", response_model=Union[AttemptedRollout, Rollout])
|
||||
async def get_rollout_by_id(rollout_id: str): # pyright: ignore[reportUnusedFunction]
|
||||
return await self.get_rollout_by_id(rollout_id)
|
||||
@@ -653,25 +597,8 @@ class LightningStoreServer(LightningStore):
|
||||
)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts", status_code=201, response_model=AttemptedRollout)
|
||||
async def start_attempt( # pyright: ignore[reportUnusedFunction]
|
||||
rollout_id: str, request: StartAttemptRequest | None = Body(None)
|
||||
):
|
||||
worker_id = request.worker_id if request else None
|
||||
return await self.start_attempt(rollout_id, worker_id=worker_id)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/search", response_model=PaginatedResult[Attempt])
|
||||
async def search_attempts( # pyright: ignore[reportUnusedFunction]
|
||||
rollout_id: str, request: QueryAttemptsRequest
|
||||
):
|
||||
_validate_paginated_request(request, Attempt)
|
||||
attempts = await self.query_attempts(
|
||||
rollout_id,
|
||||
sort_by=request.sort_by,
|
||||
sort_order=request.sort_order,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return _build_paginated_response(attempts, limit=request.limit, offset=request.offset)
|
||||
async def start_attempt(rollout_id: str): # pyright: ignore[reportUnusedFunction]
|
||||
return await self.start_attempt(rollout_id)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/rollouts/{rollout_id}/attempts/{attempt_id}", response_model=Attempt)
|
||||
async def update_attempt( # pyright: ignore[reportUnusedFunction]
|
||||
@@ -700,21 +627,6 @@ class LightningStoreServer(LightningStore):
|
||||
)
|
||||
return _build_paginated_response(workers, limit=params.limit, offset=params.offset)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/workers/search", response_model=PaginatedResult[Worker])
|
||||
async def search_workers(request: QueryWorkersRequest): # pyright: ignore[reportUnusedFunction]
|
||||
_validate_paginated_request(request, Worker)
|
||||
status_in = request.status_in if "status_in" in request.model_fields_set else None
|
||||
workers = await self.query_workers(
|
||||
status_in=status_in,
|
||||
worker_id_contains=request.worker_id_contains,
|
||||
filter_logic=request.filter_logic,
|
||||
sort_by=request.sort_by,
|
||||
sort_order=request.sort_order,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return _build_paginated_response(workers, limit=request.limit, offset=request.offset)
|
||||
|
||||
@api.get(API_AGL_PREFIX + "/workers/{worker_id}", response_model=Optional[Worker])
|
||||
async def get_worker(worker_id: str): # pyright: ignore[reportUnusedFunction]
|
||||
return await self.get_worker_by_id(worker_id)
|
||||
@@ -807,28 +719,6 @@ class LightningStoreServer(LightningStore):
|
||||
)
|
||||
return _build_paginated_response(spans, limit=params.limit, offset=params.offset)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/spans/search", response_model=PaginatedResult[Span])
|
||||
async def search_spans(request: QuerySpansRequest): # pyright: ignore[reportUnusedFunction]
|
||||
_validate_paginated_request(request, Span)
|
||||
spans = await self.query_spans(
|
||||
request.rollout_id,
|
||||
request.attempt_id,
|
||||
trace_id=request.trace_id,
|
||||
trace_id_contains=request.trace_id_contains,
|
||||
span_id=request.span_id,
|
||||
span_id_contains=request.span_id_contains,
|
||||
parent_id=request.parent_id,
|
||||
parent_id_contains=request.parent_id_contains,
|
||||
name=request.name,
|
||||
name_contains=request.name_contains,
|
||||
filter_logic=request.filter_logic,
|
||||
sort_by=request.sort_by,
|
||||
sort_order=request.sort_order,
|
||||
limit=request.limit,
|
||||
offset=request.offset,
|
||||
)
|
||||
return _build_paginated_response(spans, limit=request.limit, offset=request.offset)
|
||||
|
||||
@api.post(API_AGL_PREFIX + "/spans/next", response_model=NextSequenceIdResponse)
|
||||
async def get_next_span_sequence_id(request: NextSequenceIdRequest): # pyright: ignore[reportUnusedFunction]
|
||||
sequence_id = await self.get_next_span_sequence_id(request.rollout_id, request.attempt_id)
|
||||
@@ -847,34 +737,50 @@ class LightningStoreServer(LightningStore):
|
||||
# Finally, mount the dashboard assets
|
||||
self._setup_dashboard()
|
||||
|
||||
def _setup_metrics(self, api: APIRouter, app: FastAPI):
|
||||
def _setup_prometheus(self, api: APIRouter, app: FastAPI):
|
||||
"""Setup Prometheus metrics endpoints."""
|
||||
if self._tracker is None:
|
||||
return
|
||||
try:
|
||||
from prometheus_client import make_asgi_app # type: ignore
|
||||
from prometheus_client import (
|
||||
REGISTRY,
|
||||
CollectorRegistry,
|
||||
Counter,
|
||||
Histogram,
|
||||
multiprocess,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Prometheus client is not installed. Please either install it or set prometheus to False."
|
||||
)
|
||||
|
||||
self._tracker.register_counter(
|
||||
"agl.http.total",
|
||||
["path", "method", "status"],
|
||||
group_level=2,
|
||||
# Multi-process mode: https://prometheus.github.io/client_python/multiprocess/
|
||||
is_multiprocess = self.launcher_args.launch_mode == "mp" and self.launcher_args.n_workers > 1
|
||||
if is_multiprocess:
|
||||
registry = CollectorRegistry()
|
||||
multiprocess.MultiProcessCollector(registry)
|
||||
else:
|
||||
registry = REGISTRY
|
||||
|
||||
HTTP_REQUESTS = Counter(
|
||||
"http_requests_total",
|
||||
"Total HTTP requests",
|
||||
["method", "path", "status_code"],
|
||||
)
|
||||
self._tracker.register_histogram(
|
||||
"agl.http.latency",
|
||||
["path", "method", "status"],
|
||||
|
||||
HTTP_LATENCY = Histogram(
|
||||
"http_request_duration_seconds",
|
||||
"Latency of HTTP requests",
|
||||
["method", "path"],
|
||||
buckets=LATENCY_BUCKETS,
|
||||
group_level=2,
|
||||
)
|
||||
|
||||
def get_template_path(path: str) -> str:
|
||||
# Handle "latest" keywords BEFORE generic IDs
|
||||
if path.endswith("/attempts/latest") and "/rollouts/" in path:
|
||||
return re.sub(r"rollouts/[^/]+/attempts/latest$", "rollouts/{rollout_id}/attempts/latest", path)
|
||||
if path.endswith("/attempts/search") and "/rollouts/" in path:
|
||||
return re.sub(r"rollouts/[^/]+/attempts/search$", "rollouts/{rollout_id}/attempts/search", path)
|
||||
if path.endswith("/resources/latest"):
|
||||
elif path.endswith("/resources/latest"):
|
||||
return path
|
||||
if path.endswith("/search"):
|
||||
return path
|
||||
if "enqueue" in path or "dequeue" in path:
|
||||
elif "enqueue" in path or "dequeue" in path:
|
||||
return path
|
||||
|
||||
# Handle generic IDs
|
||||
@@ -887,55 +793,27 @@ class LightningStoreServer(LightningStore):
|
||||
return path
|
||||
|
||||
@app.middleware("http")
|
||||
async def tracking_middleware( # pyright: ignore[reportUnusedFunction]
|
||||
async def prometheus_http_middleware( # pyright: ignore[reportUnusedFunction]
|
||||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||||
) -> Response:
|
||||
if self._tracker is None:
|
||||
return await call_next(request)
|
||||
|
||||
start = time.perf_counter()
|
||||
status = 520 # Default to 520 if things crash hard
|
||||
response = await call_next(request)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
try:
|
||||
response = await call_next(request)
|
||||
status = response.status_code
|
||||
return response
|
||||
except asyncio.CancelledError:
|
||||
# Client disconnected (Timeout)
|
||||
status = 499 # Standard Nginx code for "Client Closed Request"
|
||||
server_logger.debug(f"Client disconnected (Timeout): {request.url.path}", exc_info=True)
|
||||
raise # Re-raise to let Uvicorn handle the cleanup
|
||||
except Exception as exc:
|
||||
status = resolve_error_type(exc)
|
||||
server_logger.debug(f"Server error: {request.url.path}", exc_info=True)
|
||||
raise
|
||||
finally:
|
||||
# This block executes NO MATTER WHAT happens above
|
||||
elapsed = time.perf_counter() - start
|
||||
# Strip the ID-specific URL parts
|
||||
path = get_template_path(request.url.path)
|
||||
method = request.method
|
||||
status = response.status_code
|
||||
|
||||
# Strip the ID-specific URL parts
|
||||
path = get_template_path(request.url.path)
|
||||
method = request.method
|
||||
HTTP_REQUESTS.labels(method, path, status).inc()
|
||||
HTTP_LATENCY.labels(method, path).observe(elapsed)
|
||||
|
||||
await self._tracker.inc_counter(
|
||||
"agl.http.total",
|
||||
labels={"method": method, "path": path, "status": str(status)},
|
||||
)
|
||||
await self._tracker.observe_histogram(
|
||||
"agl.http.latency",
|
||||
value=elapsed,
|
||||
labels={"method": method, "path": path, "status": str(status)},
|
||||
)
|
||||
return response
|
||||
|
||||
if self._tracker.has_prometheus():
|
||||
from prometheus_client import make_asgi_app # pyright: ignore[reportUnknownVariableType]
|
||||
metrics_app = make_asgi_app(registry=registry) # type: ignore
|
||||
|
||||
metrics_app = make_asgi_app( # pyright: ignore[reportUnknownVariableType]
|
||||
registry=get_prometheus_registry()
|
||||
)
|
||||
|
||||
# This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
|
||||
app.mount(api.prefix + "/prometheus", metrics_app) # pyright: ignore[reportUnknownArgumentType]
|
||||
# This App would need to be accessed via /v1/prometheus/ (note the trailing slash)
|
||||
app.mount(api.prefix + "/prometheus", metrics_app) # pyright: ignore[reportUnknownArgumentType]
|
||||
|
||||
def _setup_otlp(self, api: APIRouter):
|
||||
"""Setup OTLP endpoints."""
|
||||
@@ -1056,7 +934,6 @@ class LightningStoreServer(LightningStore):
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> AttemptedRollout:
|
||||
return await self._call_store_method(
|
||||
"start_rollout",
|
||||
@@ -1065,7 +942,6 @@ class LightningStoreServer(LightningStore):
|
||||
resources_id,
|
||||
config,
|
||||
metadata,
|
||||
worker_id,
|
||||
)
|
||||
|
||||
async def enqueue_rollout(
|
||||
@@ -1085,22 +961,11 @@ class LightningStoreServer(LightningStore):
|
||||
metadata,
|
||||
)
|
||||
|
||||
async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
|
||||
return await self._call_store_method("enqueue_many_rollouts", rollouts)
|
||||
|
||||
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
|
||||
return await self._call_store_method("dequeue_rollout", worker_id)
|
||||
|
||||
async def dequeue_many_rollouts(
|
||||
self,
|
||||
*,
|
||||
limit: int = 1,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> Sequence[AttemptedRollout]:
|
||||
return await self._call_store_method("dequeue_many_rollouts", limit=limit, worker_id=worker_id)
|
||||
|
||||
async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
|
||||
return await self._call_store_method("start_attempt", rollout_id, worker_id)
|
||||
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
|
||||
return await self._call_store_method("start_attempt", rollout_id)
|
||||
|
||||
async def query_rollouts(
|
||||
self,
|
||||
@@ -1520,7 +1385,7 @@ class LightningStoreClient(LightningStore):
|
||||
except aiohttp.ClientResponseError as cre:
|
||||
# Respect app-level 4xx as final
|
||||
# 4xx => application issue; do not retry (except 408 which is transient)
|
||||
client_logger.debug(f"ClientResponseError ({method} {path}): {cre.status} {cre.message}", exc_info=True)
|
||||
client_logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
|
||||
if 400 <= cre.status < 500 and cre.status != 408:
|
||||
raise
|
||||
# 5xx and others will be retried below if they raise
|
||||
@@ -1536,9 +1401,9 @@ class LightningStoreClient(LightningStore):
|
||||
asyncio.TimeoutError,
|
||||
) as net_exc:
|
||||
# Network/session issue: probe health before retrying
|
||||
client_logger.debug(f"Network/session issue ({method} {path}): {net_exc}", exc_info=True)
|
||||
client_logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
|
||||
last_exc = net_exc
|
||||
client_logger.info(f"Network/session issue: {net_exc} - will retry the request {method}: {path}")
|
||||
client_logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
|
||||
if not await self._wait_until_healthy(session):
|
||||
break # server is not healthy, do not retry
|
||||
|
||||
@@ -1574,7 +1439,6 @@ class LightningStoreClient(LightningStore):
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> AttemptedRollout:
|
||||
data = await self._request_json(
|
||||
"post",
|
||||
@@ -1585,7 +1449,6 @@ class LightningStoreClient(LightningStore):
|
||||
resources_id=resources_id,
|
||||
config=config,
|
||||
metadata=metadata,
|
||||
worker_id=worker_id,
|
||||
).model_dump(exclude_none=False),
|
||||
)
|
||||
return AttemptedRollout.model_validate(data)
|
||||
@@ -1598,64 +1461,18 @@ class LightningStoreClient(LightningStore):
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> Rollout:
|
||||
request_body = EnqueueManyRolloutsRequest(
|
||||
rollouts=[
|
||||
EnqueueRolloutRequest(
|
||||
input=input,
|
||||
mode=mode,
|
||||
resources_id=resources_id,
|
||||
config=config,
|
||||
metadata=metadata,
|
||||
)
|
||||
]
|
||||
).model_dump(exclude_none=False)
|
||||
data = await self._request_json(
|
||||
"post",
|
||||
"/queues/rollouts/enqueue",
|
||||
json=request_body,
|
||||
json=RolloutRequest(
|
||||
input=input,
|
||||
mode=mode,
|
||||
resources_id=resources_id,
|
||||
config=config,
|
||||
metadata=metadata,
|
||||
).model_dump(exclude_none=False),
|
||||
)
|
||||
if not data:
|
||||
raise RuntimeError("enqueue_rollout returned no rollouts")
|
||||
return Rollout.model_validate(data[0])
|
||||
|
||||
async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
|
||||
if not rollouts:
|
||||
return []
|
||||
request_body = EnqueueManyRolloutsRequest(rollouts=list(rollouts)).model_dump(exclude_none=False)
|
||||
data = await self._request_json(
|
||||
"post",
|
||||
"/queues/rollouts/enqueue",
|
||||
json=request_body,
|
||||
)
|
||||
return [Rollout.model_validate(entry) for entry in data]
|
||||
|
||||
async def _dequeue_batch(
|
||||
self,
|
||||
*,
|
||||
limit: int,
|
||||
worker_id: Optional[str],
|
||||
) -> List[AttemptedRollout]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
session = await self._get_session()
|
||||
url = f"{self.server_address}/queues/rollouts/dequeue"
|
||||
payload: Dict[str, Any] = {"limit": limit}
|
||||
if worker_id is not None:
|
||||
payload["worker_id"] = worker_id
|
||||
try:
|
||||
async with session.post(url, json=payload) as resp:
|
||||
resp.raise_for_status()
|
||||
data = await resp.json()
|
||||
self._dequeue_was_successful = True
|
||||
return [AttemptedRollout.model_validate(item) for item in data]
|
||||
except Exception as e:
|
||||
if self._dequeue_was_successful:
|
||||
if self._dequeue_first_unsuccessful:
|
||||
client_logger.warning(f"dequeue_rollout failed with exception: {e}")
|
||||
self._dequeue_first_unsuccessful = False
|
||||
client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
|
||||
# Else ignore the exception because the server is not ready yet
|
||||
return []
|
||||
return Rollout.model_validate(data)
|
||||
|
||||
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
|
||||
"""
|
||||
@@ -1668,23 +1485,30 @@ class LightningStoreClient(LightningStore):
|
||||
This method does NOT retry on failures. If any exception occurs (network error,
|
||||
server error, etc.), it logs the error and returns None immediately.
|
||||
"""
|
||||
attempts = await self._dequeue_batch(limit=1, worker_id=worker_id)
|
||||
return attempts[0] if attempts else None
|
||||
session = await self._get_session()
|
||||
url = f"{self.server_address}/queues/rollouts/dequeue"
|
||||
request_kwargs: Dict[str, Any] = {}
|
||||
if worker_id is not None:
|
||||
request_kwargs["json"] = {"worker_id": worker_id}
|
||||
try:
|
||||
async with session.post(url, **request_kwargs) as resp:
|
||||
resp.raise_for_status()
|
||||
data = await resp.json()
|
||||
self._dequeue_was_successful = True
|
||||
return AttemptedRollout.model_validate(data) if data else None
|
||||
except Exception as e:
|
||||
if self._dequeue_was_successful:
|
||||
if self._dequeue_first_unsuccessful:
|
||||
client_logger.warning(f"dequeue_rollout failed with exception: {e}")
|
||||
self._dequeue_first_unsuccessful = False
|
||||
client_logger.debug("dequeue_rollout failed with exception. Details:", exc_info=True)
|
||||
# Else ignore the exception because the server is not ready yet
|
||||
return None
|
||||
|
||||
async def dequeue_many_rollouts(
|
||||
self,
|
||||
*,
|
||||
limit: int = 1,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> Sequence[AttemptedRollout]:
|
||||
return await self._dequeue_batch(limit=limit, worker_id=worker_id)
|
||||
|
||||
async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
|
||||
payload = {"worker_id": worker_id} if worker_id is not None else None
|
||||
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
|
||||
data = await self._request_json(
|
||||
"post",
|
||||
f"/rollouts/{rollout_id}/attempts",
|
||||
json=payload,
|
||||
)
|
||||
return AttemptedRollout.model_validate(data)
|
||||
|
||||
@@ -1702,25 +1526,29 @@ class LightningStoreClient(LightningStore):
|
||||
status: Optional[Sequence[RolloutStatus]] = None,
|
||||
rollout_ids: Optional[Sequence[str]] = None,
|
||||
) -> PaginatedResult[Union[AttemptedRollout, Rollout]]:
|
||||
params_list: List[Tuple[str, Any]] = []
|
||||
|
||||
def _extend(key: str, values: Sequence[Any]) -> None:
|
||||
for value in values:
|
||||
params_list.append((key, value))
|
||||
|
||||
resolved_status = status_in if status_in is not None else status
|
||||
resolved_rollout_ids = rollout_id_in if rollout_id_in is not None else rollout_ids
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
if resolved_status is not None:
|
||||
payload["status_in"] = resolved_status
|
||||
_extend("status_in", resolved_status)
|
||||
if resolved_rollout_ids is not None:
|
||||
payload["rollout_id_in"] = resolved_rollout_ids
|
||||
_extend("rollout_id_in", resolved_rollout_ids)
|
||||
if rollout_id_contains is not None:
|
||||
payload["rollout_id_contains"] = rollout_id_contains
|
||||
payload["filter_logic"] = filter_logic
|
||||
params_list.append(("rollout_id_contains", rollout_id_contains))
|
||||
params_list.append(("filter_logic", filter_logic))
|
||||
if sort_by is not None:
|
||||
payload["sort_by"] = sort_by
|
||||
payload["sort_order"] = sort_order
|
||||
params_list.append(("sort_by", sort_by))
|
||||
params_list.append(("sort_order", sort_order))
|
||||
params_list.append(("limit", limit))
|
||||
params_list.append(("offset", offset))
|
||||
|
||||
data = await self._request_json("post", "/rollouts/search", json=payload)
|
||||
data = await self._request_json("get", "/rollouts", params=params_list or None)
|
||||
items = [
|
||||
(
|
||||
AttemptedRollout.model_validate(item)
|
||||
@@ -1740,14 +1568,14 @@ class LightningStoreClient(LightningStore):
|
||||
limit: int = -1,
|
||||
offset: int = 0,
|
||||
) -> PaginatedResult[Attempt]:
|
||||
payload: Dict[str, Any] = {
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
params: List[Tuple[str, Any]] = [
|
||||
("limit", limit),
|
||||
("offset", offset),
|
||||
]
|
||||
if sort_by is not None:
|
||||
payload["sort_by"] = sort_by
|
||||
payload["sort_order"] = sort_order
|
||||
data = await self._request_json("post", f"/rollouts/{rollout_id}/attempts/search", json=payload)
|
||||
params.append(("sort_by", sort_by))
|
||||
params.append(("sort_order", sort_order))
|
||||
data = await self._request_json("get", f"/rollouts/{rollout_id}/attempts", params=params)
|
||||
items = [Attempt.model_validate(item) for item in data["items"]]
|
||||
return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
|
||||
|
||||
@@ -1967,30 +1795,32 @@ class LightningStoreClient(LightningStore):
|
||||
sort_by: Optional[str] = "sequence_id",
|
||||
sort_order: Literal["asc", "desc"] = "asc",
|
||||
) -> PaginatedResult[Span]:
|
||||
payload: Dict[str, Any] = {"rollout_id": rollout_id, "limit": limit, "offset": offset}
|
||||
params: List[Tuple[str, Any]] = [("rollout_id", rollout_id)]
|
||||
if attempt_id is not None:
|
||||
payload["attempt_id"] = attempt_id
|
||||
params.append(("attempt_id", attempt_id))
|
||||
if trace_id is not None:
|
||||
payload["trace_id"] = trace_id
|
||||
params.append(("trace_id", trace_id))
|
||||
if trace_id_contains is not None:
|
||||
payload["trace_id_contains"] = trace_id_contains
|
||||
params.append(("trace_id_contains", trace_id_contains))
|
||||
if span_id is not None:
|
||||
payload["span_id"] = span_id
|
||||
params.append(("span_id", span_id))
|
||||
if span_id_contains is not None:
|
||||
payload["span_id_contains"] = span_id_contains
|
||||
params.append(("span_id_contains", span_id_contains))
|
||||
if parent_id is not None:
|
||||
payload["parent_id"] = parent_id
|
||||
params.append(("parent_id", parent_id))
|
||||
if parent_id_contains is not None:
|
||||
payload["parent_id_contains"] = parent_id_contains
|
||||
params.append(("parent_id_contains", parent_id_contains))
|
||||
if name is not None:
|
||||
payload["name"] = name
|
||||
params.append(("name", name))
|
||||
if name_contains is not None:
|
||||
payload["name_contains"] = name_contains
|
||||
payload["filter_logic"] = filter_logic
|
||||
params.append(("name_contains", name_contains))
|
||||
params.append(("filter_logic", filter_logic))
|
||||
if sort_by is not None:
|
||||
payload["sort_by"] = sort_by
|
||||
payload["sort_order"] = sort_order
|
||||
data = await self._request_json("post", "/spans/search", json=payload)
|
||||
params.append(("sort_by", sort_by))
|
||||
params.append(("sort_order", sort_order))
|
||||
params.append(("limit", limit))
|
||||
params.append(("offset", offset))
|
||||
data = await self._request_json("get", "/spans", params=params)
|
||||
items = [Span.model_validate(item) for item in data["items"]]
|
||||
return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
|
||||
|
||||
@@ -2058,17 +1888,21 @@ class LightningStoreClient(LightningStore):
|
||||
limit: int = -1,
|
||||
offset: int = 0,
|
||||
) -> PaginatedResult[Worker]:
|
||||
payload: Dict[str, Any] = {}
|
||||
params: List[Tuple[str, Any]] = [
|
||||
("limit", limit),
|
||||
("offset", offset),
|
||||
]
|
||||
if status_in is not None:
|
||||
payload["status_in"] = status_in
|
||||
for value in status_in:
|
||||
params.append(("status_in", value))
|
||||
if worker_id_contains is not None:
|
||||
payload["worker_id_contains"] = worker_id_contains
|
||||
payload["filter_logic"] = filter_logic
|
||||
params.append(("worker_id_contains", worker_id_contains))
|
||||
params.append(("filter_logic", filter_logic))
|
||||
if sort_by is not None:
|
||||
payload["sort_by"] = sort_by
|
||||
payload["sort_order"] = sort_order
|
||||
params.append(("sort_by", sort_by))
|
||||
params.append(("sort_order", sort_order))
|
||||
|
||||
data = await self._request_json("post", "/workers/search", json=payload)
|
||||
data = await self._request_json("get", "/workers", params=params)
|
||||
items = [Worker.model_validate(item) for item in data.get("items", [])]
|
||||
return PaginatedResult(items=items, limit=data["limit"], offset=data["offset"], total=data["total"])
|
||||
|
||||
|
||||
@@ -1,21 +1,9 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .base import (
|
||||
AtomicLabels,
|
||||
AtomicMode,
|
||||
Collection,
|
||||
FilterOptions,
|
||||
KeyValue,
|
||||
LightningCollections,
|
||||
PaginatedResult,
|
||||
Queue,
|
||||
SortOptions,
|
||||
)
|
||||
from .base import Collection, FilterOptions, KeyValue, LightningCollections, PaginatedResult, Queue, SortOptions
|
||||
from .memory import DequeBasedQueue, DictBasedKeyValue, InMemoryLightningCollections, ListBasedCollection
|
||||
|
||||
__all__ = [
|
||||
"AtomicLabels",
|
||||
"AtomicMode",
|
||||
"Collection",
|
||||
"Queue",
|
||||
"KeyValue",
|
||||
|
||||
@@ -2,10 +2,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from numbers import Real
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
@@ -22,14 +18,10 @@ from typing import (
|
||||
Sequence,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeGuard,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from agentlightning.store.utils import LATENCY_BUCKETS
|
||||
from agentlightning.utils.metrics import MetricsBackend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Self
|
||||
|
||||
@@ -48,146 +40,10 @@ from agentlightning.types import (
|
||||
T = TypeVar("T") # Recommended to be a BaseModel
|
||||
K = TypeVar("K")
|
||||
V = TypeVar("V")
|
||||
T_callable = TypeVar("T_callable", bound=Callable[..., Any])
|
||||
|
||||
AtomicMode = Literal["r", "w", "rw"]
|
||||
"""What is expected within the atomic context. Can be "read", "write", or "read-write"."""
|
||||
|
||||
AtomicLabels = Literal[
|
||||
"rollouts", "attempts", "spans", "resources", "workers", "rollout_queue", "span_sequence_ids", "generic"
|
||||
]
|
||||
"""Labels for atomic operations.
|
||||
|
||||
These labels are used to identify the collections that are affected by the atomic operation.
|
||||
|
||||
The `generic` label is used to identify atomic operations that are not associated with any specific collection.
|
||||
"""
|
||||
|
||||
|
||||
def resolve_error_type(exc: BaseException | None) -> str:
|
||||
if exc is None:
|
||||
return "N/A"
|
||||
|
||||
try:
|
||||
from .mongo import resolve_mongo_error_type
|
||||
|
||||
error_type = resolve_mongo_error_type(exc)
|
||||
if error_type is not None:
|
||||
return error_type
|
||||
except ImportError:
|
||||
# If the mongo backend is not available, fall back to using the exception's class name.
|
||||
pass
|
||||
|
||||
return exc.__class__.__name__
|
||||
|
||||
|
||||
def tracked(operation: str):
|
||||
"""Decorator to track the execution of the decorated method."""
|
||||
|
||||
def decorator(func: T_callable) -> T_callable:
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: TrackedCollection, *args: Any, **kwargs: Any) -> Any:
|
||||
async with self.tracking_context(operation, self.collection_name):
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
return cast(T_callable, wrapper)
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def ensure_numeric(value: Any, *, description: str) -> TypeGuard[Real]:
|
||||
"""Validate that *value* behaves like a real number.
|
||||
|
||||
Returns true or crashes.
|
||||
"""
|
||||
|
||||
if isinstance(value, bool):
|
||||
raise TypeError(f"{description} must be numeric; got bool")
|
||||
if not isinstance(value, Real):
|
||||
raise TypeError(f"{description} must be numeric; got {type(value).__name__}")
|
||||
return True
|
||||
|
||||
|
||||
class DuplicatedPrimaryKeyError(ValueError):
|
||||
"""Error raised when a duplicate key is encountered."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class TrackedCollection:
|
||||
"""An object that can be tracked by the metrics backend."""
|
||||
|
||||
def __init__(self, tracker: MetricsBackend | None = None):
|
||||
self._tracker = tracker
|
||||
|
||||
@property
|
||||
def tracker(self) -> MetricsBackend | None:
|
||||
return self._tracker
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
"""The identifier of the collection."""
|
||||
raise NotImplementedError()
|
||||
|
||||
@property
|
||||
def extra_tracking_labels(self) -> Mapping[str, Any]:
|
||||
"""Extra labels to add to the tracking context."""
|
||||
return {}
|
||||
|
||||
@asynccontextmanager
|
||||
async def tracking_context(self, operation: str, collection: str):
|
||||
"""Context manager to track the execution of the decorated method.
|
||||
|
||||
Args:
|
||||
operation: The operation to track.
|
||||
collection: The collection to track.
|
||||
"""
|
||||
if self._tracker is None:
|
||||
# no-op context manager
|
||||
yield
|
||||
|
||||
else:
|
||||
from agentlightning.store.collection_based import get_current_store_methods
|
||||
|
||||
# Enable tracking
|
||||
start_time = time.perf_counter()
|
||||
status: str = "OK"
|
||||
public_store_method, private_store_method = get_current_store_methods()
|
||||
try:
|
||||
yield
|
||||
except BaseException as exc:
|
||||
status = resolve_error_type(exc)
|
||||
raise
|
||||
finally:
|
||||
elapsed = time.perf_counter() - start_time
|
||||
await self._tracker.inc_counter( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.collections.total",
|
||||
labels={
|
||||
"store_pubmeth": public_store_method,
|
||||
"store_privmeth": private_store_method,
|
||||
"operation": operation,
|
||||
"collection": collection,
|
||||
"status": status,
|
||||
**self.extra_tracking_labels,
|
||||
},
|
||||
)
|
||||
await self._tracker.observe_histogram( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.collections.latency",
|
||||
value=elapsed,
|
||||
labels={
|
||||
"store_pubmeth": public_store_method,
|
||||
"store_privmeth": private_store_method,
|
||||
"operation": operation,
|
||||
"collection": collection,
|
||||
"status": status,
|
||||
**self.extra_tracking_labels,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class Collection(TrackedCollection, Generic[T]):
|
||||
"""Standard collection interface. Behaves like a list of items. Supporting addition, updating, and deletion of items."""
|
||||
class Collection(Generic[T]):
|
||||
"""Behaves like a list of items. Supporting addition, updating, and deletion of items."""
|
||||
|
||||
def primary_keys(self) -> Sequence[str]:
|
||||
"""Get the primary keys of the collection."""
|
||||
@@ -258,42 +114,19 @@ class Collection(TrackedCollection, Generic[T]):
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
|
||||
async def update(self, items: Sequence[T]) -> None:
|
||||
"""Update the given items in the collection.
|
||||
|
||||
Args:
|
||||
items: The items to update in the collection.
|
||||
update_fields: The fields to update. If not provided, all fields in the type will be updated.
|
||||
Only applicable if the item type is a Pydantic BaseModel.
|
||||
|
||||
Raises:
|
||||
ValueError: If an item with the primary keys does not exist.
|
||||
|
||||
Returns:
|
||||
The items that were updated.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
|
||||
async def upsert(self, items: Sequence[T]) -> None:
|
||||
"""Upsert the given items into the collection.
|
||||
|
||||
If the items with the same primary keys already exist, they will be updated.
|
||||
Otherwise, they will be inserted.
|
||||
|
||||
The operation has three semantics configurable via `update_fields`:
|
||||
|
||||
- `update_or_insert` via `collection.upsert(items, update_fields=["status", "updated_at"])`.
|
||||
If the item with the same primary keys already exists, only the specified fields will be updated.
|
||||
Otherwise, the item will be inserted.
|
||||
- `get_or_insert` via `collection.upsert(items, update_fields=[])`.
|
||||
If the item with the same primary keys already exists, the item will be left unchanged.
|
||||
Otherwise, the item will be inserted.
|
||||
- `replace_ish` via `collection.upsert(items)`.
|
||||
If the item with the same primary keys already exists, all fields from the item will be set.
|
||||
Otherwise, the item will be inserted.
|
||||
|
||||
Returns:
|
||||
The items that were upserted.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -309,7 +142,7 @@ class Collection(TrackedCollection, Generic[T]):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class Queue(TrackedCollection, Generic[T]):
|
||||
class Queue(Generic[T]):
|
||||
"""Behaves like a deque. Supporting appending items to the end and popping items from the front."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
@@ -363,7 +196,7 @@ class Queue(TrackedCollection, Generic[T]):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class KeyValue(TrackedCollection, Generic[K, V]):
|
||||
class KeyValue(Generic[K, V]):
|
||||
"""Behaves like a dictionary. Supporting addition, updating, and deletion of items."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
@@ -381,22 +214,6 @@ class KeyValue(TrackedCollection, Generic[K, V]):
|
||||
"""Set the value for the given key."""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def inc(self, key: K, amount: V) -> V:
|
||||
"""Increase the numeric value for the given key by `amount` and return the new value.
|
||||
|
||||
Raises:
|
||||
TypeError: If the existing value or `amount` is not numeric.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def chmax(self, key: K, value: V) -> V:
|
||||
"""Set the value for the given key to the maximum of the current and new value.
|
||||
|
||||
Raises:
|
||||
TypeError: If the existing value or `value` is not numeric.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def pop(self, key: K, default: V | None = None) -> V | None:
|
||||
"""Pop the value for the given key, or the default value if the key is not found."""
|
||||
raise NotImplementedError()
|
||||
@@ -406,35 +223,13 @@ class KeyValue(TrackedCollection, Generic[K, V]):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class LightningCollections(TrackedCollection):
|
||||
class LightningCollections:
|
||||
"""Collections of rollouts, attempts, spans, resources, and workers.
|
||||
|
||||
[LightningStore][agentlightning.LightningStore] implementations can use this as a storage base
|
||||
to implement the store API.
|
||||
"""
|
||||
|
||||
def __init__(self, tracker: MetricsBackend | None = None, extra_labels: Optional[Sequence[str]] = None):
|
||||
super().__init__(tracker=tracker)
|
||||
self.register_collection_metrics(extra_labels)
|
||||
|
||||
def register_collection_metrics(self, extra_labels: Optional[Sequence[str]] = None) -> None:
|
||||
if self._tracker is None:
|
||||
return
|
||||
labels = ["store_pubmeth", "operation", "collection", "store_privmeth", "status"]
|
||||
if extra_labels is not None:
|
||||
labels.extend(extra_labels)
|
||||
self._tracker.register_histogram(
|
||||
"agl.collections.latency",
|
||||
labels,
|
||||
buckets=LATENCY_BUCKETS,
|
||||
group_level=2,
|
||||
)
|
||||
self._tracker.register_counter("agl.collections.total", labels, group_level=2)
|
||||
|
||||
@property
|
||||
def tracker(self) -> MetricsBackend | None:
|
||||
return self._tracker
|
||||
|
||||
@property
|
||||
def rollouts(self) -> Collection[Rollout]:
|
||||
"""Collections of rollouts."""
|
||||
@@ -470,46 +265,20 @@ class LightningCollections(TrackedCollection):
|
||||
"""Dictionary (counter) of span sequence IDs."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def atomic(
|
||||
self,
|
||||
*,
|
||||
mode: AtomicMode = "rw",
|
||||
snapshot: bool = False,
|
||||
commit: bool = False,
|
||||
labels: Optional[Sequence[AtomicLabels]] = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncContextManager[Self]:
|
||||
def atomic(self, *args: Any, **kwargs: Any) -> AsyncContextManager[Self]:
|
||||
"""Perform a atomic operation on the collections.
|
||||
|
||||
Subclass may use args and kwargs to support multiple levels of atomicity.
|
||||
The arguments can be seen as tags. They only imply the behavior of the operation, not the implementation.
|
||||
|
||||
Args:
|
||||
mode: The mode of atomicity. See [`AtomicMode`][agentlightning.store.collection.AtomicMode].
|
||||
snapshot: Enable read snapshot for repeatable reads. Data consistency is guaranteed. The real behavior is implementation-dependent.
|
||||
commit: Enable commitment for write operations. Unsuccessful operations will be rolled back depending on the implementation.
|
||||
Recommend to use [`execute()`][agentlightning.store.collection.LightningCollections.execute] for this level to enable automatic retries.
|
||||
Remember that the real behavior is implementation-dependent.
|
||||
labels: Labels to add to the atomic operation (commonly used as lock names or collection names).
|
||||
*args: Arguments to pass to the operation.
|
||||
**kwargs: Keyword arguments to pass to the operation.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
callback: Callable[[Self], Awaitable[T]],
|
||||
*,
|
||||
mode: AtomicMode = "rw",
|
||||
snapshot: bool = False,
|
||||
commit: bool = False,
|
||||
labels: Optional[Sequence[AtomicLabels]] = None,
|
||||
**kwargs: Any,
|
||||
) -> T:
|
||||
"""Execute the given callback within an atomic operation. Retry on transient errors is implied.
|
||||
|
||||
See [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] for more details.
|
||||
"""
|
||||
async with self.atomic(mode=mode, snapshot=snapshot, commit=commit, labels=labels, **kwargs) as collections:
|
||||
async def execute(self, callback: Callable[[Self], Awaitable[T]]) -> T:
|
||||
"""Execute the given callback within an atomic operation."""
|
||||
async with self.atomic() as collections:
|
||||
return await callback(collections)
|
||||
|
||||
|
||||
|
||||
@@ -4,10 +4,9 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
import weakref
|
||||
from collections import deque
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import (
|
||||
Any,
|
||||
Deque,
|
||||
@@ -23,12 +22,8 @@ from typing import (
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import aiologic
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agentlightning.types import (
|
||||
Attempt,
|
||||
FilterField,
|
||||
@@ -40,21 +35,15 @@ from agentlightning.types import (
|
||||
Span,
|
||||
Worker,
|
||||
)
|
||||
from agentlightning.utils.metrics import MetricsBackend
|
||||
|
||||
from .base import (
|
||||
AtomicLabels,
|
||||
AtomicMode,
|
||||
Collection,
|
||||
DuplicatedPrimaryKeyError,
|
||||
FilterMap,
|
||||
KeyValue,
|
||||
LightningCollections,
|
||||
Queue,
|
||||
ensure_numeric,
|
||||
normalize_filter_options,
|
||||
resolve_sort_options,
|
||||
tracked,
|
||||
)
|
||||
|
||||
T = TypeVar("T") # Recommended to be a BaseModel, not a dict
|
||||
@@ -197,19 +186,10 @@ class ListBasedCollection(Collection[T]):
|
||||
if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
items: List[T],
|
||||
item_type: Type[T],
|
||||
primary_keys: Sequence[str],
|
||||
id: Optional[str] = None,
|
||||
tracker: Optional[MetricsBackend] = None,
|
||||
):
|
||||
super().__init__(tracker=tracker)
|
||||
def __init__(self, items: List[T], item_type: Type[T], primary_keys: Sequence[str]):
|
||||
if not primary_keys:
|
||||
raise ValueError("primary_keys must be non-empty")
|
||||
|
||||
self._id = id if id is not None else str(uuid.uuid4())
|
||||
self._items: Dict[Any, Any] = {}
|
||||
self._size: int = 0
|
||||
if issubclass(item_type, dict):
|
||||
@@ -221,10 +201,6 @@ class ListBasedCollection(Collection[T]):
|
||||
for item in items or []:
|
||||
self._mutate_single(item, mode="insert")
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
return self._id
|
||||
|
||||
def primary_keys(self) -> Sequence[str]:
|
||||
"""Return the primary key field names for this collection."""
|
||||
return self._primary_keys
|
||||
@@ -306,7 +282,7 @@ class ListBasedCollection(Collection[T]):
|
||||
# We should always return inside the loop.
|
||||
raise RuntimeError("Unreachable")
|
||||
|
||||
def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
|
||||
def _mutate_single(self, item: T, mode: MutationMode) -> None:
|
||||
"""Core mutation logic shared by insert, update, upsert, and delete."""
|
||||
self._ensure_item_type(item)
|
||||
key_values = self._extract_primary_key_values(item)
|
||||
@@ -317,43 +293,13 @@ class ListBasedCollection(Collection[T]):
|
||||
|
||||
if mode == "insert":
|
||||
if exists:
|
||||
raise DuplicatedPrimaryKeyError(
|
||||
f"Item already exists with primary key(s): {self._render_key_values(key_values)}"
|
||||
)
|
||||
raise ValueError(f"Item already exists with primary key(s): {self._render_key_values(key_values)}")
|
||||
parent[final_key] = item
|
||||
self._size += 1
|
||||
else: # upsert
|
||||
if not exists:
|
||||
self._size += 1
|
||||
parent[final_key] = item
|
||||
|
||||
elif update_fields is None:
|
||||
# update_or_insert: update all fields
|
||||
parent[final_key] = item
|
||||
|
||||
else:
|
||||
if not issubclass(self._item_type, BaseModel):
|
||||
raise TypeError(
|
||||
f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
|
||||
)
|
||||
|
||||
# Try to fetch the existing item
|
||||
existing = parent[final_key]
|
||||
if not isinstance(existing, self._item_type):
|
||||
raise ValueError(
|
||||
f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
|
||||
)
|
||||
|
||||
if not isinstance(item, self._item_type):
|
||||
raise TypeError(
|
||||
f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
|
||||
)
|
||||
|
||||
parent[final_key] = parent[final_key].model_copy(
|
||||
update={field: getattr(item, field) for field in update_fields}
|
||||
)
|
||||
|
||||
return parent[final_key]
|
||||
parent[final_key] = item
|
||||
|
||||
elif mode in ("update", "delete"):
|
||||
# For update/delete we must not create missing paths.
|
||||
@@ -368,22 +314,7 @@ class ListBasedCollection(Collection[T]):
|
||||
raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
|
||||
|
||||
if mode == "update":
|
||||
if update_fields is None:
|
||||
# replace the entire item
|
||||
parent[final_key] = item
|
||||
else:
|
||||
if not issubclass(self._item_type, BaseModel):
|
||||
raise TypeError(
|
||||
f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
|
||||
)
|
||||
if not isinstance(item, self._item_type):
|
||||
raise TypeError(
|
||||
f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
|
||||
)
|
||||
parent[final_key] = parent[final_key].model_copy(
|
||||
update={field: getattr(item, field) for field in update_fields}
|
||||
)
|
||||
return parent[final_key]
|
||||
parent[final_key] = item
|
||||
else: # delete
|
||||
del parent[final_key]
|
||||
self._size -= 1
|
||||
@@ -503,7 +434,6 @@ class ListBasedCollection(Collection[T]):
|
||||
# No items exist for this primary-key prefix.
|
||||
return ()
|
||||
|
||||
@tracked("query")
|
||||
async def query(
|
||||
self,
|
||||
filter: Optional[FilterOptions] = None,
|
||||
@@ -567,7 +497,6 @@ class ListBasedCollection(Collection[T]):
|
||||
total=total_matched,
|
||||
)
|
||||
|
||||
@tracked("get")
|
||||
async def get(
|
||||
self,
|
||||
filter: Optional[FilterOptions] = None,
|
||||
@@ -604,12 +533,11 @@ class ListBasedCollection(Collection[T]):
|
||||
|
||||
return best_item
|
||||
|
||||
@tracked("insert")
|
||||
async def insert(self, items: Sequence[T]) -> None:
|
||||
"""Insert the given items.
|
||||
|
||||
Raises:
|
||||
DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
|
||||
ValueError: If any item with the same primary keys already exists.
|
||||
"""
|
||||
seen_keys: set[Tuple[Any, ...]] = set()
|
||||
prepared: List[T] = []
|
||||
@@ -617,8 +545,8 @@ class ListBasedCollection(Collection[T]):
|
||||
self._ensure_item_type(item)
|
||||
key_values = self._extract_primary_key_values(item)
|
||||
if key_values in seen_keys:
|
||||
raise DuplicatedPrimaryKeyError(
|
||||
f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
|
||||
raise ValueError(
|
||||
f"Insert payload contains duplicate primary key(s): {self._render_key_values(key_values)}"
|
||||
)
|
||||
seen_keys.add(key_values)
|
||||
prepared.append(item)
|
||||
@@ -626,33 +554,20 @@ class ListBasedCollection(Collection[T]):
|
||||
for item in prepared:
|
||||
self._mutate_single(item, mode="insert")
|
||||
|
||||
@tracked("update")
|
||||
async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
|
||||
async def update(self, items: Sequence[T]) -> None:
|
||||
"""Update the given items.
|
||||
|
||||
Raises:
|
||||
ValueError: If any item with the given primary keys does not exist.
|
||||
"""
|
||||
updated_items: List[T] = []
|
||||
for item in items:
|
||||
updated = self._mutate_single(item, mode="update", update_fields=update_fields)
|
||||
if updated is None:
|
||||
raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
|
||||
updated_items.append(updated)
|
||||
return updated_items
|
||||
self._mutate_single(item, mode="update")
|
||||
|
||||
@tracked("upsert")
|
||||
async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
|
||||
async def upsert(self, items: Sequence[T]) -> None:
|
||||
"""Upsert the given items (insert if missing, otherwise update)."""
|
||||
upserted_items: List[T] = []
|
||||
for item in items:
|
||||
upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
|
||||
if upserted is None:
|
||||
raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
|
||||
upserted_items.append(upserted)
|
||||
return upserted_items
|
||||
self._mutate_single(item, mode="upsert")
|
||||
|
||||
@tracked("delete")
|
||||
async def delete(self, items: Sequence[T]) -> None:
|
||||
"""Delete the given items.
|
||||
|
||||
@@ -672,37 +587,23 @@ class DequeBasedQueue(Queue[T]):
|
||||
Provides O(1) amortized enqueue (append) and dequeue (popleft).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
item_type: Type[T],
|
||||
items: Optional[Sequence[T]] = None,
|
||||
id: Optional[str] = None,
|
||||
tracker: Optional[MetricsBackend] = None,
|
||||
):
|
||||
super().__init__(tracker=tracker)
|
||||
def __init__(self, item_type: Type[T], items: Optional[Sequence[T]] = None):
|
||||
self._items: Deque[T] = deque()
|
||||
self._item_type: Type[T] = item_type
|
||||
self._id = id if id is not None else str(uuid.uuid4())
|
||||
if items:
|
||||
self._items.extend(items)
|
||||
|
||||
def item_type(self) -> Type[T]:
|
||||
return self._item_type
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
return self._id
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"
|
||||
|
||||
@tracked("has")
|
||||
async def has(self, item: T) -> bool:
|
||||
if not isinstance(item, self._item_type):
|
||||
raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
|
||||
return item in self._items
|
||||
|
||||
@tracked("enqueue")
|
||||
async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
|
||||
for item in items:
|
||||
if not isinstance(item, self._item_type):
|
||||
@@ -710,7 +611,6 @@ class DequeBasedQueue(Queue[T]):
|
||||
self._items.append(item)
|
||||
return items
|
||||
|
||||
@tracked("dequeue")
|
||||
async def dequeue(self, limit: int = 1) -> Sequence[T]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
@@ -719,7 +619,6 @@ class DequeBasedQueue(Queue[T]):
|
||||
out.append(self._items.popleft())
|
||||
return out
|
||||
|
||||
@tracked("peek")
|
||||
async def peek(self, limit: int = 1) -> Sequence[T]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
@@ -731,7 +630,6 @@ class DequeBasedQueue(Queue[T]):
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
@tracked("size")
|
||||
async def size(self) -> int:
|
||||
return len(self._items)
|
||||
|
||||
@@ -739,61 +637,21 @@ class DequeBasedQueue(Queue[T]):
|
||||
class DictBasedKeyValue(KeyValue[K, V]):
|
||||
"""KeyValue implementation backed by a plain dictionary."""
|
||||
|
||||
def __init__(
|
||||
self, data: Optional[Mapping[K, V]] = None, id: Optional[str] = None, tracker: Optional[MetricsBackend] = None
|
||||
):
|
||||
super().__init__(tracker=tracker)
|
||||
def __init__(self, data: Optional[Mapping[K, V]] = None):
|
||||
self._values: Dict[K, V] = dict(data) if data else {}
|
||||
self._id = id if id is not None else str(uuid.uuid4())
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
return self._id
|
||||
|
||||
@tracked("has")
|
||||
async def has(self, key: K) -> bool:
|
||||
return key in self._values
|
||||
|
||||
@tracked("get")
|
||||
async def get(self, key: K, default: V | None = None) -> V | None:
|
||||
return self._values.get(key, default)
|
||||
|
||||
@tracked("set")
|
||||
async def set(self, key: K, value: V) -> None:
|
||||
self._values[key] = value
|
||||
|
||||
@tracked("inc")
|
||||
async def inc(self, key: K, amount: V) -> V:
|
||||
assert ensure_numeric(amount, description="amount")
|
||||
if key in self._values:
|
||||
current_value = self._values[key]
|
||||
assert ensure_numeric(current_value, description=f"value for key {key!r}")
|
||||
new_value = cast(V, current_value + amount)
|
||||
self._values[key] = new_value
|
||||
else:
|
||||
new_value = amount
|
||||
self._values[key] = new_value
|
||||
return new_value
|
||||
|
||||
@tracked("chmax")
|
||||
async def chmax(self, key: K, value: V) -> V:
|
||||
assert ensure_numeric(value, description="value")
|
||||
if key in self._values:
|
||||
current_value = self._values[key]
|
||||
assert ensure_numeric(current_value, description=f"value for key {key!r}")
|
||||
if value > current_value:
|
||||
self._values[key] = value
|
||||
return value
|
||||
return current_value
|
||||
else:
|
||||
self._values[key] = value
|
||||
return value
|
||||
|
||||
@tracked("pop")
|
||||
async def pop(self, key: K, default: V | None = None) -> V | None:
|
||||
return self._values.pop(key, default)
|
||||
|
||||
@tracked("size")
|
||||
async def size(self) -> int:
|
||||
return len(self._values)
|
||||
|
||||
@@ -804,41 +662,17 @@ class InMemoryLightningCollections(LightningCollections):
|
||||
Serves as the storage base for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
|
||||
"""
|
||||
|
||||
def __init__(self, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
|
||||
super().__init__(tracker=tracker)
|
||||
self._lock: Mapping[AtomicLabels, _LoopAwareAsyncLock | _ThreadSafeAsyncLock] = {
|
||||
"rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
"generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
|
||||
}
|
||||
self._rollouts = ListBasedCollection(
|
||||
items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
|
||||
)
|
||||
self._attempts = ListBasedCollection(
|
||||
items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"], id="attempts", tracker=tracker
|
||||
)
|
||||
def __init__(self):
|
||||
self._lock = _LoopAwareAsyncLock()
|
||||
self._rollouts = ListBasedCollection(items=[], item_type=Rollout, primary_keys=["rollout_id"])
|
||||
self._attempts = ListBasedCollection(items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"])
|
||||
self._spans = ListBasedCollection(
|
||||
items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"], id="spans", tracker=tracker
|
||||
items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"]
|
||||
)
|
||||
self._resources = ListBasedCollection(
|
||||
items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"], id="resources", tracker=tracker
|
||||
)
|
||||
self._workers = ListBasedCollection(
|
||||
items=[], item_type=Worker, primary_keys=["worker_id"], id="workers", tracker=tracker
|
||||
)
|
||||
self._rollout_queue = DequeBasedQueue(items=[], item_type=str, id="rollout_queue", tracker=tracker)
|
||||
self._span_sequence_ids = DictBasedKeyValue[str, int](
|
||||
data={}, id="span_sequence_ids", tracker=tracker
|
||||
) # rollout_id -> sequence_id
|
||||
|
||||
@property
|
||||
def collection_name(self) -> str:
|
||||
return "router"
|
||||
self._resources = ListBasedCollection(items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"])
|
||||
self._workers = ListBasedCollection(items=[], item_type=Worker, primary_keys=["worker_id"])
|
||||
self._rollout_queue = DequeBasedQueue(items=[], item_type=str)
|
||||
self._span_sequence_ids = DictBasedKeyValue[str, int](data={}) # rollout_id -> sequence_id
|
||||
|
||||
@property
|
||||
def rollouts(self) -> ListBasedCollection[Rollout]:
|
||||
@@ -869,41 +703,11 @@ class InMemoryLightningCollections(LightningCollections):
|
||||
return self._span_sequence_ids
|
||||
|
||||
@asynccontextmanager
|
||||
async def atomic(
|
||||
self,
|
||||
*,
|
||||
mode: AtomicMode = "rw",
|
||||
snapshot: bool = False,
|
||||
labels: Optional[Sequence[AtomicLabels]] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
|
||||
|
||||
Skip the locking if mode is "r" and snapshot is False.
|
||||
|
||||
This collection implementation does NOT support rollback / commit.
|
||||
"""
|
||||
if mode == "r" and not snapshot:
|
||||
async def atomic(self, *args: Any, **kwargs: Any):
|
||||
"""In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside."""
|
||||
async with self._lock:
|
||||
yield self
|
||||
return
|
||||
if not labels:
|
||||
# If no labels are provided, use all locks.
|
||||
labels = list(self._lock.keys())
|
||||
|
||||
# IMPORTANT: Sort the labels to ensure consistent locking order.
|
||||
# This is necessary to avoid deadlocks when multiple threads/coroutines
|
||||
# are trying to acquire the same locks in different orders.
|
||||
labels = sorted(labels)
|
||||
|
||||
async with self.tracking_context(operation="atomic", collection=self.collection_name):
|
||||
managers = [(label, self._lock[label]) for label in labels]
|
||||
async with AsyncExitStack() as stack:
|
||||
for label, manager in managers:
|
||||
async with self.tracking_context(operation="lock", collection=label):
|
||||
await stack.enter_async_context(manager)
|
||||
yield self
|
||||
|
||||
@tracked("evict_spans_for_rollout")
|
||||
async def evict_spans_for_rollout(self, rollout_id: str) -> None:
|
||||
"""Evict all spans for a given rollout ID.
|
||||
|
||||
@@ -950,21 +754,3 @@ class _LoopAwareAsyncLock:
|
||||
if lock is None or not lock.locked():
|
||||
raise RuntimeError("Lock released without being acquired")
|
||||
lock.release()
|
||||
|
||||
|
||||
class _ThreadSafeAsyncLock:
|
||||
"""A thread lock powered by aiologic that can be used in both async and sync contexts.
|
||||
|
||||
aiologic claims itself to be a thread-safe asyncio lock.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = aiologic.Lock()
|
||||
|
||||
async def __aenter__(self):
|
||||
await self._lock.async_acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any, **kwargs: Any):
|
||||
# .release() is non-blocking, so we can call it directly
|
||||
self._lock.async_release()
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+56
-104
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Mapping as MappingABC
|
||||
from typing import (
|
||||
@@ -18,17 +19,14 @@ from typing import (
|
||||
Optional,
|
||||
Sequence,
|
||||
Set,
|
||||
Tuple,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import aiologic
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agentlightning.types import AttemptedRollout, NamedResources, PaginatedResult, ResourcesUpdate, Rollout, Span
|
||||
from agentlightning.utils.metrics import MetricsBackend
|
||||
from agentlightning.types import AttemptedRollout, PaginatedResult, Rollout, Span
|
||||
|
||||
from .base import UNSET, LightningStoreCapabilities, LightningStoreStatistics, Unset, is_finished, is_running
|
||||
from .collection import InMemoryLightningCollections
|
||||
@@ -73,35 +71,24 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
Thread-safe and async-compatible but data is not persistent.
|
||||
|
||||
Args:
|
||||
thread_safe: Whether the store is thread-safe.
|
||||
eviction_memory_threshold: The threshold for evicting spans in bytes.
|
||||
By default, it's 70% of the total VRAM available.
|
||||
safe_memory_threshold: The threshold for safe memory usage in bytes.
|
||||
By default, it's 80% of the eviction threshold.
|
||||
span_size_estimator: A function to estimate the size of a span in bytes.
|
||||
By default, it's a simple size estimator that uses sys.getsizeof.
|
||||
tracker: The metrics tracker to use.
|
||||
scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
|
||||
Set to 0 to disable debouncing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
thread_safe: bool = False,
|
||||
eviction_memory_threshold: float | int | None = None,
|
||||
safe_memory_threshold: float | int | None = None,
|
||||
span_size_estimator: Callable[[Span], int] | None = None,
|
||||
tracker: MetricsBackend | None = None,
|
||||
scan_debounce_seconds: float = 10.0,
|
||||
prometheus: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
collections=InMemoryLightningCollections(lock_type="thread" if thread_safe else "asyncio", tracker=tracker),
|
||||
tracker=tracker,
|
||||
scan_debounce_seconds=scan_debounce_seconds,
|
||||
)
|
||||
super().__init__(collections=InMemoryLightningCollections(), prometheus=prometheus)
|
||||
|
||||
self._thread_safe = thread_safe
|
||||
self._start_time_by_rollout: Dict[str, float] = {}
|
||||
self._span_bytes_by_rollout: Dict[str, int] = Counter()
|
||||
self._total_span_bytes: int = 0
|
||||
@@ -135,7 +122,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
self._custom_span_size_estimator = span_size_estimator
|
||||
|
||||
# Completion tracking for wait_for_rollouts (cross-loop safe)
|
||||
self._completion_events: Dict[str, aiologic.Event] = {}
|
||||
self._completion_events: Dict[str, threading.Event] = {}
|
||||
|
||||
# Running rollouts cache, including preparing and running rollouts
|
||||
self._running_rollout_ids: Set[str] = set()
|
||||
@@ -147,7 +134,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
def capabilities(self) -> LightningStoreCapabilities:
|
||||
"""Return the capabilities of the store."""
|
||||
return LightningStoreCapabilities(
|
||||
thread_safe=self._thread_safe,
|
||||
thread_safe=False,
|
||||
async_safe=True,
|
||||
zero_copy=False,
|
||||
otlp_traces=False,
|
||||
@@ -166,7 +153,7 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
@tracked("wait_for_rollout")
|
||||
async def wait_for_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
|
||||
"""Wait for a specific rollout to complete with a timeout."""
|
||||
async with self.collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["rollouts"]) as collections:
|
||||
async with self.collections.atomic() as collections:
|
||||
rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
|
||||
if rollout and is_finished(rollout):
|
||||
return rollout
|
||||
@@ -194,82 +181,47 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
|
||||
# If event was set (not timeout), check if rollout is finished
|
||||
if result:
|
||||
async with self.collections.atomic(
|
||||
mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
|
||||
) as collections:
|
||||
async with self.collections.atomic() as collections:
|
||||
rollout = await collections.rollouts.get({"rollout_id": {"exact": rollout_id}})
|
||||
if rollout and is_finished(rollout):
|
||||
return rollout
|
||||
|
||||
return None
|
||||
|
||||
@tracked("add_resources_inmemory")
|
||||
async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
|
||||
ret = await super().add_resources(resources)
|
||||
async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
|
||||
self._latest_resources_id = ret.resources_id
|
||||
return ret
|
||||
|
||||
@tracked("update_resources_inmemory")
|
||||
async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
|
||||
ret = await super().update_resources(resources_id, resources)
|
||||
async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["resources"]):
|
||||
self._latest_resources_id = ret.resources_id
|
||||
return ret
|
||||
|
||||
@tracked("_post_update_rollout_inmemory")
|
||||
async def _post_update_rollout(
|
||||
self, rollouts: Sequence[Tuple[Rollout, Sequence[str]]], skip_enqueue: bool = False
|
||||
) -> None:
|
||||
@tracked("on_rollout_update")
|
||||
async def on_rollout_update(self, rollout: Rollout) -> None:
|
||||
"""Update the running rollout ids set when the rollout updates."""
|
||||
await super()._post_update_rollout(rollouts, skip_enqueue=skip_enqueue)
|
||||
async with self.collections.atomic(mode="rw", snapshot=self._read_snapshot, labels=["rollouts"]):
|
||||
for rollout, _ in rollouts:
|
||||
if is_running(rollout):
|
||||
self._running_rollout_ids.add(rollout.rollout_id)
|
||||
else:
|
||||
self._running_rollout_ids.discard(rollout.rollout_id)
|
||||
if is_running(rollout):
|
||||
self._running_rollout_ids.add(rollout.rollout_id)
|
||||
else:
|
||||
self._running_rollout_ids.discard(rollout.rollout_id)
|
||||
|
||||
if is_finished(rollout):
|
||||
self._completion_events.setdefault(rollout.rollout_id, aiologic.Event())
|
||||
self._completion_events[rollout.rollout_id].set()
|
||||
else:
|
||||
self._completion_events.setdefault(rollout.rollout_id, aiologic.Event())
|
||||
# Rollout status can never transition from finished to running (unlike attempt)
|
||||
# so we don't need to clear the completion event even in case of retrying.
|
||||
if is_finished(rollout):
|
||||
self._completion_events.setdefault(rollout.rollout_id, threading.Event())
|
||||
self._completion_events[rollout.rollout_id].set()
|
||||
else:
|
||||
self._completion_events.setdefault(rollout.rollout_id, threading.Event())
|
||||
# Rollout status can never transition from finished to running (unlike attempt)
|
||||
# so we don't need to clear the completion event even in case of retrying.
|
||||
|
||||
if rollout.rollout_id not in self._start_time_by_rollout:
|
||||
self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
|
||||
if rollout.rollout_id not in self._start_time_by_rollout:
|
||||
self._start_time_by_rollout[rollout.rollout_id] = rollout.start_time
|
||||
|
||||
@tracked("_unlocked_query_rollouts_by_rollout_ids")
|
||||
async def _unlocked_query_rollouts_by_rollout_ids(
|
||||
self, collections: InMemoryLightningCollections, rollout_ids: Sequence[str]
|
||||
) -> List[Rollout]:
|
||||
"""Always use exact. This is faster than within filter for in-memory store."""
|
||||
if len(rollout_ids) == 0:
|
||||
return []
|
||||
|
||||
rollouts = [await collections.rollouts.get({"rollout_id": {"exact": rollout_id}}) for rollout_id in rollout_ids]
|
||||
return [rollout for rollout in rollouts if rollout is not None]
|
||||
|
||||
@tracked("_unlocked_get_running_rollouts")
|
||||
async def _unlocked_get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
|
||||
"""Accelerated version of `_unlocked_get_running_rollouts` for in-memory store. Used for healthcheck."""
|
||||
async with self.collections.atomic(
|
||||
mode="r", snapshot=self._read_snapshot, labels=["rollouts", "attempts"]
|
||||
) as collections:
|
||||
rollouts = await self._unlocked_query_rollouts_by_rollout_ids(collections, list(self._running_rollout_ids))
|
||||
running_rollouts: List[AttemptedRollout] = []
|
||||
for rollout in rollouts:
|
||||
latest_attempt = await collections.attempts.get(
|
||||
filter={"rollout_id": {"exact": rollout.rollout_id}},
|
||||
sort={"name": "sequence_id", "order": "desc"},
|
||||
)
|
||||
if not latest_attempt:
|
||||
# The rollout is running but has no attempts, this should not happen
|
||||
logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
|
||||
continue
|
||||
running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
|
||||
@tracked("get_running_rollouts")
|
||||
async def get_running_rollouts(self, collections: InMemoryLightningCollections) -> List[AttemptedRollout]:
|
||||
"""Accelerated version of `get_running_rollouts` for in-memory store. Used for healthcheck."""
|
||||
rollouts = await collections.rollouts.query(filter={"rollout_id": {"within": list(self._running_rollout_ids)}})
|
||||
running_rollouts: List[AttemptedRollout] = []
|
||||
for rollout in rollouts.items:
|
||||
latest_attempt = await collections.attempts.get(
|
||||
filter={"rollout_id": {"exact": rollout.rollout_id}},
|
||||
sort={"name": "sequence_id", "order": "desc"},
|
||||
)
|
||||
if not latest_attempt:
|
||||
# The rollout is running but has no attempts, this should not happen
|
||||
logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
|
||||
continue
|
||||
running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
|
||||
return running_rollouts
|
||||
|
||||
@tracked("query_spans_inmemory") # Since this method calls super, we need to track it separately
|
||||
@@ -283,28 +235,28 @@ class InMemoryLightningStore(CollectionBasedLightningStore[InMemoryLightningColl
|
||||
raise RuntimeError(f"Spans for rollout {rollout_id} have been evicted")
|
||||
return await super().query_spans(rollout_id, attempt_id, **kwargs)
|
||||
|
||||
@tracked("_post_add_spans")
|
||||
async def _post_add_spans(self, spans: Sequence[Span], rollout_id: str, attempt_id: str) -> None:
|
||||
@tracked("_add_many_spans_unlocked_inmemory")
|
||||
async def _add_many_spans_unlocked(
|
||||
self, collections: InMemoryLightningCollections, rollout_id: str, attempt_id: str, spans: Sequence[Span]
|
||||
) -> Sequence[Span]:
|
||||
"""In-memory store needs to maintain the span data in memory, and evict spans when memory is low."""
|
||||
|
||||
await super()._post_add_spans(spans, rollout_id, attempt_id)
|
||||
async with self.collections.atomic(
|
||||
mode="rw", snapshot=self._read_snapshot, labels=["rollouts", "spans"]
|
||||
) as collections:
|
||||
for span in spans:
|
||||
await self._account_span_size(span)
|
||||
await self._maybe_evict_spans(collections)
|
||||
inserted = await super()._add_many_spans_unlocked(collections, rollout_id, attempt_id, spans)
|
||||
for span in inserted:
|
||||
await self._account_span_size(span)
|
||||
await self._maybe_evict_spans(collections)
|
||||
|
||||
@tracked("_get_latest_resources_inmemory")
|
||||
async def _get_latest_resources(self) -> Optional[ResourcesUpdate]:
|
||||
return inserted
|
||||
|
||||
@tracked("_get_latest_resources_id")
|
||||
async def _get_latest_resources_id(self, collections: InMemoryLightningCollections) -> Optional[str]:
|
||||
if isinstance(self._latest_resources_id, Unset):
|
||||
return await super()._get_latest_resources()
|
||||
if self._latest_resources_id is not None:
|
||||
async with self.collections.atomic(
|
||||
mode="r", snapshot=self._read_snapshot, labels=["resources"]
|
||||
) as collections:
|
||||
return await collections.resources.get(filter={"resources_id": {"exact": self._latest_resources_id}})
|
||||
return None
|
||||
latest_resources = await collections.resources.get(sort={"name": "update_time", "order": "desc"})
|
||||
if latest_resources:
|
||||
self._latest_resources_id = latest_resources.resources_id
|
||||
else:
|
||||
self._latest_resources_id = None
|
||||
return self._latest_resources_id
|
||||
|
||||
@staticmethod
|
||||
def _resolve_memory_threshold(
|
||||
|
||||
@@ -7,13 +7,24 @@ import hashlib
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, TypeVar, Union
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pymongo import AsyncMongoClient
|
||||
|
||||
from agentlightning.types import Attempt, AttemptedRollout, Rollout
|
||||
from agentlightning.utils.metrics import MetricsBackend
|
||||
|
||||
from .base import LightningStoreCapabilities, is_finished
|
||||
from .collection.mongo import MongoClientPool, MongoLightningCollections
|
||||
from .collection.mongo import MongoClientPool, MongoLightningCollections, MongoOperationPrometheusTracker
|
||||
from .collection_based import CollectionBasedLightningStore, healthcheck_before, tracked
|
||||
|
||||
T_callable = TypeVar("T_callable", bound=Callable[..., Any])
|
||||
@@ -31,28 +42,27 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
Data is persistent and can be shared between multiple processes.
|
||||
|
||||
Args:
|
||||
mongo_uri: MongoDB connection string (defaults to local replica set).
|
||||
mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
|
||||
database_name: The MongoDB database name. Defaults to ``agentlightning``.
|
||||
client: The MongoDB client. Could be a string URI or an instance of AsyncMongoClient.
|
||||
database: The MongoDB database. Could be a string name or an instance of AsyncDatabase.
|
||||
You must provide at least one of client or database.
|
||||
partition_id: The partition id. Useful when sharing the database among multiple Agent-lightning trainers.
|
||||
tracker: The metrics tracker to use.
|
||||
scan_debounce_seconds: The debounce time for the scan for unhealthy rollouts.
|
||||
Set to 0 to disable debouncing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
mongo_uri: str = "mongodb://localhost:27017/?replicaSet=rs0",
|
||||
mongo_client_kwargs: Mapping[str, Any] | None = None,
|
||||
client: AsyncMongoClient[Mapping[str, Any]] | str,
|
||||
database_name: str | None = None,
|
||||
partition_id: str | None = None,
|
||||
tracker: MetricsBackend | None = None,
|
||||
scan_debounce_seconds: float = 10.0,
|
||||
prometheus: bool = False,
|
||||
) -> None:
|
||||
self._mongo_uri = mongo_uri
|
||||
self._mongo_client_kwargs = dict(mongo_client_kwargs or {})
|
||||
|
||||
self._enable_prometheus = prometheus
|
||||
self._auto_created_client = False
|
||||
if isinstance(client, str):
|
||||
self._client = AsyncMongoClient[Mapping[str, Any]](client)
|
||||
self._auto_created_client = True
|
||||
else:
|
||||
self._client = client
|
||||
if database_name is None:
|
||||
database_name = "agentlightning"
|
||||
logger.info("No database name provided, using default 'agentlightning'")
|
||||
@@ -61,20 +71,16 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
partition_id = _generate_partition_id()
|
||||
logger.info("No partition id provided, generated a new one: %s", partition_id)
|
||||
|
||||
self._client_pool = MongoClientPool[Mapping[str, Any]](
|
||||
mongo_uri=self._mongo_uri,
|
||||
mongo_client_kwargs=self._mongo_client_kwargs,
|
||||
)
|
||||
self._client_pool = MongoClientPool(self._client)
|
||||
|
||||
super().__init__(
|
||||
collections=MongoLightningCollections(
|
||||
self._client_pool,
|
||||
database_name,
|
||||
partition_id,
|
||||
tracker=tracker,
|
||||
prometheus_tracker=MongoOperationPrometheusTracker(enabled=self._enable_prometheus),
|
||||
),
|
||||
tracker=tracker,
|
||||
scan_debounce_seconds=scan_debounce_seconds,
|
||||
prometheus=self._enable_prometheus,
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -90,6 +96,9 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
async def close(self) -> None:
|
||||
"""Close the store by closing the client pool."""
|
||||
await self._client_pool.close()
|
||||
# If I created the client, I should close it too.
|
||||
if self._auto_created_client:
|
||||
await self._client.close()
|
||||
|
||||
@tracked("wait_for_rollouts")
|
||||
@healthcheck_before
|
||||
@@ -106,13 +115,10 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
unfinished_rollout_ids = set(rollout_ids)
|
||||
|
||||
while deadline is None or current_time <= deadline:
|
||||
async with self.collections.atomic(
|
||||
mode="r", snapshot=self._read_snapshot, labels=["rollouts"]
|
||||
) as collections:
|
||||
# Query the rollouts that are not finished in a single query
|
||||
rollouts = await collections.rollouts.query(
|
||||
filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
|
||||
)
|
||||
# Query the rollouts that are not finished in a single query
|
||||
rollouts = await self.collections.rollouts.query(
|
||||
filter={"rollout_id": {"within": list(unfinished_rollout_ids)}}
|
||||
)
|
||||
for rollout in rollouts.items:
|
||||
if is_finished(rollout):
|
||||
finished_rollouts[rollout.rollout_id] = rollout
|
||||
@@ -127,28 +133,18 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
await asyncio.sleep(rest_time)
|
||||
current_time = time.time()
|
||||
|
||||
# Logging will help debugging when there are stuck rollouts.
|
||||
logger.debug(
|
||||
"Waiting for rollouts. Number of finished rollouts: %d; number of unfinished rollouts: %d",
|
||||
len(finished_rollouts),
|
||||
len(unfinished_rollout_ids),
|
||||
)
|
||||
if len(unfinished_rollout_ids) < 30:
|
||||
logger.debug("Unfinished rollouts: %s", unfinished_rollout_ids)
|
||||
|
||||
# Reorder the rollouts to match the input order
|
||||
return [finished_rollouts[rollout_id] for rollout_id in rollout_ids if rollout_id in finished_rollouts]
|
||||
|
||||
@tracked("_unlocked_many_rollouts_to_attempted_rollouts")
|
||||
async def _unlocked_many_rollouts_to_attempted_rollouts(
|
||||
@tracked("_many_rollouts_to_attempted_rollouts_unlocked")
|
||||
async def _many_rollouts_to_attempted_rollouts_unlocked(
|
||||
self, collections: MongoLightningCollections, rollouts: Sequence[Rollout]
|
||||
) -> List[Union[Rollout, AttemptedRollout]]:
|
||||
"""Query the latest attempts for the rollouts, and attach them to the rollout objects."""
|
||||
async with collections.atomic(mode="r", snapshot=self._read_snapshot, labels=["attempts"]) as collections:
|
||||
attempts = await collections.attempts.query(
|
||||
filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
|
||||
sort={"name": "sequence_id", "order": "desc"},
|
||||
)
|
||||
attempts = await collections.attempts.query(
|
||||
filter={"rollout_id": {"within": [rollout.rollout_id for rollout in rollouts]}},
|
||||
sort={"name": "sequence_id", "order": "desc"},
|
||||
)
|
||||
latest_attempts: Dict[str, Attempt] = {}
|
||||
for attempt in attempts:
|
||||
if attempt.rollout_id not in latest_attempts:
|
||||
|
||||
@@ -11,7 +11,6 @@ from agentlightning.types import (
|
||||
Attempt,
|
||||
AttemptedRollout,
|
||||
AttemptStatus,
|
||||
EnqueueRolloutRequest,
|
||||
NamedResources,
|
||||
ResourcesUpdate,
|
||||
Rollout,
|
||||
@@ -60,17 +59,9 @@ class LightningStoreThreaded(LightningStore):
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> AttemptedRollout:
|
||||
with self._lock:
|
||||
return await self.store.start_rollout(
|
||||
input,
|
||||
mode,
|
||||
resources_id,
|
||||
config,
|
||||
metadata,
|
||||
worker_id,
|
||||
)
|
||||
return await self.store.start_rollout(input, mode, resources_id, config, metadata)
|
||||
|
||||
async def enqueue_rollout(
|
||||
self,
|
||||
@@ -83,26 +74,13 @@ class LightningStoreThreaded(LightningStore):
|
||||
with self._lock:
|
||||
return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
|
||||
|
||||
async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
|
||||
with self._lock:
|
||||
return await self.store.enqueue_many_rollouts(rollouts)
|
||||
|
||||
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
|
||||
with self._lock:
|
||||
return await self.store.dequeue_rollout(worker_id=worker_id)
|
||||
|
||||
async def dequeue_many_rollouts(
|
||||
self,
|
||||
*,
|
||||
limit: int = 1,
|
||||
worker_id: Optional[str] = None,
|
||||
) -> Sequence[AttemptedRollout]:
|
||||
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
|
||||
with self._lock:
|
||||
return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
|
||||
|
||||
async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
|
||||
with self._lock:
|
||||
return await self.store.start_attempt(rollout_id, worker_id)
|
||||
return await self.store.start_attempt(rollout_id)
|
||||
|
||||
async def query_rollouts(
|
||||
self,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import time
|
||||
from typing import Awaitable, Callable, Dict, List, Tuple
|
||||
from typing import Awaitable, Callable, List, cast
|
||||
|
||||
from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
|
||||
|
||||
@@ -57,54 +57,66 @@ LATENCY_BUCKETS = [
|
||||
]
|
||||
|
||||
|
||||
async def rollout_status_from_attempt(
|
||||
async def propagate_status(
|
||||
update_rollout_status: UpdateRolloutStatus, # this should be unlocked
|
||||
attempt: Attempt,
|
||||
config: RolloutConfig,
|
||||
) -> RolloutStatus:
|
||||
) -> Rollout:
|
||||
"""
|
||||
Propagate the status of an attempt to the rollout.
|
||||
|
||||
Returns:
|
||||
The status of the rollout from the perspective of the attempt.
|
||||
The rollout should be made sure in a state to be outdated.
|
||||
Requeue the rollout if it should be retried.
|
||||
|
||||
This operation is completely unlocked. The caller is responsible for locking the store.
|
||||
"""
|
||||
# Propagate the status directly to the rollout
|
||||
if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
|
||||
return attempt.status
|
||||
return await update_rollout_status(
|
||||
attempt.rollout_id,
|
||||
attempt.status,
|
||||
)
|
||||
|
||||
if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
|
||||
# Check if this status should trigger a retry
|
||||
if attempt.status in config.retry_condition:
|
||||
# If we haven't exceeded max attempts, retry
|
||||
if attempt.sequence_id < config.max_attempts:
|
||||
return "requeuing"
|
||||
return await update_rollout_status(
|
||||
attempt.rollout_id,
|
||||
"requeuing",
|
||||
)
|
||||
|
||||
# If we can't retry or shouldn't retry, mark as failed
|
||||
return "failed"
|
||||
return await update_rollout_status(
|
||||
attempt.rollout_id,
|
||||
"failed",
|
||||
)
|
||||
|
||||
raise ValueError(f"Invalid attempt status: {attempt.status}")
|
||||
|
||||
|
||||
async def scan_unhealthy_rollouts(
|
||||
async def healthcheck(
|
||||
rollouts: List[AttemptedRollout],
|
||||
) -> Dict[Tuple[str, str], AttemptStatus]:
|
||||
update_rollout_status: UpdateRolloutStatus,
|
||||
update_attempt_status: UpdateAttemptStatus,
|
||||
) -> None:
|
||||
"""
|
||||
Perform health check on all running rollouts in the store.
|
||||
|
||||
This method should be called periodically to:
|
||||
|
||||
1. Check for unresponsive attempts (no heartbeat or spans for a while)
|
||||
2. Check for timed-out rollouts (running too long since start_time)
|
||||
1. Update rollout status to failed to succeeded when the attempt is done
|
||||
2. Check for unresponsive attempts (no heartbeat or spans for a while)
|
||||
3. Check for timed-out rollouts (running too long since start_time)
|
||||
4. Update attempt/rollout status accordingly
|
||||
|
||||
This operation is completely unlocked. The caller is responsible for locking the store.
|
||||
|
||||
Args:
|
||||
rollouts: The list of running rollouts to check.
|
||||
|
||||
Returns:
|
||||
A dictionary of updates to the rollouts.
|
||||
store: The LightningStore instance to check rollouts from
|
||||
"""
|
||||
current_time = time.time()
|
||||
updates: Dict[Tuple[str, str], AttemptStatus] = {}
|
||||
|
||||
for rollout in rollouts:
|
||||
config = rollout.config # policy for retry and timeout
|
||||
@@ -112,31 +124,52 @@ async def scan_unhealthy_rollouts(
|
||||
# Get the latest attempt for this rollout
|
||||
latest_attempt = rollout.attempt
|
||||
if not latest_attempt:
|
||||
# This should not happen
|
||||
continue
|
||||
|
||||
# Check if the attempt has already failed or succeeded
|
||||
if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
|
||||
await propagate_status(update_rollout_status, latest_attempt, config)
|
||||
continue
|
||||
|
||||
# Check for timeout condition (based on attempt start_time, instead of rollout start_time)
|
||||
if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
|
||||
updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
|
||||
await update_attempt_status(
|
||||
latest_attempt.rollout_id,
|
||||
latest_attempt.attempt_id,
|
||||
"timeout",
|
||||
)
|
||||
continue
|
||||
|
||||
# Check for unresponsive condition (based on last heartbeat)
|
||||
# (1) Haven't received heartbeat for a while
|
||||
if (
|
||||
latest_attempt.last_heartbeat_time
|
||||
and config.unresponsive_seconds is not None
|
||||
and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
|
||||
):
|
||||
updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
|
||||
continue
|
||||
if latest_attempt.last_heartbeat_time:
|
||||
if latest_attempt.status == "preparing":
|
||||
# If still preparing, mark it as running
|
||||
latest_attempt = await update_attempt_status(
|
||||
latest_attempt.rollout_id,
|
||||
latest_attempt.attempt_id,
|
||||
"running",
|
||||
)
|
||||
|
||||
# (2) Check if there's no last heartbeat (no spans) at all
|
||||
# Haven't received heartbeat for a while
|
||||
if (
|
||||
config.unresponsive_seconds is not None
|
||||
and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
|
||||
):
|
||||
await update_attempt_status(
|
||||
latest_attempt.rollout_id,
|
||||
latest_attempt.attempt_id,
|
||||
"unresponsive",
|
||||
)
|
||||
continue
|
||||
|
||||
# Check if there's no last heartbeat (no spans) at all
|
||||
if (
|
||||
latest_attempt.last_heartbeat_time is None
|
||||
and config.unresponsive_seconds is not None
|
||||
and current_time - latest_attempt.start_time > config.unresponsive_seconds
|
||||
):
|
||||
updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
|
||||
continue
|
||||
|
||||
return updates
|
||||
await update_attempt_status(
|
||||
latest_attempt.rollout_id,
|
||||
latest_attempt.attempt_id,
|
||||
"unresponsive",
|
||||
)
|
||||
|
||||
@@ -1,16 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .agentops import AgentOpsTracer
|
||||
from .base import Tracer, clear_active_tracer, get_active_tracer, set_active_tracer
|
||||
from .dummy import DummyTracer
|
||||
from .base import Tracer
|
||||
from .otel import OtelTracer
|
||||
|
||||
__all__ = [
|
||||
"AgentOpsTracer",
|
||||
"Tracer",
|
||||
"OtelTracer",
|
||||
"DummyTracer",
|
||||
"get_active_tracer",
|
||||
"set_active_tracer",
|
||||
"clear_active_tracer",
|
||||
]
|
||||
__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer"]
|
||||
|
||||
@@ -13,13 +13,12 @@ import agentops.sdk.core
|
||||
import opentelemetry.trace as trace_api
|
||||
from agentops.sdk.core import TracingCore
|
||||
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
|
||||
from opentelemetry.trace import get_tracer_provider
|
||||
from opentelemetry.trace.status import StatusCode
|
||||
|
||||
from agentlightning.instrumentation import instrument_all, uninstrument_all
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.utils.otel import get_span_processors, get_tracer_provider
|
||||
|
||||
from .base import with_active_tracer_context
|
||||
from .otel import LightningSpanProcessor, OtelTracer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -80,20 +79,13 @@ class AgentOpsTracer(OtelTracer):
|
||||
agentops.init(auto_start_session=False) # type: ignore
|
||||
logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
|
||||
else:
|
||||
logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized. Skip initialization.")
|
||||
logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
|
||||
|
||||
span_processors = get_span_processors(self._get_tracer_provider(), LightningSpanProcessor)
|
||||
if len(span_processors) > 0:
|
||||
logger.warning(
|
||||
"LightningSpanProcessor already present in TracerProvider. You might have called init_worker() multiple times."
|
||||
"Agent-lightning will try to reuse the existing LightningSpanProcessor."
|
||||
)
|
||||
if len(span_processors) > 1:
|
||||
logger.error("More than one LightningSpanProcessors present in TracerProvider. This should not happen.")
|
||||
self._lightning_span_processor = span_processors[0]
|
||||
else:
|
||||
self._lightning_span_processor = LightningSpanProcessor()
|
||||
self._get_tracer_provider().add_span_processor(self._lightning_span_processor) # type: ignore
|
||||
self._lightning_span_processor = LightningSpanProcessor()
|
||||
|
||||
# TODO: The span processor cannot be deleted once added.
|
||||
# This might be a problem if the tracer is entered and exited multiple times.
|
||||
self._get_tracer_provider().add_span_processor(self._lightning_span_processor) # type: ignore
|
||||
|
||||
def teardown_worker(self, worker_id: int) -> None:
|
||||
super().teardown_worker(worker_id)
|
||||
@@ -102,10 +94,6 @@ class AgentOpsTracer(OtelTracer):
|
||||
self.uninstrument(worker_id)
|
||||
logger.info(f"[Worker {worker_id}] Instrumentation removed.")
|
||||
|
||||
# NOTE: The teardown doesn't try to remove the LightningSpanProcessor from the TracerProvider.
|
||||
# Currently there is no stable way to fully restore the AgentOps state to the initial state.
|
||||
|
||||
@with_active_tracer_context
|
||||
@asynccontextmanager
|
||||
async def trace_context(
|
||||
self,
|
||||
@@ -170,6 +158,7 @@ class AgentOpsTracer(OtelTracer):
|
||||
with self._agentops_trace_context(rollout_id, attempt_id, kwargs):
|
||||
yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
|
||||
elif store is None and rollout_id is None and attempt_id is None:
|
||||
# TODO: Add tests to cover both paths
|
||||
self._disable_native_otlp_exporter()
|
||||
with self._lightning_span_processor:
|
||||
with self._agentops_trace_context(None, None, kwargs):
|
||||
|
||||
@@ -2,13 +2,14 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar
|
||||
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus
|
||||
from agentlightning.types import ParallelWorkerBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.callbacks.base import BaseCallbackHandler # type: ignore
|
||||
@@ -16,14 +17,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
_active_tracer: Optional[Tracer] = None
|
||||
|
||||
T_func = Callable[..., Awaitable[Any]]
|
||||
|
||||
|
||||
class Tracer(ParallelWorkerBase):
|
||||
"""
|
||||
An abstract base class for tracers.
|
||||
@@ -105,12 +98,12 @@ class Tracer(ParallelWorkerBase):
|
||||
"""Internal API for CI backward compatibility."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_last_trace(self) -> List[Span]:
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
"""
|
||||
Retrieves the raw list of captured spans from the most recent trace.
|
||||
|
||||
Returns:
|
||||
A list of [`Span`][agentlightning.Span] objects collected during the last trace.
|
||||
A list of OpenTelemetry `ReadableSpan` objects.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -131,48 +124,6 @@ class Tracer(ParallelWorkerBase):
|
||||
with self._trace_context_sync(name=func.__name__):
|
||||
return func(*args, **kwargs)
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
"""Notify the tracer that a span should be created here.
|
||||
|
||||
It uses a fire-and-forget approach and doesn't wait for the span to be created.
|
||||
|
||||
Args:
|
||||
name: The name of the span.
|
||||
attributes: The attributes of the span.
|
||||
timestamp: The timestamp of the span.
|
||||
status: The status of the span.
|
||||
|
||||
Returns:
|
||||
The core fields of the span.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def operation_context(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
) -> ContextManager[SpanRecordingContext]:
|
||||
"""Start to record an operation to a span.
|
||||
|
||||
Args:
|
||||
name: The name of the operation.
|
||||
attributes: The attributes of the operation.
|
||||
start_time: The start time of the operation.
|
||||
end_time: The end time of the operation.
|
||||
|
||||
Returns:
|
||||
A [`SpanRecordingContext`][agentlightning.SpanRecordingContext] for recording the operation on the span.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
|
||||
"""
|
||||
A convenience wrapper to trace the execution of a single asynchronous function.
|
||||
@@ -224,64 +175,3 @@ class Tracer(ParallelWorkerBase):
|
||||
self.teardown_worker(0)
|
||||
if has_init:
|
||||
self.teardown()
|
||||
|
||||
|
||||
def set_active_tracer(tracer: Tracer):
|
||||
"""Set the active tracer for the current process.
|
||||
|
||||
Args:
|
||||
tracer: The tracer to set as active.
|
||||
"""
|
||||
global _active_tracer
|
||||
if _active_tracer is not None:
|
||||
raise ValueError("An active tracer is already set. Cannot set a new one.")
|
||||
_active_tracer = tracer
|
||||
|
||||
|
||||
def clear_active_tracer():
|
||||
"""Clear the active tracer for the current process."""
|
||||
global _active_tracer
|
||||
_active_tracer = None
|
||||
|
||||
|
||||
def get_active_tracer() -> Optional[Tracer]:
|
||||
"""Get the active tracer for the current process.
|
||||
|
||||
Returns:
|
||||
The active tracer, or None if no tracer is active.
|
||||
"""
|
||||
global _active_tracer
|
||||
return _active_tracer
|
||||
|
||||
|
||||
class _ActiveTracerAsyncCM(AsyncContextManager[T]):
|
||||
def __init__(self, tracer: Tracer, inner: AsyncContextManager[T]):
|
||||
self._tracer = tracer
|
||||
self._inner = inner
|
||||
|
||||
async def __aenter__(self) -> T:
|
||||
set_active_tracer(self._tracer) # will raise if nested
|
||||
try:
|
||||
return await self._inner.__aenter__()
|
||||
except Exception:
|
||||
clear_active_tracer()
|
||||
raise
|
||||
|
||||
async def __aexit__(self, *args: Any, **kwargs: Any) -> Optional[bool]:
|
||||
try:
|
||||
return await self._inner.__aexit__(*args, **kwargs)
|
||||
finally:
|
||||
clear_active_tracer()
|
||||
|
||||
|
||||
def with_active_tracer_context(
|
||||
func: Callable[..., AsyncContextManager[T]],
|
||||
) -> Callable[..., AsyncContextManager[T]]:
|
||||
"""Decorate a method returning an AsyncContextManager so tracer is active for the whole `async with`."""
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(self: Tracer, *args: Any, **kwargs: Any) -> AsyncContextManager[T]:
|
||||
cm = func(self, *args, **kwargs)
|
||||
return _ActiveTracerAsyncCM(self, cm)
|
||||
|
||||
return wrapper
|
||||
|
||||
@@ -1,106 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from typing import (
|
||||
Iterator,
|
||||
Optional,
|
||||
)
|
||||
|
||||
from agentlightning.types import (
|
||||
Attributes,
|
||||
SpanCoreFields,
|
||||
SpanRecordingContext,
|
||||
StatusCode,
|
||||
TraceStatus,
|
||||
)
|
||||
from agentlightning.utils.otel import format_exception_attributes
|
||||
|
||||
from .base import Tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DummySpanRecordingContext(SpanRecordingContext):
|
||||
"""Context for recording operations on a dummy span, not dependent on any backend tracer."""
|
||||
|
||||
def __init__(self, name: str, attributes: Optional[Attributes] = None, start_time: Optional[float] = None) -> None:
|
||||
self.name = name
|
||||
self.attributes = attributes or {}
|
||||
self.start_time = start_time or time.time()
|
||||
self.end_time = None
|
||||
self.status = TraceStatus(status_code="OK")
|
||||
|
||||
def record_exception(self, exception: BaseException) -> None:
|
||||
self.record_status("ERROR", str(exception))
|
||||
self.record_attributes(format_exception_attributes(exception))
|
||||
|
||||
def record_attributes(self, attributes: Attributes) -> None:
|
||||
self.attributes.update(attributes)
|
||||
|
||||
def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
|
||||
self.status = TraceStatus(status_code=status_code, description=description)
|
||||
|
||||
def finalize(self, end_time: Optional[float] = None) -> None:
|
||||
self.end_time = end_time or time.time()
|
||||
|
||||
def get_recorded_span(self) -> SpanCoreFields:
|
||||
if self.end_time is None:
|
||||
raise ValueError("End time is not set. Call finalize() first.")
|
||||
return SpanCoreFields(
|
||||
name=self.name,
|
||||
attributes=self.attributes,
|
||||
start_time=self.start_time,
|
||||
end_time=self.end_time,
|
||||
status=self.status,
|
||||
)
|
||||
|
||||
|
||||
class DummyTracer(Tracer):
|
||||
"""A dummy tracer that does not trace anything, but it is compatible with the emitter API.
|
||||
|
||||
It doesn't rely on any backend tracer, and also doesn't use any stores.
|
||||
"""
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
if attributes is None:
|
||||
attributes = {}
|
||||
if timestamp is None:
|
||||
timestamp = time.time()
|
||||
if status is None:
|
||||
status = TraceStatus(status_code="OK")
|
||||
return SpanCoreFields(
|
||||
name=name,
|
||||
attributes=attributes,
|
||||
start_time=timestamp,
|
||||
end_time=timestamp,
|
||||
status=status,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def operation_context(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
) -> Iterator[DummySpanRecordingContext]:
|
||||
start_time = start_time or time.time()
|
||||
recording_context = DummySpanRecordingContext(name, attributes, start_time)
|
||||
try:
|
||||
yield recording_context
|
||||
except Exception as exc:
|
||||
recording_context.record_exception(exc)
|
||||
recording_context.record_status("ERROR", str(exc))
|
||||
raise
|
||||
finally:
|
||||
recording_context.finalize(end_time)
|
||||
@@ -0,0 +1,396 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import multiprocessing
|
||||
import queue
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from httpdbg.hooks.all import httprecord
|
||||
from httpdbg.records import HTTPRecords
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace import SpanKind, Status, StatusCode
|
||||
from opentelemetry.trace.span import (
|
||||
SpanContext,
|
||||
TraceFlags,
|
||||
TraceState,
|
||||
)
|
||||
|
||||
from agentlightning.store import LightningStore
|
||||
|
||||
from .base import Tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HttpTracer(Tracer):
|
||||
"""
|
||||
A tracer implementation that captures HTTP requests using httpdbg.
|
||||
|
||||
This tracer hooks into the Python HTTP libraries and captures all
|
||||
HTTP requests and responses made during the traced code execution.
|
||||
The captured requests are converted to OpenTelemetry spans for
|
||||
compatibility with the rest of the tracing ecosystem.
|
||||
|
||||
Caution: The current implementation of HttpTracer is very fragile,
|
||||
and we do not recommend using it in production.
|
||||
It is primarily for demonstration and testing purposes.
|
||||
|
||||
Deprecated: This tracer is deprecated and will be removed in a future version.
|
||||
Please use LLMProxy as an alternative.
|
||||
|
||||
Attributes:
|
||||
include_headers: Whether to include HTTP headers in the spans.
|
||||
Headers may contain sensitive information. Use with caution.
|
||||
include_body: Whether to include HTTP request and response bodies in the spans.
|
||||
Bodies may be large and contain sensitive information. Use with caution.
|
||||
include_agentlightning_requests: Whether to include requests initiated by AgentLightning itself.
|
||||
subprocess_mode: Whether to run trace_run and trace_run_async in subprocesses for isolation.
|
||||
subprocess_timeout: Timeout for subprocess execution in seconds.
|
||||
"""
|
||||
|
||||
AGENTLIGHTNING_HEADERS = {"x-agentlightning-client"}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
include_headers: bool = False,
|
||||
include_body: bool = False,
|
||||
include_agentlightning_requests: bool = False,
|
||||
subprocess_mode: bool = True,
|
||||
subprocess_timeout: float = 3600.0,
|
||||
):
|
||||
super().__init__()
|
||||
self._last_records: Optional[HTTPRecords] = None
|
||||
self.include_headers = include_headers
|
||||
self.include_body = include_body
|
||||
self.include_agentlightning_requests = include_agentlightning_requests
|
||||
self.subprocess_mode = subprocess_mode
|
||||
self.subprocess_timeout = subprocess_timeout
|
||||
|
||||
def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
|
||||
"""
|
||||
Initialize the tracer in a worker process.
|
||||
|
||||
Args:
|
||||
worker_id: The ID of the worker process.
|
||||
store: The store to add the spans to.
|
||||
"""
|
||||
super().init_worker(worker_id, store)
|
||||
logger.info(f"[Worker {worker_id}] HttpTracer initialized.")
|
||||
|
||||
@asynccontextmanager
|
||||
async def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> AsyncGenerator[HTTPRecords, None]:
|
||||
"""
|
||||
Starts a new HTTP tracing context. This should be used as a context manager.
|
||||
|
||||
Args:
|
||||
name: Optional name for the tracing context.
|
||||
"""
|
||||
with self._trace_context_sync(name=name, **kwargs) as records:
|
||||
yield records
|
||||
|
||||
@contextmanager
|
||||
def _trace_context_sync(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
|
||||
"""
|
||||
Starts a new HTTP tracing context. This should be used as a context manager.
|
||||
|
||||
Args:
|
||||
name: Optional name for the tracing context.
|
||||
|
||||
Yields:
|
||||
The HTTPRecords instance containing traced HTTP activities.
|
||||
"""
|
||||
records = HTTPRecords()
|
||||
with httprecord(records):
|
||||
self._last_records = records
|
||||
yield records
|
||||
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
"""
|
||||
Retrieves the raw list of captured spans from the most recent trace.
|
||||
|
||||
Returns:
|
||||
A list of OpenTelemetry `ReadableSpan` objects converted from HTTP records.
|
||||
"""
|
||||
if self._last_records is None:
|
||||
return []
|
||||
|
||||
return self._convert_to_spans(self._last_records)
|
||||
|
||||
def _convert_to_spans(self, records: HTTPRecords) -> List[ReadableSpan]:
|
||||
"""
|
||||
Convert HTTPRecords to OpenTelemetry spans.
|
||||
|
||||
Args:
|
||||
records: The HTTPRecords instance containing HTTP traces.
|
||||
|
||||
Returns:
|
||||
A list of ReadableSpan objects representing the HTTP activities.
|
||||
"""
|
||||
spans: List[ReadableSpan] = []
|
||||
|
||||
# Create a trace ID that will be shared by all spans in this trace
|
||||
trace_id = int(uuid.uuid4().hex[:16], 16)
|
||||
|
||||
for record in records.requests.values():
|
||||
# Skip AgentLightning requests if include_agentlightning_requests is False
|
||||
should_skip = False
|
||||
if not self.include_agentlightning_requests and record.request and record.request.headers:
|
||||
for header in record.request.headers:
|
||||
if header.name.lower() in self.AGENTLIGHTNING_HEADERS and header.value.lower() == "true":
|
||||
should_skip = True
|
||||
break
|
||||
|
||||
if should_skip:
|
||||
continue
|
||||
|
||||
# Create a span ID for this specific HTTP request
|
||||
span_id = int(uuid.uuid4().hex[:8], 16)
|
||||
|
||||
# Create a span context
|
||||
span_context = SpanContext(
|
||||
trace_id=trace_id,
|
||||
span_id=span_id,
|
||||
is_remote=False,
|
||||
trace_flags=TraceFlags(TraceFlags.SAMPLED),
|
||||
trace_state=TraceState(),
|
||||
)
|
||||
|
||||
# Extract important information from the HTTP record
|
||||
method = record.method
|
||||
url = record.url
|
||||
parsed_url = urlparse(url)
|
||||
status_code = record.status_code
|
||||
|
||||
# Create attributes dictionary
|
||||
attributes: Dict[str, Any] = {
|
||||
"http.method": method,
|
||||
"http.url": url,
|
||||
"http.target": parsed_url.path,
|
||||
"http.host": parsed_url.netloc,
|
||||
}
|
||||
|
||||
if status_code is not None and status_code > 0: # type: ignore
|
||||
attributes["http.status_code"] = status_code
|
||||
|
||||
# Calculate duration - from begin time to last update
|
||||
duration = None
|
||||
if hasattr(record, "last_update") and record.last_update and record.tbegin:
|
||||
duration = (record.last_update - record.tbegin).total_seconds()
|
||||
attributes["http.duration_ms"] = duration * 1000 # Convert to ms
|
||||
|
||||
# Optionally include headers
|
||||
if self.include_headers and record.request and record.request.headers:
|
||||
for header in record.request.headers:
|
||||
header_name = header.name.lower()
|
||||
attributes[f"http.request.header.{header_name}"] = header.value
|
||||
|
||||
if self.include_headers and record.response and record.response.headers:
|
||||
for header in record.response.headers:
|
||||
header_name = header.name.lower()
|
||||
attributes[f"http.response.header.{header_name}"] = header.value
|
||||
|
||||
# Optionally include body - preserve complete content for analysis
|
||||
if self.include_body and record.request:
|
||||
body_content = record.request.content
|
||||
if body_content:
|
||||
# Store raw body content for later parsing/analysis
|
||||
attributes["http.request.body"] = body_content
|
||||
|
||||
if self.include_body and record.response:
|
||||
body_content = record.response.content
|
||||
if body_content:
|
||||
# Store raw body content for later parsing/analysis
|
||||
attributes["http.response.body"] = body_content
|
||||
|
||||
# Determine span status
|
||||
span_status = StatusCode.OK
|
||||
if status_code and status_code >= 400 or record.exception:
|
||||
span_status = StatusCode.ERROR
|
||||
|
||||
# Create start and end timestamps in nanoseconds
|
||||
# If we have duration, use it, otherwise default to current time - 1ms
|
||||
start_time_ns = int(record.tbegin.timestamp() * 1e9)
|
||||
if duration:
|
||||
end_time_ns = int((record.tbegin.timestamp() + duration) * 1e9)
|
||||
else:
|
||||
end_time_ns = int(record.last_update.timestamp() * 1e9)
|
||||
|
||||
span = ReadableSpan(
|
||||
name=f"HTTP {method} {url}",
|
||||
context=span_context,
|
||||
parent=None,
|
||||
kind=SpanKind.CLIENT,
|
||||
status=Status(span_status),
|
||||
start_time=start_time_ns,
|
||||
end_time=end_time_ns,
|
||||
attributes=attributes,
|
||||
events=[],
|
||||
links=[],
|
||||
resource=None,
|
||||
)
|
||||
|
||||
spans.append(span)
|
||||
|
||||
return spans
|
||||
|
||||
def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
||||
"""
|
||||
A convenience wrapper to trace the execution of a single synchronous function.
|
||||
|
||||
If subprocess_mode is enabled, the function will be executed in an isolated subprocess
|
||||
to prevent HTTP hooks from affecting the parent process.
|
||||
|
||||
Args:
|
||||
func: The synchronous function to execute and trace.
|
||||
*args: Positional arguments to pass to the function.
|
||||
**kwargs: Keyword arguments to pass to the function.
|
||||
|
||||
Returns:
|
||||
The return value of the function.
|
||||
"""
|
||||
if self.subprocess_mode:
|
||||
return self._trace_run_subprocess(func, args, kwargs)
|
||||
else:
|
||||
return super().trace_run(func, *args, **kwargs)
|
||||
|
||||
async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
|
||||
"""
|
||||
A convenience wrapper to trace the execution of a single asynchronous function.
|
||||
|
||||
If subprocess_mode is enabled, the function will be executed in an isolated subprocess
|
||||
to prevent HTTP hooks from affecting the parent process.
|
||||
|
||||
Args:
|
||||
func: The asynchronous function to execute and trace.
|
||||
*args: Positional arguments to pass to the function.
|
||||
**kwargs: Keyword arguments to pass to the function.
|
||||
|
||||
Returns:
|
||||
The return value of the function.
|
||||
"""
|
||||
if self.subprocess_mode:
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(
|
||||
None, self._trace_run_subprocess, func, args, kwargs, True # True for async
|
||||
)
|
||||
else:
|
||||
return await super().trace_run_async(func, *args, **kwargs)
|
||||
|
||||
def _trace_run_subprocess(
|
||||
self,
|
||||
func: Callable[..., Any],
|
||||
args: Optional[Tuple[Any, ...]] = None,
|
||||
kwargs: Optional[Dict[str, Any]] = None,
|
||||
is_async: bool = False,
|
||||
) -> Any:
|
||||
"""
|
||||
Execute a function in a subprocess with HTTP tracing.
|
||||
|
||||
Args:
|
||||
func: The function to execute.
|
||||
args: Positional arguments to pass to the function.
|
||||
kwargs: Keyword arguments to pass to the function.
|
||||
is_async: Whether the function is asynchronous.
|
||||
|
||||
Returns:
|
||||
The return value of the function.
|
||||
"""
|
||||
if args is None:
|
||||
args = ()
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
# Create a queue to receive results from the subprocess
|
||||
result_queue = multiprocessing.Queue() # type: ignore
|
||||
|
||||
# Create and start the subprocess
|
||||
process = multiprocessing.Process(
|
||||
target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async) # type: ignore
|
||||
)
|
||||
process.start()
|
||||
|
||||
try:
|
||||
# Wait for the process to complete and get the result
|
||||
process.join(timeout=self.subprocess_timeout)
|
||||
result = result_queue.get_nowait() # type: ignore
|
||||
|
||||
if result["success"]:
|
||||
# Store the captured records for get_last_trace()
|
||||
self._last_records = result["records"]
|
||||
return result["return_value"] # type: ignore
|
||||
else:
|
||||
if "records" in result:
|
||||
self._last_records = result["records"]
|
||||
# Re-raise the exception that occurred in the subprocess
|
||||
raise result["exception"]
|
||||
|
||||
except multiprocessing.TimeoutError:
|
||||
process.terminate()
|
||||
process.join()
|
||||
raise TimeoutError(f"Subprocess execution timed out after {self.subprocess_timeout} seconds.")
|
||||
except queue.Empty:
|
||||
logger.error("Traced result is empty. This may indicate a timeout or an issue with the subprocess.")
|
||||
finally:
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join()
|
||||
|
||||
def _subprocess_worker(
|
||||
self,
|
||||
func: Callable[..., Any],
|
||||
args: Tuple[Any, ...],
|
||||
kwargs: Dict[str, Any],
|
||||
result_queue: multiprocessing.Queue, # type: ignore
|
||||
is_async: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Worker function that runs in the subprocess to execute the traced function.
|
||||
|
||||
Args:
|
||||
func: The function to execute.
|
||||
args: Positional arguments.
|
||||
kwargs: Keyword arguments.
|
||||
result_queue: Queue to send results back to parent process.
|
||||
is_async: Whether the function is asynchronous.
|
||||
"""
|
||||
# Create a new tracer instance in the subprocess (without subprocess mode to avoid recursion)
|
||||
subprocess_tracer = HttpTracer(
|
||||
include_headers=self.include_headers,
|
||||
include_body=self.include_body,
|
||||
include_agentlightning_requests=self.include_agentlightning_requests,
|
||||
subprocess_mode=False, # Disable subprocess mode in the worker
|
||||
)
|
||||
|
||||
try:
|
||||
if is_async:
|
||||
# Run async function in new event loop
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return_value = loop.run_until_complete(subprocess_tracer.trace_run_async(func, *args, **kwargs))
|
||||
finally:
|
||||
loop.close()
|
||||
else:
|
||||
# Run sync function
|
||||
return_value = subprocess_tracer.trace_run(func, *args, **kwargs)
|
||||
|
||||
# Get the captured records
|
||||
records = subprocess_tracer._last_records
|
||||
|
||||
# Send success result back to parent
|
||||
result_queue.put({"success": True, "return_value": return_value, "records": records}) # type: ignore
|
||||
|
||||
except Exception as e:
|
||||
# Log the exception
|
||||
logger.exception(f"Error in subprocess worker in http tracer: {e}")
|
||||
|
||||
# Get the captured records even when there's an exception
|
||||
records = subprocess_tracer._last_records
|
||||
# Send error result back to parent
|
||||
result_queue.put({"success": False, "exception": e, "records": records}) # type: ignore
|
||||
+19
-165
@@ -6,69 +6,27 @@ import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import warnings
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any, AsyncGenerator, Awaitable, Iterator, List, Optional
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, Awaitable, List, Optional
|
||||
|
||||
import opentelemetry.trace as trace_api
|
||||
from agentops.sdk.core import BatchSpanProcessor
|
||||
from opentelemetry.instrumentation.utils import suppress_instrumentation
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
|
||||
from agentlightning.semconv import LightningResourceAttributes
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Attributes, Span, SpanCoreFields, SpanRecordingContext, StatusCode, TraceStatus
|
||||
from agentlightning.types.tracer import convert_timestamp
|
||||
from agentlightning.utils.otel import get_tracer_provider
|
||||
from agentlightning.utils.otlp import LightningStoreOTLPExporter
|
||||
|
||||
from .base import Tracer, with_active_tracer_context
|
||||
from .base import Tracer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
STORE_WRITE_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
|
||||
def to_otel_status_code(status_code: StatusCode) -> trace_api.StatusCode:
|
||||
if status_code == "UNSET":
|
||||
return trace_api.StatusCode.UNSET
|
||||
elif status_code == "ERROR":
|
||||
return trace_api.StatusCode.ERROR
|
||||
else:
|
||||
return trace_api.StatusCode.OK
|
||||
|
||||
|
||||
class OtelSpanRecordingContext(SpanRecordingContext):
|
||||
def __init__(self, span: trace_api.Span) -> None:
|
||||
self._span = span
|
||||
|
||||
def record_exception(self, exception: BaseException) -> None:
|
||||
self._span.record_exception(exception)
|
||||
self.record_status("ERROR", str(exception))
|
||||
|
||||
def record_attributes(self, attributes: Attributes) -> None:
|
||||
self._span.set_attributes(attributes)
|
||||
|
||||
def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
|
||||
otel_status_code = to_otel_status_code(status_code)
|
||||
self._span.set_status(otel_status_code, description)
|
||||
|
||||
def get_otel_span(self) -> trace_api.Span:
|
||||
return self._span
|
||||
|
||||
def get_recorded_span(self) -> SpanCoreFields:
|
||||
if isinstance(self._span, ReadableSpan):
|
||||
return SpanCoreFields(
|
||||
name=self._span.name,
|
||||
attributes=dict(self._span.attributes) if self._span.attributes else {},
|
||||
start_time=convert_timestamp(self._span.start_time),
|
||||
end_time=convert_timestamp(self._span.end_time),
|
||||
status=TraceStatus.from_opentelemetry(self._span.status),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Span is not a ReadableSpan: {self._span}")
|
||||
|
||||
|
||||
class OtelTracer(Tracer):
|
||||
"""Tracer that provides a basic OpenTelemetry tracer provider.
|
||||
@@ -80,7 +38,7 @@ class OtelTracer(Tracer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# This provider is only initialized when the worker is initialized.
|
||||
self._tracer_provider: Optional[trace_api.TracerProvider] = None
|
||||
self._tracer_provider: Optional[TracerProvider] = None
|
||||
self._lightning_span_processor: Optional[LightningSpanProcessor] = None
|
||||
self._simple_span_processor: Optional[SimpleSpanProcessor] = None
|
||||
self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
|
||||
@@ -105,7 +63,7 @@ class OtelTracer(Tracer):
|
||||
except RuntimeError:
|
||||
logger.debug(f"[Worker {worker_id}] Tracer provider is not initialized by OtelTracer. Initializing it now.")
|
||||
|
||||
self._tracer_provider = TracerProviderImpl()
|
||||
self._tracer_provider = TracerProvider()
|
||||
trace_api.set_tracer_provider(self._tracer_provider)
|
||||
self._lightning_span_processor = LightningSpanProcessor()
|
||||
self._tracer_provider.add_span_processor(self._lightning_span_processor)
|
||||
@@ -120,7 +78,6 @@ class OtelTracer(Tracer):
|
||||
super().teardown_worker(worker_id)
|
||||
logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer does NOT remove the tracer provider.")
|
||||
|
||||
@with_active_tracer_context
|
||||
@asynccontextmanager
|
||||
async def trace_context(
|
||||
self,
|
||||
@@ -172,69 +129,12 @@ class OtelTracer(Tracer):
|
||||
else:
|
||||
raise ValueError("rollout_id and attempt_id must be either all provided or all None")
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
# Fire the span to the current active tracer provider.
|
||||
tracer_provider = self._get_tracer_provider()
|
||||
tracer = tracer_provider.get_tracer(__name__)
|
||||
span = tracer.start_span(
|
||||
name, attributes=attributes, start_time=int(timestamp * 1_000_000_000) if timestamp else None
|
||||
)
|
||||
if status is not None:
|
||||
span.set_status(to_otel_status_code(status.status_code), status.description)
|
||||
span.end(int(timestamp * 1_000_000_000) if timestamp else None)
|
||||
|
||||
# The span should have been auto-created by now.
|
||||
# Return the core fields of the span.
|
||||
if isinstance(span, ReadableSpan):
|
||||
return SpanCoreFields(
|
||||
name=name,
|
||||
attributes=dict(span.attributes) if span.attributes else {},
|
||||
start_time=convert_timestamp(span.start_time),
|
||||
end_time=convert_timestamp(span.end_time),
|
||||
status=TraceStatus.from_opentelemetry(span.status),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Span is not a ReadableSpan: {span}")
|
||||
|
||||
@contextmanager
|
||||
def operation_context(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
) -> Iterator[SpanRecordingContext]:
|
||||
if end_time is not None:
|
||||
logger.warning("OpenTelemetry doesn't support customizing the end time of a span. End time is ignored.")
|
||||
# Record the span to the current active tracer provider.
|
||||
tracer_provider = self._get_tracer_provider()
|
||||
tracer = tracer_provider.get_tracer(__name__)
|
||||
|
||||
# Activate the span as the current span within otel.
|
||||
with tracer.start_as_current_span(
|
||||
name, attributes=attributes, start_time=int(start_time * 1_000_000_000) if start_time else None
|
||||
) as span:
|
||||
recording_context = OtelSpanRecordingContext(span)
|
||||
try:
|
||||
yield recording_context
|
||||
except Exception as exc:
|
||||
recording_context.record_exception(exc)
|
||||
raise
|
||||
|
||||
# No need to retrieve the span here. It's already been sent to otel processor.
|
||||
|
||||
def get_last_trace(self) -> List[Span]:
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
"""
|
||||
Retrieves the raw list of captured spans from the most recent trace.
|
||||
|
||||
Returns:
|
||||
A list of [`Span`][agentlightning.Span] objects captured during the most recent trace.
|
||||
A list of OpenTelemetry `ReadableSpan` objects.
|
||||
"""
|
||||
if not self._lightning_span_processor:
|
||||
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
|
||||
@@ -243,8 +143,6 @@ class OtelTracer(Tracer):
|
||||
def _get_tracer_provider(self) -> TracerProviderImpl:
|
||||
if self._tracer_provider is None:
|
||||
raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
|
||||
if not isinstance(self._tracer_provider, TracerProviderImpl):
|
||||
raise TypeError(f"TracerProvider is not a opentelemetry.sdk.trace.TracerProvider: {self._tracer_provider}")
|
||||
return self._tracer_provider
|
||||
|
||||
def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
|
||||
@@ -317,20 +215,18 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
|
||||
def __init__(self, disable_store_submission: bool = False):
|
||||
self._disable_store_submission: bool = disable_store_submission
|
||||
self._spans: List[Span] = []
|
||||
self._spans: List[ReadableSpan] = []
|
||||
|
||||
# Store related context and states
|
||||
self._store: Optional[LightningStore] = None
|
||||
self._rollout_id: Optional[str] = None
|
||||
self._attempt_id: Optional[str] = None
|
||||
self._local_sequence_id: int = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# private asyncio loop running in a daemon thread
|
||||
self._loop_ready = threading.Event()
|
||||
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
self._loop_thread: Optional[threading.Thread] = None
|
||||
self._loop_init_lock = threading.Lock()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
@@ -366,19 +262,11 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
self._disable_store_submission = value
|
||||
|
||||
def _ensure_loop(self) -> None:
|
||||
# Fast path: loop already initialized
|
||||
if self._loop_thread is not None and self._loop is not None:
|
||||
return
|
||||
|
||||
with self._loop_init_lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._loop_thread is not None and self._loop is not None:
|
||||
return
|
||||
if self._loop_thread is None or self._loop is None:
|
||||
self._loop_ready.clear()
|
||||
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
|
||||
self._loop_thread.start()
|
||||
if not self._loop_ready.wait(timeout=30.0):
|
||||
raise RuntimeError("Timed out waiting for otel-loop thread to start")
|
||||
self._loop_ready.wait() # loop is ready
|
||||
|
||||
def _loop_runner(self):
|
||||
loop = asyncio.new_event_loop()
|
||||
@@ -442,13 +330,13 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return True
|
||||
|
||||
def spans(self) -> List[Span]:
|
||||
def spans(self) -> List[ReadableSpan]:
|
||||
"""
|
||||
Get the list of spans collected by this processor.
|
||||
This is useful for debugging and testing purposes.
|
||||
|
||||
Returns:
|
||||
List of [`Span`][agentlightning.Span] objects collected during tracing.
|
||||
List of ReadableSpan objects collected during tracing.
|
||||
"""
|
||||
return self._spans
|
||||
|
||||
@@ -485,46 +373,12 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
# Submit add_otel_span to the event loop and wait for it to complete
|
||||
with suppress_instrumentation():
|
||||
self._ensure_loop()
|
||||
uploaded_span = self._await_in_loop(
|
||||
self._await_in_loop(
|
||||
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
|
||||
timeout=STORE_WRITE_TIMEOUT_SECONDS,
|
||||
timeout=60.0,
|
||||
)
|
||||
if uploaded_span is not None:
|
||||
self._spans.append(uploaded_span)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"Timed out adding span %s to store after %.1f seconds. The span will be stored locally "
|
||||
"but it's not guaranteed to be persisted.",
|
||||
span.name,
|
||||
STORE_WRITE_TIMEOUT_SECONDS,
|
||||
)
|
||||
self._spans.append(
|
||||
Span.from_opentelemetry(
|
||||
span,
|
||||
rollout_id=self._rollout_id,
|
||||
attempt_id=self._attempt_id,
|
||||
sequence_id=self._local_sequence_id,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
# log; on_end MUST NOT raise
|
||||
logger.exception(f"Error adding span to store: {span.name}. The span will be store locally only.")
|
||||
self._spans.append(
|
||||
Span.from_opentelemetry(
|
||||
span,
|
||||
rollout_id=self._rollout_id,
|
||||
attempt_id=self._attempt_id,
|
||||
sequence_id=self._local_sequence_id,
|
||||
)
|
||||
)
|
||||
logger.exception(f"Error adding span to store: {span.name}")
|
||||
|
||||
else:
|
||||
# Fallback path
|
||||
created_span = Span.from_opentelemetry(
|
||||
span,
|
||||
rollout_id=self._rollout_id or "rollout-dummy",
|
||||
attempt_id=self._attempt_id or "attempt-dummy",
|
||||
sequence_id=self._local_sequence_id,
|
||||
)
|
||||
self._local_sequence_id += 1
|
||||
self._spans.append(created_span)
|
||||
self._spans.append(span)
|
||||
|
||||
@@ -1,677 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures as futures
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import weakref
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
cast,
|
||||
)
|
||||
|
||||
import weave
|
||||
from opentelemetry.semconv.attributes import exception_attributes
|
||||
from weave.trace.call import Call
|
||||
from weave.trace.settings import UserSettings
|
||||
from weave.trace.weave_client import WeaveClient
|
||||
from weave.trace_server import trace_server_interface as tsi
|
||||
from weave.wandb_interface.context import set_wandb_api_context
|
||||
|
||||
from agentlightning.instrumentation.weave import InMemoryWeaveTraceServer, instrument_weave, uninstrument_weave
|
||||
from agentlightning.semconv import LightningResourceAttributes, LightningSpanAttributes
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import (
|
||||
Attributes,
|
||||
OtelResource,
|
||||
Span,
|
||||
SpanContext,
|
||||
SpanCoreFields,
|
||||
SpanRecordingContext,
|
||||
StatusCode,
|
||||
TraceStatus,
|
||||
)
|
||||
from agentlightning.utils.id import generate_id
|
||||
from agentlightning.utils.otel import (
|
||||
filter_and_unflatten_attributes,
|
||||
flatten_attributes,
|
||||
format_exception_attributes,
|
||||
sanitize_attributes,
|
||||
)
|
||||
|
||||
from .base import Tracer, with_active_tracer_context
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def op_name_to_func_name(op_name: str) -> str:
|
||||
"""Convert a Weave operation name to a function name.
|
||||
|
||||
Weave operation names look like this: `weave:///xxx/agentlightning.tracer.weave/op/openai.chat.completions.create:019b10be-...-44d74272569c`
|
||||
"""
|
||||
match = re.search(r"/([^/:]+):", op_name)
|
||||
if match:
|
||||
return match.group(1)
|
||||
else:
|
||||
return op_name
|
||||
|
||||
|
||||
def random_project_name() -> str:
|
||||
return "agl/weave-" + generate_id(12)
|
||||
|
||||
|
||||
def get_timestamp_or_throw(date: Optional[datetime], field_name: str) -> float:
|
||||
if date is None:
|
||||
raise ValueError(f"{field_name} is required but not set")
|
||||
return date.timestamp()
|
||||
|
||||
|
||||
class WeaveSpanRecordingContext(SpanRecordingContext):
|
||||
"""Universal interface for recording operations on a Weave call."""
|
||||
|
||||
def __init__(self, call: Call) -> None:
|
||||
self._call = call
|
||||
|
||||
def record_exception(self, exception: BaseException) -> None:
|
||||
self._call.exception = str(exception)
|
||||
self.record_status("ERROR", str(exception))
|
||||
self.record_attributes(format_exception_attributes(exception))
|
||||
|
||||
def _get_input_from_attributes(self, attributes: Attributes) -> Dict[str, Any]:
|
||||
if LightningSpanAttributes.OPERATION_INPUT.value in attributes:
|
||||
# This can be a very rare case. If it happens, we can just let it throw.
|
||||
return cast(Dict[str, Any], attributes[LightningSpanAttributes.OPERATION_INPUT.value])
|
||||
else:
|
||||
filtered_attributes = filter_and_unflatten_attributes(
|
||||
attributes, LightningSpanAttributes.OPERATION_INPUT.value
|
||||
)
|
||||
if isinstance(filtered_attributes, list):
|
||||
return {str(i): v for i, v in enumerate(filtered_attributes)}
|
||||
else:
|
||||
return filtered_attributes
|
||||
|
||||
def _get_output_from_attributes(self, attributes: Attributes) -> Any:
|
||||
if LightningSpanAttributes.OPERATION_OUTPUT.value in attributes:
|
||||
return attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]
|
||||
else:
|
||||
return filter_and_unflatten_attributes(attributes, LightningSpanAttributes.OPERATION_OUTPUT.value)
|
||||
|
||||
def record_attributes(self, attributes: Attributes) -> None:
|
||||
input_attributes = self._get_input_from_attributes(attributes)
|
||||
if input_attributes:
|
||||
self._call.inputs.update(input_attributes)
|
||||
|
||||
output_attributes = self._get_output_from_attributes(attributes)
|
||||
if output_attributes:
|
||||
if self._call.output is not None:
|
||||
logger.warning(f"Output is already set. It will be overridden: {self._call.output}")
|
||||
self._call.output = output_attributes
|
||||
|
||||
if LightningSpanAttributes.OPERATION_NAME.value in attributes:
|
||||
logger.error(
|
||||
f"Cannot record operation name as an attribute. It will be skipped: {attributes[LightningSpanAttributes.OPERATION_NAME.value]}"
|
||||
)
|
||||
|
||||
# The rest of the attributes are recorded as summary.
|
||||
for key, value in attributes.items():
|
||||
if (
|
||||
not key == LightningSpanAttributes.OPERATION_INPUT.value
|
||||
and not key.startswith(LightningSpanAttributes.OPERATION_INPUT.value + ".")
|
||||
and not key == LightningSpanAttributes.OPERATION_OUTPUT.value
|
||||
and not key.startswith(LightningSpanAttributes.OPERATION_OUTPUT.value + ".")
|
||||
and not key == LightningSpanAttributes.OPERATION_NAME.value
|
||||
):
|
||||
if self._call.summary is None:
|
||||
self._call.summary = {}
|
||||
self._call.summary[key] = value
|
||||
|
||||
def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
|
||||
if status_code == "ERROR":
|
||||
if not description:
|
||||
raise ValueError("Description is required when status code is ERROR")
|
||||
self._call.exception = description
|
||||
elif status_code == "OK":
|
||||
self._call.exception = None
|
||||
# Do nothing for other status codes.
|
||||
|
||||
def finalize(self) -> None:
|
||||
# Do nothing
|
||||
pass
|
||||
|
||||
def get_recorded_span(self) -> SpanCoreFields:
|
||||
return SpanCoreFields(
|
||||
name=self._call.op_name,
|
||||
attributes=flatten_attributes(self._call.attributes or {}),
|
||||
start_time=self._call.started_at.timestamp() if self._call.started_at else None,
|
||||
end_time=self._call.ended_at.timestamp() if self._call.ended_at else None,
|
||||
status=TraceStatus(
|
||||
status_code="OK" if self._call.exception is None else "ERROR", description=self._call.exception
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class WeaveTracerManagedTraceServer(InMemoryWeaveTraceServer):
|
||||
"""A managed trace server for WeaveTracer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
partial_call_callback: Callable[[Dict[str, Any]], None],
|
||||
complete_call_callback: Callable[[tsi.CallSchema], None],
|
||||
):
|
||||
super().__init__()
|
||||
self.partial_call_callback = partial_call_callback
|
||||
self.complete_call_callback = complete_call_callback
|
||||
self._calls_already_invoked: set[str] = set()
|
||||
|
||||
def trigger_callbacks(self, call_id: str) -> None:
|
||||
with self._call_threading_lock:
|
||||
if call_id in self.calls:
|
||||
if call_id not in self._calls_already_invoked:
|
||||
self._calls_already_invoked.add(call_id)
|
||||
self.complete_call_callback(self.calls[call_id])
|
||||
else:
|
||||
logger.info(f"Call {call_id} has callback already invoked. Skipping.")
|
||||
elif call_id in self.partial_calls:
|
||||
self.partial_call_callback(self.partial_calls[call_id])
|
||||
else:
|
||||
logger.error(f"Call {call_id} not found in partial_calls or calls")
|
||||
|
||||
def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
|
||||
try:
|
||||
ret = super().call_start(req)
|
||||
self.trigger_callbacks(ret.id)
|
||||
return ret
|
||||
except Exception:
|
||||
logger.exception(f"Error calling call_start: {req}", exc_info=True)
|
||||
raise
|
||||
|
||||
def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
|
||||
try:
|
||||
ret = super().call_end(req)
|
||||
self.trigger_callbacks(req.end.id)
|
||||
return ret
|
||||
except Exception:
|
||||
logger.exception(f"Error calling call_end: {req}", exc_info=True)
|
||||
raise
|
||||
|
||||
def clear(self) -> None:
|
||||
self._calls_already_invoked.clear()
|
||||
|
||||
|
||||
class WeaveTracer(Tracer):
|
||||
"""Tracer implementation using Weave for telemetry and trace logging.
|
||||
|
||||
This replaces AgentOpsTracer with a Weave-based manual trace context. It tracks:
|
||||
|
||||
- Function/method calls
|
||||
- Input/Output data
|
||||
- Exceptions
|
||||
|
||||
and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project_name: str | None = None,
|
||||
weave_user_settings: UserSettings | None = None,
|
||||
instrument_managed: bool = True,
|
||||
):
|
||||
"""Initialize a WeaveTracer instance.
|
||||
|
||||
Args:
|
||||
project_name: Optional project name for Weave; defaults to the current module name.
|
||||
weave_user_settings: Optional UserSettings for Weave.
|
||||
instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
|
||||
"""
|
||||
super().__init__()
|
||||
self.project_name = project_name
|
||||
self.instrument_managed = instrument_managed
|
||||
self.weave_user_settings = weave_user_settings or UserSettings(use_server_cache=False)
|
||||
|
||||
self._store: Optional[LightningStore] = None
|
||||
self._server = WeaveTracerManagedTraceServer(
|
||||
partial_call_callback=self.partial_call_callback, complete_call_callback=self.complete_call_callback
|
||||
)
|
||||
|
||||
self._default_sequence_counter: int = 0
|
||||
self._calls: Dict[str, tsi.CallSchema] = {} # call_id -> call
|
||||
self._spans: List[Span] = [] # spans in the current trace
|
||||
self._rollout_id: Optional[str] = None
|
||||
self._attempt_id: Optional[str] = None
|
||||
self._partial_call_futures: Dict[str, asyncio.Future[int] | futures.Future[int]] = {}
|
||||
self._complete_call_futures: List[asyncio.Future[None] | futures.Future[None]] = []
|
||||
self._loop: weakref.ReferenceType[asyncio.AbstractEventLoop] | None = None
|
||||
|
||||
def instrument(self, worker_id: int):
|
||||
instrument_weave(self._server)
|
||||
|
||||
def uninstrument(self, worker_id: int):
|
||||
uninstrument_weave()
|
||||
|
||||
def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
|
||||
"""
|
||||
Initialize the tracer for a worker thread/process.
|
||||
|
||||
Args:
|
||||
worker_id: Identifier of the worker.
|
||||
store: Optional LightningStore for storing spans.
|
||||
"""
|
||||
super().init_worker(worker_id, store)
|
||||
logger.info(f"[Worker {worker_id}] Setting up Weave tracer...")
|
||||
self._store = store
|
||||
|
||||
# Optionally patch network calls to bypass real Weave/W&B endpoints
|
||||
if self.instrument_managed:
|
||||
self.instrument(worker_id)
|
||||
|
||||
# If WANDB_API_KEY is not set, we need to initialize Weave with a hack
|
||||
if not os.getenv("WANDB_API_KEY"):
|
||||
logger.info("WANDB_API_KEY is not set. Initializing Weave a mock context.")
|
||||
set_wandb_api_context("agl", api_key=None, headers=None, cookies=None)
|
||||
else:
|
||||
logger.debug("WANDB_API_KEY is set. Weave will be initialized automatically.")
|
||||
|
||||
weave_client = weave.get_client()
|
||||
if self.project_name is None:
|
||||
self.project_name = random_project_name()
|
||||
|
||||
if weave_client is not None:
|
||||
logger.warning("Weave client was already initialized. Reentrant calls are at your own risk.")
|
||||
if weave_client.project == self.project_name:
|
||||
logger.error(
|
||||
f"Weave client was already initialized for the same project '{self.project_name}'. It's very likely that weave won't work correctly."
|
||||
)
|
||||
|
||||
# Init no matter what
|
||||
try:
|
||||
weave.init(project_name=self.project_name, settings=self.weave_user_settings)
|
||||
logger.info(f"[Worker {worker_id}] Weave client initialized.")
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}'") from exc
|
||||
|
||||
def teardown_worker(self, worker_id: int):
|
||||
"""
|
||||
Clean up tracer resources for the worker.
|
||||
|
||||
Args:
|
||||
worker_id: Identifier of the worker.
|
||||
"""
|
||||
super().teardown_worker(worker_id)
|
||||
|
||||
if self.instrument_managed:
|
||||
self.uninstrument(worker_id)
|
||||
logger.info(f"[Worker {worker_id}] Instrumentation removed.")
|
||||
|
||||
@with_active_tracer_context
|
||||
@asynccontextmanager
|
||||
async def trace_context(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
*,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[Any]:
|
||||
"""Asynchronous implementation of the tracing context.
|
||||
|
||||
Args:
|
||||
name: Optional operation name.
|
||||
rollout_id: Optional rollout ID.
|
||||
attempt_id: Optional attempt ID.
|
||||
|
||||
Raises:
|
||||
ValueError: If store, rollout_id, and attempt_id are inconsistently provided.
|
||||
RuntimeError: If Weave is not installed or client is uninitialized.
|
||||
"""
|
||||
|
||||
if rollout_id is not None and attempt_id is not None:
|
||||
self._rollout_id = rollout_id
|
||||
self._attempt_id = attempt_id
|
||||
elif rollout_id is None and attempt_id is None:
|
||||
logger.info("No rollout_id or attempt_id provided. Skipping writing to store.")
|
||||
self._rollout_id = self._attempt_id = None
|
||||
else:
|
||||
raise ValueError("rollout_id and attempt_id must be either both provided or both None")
|
||||
|
||||
await self._init_trace_context()
|
||||
|
||||
weave_client = self._get_weave_client()
|
||||
|
||||
if weave_client.server is not self._server:
|
||||
logger.error(
|
||||
"Weave client is not using the correct trace server. You might have multiple WeaveTracer instances running in the same process. "
|
||||
f"Expected {self._server}, got {weave_client.server}"
|
||||
)
|
||||
|
||||
arg_op = name or weave_client.project
|
||||
arg_inputs: dict[str, str] = {}
|
||||
if rollout_id is not None:
|
||||
arg_inputs[LightningResourceAttributes.ROLLOUT_ID.value] = rollout_id
|
||||
if attempt_id is not None:
|
||||
arg_inputs[LightningResourceAttributes.ATTEMPT_ID.value] = attempt_id
|
||||
|
||||
try:
|
||||
# Create a new trace call object in Weave
|
||||
trace_call = weave_client.create_call( # pyright: ignore[reportUnknownMemberType]
|
||||
op=arg_op, inputs=arg_inputs
|
||||
)
|
||||
|
||||
try:
|
||||
yield trace_call
|
||||
# Finish trace even if no exception
|
||||
weave_client.finish_call(trace_call) # pyright: ignore[reportUnknownMemberType]
|
||||
except Exception as exc:
|
||||
# Finish trace and log any exception
|
||||
weave_client.finish_call(trace_call, exception=exc) # pyright: ignore[reportUnknownMemberType]
|
||||
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={exc}")
|
||||
raise
|
||||
|
||||
finally:
|
||||
try:
|
||||
weave_client.flush()
|
||||
# It's possible that the call end futures are from a dedicated Weave thread pool,
|
||||
await asyncio.gather(*[asyncio.wrap_future(future) for future in self._complete_call_futures])
|
||||
|
||||
finally:
|
||||
# Mandatory cleanup
|
||||
self._rollout_id = None
|
||||
self._attempt_id = None
|
||||
self._server.clear()
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
if timestamp is not None:
|
||||
logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
|
||||
weave_client = self._get_weave_client()
|
||||
trace_call = weave_client.create_call( # pyright: ignore[reportUnknownMemberType]
|
||||
op=name,
|
||||
attributes=attributes,
|
||||
inputs={},
|
||||
)
|
||||
# Immediately finish the call
|
||||
weave_client.finish_call(trace_call) # pyright: ignore[reportUnknownMemberType]
|
||||
# We don't wait for the call to be propagated to the server.
|
||||
start_time = trace_call.started_at.timestamp() if trace_call.started_at else None
|
||||
end_time = trace_call.ended_at.timestamp() if trace_call.ended_at else None
|
||||
trace_status = (
|
||||
TraceStatus(status_code="OK")
|
||||
if trace_call.exception is None
|
||||
else TraceStatus(status_code="ERROR", description=trace_call.exception)
|
||||
)
|
||||
return SpanCoreFields(
|
||||
name=name,
|
||||
attributes=flatten_attributes(trace_call.attributes or {}),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
status=trace_status,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def operation_context(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
) -> Iterator[SpanRecordingContext]:
|
||||
if start_time is not None:
|
||||
logger.warning("Weave doesn't support customizing the start time of a call. Timestamp is ignored.")
|
||||
if end_time is not None:
|
||||
logger.warning("Weave doesn't support customizing the end time of a call. Timestamp is ignored.")
|
||||
weave_client = self._get_weave_client()
|
||||
trace_call = weave_client.create_call( # pyright: ignore[reportUnknownMemberType]
|
||||
op=name,
|
||||
attributes=attributes,
|
||||
inputs={},
|
||||
)
|
||||
recording_context = WeaveSpanRecordingContext(trace_call)
|
||||
try:
|
||||
yield recording_context
|
||||
except Exception as exc:
|
||||
recording_context.record_exception(exc)
|
||||
raise
|
||||
finally:
|
||||
weave_client.finish_call(trace_call) # pyright: ignore[reportUnknownMemberType]
|
||||
|
||||
async def _init_trace_context(self) -> None:
|
||||
"""Initialize the trace context."""
|
||||
self._spans.clear()
|
||||
self._calls.clear()
|
||||
self._partial_call_futures.clear()
|
||||
self._complete_call_futures.clear()
|
||||
self._loop = weakref.ref(asyncio.get_running_loop())
|
||||
|
||||
def _get_weave_client(self) -> WeaveClient:
|
||||
"""Get the Weave client."""
|
||||
weave_client = weave.get_client()
|
||||
if not weave_client:
|
||||
raise RuntimeError("Weave client is not initialized. Call init_worker() first.")
|
||||
return weave_client
|
||||
|
||||
def _ensure_loop(self) -> tuple[asyncio.AbstractEventLoop, bool]:
|
||||
"""Returns a usable event loop and a boolean indicating whether it's the current running loop.
|
||||
|
||||
Prefer using the main loop if it's possible. Otherwise, use the current running loop.
|
||||
"""
|
||||
# Get the current running loop
|
||||
try:
|
||||
running_loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
running_loop = None
|
||||
|
||||
# Get the main loop, which can be a different loop
|
||||
if self._loop is not None:
|
||||
main_loop = self._loop()
|
||||
else:
|
||||
main_loop = None
|
||||
|
||||
if main_loop is not None:
|
||||
return main_loop, id(main_loop) == id(running_loop)
|
||||
elif running_loop is not None:
|
||||
return running_loop, True
|
||||
else:
|
||||
raise RuntimeError("No running event loop found. This should not happen.")
|
||||
|
||||
def get_last_trace(self) -> List[Span]:
|
||||
return self._spans
|
||||
|
||||
def partial_call_callback(self, request_content: Dict[str, Any]) -> None:
|
||||
call_id = request_content.get("id")
|
||||
if call_id is None:
|
||||
raise ValueError("Call ID is required even for partial calls")
|
||||
|
||||
if call_id in self._partial_call_futures:
|
||||
raise ValueError(f"Call {call_id} already has a start future")
|
||||
|
||||
# The callback must possibly be called from a dedicated Weave thread pool,
|
||||
# but it should be executed on the main event loop.
|
||||
try:
|
||||
loop, is_current_loop = self._ensure_loop()
|
||||
if is_current_loop:
|
||||
task = loop.create_task(self.partial_call_handler(request_content))
|
||||
else:
|
||||
# Schedule the task on the dedicated loop
|
||||
task = asyncio.run_coroutine_threadsafe(self.partial_call_handler(request_content), loop)
|
||||
self._partial_call_futures[call_id] = task
|
||||
except Exception as exc:
|
||||
logger.exception(f"Error creating call start task: {exc}", exc_info=True)
|
||||
|
||||
def complete_call_callback(self, call: tsi.CallSchema) -> None:
|
||||
try:
|
||||
loop, is_current_loop = self._ensure_loop()
|
||||
if is_current_loop:
|
||||
task = loop.create_task(self.complete_call_handler(call))
|
||||
else:
|
||||
# Schedule the task on the dedicated loop
|
||||
task = asyncio.run_coroutine_threadsafe(self.complete_call_handler(call), loop)
|
||||
self._complete_call_futures.append(task)
|
||||
except Exception as exc:
|
||||
logger.exception(f"Error creating call finish task: {exc}", exc_info=True)
|
||||
|
||||
async def _get_next_sequence_id(self) -> int:
|
||||
"""Get the next sequence ID for a span.
|
||||
|
||||
Use store to get the next sequence ID if available, otherwise use a default counter.
|
||||
"""
|
||||
if self._rollout_id and self._attempt_id and self._store:
|
||||
return await self._store.get_next_span_sequence_id(self._rollout_id, self._attempt_id)
|
||||
else:
|
||||
self._default_sequence_counter += 1
|
||||
return self._default_sequence_counter
|
||||
|
||||
async def partial_call_handler(self, request_content: Dict[str, Any]) -> int:
|
||||
"""Handler called when a Weave Call starts.
|
||||
|
||||
Args:
|
||||
request_content: The partial Weave Call object.
|
||||
|
||||
Returns:
|
||||
The sequence ID for the call.
|
||||
"""
|
||||
sequence_id = await self._get_next_sequence_id()
|
||||
return sequence_id
|
||||
|
||||
async def complete_call_handler(self, call: tsi.CallSchema) -> None:
|
||||
"""Handler called when a Weave Call finishes.
|
||||
|
||||
Converts the call (including nested children) into spans and stores them in LightningStore.
|
||||
"""
|
||||
# Make sure the corresponding call_start_future is complete
|
||||
if call.id in self._partial_call_futures:
|
||||
sequence_id = await asyncio.wrap_future(self._partial_call_futures[call.id])
|
||||
del self._partial_call_futures[call.id]
|
||||
else:
|
||||
# Fetch a new sequence ID as the call_start is somehow missing
|
||||
if call.id in self._calls:
|
||||
logger.warning(
|
||||
f"Call {call.id} is already in calls. The call is already completed. Overwriting the call."
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Call {call.id} has no start future. Fetching a new sequence ID.")
|
||||
sequence_id = await self._get_next_sequence_id()
|
||||
|
||||
self._calls[call.id] = call
|
||||
|
||||
span = await self.convert_call_to_span(call, self._rollout_id, self._attempt_id, sequence_id)
|
||||
self._spans.append(span)
|
||||
if self._store and self._rollout_id and self._attempt_id:
|
||||
try:
|
||||
await self._store.add_span(span)
|
||||
except Exception as exc:
|
||||
logger.exception(f"Error adding span to store: {exc}")
|
||||
|
||||
async def convert_call_to_span(
|
||||
self,
|
||||
call: tsi.CallSchema,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
sequence_id: Optional[int] = None,
|
||||
) -> Span:
|
||||
"""Convert a Weave Call (with nested children) into a Agent-lightning Span.
|
||||
|
||||
`rollout_id` and `attempt_id` are required to attach the spans to the store.
|
||||
|
||||
Args:
|
||||
call: The Weave Call object.
|
||||
rollout_id: Optional rollout ID to attach to spans.
|
||||
attempt_id: Optional attempt ID to attach to spans.
|
||||
sequence_id: Optional sequence ID to attach to spans.
|
||||
|
||||
Returns:
|
||||
List of converted spans.
|
||||
"""
|
||||
rollout_id = rollout_id or "rollout-dummy"
|
||||
attempt_id = attempt_id or "attempt-dummy"
|
||||
sequence_id = sequence_id or 0
|
||||
|
||||
start_ts: float = call.started_at.timestamp()
|
||||
end_ts: Optional[float] = call.ended_at.timestamp() if call.ended_at else None
|
||||
|
||||
if call.exception:
|
||||
status = TraceStatus(status_code="ERROR", description=call.exception)
|
||||
else:
|
||||
status = TraceStatus(status_code="OK")
|
||||
|
||||
attributes: Dict[str, Any] = {
|
||||
LightningSpanAttributes.OPERATION_NAME.value: call.op_name,
|
||||
# op_name can be possibly overridden by the attributes.
|
||||
**call.attributes,
|
||||
}
|
||||
if call.inputs:
|
||||
attributes[LightningSpanAttributes.OPERATION_INPUT.value] = call.inputs
|
||||
if call.output:
|
||||
attributes[LightningSpanAttributes.OPERATION_OUTPUT.value] = call.output
|
||||
if call.summary:
|
||||
# attributes can be possibly overridden by the summary.
|
||||
attributes.update(call.summary)
|
||||
if call.exception:
|
||||
attributes[exception_attributes.EXCEPTION_MESSAGE] = call.exception
|
||||
|
||||
sanitized_attributes = sanitize_attributes(flatten_attributes(attributes, expand_leaf_lists=False))
|
||||
|
||||
context = SpanContext(
|
||||
trace_id=call.trace_id,
|
||||
span_id=call.id,
|
||||
is_remote=False,
|
||||
trace_state={},
|
||||
)
|
||||
|
||||
# Get context for parent
|
||||
if call.parent_id:
|
||||
parent_call = self._calls.get(call.parent_id)
|
||||
if parent_call:
|
||||
parent_context = SpanContext(
|
||||
trace_id=parent_call.trace_id,
|
||||
span_id=parent_call.id,
|
||||
is_remote=False,
|
||||
trace_state={},
|
||||
)
|
||||
else:
|
||||
parent_context = None
|
||||
else:
|
||||
parent_context = None
|
||||
|
||||
# Build the Span object
|
||||
return Span(
|
||||
rollout_id=rollout_id,
|
||||
attempt_id=attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
trace_id=call.trace_id,
|
||||
span_id=call.id,
|
||||
parent_id=call.parent_id,
|
||||
name=op_name_to_func_name(call.op_name),
|
||||
status=status,
|
||||
attributes=sanitized_attributes,
|
||||
events=[], # Weave calls do not generate events
|
||||
links=[], # Weave calls do not generate links
|
||||
start_time=start_ts,
|
||||
end_time=end_ts,
|
||||
context=context,
|
||||
parent=parent_context,
|
||||
resource=OtelResource(
|
||||
attributes={
|
||||
LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
|
||||
LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
|
||||
LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id,
|
||||
LightningResourceAttributes.TRACER_NAME.value: "weave",
|
||||
},
|
||||
schema_url="",
|
||||
),
|
||||
)
|
||||
@@ -152,13 +152,6 @@ class Trainer(TrainerLegacy):
|
||||
# super().__init__() will call TrainerLegacy's initialization, which is not intended.
|
||||
self.worker_id: Optional[int] = None
|
||||
|
||||
if dev:
|
||||
warnings.warn(
|
||||
"Trainer(dev=True) is deprecated and will be removed in future versions. "
|
||||
"Please use Trainer.dev(...) instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self._dev = dev
|
||||
self.daemon = daemon
|
||||
self._client: AgentLightningClient | None = None # Will be initialized in fit or fit_v0
|
||||
@@ -220,6 +213,10 @@ class Trainer(TrainerLegacy):
|
||||
# We might be able to support a list of resources in future.
|
||||
self.initial_resources = initial_resources
|
||||
|
||||
# The active store for the current execution context
|
||||
self.store = self._make_store(store)
|
||||
self.runner = self._make_runner(runner)
|
||||
|
||||
self.port = port
|
||||
|
||||
self.strategy = self._make_strategy(
|
||||
@@ -227,11 +224,6 @@ class Trainer(TrainerLegacy):
|
||||
n_runners=self.n_runners,
|
||||
port=port,
|
||||
)
|
||||
|
||||
# The active store for the current execution context
|
||||
self.store = self._make_store(store, self.strategy)
|
||||
self.runner = self._make_runner(runner)
|
||||
|
||||
if hasattr(self.strategy, "n_runners"):
|
||||
strategy_runners = getattr(self.strategy, "n_runners")
|
||||
if isinstance(strategy_runners, int) and strategy_runners > 0:
|
||||
@@ -290,19 +282,13 @@ class Trainer(TrainerLegacy):
|
||||
type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
|
||||
)
|
||||
|
||||
def _make_store(self, store: ComponentSpec[LightningStore], strategy: ExecutionStrategy) -> LightningStore:
|
||||
"""Resolve the store implementation backing rollouts, attempts, spans, and resources.
|
||||
|
||||
By default, it's always a in-memory store. If using a client/server execution strategy,
|
||||
the in-memory store will be initialized in a thread-safe manner.
|
||||
"""
|
||||
is_client_server = isinstance(strategy, ClientServerExecutionStrategy)
|
||||
default_store_factory = lambda: InMemoryLightningStore(thread_safe=is_client_server)
|
||||
def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
|
||||
"""Resolve the store implementation backing rollouts, attempts, spans, and resources."""
|
||||
return build_component(
|
||||
store,
|
||||
expected_type=LightningStore,
|
||||
spec_name="store",
|
||||
default_factory=default_store_factory,
|
||||
default_factory=InMemoryLightningStore,
|
||||
invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
|
||||
type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
|
||||
)
|
||||
|
||||
@@ -28,7 +28,7 @@ from typing import (
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from .tracer import Span, SpanCoreFields
|
||||
from .tracer import Span
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.litagent import LitAgent
|
||||
@@ -53,7 +53,6 @@ __all__ = [
|
||||
"Rollout",
|
||||
"Attempt",
|
||||
"AttemptedRollout",
|
||||
"EnqueueRolloutRequest",
|
||||
"Hook",
|
||||
"Worker",
|
||||
"WorkerStatus",
|
||||
@@ -212,24 +211,6 @@ class AttemptedRollout(Rollout):
|
||||
return self
|
||||
|
||||
|
||||
class EnqueueRolloutRequest(BaseModel):
|
||||
"""Payload describing a rollout to be queued via [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout].
|
||||
|
||||
A subset of fields from [`Rollout`][agentlightning.Rollout] used for queuing new rollouts.
|
||||
"""
|
||||
|
||||
input: TaskInput
|
||||
"""Task input used to generate the rollout."""
|
||||
mode: Optional[RolloutMode] = None
|
||||
"""Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][agentlightning.RolloutMode]."""
|
||||
resources_id: Optional[str] = None
|
||||
"""Identifier of the resources required to execute the rollout."""
|
||||
config: Optional[RolloutConfig] = None
|
||||
"""Retry and timeout configuration associated with the rollout."""
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
"""Additional metadata attached to the rollout."""
|
||||
|
||||
|
||||
WorkerStatus = Literal["idle", "busy", "unknown"]
|
||||
|
||||
|
||||
@@ -307,7 +288,6 @@ RolloutRawResult = Union[
|
||||
float, # only final reward
|
||||
List[ReadableSpan], # constructed OTEL spans by user
|
||||
List[Span], # constructed Span objects by user
|
||||
List[SpanCoreFields], # constructed SpanCoreFields objects by user
|
||||
]
|
||||
"""Rollout result type.
|
||||
|
||||
|
||||
@@ -2,13 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
"""Data models that mirror OpenTelemetry spans for Agent Lightning."""
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, Union
|
||||
from typing import Any, Dict, List, Optional, Sequence, Union
|
||||
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
@@ -33,9 +31,6 @@ __all__ = [
|
||||
"SpanNames",
|
||||
"SpanAttributeNames",
|
||||
"SpanLike",
|
||||
"StatusCode",
|
||||
"SpanCoreFields",
|
||||
"SpanRecordingContext",
|
||||
]
|
||||
|
||||
|
||||
@@ -88,8 +83,6 @@ Attributes = Dict[str, AttributeValue]
|
||||
"""Mapping from attribute names to their values. Same as OpenTelemetry `Attributes` type."""
|
||||
TraceState = Dict[str, str]
|
||||
"""Mapping from trace state key to its value. Same as OpenTelemetry `TraceState` type."""
|
||||
StatusCode = Literal["UNSET", "OK", "ERROR"]
|
||||
"""The status code of the span."""
|
||||
|
||||
|
||||
class SpanContext(BaseModel):
|
||||
@@ -122,7 +115,7 @@ class SpanContext(BaseModel):
|
||||
class TraceStatus(BaseModel):
|
||||
"""Serializable variant of `opentelemetry.trace.Status`."""
|
||||
|
||||
status_code: StatusCode
|
||||
status_code: str
|
||||
"""The status code of the span. Same as OpenTelemetry `Status.status_code` type."""
|
||||
description: Optional[str] = None
|
||||
"""The description of the span. Same as OpenTelemetry `Status.description` type."""
|
||||
@@ -210,44 +203,6 @@ class OtelResource(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class SpanCoreFields(BaseModel):
|
||||
"""Core fields of a span. Used by span creators who don't care about the full span model.
|
||||
|
||||
If the spans are managed by some OTel tracer provider, it's not advised to create spans via this path.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"""The name of the span."""
|
||||
status: TraceStatus
|
||||
"""The status of the span."""
|
||||
attributes: Attributes
|
||||
"""The attributes of the span."""
|
||||
start_time: Optional[float]
|
||||
"""The start time of the span."""
|
||||
end_time: Optional[float]
|
||||
"""The end time of the span."""
|
||||
|
||||
|
||||
class SpanRecordingContext(Protocol):
|
||||
"""Context for recording operations on a span. It doesn't have to finalize the span; the caller will do it."""
|
||||
|
||||
def record_exception(self, exception: BaseException) -> None:
|
||||
"""Record an exception on the span."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def record_attributes(self, attributes: Attributes) -> None:
|
||||
"""Record attributes on the span."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def record_status(self, status_code: StatusCode, description: Optional[str] = None) -> None:
|
||||
"""Record the status of the span."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_recorded_span(self) -> SpanCoreFields:
|
||||
"""Get the recording of the span."""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class Span(BaseModel):
|
||||
"""Agent Lightning's canonical span model used for persistence and analytics.
|
||||
|
||||
@@ -385,7 +340,6 @@ class Span(BaseModel):
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
resource: Optional[OtelResource] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> "Span":
|
||||
"""Build a synthetic span from raw attributes.
|
||||
Different from the [`from_opentelemetry`][agentlightning.Span.from_opentelemetry] method,
|
||||
@@ -403,7 +357,6 @@ class Span(BaseModel):
|
||||
start_time: Span start timestamp in seconds.
|
||||
end_time: Span end timestamp in seconds.
|
||||
resource: Explicit resource information to attach to the span.
|
||||
status: Optional status of the span.
|
||||
|
||||
Returns:
|
||||
[`Span`][agentlightning.Span] populated with the provided attributes.
|
||||
@@ -431,7 +384,7 @@ class Span(BaseModel):
|
||||
name=name or AGL_VIRTUAL,
|
||||
resource=resource or OtelResource(attributes={}, schema_url=""),
|
||||
attributes=attributes,
|
||||
status=status or TraceStatus(status_code="OK"),
|
||||
status=TraceStatus(status_code="OK"),
|
||||
events=[],
|
||||
links=[],
|
||||
parent=(
|
||||
@@ -446,37 +399,6 @@ class Span(BaseModel):
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_core_fields(
|
||||
cls,
|
||||
core: SpanCoreFields,
|
||||
*,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
sequence_id: Optional[int] = None,
|
||||
) -> Span:
|
||||
"""Build a span from a core span.
|
||||
|
||||
Args:
|
||||
core: Core span to build from.
|
||||
rollout_id: Optional rollout identifier associated with the span.
|
||||
attempt_id: Optional attempt identifier associated with the span.
|
||||
sequence_id: Optional sequence number to preserve ordering.
|
||||
|
||||
Returns:
|
||||
[`Span`][agentlightning.Span] populated with the provided attributes.
|
||||
"""
|
||||
return cls.from_attributes(
|
||||
attributes=core.attributes,
|
||||
rollout_id=rollout_id,
|
||||
attempt_id=attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
name=core.name,
|
||||
start_time=core.start_time or time.time(),
|
||||
end_time=core.end_time,
|
||||
status=core.status,
|
||||
)
|
||||
|
||||
|
||||
class SpanNames(str, Enum):
|
||||
"""Enumerated span names recognised by Agent-lightning. Deprecated in favor of [semconv][agentlightning.semconv]."""
|
||||
|
||||
@@ -1,18 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import hashlib
|
||||
import uuid
|
||||
|
||||
__all__ = ["generate_id"]
|
||||
|
||||
|
||||
def generate_id(length: int) -> str:
|
||||
"""Generate a random ID of the given length.
|
||||
|
||||
Args:
|
||||
length: The length of the ID to generate.
|
||||
|
||||
Returns:
|
||||
A random ID of the given length.
|
||||
"""
|
||||
return hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:length]
|
||||
File diff suppressed because it is too large
Load Diff
+12
-152
@@ -2,25 +2,22 @@
|
||||
|
||||
"""Utilities shared for OpenTelemetry span (attributes) support."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import traceback
|
||||
from typing import Any, Dict, List, Sequence, Type, TypeVar, Union, cast
|
||||
from typing import Any, Dict, List, Sequence, Union, cast
|
||||
from warnings import filterwarnings
|
||||
|
||||
import opentelemetry.trace as trace_api
|
||||
from agentops.sdk.exporters import OTLPSpanExporter
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SpanProcessor, SynchronousMultiSpanProcessor, Tracer
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SynchronousMultiSpanProcessor, Tracer
|
||||
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
|
||||
from opentelemetry.sdk.util.instrumentation import InstrumentationInfo, InstrumentationScope
|
||||
from opentelemetry.semconv.attributes import exception_attributes
|
||||
from opentelemetry.trace import get_tracer_provider as otel_get_tracer_provider
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
|
||||
from agentlightning.semconv import LightningSpanAttributes, LinkAttributes, LinkPydanticModel
|
||||
from agentlightning.types import Attributes, AttributeValue, SpanLike
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otlp import LightningStoreOTLPExporter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -38,16 +35,8 @@ __all__ = [
|
||||
"filter_and_unflatten_attributes",
|
||||
"flatten_attributes",
|
||||
"unflatten_attributes",
|
||||
"sanitize_attribute_value",
|
||||
"sanitize_attributes",
|
||||
"sanitize_list_attribute_sanity",
|
||||
"check_attributes_sanity",
|
||||
"format_exception_attributes",
|
||||
]
|
||||
|
||||
T_SpanLike = TypeVar("T_SpanLike", bound=SpanLike)
|
||||
T_SpanProcessor = TypeVar("T_SpanProcessor", bound=SpanProcessor)
|
||||
|
||||
|
||||
def full_qualified_name(obj: type) -> str:
|
||||
if str(obj.__module__) == "builtins":
|
||||
@@ -123,25 +112,6 @@ def get_tracer_provider(inspect: bool = True) -> TracerProviderImpl:
|
||||
return tracer_provider
|
||||
|
||||
|
||||
def get_span_processors(
|
||||
tracer_provider: TracerProviderImpl, expected_type: Type[T_SpanProcessor]
|
||||
) -> List[T_SpanProcessor]:
|
||||
"""Get the span processors from the tracer provider.
|
||||
|
||||
Args:
|
||||
tracer_provider: The tracer provider to get the span processors from.
|
||||
expected_type: The type of the span processors to get.
|
||||
|
||||
Returns:
|
||||
A list of span processors of the expected type.
|
||||
"""
|
||||
processors: List[T_SpanProcessor] = []
|
||||
for processor in tracer_provider._active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
|
||||
if isinstance(processor, expected_type):
|
||||
processors.append(processor)
|
||||
return processors
|
||||
|
||||
|
||||
def get_tracer(use_active_span_processor: bool = True) -> trace_api.Tracer:
|
||||
"""Resolve the OpenTelemetry tracer configured for Agent Lightning.
|
||||
|
||||
@@ -196,7 +166,7 @@ def make_tag_attributes(tags: List[str]) -> Dict[str, Any]:
|
||||
["gen_ai.model:gpt-4", "reward.extrinsic"]
|
||||
```
|
||||
"""
|
||||
return flatten_attributes({LightningSpanAttributes.TAG.value: tags}, expand_leaf_lists=True)
|
||||
return flatten_attributes({LightningSpanAttributes.TAG.value: tags})
|
||||
|
||||
|
||||
def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
|
||||
@@ -226,10 +196,10 @@ def make_link_attributes(links: Dict[str, str]) -> Dict[str, Any]:
|
||||
if not isinstance(value, str): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
raise ValueError(f"Link value must be a string, got {type(value)} for key '{key}'")
|
||||
link_list.append({LinkAttributes.KEY_MATCH.value: key, LinkAttributes.VALUE_MATCH.value: value})
|
||||
return flatten_attributes({LightningSpanAttributes.LINK.value: link_list}, expand_leaf_lists=True)
|
||||
return flatten_attributes({LightningSpanAttributes.LINK.value: link_list})
|
||||
|
||||
|
||||
def query_linked_spans(spans: Sequence[T_SpanLike], links: List[LinkPydanticModel]) -> List[T_SpanLike]:
|
||||
def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]) -> List[SpanLike]:
|
||||
"""Query spans that are linked by the given link attributes.
|
||||
|
||||
Args:
|
||||
@@ -239,7 +209,7 @@ def query_linked_spans(spans: Sequence[T_SpanLike], links: List[LinkPydanticMode
|
||||
Returns:
|
||||
A list of spans that match the given link attributes.
|
||||
"""
|
||||
matched_spans: List[T_SpanLike] = []
|
||||
matched_spans: List[SpanLike] = []
|
||||
|
||||
for span in spans:
|
||||
span_attributes = span.attributes or {}
|
||||
@@ -324,9 +294,7 @@ def filter_and_unflatten_attributes(attributes: Dict[str, Any], prefix: str) ->
|
||||
return unflatten_attributes(stripped_attributes)
|
||||
|
||||
|
||||
def flatten_attributes(
|
||||
nested_data: Union[Dict[str, Any], List[Any]], *, expand_leaf_lists: bool = False
|
||||
) -> Dict[str, Any]:
|
||||
def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[str, Any]:
|
||||
"""Flatten a nested dictionary or list into a flat dictionary with dotted keys.
|
||||
|
||||
This function recursively traverses dictionaries and lists, producing a flat
|
||||
@@ -335,14 +303,12 @@ def flatten_attributes(
|
||||
|
||||
Example:
|
||||
|
||||
>>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}}, expand_leaf_lists=True)
|
||||
>>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}})
|
||||
{"a.b": 1, "a.c.0": 2, "a.c.1": 3}
|
||||
|
||||
Args:
|
||||
nested_data: A nested structure composed of dictionaries, lists, or primitive values.
|
||||
expand_leaf_lists: Whether to expand lists composed only of primitive values.
|
||||
When `False` (the default), lists of str/int/float/bool are treated as
|
||||
leaf values and stored without enumerating their indices.
|
||||
nested_data: A nested structure composed of dictionaries, lists, or
|
||||
primitive values.
|
||||
|
||||
Returns:
|
||||
A flat dictionary mapping dotted-string paths to primitive values.
|
||||
@@ -350,15 +316,6 @@ def flatten_attributes(
|
||||
|
||||
flat: Dict[str, Any] = {}
|
||||
|
||||
def _primitive_type(value: Any) -> Union[type[str], type[int], type[float], type[bool]]:
|
||||
if isinstance(value, bool):
|
||||
return bool
|
||||
if isinstance(value, int):
|
||||
return int
|
||||
if isinstance(value, float):
|
||||
return float
|
||||
return str
|
||||
|
||||
def _walk(value: Any, prefix: str = "") -> None:
|
||||
if isinstance(value, dict):
|
||||
for k, v in cast(Dict[Any, Any], value).items():
|
||||
@@ -369,22 +326,7 @@ def flatten_attributes(
|
||||
new_prefix = f"{prefix}.{k}" if prefix else k
|
||||
_walk(v, new_prefix)
|
||||
elif isinstance(value, list):
|
||||
maybe_list = cast(List[Any], value)
|
||||
is_leaf_candidate = bool(maybe_list) and all(
|
||||
isinstance(item, (str, int, float, bool)) for item in maybe_list
|
||||
)
|
||||
if not expand_leaf_lists and is_leaf_candidate and prefix:
|
||||
primitive_types = {_primitive_type(item) for item in maybe_list}
|
||||
if len(primitive_types) == 1:
|
||||
flat[prefix] = maybe_list
|
||||
return
|
||||
logger.warning(
|
||||
"List attribute '%s' contains mixed primitive types %s; expanding indexed keys instead.",
|
||||
prefix,
|
||||
primitive_types,
|
||||
)
|
||||
|
||||
for idx, item in enumerate(maybe_list):
|
||||
for idx, item in enumerate(cast(List[Any], value)):
|
||||
new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
|
||||
_walk(item, new_prefix)
|
||||
else:
|
||||
@@ -457,85 +399,3 @@ def unflatten_attributes(flat_data: Dict[str, Any]) -> Union[Dict[str, Any], Lis
|
||||
return node
|
||||
|
||||
return convert(root)
|
||||
|
||||
|
||||
def sanitize_attribute_value(object: Any, force: bool = True) -> AttributeValue:
|
||||
"""Sanitize an attribute value to be a valid OpenTelemetry attribute value."""
|
||||
if isinstance(object, (str, int, float, bool)):
|
||||
return object
|
||||
|
||||
if isinstance(object, list):
|
||||
try:
|
||||
return sanitize_list_attribute_sanity(cast(List[Any], object))
|
||||
except ValueError as exc:
|
||||
logger.warning(f"Failed to sanitize list attribute. Fallback to JSON serialization: {exc}")
|
||||
|
||||
try:
|
||||
# This include null, dict, etc.
|
||||
serialized = json.dumps(object, default=str if force else None)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"Object must be JSON serializable, got: {type(cast(Any, object))}.") from exc
|
||||
return serialized
|
||||
|
||||
|
||||
def sanitize_attributes(attributes: Dict[str, Any], force: bool = True) -> Attributes:
|
||||
"""Sanitize a dictionary of attributes to be a valid OpenTelemetry attributes.
|
||||
|
||||
Args:
|
||||
attributes: A dictionary of attributes to sanitize.
|
||||
force: Whether to force sanitization even when the value is not JSON serializable.
|
||||
"""
|
||||
result: Attributes = {}
|
||||
for k, v in attributes.items():
|
||||
try:
|
||||
result[k] = sanitize_attribute_value(v, force=force)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Failed to sanitize attribute '{k}': {exc}") from exc
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_list_attribute_sanity(maybe_list: List[Any]) -> AttributeValue:
|
||||
"""Try to sanitize a list of attributes to be a valid OpenTelemetry attribute value.
|
||||
|
||||
Raise error if the list contains multiple types of primitive values.
|
||||
"""
|
||||
if all(isinstance(item, str) for item in maybe_list):
|
||||
return list[str](maybe_list)
|
||||
if all(isinstance(item, bool) for item in maybe_list):
|
||||
return list[bool](maybe_list)
|
||||
if all(isinstance(item, (int, bool)) for item in maybe_list):
|
||||
return [int(item) for item in maybe_list]
|
||||
if all(isinstance(item, (float, int, bool)) for item in maybe_list):
|
||||
return [float(item) for item in maybe_list]
|
||||
|
||||
list_types: List[Any] = [type(item) for item in maybe_list]
|
||||
raise ValueError(f"List must contain only one type of primitive values, got: {set(list_types)}.")
|
||||
|
||||
|
||||
def check_attributes_sanity(attributes: Dict[Any, Any]) -> None:
|
||||
"""Check if a dictionary of attributes is a valid OpenTelemetry attributes."""
|
||||
for k, v in attributes.items():
|
||||
if not isinstance(k, str):
|
||||
raise ValueError(f"Attribute key must be a string, got {type(k)} for key '{k}'")
|
||||
if isinstance(v, list):
|
||||
try:
|
||||
sanitize_list_attribute_sanity(cast(List[Any], v))
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Failed to sanitize list attribute '{k}': {exc}") from exc
|
||||
elif not isinstance(v, (str, int, float, bool)):
|
||||
raise ValueError(
|
||||
f"Attribute value must be a string, int, float, bool, or list of these, got {type(v)} for value '{v}'"
|
||||
)
|
||||
|
||||
|
||||
def format_exception_attributes(exception: BaseException) -> Attributes:
|
||||
"""Format an exception into a dictionary of attributes."""
|
||||
stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
|
||||
span_attributes: Attributes = {
|
||||
exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
|
||||
exception_attributes.EXCEPTION_MESSAGE: str(exception),
|
||||
exception_attributes.EXCEPTION_ESCAPED: True,
|
||||
}
|
||||
if stacktrace.strip():
|
||||
span_attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
|
||||
return span_attributes
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import logging
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
|
||||
|
||||
from fastapi import Request, Response
|
||||
from google.protobuf import json_format
|
||||
@@ -39,7 +39,6 @@ from agentlightning.types.tracer import (
|
||||
OtelResource,
|
||||
Span,
|
||||
SpanContext,
|
||||
StatusCode,
|
||||
TraceStatus,
|
||||
convert_timestamp,
|
||||
)
|
||||
@@ -414,7 +413,7 @@ def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
|
||||
return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
|
||||
|
||||
|
||||
_STATUS_CODE_MAP: Mapping[ProtoStatus.StatusCode.ValueType, StatusCode] = {
|
||||
_STATUS_CODE_MAP = {
|
||||
ProtoStatus.STATUS_CODE_UNSET: "UNSET",
|
||||
ProtoStatus.STATUS_CODE_OK: "OK",
|
||||
ProtoStatus.STATUS_CODE_ERROR: "ERROR",
|
||||
|
||||
@@ -940,9 +940,9 @@ class PythonServerLauncher:
|
||||
), # Allow half the timeout for graceful shutdown
|
||||
}
|
||||
if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
|
||||
from agentlightning.utils.metrics import shutdown_metrics
|
||||
from prometheus_client import multiprocess
|
||||
|
||||
options["child_exit"] = shutdown_metrics # type: ignore
|
||||
options["child_exit"] = lambda server, worker: multiprocess.mark_process_dead(worker.pid) # type: ignore
|
||||
|
||||
self._gunicorn_app = GunicornApp(self.app, options)
|
||||
|
||||
|
||||
@@ -13,20 +13,12 @@ from gpustat import GPUStat, GPUStatCollection
|
||||
|
||||
|
||||
def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
|
||||
"""Capture a snapshot of the system's hardware and software information.
|
||||
|
||||
Args:
|
||||
include_gpu: Whether to include GPU information.
|
||||
|
||||
Returns:
|
||||
A dictionary containing the system's hardware and software information.
|
||||
"""
|
||||
# CPU
|
||||
cpu = {
|
||||
"cpu_name": platform.processor(),
|
||||
"cpu_cores": psutil.cpu_count(logical=False),
|
||||
"cpu_threads": psutil.cpu_count(logical=True),
|
||||
"cpu_usage_pct": psutil.cpu_percent(0.0),
|
||||
"cpu_usage_pct": psutil.cpu_percent(0.05),
|
||||
}
|
||||
|
||||
# Memory
|
||||
@@ -45,21 +37,20 @@ def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
|
||||
"disk_pct": du.percent,
|
||||
}
|
||||
|
||||
# GPU (only query if explicitly requested)
|
||||
# GPU
|
||||
gpus: List[Dict[str, Any]] = []
|
||||
if include_gpu:
|
||||
with suppress(Exception):
|
||||
for g in GPUStatCollection.new_query().gpus: # type: ignore
|
||||
g = cast(GPUStat, g)
|
||||
gpus.append(
|
||||
{
|
||||
"gpu": g.name, # type: ignore
|
||||
"util_pct": g.utilization,
|
||||
"mem_used_mb": g.memory_used,
|
||||
"mem_total_mb": g.memory_total,
|
||||
"temp_c": g.temperature,
|
||||
}
|
||||
)
|
||||
with suppress(Exception):
|
||||
for g in GPUStatCollection.new_query().gpus: # type: ignore
|
||||
g = cast(GPUStat, g)
|
||||
gpus.append(
|
||||
{
|
||||
"gpu": g.name, # type: ignore
|
||||
"util_pct": g.utilization,
|
||||
"mem_used_mb": g.memory_used,
|
||||
"mem_total_mb": g.memory_total,
|
||||
"temp_c": g.temperature,
|
||||
}
|
||||
)
|
||||
|
||||
# Network
|
||||
net = psutil.net_io_counters()
|
||||
|
||||
@@ -8,12 +8,6 @@ defaults:
|
||||
|
||||
agentlightning:
|
||||
port: 9999
|
||||
trace_aggregator:
|
||||
level: transition # transition or trajectory, docs refer to https://agent-lightning.github.io/posts/trajectory_level_aggregation/
|
||||
trajectory_max_prompt_length: 2048 # supported in trajectory level aggregation, suggest to set as maximum length for the prompt in first turn
|
||||
trajectory_max_response_length: 8192 # supported in trajectory level aggregation, suggest to set as maximum length for the cumulative agent responses in the full trajectory, i.e., n_turns * (max_response_length + max_prompt_length)
|
||||
debug: False # supported in trajectory level aggregation, enable to diagnose trace merging failures
|
||||
unmatch_log_dir: ./unmatch_cases # supported in trajectory level aggregation with debug=True, directory to store logs of unmatched cases
|
||||
|
||||
data:
|
||||
filter_overlong_prompts: false
|
||||
|
||||
+56
-410
@@ -2,7 +2,6 @@
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import threading
|
||||
@@ -10,7 +9,7 @@ import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple, cast
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
@@ -23,7 +22,7 @@ from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLeg
|
||||
from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
|
||||
from agentlightning.llm_proxy import LLMProxy, ModelConfig
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import EnqueueRolloutRequest, Rollout, RolloutConfig, Task
|
||||
from agentlightning.types import Rollout, RolloutConfig, Task
|
||||
|
||||
__all__ = [
|
||||
"AgentModeDaemon",
|
||||
@@ -32,85 +31,6 @@ __all__ = [
|
||||
]
|
||||
|
||||
|
||||
def ids_startswith(
|
||||
full_ids: List[int], prefix_ids: List[int], tokenizer: Any, debug: bool = False
|
||||
) -> Tuple[bool, Tuple[bool, bool, bool]]:
|
||||
is_prefix: bool
|
||||
template_mismatch, retoken_mismatch, others_mismatch = False, False, False
|
||||
if full_ids[: len(prefix_ids)] == prefix_ids:
|
||||
is_prefix = True
|
||||
return True, (template_mismatch, retoken_mismatch, others_mismatch)
|
||||
else:
|
||||
is_prefix = False
|
||||
|
||||
if not debug:
|
||||
return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
|
||||
|
||||
def _special_token_sequence(ids: List[int]) -> List[int]:
|
||||
return [id for id in ids if id in tokenizer.all_special_ids]
|
||||
|
||||
def _none_special_token_sequence(ids: List[int]) -> List[int]:
|
||||
return [id for id in ids if id not in tokenizer.all_special_ids]
|
||||
|
||||
# First, handle special tokens
|
||||
full_special_ids = _special_token_sequence(full_ids)
|
||||
prefix_special_ids = _special_token_sequence(prefix_ids)
|
||||
if sum(1 for a, b in zip(full_special_ids, prefix_special_ids) if a != b) > 0:
|
||||
template_mismatch = True
|
||||
|
||||
# Next, handle string content
|
||||
full_content_ids = _none_special_token_sequence(full_ids)
|
||||
prefix_content_ids = _none_special_token_sequence(prefix_ids)
|
||||
full_string = tokenizer.decode(full_ids, skip_special_tokens=True)
|
||||
prefix_string = tokenizer.decode(prefix_ids, skip_special_tokens=True)
|
||||
if full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and full_string.startswith(prefix_string):
|
||||
retoken_mismatch = True
|
||||
elif full_content_ids[: len(prefix_content_ids)] != prefix_content_ids and not full_string.startswith(
|
||||
prefix_string
|
||||
):
|
||||
others_mismatch = True
|
||||
return is_prefix, (template_mismatch, retoken_mismatch, others_mismatch)
|
||||
|
||||
|
||||
def log_mismatch_detail(
|
||||
diagnostic: Tuple[bool, bool, bool],
|
||||
full_ids: List[int],
|
||||
prefix_ids: List[int],
|
||||
global_steps: int,
|
||||
rollout_id: str,
|
||||
turn_id: int,
|
||||
log_dir: str | None = None,
|
||||
):
|
||||
if log_dir is None:
|
||||
return
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
template_mismatch, retoken_mismatch, others_mismatch = diagnostic
|
||||
if template_mismatch:
|
||||
with open(os.path.join(log_dir, "template_mismatch.log"), "a+") as f:
|
||||
print(
|
||||
"-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
|
||||
file=f,
|
||||
)
|
||||
print(full_ids, file=f)
|
||||
print(prefix_ids, file=f)
|
||||
if retoken_mismatch:
|
||||
with open(os.path.join(log_dir, "retoken_mismatch.log"), "a+") as f:
|
||||
print(
|
||||
"-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
|
||||
file=f,
|
||||
)
|
||||
print(full_ids, file=f)
|
||||
print(prefix_ids, file=f)
|
||||
if others_mismatch:
|
||||
with open(os.path.join(log_dir, "others_mismatch.log"), "a+") as f:
|
||||
print(
|
||||
"-" * 10 + f" Global Steps: {global_steps}, Rollout ID: {rollout_id}, Turn ID: {turn_id} " + "-" * 10,
|
||||
file=f,
|
||||
)
|
||||
print(full_ids, file=f)
|
||||
print(prefix_ids, file=f)
|
||||
|
||||
|
||||
def get_left_padded_ids_and_attention_mask(
|
||||
ids: List[int], max_length: int, pad_token_id: int
|
||||
) -> Tuple[List[int], List[int]]:
|
||||
@@ -224,9 +144,6 @@ class AgentModeDaemon:
|
||||
llm_proxy: LLMProxy | None = None,
|
||||
store: LightningStore | None = None,
|
||||
adapter: TraceToTripletBase | None = None,
|
||||
processor: Any = None,
|
||||
image_base_dir: Optional[str] = None,
|
||||
trace_aggregator: Dict[str, Any] = {"level": "transition"},
|
||||
):
|
||||
self.mode = mode
|
||||
self.llm_timeout_seconds = llm_timeout_seconds
|
||||
@@ -266,13 +183,7 @@ class AgentModeDaemon:
|
||||
self.mini_batch_size = mini_batch_size
|
||||
self.pad_token_id = pad_token_id
|
||||
self.tokenizer = tokenizer
|
||||
self.processor = processor
|
||||
self.reward_fillna_value = reward_fillna_value
|
||||
self.image_base_dir = image_base_dir
|
||||
self.trace_aggregator = trace_aggregator
|
||||
|
||||
# Check if model requires multimodal position_ids (e.g., Qwen2-VL)
|
||||
self._use_mrope = self._is_mrope_model()
|
||||
|
||||
# Internal State
|
||||
self.backend_llm_server_addresses: List[str] = []
|
||||
@@ -291,75 +202,6 @@ class AgentModeDaemon:
|
||||
loop.run_forever()
|
||||
loop.close()
|
||||
|
||||
# Multimodal utilities for M-RoPE position embeddings
|
||||
|
||||
def _is_mrope_model(self) -> bool:
|
||||
"""Check if processor requires M-RoPE position embeddings."""
|
||||
if self.processor is None or not hasattr(self.processor, "image_processor"):
|
||||
return False
|
||||
name = self.processor.image_processor.__class__.__name__
|
||||
return "Qwen2VLImageProcessor" in name or "Qwen3VLImageProcessor" in name
|
||||
|
||||
def _resolve_image_path(self, path: str) -> str:
|
||||
"""Resolve relative image path with base directory."""
|
||||
import os
|
||||
|
||||
if os.path.isabs(path):
|
||||
return path
|
||||
if self.image_base_dir is None:
|
||||
raise ValueError(f"Relative path '{path}' requires 'image_base_dir' to be set.")
|
||||
return os.path.join(self.image_base_dir, path)
|
||||
|
||||
def _get_image_grid_thw(self, image_urls: List[str]) -> Optional[torch.Tensor]:
|
||||
"""Compute image_grid_thw from image URLs for M-RoPE computation.
|
||||
|
||||
Args:
|
||||
image_urls: List of image URLs extracted from triplet prompt payload.
|
||||
URLs can be http(s):// URLs or file:// URIs, or data: URIs.
|
||||
"""
|
||||
from PIL import Image
|
||||
from verl.utils.dataset.vision_utils import process_image # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
if self.processor is None or not image_urls:
|
||||
return None
|
||||
|
||||
def to_image_uri(url: str) -> str:
|
||||
# Already a proper URI (http, https, file, data)
|
||||
if url.startswith(("http://", "https://", "file://", "data:")):
|
||||
return url
|
||||
# Treat as a file path that needs resolution
|
||||
resolved = self._resolve_image_path(url)
|
||||
return f"file://{resolved}"
|
||||
|
||||
images: List[Image.Image] = [process_image({"image": to_image_uri(url)}) for url in image_urls]
|
||||
model_inputs = self.processor(text=["dummy"], images=images, return_tensors="pt")
|
||||
return model_inputs.get("image_grid_thw")
|
||||
|
||||
def _compute_mrope_position_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
image_grid_thw: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Compute 4D position_ids for M-RoPE models."""
|
||||
from typing import Callable
|
||||
|
||||
get_rope_index: Callable[..., torch.Tensor]
|
||||
if "Qwen3VL" in self.processor.__class__.__name__:
|
||||
from verl.models.transformers.qwen3_vl import get_rope_index # pyright: ignore[reportUnknownVariableType]
|
||||
else:
|
||||
from verl.models.transformers.qwen2_vl import get_rope_index # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
vision_pos = get_rope_index(
|
||||
self.processor, input_ids=input_ids, image_grid_thw=image_grid_thw, attention_mask=attention_mask
|
||||
)
|
||||
|
||||
valid_mask = attention_mask.bool()
|
||||
text_pos = torch.zeros((1, len(input_ids)), dtype=torch.long, device=input_ids.device)
|
||||
text_pos[0, valid_mask] = torch.arange(valid_mask.sum().item(), device=input_ids.device)
|
||||
|
||||
return torch.cat([text_pos, vision_pos], dim=0)
|
||||
|
||||
def _start_proxy_server_v0(self):
|
||||
"""
|
||||
Initializes and runs a Flask-based proxy server in a separate thread.
|
||||
@@ -535,57 +377,42 @@ class AgentModeDaemon:
|
||||
num_samples = len(data[keys[0]])
|
||||
rollouts_per_sample = self.train_rollout_n if is_train else 1
|
||||
|
||||
enqueue_rollout_requests: List[EnqueueRolloutRequest] = []
|
||||
data_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
for i in range(num_samples):
|
||||
data_id = str(uuid.uuid4())
|
||||
original_sample = {key: data[key][i] for key in keys}
|
||||
original_sample["data_id"] = data_id
|
||||
data_id_to_original_sample[data_id] = original_sample
|
||||
|
||||
# For training, each sample is rolled out multiple times
|
||||
# Data ID is different from Rollout ID, as one data can have multiple rollouts.
|
||||
for _ in range(rollouts_per_sample):
|
||||
task_metadata = {"data_id": data_id, "is_train": is_train}
|
||||
|
||||
# Data ID is different from Rollout ID, as one data can have multiple rollouts.
|
||||
if self.mode == "v0":
|
||||
# Queue immediately
|
||||
rollout_id = await self.server.queue_task(
|
||||
sample=_to_native(original_sample),
|
||||
mode="train" if is_train else "val",
|
||||
resources_id=resources_id,
|
||||
metadata=task_metadata,
|
||||
)
|
||||
|
||||
# Store original sample data to reconstruct batch information later
|
||||
self._task_id_to_original_sample[rollout_id] = original_sample
|
||||
self._total_tasks_queued += 1
|
||||
else:
|
||||
# Collect tasks to enqueue in batch and queue them later
|
||||
enqueue_rollout_requests.append(
|
||||
EnqueueRolloutRequest(
|
||||
input=_to_native(original_sample),
|
||||
mode="train" if is_train else "val",
|
||||
resources_id=resources_id,
|
||||
config=RolloutConfig(
|
||||
unresponsive_seconds=self.llm_timeout_seconds,
|
||||
timeout_seconds=self.llm_timeout_seconds,
|
||||
),
|
||||
metadata=task_metadata,
|
||||
)
|
||||
rollout = await self.store.enqueue_rollout(
|
||||
input=_to_native(original_sample),
|
||||
mode="train" if is_train else "val",
|
||||
resources_id=resources_id,
|
||||
metadata=task_metadata,
|
||||
)
|
||||
await self.store.update_rollout(
|
||||
rollout_id=rollout.rollout_id,
|
||||
config=RolloutConfig(
|
||||
unresponsive_seconds=self.llm_timeout_seconds,
|
||||
timeout_seconds=self.llm_timeout_seconds,
|
||||
),
|
||||
)
|
||||
rollout_id = rollout.rollout_id
|
||||
|
||||
if self.mode == "v1":
|
||||
# Enqueue all the tasks in a single batch
|
||||
rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
|
||||
self._task_id_to_original_sample.update(
|
||||
{
|
||||
# Recover the original data and store it for later use.
|
||||
rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
|
||||
for rollout in rollouts
|
||||
}
|
||||
)
|
||||
self._total_tasks_queued += len(rollouts)
|
||||
# Store original sample data to reconstruct batch information later
|
||||
self._task_id_to_original_sample[rollout_id] = original_sample
|
||||
self._total_tasks_queued += 1
|
||||
|
||||
def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
|
||||
"""Synchronous wrapper for setting up data and server resources."""
|
||||
@@ -602,7 +429,7 @@ class AgentModeDaemon:
|
||||
raise RuntimeError("Internal loop is not running.")
|
||||
future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
|
||||
try:
|
||||
future.result(timeout=300) # Wait for completion with a timeout
|
||||
future.result(timeout=60) # Wait for completion with a timeout
|
||||
except Exception as e:
|
||||
print(f"Failed to set up data on server: {e}")
|
||||
raise
|
||||
@@ -804,9 +631,7 @@ class AgentModeDaemon:
|
||||
)
|
||||
return metric_dict
|
||||
|
||||
def get_train_data_batch(
|
||||
self, max_prompt_length: int, max_response_length: int, device: torch.device, global_steps: int
|
||||
):
|
||||
def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
|
||||
"""
|
||||
Processes completed rollouts to generate a training data batch.
|
||||
|
||||
@@ -832,14 +657,10 @@ class AgentModeDaemon:
|
||||
continue
|
||||
|
||||
# The client should report triplets that contain prompt_ids and response_ids.
|
||||
# Example triplet.prompt: {"token_ids": [...], "image_urls": [...]}
|
||||
# Example triplet.prompt: {"token_ids": [...]}
|
||||
# Example triplet.response: {"token_ids": [...]}
|
||||
trace_list = [
|
||||
{
|
||||
"prompt_ids": t.prompt.get("token_ids", []),
|
||||
"response_ids": t.response.get("token_ids", []),
|
||||
"image_urls": t.prompt.get("image_urls", []),
|
||||
}
|
||||
{"prompt_ids": t.prompt.get("token_ids", []), "response_ids": t.response.get("token_ids", [])}
|
||||
for t in rollout.triplets
|
||||
]
|
||||
info = {
|
||||
@@ -869,204 +690,60 @@ class AgentModeDaemon:
|
||||
rollout_id_list: List[str] = []
|
||||
turn_index_list: List[int] = []
|
||||
is_drop_list: List[bool] = []
|
||||
image_grid_thw_list: List[Optional[torch.Tensor]] = [] # For Qwen2-VL mrope
|
||||
n_trunc_sample_because_of_response = 0
|
||||
|
||||
if self.trace_aggregator.get("level", "transition") == "transition":
|
||||
for rollout_id, sample_info in finished_id_to_sample_info.items():
|
||||
for turn_index, trace in enumerate(sample_info["trace_list"]):
|
||||
for rollout_id, sample_info in finished_id_to_sample_info.items():
|
||||
for turn_index, trace in enumerate(sample_info["trace_list"]):
|
||||
|
||||
reward_list.append(sample_info["reward"])
|
||||
prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
|
||||
reward_list.append(sample_info["reward"])
|
||||
prompt_ids, response_ids = trace["prompt_ids"], trace["response_ids"]
|
||||
|
||||
# Mark samples with prompts exceeding max_prompt_length to be dropped later
|
||||
if len(prompt_ids) > max_prompt_length:
|
||||
prompt_ids = prompt_ids[:max_prompt_length]
|
||||
is_drop_list.append(True)
|
||||
else:
|
||||
is_drop_list.append(False)
|
||||
# Mark samples with prompts exceeding max_prompt_length to be dropped later
|
||||
if len(prompt_ids) > max_prompt_length:
|
||||
prompt_ids = prompt_ids[:max_prompt_length]
|
||||
is_drop_list.append(True)
|
||||
else:
|
||||
is_drop_list.append(False)
|
||||
|
||||
# Truncate responses that exceed max_response_length
|
||||
if len(response_ids) > max_response_length:
|
||||
response_ids = response_ids[:max_response_length]
|
||||
n_trunc_sample_because_of_response += 1
|
||||
# Truncate responses that exceed max_response_length
|
||||
if len(response_ids) > max_response_length:
|
||||
response_ids = response_ids[:max_response_length]
|
||||
n_trunc_sample_because_of_response += 1
|
||||
|
||||
# Pad prompts to the left and responses to the right
|
||||
one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
|
||||
prompt_ids, max_prompt_length, self.pad_token_id
|
||||
)
|
||||
one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
|
||||
response_ids, max_response_length, self.pad_token_id
|
||||
)
|
||||
# Pad prompts to the left and responses to the right
|
||||
one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
|
||||
prompt_ids, max_prompt_length, self.pad_token_id
|
||||
)
|
||||
one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
|
||||
response_ids, max_response_length, self.pad_token_id
|
||||
)
|
||||
|
||||
input_ids_list.append(one_input_ids)
|
||||
input_attention_mask_list.append(one_input_attention_mask)
|
||||
response_ids_list.append(one_response_ids)
|
||||
response_attention_mask_list.append(one_response_attention_mask)
|
||||
data_id_list.append(sample_info["data_id"])
|
||||
rollout_id_list.append(rollout_id)
|
||||
turn_index_list.append(turn_index)
|
||||
|
||||
# Compute image_grid_thw for this triplet using image_urls from prompt
|
||||
if self._use_mrope:
|
||||
image_urls = trace.get("image_urls", [])
|
||||
image_grid_thw_list.append(self._get_image_grid_thw(image_urls))
|
||||
|
||||
elif self.trace_aggregator.get("level", "transition") == "trajectory":
|
||||
assert not self._use_mrope, "M-RoPE is not supported in trajectory level yet."
|
||||
|
||||
response_mask_list: List[List[int]] = []
|
||||
unmerged_count: int = 0
|
||||
template_mismatch_count, retoken_mismatch_count, others_mismatch_count = 0, 0, 0
|
||||
response_per_turn_list: List[int] = []
|
||||
|
||||
for rollout_id, sample_info in finished_id_to_sample_info.items():
|
||||
merged_trace_idx: List[List[int]] = []
|
||||
|
||||
# Identify which turns can be merged based on token ids prefix matching
|
||||
current_merged_trace_idx: List[int] = []
|
||||
current_context: List[int] = []
|
||||
for turn_index, trace in enumerate(sample_info["trace_list"]):
|
||||
response_per_turn_list.append(len(trace["response_ids"]))
|
||||
is_prefix, diagnostic = ids_startswith(
|
||||
trace["prompt_ids"] + trace["response_ids"],
|
||||
current_context,
|
||||
self.tokenizer,
|
||||
self.trace_aggregator.get("debug", False),
|
||||
)
|
||||
if not is_prefix and self.trace_aggregator.get("debug", False) == True:
|
||||
template_mismatch_count += diagnostic[0]
|
||||
retoken_mismatch_count += diagnostic[1]
|
||||
others_mismatch_count += diagnostic[2]
|
||||
log_mismatch_detail(
|
||||
diagnostic,
|
||||
trace["prompt_ids"] + trace["response_ids"],
|
||||
current_context,
|
||||
global_steps,
|
||||
rollout_id,
|
||||
turn_index,
|
||||
self.trace_aggregator.get("unmatch_log_dir", None),
|
||||
)
|
||||
|
||||
if is_prefix:
|
||||
current_context = trace["prompt_ids"] + trace["response_ids"]
|
||||
current_merged_trace_idx.append(turn_index)
|
||||
else:
|
||||
merged_trace_idx.append(current_merged_trace_idx)
|
||||
current_merged_trace_idx = [turn_index]
|
||||
current_context = trace["prompt_ids"] + trace["response_ids"]
|
||||
|
||||
if current_merged_trace_idx not in merged_trace_idx:
|
||||
merged_trace_idx.append(current_merged_trace_idx)
|
||||
|
||||
if len(merged_trace_idx) > 1:
|
||||
unmerged_count += 1
|
||||
|
||||
# Merge all trace segments in merged_trace_idx into training samples
|
||||
for current_merged_trace_idx in merged_trace_idx:
|
||||
prompt_ids = sample_info["trace_list"][current_merged_trace_idx[0]]["prompt_ids"]
|
||||
|
||||
# if the merged_trace_idx doesn't start with the beginning of the prompt_ids, we need to adjust it
|
||||
if current_merged_trace_idx[0] > 0 and len(prompt_ids) > max_prompt_length:
|
||||
response_ids = prompt_ids[max_prompt_length:]
|
||||
prompt_ids = prompt_ids[:max_prompt_length]
|
||||
response_mask = [1] * len(response_ids)
|
||||
else:
|
||||
response_ids = []
|
||||
response_mask = []
|
||||
|
||||
prompt_length = len(prompt_ids)
|
||||
response_ids += sample_info["trace_list"][current_merged_trace_idx[0]]["response_ids"]
|
||||
response_mask += [1] * len(response_ids)
|
||||
for turn_index in current_merged_trace_idx[1:]:
|
||||
trace = sample_info["trace_list"][turn_index]
|
||||
new_prompt_length = len(trace["prompt_ids"]) - len(response_ids) - prompt_length
|
||||
response_ids += trace["prompt_ids"][-new_prompt_length:]
|
||||
response_ids += trace["response_ids"]
|
||||
response_mask += [0] * new_prompt_length
|
||||
response_mask += [1] * len(trace["response_ids"])
|
||||
|
||||
reward_list.append(sample_info["reward"])
|
||||
|
||||
# Mark samples with prompts exceeding max_prompt_length to be dropped later
|
||||
if len(prompt_ids) > max_prompt_length:
|
||||
prompt_ids = prompt_ids[:max_prompt_length]
|
||||
is_drop_list.append(True)
|
||||
else:
|
||||
is_drop_list.append(False)
|
||||
|
||||
# Truncate responses that exceed max_response_length
|
||||
if len(response_ids) > max_response_length:
|
||||
response_ids = response_ids[:max_response_length]
|
||||
response_mask = response_mask[:max_response_length]
|
||||
n_trunc_sample_because_of_response += 1
|
||||
|
||||
# Pad prompts to the left and responses to the right
|
||||
one_input_ids, one_input_attention_mask = get_left_padded_ids_and_attention_mask(
|
||||
prompt_ids, max_prompt_length, self.pad_token_id
|
||||
)
|
||||
one_response_ids, one_response_attention_mask = get_right_padded_ids_and_attention_mask(
|
||||
response_ids, max_response_length, self.pad_token_id
|
||||
)
|
||||
one_response_mask, _ = get_right_padded_ids_and_attention_mask(
|
||||
response_mask, max_response_length, 0
|
||||
)
|
||||
|
||||
input_ids_list.append(one_input_ids)
|
||||
input_attention_mask_list.append(one_input_attention_mask)
|
||||
response_ids_list.append(one_response_ids)
|
||||
response_attention_mask_list.append(one_response_attention_mask)
|
||||
response_mask_list.append(one_response_mask)
|
||||
data_id_list.append(sample_info["data_id"])
|
||||
rollout_id_list.append(rollout_id)
|
||||
# turn_index_list.append(current_merged_trace_idx)
|
||||
else:
|
||||
raise ValueError(f"Unknown trace_aggregator level: {self.trace_aggregator.get('level')}")
|
||||
input_ids_list.append(one_input_ids)
|
||||
input_attention_mask_list.append(one_input_attention_mask)
|
||||
response_ids_list.append(one_response_ids)
|
||||
response_attention_mask_list.append(one_response_attention_mask)
|
||||
data_id_list.append(sample_info["data_id"])
|
||||
rollout_id_list.append(rollout_id)
|
||||
turn_index_list.append(turn_index)
|
||||
|
||||
n_transition = len(input_ids_list)
|
||||
batch_input_ids = torch.LongTensor(input_ids_list).to(device)
|
||||
input_attention_mask = torch.LongTensor(input_attention_mask_list).to(device)
|
||||
batch_response_ids = torch.LongTensor(response_ids_list).to(device)
|
||||
response_attention_mask = torch.LongTensor(response_attention_mask_list).to(device)
|
||||
response_mask = (
|
||||
torch.LongTensor(response_mask_list).to(device) if self.trace_aggregator.get("level", "transition") == "trajectory" else None # type: ignore
|
||||
)
|
||||
|
||||
# Concatenate prompts and responses to form the full sequence
|
||||
batch_seq = torch.cat([batch_input_ids, batch_response_ids], dim=-1)
|
||||
attention_mask = torch.cat([input_attention_mask, response_attention_mask], dim=-1)
|
||||
|
||||
# Compute position_ids - use mrope for Qwen2-VL, standard 2D otherwise
|
||||
if self._use_mrope:
|
||||
# For Qwen2-VL: compute 4D position_ids (batch_size, 4, seq_length)
|
||||
position_ids_list: list[torch.Tensor] = []
|
||||
for i in range(n_transition):
|
||||
pos_ids = self._compute_mrope_position_ids(
|
||||
input_ids=batch_seq[i],
|
||||
attention_mask=attention_mask[i],
|
||||
image_grid_thw=image_grid_thw_list[i] if image_grid_thw_list else None,
|
||||
) # (4, seq_length)
|
||||
position_ids_list.append(pos_ids)
|
||||
# Stack to (batch_size, 4, seq_length)
|
||||
position_ids = torch.stack(position_ids_list, dim=0)
|
||||
else:
|
||||
# Standard 2D position_ids (batch_size, seq_length)
|
||||
position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
|
||||
|
||||
position_ids = torch.clamp(torch.cumsum(attention_mask, dim=-1) - 1, min=0)
|
||||
is_drop_mask = torch.BoolTensor(is_drop_list).to(device)
|
||||
scores = torch.tensor(reward_list, dtype=torch.bfloat16).to(device)
|
||||
|
||||
# Create token-level scores by placing the final reward at the last token position
|
||||
token_level_scores = torch.zeros_like(attention_mask, dtype=scores.dtype)
|
||||
# For mrope (3D position_ids), use the first dimension (text position_ids) for eos calculation
|
||||
if self._use_mrope:
|
||||
# position_ids is (batch_size, 4, seq_length), use first dim for text positions
|
||||
text_position_ids = position_ids[:, 0, :] # (batch_size, seq_length)
|
||||
eos_mask_idx = torch.argmax(text_position_ids * attention_mask, dim=-1) # (bsz,)
|
||||
else:
|
||||
eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bsz,)
|
||||
# At the eos_mask_idx position of each sample, fill in the corresponding scores.
|
||||
# torch.arange(n_transition) generates [0,1,2,...,bsz-1] as indices for the batch dimension.
|
||||
eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bsz,)
|
||||
token_level_scores[torch.arange(n_transition), eos_mask_idx] = scores
|
||||
# Only take the last response_length part of the sequence to get the token-level scores for the model's response part.
|
||||
token_level_scores = token_level_scores[:, -max_response_length:]
|
||||
@@ -1081,12 +758,7 @@ class AgentModeDaemon:
|
||||
"position_ids": position_ids,
|
||||
"is_drop_mask": is_drop_mask,
|
||||
"token_level_scores": token_level_scores.contiguous(),
|
||||
**(
|
||||
{"response_mask": response_mask}
|
||||
if self.trace_aggregator.get("level", "transition") == "trajectory"
|
||||
else {}
|
||||
),
|
||||
}, # type: ignore
|
||||
},
|
||||
batch_size=n_transition,
|
||||
)
|
||||
data_proto = DataProto(batch=batch)
|
||||
@@ -1098,38 +770,12 @@ class AgentModeDaemon:
|
||||
"training/n_rollouts_w_reward": sample_with_reward_count,
|
||||
"training/n_truncated_triplets": n_trunc_sample_because_of_response,
|
||||
"training/n_triplets": n_transition,
|
||||
# log data, only for debug testing
|
||||
**(
|
||||
{
|
||||
"training/n_unmerged_rollouts": unmerged_count, # type: ignore
|
||||
"training/n_triplets_by_turn": len(response_per_turn_list), # type: ignore
|
||||
"training/avg_response_length_by_turn": np.mean(response_per_turn_list), # type: ignore
|
||||
"training/max_response_length_by_turn": np.max(response_per_turn_list), # type: ignore
|
||||
"training/min_response_length_by_turn": np.min(response_per_turn_list), # type: ignore
|
||||
}
|
||||
if self.trace_aggregator.get("level", "transition") == "trajectory"
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"training/template_mismatch_triplets": template_mismatch_count, # type: ignore
|
||||
"training/retoken_mismatch_triplets": retoken_mismatch_count, # type: ignore
|
||||
"training/others_mismatch_triplets": others_mismatch_count, # type: ignore
|
||||
"training/template_mismatch_ratio": template_mismatch_count / len(response_per_turn_list), # type: ignore
|
||||
"training/retoken_mismatch_ratio": retoken_mismatch_count / len(response_per_turn_list), # type: ignore
|
||||
"training/others_mismatch_ratio": others_mismatch_count / len(response_per_turn_list), # type: ignore
|
||||
}
|
||||
if self.trace_aggregator.get("level", "transition") == "trajectory"
|
||||
and self.trace_aggregator.get("debug", False)
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
# Add non-tensor data for advantage calculation and logging
|
||||
data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list) # type: ignore
|
||||
data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list) # type: ignore
|
||||
if self.trace_aggregator.get("level", "transition") == "transition":
|
||||
data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list) # type: ignore
|
||||
data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list) # type: ignore
|
||||
|
||||
return data_proto, data_metrics
|
||||
|
||||
|
||||
@@ -1,16 +1,13 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# pyright: reportUnknownVariableType=false
|
||||
# pyright: reportUnknownMemberType=false
|
||||
# pyright: reportUnknownArgumentType=false
|
||||
# type: ignore
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Type
|
||||
from importlib.metadata import version
|
||||
from typing import Any
|
||||
|
||||
import hydra
|
||||
import ray
|
||||
from ray.actor import ActorClass
|
||||
from packaging import version as packaging_version
|
||||
from verl.trainer.main_ppo import create_rl_sampler
|
||||
from verl.trainer.ppo.reward import load_reward_manager
|
||||
|
||||
@@ -20,10 +17,7 @@ from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Dataset
|
||||
|
||||
from .dataset import AgentDataset, LoadedDataset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .daemon import AgentModeDaemon
|
||||
from .trainer import AgentLightningTrainer
|
||||
from .trainer import AgentLightningTrainer
|
||||
|
||||
__all__ = [
|
||||
"main",
|
||||
@@ -33,20 +27,8 @@ __all__ = [
|
||||
|
||||
|
||||
@hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
|
||||
def main(config: Any):
|
||||
from .daemon import AgentModeDaemon
|
||||
from .trainer import AgentLightningTrainer
|
||||
|
||||
run_ppo(
|
||||
config,
|
||||
train_dataset=None,
|
||||
val_dataset=None,
|
||||
store=None,
|
||||
llm_proxy=None,
|
||||
adapter=None,
|
||||
trainer_cls=AgentLightningTrainer,
|
||||
daemon_cls=AgentModeDaemon,
|
||||
)
|
||||
def main(config):
|
||||
run_ppo(config, train_dataset=None, val_dataset=None, store=None, llm_proxy=None, adapter=None)
|
||||
|
||||
|
||||
def run_ppo(
|
||||
@@ -56,8 +38,6 @@ def run_ppo(
|
||||
store: LightningStore | None,
|
||||
llm_proxy: LLMProxy | None,
|
||||
adapter: TraceAdapter[Any] | None,
|
||||
trainer_cls: Type[AgentLightningTrainer],
|
||||
daemon_cls: Type[AgentModeDaemon],
|
||||
) -> None:
|
||||
if not ray.is_initialized():
|
||||
# this is for local ray cluster
|
||||
@@ -76,15 +56,13 @@ def run_ppo(
|
||||
|
||||
runner = TaskRunner.remote()
|
||||
ray.get(
|
||||
runner.run.remote( # type: ignore
|
||||
runner.run.remote(
|
||||
config=config,
|
||||
train_dataset=train_dataset,
|
||||
val_dataset=val_dataset,
|
||||
store=store,
|
||||
llm_proxy=llm_proxy,
|
||||
adapter=adapter,
|
||||
trainer_cls=trainer_cls,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -94,13 +72,11 @@ class TaskRunner:
|
||||
def run(
|
||||
self,
|
||||
config: Any,
|
||||
train_dataset: Dataset[Any] | None,
|
||||
val_dataset: Dataset[Any] | None,
|
||||
train_dataset: Dataset | None,
|
||||
val_dataset: Dataset | None,
|
||||
store: LightningStore | None,
|
||||
llm_proxy: LLMProxy | None,
|
||||
adapter: TraceAdapter[Any] | None,
|
||||
trainer_cls: Type[AgentLightningTrainer],
|
||||
daemon_cls: Type[AgentModeDaemon],
|
||||
adapter: TraceAdapter | None,
|
||||
):
|
||||
# print initial config
|
||||
from pprint import pprint
|
||||
@@ -115,7 +91,7 @@ class TaskRunner:
|
||||
local_path = copy_to_local(config.actor_rollout_ref.model.path)
|
||||
|
||||
# instantiate tokenizer
|
||||
from verl.utils.tokenizer import hf_processor, hf_tokenizer
|
||||
from verl.utils import hf_processor, hf_tokenizer
|
||||
|
||||
trust_remote_code = config.data.get("trust_remote_code", False)
|
||||
tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)
|
||||
@@ -136,8 +112,7 @@ class TaskRunner:
|
||||
|
||||
elif config.actor_rollout_ref.actor.strategy == "megatron":
|
||||
assert config.actor_rollout_ref.actor.strategy == config.critic.strategy
|
||||
# FIXME: This import is outdated
|
||||
from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup # type: ignore
|
||||
from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup
|
||||
from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker
|
||||
|
||||
actor_rollout_cls = ActorRolloutRefWorker
|
||||
@@ -146,16 +121,9 @@ class TaskRunner:
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
from verl.trainer.ppo.ray_trainer import ResourcePoolManager
|
||||
from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role
|
||||
|
||||
try:
|
||||
# verl >= 0.6.0
|
||||
from verl.trainer.ppo.utils import Role
|
||||
except ImportError:
|
||||
# Fallback for verl <= 0.5.0
|
||||
from verl.trainer.ppo.ray_trainer import Role # type: ignore
|
||||
|
||||
role_worker_mapping: dict[Role, ActorClass[Any]] = {
|
||||
role_worker_mapping = {
|
||||
Role.ActorRollout: ray.remote(actor_rollout_cls),
|
||||
Role.Critic: ray.remote(CriticWorker),
|
||||
}
|
||||
@@ -222,7 +190,7 @@ class TaskRunner:
|
||||
val_dataset = LoadedDataset(val_dataset)
|
||||
|
||||
train_sampler = create_rl_sampler(config.data, train_dataset)
|
||||
trainer = trainer_cls(
|
||||
trainer = AgentLightningTrainer(
|
||||
config=config,
|
||||
tokenizer=tokenizer,
|
||||
processor=processor,
|
||||
@@ -238,7 +206,6 @@ class TaskRunner:
|
||||
store=store,
|
||||
llm_proxy=llm_proxy,
|
||||
adapter=adapter,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
trainer.init_workers()
|
||||
trainer.fit()
|
||||
|
||||
@@ -8,7 +8,7 @@ import random
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from pprint import pprint
|
||||
from typing import Dict, Tuple, Type
|
||||
from typing import Dict, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -174,18 +174,12 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: LightningStore | None,
|
||||
llm_proxy: LLMProxy | None,
|
||||
adapter: TraceAdapter | None,
|
||||
daemon_cls: Type[AgentModeDaemon],
|
||||
**kwargs,
|
||||
self, store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter | None, **kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.store = store
|
||||
self.llm_proxy = llm_proxy
|
||||
self.adapter = adapter
|
||||
self.daemon_cls = daemon_cls
|
||||
|
||||
def _validate(self):
|
||||
assert len(self.val_dataloader) == 1, "Please set val_batch_size to None for better throughput."
|
||||
@@ -205,37 +199,6 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
self.async_rollout_manager.sleep()
|
||||
return test_metrics
|
||||
|
||||
def _compute_reference_log_prob(self, batch: DataProto) -> DataProto:
|
||||
"""Compute reference log probability using the correct worker based on LoRA configuration.
|
||||
|
||||
In verl 0.6.0+, when LoRA is detected (indicated by ref_in_actor=True),
|
||||
the reference policy is computed by the actor rollout worker instead of a separate
|
||||
ref policy worker. This method handles both scenarios by checking the ref_in_actor flag.
|
||||
Note: verl sets ref_in_actor=True when it detects LoRA configuration (e.g., lora_rank > 0 or lora_adapter_path is set).
|
||||
|
||||
Args:
|
||||
batch: The data batch to compute reference log probabilities for.
|
||||
|
||||
Returns:
|
||||
DataProto with reference log probabilities added.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the required worker is not available.
|
||||
"""
|
||||
if getattr(self, "ref_in_actor", False):
|
||||
actor_worker = getattr(self, "actor_rollout_wg", None)
|
||||
if actor_worker is None:
|
||||
raise RuntimeError("actor_rollout_wg is required when ref_in_actor is True.")
|
||||
return actor_worker.compute_ref_log_prob(batch)
|
||||
|
||||
ref_worker = getattr(self, "ref_policy_wg", None)
|
||||
if ref_worker is None:
|
||||
raise RuntimeError(
|
||||
"Reference policy worker was not initialized. "
|
||||
"Ensure `use_reference_policy` is enabled and the VERL config exposes the ref worker."
|
||||
)
|
||||
return ref_worker.compute_ref_log_prob(batch)
|
||||
|
||||
def _train_step(self, batch_dict: dict) -> dict:
|
||||
# Isolate in a separate method to automatically recycle the variables before validation.
|
||||
batch: DataProto = DataProto.from_single_dict(batch_dict)
|
||||
@@ -255,18 +218,9 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
)
|
||||
self.agent_mode_daemon.run_until_all_finished()
|
||||
batch, agent_metrics = self.agent_mode_daemon.get_train_data_batch(
|
||||
max_prompt_length=(
|
||||
self.config.agentlightning.trace_aggregator.trajectory_max_prompt_length
|
||||
if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
|
||||
else self.config.data.max_prompt_length
|
||||
),
|
||||
max_response_length=(
|
||||
self.config.agentlightning.trace_aggregator.trajectory_max_response_length
|
||||
if self.config.agentlightning.trace_aggregator.level.startswith("trajectory")
|
||||
else self.config.data.max_response_length
|
||||
),
|
||||
max_prompt_length=self.config.data.max_prompt_length,
|
||||
max_response_length=self.config.data.max_response_length,
|
||||
device=gen_batch.batch["fake_ids"].device,
|
||||
global_steps=self.global_steps,
|
||||
)
|
||||
metrics.update(agent_metrics)
|
||||
self.agent_mode_daemon.clear_data_and_server()
|
||||
@@ -291,8 +245,7 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
# uid is used for algorithm like GRPO, should be aligned to data id
|
||||
batch.non_tensor_batch["uid"] = batch.non_tensor_batch["data_id_list"]
|
||||
|
||||
if "response_mask" not in batch.batch:
|
||||
batch.batch["response_mask"] = compute_response_mask(batch)
|
||||
batch.batch["response_mask"] = compute_response_mask(batch)
|
||||
|
||||
# compute global_valid tokens
|
||||
batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
|
||||
@@ -323,7 +276,7 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
if self.use_reference_policy:
|
||||
# compute reference log_prob
|
||||
with _timer("ref", timing_raw):
|
||||
ref_log_prob = self._compute_reference_log_prob(batch)
|
||||
ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
|
||||
batch = batch.union(ref_log_prob)
|
||||
|
||||
# compute values
|
||||
@@ -460,7 +413,7 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
else:
|
||||
# For other versions (e.g., 0.6.0), we use the full path to the model.
|
||||
model = self.config.actor_rollout_ref.model.path
|
||||
self.agent_mode_daemon = self.daemon_cls(
|
||||
self.agent_mode_daemon = AgentModeDaemon(
|
||||
self.config.agentlightning.port,
|
||||
self.config.actor_rollout_ref.rollout.n,
|
||||
train_information={
|
||||
@@ -474,9 +427,6 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
store=self.store,
|
||||
llm_proxy=self.llm_proxy,
|
||||
adapter=self.adapter,
|
||||
processor=self.processor, # For Qwen2-VL mrope position_ids
|
||||
image_base_dir=getattr(self.config.data, "image_base_dir", None),
|
||||
trace_aggregator=self.config.agentlightning.trace_aggregator,
|
||||
)
|
||||
self.agent_mode_daemon.start()
|
||||
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
# Put contrib-related gitignore files here.
|
||||
@@ -1,4 +0,0 @@
|
||||
# Put code owner definitions here.
|
||||
|
||||
# Recipes
|
||||
recipes/search_r1 @SiyunZhao @JiahangXu
|
||||
@@ -1,21 +0,0 @@
|
||||
# Contrib Area
|
||||
|
||||
This tree hosts experimental integrations, third-party recipes, and curated recipes that are not ready for the main `agentlightning/`, `examples/`, or `docs/` trees. Treat it as an incubator: keep contributions self-contained, clearly owned, and reproducible so downstream users can vendor them without guesswork.
|
||||
|
||||
## When to add something here
|
||||
|
||||
- You are iterating on a runtime extension that would bloat the primary `agentlightning/` namespace.
|
||||
- You want to share a recipe that assembles existing components for a focused agent training or optimization workflow and needs more context than the main examples directory allows.
|
||||
- You need automation scripts or download helpers that will help the community but should not live under `scripts/` at the repo root.
|
||||
|
||||
If a contribution starts depending on core release cadence, tight CI guarantees, or repo-wide infrastructure, talk to maintainers about graduating it out of `contrib/`.
|
||||
|
||||
## Directory map
|
||||
|
||||
- `agentlightning/` — Namespace packages, utilities, and adapters that extend the published wheel. Place new code under `agentlightning/contrib/<feature>/` so `import agentlightning.contrib.<feature>` works for downstream users.
|
||||
- `recipes/` — Task-focused example bundles that solve a specific problem and derive certain results. Each recipe belongs in its own directory with a README that documents usage, result reports, and ownership.
|
||||
- `scripts/` — Shared automation, dataset download steps, or reproducibility helpers that support the contrib modules above.
|
||||
|
||||
When adding folders, document the intent in a local README, link to companion docs or examples, and update `CODEOWNERS` so future fixes reach the right reviewers quickly.
|
||||
|
||||
Questions or proposals for new subtrees can be discussed in Discord, GitHub issues, or GitHub Discussions before opening a PR. For the canonical requirements and review checklist, see the “Agent-lightning Contrib” section of [`docs/community/contributing.md`](../docs/community/contributing.md).
|
||||
@@ -1,3 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# Namespace package for agentlightning.contrib.
|
||||
@@ -1,171 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from copy import deepcopy
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict
|
||||
|
||||
import pandas as pd
|
||||
from search_r1_agent import SearchR1Agent
|
||||
|
||||
import agentlightning as agl
|
||||
|
||||
RL_TRAINING_CONFIG: Dict[str, Any] = {
|
||||
"algorithm": {
|
||||
"adv_estimator": "grpo",
|
||||
"use_kl_in_reward": False,
|
||||
},
|
||||
"data": {
|
||||
"train_files": "data/train.parquet",
|
||||
"val_files": "data/test.parquet",
|
||||
"train_batch_size": 512,
|
||||
"max_prompt_length": 6000,
|
||||
"max_response_length": 4096,
|
||||
"truncation": "error",
|
||||
},
|
||||
"actor_rollout_ref": {
|
||||
"rollout": {
|
||||
"tensor_model_parallel_size": 1,
|
||||
"n": 5,
|
||||
"log_prob_micro_batch_size_per_gpu": 4,
|
||||
"multi_turn": {"format": "hermes"},
|
||||
"name": "vllm",
|
||||
"gpu_memory_utilization": 0.5,
|
||||
"engine_kwargs": {
|
||||
"vllm": {
|
||||
"enable_auto_tool_choice": True,
|
||||
"tool_call_parser": "hermes",
|
||||
}
|
||||
},
|
||||
},
|
||||
"actor": {
|
||||
"ppo_mini_batch_size": 256,
|
||||
"ppo_micro_batch_size_per_gpu": 4,
|
||||
"optim": {"lr": 1e-6, "lr_warmup_steps_ratio": 0.95},
|
||||
"use_kl_loss": True,
|
||||
"kl_loss_type": "low_var_kl",
|
||||
"kl_loss_coef": 0.001,
|
||||
"entropy_coeff": 0,
|
||||
"clip_ratio_low": 0.2,
|
||||
"clip_ratio_high": 0.3,
|
||||
"fsdp_config": {
|
||||
"param_offload": True,
|
||||
"optimizer_offload": True,
|
||||
},
|
||||
},
|
||||
"ref": {
|
||||
"log_prob_micro_batch_size_per_gpu": 4,
|
||||
"fsdp_config": {"param_offload": True},
|
||||
},
|
||||
"model": {
|
||||
"path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
|
||||
"use_remove_padding": True,
|
||||
"enable_gradient_checkpointing": True,
|
||||
},
|
||||
},
|
||||
"trainer": {
|
||||
"n_gpus_per_node": 8,
|
||||
"val_before_train": True,
|
||||
"critic_warmup": 0,
|
||||
"logger": ["console", "wandb"],
|
||||
"project_name": "AgentLightning",
|
||||
"experiment_name": "searchr1",
|
||||
"nnodes": 1,
|
||||
"test_freq": 10,
|
||||
"save_freq": 10,
|
||||
"total_epochs": 15,
|
||||
"total_training_steps": 300,
|
||||
"default_local_dir": "checkpoints/searchr1_checkpoints/",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def config_train_fast() -> Dict[str, Any]:
|
||||
"""A fast training run for CI testing purposes."""
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
EXPERIMENT_NAME = f"searchr1_{timestamp}"
|
||||
PROJECT_NAME = "AgentLightningCI"
|
||||
|
||||
# Simulate writing to $GITHUB_OUTPUT if it’s set
|
||||
github_output = os.getenv("GITHUB_OUTPUT")
|
||||
if github_output:
|
||||
with open(github_output, "a") as f:
|
||||
f.write(f"project_name={PROJECT_NAME}\n")
|
||||
f.write(f"run_name={EXPERIMENT_NAME}\n")
|
||||
|
||||
print("Set environment variables:")
|
||||
print(f"PROJECT_NAME={PROJECT_NAME}")
|
||||
print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
|
||||
|
||||
config = deepcopy(RL_TRAINING_CONFIG)
|
||||
config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
|
||||
config["actor_rollout_ref"]["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
|
||||
config["data"]["val_files"] = "data/test_dev.parquet"
|
||||
config["trainer"]["total_epochs"] = 1
|
||||
config["trainer"]["total_training_steps"] = 1
|
||||
config["trainer"]["experiment_name"] = EXPERIMENT_NAME
|
||||
config["trainer"]["project_name"] = PROJECT_NAME
|
||||
config["trainer"]["test_freq"] = 1
|
||||
return config
|
||||
|
||||
|
||||
def config_train_qwen() -> Dict[str, Any]:
|
||||
"""A configuration for training with Qwen-2.5."""
|
||||
|
||||
config = deepcopy(RL_TRAINING_CONFIG)
|
||||
return config
|
||||
|
||||
|
||||
def config_train_llama() -> Dict[str, Any]:
|
||||
"""A configuration for training with LLaMA-3.2-3B-Instruct.
|
||||
|
||||
You will need a `HF_TOKEN` set to run with this config.
|
||||
"""
|
||||
|
||||
config = deepcopy(RL_TRAINING_CONFIG)
|
||||
config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
|
||||
config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"] = "llama3_json"
|
||||
config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-3B-Instruct"
|
||||
return config
|
||||
|
||||
|
||||
def train(config: Dict[str, Any]) -> None:
|
||||
|
||||
agent = SearchR1Agent()
|
||||
algorithm = agl.VERL(config)
|
||||
trainer = agl.Trainer(n_runners=32, algorithm=algorithm)
|
||||
|
||||
train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records") # type: ignore
|
||||
val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records") # type: ignore
|
||||
trainer.fit(agent, train_dataset=train_data, val_dataset=val_data) # type: ignore
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Main function to parse arguments and run training."""
|
||||
parser = argparse.ArgumentParser(description="Train a Search-R1 agent using different model configurations")
|
||||
|
||||
parser.add_argument(
|
||||
"config",
|
||||
choices=["fast", "qwen", "llama"],
|
||||
help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B-Instruct)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get the appropriate configuration
|
||||
config_functions = {"fast": config_train_fast, "qwen": config_train_qwen, "llama": config_train_llama}
|
||||
|
||||
config = config_functions[args.config]()
|
||||
|
||||
print(f"Starting training with '{args.config}' configuration...")
|
||||
|
||||
train(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -130,6 +130,3 @@ dist
|
||||
.pnp.*
|
||||
|
||||
.DS_Store
|
||||
|
||||
# Storybook build output
|
||||
storybook-static
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
import { useMemo, useState } from 'react';
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { IconSearch } from '@tabler/icons-react';
|
||||
import { waitFor, within } from '@testing-library/dom';
|
||||
import userEvent from '@testing-library/user-event';
|
||||
import { Box, Stack, TextInput, Title } from '@mantine/core';
|
||||
import type { Span } from '@/types';
|
||||
import { compareRecords } from '@/utils/table';
|
||||
@@ -367,224 +365,3 @@ export const NestedSpans: Story = {
|
||||
<TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.traceId === 'trace-nested456')} />
|
||||
),
|
||||
};
|
||||
|
||||
// Test data with sequence IDs that would sort incorrectly if treated as strings
|
||||
const sequenceSortTestSpans: Span[] = [
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 2,
|
||||
traceId: 'trace-seq-002',
|
||||
spanId: 'span-seq-002',
|
||||
parentId: null,
|
||||
name: 'task_sequence_2',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 200,
|
||||
endTime: now - 190,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 10,
|
||||
traceId: 'trace-seq-010',
|
||||
spanId: 'span-seq-010',
|
||||
parentId: null,
|
||||
name: 'task_sequence_10',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 180,
|
||||
endTime: now - 170,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 3,
|
||||
traceId: 'trace-seq-003',
|
||||
spanId: 'span-seq-003',
|
||||
parentId: null,
|
||||
name: 'task_sequence_3',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 160,
|
||||
endTime: now - 150,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 11,
|
||||
traceId: 'trace-seq-011',
|
||||
spanId: 'span-seq-011',
|
||||
parentId: null,
|
||||
name: 'task_sequence_11',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 140,
|
||||
endTime: now - 130,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 9,
|
||||
traceId: 'trace-seq-009',
|
||||
spanId: 'span-seq-009',
|
||||
parentId: null,
|
||||
name: 'task_sequence_9',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 120,
|
||||
endTime: now - 110,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 12,
|
||||
traceId: 'trace-seq-012',
|
||||
spanId: 'span-seq-012',
|
||||
parentId: null,
|
||||
name: 'task_sequence_12',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 100,
|
||||
endTime: now - 90,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 4,
|
||||
traceId: 'trace-seq-004',
|
||||
spanId: 'span-seq-004',
|
||||
parentId: null,
|
||||
name: 'task_sequence_4',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 80,
|
||||
endTime: now - 70,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 6,
|
||||
traceId: 'trace-seq-006',
|
||||
spanId: 'span-seq-006',
|
||||
parentId: null,
|
||||
name: 'task_sequence_6',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 60,
|
||||
endTime: now - 50,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 7,
|
||||
traceId: 'trace-seq-007',
|
||||
spanId: 'span-seq-007',
|
||||
parentId: null,
|
||||
name: 'task_sequence_7',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 40,
|
||||
endTime: now - 30,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 13,
|
||||
traceId: 'trace-seq-013',
|
||||
spanId: 'span-seq-013',
|
||||
parentId: null,
|
||||
name: 'task_sequence_13',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 20,
|
||||
endTime: now - 10,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-seq-test',
|
||||
attemptId: 'at-seq-test',
|
||||
sequenceId: 14,
|
||||
traceId: 'trace-seq-014',
|
||||
spanId: 'span-seq-014',
|
||||
parentId: null,
|
||||
name: 'task_sequence_14',
|
||||
status: { status_code: 'UNSET', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 5,
|
||||
endTime: now,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
];
|
||||
|
||||
export const SequenceIdSortTest: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={1200} spans={sequenceSortTestSpans} />,
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
const seqHeader = await canvas.findByRole('button', { name: /seq\./i });
|
||||
|
||||
await userEvent.click(seqHeader);
|
||||
|
||||
await waitFor(() => {
|
||||
const rows = canvas.getAllByRole('row');
|
||||
const firstRow = rows[1];
|
||||
if (!firstRow) {
|
||||
throw new Error('Expected at least one data row after sorting by sequence ID');
|
||||
}
|
||||
within(firstRow).getByText('task_sequence_14');
|
||||
});
|
||||
},
|
||||
};
|
||||
|
||||
@@ -23,7 +23,6 @@ export const selectTracesViewMode = (state: RootState) => selectTracesState(stat
|
||||
|
||||
const TRACES_SORT_FIELD_MAP: Record<string, string> = {
|
||||
name: 'name',
|
||||
sequenceId: 'sequence_id',
|
||||
traceId: 'trace_id',
|
||||
spanId: 'span_id',
|
||||
parentId: 'parent_id',
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
services:
|
||||
prometheus:
|
||||
image: prom/prometheus:latest
|
||||
command:
|
||||
- "--config.file=/etc/prometheus/prometheus.yml"
|
||||
- "--storage.tsdb.path=/prometheus"
|
||||
volumes:
|
||||
- ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
|
||||
- ${AGL_MONITORING_DATA_PATH:?Set AGL_MONITORING_DATA_PATH to the metrics directory}/prometheus:/prometheus
|
||||
ports:
|
||||
- "9090:9090"
|
||||
|
||||
grafana:
|
||||
image: grafana/grafana:latest
|
||||
depends_on:
|
||||
- prometheus
|
||||
ports:
|
||||
- "9091:3000"
|
||||
volumes:
|
||||
- ./data/grafana:/var/lib/grafana
|
||||
- ./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml
|
||||
- ./grafana/dashboard-provider.yml:/etc/grafana/provisioning/dashboards/provider.yml
|
||||
- ./grafana/dashboards:/var/lib/grafana/dashboards
|
||||
environment:
|
||||
- GF_INSTALL_PLUGINS=grafana-piechart-panel
|
||||
- GF_AUTH_ANONYMOUS_ENABLED=true
|
||||
- GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
|
||||
- GF_AUTH_DISABLE_LOGIN_FORM=true
|
||||
- GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
|
||||
@@ -4,27 +4,7 @@ services:
|
||||
file: compose.store.yml
|
||||
service: app
|
||||
|
||||
depends_on:
|
||||
- app-exporter # Wait for the exporter to be ready first
|
||||
|
||||
command: agl store --host 0.0.0.0 --port 4747 --tracker console prometheus --backend memory
|
||||
environment:
|
||||
- PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
|
||||
volumes:
|
||||
- prometheus_multiproc:/tmp/prometheus_multiproc
|
||||
|
||||
app-exporter:
|
||||
build:
|
||||
context: ../
|
||||
dockerfile: docker/Dockerfile.dev
|
||||
|
||||
command: agl prometheus --host 0.0.0.0 --port 4748
|
||||
ports:
|
||||
- "4748:4748"
|
||||
environment:
|
||||
- PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
|
||||
volumes:
|
||||
- prometheus_multiproc:/tmp/prometheus_multiproc
|
||||
command: agl store --host 0.0.0.0 --port 4747 --prometheus --backend memory
|
||||
|
||||
node-exporter:
|
||||
image: prom/node-exporter:latest
|
||||
@@ -42,11 +22,10 @@ services:
|
||||
- "--storage.tsdb.path=/prometheus"
|
||||
- "--storage.tsdb.retention.time=1h"
|
||||
volumes:
|
||||
- ./prometheus/prometheus.base.yml:/etc/prometheus/prometheus.yml:ro
|
||||
- ./prometheus.memory-store.yml:/etc/prometheus/prometheus.yml:ro
|
||||
- ./data/prometheus:/prometheus
|
||||
depends_on:
|
||||
- app
|
||||
- app-exporter
|
||||
- node-exporter
|
||||
ports:
|
||||
- "9090:9090"
|
||||
@@ -72,10 +51,3 @@ services:
|
||||
- GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
|
||||
- GF_AUTH_DISABLE_LOGIN_FORM=true
|
||||
- GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
|
||||
|
||||
volumes:
|
||||
prometheus_multiproc:
|
||||
driver: local
|
||||
driver_opts:
|
||||
type: tmpfs
|
||||
device: tmpfs
|
||||
|
||||
@@ -21,33 +21,18 @@ services:
|
||||
|
||||
depends_on:
|
||||
- mongo
|
||||
- app-exporter # Wait for the exporter to be ready first
|
||||
|
||||
command:
|
||||
- /bin/bash
|
||||
- -c
|
||||
- |
|
||||
mkdir -p /tmp/prometheus &&
|
||||
agl store --host 0.0.0.0 --port 4747 \
|
||||
--tracker console prometheus --backend mongo \
|
||||
--prometheus --backend mongo \
|
||||
--mongo-uri mongodb://mongo:27017/?replicaSet=rs0 \
|
||||
--n-workers ${AGL_STORE_N_WORKERS:-32}
|
||||
environment:
|
||||
- PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
|
||||
volumes:
|
||||
- prometheus_multiproc:/tmp/prometheus_multiproc
|
||||
|
||||
app-exporter:
|
||||
build:
|
||||
context: ../
|
||||
dockerfile: docker/Dockerfile.dev
|
||||
|
||||
command: agl prometheus --host 0.0.0.0 --port 4748
|
||||
ports:
|
||||
- "4748:4748"
|
||||
environment:
|
||||
- PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc
|
||||
volumes:
|
||||
- prometheus_multiproc:/tmp/prometheus_multiproc
|
||||
- PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus
|
||||
|
||||
mongodb-exporter:
|
||||
image: percona/mongodb_exporter:0.47.1
|
||||
@@ -76,11 +61,10 @@ services:
|
||||
- "--storage.tsdb.path=/prometheus"
|
||||
- "--storage.tsdb.retention.time=1h"
|
||||
volumes:
|
||||
- ./prometheus/prometheus.mongo.yml:/etc/prometheus/prometheus.yml:ro
|
||||
- ./prometheus.mongo-store.yml:/etc/prometheus/prometheus.yml:ro
|
||||
- ./data/prometheus:/prometheus
|
||||
depends_on:
|
||||
- app
|
||||
- app-exporter
|
||||
- mongodb-exporter
|
||||
- node-exporter
|
||||
ports:
|
||||
@@ -107,10 +91,3 @@ services:
|
||||
- GF_AUTH_ANONYMOUS_ORG_ROLE=Admin
|
||||
- GF_AUTH_DISABLE_LOGIN_FORM=true
|
||||
- GF_DASHBOARDS_DEFAULT_HOME_DASHBOARD_PATH=/var/lib/grafana/dashboards/agentlightning.json
|
||||
|
||||
volumes:
|
||||
prometheus_multiproc:
|
||||
driver: local
|
||||
driver_opts:
|
||||
type: tmpfs
|
||||
device: tmpfs
|
||||
|
||||
@@ -9,11 +9,6 @@ services:
|
||||
|
||||
command: agl store --host 0.0.0.0 --port 4747
|
||||
|
||||
ulimits:
|
||||
nofile:
|
||||
soft: 65535
|
||||
hard: 65535
|
||||
|
||||
develop:
|
||||
watch:
|
||||
# Sync the working directory with the `/app` directory in the container
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,11 +1,11 @@
|
||||
global:
|
||||
scrape_interval: 5s
|
||||
evaluation_interval: 5s
|
||||
scrape_interval: 2s
|
||||
evaluation_interval: 2s
|
||||
|
||||
scrape_configs:
|
||||
- job_name: app
|
||||
static_configs:
|
||||
- targets: ["app-exporter:4748"]
|
||||
- targets: ["app:4747"]
|
||||
metrics_path: /v1/prometheus/
|
||||
|
||||
- job_name: node
|
||||
@@ -1,11 +1,11 @@
|
||||
global:
|
||||
scrape_interval: 5s
|
||||
evaluation_interval: 5s
|
||||
scrape_interval: 2s
|
||||
evaluation_interval: 2s
|
||||
|
||||
scrape_configs:
|
||||
- job_name: app
|
||||
static_configs:
|
||||
- targets: ["app-exporter:4748"]
|
||||
- targets: ["app:4747"]
|
||||
metrics_path: /v1/prometheus/
|
||||
|
||||
- job_name: node
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 275 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 247 KiB |
File diff suppressed because one or more lines are too long
@@ -28,7 +28,7 @@ Documentation improvements are the easiest way to get started. You can find more
|
||||
|
||||
Bug fixes are the fastest way to get familiar with the codebase. To get started, you can:
|
||||
|
||||
- Browse the ["help wanted"](https://github.com/microsoft/agent-lightning/labels/help%20wanted) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
|
||||
- Browse the ["good first issue"](https://github.com/microsoft/agent-lightning/labels/good%20first%20issue) and ["bug"](https://github.com/microsoft/agent-lightning/labels/bug) labels; drop a comment before you start so we can mark it as taken.
|
||||
- For fresh bugs, open an issue with reproduction steps, logs, and expected behavior before submitting a fix.
|
||||
- Keep each pull request focused, ideally avoiding breaking API changes. Larger refactors should be discussed via RFC or maintainer sync.
|
||||
|
||||
@@ -46,7 +46,6 @@ Bonus points for examples that:
|
||||
- Ship CI or self-test coverage so we know they still work as the core evolves. **Otherwise, we would have to mark the example as unmaintained because we won't be able to test the examples manually before each release.**
|
||||
- Include a [`docs/how-to/`]({{ src("docs/how-to/") }}) guide (or a detailed README if no how-to exists) without duplicating content in multiple places.
|
||||
- Favor simple, dependency-light code over heavy abstractions.
|
||||
- Ship a README that documents smoke-test instructions and includes an "Included Files" section summarizing every file and its role; keep the runnable module self-contained with a module-level docstring explaining CLI usage, plus targeted docstrings or inline comments for educational functions/classes.
|
||||
|
||||
!!! warning "Please discuss first"
|
||||
|
||||
@@ -75,27 +74,6 @@ Most brand-new algorithms ultimately land as “new examples,” so read that se
|
||||
|
||||
Have a project that builds on Agent-lightning but does not belong in the main repo? Fork it or depend on it externally, then let us know. We can showcase notable projects in [Community Projects](../index.md) and the main [README]({{ src("README.md") }}).
|
||||
|
||||
### Agent-lightning Contrib
|
||||
|
||||
[`contrib/`]({{ src("contrib") }}) is where work-in-progress or third-party integrations, and curated recipes live before they are hardened enough for the core runtime tree. Think of it as an incubator: additions should remain easy to consume, clearly owned, and scoped so downstream users can vendor them with minimal risk.
|
||||
|
||||
The following types of contributions are welcome in the contrib area:
|
||||
|
||||
- **Recipes** that assemble multiple Agent Lightning components for a narrow task (`contrib/recipes/<topic>/`). Each recipe must be self-contained, include running instructions and result reports.
|
||||
- **Runtime extensions** that would bloat the primary `agentlightning/` namespace (`contrib/agentlightning/contrib/<feature>/`). These should mirror the published wheel layout so that `import agentlightning.contrib.<feature>` works out of the box.
|
||||
- **Supporting scripts and assets** (`contrib/scripts/`) that automate dataset downloads, environment preparation, or benchmarks required by contrib modules.
|
||||
|
||||
If you are unsure where a contribution should live, start a thread in Discord or open an issue before writing code. The [contrib README]({{ src("contrib/README.md") }}) also lists the directory expectations.
|
||||
|
||||
A quick checklist for contributions to be accepted:
|
||||
|
||||
1. **Document everything.** Include configuration steps, environment variables, and sample commands so contributors can reproduce the results without guesswork. Pin to a specific version of Agent-lightning and other dependencies to avoid unexpected changes if you don't want to update the recipe frequently.
|
||||
2. **Keep quality predictable.** Match the repo’s style guide, apply exhaustive type hints, and run `uv run --no-sync pyright` plus targeted `pytest` suites for any Python module you touch.
|
||||
3. **Ship reproducibility artifacts.** Store only scripts or instructions for downloading datasets, weights, or binaries. Never upload large artifacts or credentials directly.
|
||||
4. **Update ownership.** Add `CODEOWNERS` entries when new directories appear so maintainers know who can review follow-up fixes.
|
||||
|
||||
Contrib entries do not need the same maturity level as core code, but they must still meet the baseline above. Submissions that lack documentation, hide ownership, or depend on untracked assets are typically rejected until those gaps are resolved.
|
||||
|
||||
### Other Contribution Ideas
|
||||
|
||||
- **Tests.** Add or improve cases in [`tests/`]({{ src("tests") }}) (unit, integration, or end-to-end).
|
||||
@@ -148,13 +126,13 @@ After `uv sync`, run commands via `uv run ...` (add `--no-sync` once the environ
|
||||
Formatting and linting are enforced through [pre-commit](https://pre-commit.com/). Install once, then run before each push:
|
||||
|
||||
```bash
|
||||
uv run --no-sync pre-commit install
|
||||
uv run --no-sync pre-commit run --all-files --show-diff-on-failure --color=always
|
||||
uv run pre-commit install
|
||||
uv run pre-commit run --all-files --show-diff-on-failure --color=always
|
||||
```
|
||||
|
||||
Once installed, the hooks run automatically on every `git commit`. Running the pre-commit hooks locally keeps CI green and diffs manageable.
|
||||
|
||||
### 3. Branch from Fresh `main` and Code
|
||||
### 3. Branch From a Fresh `main`
|
||||
|
||||
Start all work from the latest upstream state:
|
||||
|
||||
@@ -187,28 +165,20 @@ Use lowercase with hyphens, e.g., `feature/async-runner-hooks`.
|
||||
|
||||
Remember to register new docs in [`mkdocs.yml`]({{ src("mkdocs.yml") }}), add examples to [examples/README]({{ src("examples/README.md") }}), and update the [Examples Catalog](../how-to/examples-catalog.md).
|
||||
|
||||
Before you start coding, bring the shared coding conventions with you:
|
||||
|
||||
- Target `requires-python >= 3.10`, four-space indentation, ~120-character lines (docstrings may run longer), and formatter-owned diffs (Black + isort with the `black` profile).
|
||||
- Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
|
||||
- Maintain exhaustive type hints (pyright enforces them), write succinct Google-style docstrings (with `[][]` cross-references).
|
||||
- Prefer dataclasses or Pydantic models from `agentlightning.types`.
|
||||
- Log via `logging.getLogger(__name__)` with targeted DEBUG/INFO/WARNING/ERROR calls—especially for long multi-step functions or broad `try/except` blocks.
|
||||
|
||||
### 4. Test and Validate
|
||||
|
||||
Most contributions require automated checks. Once `uv sync` locks dependencies, prefix commands with `uv run --no-sync ...` so they share the same environment as CI.
|
||||
Most contributions require automated checks. Prefix commands with `uv run` so they use the project environment.
|
||||
|
||||
**Full test suite**
|
||||
|
||||
```bash
|
||||
uv run --no-sync pytest -v
|
||||
uv run pytest -v
|
||||
```
|
||||
|
||||
**Targeted tests**
|
||||
|
||||
```bash
|
||||
uv run --no-sync pytest tests/path/to/test_file.py -k test_name
|
||||
uv run pytest tests/path/to/test_file.py -k test_name
|
||||
```
|
||||
|
||||
**Optional/gated tests:** GPU-specific suites or API-dependent tests run automatically when the required hardware or environment variables (such as `OPENAI_API_KEY`) are present.
|
||||
@@ -216,7 +186,7 @@ uv run --no-sync pytest tests/path/to/test_file.py -k test_name
|
||||
**Static analysis:**
|
||||
|
||||
```bash
|
||||
uv run --no-sync pyright
|
||||
uv run pyright
|
||||
```
|
||||
|
||||
If you have touched code under `examples/`, you should run the example-specific smoke tests. Each directory includes a README with example-specific smoke tests—run those too.
|
||||
@@ -226,8 +196,8 @@ If you have touched code under `examples/`, you should run the example-specific
|
||||
Keep API references under [docs/reference]({{ src("docs/reference/") }}) up to date. Doc-only changes should still build cleanly:
|
||||
|
||||
```bash
|
||||
uv run --no-sync mkdocs serve --strict # live reload
|
||||
uv run --no-sync mkdocs build --strict # CI-equivalent
|
||||
uv run mkdocs serve --strict # live reload
|
||||
uv run mkdocs build --strict # CI-equivalent
|
||||
```
|
||||
|
||||
`--strict` elevates warnings to errors so you catch issues before CI.
|
||||
@@ -235,7 +205,7 @@ If you have touched code under `examples/`, you should run the example-specific
|
||||
Before opening a PR, double-check the basics:
|
||||
|
||||
- Run `uv lock` if you changed dependencies.
|
||||
- Run `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
|
||||
- Run `uv run pre-commit run --all-files` (hooks installed via `pre-commit install` run automatically on `git commit`, but rerun them if you amended history).
|
||||
- Execute the relevant commands from the test list above.
|
||||
- Validate each affected example via its README instructions.
|
||||
|
||||
|
||||
+21
-209
@@ -1,8 +1,8 @@
|
||||
# Understanding Store
|
||||
|
||||
The **[`LightningStore`][agentlightning.LightningStore]** is the central coordination point for Agent-lightning. It holds the task queue, rollouts, attempts, spans, and versioned resources, and exposes a small API both Runners and Algorithms use to communicate. This document explains what's in the store, how statuses transition, how spans are recorded, and the concurrency model (threads & processes).
|
||||
The **[`LightningStore`][agentlightning.LightningStore]** is the central coordination point for Agent-lightning. It holds the task queue, rollouts, attempts, spans, and versioned resources, and exposes a small API both Runners and Algorithms use to communicate. This document explains what’s in the store, how statuses transition, how spans are recorded, and the concurrency model (threads & processes).
|
||||
|
||||
## What's in the Store?
|
||||
## What’s in the Store?
|
||||
|
||||
{ .center }
|
||||
|
||||
@@ -13,11 +13,12 @@ At a high level:
|
||||
* **Attempts** – Each rollout can have multiple executions (retries). Attempts track [`status`][agentlightning.Attempt.status], [`start_time`][agentlightning.Attempt.start_time], [`end_time`][agentlightning.Attempt.end_time], [`last_heartbeat_time`][agentlightning.Attempt.last_heartbeat_time] and link to spans. Valid [AttemptStatus][agentlightning.AttemptStatus] are `preparing`, `running`, `succeeded`, `failed`, `requeuing`, `cancelled`.
|
||||
* **Spans** – Structured trace events produced by the Tracer during an attempt. Spans are ordered by a **monotonic sequence id** per `(rollout_id, attempt_id)`.
|
||||
* **Resources** – Versioned, named bundles (e.g., prompt templates) referenced by rollouts.
|
||||
* **Workers** – Metadata about runner instances: heartbeat timestamps, current assignment, and status.
|
||||
|
||||
Rollout and Task share the same surface in practice: [`Rollout.input`][agentlightning.types.Rollout] is the task input. The queue stores rollouts that are not yet running; [Runners][agentlightning.Runner] dequeue them and update the same rollout's status as work progresses.
|
||||
Rollout and Task share the same surface in practice: [`Rollout.input`][agentlightning.types.Rollout] is the task input. The queue stores rollouts that are not yet running; [Runners][agentlightning.Runner] dequeue them and update the same rollout’s status as work progresses.
|
||||
|
||||
Before we look at status transitions, it helps to keep in mind that rollouts are the "outside view," while attempts are the "inside view." Attempts are what actually run; rollouts summarize the latest attempt plus a small set of control actions like queueing and cancellation.
|
||||
All [`LightningStore`][agentlightning.LightningStore] implementations must inherit from [`LightningStore`][agentlightning.LightningStore] and override the methods to implement the storage logic.
|
||||
|
||||
Before we look at status transitions, it helps to keep in mind that rollouts are the “outside view,” while attempts are the “inside view.” Attempts are what actually run; rollouts summarize the latest attempt plus a small set of control actions like queueing and cancellation.
|
||||
|
||||
## Attempt Status Transitions
|
||||
|
||||
@@ -116,7 +117,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
|
||||
| ------------------------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------- |
|
||||
| N/A | `queuing` | Created by `enqueue_rollout()`. |
|
||||
| `preparing` | `queuing/requeuing` → `preparing` | Typically `dequeue_rollout()` or `start_rollout()`/`start_attempt()` creates a new attempt. |
|
||||
| `running` | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `rollout_status_from_attempt`. |
|
||||
| `running` | `preparing/queuing/requeuing` → `running` | First `add_[otel_]span()` flips the attempt to `running`; rollout follows via `propagate_status`. |
|
||||
| `succeeded` | `*` → `succeeded` | Terminal. Rollout `end_time` set. |
|
||||
| `failed` / `timeout` / `unresponsive` | `*` → `requeuing` | **Only if** `status ∈ retry_condition ∧ sequence_id < max_attempts`. |
|
||||
| `failed` / `timeout` / `unresponsive` | `*` → `failed` | Otherwise (no retries left or retries disabled). |
|
||||
@@ -124,7 +125,7 @@ rollout = await store.enqueue_rollout(input, config=cfg)
|
||||
|
||||
!!! note "Why aggregation?"
|
||||
|
||||
In code, we use `rollout_status_from_attempt()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
|
||||
In code, we use `propagate_status()` which actively updates the rollout based on the latest attempt. Reading the table above is usually easier than reverse-engineering the propagation logic in the code: think of the rollout’s transitions as *callbacks* on attempt state changes, plus queue/cancel paths.
|
||||
|
||||
## Spans
|
||||
|
||||
@@ -151,220 +152,31 @@ Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span
|
||||
|
||||
[`add_span`][agentlightning.LightningStore.add_span] or [`add_otel_span`][agentlightning.LightningStore.add_otel_span] both appends a span *and* acts as a heartbeat that can revive `unresponsive` → `running`.
|
||||
|
||||
### OTLP Compatibility
|
||||
## OTLP Compatibility
|
||||
|
||||
Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
|
||||
Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence
|
||||
numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
|
||||
|
||||
Because the interface speaks standard OTLP, any OpenTelemetry-compatible SDK or collector can emit spans directly to a LightningStore OTLP endpoint without custom shims. The server responds according to the OTLP contract (status code, encoding, and error payloads), which keeps Agent-lightning interoperable with existing observability tooling. This compatibility serves as a strong complement to the OpenTelemetry conversion discussed above.
|
||||
|
||||
Check whether the store supports OTLP traces via the [`capabilities["otlp_traces"]`][agentlightning.LightningStore.capabilities] property.
|
||||
## Store Implementations
|
||||
|
||||
## Implementation Overview
|
||||
Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
|
||||
|
||||
The `agentlightning.store` module is organized into two distinct layers plus optional wrappers:
|
||||
- Fast startup, zero external dependencies, and ideal for local development, CI, and unit tests.
|
||||
- Fully asyncio-safe for writes; most reader operations can iterate without locks, except those that need to perform multiple queries.
|
||||
- Includes a best-effort span eviction policy once memory crosses a configured watermark; querying evicted spans raises a clear error so callers can fall back.
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
direction TB
|
||||
For production you will likely want persistence. We’re actively building a SQLite-backed store that keeps the same API surface while adding durability, crash recovery, and better historical span queries. If you need something sooner, implement your own store by subclassing [`LightningStore`][agentlightning.LightningStore] and providing concrete storage for the small set of abstract methods (`enqueue_rollout`, `dequeue_rollout`, `update_attempt`, `add_span`, etc.). This document plus the tests in `tests/store/` illustrate the expected behavior.
|
||||
|
||||
class LightningStore {
|
||||
<<abstract>>
|
||||
+enqueue_rollout()
|
||||
+dequeue_rollout()
|
||||
+update_attempt()
|
||||
+add_span()
|
||||
+query_rollouts()
|
||||
...
|
||||
}
|
||||
|
||||
class LightningCollections {
|
||||
<<abstract>>
|
||||
+rollouts: Collection
|
||||
+attempts: Collection
|
||||
+spans: Collection
|
||||
+resources: Collection
|
||||
+workers: Collection
|
||||
+rollout_queue: Queue
|
||||
+span_sequence_ids: KeyValue
|
||||
+atomic()
|
||||
}
|
||||
|
||||
class CollectionBasedLightningStore~T~ {
|
||||
+collections: T
|
||||
-healthcheck_before()
|
||||
-tracked()
|
||||
}
|
||||
|
||||
class InMemoryLightningStore
|
||||
class MongoLightningStore
|
||||
class InMemoryLightningCollections
|
||||
class MongoLightningCollections
|
||||
|
||||
class LightningStoreServer {
|
||||
+store: LightningStore
|
||||
+start()
|
||||
+stop()
|
||||
}
|
||||
class LightningStoreClient {
|
||||
+server_address: str
|
||||
}
|
||||
class LightningStoreThreaded {
|
||||
+store: LightningStore
|
||||
}
|
||||
|
||||
LightningStore <|-- CollectionBasedLightningStore
|
||||
LightningStore <|-- LightningStoreServer
|
||||
LightningStore <|-- LightningStoreClient
|
||||
LightningStore <|-- LightningStoreThreaded
|
||||
|
||||
CollectionBasedLightningStore <|-- InMemoryLightningStore
|
||||
CollectionBasedLightningStore <|-- MongoLightningStore
|
||||
|
||||
LightningCollections <|-- InMemoryLightningCollections
|
||||
LightningCollections <|-- MongoLightningCollections
|
||||
|
||||
InMemoryLightningStore ..> InMemoryLightningCollections : uses
|
||||
MongoLightningStore ..> MongoLightningCollections : uses
|
||||
|
||||
LightningStoreServer o-- LightningStore : wraps
|
||||
LightningStoreThreaded o-- LightningStore : wraps
|
||||
```
|
||||
|
||||
1. **Collections Layer** – Low-level storage primitives ([`LightningCollections`][agentlightning.store.collection.LightningCollections]) providing CRUD operations via [`Collection`][agentlightning.store.collection.Collection], [`Queue`][agentlightning.store.collection.Queue], and [`KeyValue`][agentlightning.store.collection.KeyValue] interfaces. Each backend (in-memory, MongoDB) implements these primitives.
|
||||
|
||||
2. **Store Layer** – All [`LightningStore`][agentlightning.LightningStore] implementations must inherit from [`LightningStore`][agentlightning.LightningStore] and override the methods to implement the storage logic. [`CollectionBasedLightningStore`][agentlightning.CollectionBasedLightningStore] builds on collections to implement the full [`LightningStore`][agentlightning.LightningStore] API, including business logic like status transitions, watchdog health checks, and retry policies.
|
||||
|
||||
3. **Wrappers** – Cross-cutting concerns live in thin wrappers:
|
||||
- [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded] adds mutex-based thread safety.
|
||||
- [`LightningStoreServer`][agentlightning.LightningStoreServer] / [`LightningStoreClient`][agentlightning.LightningStoreClient] enable multi-process access over HTTP.
|
||||
|
||||
## Collections
|
||||
|
||||
The collections layer provides storage primitives that [`CollectionBasedLightningStore`][agentlightning.CollectionBasedLightningStore] builds upon. This separation keeps business logic (status transitions, watchdog, retries) in the store layer while allowing different backends to focus purely on persistence.
|
||||
|
||||
The off-the-shelf implementations are [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] and [`MongoLightningCollections`][agentlightning.store.collection.mongo.MongoLightningCollections], which are the underlying collections for [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] and [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], respectively.
|
||||
|
||||
### Collection Primitives
|
||||
|
||||
[`LightningCollections`][agentlightning.store.collection.LightningCollections] bundles three primitive types:
|
||||
|
||||
| Primitive | Purpose | Methods |
|
||||
|-----------|---------|---------|
|
||||
| [`Collection[T]`][agentlightning.store.collection.Collection] | Indexed storage with primary keys | [`query()`][agentlightning.store.collection.Collection.query], [`get()`][agentlightning.store.collection.Collection.get], [`insert()`][agentlightning.store.collection.Collection.insert], [`update()`][agentlightning.store.collection.Collection.update], [`upsert()`][agentlightning.store.collection.Collection.upsert], [`delete()`][agentlightning.store.collection.Collection.delete] |
|
||||
| [`Queue[T]`][agentlightning.store.collection.Queue] | FIFO queue for task scheduling | [`enqueue()`][agentlightning.store.collection.Queue.enqueue], [`dequeue()`][agentlightning.store.collection.Queue.dequeue], [`peek()`][agentlightning.store.collection.Queue.peek], [`size()`][agentlightning.store.collection.Queue.size] |
|
||||
| [`KeyValue[K, V]`][agentlightning.store.collection.KeyValue] | Simple key-value store | [`get()`][agentlightning.store.collection.KeyValue.get], [`set()`][agentlightning.store.collection.KeyValue.set], [`inc()`][agentlightning.store.collection.KeyValue.inc], [`chmax()`][agentlightning.store.collection.KeyValue.chmax], [`pop()`][agentlightning.store.collection.KeyValue.pop] |
|
||||
|
||||
Every [`LightningCollections`][agentlightning.store.collection.LightningCollections] instance exposes these named collections:
|
||||
|
||||
- `rollouts` – [`Collection[Rollout]`][agentlightning.store.collection.Collection] keyed by `rollout_id`
|
||||
- `attempts` – [`Collection[Attempt]`][agentlightning.store.collection.Collection] keyed by `(rollout_id, attempt_id)`
|
||||
- `spans` – [`Collection[Span]`][agentlightning.store.collection.Collection] keyed by `(rollout_id, attempt_id, span_id)`
|
||||
- `resources` – [`Collection[ResourcesUpdate]`][agentlightning.store.collection.Collection] keyed by `resources_id`
|
||||
- `workers` – [`Collection[Worker]`][agentlightning.store.collection.Collection] keyed by `worker_id`
|
||||
- `rollout_queue` – [`Queue[str]`][agentlightning.store.collection.Queue] holding rollout IDs awaiting execution
|
||||
- `span_sequence_ids` – [`KeyValue[str, int]`][agentlightning.store.collection.KeyValue] tracking monotonic sequence counters
|
||||
|
||||
### Atomic Operations
|
||||
|
||||
Collections support atomic operations through the [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] context manager:
|
||||
|
||||
```python
|
||||
async with collections.atomic(mode="rw", labels=["rollouts", "attempts"]) as ctx:
|
||||
rollout = await ctx.rollouts.get(filter={"rollout_id": {"exact": rollout_id}})
|
||||
# modify and update within the same transaction
|
||||
await ctx.rollouts.update([updated_rollout])
|
||||
```
|
||||
|
||||
The arguments passed to [`atomic()`][agentlightning.store.collection.LightningCollections.atomic] are quite arbitrary and flexible. Different implementations may have different interpretations of the arguments. For example, to [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections], the `mode` parameter controls locking behavior (`"r"` for read-only, `"rw"` for read-write), while `labels` specifies which collections to lock. Acquiring locks in sorted order prevents deadlocks when multiple operations run concurrently.
|
||||
|
||||
### Implementing a Custom Backend
|
||||
|
||||
To add a new storage backend, implement [`LightningCollections`][agentlightning.store.collection.LightningCollections]:
|
||||
|
||||
```python
|
||||
from agentlightning.store.collection import LightningCollections, Collection, Queue, KeyValue
|
||||
|
||||
class MyLightningCollections(LightningCollections):
|
||||
@property
|
||||
def rollouts(self) -> Collection[Rollout]:
|
||||
return self._rollouts # your implementation
|
||||
|
||||
@property
|
||||
def rollout_queue(self) -> Queue[str]:
|
||||
return self._queue # your implementation
|
||||
|
||||
# ... implement remaining properties
|
||||
|
||||
async def atomic(self, *, mode, snapshot=False, labels=None, **kwargs):
|
||||
# provide transaction / locking semantics
|
||||
...
|
||||
```
|
||||
|
||||
Then instantiate your store:
|
||||
|
||||
```python
|
||||
from agentlightning.store.collection_based import CollectionBasedLightningStore
|
||||
|
||||
store = CollectionBasedLightningStore(collections=MyLightningCollections())
|
||||
```
|
||||
|
||||
The store layer handles all business logic; your collections just need to provide correct CRUD semantics.
|
||||
|
||||
## Collection-based Store Implementations
|
||||
|
||||
Agent-lightning ships with two collection-based store implementations:
|
||||
|
||||
### InMemoryLightningStore
|
||||
|
||||
[`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] uses [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] backed by Python data structures. It supports **fast startup** with zero external dependencies—ideal for local development, CI, and unit tests. It also provides two lock modes, configurable between `"asyncio"` (single-thread, multiple coroutines) and `"thread"` (multi-threaded via [aiologic](https://github.com/x42005e1f/aiologic)).
|
||||
[`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] use nested dictionaries for O(1) primary-key lookup and `deque` for the task queue.
|
||||
|
||||
### MongoLightningStore
|
||||
|
||||
[`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] uses [`MongoLightningCollections`][agentlightning.store.collection.mongo.MongoLightningCollections] backed by MongoDB. It supports **persistent storage** suitable for production deployments and **multi-process safe** via database-level atomicity. It also supports **partition support** via `partition_id` for running multiple trainers against the same database.
|
||||
|
||||
```python
|
||||
from agentlightning.store.mongo import MongoLightningStore
|
||||
|
||||
store = MongoLightningStore(
|
||||
mongo_uri="mongodb://localhost:27017/?replicaSet=rs0",
|
||||
database_name="agentlightning",
|
||||
partition_id="trainer-1", # optional: isolate data per trainer
|
||||
)
|
||||
```
|
||||
|
||||
!!! note
|
||||
|
||||
[`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] requires the `mongo` optional dependency. Install with `pip install agentlightning[mongo]`.
|
||||
|
||||
### Capabilities
|
||||
|
||||
[](){ #store-capabilities }
|
||||
|
||||
Different stores have different capabilities. Check the [`capabilities`][agentlightning.LightningStore.capabilities] property to understand what a store supports:
|
||||
|
||||
| Capability | Description | InMemory | Mongo | Server | Client |
|
||||
|------------|-------------|----------|-------|--------|--------|
|
||||
| `thread_safe` | Safe for concurrent access from multiple threads | configurable | ✓ | ✓ | ✓ |
|
||||
| `async_safe` | Safe for concurrent access from multiple coroutines | ✓ | ✓ | ✓ | ✓ |
|
||||
| `zero_copy` | Can be shared across processes without serialization | ✗ | ✓ | ✓ | ✓ |
|
||||
| `otlp_traces` | Exposes an OTLP-compatible `/v1/traces` endpoint | ✗ | ✗ | ✓ | ✓ |
|
||||
Different store implementations may have different capabilities. For example, [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] does not support exporting traces via OTLP. Try to distinguish the capabilities of a store implementation by checking the [`capabilities`][agentlightning.LightningStore.capabilities] property.
|
||||
|
||||
## Thread Safety
|
||||
|
||||
Thread safety can be achieved at different layers:
|
||||
|
||||
**At the collections layer**: [`InMemoryLightningCollections`][agentlightning.store.collection.InMemoryLightningCollections] accepts a `lock_type` parameter:
|
||||
|
||||
- `"asyncio"` – Uses per-event-loop `asyncio.Lock` for single-threaded, multi-coroutine scenarios.
|
||||
- `"thread"` – Uses `aiologic.Lock` for true multi-threaded access.
|
||||
|
||||
**At the store layer**: [`LightningStoreThreaded`][agentlightning.LightningStoreThreaded] wraps any [`LightningStore`][agentlightning.LightningStore] to add mutex-based thread safety:
|
||||
**[`LightningStoreThreaded`][agentlightning.LightningStoreThreaded]** is a subclass of [`LightningStore`][agentlightning.LightningStore] that wraps another underlying store to make a store instance safe for multi-threaded callers. It wraps every state-mutating call in a mutex. Specifically:
|
||||
|
||||
* Methods like [`start_rollout`][agentlightning.LightningStore.start_rollout], [`enqueue_rollout`][agentlightning.LightningStore.enqueue_rollout], [`update_attempt`][agentlightning.LightningStore.update_attempt], [`add_span`][agentlightning.LightningStore.add_span], etc. are guarded by a lock.
|
||||
* Non-mutating, potentially blocking calls remain pass-through by design (e.g., [`wait_for_rollouts`][agentlightning.LightningStore.wait_for_rollouts]), as they don't modify shared state and should not hold the lock for long periods.
|
||||
|
||||
Database-based stores like [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] are inherently thread-safe through database atomicity guarantees.
|
||||
* Non-mutating, potentially blocking calls remain pass-through by design (e.g., [`wait_for_rollouts`][agentlightning.LightningStore.wait_for_rollouts]), as they don’t modify shared state and should not hold the lock for long periods.
|
||||
|
||||
## Process Safety and Client-server Store
|
||||
|
||||
@@ -376,7 +188,7 @@ Database-based stores like [`MongoLightningStore`][agentlightning.store.mongo.Mo
|
||||
|
||||
The server tracks the creator PID. In the owner process it delegates directly to the in-memory store; in other processes it lazily constructs a [`LightningStoreClient`][agentlightning.LightningStoreClient] to talk to the HTTP API. This prevents accidental cross-process mutation of the wrong memory image. When the server is pickled (e.g., via `multiprocessing`), only the minimal fields are serialized, but **NOT** the FastAPI/uvicorn objects. Subprocesses won’t accidentally carry live server state. Forked subprocess should also use [`LightningStoreClient`][agentlightning.LightningStoreClient] to communicate with the server in the main process.
|
||||
|
||||
On the client side, the client retries network/5xx failures using a small backoff, and probes `/v1/agl/health` between attempts. Application exceptions inside the server are wrapped as HTTP 400 with a traceback—these are **not retried**. The client also maintains a **per-event-loop** `aiohttp.ClientSession` map so that tracer callbacks (often on separate loops/threads) don’t hang by reusing a session from another loop.
|
||||
On the client side, the client retries network/5xx failures using a small backoff, and probes `/health` between attempts. Application exceptions inside the server are wrapped as HTTP 400 with a traceback—these are **not retried**. The client also maintains a **per-event-loop** `aiohttp.ClientSession` map so that tracer callbacks (often on separate loops/threads) don’t hang by reusing a session from another loop.
|
||||
|
||||
Minimal lifecycle:
|
||||
|
||||
|
||||
@@ -30,14 +30,6 @@
|
||||
|
||||
[:octicons-repo-24: Browse source]({{ src("examples/calc_x") }})
|
||||
|
||||
- :material-chart-box:{ .lg .middle } __ChartQA vision-language RL__
|
||||
|
||||
---
|
||||
|
||||
LangGraph-powered workflow for answering chart questions end to end: rollout the multi-modality agent with GPT or vLLM, and train with VERL/GRPO plus self-refinement loops.
|
||||
|
||||
[:octicons-repo-24: Browse source]({{ src("examples/chartqa") }})
|
||||
|
||||
- :material-code-braces:{ .lg .middle } __Claude Code SWE-bench__
|
||||
|
||||
---
|
||||
@@ -62,6 +54,14 @@
|
||||
|
||||
[:octicons-repo-24: Browse source]({{ src("examples/rag") }})
|
||||
|
||||
- :material-magnify:{ .lg .middle } __Search-R1 RL__
|
||||
|
||||
---
|
||||
|
||||
Reproduction of the Search-R1 workflow that prepares its own retrieval backend, runs the rollout script, and coordinates GRPO-style training without extra orchestration layers (last validated on v0.1.x).
|
||||
|
||||
[:octicons-repo-24: Browse source]({{ src("examples/search_r1") }})
|
||||
|
||||
- :material-database:{ .lg .middle } __Spider SQL agent__
|
||||
|
||||
---
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Train SQL Agent with Agent-lightning and VERL
|
||||
|
||||
This walkthrough builds upon the **Agent-lightning SQL Agent** example and explains how the system components integrate: a **LangGraph-based SQL agent** wrapped as a [`LitAgent`][agentlightning.LitAgent], the **[`VERL`][agentlightning.algorithm.verl.VERL] reinforcement learning (RL) algorithm**, and the **[`Trainer`][agentlightning.Trainer]**, which coordinates both training and debugging.
|
||||
This walkthrough builds upon the **Agent-lightning v0.2 SQL Agent** example and explains how the system components integrate: a **LangGraph-based SQL agent** wrapped as a [`LitAgent`][agentlightning.LitAgent], the **[`VERL`][agentlightning.algorithm.verl.VERL] reinforcement learning (RL) algorithm**, and the **[`Trainer`][agentlightning.Trainer]**, which coordinates both training and debugging.
|
||||
|
||||
The command-line interface in [`examples/spider/train_sql_agent.py`]({{ src("examples/spider/train_sql_agent.py") }}) provides a complete runnable example. However, this document focuses on understanding the underlying architecture so you can effectively adapt the workflow to your own agents.
|
||||
|
||||
|
||||
@@ -35,8 +35,6 @@ This documentation is organized into the following parts:
|
||||
|
||||
- [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
|
||||
- [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
|
||||
- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl) and their blog [*Stop Wrestling with Your Agent RL: How Youtu-Agent Achieved Stable, 128-GPU Scaling Without Breaking a Sweat*](https://spotted-coconut-df8.notion.site/Stop-Wrestling-with-Your-Agent-RL-How-Youtu-Agent-Achieved-Stable-128-GPU-Scaling-Without-Breaking-2ca5e8f089ba80539a98c582b65e0233).
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
|
||||
@@ -22,8 +22,6 @@
|
||||
|
||||
## Emitter
|
||||
|
||||
::: agentlightning.operation
|
||||
|
||||
::: agentlightning.emit_annotation
|
||||
|
||||
::: agentlightning.emit_reward
|
||||
|
||||
+17
-46
@@ -1,5 +1,7 @@
|
||||
# Command Line Interface
|
||||
|
||||
<!-- TODO: This document should be auto-generated. -->
|
||||
|
||||
!!! warning
|
||||
|
||||
This document is a work in progress and might not be updated with the latest changes.
|
||||
@@ -12,18 +14,17 @@
|
||||
## agl
|
||||
|
||||
```text
|
||||
usage: agl [-h] {vllm,store,prometheus,agentops}
|
||||
usage: agl [-h] {vllm,store,agentops}
|
||||
|
||||
Agent Lightning CLI entry point.
|
||||
|
||||
Available subcommands:
|
||||
vllm Run the vLLM CLI with Agent Lightning instrumentation.
|
||||
store Run a LightningStore server.
|
||||
prometheus Serve Prometheus metrics from the multiprocess registry.
|
||||
agentops Start the AgentOps server manager.
|
||||
vllm Run the vLLM CLI with Agent Lightning instrumentation.
|
||||
store Run a LightningStore server.
|
||||
agentops Start the AgentOps server manager.
|
||||
|
||||
positional arguments:
|
||||
{vllm,store,prometheus,agentops}
|
||||
{vllm,store,agentops}
|
||||
Subcommand to run.
|
||||
|
||||
options:
|
||||
@@ -63,56 +64,26 @@ Agent-lightning's LightningStore CLI. Use it to start an independent LightningSt
|
||||
Currently the store data are stored in memory and will be lost when the server is stopped.
|
||||
|
||||
```text
|
||||
usage: agl store [-h] [--host HOST] [--port PORT] [--cors-origin CORS_ORIGINS] [--log-level {DEBUG,INFO,WARNING,ERROR}] [--tracker {prometheus,console} [{prometheus,console} ...]] [--n-workers N_WORKERS] [--backend {memory,mongo}]
|
||||
[--mongo-uri MONGO_URI]
|
||||
usage: agl store [-h] [--port PORT]
|
||||
|
||||
Run a LightningStore server
|
||||
|
||||
options:
|
||||
-h, --help show this help message and exit
|
||||
--host HOST Host to bind the server to
|
||||
--port PORT Port to run the server on
|
||||
--cors-origin CORS_ORIGINS
|
||||
Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.
|
||||
--log-level {DEBUG,INFO,WARNING,ERROR}
|
||||
Configure the logging level for the store.
|
||||
--tracker {prometheus,console} [{prometheus,console} ...]
|
||||
Enable metrics tracking. Repeat for multiple trackers.
|
||||
--n-workers N_WORKERS
|
||||
Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. Only applicable for zero-copy stores such as MongoDB backend.
|
||||
--backend {memory,mongo}
|
||||
Backend to use for the store.
|
||||
--mongo-uri MONGO_URI
|
||||
MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.
|
||||
-h, --help show this help message and exit
|
||||
--port PORT Port to run the server on
|
||||
```
|
||||
|
||||
!!! tip
|
||||
## agl agentops
|
||||
|
||||
After launching the store via CLI, you can tell the [`Trainer`][agentlightning.Trainer] to use the store by passing the store address to the trainer.
|
||||
|
||||
```python
|
||||
store_client = agl.LightningStoreClient("http://localhost:4747")
|
||||
trainer = agl.Trainer(store=store_client, ...)
|
||||
```
|
||||
|
||||
See [using external store][debug-with-external-store] for more details.
|
||||
|
||||
## agl prometheus
|
||||
|
||||
Expose the Prometheus multiprocess registry on a dedicated FastAPI server. This is useful when the main LightningStore service is under heavy load; exporters can scrape this auxiliary endpoint instead.
|
||||
Start a mock AgentOps server to bypass the online service of AgentOps.
|
||||
|
||||
```text
|
||||
usage: agl prometheus [-h] [--host HOST] [--port PORT] [--metrics-path METRICS_PATH] [--log-level {DEBUG,INFO,WARNING,ERROR}] [--access-log]
|
||||
usage: agl agentops [-h] [--daemon] [--port PORT]
|
||||
|
||||
Serve Prometheus metrics outside the LightningStore server.
|
||||
Start AgentOps server
|
||||
|
||||
options:
|
||||
-h, --help show this help message and exit
|
||||
--host HOST Host to bind the metrics server to.
|
||||
--port PORT Port to expose the Prometheus metrics on.
|
||||
--metrics-path METRICS_PATH
|
||||
HTTP path used to expose metrics. Must start with '/' and not be the root path.
|
||||
--log-level {DEBUG,INFO,WARNING,ERROR}
|
||||
Configure the logging level for the metrics server.
|
||||
--access-log Enable uvicorn access logs. Disabled by default to reduce noise.
|
||||
-h, --help show this help message and exit
|
||||
--daemon Run server as a daemon
|
||||
--port PORT Port to run the server on
|
||||
```
|
||||
|
||||
@@ -24,8 +24,6 @@
|
||||
|
||||
::: agentlightning.litagent.decorator.prompt_rollout
|
||||
|
||||
::: agentlightning.emitter.annotation.OperationContext
|
||||
|
||||
## LLM Proxy
|
||||
|
||||
::: agentlightning.llm_proxy.ModelConfig
|
||||
@@ -46,14 +44,48 @@
|
||||
|
||||
::: agentlightning.store.base.UNSET
|
||||
|
||||
::: agentlightning.store.utils.rollout_status_from_attempt
|
||||
|
||||
::: agentlightning.store.utils.scan_unhealthy_rollouts
|
||||
::: agentlightning.store.utils.propagate_status
|
||||
|
||||
## Tracing and OpenTelemetry
|
||||
|
||||
::: agentlightning.tracer.otel.LightningSpanProcessor
|
||||
|
||||
## Utilities
|
||||
|
||||
::: agentlightning.utils.server_launcher.PythonServerLauncher
|
||||
|
||||
::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
|
||||
|
||||
::: agentlightning.utils.server_launcher.LaunchMode
|
||||
|
||||
::: agentlightning.utils.otel.full_qualified_name
|
||||
|
||||
::: agentlightning.utils.otel.get_tracer_provider
|
||||
|
||||
::: agentlightning.utils.otel.get_tracer
|
||||
|
||||
::: agentlightning.utils.otel.make_tag_attributes
|
||||
|
||||
::: agentlightning.utils.otel.extract_tags_from_attributes
|
||||
|
||||
::: agentlightning.utils.otel.make_link_attributes
|
||||
|
||||
::: agentlightning.utils.otel.query_linked_spans
|
||||
|
||||
::: agentlightning.utils.otel.extract_links_from_attributes
|
||||
|
||||
::: agentlightning.utils.otel.filter_attributes
|
||||
|
||||
::: agentlightning.utils.otel.filter_and_unflatten_attributes
|
||||
|
||||
::: agentlightning.utils.otel.flatten_attributes
|
||||
|
||||
::: agentlightning.utils.otel.unflatten_attributes
|
||||
|
||||
::: agentlightning.utils.otlp.handle_otlp_export
|
||||
|
||||
::: agentlightning.utils.otlp.spans_from_proto
|
||||
|
||||
## Deprecated APIs
|
||||
|
||||
::: agentlightning.emitter.reward.reward
|
||||
|
||||
@@ -17,15 +17,3 @@
|
||||
::: agentlightning.OtelTracer
|
||||
|
||||
::: agentlightning.Tracer
|
||||
|
||||
::: agentlightning.tracer.weave.WeaveTracer
|
||||
|
||||
::: agentlightning.DummyTracer
|
||||
|
||||
::: agentlightning.set_active_tracer
|
||||
|
||||
::: agentlightning.get_active_tracer
|
||||
|
||||
::: agentlightning.clear_active_tracer
|
||||
|
||||
::: agentlightning.tracer.weave.WeaveTracer
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
# Semantic Conventions
|
||||
|
||||
::: agentlightning.semconv
|
||||
@@ -8,8 +8,6 @@
|
||||
|
||||
::: agentlightning.InMemoryLightningStore
|
||||
|
||||
::: agentlightning.store.mongo.MongoLightningStore
|
||||
|
||||
::: agentlightning.CollectionBasedLightningStore
|
||||
|
||||
## Client-Server and Thread-safe Wrappers
|
||||
@@ -22,10 +20,6 @@
|
||||
|
||||
## Collections and Collection Implementations
|
||||
|
||||
::: agentlightning.store.collection.AtomicMode
|
||||
|
||||
::: agentlightning.store.collection.AtomicLabels
|
||||
|
||||
::: agentlightning.store.collection.Collection
|
||||
|
||||
::: agentlightning.store.collection.Queue
|
||||
@@ -41,13 +35,3 @@
|
||||
::: agentlightning.store.collection.DictBasedKeyValue
|
||||
|
||||
::: agentlightning.store.collection.InMemoryLightningCollections
|
||||
|
||||
::: agentlightning.store.collection.mongo.MongoBasedCollection
|
||||
|
||||
::: agentlightning.store.collection.mongo.MongoBasedQueue
|
||||
|
||||
::: agentlightning.store.collection.mongo.MongoBasedKeyValue
|
||||
|
||||
::: agentlightning.store.collection.mongo.MongoClientPool
|
||||
|
||||
::: agentlightning.store.collection.mongo.MongoLightningCollections
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user