Compare commits
46 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 901d041d56 | |||
| 2acb0b9864 | |||
| 23e3336016 | |||
| 88555bafdb | |||
| 9dbf3d362c | |||
| b6b16c1827 | |||
| 66ea11b9bb | |||
| d0a19cf5c9 | |||
| 23f7edb6d7 | |||
| 69e9b1c90e | |||
| 8ebc28702d | |||
| ccf5e3337d | |||
| 8c336201a4 | |||
| cb7edf48d3 | |||
| 8c236b89ad | |||
| 9c7a6316f1 | |||
| 01de392b29 | |||
| 35d5763c1f | |||
| 779a971a95 | |||
| bfb94a8750 | |||
| 25eda47a29 | |||
| a214474402 | |||
| 3b5d733861 | |||
| 158f5df28e | |||
| 40dc59205b | |||
| c1a43b6c3a | |||
| 4235731a0d | |||
| 22b80b38bf | |||
| 9f178accaf | |||
| 68a47d5087 | |||
| 4b36b25aad | |||
| a13e09fc6c | |||
| e63c340ebd | |||
| e62b7ca252 | |||
| f66d87745f | |||
| 52090e9dd5 | |||
| fdaf3f1777 | |||
| 087c7d350a | |||
| 2203070ef0 | |||
| 28b0895fe2 | |||
| c39a84d12c | |||
| 2bf90cda51 | |||
| 26cc501233 | |||
| 5706da6908 | |||
| 2477c4bca1 | |||
| 03f08b422c |
+517
-232
@@ -3,10 +3,13 @@ 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 }})
|
||||
name: ${{ matrix.workload.kind }} (${{ matrix.backend.id }}, ${{ matrix.workload.display }}, ${{ matrix.trace_sink }})
|
||||
runs-on: ${{ matrix.workload.runner }}
|
||||
timeout-minutes: ${{ matrix.workload.timeout }}
|
||||
strategy:
|
||||
@@ -17,105 +20,107 @@ jobs:
|
||||
compose_file: compose.prometheus-memory-store.yml
|
||||
- id: mongo
|
||||
compose_file: compose.prometheus-mongo-store.yml
|
||||
trace_sink: [store, kafka, clickhouse]
|
||||
# trace_sink: [clickhouse]
|
||||
workload:
|
||||
- id: scenario-minimal-scale
|
||||
display: Minimal production scale
|
||||
kind: scenario
|
||||
store_workers: 4
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 4096
|
||||
--batch-size 256
|
||||
--n-runners 32
|
||||
--max-rounds 6
|
||||
--sleep-seconds 0.5
|
||||
- id: scenario-medium-scale
|
||||
display: Medium production scale
|
||||
kind: scenario
|
||||
store_workers: 16
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 10000
|
||||
--batch-size 1000
|
||||
--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
|
||||
display: Large batch waves
|
||||
kind: scenario
|
||||
store_workers: 32
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch
|
||||
--total-tasks 100000
|
||||
--batch-size 8192
|
||||
--n-runners 256
|
||||
--max-rounds 6
|
||||
--sleep-seconds 0.1
|
||||
# - id: scenario-minimal-scale
|
||||
# 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
|
||||
# --batch-size 256
|
||||
# --n-runners 32
|
||||
# --max-rounds 6
|
||||
# --sleep-seconds 0.5
|
||||
# - id: scenario-medium-scale
|
||||
# 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
|
||||
# --batch-size 1000
|
||||
# --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 300
|
||||
# --max-rounds 6
|
||||
# --sleep-seconds 0.1
|
||||
# - id: scenario-large-batch
|
||||
# display: Large batch waves
|
||||
# kind: scenario
|
||||
# store_workers: 96
|
||||
# runner:
|
||||
# - self-hosted
|
||||
# - 1ES.Pool=agl-runner-cpu-high
|
||||
# timeout: 120
|
||||
# args: >-
|
||||
# --mode batch
|
||||
# --total-tasks 50000
|
||||
# --batch-size 8192
|
||||
# --n-runners 1000
|
||||
# --max-rounds 3
|
||||
# --sleep-seconds 0.1
|
||||
- id: scenario-long-queues
|
||||
display: Long rollout queues
|
||||
kind: scenario
|
||||
store_workers: 32
|
||||
store_workers: 48
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
timeout: 120
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
--total-tasks 100000
|
||||
--total-tasks 50000
|
||||
--batch-size 1024
|
||||
--n-runners 256
|
||||
--remaining-tasks 4096
|
||||
--max-rounds 4
|
||||
--sleep-seconds 0.1
|
||||
- id: scenario-high-concurrency
|
||||
display: High-throughput concurrent requests
|
||||
kind: scenario
|
||||
store_workers: 32
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode single
|
||||
--total-tasks 100000
|
||||
--concurrency 2048
|
||||
--n-runners 256
|
||||
--max-rounds 2
|
||||
--sleep-seconds 0.1
|
||||
# - id: scenario-high-concurrency
|
||||
# display: High-throughput concurrent requests
|
||||
# kind: scenario
|
||||
# store_workers: 96
|
||||
# runner:
|
||||
# - self-hosted
|
||||
# - 1ES.Pool=agl-runner-cpu-high
|
||||
# timeout: 120
|
||||
# args: >-
|
||||
# --mode single
|
||||
# --total-tasks 50000
|
||||
# --concurrency 2048
|
||||
# --n-runners 256
|
||||
# --max-rounds 2
|
||||
# --sleep-seconds 0.1
|
||||
- id: scenario-heavy-traces
|
||||
display: Heavy rollouts with deep traces
|
||||
kind: scenario
|
||||
store_workers: 64
|
||||
store_workers: 96
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
- 1ES.Pool=agl-runner-cpu-high
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
@@ -126,62 +131,65 @@ jobs:
|
||||
--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
|
||||
# - 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 }}
|
||||
BACKEND_ID: ${{ matrix.backend.id }}
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
TRACE_SINK_ID: ${{ matrix.trace_sink }}
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}-{2}', matrix.workload.id, matrix.backend.id, matrix.trace_sink) }}
|
||||
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) }}
|
||||
ANALYSIS_FILE: ${{ format('analysis-{0}-{1}.log', matrix.workload.id, matrix.trace_sink) }}
|
||||
SUMMARY_FILE: ${{ format('summary-{0}-{1}.log', matrix.workload.id, matrix.trace_sink) }}
|
||||
PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
ARTIFACT_NAME: ${{ format('{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
ARTIFACT_NAME: ${{ format('{0}-{1}-{2}', matrix.workload.id, matrix.backend.id, matrix.trace_sink) }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
@@ -227,9 +235,217 @@ jobs:
|
||||
cd docker && docker compose -f "$COMPOSE_FILE" logs app
|
||||
exit 1
|
||||
|
||||
|
||||
- name: Configure trace sink (store vs otlp)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ "${{ matrix.trace_sink }}" = "kafka" ] || [ "${{ matrix.trace_sink }}" = "clickhouse" ]; then
|
||||
echo "AGL_OTLP_ENDPOINT=http://localhost:4318/v1/traces" >> "$GITHUB_ENV"
|
||||
fi
|
||||
|
||||
- name: Launch Kafka + OTel Collector (OTLP -> Kafka)
|
||||
if: ${{ matrix.trace_sink == 'kafka' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
|
||||
# Generate OTel Collector config (OTLP/HTTP receiver -> Kafka exporter)
|
||||
cat > otelcol-kafka.yml <<'YAML'
|
||||
receivers:
|
||||
otlp:
|
||||
protocols:
|
||||
http:
|
||||
endpoint: 0.0.0.0:4318
|
||||
processors:
|
||||
batch: {}
|
||||
exporters:
|
||||
kafka:
|
||||
brokers: ["kafka:9092"]
|
||||
topic: "agl-otlp-spans"
|
||||
encoding: otlp_proto
|
||||
service:
|
||||
pipelines:
|
||||
traces:
|
||||
receivers: [otlp]
|
||||
processors: [batch]
|
||||
exporters: [kafka]
|
||||
YAML
|
||||
|
||||
# Launch Kafka + Zookeeper + OTel Collector
|
||||
cat > compose.kafka-otel.yml <<'YAML'
|
||||
services:
|
||||
zookeeper:
|
||||
image: confluentinc/cp-zookeeper:7.6.1
|
||||
environment:
|
||||
ZOOKEEPER_CLIENT_PORT: 2181
|
||||
ZOOKEEPER_TICK_TIME: 2000
|
||||
|
||||
kafka:
|
||||
image: confluentinc/cp-kafka:7.6.1
|
||||
depends_on: [zookeeper]
|
||||
environment:
|
||||
KAFKA_BROKER_ID: 1
|
||||
KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
|
||||
KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
|
||||
KAFKA_LISTENERS: PLAINTEXT://0.0.0.0:9092
|
||||
KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:9092
|
||||
KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT
|
||||
# Enlarge max message size to accommodate large spans
|
||||
KAFKA_MESSAGE_MAX_BYTES: "10000000"
|
||||
KAFKA_REPLICA_FETCH_MAX_BYTES: "10000000"
|
||||
KAFKA_SOCKET_REQUEST_MAX_BYTES: "10000000"
|
||||
|
||||
otelcol:
|
||||
image: otel/opentelemetry-collector-contrib:latest
|
||||
depends_on: [kafka]
|
||||
command: ["--config=/etc/otelcol/config.yml"]
|
||||
# command:
|
||||
# - "--config=/etc/otelcol/config.yml"
|
||||
# - "--set=service.telemetry.logs.level=debug"
|
||||
volumes:
|
||||
- ./otelcol-kafka.yml:/etc/otelcol/config.yml:ro
|
||||
ports:
|
||||
- "4318:4318"
|
||||
YAML
|
||||
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml down -v || true
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml up -d --quiet-pull
|
||||
|
||||
# Create topic (idempotent)
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml exec -T kafka \
|
||||
kafka-topics --bootstrap-server kafka:9092 \
|
||||
--create --if-not-exists \
|
||||
--topic agl-otlp-spans --partitions 3 --replication-factor 1
|
||||
|
||||
# Wait for OTLP/HTTP port to be reachable on the host
|
||||
for attempt in {1..30}; do
|
||||
if (echo > /dev/tcp/127.0.0.1/4318) >/dev/null 2>&1; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "OTel Collector port 4318 not reachable in time" >&2
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml logs otelcol || true
|
||||
exit 1
|
||||
|
||||
- name: Launch ClickHouse + OTel Collector (OTLP -> ClickHouse)
|
||||
if: ${{ matrix.trace_sink == 'clickhouse' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
|
||||
# Generate OTel Collector config (OTLP/HTTP receiver -> ClickHouse exporter)
|
||||
cat > otelcol-clickhouse.yml <<'YAML'
|
||||
receivers:
|
||||
otlp:
|
||||
protocols:
|
||||
http:
|
||||
endpoint: 0.0.0.0:4318
|
||||
processors:
|
||||
batch: {}
|
||||
exporters:
|
||||
clickhouse:
|
||||
endpoint: tcp://clickhouse:9000?dial_timeout=10s&compress=lz4
|
||||
database: otel
|
||||
traces_table_name: otel_traces
|
||||
username: otel
|
||||
password: ${env:CLICKHOUSE_PASSWORD}
|
||||
service:
|
||||
pipelines:
|
||||
traces:
|
||||
receivers: [otlp]
|
||||
processors: [batch]
|
||||
exporters: [clickhouse]
|
||||
YAML
|
||||
|
||||
# Launch ClickHouse + OTel Collector
|
||||
cat > compose.clickhouse-otel.yml <<'YAML'
|
||||
services:
|
||||
clickhouse:
|
||||
image: clickhouse/clickhouse-server:latest
|
||||
environment:
|
||||
CLICKHOUSE_USER: "otel"
|
||||
CLICKHOUSE_PASSWORD: "changeme"
|
||||
CLICKHOUSE_DB: "otel"
|
||||
ulimits:
|
||||
nofile:
|
||||
soft: 262144
|
||||
hard: 262144
|
||||
ports:
|
||||
- "8123:8123"
|
||||
- "9000:9000"
|
||||
volumes:
|
||||
- ch_data:/var/lib/clickhouse
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "wget -qO- http://localhost:8123/ping | grep -q Ok"]
|
||||
interval: 2s
|
||||
timeout: 2s
|
||||
retries: 30
|
||||
start_period: 10s
|
||||
|
||||
otelcol:
|
||||
image: otel/opentelemetry-collector-contrib:latest
|
||||
depends_on:
|
||||
clickhouse:
|
||||
condition: service_healthy
|
||||
restart: unless-stopped
|
||||
command:
|
||||
- "--config=/etc/otelcol/config.yml"
|
||||
- "--set=service.telemetry.logs.level=debug"
|
||||
environment:
|
||||
CLICKHOUSE_PASSWORD: "changeme"
|
||||
volumes:
|
||||
- ./otelcol-clickhouse.yml:/etc/otelcol/config.yml:ro
|
||||
ports:
|
||||
- "4318:4318"
|
||||
|
||||
volumes:
|
||||
ch_data:
|
||||
|
||||
YAML
|
||||
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml down -v || true
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml up -d --quiet-pull
|
||||
|
||||
# Wait for OTLP/HTTP port to be reachable on the host
|
||||
for attempt in {1..30}; do
|
||||
if (echo > /dev/tcp/127.0.0.1/4318) >/dev/null 2>&1; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "OTel Collector port 4318 not reachable in time" >&2
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml logs otelcol || true
|
||||
exit 1
|
||||
|
||||
- name: Prepare artifact directory
|
||||
run: mkdir -p "$ARTIFACT_DIR"
|
||||
|
||||
- name: Kafka topic offsets (before workload)
|
||||
if: ${{ matrix.trace_sink == 'kafka' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
# Print offsets for common topic spellings; at least one should exist.
|
||||
{
|
||||
echo "== Kafka offsets BEFORE workload =="
|
||||
# docker compose -p agl-kafka -f compose.kafka-otel.yml exec -T kafka \
|
||||
# kafka-run-class kafka.tools.GetOffsetShell --broker-list kafka:9092 --topic agl_otlp_spans 2>/dev/null || true
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml exec -T kafka \
|
||||
kafka-run-class kafka.tools.GetOffsetShell --broker-list kafka:9092 --topic agl-otlp-spans 2>/dev/null || true
|
||||
} | tee "$GITHUB_WORKSPACE/$ARTIFACT_DIR/kafka-offsets-before.txt"
|
||||
|
||||
|
||||
- name: ClickHouse row count (before workload)
|
||||
if: ${{ matrix.trace_sink == 'clickhouse' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
{
|
||||
echo "== ClickHouse rows BEFORE workload =="
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml exec -T clickhouse clickhouse-client -q "SELECT count() FROM otel.otel_traces" 2>/dev/null || true
|
||||
} | tee "$GITHUB_WORKSPACE/$ARTIFACT_DIR/clickhouse-rows-before.txt"
|
||||
|
||||
- name: Record workload start
|
||||
run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
@@ -255,6 +471,30 @@ jobs:
|
||||
if: ${{ always() }}
|
||||
run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Kafka topic offsets (after workload)
|
||||
if: ${{ always() && matrix.trace_sink == 'kafka' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
{
|
||||
echo "== Kafka offsets AFTER workload =="
|
||||
# docker compose -p agl-kafka -f compose.kafka-otel.yml exec -T kafka \
|
||||
# kafka-run-class kafka.tools.GetOffsetShell --broker-list kafka:9092 --topic agl_otlp_spans 2>/dev/null || true
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml exec -T kafka \
|
||||
kafka-run-class kafka.tools.GetOffsetShell --broker-list kafka:9092 --topic agl-otlp-spans 2>/dev/null || true
|
||||
} | tee "$GITHUB_WORKSPACE/$ARTIFACT_DIR/kafka-offsets-after.txt"
|
||||
|
||||
|
||||
- name: ClickHouse row count (after workload)
|
||||
if: ${{ always() && matrix.trace_sink == 'clickhouse' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
{
|
||||
echo "== ClickHouse rows AFTER workload =="
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml exec -T clickhouse clickhouse-client -q "SELECT count() FROM otel.otel_traces" 2>/dev/null || true
|
||||
} | tee "$GITHUB_WORKSPACE/$ARTIFACT_DIR/clickhouse-rows-after.txt"
|
||||
|
||||
- name: Show micro benchmark summary
|
||||
if: ${{ always() && matrix.workload.kind == 'micro' }}
|
||||
run: |
|
||||
@@ -298,6 +538,51 @@ jobs:
|
||||
docker compose -f "$COMPOSE_FILE" logs "$service" > "../$ARTIFACT_DIR/docker-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
|
||||
done
|
||||
|
||||
|
||||
- name: Collect Kafka + OTel Collector logs
|
||||
if: ${{ always() && matrix.trace_sink == 'kafka' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
cd docker
|
||||
if [ -f compose.kafka-otel.yml ]; then
|
||||
for service in zookeeper kafka otelcol; do
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml logs "$service" \
|
||||
> "../$ARTIFACT_DIR/docker-kafka-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
|
||||
done
|
||||
fi
|
||||
|
||||
|
||||
- name: Collect ClickHouse + OTel Collector logs
|
||||
if: ${{ always() && matrix.trace_sink == 'clickhouse' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
cd docker
|
||||
if [ -f compose.clickhouse-otel.yml ]; then
|
||||
for service in clickhouse otelcol; do
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml logs "$service" > "../$ARTIFACT_DIR/docker-clickhouse-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
|
||||
done
|
||||
fi
|
||||
|
||||
- name: Stop Kafka + OTel Collector
|
||||
if: ${{ always() && matrix.trace_sink == 'kafka' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
if [ -f compose.kafka-otel.yml ]; then
|
||||
docker compose -p agl-kafka -f compose.kafka-otel.yml down -v || true
|
||||
fi
|
||||
|
||||
- name: Stop ClickHouse + OTel Collector
|
||||
if: ${{ always() && matrix.trace_sink == 'clickhouse' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
if [ -f compose.clickhouse-otel.yml ]; then
|
||||
docker compose -p agl-clickhouse -f compose.clickhouse-otel.yml down -v || true
|
||||
fi
|
||||
|
||||
- name: Stop ${{ matrix.backend.id }} Prometheus stack
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
@@ -322,115 +607,115 @@ jobs:
|
||||
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: 100000
|
||||
concurrency: 2048
|
||||
type: insert
|
||||
- id: medium-insert
|
||||
total_tasks: 100000
|
||||
concurrency: 128
|
||||
type: insert
|
||||
- id: low-insert
|
||||
total_tasks: 100000
|
||||
concurrency: 4
|
||||
type: insert
|
||||
- id: high-dequeue
|
||||
total_tasks: 100000
|
||||
concurrency: 2048
|
||||
type: dequeue
|
||||
- id: medium-dequeue
|
||||
total_tasks: 100000
|
||||
concurrency: 128
|
||||
type: dequeue
|
||||
- id: low-dequeue
|
||||
total_tasks: 100000
|
||||
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
|
||||
# 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'
|
||||
# - 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: 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: 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: 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: 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: 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 }}
|
||||
path: ${{ env.ARTIFACT_DIR }}
|
||||
if-no-files-found: error
|
||||
# - name: Upload collection artifacts
|
||||
# if: ${{ always() }}
|
||||
# uses: actions/upload-artifact@v4
|
||||
# with:
|
||||
# name: ${{ env.ARTIFACT_NAME }}
|
||||
# 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 --group torch-gpu-stable
|
||||
--group dev --group experiment --group agents --extra weave --extra mongo --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 torch-gpu-${{ matrix.setup-script }}
|
||||
--group dev --group experiment --group agents --extra weave --extra mongo --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
@@ -270,6 +270,33 @@ 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
|
||||
@@ -295,6 +322,52 @@ jobs:
|
||||
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
|
||||
|
||||
@@ -2,8 +2,25 @@ 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: >
|
||||
|
||||
@@ -45,6 +45,12 @@ jobs:
|
||||
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
|
||||
@@ -54,7 +60,7 @@ jobs:
|
||||
# Other uncovered tests
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not gpu and not llmproxy'
|
||||
pytest-mark: 'not store and not agentops and not weave and not gpu and not llmproxy'
|
||||
runs-on: ubuntu-latest
|
||||
has-gpu: false
|
||||
env:
|
||||
@@ -83,24 +89,24 @@ jobs:
|
||||
|
||||
- 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 mongo --group dev --group agents --group langchain --group torch-gpu-stable
|
||||
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 mongo --group dev --group agents --group langchain --group core-stable
|
||||
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 mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.env.setup-script }}
|
||||
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 mongo --group dev --group agents --group langchain --group core-stable
|
||||
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 mongo --group dev --group agents --group torch-gpu-legacy
|
||||
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 mongo --group dev --group agents --group core-legacy
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group core-legacy
|
||||
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
@@ -126,49 +132,7 @@ jobs:
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Setup Docker environments
|
||||
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
|
||||
run: ./scripts/mongodb_docker_run.sh
|
||||
shell: bash
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
@@ -181,7 +145,7 @@ jobs:
|
||||
# mongo, openai, gpu, all enabled by default
|
||||
- name: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}"
|
||||
uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}${{ matrix.env.setup-script == 'legacy' && ' and not langchain' || '' }}"
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
@@ -270,6 +234,14 @@ 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
|
||||
|
||||
@@ -19,7 +19,7 @@ jobs:
|
||||
lint:
|
||||
strategy:
|
||||
matrix:
|
||||
setup: [fast, slow]
|
||||
setup: [fast, slow, next]
|
||||
fail-fast: false
|
||||
name: Lint - ${{ matrix.setup }}
|
||||
runs-on: ubuntu-latest
|
||||
@@ -33,10 +33,14 @@ 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 \
|
||||
@@ -47,7 +51,7 @@ jobs:
|
||||
--group agents \
|
||||
--group langchain \
|
||||
--no-default-groups
|
||||
if: matrix.setup == 'slow'
|
||||
if: matrix.setup != 'fast'
|
||||
# This pre-commit skips JavaScript on purpose.
|
||||
- name: Run pre-commit
|
||||
uses: pre-commit/action@v3.0.1
|
||||
@@ -62,7 +66,7 @@ jobs:
|
||||
if: matrix.setup == 'fast'
|
||||
- name: Run pyright (slow)
|
||||
run: uv run --locked --no-sync pyright -p pyrightconfig.json
|
||||
if: matrix.setup == 'slow'
|
||||
if: matrix.setup != 'fast'
|
||||
|
||||
lint-js:
|
||||
name: Lint - JavaScript
|
||||
@@ -107,6 +111,10 @@ 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
|
||||
@@ -128,6 +136,10 @@ jobs:
|
||||
- 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
|
||||
@@ -139,7 +151,7 @@ jobs:
|
||||
# unmarked tests: adapter, execution engine, etc.
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not llmproxy and not utils'
|
||||
pytest-mark: 'not store and not agentops and not weave and not llmproxy and not utils'
|
||||
env:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
@@ -164,10 +176,10 @@ jobs:
|
||||
run: uv lock --upgrade
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-stable
|
||||
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'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
|
||||
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'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
|
||||
+2
-1
@@ -1,7 +1,8 @@
|
||||
# Agentlightning specific files
|
||||
verl_old
|
||||
meta-llama/**
|
||||
debug/*.png
|
||||
**/debug/**/*.png
|
||||
**/debug/**/*.json
|
||||
requirements-freeze*.txt
|
||||
/playground
|
||||
|
||||
|
||||
@@ -3,13 +3,14 @@ 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$)
|
||||
exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)|(.*store-openapi\.json$)
|
||||
- id: check-shebang-scripts-are-executable
|
||||
- id: detect-private-key
|
||||
- repo: https://github.com/pycqa/isort
|
||||
|
||||
@@ -37,7 +37,7 @@ pip install agentlightning
|
||||
For the latest nightly build (cutting-edge features), you can install from Test PyPI:
|
||||
|
||||
```bash
|
||||
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
|
||||
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre agentlightning
|
||||
```
|
||||
|
||||
Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
|
||||
@@ -46,6 +46,7 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.
|
||||
|
||||
## ⚡ Articles
|
||||
|
||||
- 12/17/2025 [Adopting the Trajectory Level Aggregation for Faster Training](https://agent-lightning.github.io/posts/trajectory_level_aggregation/) Agent-lightning blog.
|
||||
- 11/4/2025 [Tuning ANY AI agent with Tinker ✕ Agent-lightning](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-1-1d8c9a397f0e) Medium. See also [Part 2](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-2-332c5437f0dc).
|
||||
- 10/22/2025 [No More Retokenization Drift: Returning Token IDs via the OpenAI Compatible API Matters in Agent RL](https://blog.vllm.ai/2025/10/22/agent-lightning.html) vLLM blog. See also [Zhihu writeup](https://zhuanlan.zhihu.com/p/1965067274642785725).
|
||||
- 8/11/2025 [Training AI Agents to Write and Self-correct SQL with Reinforcement Learning](https://medium.com/@yugez/training-ai-agents-to-write-and-self-correct-sql-with-reinforcement-learning-571ed31281ad) Medium.
|
||||
@@ -57,7 +58,7 @@ 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).
|
||||
- [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
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
__version__ = "0.3.0"
|
||||
__version__ = "0.3.1"
|
||||
|
||||
from .adapter import *
|
||||
from .algorithm import *
|
||||
|
||||
@@ -12,6 +12,7 @@ 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
|
||||
|
||||
@@ -20,6 +21,47 @@ 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.
|
||||
|
||||
@@ -132,7 +174,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
|
||||
@@ -309,6 +351,19 @@ 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.
|
||||
|
||||
@@ -328,7 +383,17 @@ class TraceTree:
|
||||
`True` when the span payload describes a reward, otherwise `False`.
|
||||
"""
|
||||
maybe_reward = self.maybe_reward_dict()
|
||||
return maybe_reward and maybe_reward.get("type") == "reward" # type: ignore
|
||||
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
|
||||
|
||||
def find_llm_calls(
|
||||
self,
|
||||
@@ -365,7 +430,9 @@ class TraceTree:
|
||||
is_llm_call = False
|
||||
if is_llm_call:
|
||||
# Check the response id
|
||||
response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id") # type: ignore
|
||||
response_id = _attributes_get_multiple(
|
||||
self.span.attributes, ["gen_ai.response.id", "agentlightning.operation.output.id"]
|
||||
)
|
||||
if response_id is None and within_llm_call is True:
|
||||
is_llm_call = False
|
||||
if (
|
||||
@@ -547,7 +614,7 @@ class TraceTree:
|
||||
try:
|
||||
content = json.loads(content) # This content should now be a list
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse message content as JSON: {content}")
|
||||
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):
|
||||
@@ -567,18 +634,57 @@ class TraceTree:
|
||||
Subclass can override this method to add more fields to the triplet,
|
||||
such as chat messages and tool calls.
|
||||
"""
|
||||
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
|
||||
request_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.request")
|
||||
response_metadata = filter_and_unflatten_attributes(span.attributes, "gen_ai.response")
|
||||
prompt_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.prompt")
|
||||
completion_raw_content = filter_and_unflatten_attributes(span.attributes, "gen_ai.completion")
|
||||
image_urls = self.extract_prompt_image_urls(prompt_raw_content)
|
||||
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 []
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
@@ -112,6 +112,8 @@ class APO(Algorithm, Generic[T_task]):
|
||||
beam_rounds: int = 3,
|
||||
rollout_batch_timeout: float = 3600.0,
|
||||
run_initial_validation: bool = True,
|
||||
gradient_prompt_files: Optional[List[Path]] = None,
|
||||
apply_edit_prompt_files: Optional[List[Path]] = None,
|
||||
# Internal flags for debugging
|
||||
_poml_trace: bool = False,
|
||||
):
|
||||
@@ -132,6 +134,8 @@ class APO(Algorithm, Generic[T_task]):
|
||||
rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
|
||||
run_initial_validation: If True, runs validation on the seed prompt before starting
|
||||
optimization to establish a baseline score. Defaults to True.
|
||||
gradient_prompt_files: Prompt templates used to compute textual gradients (critiques).
|
||||
apply_edit_prompt_files: Prompt templates used to apply edits based on critiques.
|
||||
"""
|
||||
self.async_openai_client = async_openai_client
|
||||
self.gradient_model = gradient_model
|
||||
@@ -144,6 +148,8 @@ class APO(Algorithm, Generic[T_task]):
|
||||
self.beam_rounds = beam_rounds
|
||||
self.rollout_batch_timeout = rollout_batch_timeout
|
||||
self.run_initial_validation = run_initial_validation
|
||||
self.gradient_prompt_files = gradient_prompt_files or GRADIENT_PROMPT_FILES
|
||||
self.apply_edit_prompt_files = apply_edit_prompt_files or APPLY_EDIT_PROMPT_FILES
|
||||
|
||||
self._history_best_prompt: Optional[PromptTemplate] = None
|
||||
self._history_best_score: float = float("-inf")
|
||||
@@ -270,7 +276,7 @@ class APO(Algorithm, Generic[T_task]):
|
||||
Returns:
|
||||
A textual critique generated by the LLM, or None if generation fails.
|
||||
"""
|
||||
tg_template = random.choice(GRADIENT_PROMPT_FILES)
|
||||
tg_template = random.choice(self.gradient_prompt_files)
|
||||
|
||||
if len(rollout_results) < self.gradient_batch_size:
|
||||
self._log(
|
||||
@@ -352,7 +358,7 @@ class APO(Algorithm, Generic[T_task]):
|
||||
return current_prompt.prompt_template.template
|
||||
|
||||
# 2) Apply edit
|
||||
ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
|
||||
ae_template = random.choice(self.apply_edit_prompt_files)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",
|
||||
|
||||
@@ -32,6 +32,30 @@ class VERL(Algorithm):
|
||||
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": 4096,
|
||||
"trajectory_max_response_length": 34384,
|
||||
}
|
||||
```
|
||||
|
||||
Keep conversations structured (message lists rather than manual string
|
||||
concatenation) so prefix matching can stitch traces. `trajectory_max_prompt_length`
|
||||
should be set to the maximum length of the prompt for the first turn, and
|
||||
`trajectory_max_response_length` should be set to the maximum cumulative
|
||||
length of agent responses in the full trajectory.
|
||||
Toggle `debug=True` plus `mismatch_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
|
||||
from agentlightning.algorithm.verl import VERL
|
||||
|
||||
@@ -1,5 +1,16 @@
|
||||
# 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 .exception import emit_exception
|
||||
from .message import emit_message, get_message_value
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
import asyncio
|
||||
import functools
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
@@ -22,19 +21,18 @@ from typing import (
|
||||
overload,
|
||||
)
|
||||
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
|
||||
from agentlightning.utils.otel import flatten_attributes, get_tracer
|
||||
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
|
||||
|
||||
_FnType = TypeVar("_FnType", bound=Callable[..., Any])
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> ReadableSpan:
|
||||
def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanCoreFields:
|
||||
"""Emit a new annotation span.
|
||||
|
||||
This is the underlying implementation of [`emit_reward`][agentlightning.emit_reward].
|
||||
@@ -48,62 +46,46 @@ def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> Reada
|
||||
Args:
|
||||
annotation: Dictionary containing annotation key-value pairs.
|
||||
Representatives are rewards, tags, and metadata.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
propagate: Whether to propagate the span to tracers automatically.
|
||||
"""
|
||||
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)")
|
||||
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())
|
||||
|
||||
# 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,
|
||||
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"),
|
||||
)
|
||||
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}")
|
||||
|
||||
return span
|
||||
|
||||
|
||||
def _safe_json_dump(obj: Any) -> str:
|
||||
"""Serialize an object to JSON, falling back to ``str(obj)`` if needed.
|
||||
|
||||
Args:
|
||||
obj: Object to be serialized.
|
||||
|
||||
Returns:
|
||||
The JSON-encoded string representation of the object, or its string
|
||||
representation if JSON encoding fails.
|
||||
"""
|
||||
try:
|
||||
return json.dumps(obj, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(obj)
|
||||
|
||||
|
||||
class OperationContext:
|
||||
"""Context manager and decorator for tracing operations.
|
||||
|
||||
This class manages an OpenTelemetry span for a logical unit of work. It can
|
||||
be used either:
|
||||
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 :meth:`set_input` and :meth:`set_output`.
|
||||
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: OpenTelemetry tracer used to create spans.
|
||||
span: The currently active span, if any.
|
||||
tracer: Tracer implementation used to create spans.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, attributes: Dict[str, Any], *, propagate: bool = True) -> None:
|
||||
def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
|
||||
"""Initialize a new operation context.
|
||||
|
||||
Args:
|
||||
@@ -112,12 +94,19 @@ class OperationContext:
|
||||
JSON-serialized where necessary.
|
||||
propagate: Whether the span should be sent to active exporters.
|
||||
"""
|
||||
self.name: str = name
|
||||
self.initial_attributes: Dict[str, Any] = attributes
|
||||
self.propagate: bool = propagate
|
||||
self.tracer: trace.Tracer = get_tracer(use_active_span_processor=propagate)
|
||||
self.span: Optional[trace.Span] = None
|
||||
self._ctx_token: Optional[ContextManager[Any]] = None
|
||||
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.
|
||||
@@ -125,15 +114,10 @@ class OperationContext:
|
||||
Returns:
|
||||
The current :class:`OperationContext` instance with an active span.
|
||||
"""
|
||||
# 1. Start the span with initial attributes (JSON serialized)
|
||||
sanitized_attrs = {
|
||||
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
|
||||
for k, v in self.initial_attributes.items()
|
||||
}
|
||||
|
||||
self.span = self.tracer.start_span(self.name, attributes=sanitized_attrs)
|
||||
self._ctx_token = trace.use_span(self.span, end_on_exit=True)
|
||||
self._ctx_token.__enter__()
|
||||
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__(
|
||||
@@ -142,57 +126,63 @@ class OperationContext:
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> None:
|
||||
"""Exit the context manager and finish the span.
|
||||
"""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
|
||||
|
||||
Any exception raised inside the context is recorded on the span and the
|
||||
span status is set to error.
|
||||
|
||||
Args:
|
||||
exc_type: Exception type, if an exception occurred.
|
||||
exc_val: Exception instance, if an exception occurred.
|
||||
exc_tb: Traceback object, if an exception occurred.
|
||||
"""
|
||||
# 1. Record Exception if present
|
||||
if exc_val and self.span:
|
||||
self.span.record_exception(exc_val)
|
||||
self.span.set_status(Status(StatusCode.ERROR, str(exc_val)))
|
||||
|
||||
# 2. Close span
|
||||
if self._ctx_token:
|
||||
self._ctx_token.__exit__(exc_type, exc_val, exc_tb)
|
||||
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`` attribute,
|
||||
and keyword arguments are stored under ``input.<name>`` attributes.
|
||||
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.
|
||||
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.span:
|
||||
return
|
||||
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:
|
||||
self.span.set_attribute("input.args", _safe_json_dump(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 k, v in kwargs.items():
|
||||
self.span.set_attribute(f"input.{k}", _safe_json_dump(v))
|
||||
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.
|
||||
This is intended for use inside a `with operation(...) as op` block.
|
||||
|
||||
Args:
|
||||
output: The output value to record.
|
||||
"""
|
||||
if not self.span:
|
||||
return
|
||||
self.span.set_attribute("output", _safe_json_dump(output))
|
||||
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.
|
||||
@@ -212,60 +202,60 @@ class OperationContext:
|
||||
|
||||
sig = inspect.signature(fn)
|
||||
|
||||
def _record_auto_inputs(span: trace.Span, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None:
|
||||
"""Bind arguments to signature and log them on the span.
|
||||
sanitized_init_attrs = sanitize_attributes(
|
||||
{LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
|
||||
)
|
||||
|
||||
Args:
|
||||
span: Span on which to record attributes.
|
||||
args: Positional arguments passed to the wrapped callable.
|
||||
kwargs: Keyword arguments passed to the wrapped callable.
|
||||
"""
|
||||
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 k, v in bound.arguments.items():
|
||||
span.set_attribute(
|
||||
f"{LightningSpanAttributes.OPERATION_INPUT.value}.{k}",
|
||||
_safe_json_dump(v),
|
||||
)
|
||||
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:
|
||||
span.set_attribute(
|
||||
f"{LightningSpanAttributes.OPERATION_INPUT.value}.args",
|
||||
_safe_json_dump(args),
|
||||
)
|
||||
span.set_attribute(
|
||||
f"{LightningSpanAttributes.OPERATION_INPUT.value}.kwargs",
|
||||
_safe_json_dump(kwargs),
|
||||
)
|
||||
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."""
|
||||
# Reuse __enter__ logic via 'with self' would share state incorrectly
|
||||
# across concurrent calls. We must create a new span per call.
|
||||
# So we manually reimplement the span logic for the wrapper here.
|
||||
|
||||
sanitized_attrs = {
|
||||
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
|
||||
for k, v in self.initial_attributes.items()
|
||||
}
|
||||
|
||||
with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
|
||||
span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
|
||||
_record_auto_inputs(span, args, kwargs)
|
||||
try:
|
||||
result = await fn(*args, **kwargs)
|
||||
span.set_attribute(
|
||||
LightningSpanAttributes.OPERATION_OUTPUT.value,
|
||||
_safe_json_dump(result),
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
span.record_exception(e)
|
||||
span.set_status(Status(StatusCode.ERROR, str(e)))
|
||||
raise
|
||||
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)
|
||||
|
||||
@@ -274,41 +264,48 @@ class OperationContext:
|
||||
@functools.wraps(fn)
|
||||
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Sync wrapper that traces the wrapped callable."""
|
||||
sanitized_attrs = {
|
||||
k: _safe_json_dump(v) if not isinstance(v, (str, int, float, bool)) else v
|
||||
for k, v in self.initial_attributes.items()
|
||||
}
|
||||
|
||||
with self.tracer.start_as_current_span(self.name, attributes=sanitized_attrs) as span:
|
||||
span.set_attribute(LightningSpanAttributes.OPERATION_NAME.value, function_name)
|
||||
_record_auto_inputs(span, args, kwargs)
|
||||
try:
|
||||
result = fn(*args, **kwargs)
|
||||
span.set_attribute(
|
||||
LightningSpanAttributes.OPERATION_OUTPUT.value,
|
||||
_safe_json_dump(result),
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
span.record_exception(e)
|
||||
span.set_status(Status(StatusCode.ERROR, str(e)))
|
||||
raise
|
||||
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, **additional_attributes: Any) -> _FnType: ...
|
||||
def operation(
|
||||
fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
|
||||
) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(*, propagate: bool = True, **additional_attributes: Any) -> OperationContext: ...
|
||||
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.
|
||||
@@ -344,6 +341,9 @@ def operation(
|
||||
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.
|
||||
|
||||
@@ -352,6 +352,12 @@ def operation(
|
||||
[`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
|
||||
|
||||
@@ -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.utils.otel import get_tracer
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -32,25 +32,23 @@ def emit_exception(
|
||||
"""
|
||||
if not isinstance(exception, BaseException): # type: ignore
|
||||
raise TypeError(f"Expected a BaseException instance, got: {type(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
|
||||
span_attributes = format_exception_attributes(exception)
|
||||
|
||||
if attributes:
|
||||
span_attributes.update(attributes)
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
|
||||
span = tracer.start_span(
|
||||
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(
|
||||
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,8 +4,10 @@ import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otel import get_tracer
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -27,17 +29,21 @@ 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)}.")
|
||||
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
span_attributes = {LightningSpanAttributes.MESSAGE_BODY.value: 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}
|
||||
if attributes:
|
||||
span_attributes.update(attributes)
|
||||
span = tracer.start_span(
|
||||
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(
|
||||
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,13 +6,15 @@ import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
|
||||
from agentlightning.types import SpanLike
|
||||
from agentlightning.utils.otel import full_qualified_name, get_tracer
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
|
||||
def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> SpanCoreFields:
|
||||
"""Emit an object's serialized representation as an OpenTelemetry span.
|
||||
|
||||
Args:
|
||||
@@ -25,20 +27,29 @@ def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propag
|
||||
"""
|
||||
span_attributes = encode_object(object)
|
||||
if attributes:
|
||||
span_attributes.update(attributes)
|
||||
tracer = get_tracer(use_active_span_processor=propagate)
|
||||
span = tracer.start_span(
|
||||
AGL_OBJECT,
|
||||
attributes=span_attributes,
|
||||
)
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
|
||||
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)
|
||||
with span:
|
||||
pass
|
||||
|
||||
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"),
|
||||
)
|
||||
|
||||
|
||||
def encode_object(object: Any) -> Dict[str, Any]:
|
||||
|
||||
@@ -20,13 +20,10 @@ 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 SpanLike
|
||||
from agentlightning.types import SpanCoreFields, SpanLike
|
||||
from agentlightning.utils.otel import filter_and_unflatten_attributes
|
||||
|
||||
from .annotation import emit_annotation
|
||||
@@ -61,6 +58,8 @@ _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
|
||||
|
||||
|
||||
@@ -81,6 +80,8 @@ 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:
|
||||
@@ -146,7 +147,7 @@ def emit_reward(
|
||||
primary_key: str | None = None,
|
||||
attributes: Dict[str, Any] | None = None,
|
||||
propagate: bool = True,
|
||||
) -> ReadableSpan:
|
||||
) -> SpanCoreFields:
|
||||
"""Emit a reward value as an OpenTelemetry span.
|
||||
|
||||
Examples:
|
||||
@@ -172,11 +173,7 @@ def emit_reward(
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
Returns:
|
||||
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.
|
||||
Span core fields capturing the recorded reward.
|
||||
"""
|
||||
logger.debug(f"Emitting reward: {reward}")
|
||||
reward_dimensions: List[RewardDimension] = []
|
||||
|
||||
@@ -39,13 +39,6 @@ try:
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from . import weave # type: ignore
|
||||
|
||||
WEAVE_INSTALLED = True # type: ignore
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def instrument_all():
|
||||
"""Instrument all the instrumentation libraries."""
|
||||
@@ -119,20 +112,3 @@ def uninstrument_all():
|
||||
warnings.warn("agentops_langchain is installed but uninstrument_agentops_langchain could not be imported.")
|
||||
else:
|
||||
warnings.warn("Agentops-langchain integration is not installed. It's therefore not uninstrumented.")
|
||||
|
||||
|
||||
def instrument_weave():
|
||||
if WEAVE_INSTALLED:
|
||||
from .weave import instrument_weave
|
||||
|
||||
instrument_weave()
|
||||
|
||||
|
||||
def uninstrument_weave():
|
||||
if WEAVE_INSTALLED:
|
||||
try:
|
||||
from .weave import uninstrument_weave
|
||||
|
||||
uninstrument_weave()
|
||||
except ImportError:
|
||||
warnings.warn("weave is installed but uninstrument_weave could not be imported.")
|
||||
|
||||
@@ -1,139 +1,500 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Callable, Optional
|
||||
from __future__ import annotations
|
||||
|
||||
import requests
|
||||
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_default_entity_name_getter: Callable[..., Any] | None = None
|
||||
_original_upsert_project_getter: Callable[..., Any] | None = None
|
||||
_original_weave_get = False
|
||||
_original_weave_post = False
|
||||
_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 instrument_weave():
|
||||
"""
|
||||
Patch the Weave/W&B integration to bypass actual network calls for testing.
|
||||
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
|
||||
|
||||
- Mocks HTTP POST/GET requests
|
||||
- Patches wandb.Api methods
|
||||
- Silences Weave logging
|
||||
- Sets dummy WANDB_API_KEY if not provided
|
||||
"""
|
||||
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:
|
||||
import weave
|
||||
from weave.compat import wandb # type: ignore
|
||||
except ImportError:
|
||||
logger.warning("Weave or wandb not installed; cannot uninstrument.")
|
||||
return
|
||||
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"
|
||||
|
||||
_weave_tracer_entity_name = "weave_tracer_entity"
|
||||
|
||||
def default_entity_name_getter(_self) -> str: # type: ignore
|
||||
return _weave_tracer_entity_name
|
||||
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 upsert_project_getter(
|
||||
_self, project: str, description: Optional[str] = None, entity: Optional[str] = None # type: ignore
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"upsertModel": {
|
||||
"model": {
|
||||
"name": project,
|
||||
"description": description or "",
|
||||
"entity": entity or _weave_tracer_entity_name,
|
||||
}
|
||||
},
|
||||
"project": "weave_tracer_project",
|
||||
}
|
||||
|
||||
# Mock network requests to avoid real HTTP calls
|
||||
def post(url: str, *args: Any, **kwargs: Any) -> requests.Response:
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response._content = b'{"digest": "mocked_digest"}'
|
||||
return response
|
||||
def instrument_weave(server: InMemoryWeaveTraceServer):
|
||||
"""Patch the Weave/W&B integration to bypass actual network calls for testing."""
|
||||
|
||||
def get(url: str, *args: Any, **kwargs: Any) -> requests.Response:
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response._content = b'{"min_required_weave_python_version": "0.52.14"}'
|
||||
return response
|
||||
|
||||
# Patch API methods and HTTP requests
|
||||
global _original_default_entity_name_getter
|
||||
global _original_upsert_project_getter
|
||||
global _original_weave_post
|
||||
global _original_weave_get
|
||||
_original_default_entity_name_getter = wandb.Api.default_entity_name # type: ignore
|
||||
_original_upsert_project_getter = wandb.Api.upsert_project # type: ignore
|
||||
_original_weave_post = weave.utils.http_requests.session.post # type: ignore
|
||||
_original_weave_get = weave.utils.http_requests.session.get # type: ignore
|
||||
|
||||
# Patch API methods and HTTP requests
|
||||
wandb.Api.default_entity_name = default_entity_name_getter # type: ignore
|
||||
wandb.Api.upsert_project = upsert_project_getter # type: ignore
|
||||
weave.utils.http_requests.session.post = post # type: ignore
|
||||
weave.utils.http_requests.session.get = get # type: ignore
|
||||
|
||||
# Silence Weave logging
|
||||
for name in logging.root.manager.loggerDict:
|
||||
if name.startswith("weave"):
|
||||
logging.getLogger(name).disabled = True
|
||||
|
||||
# Set dummy API key if missing
|
||||
if not os.environ.get("WANDB_API_KEY"):
|
||||
os.environ["WANDB_API_KEY"] = "dumped_api_key_for_weave_tracer"
|
||||
|
||||
# if needed in future tests, enable this and replace WF_TRACE_SERVER_URL to local server
|
||||
# full_url = f"http://127.0.0.1:{_port}"
|
||||
# os.environ["WF_TRACE_SERVER_URL"] = full_url
|
||||
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.
|
||||
"""
|
||||
try:
|
||||
import weave
|
||||
from weave.compat import wandb # type: ignore
|
||||
except ImportError:
|
||||
logger.warning("Weave or wandb not installed; cannot uninstrument.")
|
||||
return
|
||||
"""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
|
||||
|
||||
global _original_default_entity_name_getter
|
||||
if _original_default_entity_name_getter is not None:
|
||||
wandb.Api.default_entity_name = _original_default_entity_name_getter # type: ignore
|
||||
_original_default_entity_name_getter = None
|
||||
logger.info("restored wandb.Api.default_entity_name")
|
||||
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.")
|
||||
|
||||
global _original_upsert_project_getter
|
||||
if _original_upsert_project_getter is not None:
|
||||
wandb.Api.upsert_project = _original_upsert_project_getter # type: ignore
|
||||
_original_upsert_project_getter = None
|
||||
logger.info("restored wandb.Api.upsert_project")
|
||||
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.")
|
||||
|
||||
global _original_weave_post
|
||||
if _original_weave_post is not None:
|
||||
weave.utils.http_requests.session.post = _original_weave_post # type: ignore
|
||||
_original_weave_post = None
|
||||
logger.info("restored weave.utils.http_requests.session.post")
|
||||
|
||||
global _original_weave_get
|
||||
if _original_weave_get is not None:
|
||||
weave.utils.http_requests.session.get = _original_weave_get # type: ignore
|
||||
_original_weave_get = None
|
||||
logger.info("restored weave.utils.http_requests.session.get")
|
||||
|
||||
# Restore Weave logging
|
||||
for name in logging.root.manager.loggerDict:
|
||||
if name.startswith("weave"):
|
||||
logging.getLogger(name).disabled = False
|
||||
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,6 +198,7 @@ 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.")
|
||||
|
||||
|
||||
+240
-57
@@ -43,6 +43,7 @@ from agentlightning.types import (
|
||||
RolloutMode,
|
||||
RolloutRawResult,
|
||||
Span,
|
||||
SpanCoreFields,
|
||||
)
|
||||
from agentlightning.utils.system_snapshot import system_snapshot
|
||||
|
||||
@@ -74,7 +75,8 @@ class LitAgentRunner(Runner[T_task]):
|
||||
poll_interval: float = 5.0,
|
||||
heartbeat_interval: float = 10.0,
|
||||
interval_jitter: float = 0.5,
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
|
||||
heartbeat_include_gpu: bool = False,
|
||||
) -> None:
|
||||
"""Initialize the agent runner.
|
||||
|
||||
@@ -88,7 +90,10 @@ 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".
|
||||
"asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
|
||||
"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.
|
||||
"""
|
||||
super().__init__()
|
||||
self._tracer = tracer
|
||||
@@ -97,6 +102,7 @@ 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
|
||||
@@ -276,7 +282,7 @@ class LitAgentRunner(Runner[T_task]):
|
||||
"""
|
||||
store = self.get_store()
|
||||
|
||||
trace_spans: list[ReadableSpan] | list[Span] = []
|
||||
trace_spans: list[Span] = []
|
||||
result_recognized: bool = False
|
||||
|
||||
# Case 0: result is None
|
||||
@@ -295,31 +301,38 @@ class LitAgentRunner(Runner[T_task]):
|
||||
# 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 = emit_reward(raw_result, propagate=False)
|
||||
reward_span_core_fields = emit_reward(raw_result, propagate=False)
|
||||
# We add it to the store manually
|
||||
await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
|
||||
trace_spans.append(reward_span)
|
||||
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
|
||||
|
||||
# Case 2-3: result is a list
|
||||
# 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 not isinstance(self._tracer, OtelTracer):
|
||||
for span in raw_result:
|
||||
await store.add_otel_span(
|
||||
rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
|
||||
)
|
||||
else:
|
||||
if isinstance(self._tracer, OtelTracer):
|
||||
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. "
|
||||
"No need to return anything from rollout."
|
||||
"Returning the traces from the rollout will result in duplicate spans."
|
||||
)
|
||||
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)
|
||||
@@ -327,7 +340,25 @@ class LitAgentRunner(Runner[T_task]):
|
||||
# Add the spans directly to the store
|
||||
for span in raw_result:
|
||||
await store.add_span(cast(Span, span))
|
||||
trace_spans = raw_result
|
||||
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
|
||||
|
||||
# Left over cases for list
|
||||
@@ -336,7 +367,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 = raw_result
|
||||
trace_spans = []
|
||||
result_recognized = True
|
||||
|
||||
else:
|
||||
@@ -355,14 +386,46 @@ class LitAgentRunner(Runner[T_task]):
|
||||
return trace_spans
|
||||
|
||||
async def _emit_heartbeat(self, store: LightningStore) -> None:
|
||||
"""Send a heartbeat tick to the store."""
|
||||
"""Send a heartbeat tick to the store.
|
||||
|
||||
Args:
|
||||
store: The lightning store to update.
|
||||
"""
|
||||
logger.debug(f"{self._log_prefix()} Preparing to emit heartbeat.")
|
||||
worker_id = self.get_worker_id()
|
||||
|
||||
try:
|
||||
await store.update_worker(worker_id, system_snapshot())
|
||||
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
|
||||
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())
|
||||
|
||||
@@ -377,51 +440,161 @@ class LitAgentRunner(Runner[T_task]):
|
||||
return None
|
||||
|
||||
if self._heartbeat_launch_mode == "asyncio":
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
async def heartbeat_loop() -> None:
|
||||
while not stop_event.is_set():
|
||||
await self._emit_heartbeat(store)
|
||||
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
|
||||
|
||||
return self._start_heartbeat_asyncio_loop(store)
|
||||
if self._heartbeat_launch_mode == "thread":
|
||||
stop_evt = threading.Event()
|
||||
return self._start_heartbeat_thread_loop(store)
|
||||
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
|
||||
|
||||
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))
|
||||
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.
|
||||
await self._emit_heartbeat(store)
|
||||
except Exception:
|
||||
logger.exception("%s Heartbeat failed.", self._log_prefix())
|
||||
with suppress(asyncio.TimeoutError):
|
||||
interval = self._heartbeat_interval + self._random_state.uniform(
|
||||
-self._interval_jitter, self._interval_jitter
|
||||
)
|
||||
interval = max(interval, 0.01)
|
||||
stop_evt.wait(interval)
|
||||
await asyncio.wait_for(stop_event.wait(), timeout=interval)
|
||||
|
||||
thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
|
||||
thread.start()
|
||||
task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
|
||||
|
||||
async def stop() -> None:
|
||||
stop_evt.set()
|
||||
await asyncio.to_thread(thread.join)
|
||||
async def stop() -> None:
|
||||
stop_event.set()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
return stop
|
||||
return stop
|
||||
|
||||
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
|
||||
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
|
||||
|
||||
async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
|
||||
"""Sleep until the next poll interval, with optional event-based interruption.
|
||||
@@ -477,6 +650,8 @@ 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
|
||||
|
||||
@@ -484,9 +659,11 @@ 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
|
||||
)
|
||||
@@ -498,21 +675,27 @@ 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)
|
||||
|
||||
@@ -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, Triplet
|
||||
from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, 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[ReadableSpan]] = None
|
||||
trace_spans: Optional[List[SpanLike]] = None
|
||||
|
||||
# Handle different types of results from the agent
|
||||
# Case 1: result is a float (final reward)
|
||||
@@ -108,10 +108,14 @@ 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: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if isinstance(result, list) and all(isinstance(t, ReadableSpan) for t in result):
|
||||
# Case 3.1: 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
|
||||
@@ -123,10 +127,9 @@ 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):
|
||||
spans = self.tracer.get_last_trace()
|
||||
if spans:
|
||||
trace = [json.loads(readable_span.to_json()) for readable_span in spans]
|
||||
trace_spans = spans
|
||||
trace_spans = self.tracer.get_last_trace() # type: ignore
|
||||
if trace_spans:
|
||||
trace = [cast(Span, span).model_dump() for span in trace_spans]
|
||||
|
||||
# Always extract triplets from the trace using TracerTraceToTriplet
|
||||
if trace_spans:
|
||||
|
||||
@@ -34,6 +34,9 @@ AGL_OPERATION = "agentlightning.operation"
|
||||
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.
|
||||
|
||||
@@ -53,6 +56,9 @@ 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.
|
||||
|
||||
@@ -903,9 +903,11 @@ class LightningStoreServer(LightningStore):
|
||||
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
|
||||
@@ -1518,7 +1520,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: {cre.status} {cre.message}", exc_info=True)
|
||||
client_logger.debug(f"ClientResponseError ({method} {path}): {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
|
||||
@@ -1534,9 +1536,9 @@ class LightningStoreClient(LightningStore):
|
||||
asyncio.TimeoutError,
|
||||
) as net_exc:
|
||||
# Network/session issue: probe health before retrying
|
||||
client_logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
|
||||
client_logger.debug(f"Network/session issue ({method} {path}): {net_exc}", exc_info=True)
|
||||
last_exc = net_exc
|
||||
client_logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
|
||||
client_logger.info(f"Network/session issue: {net_exc} - will retry the request {method}: {path}")
|
||||
if not await self._wait_until_healthy(session):
|
||||
break # server is not healthy, do not retry
|
||||
|
||||
|
||||
@@ -148,12 +148,12 @@ class TrackedCollection:
|
||||
yield
|
||||
|
||||
else:
|
||||
from agentlightning.store.collection_based import nearest_lightning_store_method_from_stack
|
||||
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 = nearest_lightning_store_method_from_stack()
|
||||
public_store_method, private_store_method = get_current_store_methods()
|
||||
try:
|
||||
yield
|
||||
except BaseException as exc:
|
||||
|
||||
@@ -1077,7 +1077,7 @@ class MongoBasedKeyValue(KeyValue[K, V], Generic[K, V]):
|
||||
class MongoLightningCollections(LightningCollections):
|
||||
"""Mongo implementation of LightningCollections using MongoDB collections.
|
||||
|
||||
Serves as the storage base for [`MongoLightningStore`][agentlightning.store.MongoLightningStore].
|
||||
Serves as the storage base for [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore].
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -15,13 +15,11 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import hashlib
|
||||
import inspect
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
from contextvars import ContextVar
|
||||
from types import CoroutineType
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -61,6 +59,7 @@ from agentlightning.types import (
|
||||
Worker,
|
||||
WorkerStatus,
|
||||
)
|
||||
from agentlightning.utils.id import generate_id
|
||||
from agentlightning.utils.metrics import MetricsBackend
|
||||
|
||||
from .base import (
|
||||
@@ -88,6 +87,12 @@ SelfT = TypeVar("SelfT", bound="CollectionBasedLightningStore[Any]")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ContextVars for tracking the current store method without expensive stack introspection.
|
||||
# These are set by the @tracked decorator and read by tracking_context in collection/base.py.
|
||||
_UNKNOWN_STORE_METHOD = "unknown"
|
||||
_current_public_store_method: ContextVar[str] = ContextVar("public_store_method", default=_UNKNOWN_STORE_METHOD)
|
||||
_current_private_store_method: ContextVar[str] = ContextVar("private_store_method", default=_UNKNOWN_STORE_METHOD)
|
||||
|
||||
|
||||
def _with_collections_execute(labels: Sequence[AtomicLabels]):
|
||||
"""Hands over the function execution to the collections.execute method.
|
||||
@@ -125,38 +130,47 @@ def tracked(name: str):
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: CollectionBasedLightningStore[T_collections], *args: Any, **kwargs: Any) -> Any:
|
||||
# Backtracking where this method comes from
|
||||
public_meth_in_stack, _ = nearest_lightning_store_method_from_stack()
|
||||
# Get the current public method from ContextVar (set by outer tracked methods)
|
||||
public_meth_in_stack = _current_public_store_method.get()
|
||||
|
||||
# For backtracking in collection methods.
|
||||
# Only track the public methods (+healthcheck)
|
||||
# Set ContextVars for nested calls to read. Use tokens for proper cleanup.
|
||||
pub_token = None
|
||||
priv_token = None
|
||||
if name in COLLECTION_STORE_PUBLIC_METHODS:
|
||||
public_method_name = name # pyright: ignore[reportUnusedVariable]
|
||||
pub_token = _current_public_store_method.set(name)
|
||||
public_meth_in_stack = name # We are in a public method already.
|
||||
if name in COLLECTION_STORE_ALL_METHODS:
|
||||
private_method_name = name # pyright: ignore[reportUnusedVariable]
|
||||
priv_token = _current_private_store_method.set(name)
|
||||
|
||||
if self._tracker is None: # pyright: ignore[reportPrivateUsage]
|
||||
# Skip the tracking because tracking is not configured
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
status: str = "OK"
|
||||
try:
|
||||
return await func(self, *args, **kwargs)
|
||||
except BaseException as exc:
|
||||
status = exc.__class__.__name__
|
||||
raise
|
||||
if self._tracker is None: # pyright: ignore[reportPrivateUsage]
|
||||
# Skip the tracking because tracking is not configured
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
status: str = "OK"
|
||||
try:
|
||||
return await func(self, *args, **kwargs)
|
||||
except BaseException as exc:
|
||||
status = exc.__class__.__name__
|
||||
raise
|
||||
finally:
|
||||
elapsed = time.perf_counter() - start_time
|
||||
await self._tracker.inc_counter( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.store.total",
|
||||
labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
|
||||
)
|
||||
await self._tracker.observe_histogram( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.store.latency",
|
||||
value=elapsed,
|
||||
labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
|
||||
)
|
||||
finally:
|
||||
elapsed = time.perf_counter() - start_time
|
||||
await self._tracker.inc_counter( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.store.total", labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status}
|
||||
)
|
||||
await self._tracker.observe_histogram( # pyright: ignore[reportPrivateUsage]
|
||||
"agl.store.latency",
|
||||
value=elapsed,
|
||||
labels={"method": name, "store_pubmeth": public_meth_in_stack, "status": status},
|
||||
)
|
||||
# Reset ContextVars to their previous values
|
||||
if pub_token is not None:
|
||||
_current_public_store_method.reset(pub_token)
|
||||
if priv_token is not None:
|
||||
_current_private_store_method.reset(priv_token)
|
||||
|
||||
return cast(T_callable, wrapper)
|
||||
|
||||
@@ -195,19 +209,16 @@ def healthcheck_before(func: T_callable) -> T_callable:
|
||||
|
||||
|
||||
def _generate_resources_id() -> str:
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
|
||||
return "rs-" + short_id
|
||||
return "rs-" + generate_id(12)
|
||||
|
||||
|
||||
def _generate_rollout_id() -> str:
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
|
||||
return "ro-" + short_id
|
||||
return "ro-" + generate_id(12)
|
||||
|
||||
|
||||
def _generate_attempt_id() -> str:
|
||||
"""We don't need that long because attempts are limited to rollouts."""
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
|
||||
return "at-" + short_id
|
||||
return "at-" + generate_id(8)
|
||||
|
||||
|
||||
class CollectionBasedLightningStore(LightningStore, Generic[T_collections]):
|
||||
@@ -1752,41 +1763,14 @@ COLLECTION_STORE_PUBLIC_METHODS = frozenset(
|
||||
|
||||
COLLECTION_STORE_ALL_METHODS = frozenset([name for name in CollectionBasedLightningStore.__dict__])
|
||||
|
||||
_UNKNOWN_STORE_METHOD = "unknown"
|
||||
|
||||
def get_current_store_methods() -> Tuple[str, str]:
|
||||
"""Get the current store method names from ContextVars.
|
||||
|
||||
def nearest_lightning_store_method_from_stack() -> Tuple[str, str]:
|
||||
"""Stack introspection so that we capture the nearest public API method from the
|
||||
call stack whenever metrics are recorded.
|
||||
This is a fast O(1) replacement for stack introspection. The ContextVars are
|
||||
set by the @tracked decorator when entering store methods.
|
||||
|
||||
Returns:
|
||||
A tuple of public method name and nearest private method name.
|
||||
A tuple of (public_method_name, private_method_name).
|
||||
"""
|
||||
frame = inspect.currentframe()
|
||||
final_public_method_name = final_private_method_name = _UNKNOWN_STORE_METHOD
|
||||
try:
|
||||
if frame is not None:
|
||||
frame = frame.f_back
|
||||
while frame is not None:
|
||||
self_obj = frame.f_locals.get("self")
|
||||
public_method_name = frame.f_locals.get("public_method_name")
|
||||
private_method_name = frame.f_locals.get("private_method_name")
|
||||
if (
|
||||
final_public_method_name == _UNKNOWN_STORE_METHOD
|
||||
and public_method_name in COLLECTION_STORE_PUBLIC_METHODS
|
||||
and isinstance(self_obj, LightningStore)
|
||||
):
|
||||
final_public_method_name = public_method_name
|
||||
if (
|
||||
final_private_method_name == _UNKNOWN_STORE_METHOD
|
||||
and private_method_name in COLLECTION_STORE_ALL_METHODS
|
||||
and isinstance(self_obj, LightningStore)
|
||||
):
|
||||
final_private_method_name = private_method_name
|
||||
frame = frame.f_back
|
||||
except Exception as exc:
|
||||
logger.debug("Error during stack introspection for LightningStore method: %s", exc)
|
||||
finally:
|
||||
del frame
|
||||
|
||||
return final_public_method_name, final_private_method_name
|
||||
return _current_public_store_method.get(), _current_private_store_method.get()
|
||||
|
||||
@@ -33,7 +33,7 @@ class MongoLightningStore(CollectionBasedLightningStore[MongoLightningCollection
|
||||
Args:
|
||||
mongo_uri: MongoDB connection string (defaults to local replica set).
|
||||
mongo_client_kwargs: Extra keyword arguments forwarded to `AsyncMongoClient`.
|
||||
database: The MongoDB database name. Defaults to ``agentlightning``.
|
||||
database_name: The MongoDB database name. Defaults to ``agentlightning``.
|
||||
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.
|
||||
|
||||
@@ -1,8 +1,16 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .agentops import AgentOpsTracer
|
||||
from .base import Tracer
|
||||
from .base import Tracer, clear_active_tracer, get_active_tracer, set_active_tracer
|
||||
from .dummy import DummyTracer
|
||||
from .otel import OtelTracer
|
||||
from .weave import WeaveTracer
|
||||
|
||||
__all__ = ["AgentOpsTracer", "Tracer", "OtelTracer", "WeaveTracer"]
|
||||
__all__ = [
|
||||
"AgentOpsTracer",
|
||||
"Tracer",
|
||||
"OtelTracer",
|
||||
"DummyTracer",
|
||||
"get_active_tracer",
|
||||
"set_active_tracer",
|
||||
"clear_active_tracer",
|
||||
]
|
||||
|
||||
@@ -13,12 +13,13 @@ 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:
|
||||
@@ -79,13 +80,20 @@ 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.")
|
||||
logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized. Skip initialization.")
|
||||
|
||||
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
|
||||
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
|
||||
|
||||
def teardown_worker(self, worker_id: int) -> None:
|
||||
super().teardown_worker(worker_id)
|
||||
@@ -94,6 +102,10 @@ 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,
|
||||
@@ -158,7 +170,6 @@ 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,14 +2,13 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional, TypeVar
|
||||
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import ParallelWorkerBase
|
||||
from agentlightning.types import Attributes, ParallelWorkerBase, Span, SpanCoreFields, SpanRecordingContext, TraceStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.callbacks.base import BaseCallbackHandler # type: ignore
|
||||
@@ -17,6 +16,14 @@ 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.
|
||||
@@ -98,12 +105,12 @@ class Tracer(ParallelWorkerBase):
|
||||
"""Internal API for CI backward compatibility."""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
def get_last_trace(self) -> List[Span]:
|
||||
"""
|
||||
Retrieves the raw list of captured spans from the most recent trace.
|
||||
|
||||
Returns:
|
||||
A list of OpenTelemetry `ReadableSpan` objects.
|
||||
A list of [`Span`][agentlightning.Span] objects collected during the last trace.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@@ -124,6 +131,48 @@ 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.
|
||||
@@ -175,3 +224,64 @@ 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
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
# 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)
|
||||
+199
-22
@@ -4,29 +4,72 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import warnings
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, Awaitable, List, Optional
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import Any, AsyncGenerator, Awaitable, Iterator, 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 SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor, 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
|
||||
from .base import Tracer, with_active_tracer_context
|
||||
|
||||
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.
|
||||
@@ -38,7 +81,7 @@ class OtelTracer(Tracer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# This provider is only initialized when the worker is initialized.
|
||||
self._tracer_provider: Optional[TracerProvider] = None
|
||||
self._tracer_provider: Optional[trace_api.TracerProvider] = None
|
||||
self._lightning_span_processor: Optional[LightningSpanProcessor] = None
|
||||
self._simple_span_processor: Optional[SimpleSpanProcessor] = None
|
||||
self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
|
||||
@@ -63,7 +106,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 = TracerProvider()
|
||||
self._tracer_provider = TracerProviderImpl()
|
||||
trace_api.set_tracer_provider(self._tracer_provider)
|
||||
self._lightning_span_processor = LightningSpanProcessor()
|
||||
self._tracer_provider.add_span_processor(self._lightning_span_processor)
|
||||
@@ -78,8 +121,9 @@ 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(
|
||||
async def trace_context( # kh: runner.step_impl에서 옴
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
*,
|
||||
@@ -117,6 +161,7 @@ class OtelTracer(Tracer):
|
||||
if store.capabilities.get("otlp_traces", False) is True:
|
||||
logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
|
||||
self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
|
||||
# kh: otlp traces 지원하는 store인 경우, otlp exporter 사용
|
||||
else:
|
||||
self._disable_native_otlp_exporter()
|
||||
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
|
||||
@@ -129,12 +174,69 @@ class OtelTracer(Tracer):
|
||||
else:
|
||||
raise ValueError("rollout_id and attempt_id must be either all provided or all None")
|
||||
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
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]:
|
||||
"""
|
||||
Retrieves the raw list of captured spans from the most recent trace.
|
||||
|
||||
Returns:
|
||||
A list of OpenTelemetry `ReadableSpan` objects.
|
||||
A list of [`Span`][agentlightning.Span] objects captured during the most recent trace.
|
||||
"""
|
||||
if not self._lightning_span_processor:
|
||||
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
|
||||
@@ -143,6 +245,8 @@ 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):
|
||||
@@ -160,16 +264,33 @@ class OtelTracer(Tracer):
|
||||
)
|
||||
instrumented = False
|
||||
candidates: List[str] = []
|
||||
# kh: processor 별로 조건 check 하는 것임!!
|
||||
for processor in active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
|
||||
if isinstance(processor, LightningSpanProcessor):
|
||||
# We don't need the LightningSpanProcessor any more.
|
||||
logger.debug("LightningSpanProcessor already present in TracerProvider, disabling it.")
|
||||
processor.disable_store_submission = True
|
||||
|
||||
# kh added
|
||||
logger.debug(f"disable_store_submission set to True for rollout={rollout_id} attempt={attempt_id}")
|
||||
elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
|
||||
# Instead, we rely on the OTLPSpanExporter to send spans to the store.
|
||||
if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
|
||||
processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
|
||||
logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
|
||||
# option 1
|
||||
# kh 기존: store로 보낸다는 설정;
|
||||
# processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
|
||||
# logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
|
||||
|
||||
# option 2 (hardcoded for collector)
|
||||
# collector 실험: otel collector로 설정?
|
||||
# processor.span_exporter.enable_store_otlp("http://localhost:4318/v1/traces", rollout_id, attempt_id)
|
||||
# logger.info(f"Set LightningStoreOTLPExporter endpoint to http://localhost:4318/v1/traces")
|
||||
|
||||
# option 3 (switchable via env var)
|
||||
endpoint = os.getenv("AGL_OTLP_ENDPOINT") or store.otlp_traces_endpoint()
|
||||
processor.span_exporter.enable_store_otlp(endpoint, rollout_id, attempt_id)
|
||||
logger.debug(f"[SET] enabling span export to {endpoint} rollout={rollout_id} attempt={attempt_id}")
|
||||
|
||||
instrumented = True
|
||||
else:
|
||||
candidates.append(
|
||||
@@ -186,6 +307,9 @@ class OtelTracer(Tracer):
|
||||
)
|
||||
|
||||
def _disable_native_otlp_exporter(self):
|
||||
# kh added
|
||||
# logger.info("disable_store_submission reset to False")
|
||||
|
||||
tracer_provider = self._get_tracer_provider()
|
||||
active_span_processor = tracer_provider._active_span_processor # pyright: ignore[reportPrivateUsage]
|
||||
tracer_provider._resource = tracer_provider._resource.merge( # pyright: ignore[reportPrivateUsage]
|
||||
@@ -215,18 +339,20 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
|
||||
def __init__(self, disable_store_submission: bool = False):
|
||||
self._disable_store_submission: bool = disable_store_submission
|
||||
self._spans: List[ReadableSpan] = []
|
||||
self._spans: List[Span] = []
|
||||
|
||||
# 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 (
|
||||
@@ -262,11 +388,19 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
self._disable_store_submission = value
|
||||
|
||||
def _ensure_loop(self) -> None:
|
||||
if self._loop_thread is None or self._loop is 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
|
||||
self._loop_ready.clear()
|
||||
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
|
||||
self._loop_thread.start()
|
||||
self._loop_ready.wait() # loop is ready
|
||||
if not self._loop_ready.wait(timeout=30.0):
|
||||
raise RuntimeError("Timed out waiting for otel-loop thread to start")
|
||||
|
||||
def _loop_runner(self):
|
||||
loop = asyncio.new_event_loop()
|
||||
@@ -330,13 +464,13 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
def force_flush(self, timeout_millis: int = 30000) -> bool:
|
||||
return True
|
||||
|
||||
def spans(self) -> List[ReadableSpan]:
|
||||
def spans(self) -> List[Span]:
|
||||
"""
|
||||
Get the list of spans collected by this processor.
|
||||
This is useful for debugging and testing purposes.
|
||||
|
||||
Returns:
|
||||
List of ReadableSpan objects collected during tracing.
|
||||
List of [`Span`][agentlightning.Span] objects collected during tracing.
|
||||
"""
|
||||
return self._spans
|
||||
|
||||
@@ -368,17 +502,60 @@ class LightningSpanProcessor(SpanProcessor):
|
||||
if not span.context or not span.context.trace_flags.sampled:
|
||||
return
|
||||
|
||||
# kh: _disable_store_submission 조건에 안 맞아서 안 들어갈 가능성 큼
|
||||
if not self._disable_store_submission and self._store and self._rollout_id and self._attempt_id:
|
||||
try:
|
||||
# Submit add_otel_span to the event loop and wait for it to complete
|
||||
with suppress_instrumentation():
|
||||
self._ensure_loop()
|
||||
self._await_in_loop(
|
||||
|
||||
# kh added
|
||||
# logger.info(
|
||||
# "[SET] Exporting spans to STORE; rollout=%s attempt=%s",
|
||||
# span.name, self._rollout_id, self._attempt_id
|
||||
# )
|
||||
|
||||
uploaded_span = self._await_in_loop(
|
||||
# kh: store에 span 저장 시도 (이건 otlp collector로 보내는 게 아님, store에 직접!)
|
||||
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
|
||||
timeout=60.0,
|
||||
timeout=STORE_WRITE_TIMEOUT_SECONDS,
|
||||
)
|
||||
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}")
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
self._spans.append(span)
|
||||
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)
|
||||
|
||||
+540
-170
@@ -2,63 +2,260 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures as futures
|
||||
import logging
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Tuple, Union
|
||||
import re
|
||||
import weakref
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
cast,
|
||||
)
|
||||
|
||||
from agentlightning.instrumentation import instrument_weave, uninstrument_weave
|
||||
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.tracer import OtelResource, Span, SpanContext, TraceStatus
|
||||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from weave.trace.call import Call # type: ignore
|
||||
|
||||
JSONPrimitive = Union[str, int, float, bool, None]
|
||||
from .base import Tracer, with_active_tracer_context
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WeaveTracer(Tracer):
|
||||
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`
|
||||
"""
|
||||
Tracer implementation using Weave for telemetry and trace logging.
|
||||
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.
|
||||
|
||||
Attributes:
|
||||
project_name: Name of the Weave project. Used to initialize the Weave client.
|
||||
_store: Optional LightningStore instance for storing collected spans.
|
||||
instrument_managed: Whether to patch the Weave/W&B integration to bypass actual network calls for testing.
|
||||
and logs them to Weave Cloud (W&B backend) or optionally bypasses the network for testing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, *, project_name: str | None = None, wandb_api_key: str | None = None, instrument_managed: bool = True
|
||||
self,
|
||||
*,
|
||||
project_name: str | None = None,
|
||||
weave_user_settings: UserSettings | None = None,
|
||||
instrument_managed: bool = True,
|
||||
):
|
||||
"""
|
||||
Initialize a WeaveTracer instance.
|
||||
"""Initialize a WeaveTracer instance.
|
||||
|
||||
Args:
|
||||
project_name: Optional project name for Weave; defaults to the current module name.
|
||||
wandb_api_key: Optional W&B API key; sets environment variable if provided.
|
||||
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 or __name__
|
||||
self.sequence_id = 0
|
||||
self._store: Optional[LightningStore] = None
|
||||
self.project_name = project_name
|
||||
self.instrument_managed = instrument_managed
|
||||
self.weave_user_settings = weave_user_settings or UserSettings(use_server_cache=False)
|
||||
|
||||
if wandb_api_key:
|
||||
os.environ["WANDB_API_KEY"] = wandb_api_key
|
||||
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()
|
||||
instrument_weave(self._server)
|
||||
|
||||
def uninstrument(self, worker_id: int):
|
||||
uninstrument_weave()
|
||||
@@ -75,22 +272,34 @@ class WeaveTracer(Tracer):
|
||||
logger.info(f"[Worker {worker_id}] Setting up Weave tracer...")
|
||||
self._store = store
|
||||
|
||||
try:
|
||||
import weave
|
||||
except ImportError:
|
||||
raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
|
||||
|
||||
# Optionally patch network calls to bypass real Weave/W&B endpoints
|
||||
if self.instrument_managed:
|
||||
self.instrument(worker_id)
|
||||
|
||||
# Initialize the Weave client if not already initialized
|
||||
if weave.get_client() is None: # type: ignore
|
||||
try:
|
||||
weave.init(project_name=self.project_name) # type: ignore
|
||||
logger.info(f"[Worker {worker_id}] Weave client initialized.")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to initialize Weave for project '{self.project_name}': {e}")
|
||||
# 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):
|
||||
"""
|
||||
@@ -105,21 +314,20 @@ class WeaveTracer(Tracer):
|
||||
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,
|
||||
*,
|
||||
store: Optional[LightningStore] = None,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[Any]:
|
||||
"""
|
||||
Synchronous implementation of the tracing context.
|
||||
"""Asynchronous implementation of the tracing context.
|
||||
|
||||
Args:
|
||||
name: Optional operation name.
|
||||
store: Optional LightningStore instance.
|
||||
rollout_id: Optional rollout ID.
|
||||
attempt_id: Optional attempt ID.
|
||||
|
||||
@@ -127,181 +335,343 @@ class WeaveTracer(Tracer):
|
||||
ValueError: If store, rollout_id, and attempt_id are inconsistently provided.
|
||||
RuntimeError: If Weave is not installed or client is uninitialized.
|
||||
"""
|
||||
arg_op = name or self.project_name
|
||||
arg_inputs: dict[str, str] | None = {"rollout_id": rollout_id or "", "attempt_id": attempt_id or ""}
|
||||
|
||||
if store is not None and rollout_id is not None and attempt_id is not None:
|
||||
if rollout_id is not None and attempt_id is not None:
|
||||
self._rollout_id = rollout_id
|
||||
self._attempt_id = attempt_id
|
||||
self._store = store
|
||||
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("store, rollout_id, and attempt_id must be either all provided")
|
||||
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:
|
||||
import datetime
|
||||
# Create a new trace call object in Weave
|
||||
trace_call = weave_client.create_call( # pyright: ignore[reportUnknownMemberType]
|
||||
op=arg_op, inputs=arg_inputs
|
||||
)
|
||||
|
||||
import weave
|
||||
except ImportError:
|
||||
raise RuntimeError("Weave is not installed. Install it to use WeaveTracer.")
|
||||
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
|
||||
|
||||
weave_client = weave.get_client() # type: ignore
|
||||
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
|
||||
|
||||
# Create a new trace call object in Weave
|
||||
trace_call = weave_client.create_call(op=arg_op, inputs=arg_inputs) # type: ignore
|
||||
trace_call.started_at = datetime.datetime.now(tz=datetime.timezone.utc)
|
||||
def _ensure_loop(self) -> tuple[asyncio.AbstractEventLoop, bool]:
|
||||
"""Returns a usable event loop and a boolean indicating whether it's the current running loop.
|
||||
|
||||
try:
|
||||
yield trace_call
|
||||
except Exception as e:
|
||||
# Finish trace and log any exception
|
||||
weave_client.finish_call(trace_call, exception=e) # type: ignore
|
||||
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
|
||||
finally:
|
||||
# Finish trace even if no exception
|
||||
weave_client.finish_call(trace_call) # type: ignore
|
||||
await self._on_finish_handler(trace_call) # type: ignore
|
||||
|
||||
async def _on_finish_handler(self, call: "Call", *args: Any, **kwargs: Any) -> None: # type: ignore
|
||||
Prefer using the main loop if it's possible. Otherwise, use the current running loop.
|
||||
"""
|
||||
Handler called when a Weave Call finishes.
|
||||
# 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.
|
||||
"""
|
||||
spans, self.sequence_id = self.convert_call_to_spans(call, self._rollout_id, self._attempt_id, self.sequence_id) # type: ignore
|
||||
# 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_many_spans(spans)
|
||||
except Exception as e:
|
||||
logger.exception(f"Error adding span to store: {e}")
|
||||
await self._store.add_span(span)
|
||||
except Exception as exc:
|
||||
logger.exception(f"Error adding span to store: {exc}")
|
||||
|
||||
def convert_call_to_spans(
|
||||
async def convert_call_to_span(
|
||||
self,
|
||||
call: "Call", # type: ignore
|
||||
call: tsi.CallSchema,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
seq_start: int = 0,
|
||||
) -> tuple[List[Span], int]:
|
||||
"""
|
||||
Recursively convert a Weave Call (with nested children) into a flat list of Agent Lightning Spans.
|
||||
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.
|
||||
seq_start: Sequence number to start from.
|
||||
sequence_id: Optional sequence ID to attach to spans.
|
||||
|
||||
Returns:
|
||||
Tuple of (list_of_spans, next_sequence_id).
|
||||
List of converted spans.
|
||||
"""
|
||||
spans: List[Span] = []
|
||||
sequence_id = seq_start
|
||||
rollout_id = rollout_id or "rollout-dummy"
|
||||
attempt_id = attempt_id or "attempt-dummy"
|
||||
sequence_id = sequence_id or 0
|
||||
|
||||
rollout_id = rollout_id or "" # type: ignore
|
||||
attempt_id = attempt_id or "" # type: ignore
|
||||
start_ts: float = call.started_at.timestamp()
|
||||
end_ts: Optional[float] = call.ended_at.timestamp() if call.ended_at else None
|
||||
|
||||
start_dt = getattr(call, "started_at", None) # type: ignore
|
||||
start_ts: Optional[float] = start_dt.timestamp() if start_dt else None
|
||||
if call.exception:
|
||||
status = TraceStatus(status_code="ERROR", description=call.exception)
|
||||
else:
|
||||
status = TraceStatus(status_code="OK")
|
||||
|
||||
end_dt = getattr(call, "ended_at", None) # type: ignore
|
||||
end_ts: Optional[float] = end_dt.timestamp() if end_dt else None
|
||||
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
|
||||
|
||||
trace_id = str(getattr(call, "trace_id", None)) # type: ignore
|
||||
span_id = str(getattr(call, "id", None)) # type: ignore
|
||||
parent_id = str(getattr(call, "parent_id", None)) if getattr(call, "parent_id", None) else None # type: ignore
|
||||
|
||||
exception = getattr(call, "exception", None) # type: ignore
|
||||
status_code = "ERROR" if exception else "OK"
|
||||
|
||||
def sanitize(
|
||||
inputs: Dict[str, Any],
|
||||
output: Dict[str, Any],
|
||||
) -> Dict[str, str | JSONPrimitive]:
|
||||
stack: List[Tuple[Any, str]] = [
|
||||
(inputs or {}, "input"),
|
||||
(output or {}, "output"),
|
||||
]
|
||||
|
||||
attributes: Dict[str, str | JSONPrimitive] = {}
|
||||
|
||||
while stack:
|
||||
value, key = stack.pop()
|
||||
|
||||
if isinstance(value, dict):
|
||||
for k, v in value.items(): # type: ignore
|
||||
stack.append((v, f"{key}.{k}")) # type: ignore
|
||||
elif isinstance(value, (list, tuple)):
|
||||
for i, v in enumerate(value): # type: ignore
|
||||
stack.append((v, f"{key}.{i}")) # type: ignore
|
||||
else:
|
||||
if value is None:
|
||||
attributes[key] = "None"
|
||||
elif isinstance(value, (str, int, float, bool)):
|
||||
attributes[key] = value
|
||||
else:
|
||||
try:
|
||||
attributes[key] = str(value)
|
||||
except Exception:
|
||||
attributes[key] = "None"
|
||||
|
||||
return attributes
|
||||
|
||||
inputs = getattr(call, "inputs", {}) # type: ignore
|
||||
output = getattr(call, "output", {}) # type: ignore
|
||||
attributes = sanitize(inputs, output)
|
||||
sanitized_attributes = sanitize_attributes(flatten_attributes(attributes, expand_leaf_lists=False))
|
||||
|
||||
context = SpanContext(
|
||||
trace_id=trace_id,
|
||||
span_id=span_id,
|
||||
trace_id=call.trace_id,
|
||||
span_id=call.id,
|
||||
is_remote=False,
|
||||
trace_state={},
|
||||
)
|
||||
|
||||
parent_context = (
|
||||
SpanContext(
|
||||
trace_id=trace_id,
|
||||
span_id=parent_id,
|
||||
is_remote=False,
|
||||
trace_state={},
|
||||
)
|
||||
if parent_id
|
||||
else None
|
||||
)
|
||||
# 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
|
||||
span = Span(
|
||||
rollout_id=rollout_id or "",
|
||||
attempt_id=attempt_id or "",
|
||||
return Span(
|
||||
rollout_id=rollout_id,
|
||||
attempt_id=attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
trace_id=trace_id,
|
||||
span_id=span_id,
|
||||
parent_id=parent_id,
|
||||
name=getattr(call, "func_name", "unknown"), # type: ignore
|
||||
status=TraceStatus(status_code=status_code),
|
||||
attributes=attributes, # type: ignore
|
||||
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={}, schema_url=""),
|
||||
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="",
|
||||
),
|
||||
)
|
||||
|
||||
spans.append(span)
|
||||
sequence_id += 1
|
||||
|
||||
children: List["Call"] = getattr(call, "_children", []) # type: ignore
|
||||
# Recursively process child calls
|
||||
for child in children: # type: ignore
|
||||
child_spans, sequence_id = self.convert_call_to_spans( # type: ignore
|
||||
child, # type: ignore
|
||||
rollout_id=rollout_id,
|
||||
attempt_id=attempt_id,
|
||||
seq_start=sequence_id,
|
||||
)
|
||||
spans.extend(child_spans)
|
||||
|
||||
return spans, sequence_id
|
||||
|
||||
@@ -28,7 +28,7 @@ from typing import (
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from .tracer import Span
|
||||
from .tracer import Span, SpanCoreFields
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.litagent import LitAgent
|
||||
@@ -307,6 +307,7 @@ 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,11 +2,13 @@
|
||||
|
||||
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, Optional, Sequence, Union
|
||||
from typing import Any, Dict, List, Literal, Optional, Protocol, Sequence, Union
|
||||
|
||||
from opentelemetry import trace as trace_api
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
@@ -31,6 +33,9 @@ __all__ = [
|
||||
"SpanNames",
|
||||
"SpanAttributeNames",
|
||||
"SpanLike",
|
||||
"StatusCode",
|
||||
"SpanCoreFields",
|
||||
"SpanRecordingContext",
|
||||
]
|
||||
|
||||
|
||||
@@ -83,6 +88,8 @@ 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):
|
||||
@@ -115,7 +122,7 @@ class SpanContext(BaseModel):
|
||||
class TraceStatus(BaseModel):
|
||||
"""Serializable variant of `opentelemetry.trace.Status`."""
|
||||
|
||||
status_code: str
|
||||
status_code: StatusCode
|
||||
"""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."""
|
||||
@@ -203,6 +210,44 @@ 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.
|
||||
|
||||
@@ -340,6 +385,7 @@ 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,
|
||||
@@ -357,6 +403,7 @@ 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.
|
||||
@@ -384,7 +431,7 @@ class Span(BaseModel):
|
||||
name=name or AGL_VIRTUAL,
|
||||
resource=resource or OtelResource(attributes={}, schema_url=""),
|
||||
attributes=attributes,
|
||||
status=TraceStatus(status_code="OK"),
|
||||
status=status or TraceStatus(status_code="OK"),
|
||||
events=[],
|
||||
links=[],
|
||||
parent=(
|
||||
@@ -399,6 +446,37 @@ 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]."""
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# 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]
|
||||
+152
-12
@@ -2,22 +2,25 @@
|
||||
|
||||
"""Utilities shared for OpenTelemetry span (attributes) support."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, List, Sequence, Union, cast
|
||||
import traceback
|
||||
from typing import Any, Dict, List, Sequence, Type, TypeVar, 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, SynchronousMultiSpanProcessor, Tracer
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits, SpanProcessor, 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 SpanLike
|
||||
from agentlightning.types import Attributes, AttributeValue, SpanLike
|
||||
from agentlightning.utils.otlp import LightningStoreOTLPExporter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -35,8 +38,16 @@ __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":
|
||||
@@ -112,6 +123,25 @@ 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.
|
||||
|
||||
@@ -166,7 +196,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})
|
||||
return flatten_attributes({LightningSpanAttributes.TAG.value: tags}, expand_leaf_lists=True)
|
||||
|
||||
|
||||
def extract_tags_from_attributes(attributes: Dict[str, Any]) -> List[str]:
|
||||
@@ -196,10 +226,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})
|
||||
return flatten_attributes({LightningSpanAttributes.LINK.value: link_list}, expand_leaf_lists=True)
|
||||
|
||||
|
||||
def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]) -> List[SpanLike]:
|
||||
def query_linked_spans(spans: Sequence[T_SpanLike], links: List[LinkPydanticModel]) -> List[T_SpanLike]:
|
||||
"""Query spans that are linked by the given link attributes.
|
||||
|
||||
Args:
|
||||
@@ -209,7 +239,7 @@ def query_linked_spans(spans: Sequence[SpanLike], links: List[LinkPydanticModel]
|
||||
Returns:
|
||||
A list of spans that match the given link attributes.
|
||||
"""
|
||||
matched_spans: List[SpanLike] = []
|
||||
matched_spans: List[T_SpanLike] = []
|
||||
|
||||
for span in spans:
|
||||
span_attributes = span.attributes or {}
|
||||
@@ -294,7 +324,9 @@ 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]]) -> Dict[str, Any]:
|
||||
def flatten_attributes(
|
||||
nested_data: Union[Dict[str, Any], List[Any]], *, expand_leaf_lists: bool = False
|
||||
) -> 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
|
||||
@@ -303,12 +335,14 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
|
||||
|
||||
Example:
|
||||
|
||||
>>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}})
|
||||
>>> flatten_attributes({"a": {"b": 1, "c": [2, 3]}}, expand_leaf_lists=True)
|
||||
{"a.b": 1, "a.c.0": 2, "a.c.1": 3}
|
||||
|
||||
Args:
|
||||
nested_data: A nested structure composed of dictionaries, lists, or
|
||||
primitive values.
|
||||
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.
|
||||
|
||||
Returns:
|
||||
A flat dictionary mapping dotted-string paths to primitive values.
|
||||
@@ -316,6 +350,15 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
|
||||
|
||||
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():
|
||||
@@ -326,7 +369,22 @@ def flatten_attributes(nested_data: Union[Dict[str, Any], List[Any]]) -> Dict[st
|
||||
new_prefix = f"{prefix}.{k}" if prefix else k
|
||||
_walk(v, new_prefix)
|
||||
elif isinstance(value, list):
|
||||
for idx, item in enumerate(cast(List[Any], value)):
|
||||
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):
|
||||
new_prefix = f"{prefix}.{idx}" if prefix else str(idx)
|
||||
_walk(item, new_prefix)
|
||||
else:
|
||||
@@ -399,3 +457,85 @@ 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, Optional, Sequence, Tuple, Type, TypeVar
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar
|
||||
|
||||
from fastapi import Request, Response
|
||||
from google.protobuf import json_format
|
||||
@@ -39,6 +39,7 @@ from agentlightning.types.tracer import (
|
||||
OtelResource,
|
||||
Span,
|
||||
SpanContext,
|
||||
StatusCode,
|
||||
TraceStatus,
|
||||
convert_timestamp,
|
||||
)
|
||||
@@ -413,7 +414,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 = {
|
||||
_STATUS_CODE_MAP: Mapping[ProtoStatus.StatusCode.ValueType, StatusCode] = {
|
||||
ProtoStatus.STATUS_CODE_UNSET: "UNSET",
|
||||
ProtoStatus.STATUS_CODE_OK: "OK",
|
||||
ProtoStatus.STATUS_CODE_ERROR: "ERROR",
|
||||
|
||||
@@ -13,12 +13,20 @@ 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.05),
|
||||
"cpu_usage_pct": psutil.cpu_percent(0.0),
|
||||
}
|
||||
|
||||
# Memory
|
||||
@@ -37,20 +45,21 @@ def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
|
||||
"disk_pct": du.percent,
|
||||
}
|
||||
|
||||
# GPU
|
||||
# GPU (only query if explicitly requested)
|
||||
gpus: List[Dict[str, Any]] = []
|
||||
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,
|
||||
}
|
||||
)
|
||||
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,
|
||||
}
|
||||
)
|
||||
|
||||
# Network
|
||||
net = psutil.net_io_counters()
|
||||
|
||||
@@ -8,6 +8,12 @@ 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
|
||||
mismatch_log_dir: ./mismatch_cases # supported in trajectory level aggregation with debug=True, directory to store logs of mismatch cases
|
||||
|
||||
data:
|
||||
filter_overlong_prompts: false
|
||||
|
||||
+266
-36
@@ -2,6 +2,7 @@
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import threading
|
||||
@@ -31,6 +32,85 @@ __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]]:
|
||||
@@ -146,6 +226,7 @@ class AgentModeDaemon:
|
||||
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
|
||||
@@ -188,6 +269,7 @@ class AgentModeDaemon:
|
||||
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()
|
||||
@@ -520,7 +602,7 @@ class AgentModeDaemon:
|
||||
raise RuntimeError("Internal loop is not running.")
|
||||
future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
|
||||
try:
|
||||
future.result(timeout=60) # Wait for completion with a timeout
|
||||
future.result(timeout=300) # Wait for completion with a timeout
|
||||
except Exception as e:
|
||||
print(f"Failed to set up data on server: {e}")
|
||||
raise
|
||||
@@ -722,7 +804,9 @@ class AgentModeDaemon:
|
||||
)
|
||||
return metric_dict
|
||||
|
||||
def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
|
||||
def get_train_data_batch(
|
||||
self, max_prompt_length: int, max_response_length: int, device: torch.device, global_steps: int
|
||||
):
|
||||
"""
|
||||
Processes completed rollouts to generate a training data batch.
|
||||
|
||||
@@ -788,50 +872,165 @@ class AgentModeDaemon:
|
||||
image_grid_thw_list: List[Optional[torch.Tensor]] = [] # For Qwen2-VL mrope
|
||||
n_trunc_sample_because_of_response = 0
|
||||
|
||||
for rollout_id, sample_info in finished_id_to_sample_info.items():
|
||||
for turn_index, trace in enumerate(sample_info["trace_list"]):
|
||||
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"]):
|
||||
|
||||
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)
|
||||
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))
|
||||
# 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("mismatch_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')}")
|
||||
|
||||
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)
|
||||
@@ -882,7 +1081,12 @@ 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)
|
||||
@@ -894,12 +1098,38 @@ 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
|
||||
data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_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
|
||||
|
||||
return data_proto, data_metrics
|
||||
|
||||
|
||||
@@ -34,6 +34,9 @@ __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,
|
||||
|
||||
@@ -255,9 +255,18 @@ 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.data.max_prompt_length,
|
||||
max_response_length=self.config.data.max_response_length,
|
||||
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
|
||||
),
|
||||
device=gen_batch.batch["fake_ids"].device,
|
||||
global_steps=self.global_steps,
|
||||
)
|
||||
metrics.update(agent_metrics)
|
||||
self.agent_mode_daemon.clear_data_and_server()
|
||||
@@ -282,7 +291,8 @@ 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"]
|
||||
|
||||
batch.batch["response_mask"] = compute_response_mask(batch)
|
||||
if "response_mask" not in batch.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()
|
||||
@@ -466,6 +476,7 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
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 +1,4 @@
|
||||
# Put code owner definitions here.
|
||||
|
||||
# Recipes
|
||||
recipes/search_r1 @SiyunZhao @JiahangXu
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
## Overview
|
||||
|
||||
This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.1.x**.
|
||||
This example implements **Search R1** within Agent Lightning. It also serves as a demonstration of a **framework-free agent training pipeline**, showing how to run end-to-end RL training without relying on specialized frameworks. **It's tested and compatible with Agent-lightning v0.2.x**.
|
||||
|
||||
The example is designed to run on a single node with 8 GPUs, each having at least 40 GB of memory.
|
||||
|
||||
@@ -14,7 +14,7 @@ The example is designed to run on a single node with 8 GPUs, each having at leas
|
||||
| `retrieval_launch.sh` | Launches the retrieval service backed by the processed corpus |
|
||||
| `retrieval_server.py` | FastAPI server that powers document retrieval during training |
|
||||
| `search_r1_agent.py` | Agent-Lightning rollout script implementing the Search-R1 workflow |
|
||||
| `train.sh` | Starts the RL training server that coordinates GRPO optimization |
|
||||
| `train_search_r1_agent.py` | RL training script that coordinates GRPO optimization |
|
||||
| `qa_em.py` | Exact-match evaluation utilities for validating model predictions |
|
||||
|
||||
---
|
||||
@@ -54,7 +54,7 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv
|
||||
|
||||
---
|
||||
|
||||
## Run RL Training (GRPO) with Llama-3.2-3b-base
|
||||
## Run RL Training (GRPO) with Llama-3.2-3B-Instruct
|
||||
|
||||
1. **Start Ray**
|
||||
|
||||
@@ -65,26 +65,28 @@ The retrieval server implementation is based on `search_r1/search/retrieval_serv
|
||||
> If you plan to use WandB for experiment tracking, set the environment variable
|
||||
> `WANDB_API_KEY` before starting Ray.
|
||||
|
||||
2. **Launch the Agent**
|
||||
|
||||
```bash
|
||||
python search_r1_agent.py
|
||||
```
|
||||
|
||||
This script automatically launches **128 agent workers** by default. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.
|
||||
|
||||
|
||||
3. **Start the Training Server**
|
||||
2. **Start the Training Server**
|
||||
In another terminal, run:
|
||||
|
||||
```bash
|
||||
bash train.sh
|
||||
python train_search_r1_agent.py llama
|
||||
```
|
||||
|
||||
This script starts the RL training server.
|
||||
This script starts the RL training. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.
|
||||
|
||||
---
|
||||
|
||||
## Evaluation
|
||||
## Benchmark Results
|
||||
|
||||
Evaluation scripts and benchmark results will be released soon.
|
||||
We evaluated Search-R1 across seven diverse question-answering benchmarks, covering both General QA (NQ, TriviaQA, PopQA) and complex multi-hop reasoning tasks (HotpotQA, 2WikiMultiHopQA, Musique, and Bamboogle).
|
||||
|
||||
The following tables compare the performance of the original Search-R1 implementation and the Agent-Lightning version across various base models.
|
||||
|
||||
| Model | Source | NQ | TriviaQA | PopQA | HotpotQA | 2Wiki | Musique | Bamboogle |
|
||||
| :--- | :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
||||
| **Qwen2.5-3B-Instruct** | **Search-R1 (Original)** | 34.1 | 54.5 | 37.8 | 32.4 | 31.9 | 10.3 | 26.4 |
|
||||
| | **Agent-Lightning** | **45.3** | **61.7** | **43.8** | **42.6** | **36.4** | **17.1** | **37.6** |
|
||||
| **Qwen2.5-7B-Instruct** | **Search-R1 (Original)** | 39.3 | 61.0 | 39.7 | 37.0 | 41.4 | 14.6 | 36.8 |
|
||||
| | **Agent-Lightning** | **46.5** | **65.9** | **46.8** | **43.7** | **46.2** | **20.3** | **47.2** |
|
||||
| **Llama-3.2-3B** | **Search-R1 (Reproduced)** | 26.3 | 49.0 | 23.0 | 21.6 | 27.3 | 4.5 | 9.7 |
|
||||
| | **Agent-Lightning** | **29.6** | **51.9** | **25.7** | **23.2** | **28.3** | **5.8** | 9.6 |
|
||||
@@ -75,7 +75,7 @@ def extract_solution(solution_str: str) -> Optional[str]:
|
||||
matches = list(match_iter)
|
||||
|
||||
# If there are 0 or exactly 1 matches, return None
|
||||
if len(matches) <= 1:
|
||||
if len(matches) == 0:
|
||||
return None
|
||||
|
||||
# If there are 2 or more matches, return the last one
|
||||
+90
-39
@@ -1,16 +1,21 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple, TypedDict, cast
|
||||
|
||||
import pandas as pd
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
from qa_em import compute_score_em
|
||||
|
||||
from agentlightning import LLM, LitAgent, NamedResources, Trainer, reward, setup_logging
|
||||
from agentlightning import LLM, LitAgent, NamedResources, Rollout, Trainer, configure_logger, setup_logging
|
||||
|
||||
setup_logging()
|
||||
logger = configure_logger(name=__name__)
|
||||
|
||||
# Copied and adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/data_process/nq_search.py
|
||||
INSTRUCTION_FORMAT = """Answer the given question. You must conduct reasoning inside <think> and </think> first every time you get new information. After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search> and it will return the top searched results between <information> and </information>. You can search as many times as your want. If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. For example, <answer> Beijing </answer>. Question: """
|
||||
@@ -24,8 +29,7 @@ class RetrievalItem(TypedDict):
|
||||
document: Document
|
||||
|
||||
|
||||
@reward
|
||||
async def eval(prediction: str, ground_truth: List[str]) -> float:
|
||||
def eval(prediction: str, ground_truth: List[str]) -> float:
|
||||
reward_score = float(compute_score_em(prediction, ground_truth))
|
||||
print(f"pred: {prediction} | {type(ground_truth)} gold_answer: {ground_truth} | res: {reward_score}")
|
||||
return reward_score
|
||||
@@ -106,62 +110,109 @@ def call_llm(
|
||||
return response.choices[0].message.content or ""
|
||||
|
||||
|
||||
class Searchr1Agent(LitAgent[Any]):
|
||||
async def training_rollout_async(
|
||||
class SearchR1Agent(LitAgent[Dict[str, Any]]):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
task: Any,
|
||||
val_temperature: Optional[float] = 0.0,
|
||||
max_turns: int = 4,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.val_temperature = val_temperature
|
||||
self.data_dir = os.environ.get("VERL_SEARCHR1_DATA_DIR", "data")
|
||||
self.max_turns = max_turns
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
task: Dict[str, Any],
|
||||
resources: NamedResources,
|
||||
rollout: Any,
|
||||
temperature: float = 1.0,
|
||||
) -> Any:
|
||||
rollout: Rollout,
|
||||
) -> float | None:
|
||||
prompt = INSTRUCTION_FORMAT + task["question"]
|
||||
answer_list: List[str] = cast(List[str], task["golden_answers"])
|
||||
llm: LLM = cast(LLM, resources.get("main_llm"))
|
||||
rollout_id = rollout.rollout_id
|
||||
logger.info(f"[Rollout {rollout_id}] Question: {task['question']}")
|
||||
logger.info(f"[Rollout {rollout_id}] Ground Truth: {answer_list}")
|
||||
|
||||
start_time = time.time()
|
||||
llm: LLM = cast(LLM, resources["main_llm"])
|
||||
client = OpenAI(
|
||||
base_url=llm.endpoint,
|
||||
base_url=llm.get_base_url(rollout_id, rollout.attempt.attempt_id), # type: ignore
|
||||
api_key=os.environ.get("OPENAI_API_KEY", "token-abc123"),
|
||||
)
|
||||
|
||||
if rollout.mode == "train":
|
||||
temperature = llm.sampling_parameters.get("temperature", 1.0)
|
||||
else:
|
||||
temperature = self.val_temperature if self.val_temperature is not None else 0.0
|
||||
|
||||
turn_id = 0
|
||||
finished_flag = False
|
||||
rollout_content: str = ""
|
||||
|
||||
while turn_id < 4 and not finished_flag:
|
||||
turn_id += 1
|
||||
turn_response = call_llm(
|
||||
client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
|
||||
)
|
||||
valid_turn_response = postprocess_response(turn_response)
|
||||
turn_env_feedback = execute_response(valid_turn_response)
|
||||
if len(turn_env_feedback) == 0:
|
||||
finished_flag = True
|
||||
print(f"TURN ID {turn_id} | RESP: {turn_response} | ENV FEEDBACK: {turn_env_feedback}")
|
||||
rollout_content += turn_response + turn_env_feedback
|
||||
try:
|
||||
while turn_id < self.max_turns and not finished_flag:
|
||||
turn_id += 1
|
||||
turn_response = call_llm(
|
||||
client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
|
||||
)
|
||||
valid_turn_response = postprocess_response(turn_response)
|
||||
rollout_content += valid_turn_response
|
||||
turn_env_feedback = execute_response(valid_turn_response)
|
||||
if len(turn_env_feedback) == 0:
|
||||
finished_flag = True
|
||||
else:
|
||||
rollout_content += turn_env_feedback
|
||||
logger.info(f"TURN ID {turn_id} | RESP: {turn_response} | ENV FEEDBACK: {turn_env_feedback}")
|
||||
|
||||
if not finished_flag:
|
||||
turn_response = call_llm(
|
||||
client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
|
||||
)
|
||||
rollout_content += turn_response
|
||||
print(f"LAST TURN GENERATE | RESP: {turn_response}")
|
||||
if not finished_flag:
|
||||
turn_response = call_llm(
|
||||
client, llm.model, prompt + rollout_content, temperature=temperature, max_tokens=500
|
||||
)
|
||||
rollout_content += turn_response
|
||||
logger.info(f"LAST TURN GENERATE | RESP: {turn_response}")
|
||||
|
||||
reward_score = await eval(rollout_content, answer_list) # reward is tracked with the decorator
|
||||
print(
|
||||
except Exception as e:
|
||||
logger.exception(f"[Rollout {rollout_id}] Error during rollout: {e}")
|
||||
return None
|
||||
|
||||
end_time_rollout = time.time()
|
||||
reward_score = eval(rollout_content, answer_list)
|
||||
logger.info("[Rollout %s] Reward: %s", rollout_id, reward_score)
|
||||
end_time_eval = time.time()
|
||||
|
||||
logger.info("[Rollout %s] Time taken for rollout: %.2f seconds", rollout_id, end_time_rollout - start_time)
|
||||
logger.info(
|
||||
"[Rollout %s] Time taken for evaluation: %.2f seconds", rollout_id, end_time_eval - end_time_rollout
|
||||
)
|
||||
logger.info(
|
||||
"question: {} answer: {} ground_truth: {} reward: {}".format(
|
||||
task["question"], rollout_content, answer_list, reward_score
|
||||
)
|
||||
)
|
||||
return reward_score
|
||||
|
||||
async def validation_rollout_async(
|
||||
self,
|
||||
task: Any,
|
||||
resources: NamedResources,
|
||||
rollout: Any,
|
||||
) -> Any:
|
||||
# Use the same resources; set temperature to 0.0 for deterministic validation.
|
||||
return await self.training_rollout_async(task, resources, rollout, temperature=0.0)
|
||||
|
||||
def debug_search_r1_agent():
|
||||
searchr1_dev_data_path = os.path.join(os.environ.get("VERL_SEARCHR1_DATA_DIR", "data"), "test.parquet")
|
||||
if not os.path.exists(searchr1_dev_data_path):
|
||||
raise FileNotFoundError(f"Search_R1 dev data file {searchr1_dev_data_path} does not exist.")
|
||||
df = pd.read_parquet(searchr1_dev_data_path).head(10) # type: ignore
|
||||
df = cast(List[Dict[str, Any]], df.to_dict(orient="records")) # type: ignore
|
||||
print("Debug data:", df)
|
||||
|
||||
trainer = Trainer(
|
||||
n_workers=1,
|
||||
initial_resources={
|
||||
"main_llm": LLM(
|
||||
endpoint=os.environ["OPENAI_API_BASE"],
|
||||
model="gpt-4.1-nano",
|
||||
sampling_parameters={"temperature": 0.0},
|
||||
)
|
||||
},
|
||||
)
|
||||
trainer.dev(SearchR1Agent(), df)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
Trainer(n_workers=128).fit(Searchr1Agent(), "http://localhost:9999/")
|
||||
debug_search_r1_agent()
|
||||
@@ -0,0 +1,171 @@
|
||||
# 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()
|
||||
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "agent-lightning-dashboard",
|
||||
"version": "0.3.0",
|
||||
"version": "0.3.1",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "agent-lightning-dashboard",
|
||||
"version": "0.3.0",
|
||||
"version": "0.3.1",
|
||||
"dependencies": {
|
||||
"@mantine/core": "8.3.5",
|
||||
"@mantine/hooks": "8.3.5",
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"name": "agent-lightning-dashboard",
|
||||
"type": "module",
|
||||
"version": "0.3.0",
|
||||
"version": "0.3.1",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "tsc && vite build",
|
||||
|
||||
@@ -9,6 +9,11 @@ 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
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 275 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 247 KiB |
File diff suppressed because one or more lines are too long
@@ -1,5 +1,97 @@
|
||||
# Changelog
|
||||
|
||||
## Agent-lightning v0.3.0 (12/24/2025)
|
||||
|
||||
Agent-lightning v0.3.0 is a major release that introduces several new features and bug fixes. The release is a collaborative effort between Agent-lightning core teams and the community. Thanks to all the contributors who made this release possible.
|
||||
|
||||
### Highlights
|
||||
|
||||
* **Tinker integration**: Support Tinker as an alternative backend for Reinforcement Learning (#226 #245 #264 #269 #327). See [example code](https://github.com/microsoft/agent-lightning/tree/v0.3.0/examples/tinker), [blog 1](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-1-1d8c9a397f0e) and [blog 2](https://medium.com/@yugez/tuning-any-ai-agent-with-tinker-agent-lightning-part-2-332c5437f0dc).
|
||||
* **Azure OpenAI integration**: Support Azure OpenAI as a backend for LLM inference and supervised fine-tuning (#256 #327). [Example code](https://github.com/microsoft/agent-lightning/tree/v0.3.0/examples/azure).
|
||||
* **MongoDB-based Lightning Store** is added as an alternative backend for Lightning Store (#323). [Documentation](https://microsoft.github.io/agent-lightning/0.3.0/tutorials/parallelize/#parallelizing-lightningstore).
|
||||
* **Contrib package**: Add contrib package for community projects. Search-R1 is integrated as a contrib recipe. More coming. (#239 #396 #410 #412 #417).
|
||||
* **RESTful API**: Stabilize and document RESTful API for Lightning Store (#241 #275). [Documentation](https://microsoft.github.io/agent-lightning/0.3.0/reference/restful/).
|
||||
* **OTel Semantic Conventions** that are specifically designed for Agent-optimization areas (#340). [Documentation](https://microsoft.github.io/agent-lightning/0.3.0/reference/semconv/).
|
||||
* *[Preview]* **Agent-lightning Dashboard** is now available (#288 #289 #291 #296 #371 #375). It's the official web application for inspecting and debugging Agent-lightning experiments. See details [here](https://microsoft.github.io/agent-lightning/0.3.0/tutorials/debug/).
|
||||
* *[Preview]* **Multi-modality example** featuring VERL and a LangGraph agent on ChartQA dataset (#379). [Example code](https://github.com/microsoft/agent-lightning/tree/v0.3.0/examples/chartqa).
|
||||
* *[Preview]* Integrate **Claude Code** as a LitAgent and support training on SWE-Bench (#332 #346 #348). [Example code](https://github.com/microsoft/agent-lightning/tree/v0.3.0/examples/claude_code).
|
||||
* *[Preview]* **Weave tracer** as a substitute for AgentOps tracer (#277 #411 #420 #423). [Documentation](https://microsoft.github.io/agent-lightning/0.3.0/tutorials/traces/#weave-tracer-experimental).
|
||||
* *[Preview]* **Trajectory Level Aggregation** for more efficient training with VERL. See [blog](https://agent-lightning.github.io/posts/trajectory_level_aggregation/) and [documentation](https://microsoft.github.io/agent-lightning/0.3.0/algorithm-zoo/verl/).
|
||||
|
||||
### Store Benchmark
|
||||
|
||||
In this release, the Lightning Store core was redesigned for significantly greater efficiency and scalability (#315 #318 #328 #342 #344 #356 #380 #388 #418 #421). The benchmark results below demonstrate the impact: with large numbers of concurrent runners, v0.3.0 delivers up to a 15x increase in throughput compared to v0.2.2.
|
||||
|
||||
| Throughput (\#rollout/sec) | v0.2.2 | v0.3.0 (in-memory) | v0.3.0 (Mongo) |
|
||||
| :---- | :---- | :---- | :---- |
|
||||
| Minimal (batch, #runner=32, #turns=6) | 8.73 | 9.06 | 8.71 |
|
||||
| Medium (batch, #runners=100, #turns=10) | 12.03 | 23.26 | 32.79 |
|
||||
| Mid-high (batch, #runners=300, #turns=6) | 10.61 | 24.42 | 40.24 |
|
||||
| Large (batch, #runners=1000, #turns=3) | 3.36 | 14.60 | 50.05 |
|
||||
| Long queue (queue, #runners=256, #turns=4) | 7.42 | 30.86 | 57.01 |
|
||||
| Heavy trace (queue, #runners=512, #turns=20) | 5.93 | 13.28 | 29.41 |
|
||||
|
||||
*Notes:*
|
||||
|
||||
1. Benchmarks were run on a single Standard_D32as_v4 Azure VM (Large and heavy trace tests used Standard_D64ads_v5), executed via GitHub Actions.
|
||||
2. Two algorithm patterns are evaluated: the batch pattern submits a group of rollouts and waits for all to finish before starting the next group, while the queue pattern maintains a set number of in-flight rollouts, submitting new ones as soon as capacity frees up. Configuration details are available [here](https://github.com/microsoft/agent-lightning/blob/v0.3.0/.github/workflows/benchmark.yml).
|
||||
3. The number of turns is directly proportional to the number of spans each rollout generates.
|
||||
|
||||
### Maintenance and Bug fixes
|
||||
|
||||
#### Core (Store, Interfaces, etc.)
|
||||
|
||||
* Add Trainer port option for client-server strategies (#198)
|
||||
* Fix store port conflict handling (#227)
|
||||
* Unified PythonServerLauncher (#286 #292 #303)
|
||||
* Make health timeout configurable (#305)
|
||||
* Refactor logging (#306)
|
||||
* Support OTLP in LightningStore (#313)
|
||||
* Centralized metrics helper (#368)
|
||||
* Fix redundant cancel tracebacks on Ctrl+C (#370)
|
||||
|
||||
#### Proxy, Adapters and Algorithms
|
||||
|
||||
* Fix training metrics before and after processing in VERL (#145)
|
||||
* Forward streaming requests for Anthropic and OpenAI APIs (as non-streaming requests) (#299)
|
||||
* Check traces with reward for VERL (#317)
|
||||
* Patch LiteLLM root span (#341)
|
||||
* Handle ref_in_actor flag for LoRA compatibility (#386)
|
||||
* Support `with_llm_proxy` and `with_store` in algorithms (#398)
|
||||
* Support image URL export in TracerTraceToTriplets (#400)
|
||||
* Fix match_rewards assign_to elements in TraceTree (#403)
|
||||
* Support customizing trainer and daemon in VERL (#407)
|
||||
|
||||
#### Runners, Tracers and Agents
|
||||
|
||||
* Refactor tracer initialization (#321)
|
||||
* Fix OpenAI Agents 0.6 compatibility (#322)
|
||||
* `emit_operation`, `emit_annotation`, tags and links (#359)
|
||||
* Sunset HTTP tracer (#402)
|
||||
|
||||
#### Examples
|
||||
|
||||
* Fix typos in train-first-agent.md (#263)
|
||||
* Fix room_selector example which always runs the first task (#270)
|
||||
* Fix typo in SQL agent example (#285)
|
||||
* Add the README and script files for training SQL agent on NPU (#272)
|
||||
* Examples Catalog and Refine Contribution Guide (#331)
|
||||
* Upgrade LangChain to 1.x (#364)
|
||||
* Update RAG example to Agent-lightning v0.2.x (#349)
|
||||
|
||||
#### Miscellaneous
|
||||
|
||||
* DeepWiki Badge (#263)
|
||||
* Add AGENTS.md (#374)
|
||||
|
||||
### New Contributors
|
||||
|
||||
Warm welcome to our first-time contributors: @cptnm3, @TerryChan, @genji970, @zxgx, @xiaochulaoban, @lspinheiro, @Kwanghoon-Choi, @Vasuk12, @totoluo, @jinghuan-Chen 🎉
|
||||
|
||||
**Full Changelog**: https://github.com/microsoft/agent-lightning/compare/v0.2.0...v0.3.0
|
||||
|
||||
---
|
||||
|
||||
## Agent-lightning v0.2.2 (11/12/2025)
|
||||
|
||||
Agent-lightning v0.2.2 is a stabilization release for v0.2.1. It introduces several bug fixes.
|
||||
|
||||
+207
-19
@@ -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,12 +13,11 @@ 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.
|
||||
|
||||
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.
|
||||
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
|
||||
|
||||
@@ -152,31 +151,220 @@ 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.
|
||||
|
||||
## Store Implementations
|
||||
Check whether the store supports OTLP traces via the [`capabilities["otlp_traces"]`][agentlightning.LightningStore.capabilities] property.
|
||||
|
||||
Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
|
||||
## Implementation Overview
|
||||
|
||||
- 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.
|
||||
The `agentlightning.store` module is organized into two distinct layers plus optional wrappers:
|
||||
|
||||
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.
|
||||
```mermaid
|
||||
classDiagram
|
||||
direction TB
|
||||
|
||||
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.
|
||||
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 | ✗ | ✗ | ✓ | ✓ |
|
||||
|
||||
## 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:
|
||||
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:
|
||||
|
||||
* 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.
|
||||
* 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.
|
||||
|
||||
## Process Safety and Client-server Store
|
||||
|
||||
@@ -188,7 +376,7 @@ Different store implementations may have different capabilities. For example, [`
|
||||
|
||||
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 `/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 `/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.
|
||||
|
||||
Minimal lifecycle:
|
||||
|
||||
|
||||
@@ -62,14 +62,6 @@
|
||||
|
||||
[: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 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.
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
|
||||
+1
-1
@@ -35,7 +35,7 @@ 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).
|
||||
- [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
|
||||
|
||||
+28
-20
@@ -1,7 +1,5 @@
|
||||
# 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.
|
||||
@@ -65,15 +63,40 @@ 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] [--port PORT]
|
||||
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]
|
||||
|
||||
Run a LightningStore server
|
||||
|
||||
options:
|
||||
-h, --help show this help message and exit
|
||||
--port PORT Port to run the server on
|
||||
-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'.
|
||||
```
|
||||
|
||||
!!! tip
|
||||
|
||||
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.
|
||||
@@ -93,18 +116,3 @@ options:
|
||||
Configure the logging level for the metrics server.
|
||||
--access-log Enable uvicorn access logs. Disabled by default to reduce noise.
|
||||
```
|
||||
|
||||
## agl agentops
|
||||
|
||||
Start a mock AgentOps server to bypass the online service of AgentOps.
|
||||
|
||||
```text
|
||||
usage: agl agentops [-h] [--daemon] [--port PORT]
|
||||
|
||||
Start AgentOps server
|
||||
|
||||
options:
|
||||
-h, --help show this help message and exit
|
||||
--daemon Run server as a daemon
|
||||
--port PORT Port to run the server on
|
||||
```
|
||||
|
||||
@@ -54,56 +54,6 @@
|
||||
|
||||
::: agentlightning.tracer.otel.LightningSpanProcessor
|
||||
|
||||
## Utilities
|
||||
|
||||
::: agentlightning.utils.metrics.MetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.ConsoleMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.PrometheusMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.MultiMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.setup_multiprocess_prometheus
|
||||
|
||||
::: agentlightning.utils.metrics.get_prometheus_registry
|
||||
|
||||
::: agentlightning.utils.metrics.shutdown_metrics
|
||||
|
||||
::: 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
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
# RESTful API References
|
||||
|
||||
!!! warning
|
||||
!!! note
|
||||
|
||||
The following contents are still under construction.
|
||||
|
||||
## Store RESTful API
|
||||
Shown in the following is the RESTful API for Lightning Store.
|
||||
|
||||
<div id="swagger-ui"></div>
|
||||
<link rel="stylesheet" href="https://unpkg.com/swagger-ui-dist/swagger-ui.css" />
|
||||
|
||||
@@ -17,3 +17,15 @@
|
||||
::: 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
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# Semantic Conventions
|
||||
|
||||
::: agentlightning.semconv
|
||||
@@ -8,6 +8,8 @@
|
||||
|
||||
::: agentlightning.InMemoryLightningStore
|
||||
|
||||
::: agentlightning.store.mongo.MongoLightningStore
|
||||
|
||||
::: agentlightning.CollectionBasedLightningStore
|
||||
|
||||
## Client-Server and Thread-safe Wrappers
|
||||
@@ -39,3 +41,13 @@
|
||||
::: 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
|
||||
|
||||
@@ -82,9 +82,9 @@
|
||||
|
||||
::: agentlightning.SpanLike
|
||||
|
||||
## Semantic Conventions
|
||||
::: agentlightning.SpanCoreFields
|
||||
|
||||
::: agentlightning.semconv
|
||||
::: agentlightning.SpanRecordingContext
|
||||
|
||||
## Environment Variables
|
||||
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# Utility References
|
||||
|
||||
## ID
|
||||
|
||||
::: agentlightning.utils.id.generate_id
|
||||
|
||||
## Metrics
|
||||
|
||||
::: agentlightning.utils.metrics.MetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.ConsoleMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.PrometheusMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.MultiMetricsBackend
|
||||
|
||||
::: agentlightning.utils.metrics.setup_multiprocess_prometheus
|
||||
|
||||
::: agentlightning.utils.metrics.get_prometheus_registry
|
||||
|
||||
::: agentlightning.utils.metrics.shutdown_metrics
|
||||
|
||||
## Server Launcher
|
||||
|
||||
::: agentlightning.utils.server_launcher.PythonServerLauncher
|
||||
|
||||
::: agentlightning.utils.server_launcher.PythonServerLauncherArgs
|
||||
|
||||
::: agentlightning.utils.server_launcher.LaunchMode
|
||||
|
||||
## OpenTelemetry
|
||||
|
||||
::: 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.otel.sanitize_attribute_value
|
||||
|
||||
::: agentlightning.utils.otel.sanitize_attributes
|
||||
|
||||
::: agentlightning.utils.otel.sanitize_list_attribute_sanity
|
||||
|
||||
::: agentlightning.utils.otel.check_attributes_sanity
|
||||
|
||||
::: agentlightning.utils.otel.format_exception_attributes
|
||||
|
||||
## OTLP
|
||||
|
||||
::: agentlightning.utils.otlp.handle_otlp_export
|
||||
|
||||
::: agentlightning.utils.otlp.spans_from_proto
|
||||
|
||||
## System Snapshot
|
||||
|
||||
::: agentlightning.utils.system_snapshot.system_snapshot
|
||||
@@ -2,6 +2,42 @@
|
||||
|
||||
When you train your own agent with Agent-lightning, most failures surface because the agent logic is brittle or simply incorrect. Debugging becomes easier when you peel back the stack: start by driving the rollout logic on its own, dry-run the trainer loop, and only then bring the full algorithm and runner topology online. The [`examples/apo/apo_debug.py`]({{ src("examples/apo/apo_debug.py") }}) script demonstrates these techniques; this guide expands on each approach and helps you decide when to reach for them.
|
||||
|
||||
## Debugging with Dashboard
|
||||
|
||||
When you launch an experiment with [`Trainer.fit`][agentlightning.Trainer.fit] or start an isolated store via [`agl store`](../reference/cli.md), the terminal prints a message similar to:
|
||||
|
||||
```text
|
||||
INFO Agent-lightning dashboard will be available at http://192.168.0.107:4747
|
||||
```
|
||||
|
||||
Visit that URL, and you will see the Agent-lightning dashboard:
|
||||
|
||||

|
||||
|
||||
The dashboard surfaces everything stored inside [the store](../deep-dive/store.md). Because the store mediates interactions between algorithms and runners, inspecting it often reveals which side is causing issues such as stale rollouts, unresponsive workers, or empty traces.
|
||||
|
||||
For example, the VERL algorithm may receive no token IDs and emit `cannot reshape tensor of 0 elements into shape [1, 0, -1, 128] because the unspecified dimension size -1 can be any value and is ambiguous` ([Issue #50](https://github.com/microsoft/agent-lightning/issues/50), [Issue #76](https://github.com/microsoft/agent-lightning/issues/76)). Several scenarios can produce that error: the runner might not produce trace spans at all, it might produce spans without token IDs, or the IDs may be present but formatted incorrectly. Inspecting the dashboard traces helps you pinpoint which condition applies.
|
||||
|
||||

|
||||
|
||||
By checking whether the trace span is empty and whether token IDs appear in the span attributes, you can narrow the issue to either the runner (agent) side or the algorithm side. Then apply the techniques below to debug the faulty component.
|
||||
|
||||
## Debug-level Logging
|
||||
|
||||
Starting from v0.3, detailed signals such as store server access logs, runner lifecycle logs, and span payloads only appear when the log level is `DEBUG` so the default output stays readable. Enable debug-level logging by adding the following snippet near the top of your script:
|
||||
|
||||
```python
|
||||
import agentlightning as agl
|
||||
|
||||
agl.setup_logging("DEBUG")
|
||||
```
|
||||
|
||||
Set the log level on every process if your setup involves multiple workers. For example, when [running stores in isolation][debug-with-external-store], configure the store process explicitly:
|
||||
|
||||
```bash
|
||||
agl store --port 4747 --log-level DEBUG
|
||||
```
|
||||
|
||||
## Using [`Runner`][agentlightning.Runner] in Isolation
|
||||
|
||||
[`Runner`][agentlightning.Runner] is a long-lived worker that wraps your [`LitAgent`][agentlightning.LitAgent], coordinates tracing, and talks to the [`LightningStore`][agentlightning.LightningStore]. In typical training flows the trainer manages runners for you, but being able to spin one up manually is invaluable while debugging.
|
||||
@@ -255,6 +291,8 @@ In a separate terminal, start the store:
|
||||
agl store --port 4747
|
||||
```
|
||||
|
||||
Add `--log-level DEBUG` to the command to see the detailed logs.
|
||||
|
||||
Then, in your training script, create a [`LightningStoreClient`][agentlightning.LightningStoreClient] and pass it to the trainer:
|
||||
|
||||
```python
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
# Using Emitters
|
||||
|
||||
[](){ #using-emitter }
|
||||
|
||||
While returning a single float for the final reward is sufficient for many algorithm-agent combinations, some advanced scenarios require richer feedback. For instance, an algorithm might learn more effectively if it receives intermediate rewards throughout a multi-step task, or if the agent needs to emit additional spans for debugging or analysis.
|
||||
|
||||
Agent-lightning provides an **emitter** module for recording custom spans inside your agent logic. Just as [Tracer][agentlightning.Tracer] automatically instruments common operations (for example, LLM calls), each emitter helper sends a [Span][agentlightning.Span] that captures Agent-lightning-specific work so downstream algorithms can query it later. See [Working with Traces](./traces.md) for more details.
|
||||
|
||||
For multi-step routines such as function calls, tools, or adapters, wrap code with [`operation`][agentlightning.operation] — either as a decorator or a context manager — to capture inputs, outputs, and metadata on a dedicated [`operation`][agentlightning.operation] span. This makes it easier to correlate downstream annotations (like rewards or messages) with the higher-level work that produced them.
|
||||
|
||||
You can find the emitter functions in [`agentlightning.emitter`](../reference/agent.md).
|
||||
|
||||
## Emitting Rewards, Messages, and More
|
||||
|
||||
Here are the primary emitter functions:
|
||||
|
||||
* [`emit_reward(value: float)`][agentlightning.emit_reward]: Records an intermediate/final reward, which is a convenient wrapper of [`emit_annotation`][agentlightning.emit_annotation].
|
||||
* [`emit_annotation(attributes: Dict[str, Any])`][agentlightning.emit_annotation]: Records arbitrary metadata as a span.
|
||||
* [`emit_message(message: str)`][agentlightning.emit_message]: Records a simple log message as a span.
|
||||
* [`emit_exception(exception: BaseException)`][agentlightning.emit_exception]: Records a Python exception, including its type, message, and stack trace.
|
||||
* [`emit_object(obj: Any)`][agentlightning.emit_object]: Records any JSON-serializable object, perfect for structured data.
|
||||
|
||||
Let's first see an example of an agent using these emitters to provide detailed feedback.
|
||||
|
||||
```python
|
||||
import agentlightning as agl
|
||||
|
||||
@agl.rollout
|
||||
def multi_step_agent(task: dict, prompt_template: PromptTemplate) -> float:
|
||||
try:
|
||||
# Step 1: Initial planning
|
||||
agl.emit_message("Starting planning phase.")
|
||||
plan = generate_plan(task, prompt_template)
|
||||
agl.emit_object({"plan_steps": len(plan), "first_step": plan[0]})
|
||||
|
||||
# Award a small reward for a valid plan
|
||||
plan_reward = grade_plan(plan)
|
||||
agl.emit_reward(plan_reward)
|
||||
|
||||
# Step 2: Execute the plan
|
||||
agl.emit_message(f"Executing {len(plan)}-step plan.")
|
||||
execution_result = execute_plan(plan)
|
||||
|
||||
# Step 3: Final evaluation
|
||||
final_reward = custom_grade_final_result(execution_result, task["expected_output"])
|
||||
|
||||
# The return value is treated as the final reward for the rollout
|
||||
return final_reward
|
||||
|
||||
except ValueError as e:
|
||||
# Record the specific error and return a failure reward
|
||||
agl.emit_exception(e)
|
||||
return 0.0
|
||||
```
|
||||
|
||||
Each helper accepts nested `attributes` (or keyword arguments for [`operation`][agentlightning.operation]) and automatically flattens/sanitizes them into dotted OpenTelemetry keys. This means you can pass ordinary dictionaries/lists without pre-processing and still get consistent attribute names such as `meta.any_attribute` across all emitter operations. Agent-lightning does not restrict the attributes you supply, but it is best to consult [OpenTelemetry's semantic conventions](https://opentelemetry.io/docs/specs/semconv/) for recommended names. Agent-lightning also defines [specific semconv](../reference/semconv.md) for its own use cases. The pattern looks like this:
|
||||
|
||||
```python
|
||||
from opentelemetry.semconv.attributes import server_attributes
|
||||
from agentlightning import emit_object
|
||||
|
||||
emit_object({
|
||||
"name": "John Doe",
|
||||
"age": 30,
|
||||
"email": "john.doe@example.com",
|
||||
}, attributes={
|
||||
server_attributes.SERVER_ADDRESS: "127.0.0.1",
|
||||
server_attributes.SERVER_PORT: 8080,
|
||||
})
|
||||
```
|
||||
|
||||
Running the above code sends the following span to the backend if you have a tracer active:
|
||||
|
||||
```text
|
||||
Span(
|
||||
name='agentlightning.object',
|
||||
attributes={
|
||||
'agentlightning.object.type': 'dict',
|
||||
'agentlightning.object.json': '{"name": "John Doe", "age": 30, "email": "john.doe@example.com"}',
|
||||
'server.address': '127.0.0.1',
|
||||
'server.port': 8080
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
!!! tip
|
||||
|
||||
If you don't have a tracer active, the above code will raise the following error:
|
||||
|
||||
```text
|
||||
RuntimeError: No active tracer found. Cannot emit object span.
|
||||
```
|
||||
|
||||
By default, emitter helpers delegate to the active tracer to create and export spans (specifically via [`Tracer.create_span`][agentlightning.Tracer.create_span]). If you want to emit spans without an active tracer, set `propagate=False` to keep the span local — a useful option for offline tests. The default `True` streams spans through the active tracer/exporters.
|
||||
|
||||
When working with [agentlightning.semconv](../reference/semconv.md), you typically use utilities such as [`make_tag_attributes`][agentlightning.utils.otel.make_tag_attributes] and [`make_link_attributes`][agentlightning.utils.otel.make_link_attributes] to build the attributes dictionary. For example:
|
||||
|
||||
```python
|
||||
from agentlightning.utils.otel import make_tag_attributes
|
||||
|
||||
emit_annotation(make_tag_attributes(["tool", "calculator", "fast", "good"]))
|
||||
```
|
||||
|
||||
The above code will send a span with the following attributes to the backend:
|
||||
|
||||
```json
|
||||
{
|
||||
"agentlightning.tag.0": "tool",
|
||||
"agentlightning.tag.1": "calculator",
|
||||
"agentlightning.tag.2": "fast",
|
||||
"agentlightning.tag.3": "good"
|
||||
}
|
||||
```
|
||||
|
||||
A counterpart utility function [`extract_tags_from_attributes`][agentlightning.utils.otel.extract_tags_from_attributes] is also available to extract the tags from the attributes dictionary.
|
||||
|
||||
## Operations
|
||||
|
||||
The [`operation`][agentlightning.operation] helper tracks logical units of work within your agent, capturing inputs, outputs, timing, and success/failure status. Unlike point-in-time emitters, operations create a span representing a time interval. Use operations for tool calls, multi-step workflows, debugging, and performance monitoring. [`operation`][agentlightning.operation] works as either a decorator or a context manager.
|
||||
|
||||
The decorator automatically captures function arguments as inputs and the return value as output:
|
||||
|
||||
```python
|
||||
import agentlightning as agl
|
||||
|
||||
@agl.operation
|
||||
def search_documents(query: str, max_results: int = 10) -> list[dict]:
|
||||
results = perform_search(query, max_results)
|
||||
return results
|
||||
|
||||
@agl.operation(category="tool", priority="high")
|
||||
def execute_calculation(expression: str) -> float:
|
||||
return eval_safely(expression)
|
||||
```
|
||||
|
||||
The example above emits a span with `{"category": "tool", "priority": "high"}` attributes. It also records the function input and output via [OPERATION_INPUT][agentlightning.semconv.LightningSpanAttributes.OPERATION_INPUT] and [OPERATION_OUTPUT][agentlightning.semconv.LightningSpanAttributes.OPERATION_OUTPUT]. It works with async functions too:
|
||||
|
||||
```python
|
||||
@agl.operation
|
||||
async def async_api_call(endpoint: str, payload: dict) -> dict:
|
||||
response = await http_client.post(endpoint, json=payload)
|
||||
return response.json()
|
||||
```
|
||||
|
||||
Override the operation name if needed:
|
||||
|
||||
```python
|
||||
@agl.operation(name="custom-name")
|
||||
def any_weird_name_i_dont_want():
|
||||
pass
|
||||
```
|
||||
|
||||
For more control, [`operation`][agentlightning.operation] can also be used as a context manager to explicitly record inputs and outputs:
|
||||
|
||||
```python
|
||||
with agl.operation(tool_name="web_search") as op:
|
||||
op.set_input(query="latest AI research", filters={"date": "2024"})
|
||||
results = search_web("latest AI research", {"date": "2024"})
|
||||
op.set_output({"result_count": len(results), "top_result": results[0]})
|
||||
```
|
||||
|
||||
The `propagate=False` flag also applies to [`operation`][agentlightning.operation] when you want to keep operations local without requiring an active tracer:
|
||||
|
||||
```python
|
||||
@agl.operation(propagate=False)
|
||||
def local_test():
|
||||
return "Not sent to backend"
|
||||
```
|
||||
|
||||
## Linking to Other Spans
|
||||
|
||||
Sometimes a span should explicitly point back to another span that produced the input it is working on (for example, linking a reward annotation to the [`agentlightning.operation`][agentlightning.operation] span that generated a response). Agent-lightning encodes these relationships through flattened link attributes. The helper [`make_link_attributes`][agentlightning.utils.otel.make_link_attributes] converts a dictionary of keys such as `trace_id`, `span_id`, or any custom attribute into the `"agentlightning.link.*"` ([LightningSpanAttributes.LINK][agentlightning.semconv.LightningSpanAttributes.LINK]) fields expected by the backend. Later, [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] can recover the original span(s) from those link descriptors.
|
||||
|
||||
```python
|
||||
import opentelemetry.trace as trace_api
|
||||
from agentlightning import emit_annotation, operation
|
||||
from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
|
||||
|
||||
with operation(conversation_id="chat-42") as op:
|
||||
# ... perform the work ...
|
||||
link_attrs = make_link_attributes({
|
||||
"conversation_id": "chat-42",
|
||||
})
|
||||
|
||||
emit_annotation(
|
||||
{
|
||||
**link_attrs,
|
||||
**make_tag_attributes(["reward", "good"]),
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
When analyzing in adapters, pass the extracted link models to [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] to retrieve the matching span(s):
|
||||
|
||||
```python
|
||||
from agentlightning.utils.otel import extract_links_from_attributes, query_linked_spans
|
||||
|
||||
annotation_span = ... # Span from your trace store
|
||||
operation_spans = [...] # list of spans you want to search
|
||||
|
||||
link_models = extract_links_from_attributes(annotation_span.attributes)
|
||||
matches = query_linked_spans(operation_spans, link_models)
|
||||
assert matches # Contains the original operation span
|
||||
```
|
||||
|
||||
!!! tip "Correlating Rewards with LLM Requests"
|
||||
|
||||
[Tracer](./traces.md) instruments each request/response as its own span. You can link to the [`gen_ai.response.id`](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/) attribute, which comes from the LLM response ID.
|
||||
|
||||
```python
|
||||
from agentlightning import emit_reward
|
||||
from agentlightning.utils.otel import make_link_attributes
|
||||
|
||||
result = call_llm(prompt)
|
||||
reward_links = make_link_attributes({"gen_ai.response.id": result.id})
|
||||
emit_reward(0.9, attributes=reward_links)
|
||||
```
|
||||
|
||||
Later, use the same `gen_ai.response.id` key inside `query_linked_spans` to find the reward(s) that reference that specific LLM request span.
|
||||
@@ -34,7 +34,7 @@ This installs or upgrades Agent-Lightning to the newest stable version.
|
||||
Agent-Lightning also publishes **nightly builds**, which contain the latest experimental features and improvements from the main branch. These are available via **Test PyPI**.
|
||||
|
||||
```bash
|
||||
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
|
||||
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ --pre agentlightning
|
||||
```
|
||||
|
||||
!!! warning
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# Scaling out Algorithms and Rollouts
|
||||
# Scaling out Agent-lightning
|
||||
|
||||
Agent-lightning splits training into an **algorithm bundle** and a **runner bundle** that exchange work through the [`LightningStore`][agentlightning.LightningStore]. This tutorial shows how to increase rollout throughput, place bundles across processes or machines, and keep the algorithm side scalable with external frameworks.
|
||||
|
||||
@@ -226,3 +226,109 @@ Agent-lightning strives to make algorithms’ own parallelization work well unde
|
||||
!!! note
|
||||
|
||||
The [birds' eye view][birds-eye-view-client-server-strategy] illustrates how adapters, proxies, and stores interact when the algorithm spawns additional workers. Use that diagram as a checklist when introducing new distributed components.
|
||||
|
||||
## Parallelizing [`LightningStore`][agentlightning.LightningStore]
|
||||
|
||||
By default, Agent-lightning persists rollouts and spans in an in-memory store. [`Trainer.fit`][agentlightning.Trainer.fit] spins it up automatically, or you can launch it yourself via the [`agl store` command](../reference/cli.md). [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] keeps all state inside the current process, which makes local iteration fast but introduces two production constraints:
|
||||
|
||||
1. Spans are evicted once the process crosses its memory cap, so long runs risk data loss unless the host has abundant RAM.
|
||||
2. Although the store is well optimized via asynchronous programming, the store lives in a single process and remains bound by the GIL, preventing it from saturating multi-core machines.
|
||||
|
||||
!!! note "General note for all server-client stores"
|
||||
|
||||
If your algorithm and runners communicate through HTTP protocol (which should be the default for 99% of the cases), you need to ensure the file limit is sufficiently large to avoid the "Too many open files" error. You can set the file limit by running the following command:
|
||||
|
||||
```bash
|
||||
ulimit -n 100000
|
||||
```
|
||||
|
||||
For resilient runs, switch to a persistent backend such as [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], which writes data to MongoDB instead of local RAM. Agent-lightning relies on [pymongo](https://pymongo.readthedocs.io/en/stable/) to interact with MongoDB, which can be installed via:
|
||||
|
||||
```bash
|
||||
pip install agentlightning[mongo]
|
||||
```
|
||||
|
||||
To use the MongoDB store, you need to pass the MongoDB URI to the store constructor. The URI should be in the format of `mongodb://<host>:<port>/<database>?replicaSet=<replicaSet>`.
|
||||
|
||||
```python
|
||||
from agentlightning.store.mongo import MongoLightningStore
|
||||
|
||||
trainer = agl.Trainer(
|
||||
algorithm=algorithm,
|
||||
store=MongoLightningStore(mongo_uri="mongodb://localhost:27017/?replicaSet=rs0"),
|
||||
)
|
||||
```
|
||||
|
||||
!!! tip "Setting up MongoDB"
|
||||
|
||||
MongoDB is a popular document-oriented database. Before running Agent-lightning with [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore], make sure that you've already had a MongoDB instance running. Setting up can be conveniently done via Docker Compose via [compose.mongo.yml]({{ src("docker/compose.mongo.yml") }}). Unless targeting serious production use, we recommend creating the data folders and setting them to `777` permission to avoid permission issues.
|
||||
|
||||
```bash
|
||||
mkdir -p data/mongo-host
|
||||
chmod 777 data/mongo-host
|
||||
docker compose -f compose.mongo.yml up -d
|
||||
```
|
||||
|
||||
Alternatively, you can also install MongoDB manually following the [official documentation](https://www.mongodb.com/docs/manual/installation/). If you installed MongoDB manually, an important note is that you need to ensure that the MongoDB instance has enabled replica set feature, since Agent-lightning uses the transactional operations internally. The simplest approach is to use the following script (executed in the MongoDB shell) to initialize the replica set:
|
||||
|
||||
```javascript
|
||||
rs.initiate({
|
||||
_id: "rs0",
|
||||
members: [{ _id: 0, host: "localhost:27017" }],
|
||||
});
|
||||
```
|
||||
|
||||
To scale out further, launch the store server via [`agl store --backend mongo`](../reference/cli.md) (see [Debugging with External Store][debug-with-external-store]). The CLI accepts `--n-workers`, which starts the server under `gunicorn` with multiple worker processes so concurrent runners can push and pull at higher throughput. This option applies only to persistent backends; an in-memory store, on the other hand, cannot be sharded across workers because its state lives inside one process.
|
||||
|
||||
!!! note
|
||||
|
||||
The `--n-workers` here is the number of worker processes for the store server, NOT related to the number of rollout runners.
|
||||
|
||||
## Increasing Throughput of LLM Proxy
|
||||
|
||||
Agent-lightning includes an optional [`LLMProxy`][agentlightning.LLMProxy] that wraps [LiteLLM](https://docs.litellm.ai/) to provide a unified OpenAI-compatible endpoint for your agents. When rollout throughput increases, the proxy can become a bottleneck. You can scale it out using the same pattern as the store server.
|
||||
|
||||
To increase proxy throughput, pass `num_workers` when constructing the proxy:
|
||||
|
||||
```python
|
||||
import agentlightning as agl
|
||||
|
||||
proxy = agl.LLMProxy(
|
||||
port=4000,
|
||||
launch_mode="mp", # multiprocessing mode
|
||||
num_workers=4, # four gunicorn workers handle concurrent requests
|
||||
)
|
||||
```
|
||||
|
||||
You can also configure the proxy through [`Trainer`][agentlightning.Trainer]:
|
||||
|
||||
```python
|
||||
trainer = agl.Trainer(
|
||||
algorithm=algorithm,
|
||||
n_runners=8, # The runners here is the rollout runners, not related to LLM proxy replicas
|
||||
llm_proxy={"port": 4000, "num_workers": 4}, # launch mode is actually mp by default
|
||||
)
|
||||
```
|
||||
|
||||
When `num_workers > 1`, the launcher starts gunicorn with the specified number of worker processes. Each worker runs its own event loop, allowing the proxy to handle many concurrent LLM requests without being blocked by Python's GIL.
|
||||
|
||||
!!! tip
|
||||
|
||||
When using `mp` launch mode, [`LLMProxy`][agentlightning.LLMProxy] will start the server in a separate process. To make sure the proxy is still accessing the same store as the main process, you need to set the store to be [zero-copy compatible][store-capabilities], which means, either the store is a native zero-copy store like [`MongoLightningStore`][agentlightning.store.mongo.MongoLightningStore] or the store is wrapped via [`LightningStoreServer`][agentlightning.LightningStoreServer] or [`LightningStoreClient`][agentlightning.LightningStoreClient].
|
||||
|
||||
!!! note "Shared Server Infrastructure"
|
||||
|
||||
Both [`LightningStoreServer`][agentlightning.LightningStoreServer] and [`LLMProxy`][agentlightning.LLMProxy] rely on a common utility called [`PythonServerLauncherArgs`][agentlightning.utils.server_launcher.PythonServerLauncherArgs]. This dataclass captures the settings needed to launch a FastAPI application:
|
||||
|
||||
```python
|
||||
from agentlightning.utils import PythonServerLauncherArgs
|
||||
|
||||
args = PythonServerLauncherArgs(
|
||||
port=8000,
|
||||
host="0.0.0.0",
|
||||
n_workers=4, # spawn 4 gunicorn workers
|
||||
launch_mode="thread", # or "mp" for multiprocessing, "asyncio" for in-loop
|
||||
)
|
||||
```
|
||||
|
||||
Under the hood, [`PythonServerLauncher`][agentlightning.utils.server_launcher.PythonServerLauncher] reads these arguments and chooses between uvicorn (single worker) and gunicorn (multiple workers) automatically.
|
||||
|
||||
@@ -48,6 +48,16 @@ If a vendor integration behaves unexpectedly, users are encouraged to combine th
|
||||
|
||||
Inside your agent you can call `opentelemetry.trace.get_trace_provider().get_tracer("my-agent")` and use that tracer to [create spans](https://opentelemetry.io/docs/languages/python/cookbook/) exactly as you would in any OpenTelemetry application. The Lightning span processor attached by [`OtelTracer`][agentlightning.OtelTracer] guarantees that every span is sequenced, converted, and written to the store. The same applies for emitted rewards ([`emit_reward`][agentlightning.emit_reward]) and other emitter signals, which are just a special case of manually-created spans.
|
||||
|
||||
### Weave Tracer (Experimental)
|
||||
|
||||
[`WeaveTracer`][agentlightning.tracer.weave.WeaveTracer] is an experimental tracer that integrates with the [Weave Python SDK](https://docs.wandb.ai/weave). Use it as a substitute for [`AgentOpsTracer`][agentlightning.AgentOpsTracer] when the AgentOps SDK does not fit your environment.
|
||||
|
||||
The Weave SDK instruments LLM calls and agent libraries directly. Unlike [`AgentOpsTracer`][agentlightning.AgentOpsTracer], Weave does not rely on OpenTelemetry to export spans; it routes everything through a dedicated Weave Trace Server. Agent-lightning implements a custom Weave Trace Server so every call captured by the Weave SDK can be persisted to the [`LightningStore`][agentlightning.LightningStore].
|
||||
|
||||
!!! warning
|
||||
|
||||
[`WeaveTracer`][agentlightning.tracer.weave.WeaveTracer] remains experimental and has not been tested as thoroughly as [`AgentOpsTracer`][agentlightning.AgentOpsTracer]. It may conflict with libraries that ship OpenTelemetry instrumentation by default (for example, LiteLLM-based LLM proxies). Use the tracer with caution and report any issues to the Agent-lightning team.
|
||||
|
||||
### LLM Proxy
|
||||
|
||||
Sometimes the runner can’t observe the agent directly — because it’s in another language or running remotely. [`LLMProxy`][agentlightning.LLMProxy] bridges that gap by instrumenting the server side of LLM calls. It wraps [LiteLLM](https://docs.litellm.ai/) and adds middleware that accepts prefixed routes like `/rollout/{rid}/attempt/{aid}/v1/chat/completions`. Before forwarding, the middleware rewrites the path to `/v1/chat/completions`, fetches a monotonic `sequence_id` from the `LightningStore`, injects `x-rollout-id`, `x-attempt-id`, and `x-sequence-id` into the request headers, and then forwards the request to the backend LLM endpoint.
|
||||
|
||||
@@ -115,17 +115,17 @@ The value your agent function returns (i.e., the return value of the function de
|
||||
|
||||
!!! important "Emitting the Final Reward"
|
||||
|
||||
When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Emitter section][using-emitter] below). Wrapping your reward calculation function with the `@reward` decorator is NOT the recommended approach any more.
|
||||
When returning `None`, you must still ensure a final reward is logged. You can do this by using the [`emit_reward`][agentlightning.emit_reward] function (covered in the [Use Emitters](./emitter.md) documentation). Wrapping your reward calculation function with the `@reward` decorator is NOT the recommended approach any more.
|
||||
|
||||
* **`list[ReadableSpan]`** or **`list[Span]`**: For advanced use cases, you can manually construct and return a complete list of all spans for the rollout. This gives you full control over the trace data. You can return either a list of OpenTelemetry `ReadableSpan` objects or Agent-lightning's native `Span` objects.
|
||||
* **`list[ReadableSpan]`**, **`list[SpanCoreFields]`**, or **`list[Span]`**: For advanced use cases, you can manually construct and return a complete list of all spans for the rollout. This gives you full control over the trace data. You can return either a list of OpenTelemetry `ReadableSpan` objects or Agent-lightning's native `Span` objects.
|
||||
|
||||
For most users, returning a **`float`** for simple agents or returning **`None`** and using the emitter for more complex ones are the recommended approaches.
|
||||
|
||||
## Class-based Agents
|
||||
|
||||
For more complex agents that require state, helper methods, or distinct logic for training versus validation, you can create a class that inherits from `LitAgent`. This object-oriented approach provides more structure and control over the agent's lifecycle.
|
||||
For more complex agents that require state, helper methods, or distinct logic for training versus validation, you can create a class that inherits from [`LitAgent`][agentlightning.LitAgent]. This object-oriented approach provides more structure and control over the agent's lifecycle.
|
||||
|
||||
To create a class-based agent, you subclass [agentlightning.LitAgent][] and implement its `rollout` method.
|
||||
To create a class-based agent, you subclass [agentlightning.LitAgent][] and implement its [`rollout`][agentlightning.LitAgent.rollout] method.
|
||||
|
||||
[](){ #introduction-to-named-resources }
|
||||
|
||||
@@ -156,11 +156,11 @@ class RoomSelectorAgent(agl.LitAgent[RoomSelectionTask]):
|
||||
# trainer.fit(agent=agent, ...)
|
||||
```
|
||||
|
||||
The `LitAgent` class provides several methods you can override for more fine-grained control:
|
||||
The [`LitAgent`][agentlightning.LitAgent] class provides several methods you can override for more fine-grained control:
|
||||
|
||||
* `rollout()`: The primary method for the agent's logic. It's called for both training and validation by default.
|
||||
* `training_rollout()` / `validation_rollout()`: Implement these if you need different behavior during training (e.g., with exploration) and validation (e.g., with deterministic choices).
|
||||
* `rollout_async()` / `training_rollout_async()` / `validation_rollout_async()`: Implement the asynchronous versions of these methods if your agent uses `asyncio`.
|
||||
* [`rollout()`][agentlightning.LitAgent.rollout]: The primary method for the agent's logic. It's called for both training and validation by default.
|
||||
* [`training_rollout()`][agentlightning.LitAgent.training_rollout] / [`validation_rollout()`][agentlightning.LitAgent.validation_rollout]: Implement these if you need different behavior during training (e.g., with exploration) and validation (e.g., with deterministic choices).
|
||||
* [`rollout_async()`][agentlightning.LitAgent.rollout_async] / [`training_rollout_async()`][agentlightning.LitAgent.training_rollout_async] / [`validation_rollout_async()`][agentlightning.LitAgent.validation_rollout_async]: Implement the asynchronous versions of these methods if your agent uses `asyncio`.
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -202,110 +202,3 @@ The `LitAgent` class provides several methods you can override for more fine-gra
|
||||
return payload
|
||||
raise payload
|
||||
```
|
||||
|
||||
## Using the Emitter
|
||||
|
||||
[](){ #using-emitter }
|
||||
|
||||
While returning a single float for the final reward is sufficient for many algorithms, some advanced scenarios require richer feedback. For instance, an algorithm might learn more effectively if it receives intermediate rewards throughout a multi-step task.
|
||||
|
||||
Agent-lightning provides an **emitter** module that allows you to record custom spans from within your agent's logic. Like many common operations (like LLM calls) that are automatically instrumented by [Tracer][agentlightning.Tracer], the emitter will also send a [Span][agentlightning.Span] that records an Agent-lightning-specific operation. Then algorithms can query and read those spans later. See [Working with Traces](./traces.md) for more details.
|
||||
|
||||
For multi-step routines (function calls, tools, or adapters) you can wrap code with [`operation`][agentlightning.operation], either as a decorator or a context manager,to capture inputs, outputs, and metadata on a dedicated `"agentlightning.operation"` span. This makes it easier to correlate downstream annotations (like rewards or messages) with the higher-level work that produced them.
|
||||
|
||||
You can find the emitter functions from [agentlightning.emitter](../reference/agent.md).
|
||||
|
||||
### Emitting Rewards, Messages, and More
|
||||
|
||||
Here are the primary emitter functions:
|
||||
|
||||
* [`emit_reward(value: float)`][agentlightning.emit_reward]: Records an intermediate reward.
|
||||
* [`emit_message(message: str)`][agentlightning.emit_message]: Records a simple log message as a span.
|
||||
* [`emit_exception(exception: BaseException)`][agentlightning.emit_exception]: Records a Python exception, including its type, message, and stack trace.
|
||||
* [`emit_object(obj: Any)`][agentlightning.emit_object]: Records any JSON-serializable object, perfect for structured data.
|
||||
Let's see an example of an agent using these emitters to provide detailed feedback.
|
||||
|
||||
```python
|
||||
import agentlightning as agl
|
||||
|
||||
@agl.rollout
|
||||
def multi_step_agent(task: dict, prompt_template: PromptTemplate) -> float:
|
||||
try:
|
||||
# Step 1: Initial planning
|
||||
agl.emit_message("Starting planning phase.")
|
||||
plan = generate_plan(task, prompt_template)
|
||||
agl.emit_object({"plan_steps": len(plan), "first_step": plan[0]})
|
||||
|
||||
# Award a small reward for a valid plan
|
||||
plan_reward = grade_plan(plan)
|
||||
agl.emit_reward(plan_reward)
|
||||
|
||||
# Step 2: Execute the plan
|
||||
agl.emit_message(f"Executing {len(plan)}-step plan.")
|
||||
execution_result = execute_plan(plan)
|
||||
|
||||
# Step 3: Final evaluation
|
||||
final_reward = custom_grade_final_result(execution_result, task["expected_output"])
|
||||
|
||||
# The return value is treated as the final reward for the rollout
|
||||
return final_reward
|
||||
|
||||
except ValueError as e:
|
||||
# Record the specific error and return a failure reward
|
||||
agl.emit_exception(e)
|
||||
return 0.0
|
||||
```
|
||||
|
||||
By using the emitter, you create a rich, detailed trace of your agent's execution. This data can be invaluable for debugging and is essential for advanced algorithms that can learn from more than just a single final score.
|
||||
|
||||
### Linking to Other Spans
|
||||
|
||||
Sometimes a span should explicitly point back to another span that produced the input it is working on (for example, linking a reward annotation to the `"agentlightning.operation"` span that generated a response). Agent-lightning encodes these relationships through flattened link attributes. The helper [`make_link_attributes`][agentlightning.utils.otel.make_link_attributes] converts a dictionary of keys—such as `trace_id`, `span_id`, or any custom attribute—into the `"agentlightning.link.*"` fields expected by the backend. Later on, [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] can be used to recover the original span(s) from those link descriptors.
|
||||
|
||||
```python
|
||||
import opentelemetry.trace as trace_api
|
||||
from agentlightning import emit_annotation, operation
|
||||
from agentlightning.utils.otel import make_link_attributes, make_tag_attributes
|
||||
|
||||
with operation(conversation_id="chat-42") as op:
|
||||
# ... perform the work ...
|
||||
span_ctx = op.span.get_span_context()
|
||||
link_attrs = make_link_attributes({
|
||||
"conversation_id": "chat-42",
|
||||
})
|
||||
|
||||
emit_annotation(
|
||||
{
|
||||
**link_attrs,
|
||||
**make_tag_attributes(["reward", "good"]),
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
When analyzing in adapters, pass the extracted link models to [`query_linked_spans`][agentlightning.utils.otel.query_linked_spans] to retrieve the matching span(s):
|
||||
|
||||
```python
|
||||
from agentlightning.utils.otel import extract_links_from_attributes, query_linked_spans
|
||||
|
||||
annotation_span = ... # Span from your trace store
|
||||
operation_spans = [...] # list of spans you want to search
|
||||
|
||||
link_models = extract_links_from_attributes(annotation_span.attributes)
|
||||
matches = query_linked_spans(operation_spans, link_models)
|
||||
assert matches # Contains the original operation span
|
||||
```
|
||||
|
||||
!!! tip "Correlating Rewards with LLM Requests"
|
||||
|
||||
[Tracer](./traces.md) instruments each request/response as its own span. You can link to the [`gen_ai.response.id`](https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/) attribute, which comes from the LLM response ID.
|
||||
|
||||
```python
|
||||
from agentlightning import emit_reward
|
||||
from agentlightning.utils.otel import make_link_attributes
|
||||
|
||||
result = call_llm(prompt)
|
||||
reward_links = make_link_attributes({"gen_ai.response.id": result.id})
|
||||
emit_reward(0.9, attributes=reward_links)
|
||||
```
|
||||
|
||||
Later, use the same `gen_ai.response.id` key inside `query_linked_spans` to find the reward(s) that reference that specific LLM request span.
|
||||
|
||||
+18
-2
@@ -2,6 +2,8 @@
|
||||
|
||||
This catalog highlights the examples shipped with Agent-lightning.
|
||||
|
||||
Community-contributed examples and recipes are available in the [contrib](../contrib) directory.
|
||||
|
||||
| Example | Description | CI Maintenance |
|
||||
|---------|-------------|----------------|
|
||||
| [apo](./apo) | Automatic Prompt Optimization tutorials covering built-in, custom, and debugging workflows. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml) |
|
||||
@@ -11,9 +13,23 @@ This catalog highlights the examples shipped with Agent-lightning.
|
||||
| [claude_code](./claude_code) | Claude Code SWE-bench harness that records Agent-lightning traces across Anthropic, vLLM, and OpenAI-compatible backends. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml) |
|
||||
| [minimal](./minimal) | Bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation. | [](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
|
||||
| [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml) |
|
||||
| [search_r1](./search_r1) | Framework-free Search-R1 reinforcement learning training workflow with a retrieval backend. | **Last verified with Agent-lightning v0.1.2** |
|
||||
| [spider](./spider) | Text-to-SQL reinforcement learning training on the Spider dataset using LangGraph. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml) |
|
||||
| [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml) |
|
||||
| [unsloth](./unsloth) | Supervised fine-tuning example powered by Unsloth with 4-bit quantization and LoRA. | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-unsloth.yml) |
|
||||
|
||||
*NOTE: CI status avoids taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
|
||||
## `examples-*` workflow status
|
||||
|
||||
CI status above avoids taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. The following table displays the raw `examples-*` workflow status whenever the project is maintained by CI.
|
||||
|
||||
| Workflow | Status |
|
||||
|----------|--------|
|
||||
| `examples-apo.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-apo.yml) |
|
||||
| `examples-azure.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-azure.yml) |
|
||||
| `examples-calc-x.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-calc-x.yml) |
|
||||
| `examples-chartqa.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml) |
|
||||
| `examples-claude-code.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-claude-code.yml) |
|
||||
| `examples-compat.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml) |
|
||||
| `examples-rag.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml) |
|
||||
| `examples-spider.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml) |
|
||||
| `examples-tinker.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-tinker.yml) |
|
||||
| `examples-unsloth.yml` | [](https://github.com/microsoft/agent-lightning/actions/workflows/examples-unsloth.yml) |
|
||||
|
||||
@@ -120,6 +120,9 @@ def train(
|
||||
lora: bool,
|
||||
lora_rank: int,
|
||||
lora_adapter_path: Optional[str],
|
||||
trajectory_level: bool = False,
|
||||
weave: bool,
|
||||
mongo_uri: Optional[str],
|
||||
):
|
||||
"""The training entrypoint function for Calc-X agent with VERL algorithm.
|
||||
|
||||
@@ -135,6 +138,9 @@ def train(
|
||||
lora: Whether to enable LoRA training.
|
||||
lora_rank: LoRA rank to use when LoRA is enabled.
|
||||
lora_adapter_path: Optional path to a pre-trained LoRA adapter to load.
|
||||
trajectory_level: Whether to enable trajectory level in trace aggregator.
|
||||
weave: Whether to enable Weave tracing.
|
||||
mongo_uri: MongoDB URI to use for the store.
|
||||
"""
|
||||
# Load datasets (respect CLI file paths)
|
||||
train_dataset = cast(agl.Dataset[MathProblem], HuggingFaceDataset.from_parquet(train_file).to_list()) # type: ignore
|
||||
@@ -159,6 +165,16 @@ def train(
|
||||
print(f"Loading LoRA adapter from: {lora_adapter_path}")
|
||||
print("LoRA configuration will trigger verl to set ref_in_actor=True (LoRA mode)")
|
||||
|
||||
if trajectory_level:
|
||||
config["agentlightning"] = {
|
||||
"trace_aggregator": {
|
||||
"level": "trajectory",
|
||||
"trajectory_max_prompt_length": 2048,
|
||||
"trajectory_max_response_length": 8192,
|
||||
}
|
||||
}
|
||||
print("Trajectory level enabled in trace aggregator.")
|
||||
|
||||
# CI toggle keeps everything else the same but you can tweak the lightweight bits here if desired
|
||||
if ci or ci_fast:
|
||||
# Config the experiment name and project name so that they are available to CI
|
||||
@@ -202,6 +218,10 @@ def train(
|
||||
|
||||
if external_store_address:
|
||||
store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address)
|
||||
elif mongo_uri:
|
||||
from agentlightning.store.mongo import MongoLightningStore
|
||||
|
||||
store = MongoLightningStore(mongo_uri=mongo_uri)
|
||||
else:
|
||||
store = None
|
||||
|
||||
@@ -209,6 +229,14 @@ def train(
|
||||
tracer = agl.OtelTracer() # dummy tracer for LLM Proxy
|
||||
adapter = agl.LlmProxyTraceToTriplet()
|
||||
trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store, tracer=tracer, adapter=adapter)
|
||||
elif weave:
|
||||
# NOTE: Don't import WeaveTracer at the module level or in __init__.py files.
|
||||
# Always import it lazily/conditionally (behind a feature flag) to avoid interfering
|
||||
# with other libraries like LiteLLM/OpenTelemetry when weave is not explicitly enabled.
|
||||
from agentlightning.tracer.weave import WeaveTracer
|
||||
|
||||
tracer = WeaveTracer()
|
||||
trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store, tracer=tracer)
|
||||
else:
|
||||
trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store)
|
||||
|
||||
@@ -221,6 +249,7 @@ def main():
|
||||
parser.add_argument("--val-file", type=str, default="data/test.parquet", help="Path to val parquet file")
|
||||
parser.add_argument("--model", type=str, default=None, help="HF model id or path (optional)")
|
||||
parser.add_argument("--llm-proxy", action="store_true", help="Enable LLM Proxy tracing/adapter")
|
||||
parser.add_argument("--weave", action="store_true", help="Enable Weave tracing")
|
||||
parser.add_argument("--ci", action="store_true", help="Run a minimal CI-style training loop")
|
||||
parser.add_argument(
|
||||
"--ci-fast", action="store_true", help="Limit the training loop to a single step (implies --ci)"
|
||||
@@ -250,6 +279,17 @@ def main():
|
||||
default=None,
|
||||
help="Optional path to a pre-trained LoRA adapter to load when --lora is enabled",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trajectory-level",
|
||||
action="store_true",
|
||||
help="Enable trajectory level in trace aggregator.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mongo-uri",
|
||||
type=str,
|
||||
default=None,
|
||||
help="MongoDB URI to use for the store.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -278,6 +318,9 @@ def main():
|
||||
lora=args.lora,
|
||||
lora_rank=args.lora_rank,
|
||||
lora_adapter_path=args.lora_adapter_path,
|
||||
trajectory_level=args.trajectory_level,
|
||||
weave=args.weave,
|
||||
mongo_uri=args.mongo_uri,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
[](https://github.com/microsoft/agent-lightning/actions/workflows/examples-chartqa.yml)
|
||||
|
||||
This example demonstrates training a visual reasoning agent on the ChartQA dataset using Agent-Lightning with the VERL algorithm and LangGraph framework. The agent answers questions about charts through a multi-step workflow with self-refinement. It's compatible with Agent-lightning v0.3.0 or later.
|
||||
This example demonstrates training a visual reasoning agent on the ChartQA dataset using Agent-Lightning with the VERL algorithm and LangGraph framework. The agent answers questions about charts through a multi-step workflow with self-refinement.
|
||||
|
||||
## Requirements
|
||||
|
||||
|
||||
@@ -10,19 +10,37 @@ Prior to running this example with `--use-client` flag, please start a Lightning
|
||||
```bash
|
||||
agl store --port 45993 --log-level DEBUG
|
||||
```
|
||||
|
||||
The CLI also ships an `operation` mode showing how to record a synthetic operation span with
|
||||
[`operation`][agentlightning.operation], build link attributes via
|
||||
[`make_link_attributes`][agentlightning.utils.otel.make_link_attributes], tag the
|
||||
follow-up reward with [`make_tag_attributes`][agentlightning.utils.otel.make_tag_attributes],
|
||||
emit a reward span tied back to that operation, and then verify the recorded spans by
|
||||
extracting rewards, tags, and links from the store using `agentlightning.utils.otel` helpers.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import random
|
||||
import time
|
||||
from typing import Sequence
|
||||
from typing import Any, Dict, List, Sequence
|
||||
from uuid import uuid4
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from rich.console import Console
|
||||
|
||||
from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, setup_logging
|
||||
from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, operation, setup_logging
|
||||
from agentlightning.semconv import AGL_OPERATION, LightningSpanAttributes
|
||||
from agentlightning.store import InMemoryLightningStore
|
||||
from agentlightning.utils.otel import get_tracer_provider
|
||||
from agentlightning.utils.otel import (
|
||||
extract_links_from_attributes,
|
||||
extract_tags_from_attributes,
|
||||
filter_and_unflatten_attributes,
|
||||
get_tracer_provider,
|
||||
make_link_attributes,
|
||||
make_tag_attributes,
|
||||
query_linked_spans,
|
||||
)
|
||||
|
||||
console = Console()
|
||||
|
||||
@@ -173,10 +191,111 @@ async def _verify_agentops_traces(spans: Sequence[Span], use_client: bool = Fals
|
||||
assert span.attributes["agentops.span.kind"] == "session"
|
||||
|
||||
|
||||
async def send_operation_links(use_client: bool = False) -> None:
|
||||
"""Demonstrate operation spans wired to reward annotations and verify the stored spans."""
|
||||
|
||||
tracer = OtelTracer()
|
||||
if not use_client:
|
||||
store = InMemoryLightningStore()
|
||||
else:
|
||||
store = LightningStoreClient("http://localhost:45993")
|
||||
conversation_id = "chat-42"
|
||||
tags: Sequence[str] = ("demo.operation", "reward.positive")
|
||||
reward_value = 0.9
|
||||
operation_id = f"{conversation_id}-{uuid4().hex[:8]}"
|
||||
rollout = await store.start_rollout(input={"origin": "write_traces_operation"})
|
||||
|
||||
with tracer.lifespan(store):
|
||||
async with tracer.trace_context(
|
||||
"operation-demo", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
|
||||
):
|
||||
console.print(f"[operation] recording span conversation={conversation_id} operation_id={operation_id}")
|
||||
with operation(conversation_id=conversation_id, operation_id=operation_id) as op_ctx:
|
||||
op_ctx.set_input(
|
||||
task={"conversation_id": conversation_id},
|
||||
metadata={"operation_id": operation_id},
|
||||
)
|
||||
synthetic_payload = {
|
||||
"operation_id": operation_id,
|
||||
"status": "ok",
|
||||
"latency_seconds": round(random.uniform(0.05, 0.2), 3),
|
||||
}
|
||||
await asyncio.sleep(0.05)
|
||||
op_ctx.set_output(synthetic_payload)
|
||||
|
||||
link_attrs = make_link_attributes({"conversation_id": conversation_id, "operation_id": operation_id})
|
||||
tag_attrs = make_tag_attributes(list(tags))
|
||||
emit_reward(
|
||||
reward_value,
|
||||
attributes={**link_attrs, **tag_attrs},
|
||||
)
|
||||
|
||||
spans = await store.query_spans(rollout_id=rollout.rollout_id)
|
||||
console.print(spans)
|
||||
_verify_operation_spans(spans, conversation_id, operation_id, tags, reward_value)
|
||||
|
||||
if isinstance(store, LightningStoreClient):
|
||||
await store.close()
|
||||
|
||||
|
||||
def _verify_operation_spans(
|
||||
spans: Sequence[Span],
|
||||
conversation_id: str,
|
||||
operation_id: str,
|
||||
tags: Sequence[str],
|
||||
expected_reward: float,
|
||||
) -> None:
|
||||
"""Verify spans recorded by the operation demo using OTEL helpers."""
|
||||
|
||||
operation_spans = [span for span in spans if span.name == AGL_OPERATION]
|
||||
if not operation_spans:
|
||||
raise RuntimeError("No operation spans recorded.")
|
||||
console.print(f"[verify] found {len(operation_spans)} operation spans")
|
||||
|
||||
reward_span: Span | None = None
|
||||
reward_payload: List[Dict[str, Any]] = []
|
||||
for span in spans:
|
||||
flattened = dict(span.attributes or {})
|
||||
reward_section = filter_and_unflatten_attributes(flattened, LightningSpanAttributes.REWARD.value)
|
||||
if reward_section:
|
||||
reward_span = span
|
||||
if isinstance(reward_section, list):
|
||||
reward_payload = [dict(item) for item in reward_section] # type: ignore[arg-type]
|
||||
else:
|
||||
reward_payload = [dict(reward_section)] # type: ignore[arg-type]
|
||||
break
|
||||
|
||||
if reward_span is None or not reward_payload:
|
||||
raise RuntimeError("No reward span recorded for operation demo.")
|
||||
|
||||
primary_reward = reward_payload[0].get("value")
|
||||
console.print(f"[verify] reward dimensions: {reward_payload}")
|
||||
if primary_reward != expected_reward:
|
||||
raise AssertionError(f"Expected reward {expected_reward}, observed {primary_reward}")
|
||||
|
||||
reward_attributes = dict(reward_span.attributes or {})
|
||||
extracted_tags = extract_tags_from_attributes(reward_attributes)
|
||||
console.print(f"[verify] reward tags: {extracted_tags}")
|
||||
for tag in tags:
|
||||
if tag not in extracted_tags:
|
||||
raise AssertionError(f"Missing tag '{tag}' on reward span")
|
||||
|
||||
link_models = extract_links_from_attributes(reward_attributes)
|
||||
matches = query_linked_spans(operation_spans, link_models)
|
||||
if not matches:
|
||||
raise AssertionError("No operation span matched the reward links")
|
||||
console.print(f"[verify] reward links resolved spans: {[span.span_id for span in matches]}")
|
||||
|
||||
linked_attrs = dict(matches[0].attributes or {})
|
||||
if linked_attrs.get("conversation_id") != conversation_id or linked_attrs.get("operation_id") != operation_id:
|
||||
raise AssertionError("Linked operation span attributes do not match expected identifiers")
|
||||
console.print("[verify] linked operation span attributes validated")
|
||||
|
||||
|
||||
def main():
|
||||
setup_logging("DEBUG")
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("mode", choices=["otel", "agentops"])
|
||||
parser.add_argument("mode", choices=["otel", "agentops", "operation"])
|
||||
parser.add_argument("--use-client", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -184,6 +303,8 @@ def main():
|
||||
asyncio.run(send_traces_via_otel(use_client=args.use_client))
|
||||
elif args.mode == "agentops":
|
||||
asyncio.run(send_traces_via_agentops(use_client=args.use_client))
|
||||
elif args.mode == "operation":
|
||||
asyncio.run(send_operation_links(use_client=args.use_client))
|
||||
else:
|
||||
raise ValueError(f"Invalid mode: {args.mode}")
|
||||
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -e
|
||||
|
||||
export N_GPUS=8
|
||||
export BASE_MODEL=meta-llama/Llama-3.2-3B
|
||||
export ROLLOUT_TP_SIZE=1
|
||||
export DATA_DIR=data
|
||||
export EXPERIMENT_NAME=searchr1
|
||||
export PROJECT_NAME=AgentLightning-searchr1
|
||||
echo "Starting training script..."
|
||||
|
||||
python -m agentlightning.verl \
|
||||
algorithm.adv_estimator=grpo \
|
||||
data.train_files=${DATA_DIR}/train.parquet \
|
||||
data.val_files=${DATA_DIR}/test.parquet \
|
||||
actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP_SIZE} \
|
||||
trainer.n_gpus_per_node=${N_GPUS} \
|
||||
data.train_batch_size=512 \
|
||||
actor_rollout_ref.rollout.n=5 \
|
||||
actor_rollout_ref.actor.ppo_mini_batch_size=128 \
|
||||
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
|
||||
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
|
||||
actor_rollout_ref.rollout.multi_turn.format=hermes \
|
||||
actor_rollout_ref.model.path=${BASE_MODEL} \
|
||||
data.max_prompt_length=4096 \
|
||||
data.max_response_length=4096 \
|
||||
data.truncation='error' \
|
||||
trainer.val_before_train=True \
|
||||
actor_rollout_ref.actor.optim.lr=1e-6 \
|
||||
actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.95 \
|
||||
actor_rollout_ref.model.use_remove_padding=True \
|
||||
actor_rollout_ref.actor.use_kl_loss=true \
|
||||
actor_rollout_ref.actor.kl_loss_type=low_var_kl \
|
||||
actor_rollout_ref.actor.kl_loss_coef=0.001 \
|
||||
actor_rollout_ref.actor.entropy_coeff=0 \
|
||||
actor_rollout_ref.actor.clip_ratio_low=0.2 \
|
||||
actor_rollout_ref.actor.clip_ratio_high=0.3 \
|
||||
actor_rollout_ref.model.enable_gradient_checkpointing=True \
|
||||
actor_rollout_ref.actor.fsdp_config.param_offload=True \
|
||||
actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
|
||||
actor_rollout_ref.rollout.name=vllm \
|
||||
actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
|
||||
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
|
||||
actor_rollout_ref.ref.fsdp_config.param_offload=True \
|
||||
algorithm.use_kl_in_reward=False \
|
||||
trainer.critic_warmup=0 \
|
||||
trainer.logger=['console','wandb'] \
|
||||
trainer.default_local_dir=checkpoints/searchr1_checkpoints/$EXPERIMENT_NAME \
|
||||
trainer.project_name=${PROJECT_NAME} \
|
||||
trainer.experiment_name=${EXPERIMENT_NAME} \
|
||||
trainer.nnodes=1 \
|
||||
trainer.save_freq=10 \
|
||||
trainer.test_freq=20 \
|
||||
trainer.total_epochs=15 \
|
||||
trainer.total_training_steps=300
|
||||
@@ -60,22 +60,19 @@ def unsloth_training(model_path: str, sft_dataset: HuggingFaceDataset, next_mode
|
||||
loftq_config=None, # And LoftQ
|
||||
)
|
||||
|
||||
sft_config = (
|
||||
SFTConfig(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4, # Use GA to mimic batch size!
|
||||
warmup_steps=5,
|
||||
max_steps=60, # Maximum number of steps to train for
|
||||
# num_train_epochs = 1, # Set this for 1 full training run
|
||||
learning_rate=2e-4, # Reduce to 2e-5 for long training runs
|
||||
logging_steps=1,
|
||||
optim="adamw_8bit",
|
||||
weight_decay=0.01,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
# FIXME: For some reason, report_to="none" still tries to report to W&B when it's installed.
|
||||
report_to="none", # Use this for W&B etc
|
||||
),
|
||||
sft_config = SFTConfig(
|
||||
per_device_train_batch_size=2,
|
||||
gradient_accumulation_steps=4, # Use GA to mimic batch size!
|
||||
warmup_steps=5,
|
||||
max_steps=60, # Maximum number of steps to train for
|
||||
# num_train_epochs = 1, # Set this for 1 full training run
|
||||
learning_rate=2e-4, # Reduce to 2e-5 for long training runs
|
||||
logging_steps=1,
|
||||
optim="adamw_8bit",
|
||||
weight_decay=0.01,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
report_to="none", # Use this for W&B etc
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
|
||||
+4
-1
@@ -108,7 +108,8 @@ nav:
|
||||
- Learning More:
|
||||
- Write Agents: tutorials/write-agents.md
|
||||
- Debugging: tutorials/debug.md
|
||||
- Working with Traces: tutorials/traces.md
|
||||
- Work with Traces: tutorials/traces.md
|
||||
- Use Emitters: tutorials/emitter.md
|
||||
- Parallelize: tutorials/parallelize.md
|
||||
- Algorithm Zoo:
|
||||
- Overview: algorithm-zoo/index.md
|
||||
@@ -124,10 +125,12 @@ nav:
|
||||
- Command Line: reference/cli.md
|
||||
- Instrumentation: reference/instrumentation.md
|
||||
- Runner: reference/runner.md
|
||||
- Semantic Conventions: reference/semconv.md
|
||||
- Store: reference/store.md
|
||||
- Trainer: reference/trainer.md
|
||||
- Types: reference/types.md
|
||||
- RESTful: reference/restful.md
|
||||
- Utilities: reference/utilities.md
|
||||
- Internal: reference/internal.md
|
||||
- Miscellaneous:
|
||||
- Contributing Guide: community/contributing.md
|
||||
|
||||
+20
-11
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "agentlightning"
|
||||
version = "0.3.0"
|
||||
version = "0.3.1"
|
||||
description = "Agent-lightning is the absolute trainer to light up AI agents."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -39,7 +39,7 @@ verl = [
|
||||
]
|
||||
|
||||
weave = [
|
||||
"weave",
|
||||
"weave>=0.52.22",
|
||||
]
|
||||
|
||||
# Store-related dependencies.
|
||||
@@ -288,6 +288,8 @@ override-dependencies = [
|
||||
# verl's numpy<2.0.0 constraint is related to Docker images, not code incompatibility
|
||||
# Lock to 2.3.0 because numba relies on numpy 2.2
|
||||
"numpy>=2.0.0,<2.3.0",
|
||||
# vllm relies on setuptools<80, but polyfile-weave depends on setuptools>=80.9.0
|
||||
"setuptools>=80.9.0",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
@@ -326,19 +328,24 @@ include = [
|
||||
"agentlightning/**/*.yaml",
|
||||
"agentlightning/**/*.yml",
|
||||
"agentlightning/**/*.poml",
|
||||
"agentlightning/**/*.html",
|
||||
"agentlightning/**/*.js",
|
||||
"agentlightning/**/*.css",
|
||||
"agentlightning/**/*.svg",
|
||||
]
|
||||
artifacts = [
|
||||
"agentlightning/dashboard/**",
|
||||
]
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
exclude = [
|
||||
"examples/**",
|
||||
"tests/**",
|
||||
"docs/**",
|
||||
"scripts/**",
|
||||
"dashboard/**",
|
||||
"/examples/**",
|
||||
"/tests/**",
|
||||
"/docs/**",
|
||||
"/scripts/**",
|
||||
"/dashboard/**",
|
||||
"/docker/**",
|
||||
"/contrib/**",
|
||||
"/.github/**",
|
||||
]
|
||||
artifacts = [
|
||||
"agentlightning/dashboard/**",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
@@ -347,11 +354,13 @@ markers = [
|
||||
"openai: tests that require OpenAI API",
|
||||
"gpu: tests that require GPU",
|
||||
"agentops: tests that require AgentOps",
|
||||
"weave: tests that require Weave",
|
||||
"llmproxy: tests that require LiteLLM",
|
||||
"mongo: tests that require MongoDB",
|
||||
"store: tests for agentlightning.store module",
|
||||
"prometheus: tests that require Prometheus",
|
||||
"utils: tests for utility functions",
|
||||
"langchain: tests that require LangChain",
|
||||
]
|
||||
|
||||
[tool.black]
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
# MongoDB Development Setup
|
||||
|
||||
This script is used to setup MongoDB for development.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
@@ -1,10 +0,0 @@
|
||||
services:
|
||||
mongo:
|
||||
image: mongo:latest
|
||||
container_name: mongo-dev
|
||||
ports:
|
||||
- "27017:27017"
|
||||
command: ["mongod", "--bind_ip_all", "--replSet", "rs0"]
|
||||
volumes:
|
||||
- ./data:/data/db
|
||||
- ./init-rs.js:/docker-entrypoint-initdb.d/init-rs.js:ro
|
||||
@@ -1,6 +0,0 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
rs.initiate({
|
||||
_id: "rs0",
|
||||
members: [{ _id: 0, host: "localhost:27017" }],
|
||||
});
|
||||
Executable
+44
@@ -0,0 +1,44 @@
|
||||
#!/bin/bash
|
||||
|
||||
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
|
||||
@@ -44,7 +44,9 @@ else:
|
||||
print(f"::error::Run with name '{run_name}' not found in project '{project}'.")
|
||||
sys.exit(1)
|
||||
|
||||
hist = run.history(keys=["val/reward", "val/n_rollouts_w_reward", "val/n_rollouts_w_trace"], pandas=True)
|
||||
hist = run.history(
|
||||
keys=["val/reward", "val/n_rollouts_w_reward", "val/n_rollouts_w_trace", "val/mean_response_length"], pandas=True
|
||||
)
|
||||
print("History:", hist)
|
||||
if hist.empty:
|
||||
print("::error::No history found for the run.")
|
||||
@@ -88,6 +90,11 @@ else:
|
||||
f"{first_trace_rollouts} -> {last_trace_rollouts}"
|
||||
)
|
||||
|
||||
val_mean_response = last_row["val/mean_response_length"]
|
||||
if val_mean_response < 1:
|
||||
print(f"::error::Mean response length is too short: {val_mean_response} (expected >= 1)")
|
||||
sys.exit(1)
|
||||
|
||||
first_reward, last_reward = first_row["val/reward"], last_row["val/reward"]
|
||||
if last_reward <= first_reward:
|
||||
print(
|
||||
|
||||
@@ -714,7 +714,7 @@ def fmt_rate(value: Optional[float]) -> str:
|
||||
def fmt_latency(value: Optional[float]) -> str:
|
||||
if value is None or math.isnan(value):
|
||||
return "-"
|
||||
if abs(value) < 0.5:
|
||||
if abs(value) < 10:
|
||||
return f"{value * 1e3:.2f} ms"
|
||||
return f"{value:.2f} s"
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import os
|
||||
import random
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Dict, List, Literal, Optional, Sequence, Set, Tuple, cast
|
||||
|
||||
from rich.console import Console
|
||||
@@ -17,9 +18,48 @@ from agentlightning.utils.otel import get_tracer
|
||||
|
||||
from .utils import flatten_dict, random_dict
|
||||
|
||||
console = Console()
|
||||
console = Console(width=200)
|
||||
|
||||
MAX_RUNTIME_SECONDS = 30 * 60
|
||||
# Minus 10 to leave time for setting up env.
|
||||
MAX_RUNTIME_SECONDS = (int(os.getenv("GITHUB_ACTIONS_TIMEOUT_MINUTES", "30")) - 10) * 60
|
||||
MAX_STALE_SECONDS = 300
|
||||
|
||||
|
||||
class RolloutProgressTracker:
|
||||
"""Helper for tracking rollout progress and surfacing stale worker states."""
|
||||
|
||||
def __init__(self, max_stale_seconds: float = MAX_STALE_SECONDS) -> None:
|
||||
self._max_stale_seconds = max_stale_seconds
|
||||
self._last_progress = time.perf_counter()
|
||||
|
||||
def record_progress(self) -> None:
|
||||
self._last_progress = time.perf_counter()
|
||||
|
||||
async def handle_progress(
|
||||
self,
|
||||
*,
|
||||
progress_made: bool,
|
||||
pending_rollout_ids: Sequence[str],
|
||||
store: agl.LightningStore,
|
||||
) -> None:
|
||||
if progress_made:
|
||||
self.record_progress()
|
||||
return
|
||||
await self._check_for_stale(pending_rollout_ids=pending_rollout_ids, store=store)
|
||||
|
||||
async def _check_for_stale(self, *, pending_rollout_ids: Sequence[str], store: agl.LightningStore) -> None:
|
||||
if not pending_rollout_ids:
|
||||
return
|
||||
elapsed = time.perf_counter() - self._last_progress
|
||||
if elapsed <= self._max_stale_seconds / 2:
|
||||
return
|
||||
console.print(f"Stale rollouts: {pending_rollout_ids}")
|
||||
if elapsed > self._max_stale_seconds:
|
||||
current_workers = await store.query_workers()
|
||||
console.print("Stalled. Current worker status shown below:")
|
||||
for worker in current_workers:
|
||||
console.print(f" Worker: {worker}", no_wrap=True, overflow="ignore", crop=False)
|
||||
raise RuntimeError("Rollout progress has stalled for too long")
|
||||
|
||||
|
||||
def _abort_due_to_timeout() -> None:
|
||||
@@ -53,6 +93,8 @@ def make_agent(max_rounds: int, sleep_seconds: float) -> agl.LitAgent[str]:
|
||||
rounds = random.randint(1, max_rounds)
|
||||
selected_round = random.randint(0, rounds - 1)
|
||||
|
||||
# kh: 각 span이 with block에서 종료되면, otlp SDK가 SpanProcessor.on_end(span) 호출,
|
||||
# kh: 내부적으로 span_exporter.export, 미리 지정한 endpoint(collector)로 span 전송 시도.
|
||||
for i in range(rounds):
|
||||
with tracer.start_as_current_span(f"agent{i}") as span:
|
||||
# Nested Span
|
||||
@@ -135,6 +177,7 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
After that, the algorithm will enqueue a new batch of new tasks, until the total number of tasks is reached.
|
||||
"""
|
||||
store = self.get_store()
|
||||
tracker = RolloutProgressTracker()
|
||||
submitted = 0
|
||||
|
||||
while submitted < total_tasks:
|
||||
@@ -157,7 +200,7 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
|
||||
pending = {rollout_id: task_name for rollout_id, task_name in batch_rollouts}
|
||||
completed_ids: Set[str] = set()
|
||||
completed_ids_last_updated: int = 0
|
||||
tracker.record_progress()
|
||||
while len(completed_ids) < len(batch_rollouts):
|
||||
finished_rollouts = await store.wait_for_rollouts(
|
||||
rollout_ids=[rollout_id for rollout_id, _ in batch_rollouts],
|
||||
@@ -170,20 +213,20 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
continue
|
||||
if rollout.status != "succeeded":
|
||||
raise RuntimeError(f"Rollout {rollout_id} finished with status {rollout.status}")
|
||||
spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
|
||||
check_spans(spans, pending[rollout_id])
|
||||
|
||||
# kh: Skipping reading from store
|
||||
# spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
|
||||
# check_spans(spans, pending[rollout_id])
|
||||
|
||||
completed_ids.add(rollout_id)
|
||||
complete_ids_updated = True
|
||||
|
||||
# Check and warn for stale rollouts
|
||||
if complete_ids_updated:
|
||||
completed_ids_last_updated = 0
|
||||
else:
|
||||
completed_ids_last_updated += 1
|
||||
if completed_ids_last_updated >= 10:
|
||||
unfinished_ids = set(rollout_id for rollout_id, _ in batch_rollouts) - completed_ids
|
||||
print(f"Stale rollouts: {unfinished_ids}")
|
||||
completed_ids_last_updated = 0
|
||||
unfinished_ids = [rollout_id for rollout_id, _ in batch_rollouts if rollout_id not in completed_ids]
|
||||
await tracker.handle_progress(
|
||||
progress_made=complete_ids_updated,
|
||||
pending_rollout_ids=unfinished_ids,
|
||||
store=store,
|
||||
)
|
||||
|
||||
await asyncio.sleep(5.0)
|
||||
|
||||
@@ -192,6 +235,7 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
It will enqueue a new batch of new tasks when the number of running rollouts is less than the remaining tasks threshold.
|
||||
"""
|
||||
store = self.get_store()
|
||||
tracker = RolloutProgressTracker()
|
||||
submitted = 0
|
||||
completed = 0
|
||||
active_rollouts: Dict[str, str] = {}
|
||||
@@ -234,6 +278,12 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
completed += 1
|
||||
newly_completed += 1
|
||||
|
||||
await tracker.handle_progress(
|
||||
progress_made=newly_completed > 0,
|
||||
pending_rollout_ids=list(active_rollouts.keys()),
|
||||
store=store,
|
||||
)
|
||||
|
||||
if newly_completed == 0:
|
||||
await asyncio.sleep(5.0)
|
||||
|
||||
@@ -245,6 +295,19 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
"""
|
||||
store = self.get_store()
|
||||
semaphore = asyncio.Semaphore(concurrency)
|
||||
tracker = RolloutProgressTracker()
|
||||
active_rollouts: Set[str] = set()
|
||||
active_lock = asyncio.Lock()
|
||||
|
||||
async def emit_progress(progress_made: bool) -> None:
|
||||
if progress_made:
|
||||
async with active_lock:
|
||||
pending_ids = list(active_rollouts)
|
||||
await tracker.handle_progress(progress_made=True, pending_rollout_ids=pending_ids, store=store)
|
||||
return
|
||||
async with active_lock:
|
||||
pending_ids = list(active_rollouts)
|
||||
await tracker.handle_progress(progress_made=False, pending_rollout_ids=pending_ids, store=store)
|
||||
|
||||
async def handle_single(task_index: int) -> None:
|
||||
task_name = f"task-{task_index}"
|
||||
@@ -260,15 +323,23 @@ class AlgorithmBatch(agl.Algorithm):
|
||||
)
|
||||
rollout = await store.enqueue_rollout(input=task_name, mode="train")
|
||||
rollout_id = rollout.rollout_id
|
||||
while True:
|
||||
current = await store.get_rollout_by_id(rollout_id)
|
||||
if current is not None and current.status in ("failed", "succeeded", "cancelled"):
|
||||
if current.status != "succeeded":
|
||||
raise RuntimeError(f"Rollout {rollout_id} finished with status {current.status}")
|
||||
break
|
||||
await asyncio.sleep(5.0)
|
||||
spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
|
||||
check_spans(spans, task_name)
|
||||
async with active_lock:
|
||||
active_rollouts.add(rollout_id)
|
||||
try:
|
||||
while True:
|
||||
current = await store.get_rollout_by_id(rollout_id)
|
||||
if current is not None and current.status in ("failed", "succeeded", "cancelled"):
|
||||
if current.status != "succeeded":
|
||||
raise RuntimeError(f"Rollout {rollout_id} finished with status {current.status}")
|
||||
break
|
||||
await emit_progress(progress_made=False)
|
||||
await asyncio.sleep(5.0)
|
||||
spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
|
||||
check_spans(spans, task_name)
|
||||
await emit_progress(progress_made=True)
|
||||
finally:
|
||||
async with active_lock:
|
||||
active_rollouts.discard(rollout_id)
|
||||
|
||||
all_tasks = [handle_single(i) for i in range(total_tasks)]
|
||||
await asyncio.gather(*all_tasks)
|
||||
@@ -295,6 +366,8 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
|
||||
parser.add_argument("--n-runners", type=int, default=32, help="Number of runner processes to launch.")
|
||||
parser.add_argument("--max-rounds", type=int, default=10, help="Maximum number of rounds for each rollout.")
|
||||
parser.add_argument("--sleep-seconds", type=float, default=1.0, help="Sleep seconds for each rollout.")
|
||||
parser.add_argument("--debug", action="store_true", help="Enable verbose debug logging.")
|
||||
parser.add_argument("--debug-otel", action="store_true", help="Enable verbose debug logging for OTel.")
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.total_tasks <= 0:
|
||||
@@ -317,7 +390,10 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> None:
|
||||
args = parse_args(argv)
|
||||
agl.setup_logging()
|
||||
agl.setup_logging(
|
||||
"DEBUG" if args.debug else "INFO",
|
||||
submodule_levels={"agentlightning.utils.otel": "DEBUG" if args.debug_otel else "INFO"},
|
||||
)
|
||||
store = agl.LightningStoreClient(args.store_url)
|
||||
timeout_guard = _start_timeout_guard(MAX_RUNTIME_SECONDS)
|
||||
try:
|
||||
@@ -336,6 +412,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
|
||||
"managed_store": False,
|
||||
},
|
||||
)
|
||||
# kh: make_agent에서 정의한 agent가 들어감. 그리고 이 agent는 매 rollout 마다 round, span, attribute 생성하고 sleep도 호출.
|
||||
trainer.fit(make_agent(max_rounds=args.max_rounds, sleep_seconds=args.sleep_seconds))
|
||||
finally:
|
||||
timeout_guard.cancel()
|
||||
|
||||
@@ -1,9 +1,16 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import agentops
|
||||
import agentops.sdk.core as agentops_core
|
||||
import opentelemetry.trace as trace_api
|
||||
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import Attributes, SpanCoreFields, TraceStatus
|
||||
|
||||
|
||||
# pyright: reportPrivateUsage=false
|
||||
def clear_tracer_provider() -> None:
|
||||
@@ -33,3 +40,22 @@ def clear_agentops_init() -> None:
|
||||
"""Make agentops.init() runnable again."""
|
||||
agentops.get_client().initialized = False
|
||||
agentops_core.tracer._initialized = False
|
||||
|
||||
|
||||
class RecordingDummyTracer(DummyTracer):
|
||||
"""Dummy tracer that captures the most recent span request for assertions."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.last_span: Optional[SpanCoreFields] = None
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
span = super().create_span(name, attributes, timestamp, status)
|
||||
self.last_span = span
|
||||
return span
|
||||
|
||||
@@ -2,71 +2,48 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
from typing import Dict
|
||||
|
||||
import opentelemetry.trace as trace_api
|
||||
import pytest
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace import TraceFlags
|
||||
|
||||
from agentlightning.emitter import annotation as annotation_module
|
||||
from agentlightning.emitter.annotation import emit_annotation
|
||||
from agentlightning.semconv import AGL_ANNOTATION
|
||||
|
||||
|
||||
class DummyReadableSpan(ReadableSpan):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(
|
||||
name="dummy",
|
||||
context=trace_api.SpanContext(
|
||||
trace_id=0x1,
|
||||
span_id=0x2,
|
||||
is_remote=False,
|
||||
trace_flags=TraceFlags(TraceFlags.SAMPLED),
|
||||
trace_state=trace_api.TraceState(),
|
||||
),
|
||||
resource=Resource.create({}),
|
||||
)
|
||||
|
||||
def __enter__(self) -> "DummyReadableSpan":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> bool:
|
||||
return False
|
||||
from ..common.tracer import RecordingDummyTracer
|
||||
|
||||
|
||||
class DummyTracer:
|
||||
def __init__(self, span: DummyReadableSpan) -> None:
|
||||
self._span = span
|
||||
self.last_name: str | None = None
|
||||
self.last_attributes: Dict[str, Any] | None = None
|
||||
|
||||
def start_span(self, name: str, attributes: Dict[str, Any] | None = None) -> DummyReadableSpan:
|
||||
self.last_name = name
|
||||
self.last_attributes = attributes or {}
|
||||
return self._span
|
||||
def _install_tracer(monkeypatch: pytest.MonkeyPatch) -> RecordingDummyTracer:
|
||||
tracer = RecordingDummyTracer()
|
||||
monkeypatch.setattr(annotation_module, "get_active_tracer", lambda: tracer)
|
||||
return tracer
|
||||
|
||||
|
||||
def test_emit_annotation_flattens_and_respects_propagation(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = DummyReadableSpan()
|
||||
tracer = DummyTracer(span)
|
||||
captured: Dict[str, Any] = {}
|
||||
def test_emit_annotation_flattens_and_sanitizes_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_tracer(monkeypatch)
|
||||
|
||||
def fake_get_tracer(*_: Any, **kwargs: Any) -> DummyTracer:
|
||||
captured["propagate"] = kwargs.get("use_active_span_processor")
|
||||
return tracer
|
||||
result = emit_annotation({"meta": {"tag": "foo"}, "score": 1.5})
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
|
||||
assert result.name == AGL_ANNOTATION
|
||||
assert tracer.last_span is not None
|
||||
assert tracer.last_span.attributes == {"meta.tag": "foo", "score": 1.5}
|
||||
|
||||
result = emit_annotation({"meta": {"tag": "foo"}, "score": 1.5}, propagate=False)
|
||||
|
||||
assert result is span
|
||||
assert captured["propagate"] is False
|
||||
assert tracer.last_name == AGL_ANNOTATION
|
||||
assert tracer.last_attributes == {"meta.tag": "foo", "score": 1.5}
|
||||
def test_emit_annotation_propagate_false_bypasses_active_tracer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: Dict[str, bool] = {"called": False}
|
||||
|
||||
def fail_get_active_tracer() -> RecordingDummyTracer:
|
||||
captured["called"] = True
|
||||
raise AssertionError("Should not resolve active tracer when propagate is False")
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_active_tracer", fail_get_active_tracer)
|
||||
|
||||
result = emit_annotation({"score": 1}, propagate=False)
|
||||
|
||||
assert result.name == AGL_ANNOTATION
|
||||
assert captured["called"] is False
|
||||
|
||||
|
||||
def test_emit_annotation_rejects_non_primitive_values() -> None:
|
||||
with pytest.raises(TypeError):
|
||||
with pytest.raises(ValueError):
|
||||
emit_annotation({"bad": {"set": {1}}})
|
||||
|
||||
@@ -2,8 +2,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import pytest
|
||||
from opentelemetry.semconv.attributes import exception_attributes
|
||||
|
||||
@@ -11,62 +9,47 @@ from agentlightning.emitter import emit_exception
|
||||
from agentlightning.emitter import exception as exception_module
|
||||
from agentlightning.semconv import AGL_EXCEPTION
|
||||
|
||||
|
||||
class DummySpan:
|
||||
def __init__(self) -> None:
|
||||
self.recorded_exception: Optional[Exception] = None
|
||||
|
||||
def __enter__(self) -> "DummySpan":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
|
||||
return False
|
||||
|
||||
def record_exception(self, exception: Exception) -> None:
|
||||
self.recorded_exception = exception
|
||||
from ..common.tracer import RecordingDummyTracer
|
||||
|
||||
|
||||
class DummyTracer:
|
||||
def __init__(self, span: DummySpan) -> None:
|
||||
self._span = span
|
||||
self.last_name: Optional[str] = None
|
||||
self.last_attributes: Optional[Dict[str, Any]] = None
|
||||
|
||||
def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
|
||||
self.last_name = name
|
||||
self.last_attributes = attributes or {}
|
||||
return self._span
|
||||
|
||||
|
||||
def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
|
||||
tracer = DummyTracer(span)
|
||||
|
||||
def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
|
||||
return tracer
|
||||
|
||||
monkeypatch.setattr(exception_module, "get_tracer", fake_get_tracer)
|
||||
def _install_tracer(monkeypatch: pytest.MonkeyPatch) -> RecordingDummyTracer:
|
||||
tracer = RecordingDummyTracer()
|
||||
monkeypatch.setattr(exception_module, "get_active_tracer", lambda: tracer)
|
||||
return tracer
|
||||
|
||||
|
||||
def test_emit_exception_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = DummySpan()
|
||||
tracer = _stub_tracer(monkeypatch, span)
|
||||
tracer = _install_tracer(monkeypatch)
|
||||
err = ValueError("boom")
|
||||
|
||||
exc: Optional[Exception] = None
|
||||
try:
|
||||
raise ValueError("boom")
|
||||
except ValueError as err:
|
||||
emit_exception(err)
|
||||
exc = err
|
||||
emit_exception(err)
|
||||
|
||||
assert tracer.last_name == AGL_EXCEPTION
|
||||
assert tracer.last_attributes is not None
|
||||
assert tracer.last_attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
|
||||
assert tracer.last_attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
|
||||
assert tracer.last_attributes[exception_attributes.EXCEPTION_ESCAPED] is True
|
||||
assert span.recorded_exception is exc
|
||||
assert tracer.last_span is not None
|
||||
assert tracer.last_span.name == AGL_EXCEPTION
|
||||
assert tracer.last_span.attributes[exception_attributes.EXCEPTION_TYPE] == "ValueError"
|
||||
assert tracer.last_span.attributes[exception_attributes.EXCEPTION_MESSAGE] == "boom"
|
||||
assert tracer.last_span.attributes[exception_attributes.EXCEPTION_ESCAPED] is True
|
||||
|
||||
|
||||
def test_emit_exception_requires_exception_instance() -> None:
|
||||
with pytest.raises(TypeError):
|
||||
emit_exception("boom") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_emit_exception_flattens_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_tracer(monkeypatch)
|
||||
|
||||
emit_exception(ValueError("boom"), attributes={"meta": {"tag": "foo"}, "labels": ["x"]})
|
||||
|
||||
assert tracer.last_span is not None
|
||||
assert tracer.last_span.attributes["meta.tag"] == "foo"
|
||||
assert tracer.last_span.attributes["labels"] == ["x"]
|
||||
|
||||
|
||||
def test_emit_exception_propagate_false(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fail_get_active_tracer() -> RecordingDummyTracer:
|
||||
raise AssertionError("Should not resolve tracer when propagate=False")
|
||||
|
||||
monkeypatch.setattr(exception_module, "get_active_tracer", fail_get_active_tracer)
|
||||
|
||||
emit_exception(ValueError("boom"), propagate=False)
|
||||
|
||||
@@ -13,39 +13,17 @@ from agentlightning.emitter.message import get_message_value
|
||||
from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
|
||||
from agentlightning.types.tracer import SpanLike
|
||||
|
||||
from ..common.tracer import RecordingDummyTracer
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeSpan:
|
||||
attributes: Optional[Dict[str, Any]]
|
||||
|
||||
|
||||
class DummySpan:
|
||||
def __enter__(self) -> "DummySpan":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class DummyTracer:
|
||||
def __init__(self, span: DummySpan) -> None:
|
||||
self._span = span
|
||||
self.last_name: Optional[str] = None
|
||||
self.last_attributes: Optional[Dict[str, Any]] = None
|
||||
|
||||
def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
|
||||
self.last_name = name
|
||||
self.last_attributes = attributes or {}
|
||||
return self._span
|
||||
|
||||
|
||||
def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTracer:
|
||||
tracer = DummyTracer(span)
|
||||
|
||||
def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
|
||||
return tracer
|
||||
|
||||
monkeypatch.setattr(message_module, "get_tracer", fake_get_tracer)
|
||||
def _stub_tracer(monkeypatch: pytest.MonkeyPatch) -> RecordingDummyTracer:
|
||||
tracer = RecordingDummyTracer()
|
||||
monkeypatch.setattr(message_module, "get_active_tracer", lambda: tracer)
|
||||
return tracer
|
||||
|
||||
|
||||
@@ -69,15 +47,34 @@ def test_get_message_value_rejects_non_string() -> None:
|
||||
|
||||
|
||||
def test_emit_message_valid(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = DummySpan()
|
||||
tracer = _stub_tracer(monkeypatch, span)
|
||||
tracer = _stub_tracer(monkeypatch)
|
||||
|
||||
emit_message("hello world")
|
||||
|
||||
assert tracer.last_name == AGL_MESSAGE
|
||||
assert tracer.last_attributes == {LightningSpanAttributes.MESSAGE_BODY.value: "hello world"}
|
||||
assert tracer.last_span is not None
|
||||
assert tracer.last_span.name == AGL_MESSAGE
|
||||
assert tracer.last_span.attributes == {LightningSpanAttributes.MESSAGE_BODY.value: "hello world"}
|
||||
|
||||
|
||||
def test_emit_message_requires_string() -> None:
|
||||
with pytest.raises(TypeError):
|
||||
emit_message(123) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_emit_message_flattens_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _stub_tracer(monkeypatch)
|
||||
|
||||
emit_message("hello", attributes={"meta": {"tag": "foo"}, "labels": ["a", "b"]})
|
||||
|
||||
assert tracer.last_span is not None
|
||||
assert tracer.last_span.attributes["meta.tag"] == "foo"
|
||||
assert tracer.last_span.attributes["labels"] == ["a", "b"]
|
||||
|
||||
|
||||
def test_emit_message_propagate_false(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fail_get_active_tracer() -> RecordingDummyTracer:
|
||||
raise AssertionError("Should not resolve tracer when propagate=False")
|
||||
|
||||
monkeypatch.setattr(message_module, "get_active_tracer", fail_get_active_tracer)
|
||||
|
||||
emit_message("local", propagate=False)
|
||||
|
||||
@@ -35,7 +35,7 @@ class DummyTracer:
|
||||
self.last_name: Optional[str] = None
|
||||
self.last_attributes: Optional[Dict[str, Any]] = None
|
||||
|
||||
def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> DummySpan:
|
||||
def create_span(self, name: str, attributes: Optional[Dict[str, Any]] = None, **kwargs: Any) -> DummySpan:
|
||||
self.last_name = name
|
||||
self.last_attributes = attributes or {}
|
||||
return self._span
|
||||
@@ -47,7 +47,7 @@ def _stub_tracer(monkeypatch: pytest.MonkeyPatch, span: DummySpan) -> DummyTrace
|
||||
def fake_get_tracer(*_: Any, **__: Any) -> DummyTracer:
|
||||
return tracer
|
||||
|
||||
monkeypatch.setattr(object_module, "get_tracer", fake_get_tracer)
|
||||
monkeypatch.setattr(object_module, "get_active_tracer", fake_get_tracer)
|
||||
return tracer
|
||||
|
||||
|
||||
@@ -175,3 +175,23 @@ def test_get_object_value_raises_for_unknown_literal_type() -> None:
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
get_object_value(cast(SpanLike, span))
|
||||
|
||||
|
||||
def test_emit_object_flattens_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = DummySpan()
|
||||
tracer = _stub_tracer(monkeypatch, span)
|
||||
|
||||
emit_object({"foo": "bar"}, attributes={"meta": {"tag": "foo"}, "labels": ["x", "y"]})
|
||||
|
||||
assert tracer.last_attributes is not None
|
||||
assert tracer.last_attributes["meta.tag"] == "foo"
|
||||
assert tracer.last_attributes["labels"] == ["x", "y"]
|
||||
|
||||
|
||||
def test_emit_object_propagate_false(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fail_get_active_tracer() -> DummyTracer:
|
||||
raise AssertionError("Should not resolve tracer when propagate=False")
|
||||
|
||||
monkeypatch.setattr(object_module, "get_active_tracer", fail_get_active_tracer)
|
||||
|
||||
emit_object({"foo": "bar"}, propagate=False)
|
||||
|
||||
+259
-201
@@ -3,9 +3,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from types import TracebackType
|
||||
from typing import Any, Dict, List, Optional, Tuple, Type
|
||||
from typing import Any, ContextManager, Dict, Iterator, List, Optional, Tuple, Type, cast
|
||||
|
||||
import opentelemetry.trace as trace_api
|
||||
import pytest
|
||||
@@ -15,96 +17,141 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
import agentlightning.emitter.annotation as annotation_module
|
||||
from agentlightning.emitter.annotation import _safe_json_dump # pyright: ignore[reportPrivateUsage]
|
||||
from agentlightning.emitter.annotation import (
|
||||
OperationContext,
|
||||
emit_annotation,
|
||||
operation,
|
||||
)
|
||||
from agentlightning.emitter.annotation import OperationContext, emit_annotation, operation
|
||||
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
|
||||
from agentlightning.utils.otel import extract_links_from_attributes, make_link_attributes, query_linked_spans
|
||||
from agentlightning.tracer.dummy import DummySpanRecordingContext, DummyTracer
|
||||
from agentlightning.types import SpanCoreFields, TraceStatus
|
||||
from agentlightning.types.tracer import Attributes
|
||||
from agentlightning.utils.otel import (
|
||||
extract_links_from_attributes,
|
||||
filter_and_unflatten_attributes,
|
||||
make_link_attributes,
|
||||
query_linked_spans,
|
||||
)
|
||||
|
||||
|
||||
class RecordingSpan:
|
||||
class RecordingTracer:
|
||||
def __init__(self) -> None:
|
||||
self.attributes: Dict[str, Any] = {}
|
||||
self.recorded_exceptions: List[BaseException] = []
|
||||
self.statuses: List[Status] = []
|
||||
self._delegate = DummyTracer()
|
||||
self.recordings: List[DummySpanRecordingContext] = []
|
||||
|
||||
def set_attribute(self, key: str, value: Any) -> None:
|
||||
self.attributes[key] = value
|
||||
|
||||
def record_exception(self, exc: BaseException) -> None:
|
||||
self.recorded_exceptions.append(exc)
|
||||
|
||||
def set_status(self, status: Status) -> None:
|
||||
self.statuses.append(status)
|
||||
|
||||
|
||||
class DummySpanContextManager:
|
||||
def __init__(self, span: RecordingSpan) -> None:
|
||||
self.span = span
|
||||
self.exit_calls: List[
|
||||
Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
|
||||
] = []
|
||||
|
||||
def __enter__(self) -> RecordingSpan:
|
||||
return self.span
|
||||
|
||||
def __exit__(
|
||||
def operation_context(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> bool:
|
||||
self.exit_calls.append((exc_type, exc_val, exc_tb))
|
||||
return False
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
) -> ContextManager[DummySpanRecordingContext]:
|
||||
parent_ctx = self._delegate.operation_context(name, attributes, start_time, end_time)
|
||||
|
||||
@contextmanager
|
||||
def _wrapper() -> Iterator[DummySpanRecordingContext]:
|
||||
with parent_ctx as recording:
|
||||
self.recordings.append(recording)
|
||||
yield recording
|
||||
|
||||
return _wrapper()
|
||||
|
||||
def create_span(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Attributes] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
return self._delegate.create_span(name, attributes, timestamp, status)
|
||||
|
||||
|
||||
class DummyTracer:
|
||||
def __init__(self, start_span_instance: Optional[RecordingSpan] = None) -> None:
|
||||
self._start_span_instance = start_span_instance
|
||||
self.start_span_calls: List[Tuple[str, Dict[str, Any]]] = []
|
||||
self.start_as_current_span_calls: List[Tuple[str, Dict[str, Any], RecordingSpan]] = []
|
||||
class OtelSpanRecordingContext:
|
||||
def __init__(self, span: trace_api.Span) -> None:
|
||||
self._span = span
|
||||
|
||||
def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> RecordingSpan:
|
||||
span = self._start_span_instance or RecordingSpan()
|
||||
self.start_span_calls.append((name, dict(attributes or {})))
|
||||
return span
|
||||
def record_exception(self, exception: BaseException) -> None:
|
||||
self._span.record_exception(exception)
|
||||
self.record_status("ERROR", str(exception))
|
||||
|
||||
def start_as_current_span(
|
||||
def record_attributes(self, attributes: Dict[str, Any]) -> None:
|
||||
for key, value in attributes.items():
|
||||
self._span.set_attribute(key, value)
|
||||
|
||||
def record_status(self, status_code: str, description: Optional[str] = None) -> None:
|
||||
self._span.set_status(Status(StatusCode[status_code], description)) # type: ignore[index]
|
||||
|
||||
def get_recorded_span(self) -> None:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class OtelTracerAdapter:
|
||||
def __init__(self, tracer: trace_api.Tracer) -> None:
|
||||
self._tracer = tracer
|
||||
|
||||
def operation_context(
|
||||
self,
|
||||
name: str,
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> DummySpanContextManager:
|
||||
span = RecordingSpan()
|
||||
self.start_as_current_span_calls.append((name, dict(attributes or {}), span))
|
||||
return DummySpanContextManager(span)
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
):
|
||||
ctx = self._tracer.start_as_current_span(name, attributes=attributes)
|
||||
|
||||
class _ContextManager:
|
||||
def __enter__(self) -> OtelSpanRecordingContext:
|
||||
span = ctx.__enter__()
|
||||
return OtelSpanRecordingContext(span)
|
||||
|
||||
class DummyUseSpan:
|
||||
def __init__(self) -> None:
|
||||
self.calls: List[Tuple[RecordingSpan, bool]] = []
|
||||
self.exit_calls: List[
|
||||
Tuple[Optional[Type[BaseException]], Optional[BaseException], Optional[TracebackType]]
|
||||
] = []
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> bool:
|
||||
result = ctx.__exit__(exc_type, exc_val, exc_tb)
|
||||
return bool(result)
|
||||
|
||||
def __call__(self, span: RecordingSpan, end_on_exit: bool) -> DummyUseSpan:
|
||||
self.calls.append((span, end_on_exit))
|
||||
self._span = span
|
||||
return self
|
||||
return _ContextManager()
|
||||
|
||||
def __enter__(self) -> None:
|
||||
return None
|
||||
|
||||
def __exit__(
|
||||
def create_span(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> bool:
|
||||
self.exit_calls.append((exc_type, exc_val, exc_tb))
|
||||
return False
|
||||
name: str,
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
timestamp: Optional[float] = None,
|
||||
status: Optional[TraceStatus] = None,
|
||||
) -> SpanCoreFields:
|
||||
span = self._tracer.start_span(name, attributes=attributes)
|
||||
if status:
|
||||
span.set_status(Status(StatusCode[status.status_code], status.description)) # type: ignore[index]
|
||||
span.end()
|
||||
start = timestamp or time.time()
|
||||
return SpanCoreFields(
|
||||
name=name,
|
||||
attributes=attributes or {},
|
||||
start_time=start,
|
||||
end_time=start,
|
||||
status=status or TraceStatus(status_code="OK"),
|
||||
)
|
||||
|
||||
|
||||
def _install_recording_tracer(monkeypatch: pytest.MonkeyPatch) -> RecordingTracer:
|
||||
tracer = RecordingTracer()
|
||||
|
||||
def fake_get_active_tracer() -> RecordingTracer:
|
||||
return tracer
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_active_tracer", fake_get_active_tracer)
|
||||
return tracer
|
||||
|
||||
|
||||
def _resolve_attr(recording: DummySpanRecordingContext, key: str) -> Any:
|
||||
if key in recording.attributes:
|
||||
value = recording.attributes[key]
|
||||
else:
|
||||
value = filter_and_unflatten_attributes(recording.attributes, key)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -113,20 +160,8 @@ class ComplexResult:
|
||||
marker: str
|
||||
|
||||
|
||||
def test_safe_json_dump_handles_recursive_structures() -> None:
|
||||
payload: List[Any] = []
|
||||
payload.append(payload)
|
||||
|
||||
assert _safe_json_dump(payload) == "[[...]]"
|
||||
|
||||
|
||||
def test_operation_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
def test_operation_context_serializes_inputs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
ctx = OperationContext("custom-span", {"meta": {"foo": 1}, "count": 2})
|
||||
|
||||
@@ -134,63 +169,65 @@ def test_operation_context_records_inputs_and_outputs(monkeypatch: pytest.Monkey
|
||||
op.set_input({"payload": 1}, flag=True)
|
||||
op.set_output({"success": True})
|
||||
|
||||
assert tracer.start_span_calls
|
||||
start_name, start_attributes = tracer.start_span_calls[0]
|
||||
assert start_name == "custom-span"
|
||||
assert json.loads(start_attributes["meta"]) == {"foo": 1}
|
||||
assert start_attributes["count"] == 2
|
||||
|
||||
assert json.loads(span.attributes["input.args"]) == [{"payload": 1}]
|
||||
assert span.attributes["input.flag"] == "true"
|
||||
assert json.loads(span.attributes["output"]) == {"success": True}
|
||||
assert use_span.calls == [(span, True)]
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.name == "custom-span"
|
||||
assert recording.attributes["meta.foo"] == 1
|
||||
assert recording.attributes["count"] == 2
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert _resolve_attr(recording, f"{input_prefix}.args") == [{"payload": 1}]
|
||||
assert recording.attributes[f"{input_prefix}.flag"] is True
|
||||
assert _resolve_attr(recording, LightningSpanAttributes.OPERATION_OUTPUT.value) == {"success": True}
|
||||
|
||||
|
||||
def test_operation_context_set_input_supports_multiple_values(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
ctx = OperationContext("ctx", {})
|
||||
|
||||
with ctx as op:
|
||||
op.set_input(1, 2, data={"foo": ["bar"]}, flags=[True, False])
|
||||
|
||||
assert json.loads(span.attributes["input.args"]) == [1, 2]
|
||||
assert json.loads(span.attributes["input.data"]) == {"foo": ["bar"]}
|
||||
assert json.loads(span.attributes["input.flags"]) == [True, False]
|
||||
recording = tracer.recordings[-1]
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert _resolve_attr(recording, f"{input_prefix}.args") == [1, 2]
|
||||
assert _resolve_attr(recording, f"{input_prefix}.data") == {"foo": ["bar"]}
|
||||
assert _resolve_attr(recording, f"{input_prefix}.flags") == [True, False]
|
||||
|
||||
|
||||
def test_operation_context_records_non_serializable_output(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
class Unserializable:
|
||||
def __str__(self) -> str:
|
||||
return "<Unserializable>"
|
||||
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
def test_operation_context_set_input_expands_positional_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
ctx = OperationContext("ctx", {})
|
||||
|
||||
with ctx as op:
|
||||
op.set_output(Unserializable())
|
||||
op.set_input("alpha", "beta")
|
||||
|
||||
assert json.loads(span.attributes["output"]) == "<Unserializable>"
|
||||
recording = tracer.recordings[-1]
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert recording.attributes[f"{input_prefix}.args.0"] == "alpha"
|
||||
assert recording.attributes[f"{input_prefix}.args.1"] == "beta"
|
||||
|
||||
|
||||
def test_operation_context_serializes_non_serializable_output(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
class CustomObject:
|
||||
def __str__(self) -> str:
|
||||
return "custom-output"
|
||||
|
||||
ctx = OperationContext("ctx", {})
|
||||
|
||||
with ctx as op:
|
||||
op.set_output(CustomObject())
|
||||
|
||||
recording = tracer.recordings[-1]
|
||||
assert (
|
||||
json.loads(cast(str, recording.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value])) == "custom-output"
|
||||
)
|
||||
|
||||
|
||||
def test_operation_context_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
ctx = OperationContext("custom-span", {})
|
||||
|
||||
@@ -198,48 +235,47 @@ def test_operation_context_records_exceptions(monkeypatch: pytest.MonkeyPatch) -
|
||||
with ctx:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
assert isinstance(span.recorded_exceptions[0], RuntimeError)
|
||||
status = span.statuses[-1]
|
||||
assert status.status_code == StatusCode.ERROR
|
||||
assert status.description == "boom"
|
||||
assert use_span.exit_calls[-1][1].args == ("boom",) # type: ignore
|
||||
recording = tracer.recordings[-1]
|
||||
assert "exception.type" in recording.attributes
|
||||
assert recording.status.status_code == "ERROR"
|
||||
assert recording.status.description == "boom"
|
||||
|
||||
|
||||
def test_operation_factory_context_records_inputs_and_outputs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
with operation(tags=["one", "two"]) as ctx:
|
||||
ctx.set_input("alpha", meta={"score": 0.5})
|
||||
ctx.set_output(["beta", "gamma"])
|
||||
|
||||
start_name, attrs = tracer.start_span_calls[0]
|
||||
assert start_name == AGL_OPERATION
|
||||
assert json.loads(attrs["tags"]) == ["one", "two"]
|
||||
assert json.loads(span.attributes["input.args"]) == ["alpha"]
|
||||
assert json.loads(span.attributes["input.meta"]) == {"score": 0.5}
|
||||
assert json.loads(span.attributes["output"]) == ["beta", "gamma"]
|
||||
recording = tracer.recordings[-1]
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert recording.name == AGL_OPERATION
|
||||
assert recording.attributes["tags"] == ["one", "two"]
|
||||
assert _resolve_attr(recording, f"{input_prefix}.args") == ["alpha"]
|
||||
assert _resolve_attr(recording, f"{input_prefix}.meta") == {"score": 0.5}
|
||||
assert _resolve_attr(recording, LightningSpanAttributes.OPERATION_OUTPUT.value) == ["beta", "gamma"]
|
||||
|
||||
|
||||
def test_operation_factory_aliases_name_attribute(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
with operation(name="custom-operation") as ctx:
|
||||
ctx.set_output("done")
|
||||
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "custom-operation"
|
||||
|
||||
|
||||
def test_operation_factory_uses_standard_span_name(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
span = RecordingSpan()
|
||||
tracer = DummyTracer(start_span_instance=span)
|
||||
use_span = DummyUseSpan()
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
with operation(user={"id": 5}) as ctx:
|
||||
ctx.set_output("done")
|
||||
|
||||
assert tracer.start_span_calls
|
||||
start_name, attrs = tracer.start_span_calls[0]
|
||||
assert start_name == AGL_OPERATION
|
||||
assert json.loads(attrs["user"]) == {"id": 5}
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.name == AGL_OPERATION
|
||||
assert recording.attributes["user.id"] == 5
|
||||
|
||||
|
||||
def test_operation_rejects_custom_span_names() -> None:
|
||||
@@ -248,8 +284,7 @@ def test_operation_rejects_custom_span_names() -> None:
|
||||
|
||||
|
||||
def test_operation_decorator_sync_records_span_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = DummyTracer()
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
@operation(category={"kind": "combine"})
|
||||
def combine(data: Dict[str, int], *, meta: Dict[str, str]) -> Dict[str, Any]:
|
||||
@@ -258,21 +293,31 @@ def test_operation_decorator_sync_records_span_attributes(monkeypatch: pytest.Mo
|
||||
result = combine({"value": 1}, meta={"source": "unit"})
|
||||
|
||||
assert result == {"joined": {"value": 1, "source": "unit"}}
|
||||
assert tracer.start_as_current_span_calls
|
||||
span_name, span_attributes, span = tracer.start_as_current_span_calls[0]
|
||||
assert span_name == AGL_OPERATION
|
||||
assert json.loads(span_attributes["category"]) == {"kind": "combine"}
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.name == AGL_OPERATION
|
||||
assert recording.attributes["category.kind"] == "combine"
|
||||
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert json.loads(span.attributes[f"{input_prefix}.data"]) == {"value": 1}
|
||||
assert json.loads(span.attributes[f"{input_prefix}.meta"]) == {"source": "unit"}
|
||||
assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "combine"
|
||||
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
|
||||
assert _resolve_attr(recording, f"{input_prefix}.data") == {"value": 1}
|
||||
assert _resolve_attr(recording, f"{input_prefix}.meta") == {"source": "unit"}
|
||||
assert recording.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "combine"
|
||||
assert _resolve_attr(recording, LightningSpanAttributes.OPERATION_OUTPUT.value) == result
|
||||
|
||||
|
||||
def test_operation_decorator_aliases_operation_name_attribute(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
@operation(name="explicit-name")
|
||||
def compute(value: int) -> int:
|
||||
return value * 2
|
||||
|
||||
assert compute(3) == 6
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "explicit-name"
|
||||
|
||||
|
||||
def test_operation_decorator_handles_complex_signature(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = DummyTracer()
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
@operation()
|
||||
def complicated(
|
||||
@@ -289,23 +334,38 @@ def test_operation_decorator_handles_complex_signature(monkeypatch: pytest.Monke
|
||||
|
||||
result = complicated(1, "req", 7, 8, 9, kwonly="x", kwdefault="y", tag="value")
|
||||
|
||||
span = tracer.start_as_current_span_calls[0][2]
|
||||
recording = tracer.recordings[-1]
|
||||
input_prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
|
||||
assert json.loads(span.attributes[f"{input_prefix}.first"]) == 1
|
||||
assert json.loads(span.attributes[f"{input_prefix}.required"]) == "req"
|
||||
assert json.loads(span.attributes[f"{input_prefix}.default"]) == 7
|
||||
assert json.loads(span.attributes[f"{input_prefix}.extra"]) == [8, 9]
|
||||
assert json.loads(span.attributes[f"{input_prefix}.kwonly"]) == "x"
|
||||
assert json.loads(span.attributes[f"{input_prefix}.kwdefault"]) == "y"
|
||||
assert json.loads(span.attributes[f"{input_prefix}.rest"]) == {"tag": "value"}
|
||||
assert span.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "complicated"
|
||||
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == str(result)
|
||||
assert _resolve_attr(recording, f"{input_prefix}.first") == 1
|
||||
assert _resolve_attr(recording, f"{input_prefix}.required") == "req"
|
||||
assert _resolve_attr(recording, f"{input_prefix}.default") == 7
|
||||
assert _resolve_attr(recording, f"{input_prefix}.extra") == [8, 9]
|
||||
assert _resolve_attr(recording, f"{input_prefix}.kwonly") == "x"
|
||||
assert _resolve_attr(recording, f"{input_prefix}.kwdefault") == "y"
|
||||
assert _resolve_attr(recording, f"{input_prefix}.rest") == {"tag": "value"}
|
||||
assert recording.attributes[LightningSpanAttributes.OPERATION_NAME.value] == "complicated"
|
||||
assert _resolve_attr(recording, LightningSpanAttributes.OPERATION_OUTPUT.value) == (
|
||||
"ComplexResult(values=(1, 2, 1), marker='xyreq')"
|
||||
)
|
||||
assert isinstance(result, ComplexResult)
|
||||
assert recording.status.status_code == "OK"
|
||||
|
||||
|
||||
def test_operation_name_alias_does_not_override_explicit_attribute(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_install_recording_tracer(monkeypatch)
|
||||
|
||||
attrs = {
|
||||
LightningSpanAttributes.OPERATION_NAME.value: "explicit-name",
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="specify both"):
|
||||
with operation(name="alias-name", **attrs) as ctx:
|
||||
ctx.set_output("done")
|
||||
|
||||
|
||||
def test_operation_decorator_records_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = DummyTracer()
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
@operation()
|
||||
def fail(value: int) -> int:
|
||||
@@ -314,17 +374,14 @@ def test_operation_decorator_records_exceptions(monkeypatch: pytest.MonkeyPatch)
|
||||
with pytest.raises(ValueError):
|
||||
fail(1)
|
||||
|
||||
span = tracer.start_as_current_span_calls[0][2]
|
||||
assert isinstance(span.recorded_exceptions[0], ValueError)
|
||||
status = span.statuses[-1]
|
||||
assert status.status_code == StatusCode.ERROR
|
||||
assert status.description == "bad input"
|
||||
recording = tracer.recordings[-1]
|
||||
assert recording.status.status_code == "ERROR"
|
||||
assert recording.status.description == "bad input"
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_operation_async_wrapper_records_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = DummyTracer()
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
tracer = _install_recording_tracer(monkeypatch)
|
||||
|
||||
@operation()
|
||||
async def echo(payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -333,20 +390,19 @@ async def test_operation_async_wrapper_records_attributes(monkeypatch: pytest.Mo
|
||||
result = await echo({"value": 3})
|
||||
|
||||
assert result == {"payload": {"value": 3}}
|
||||
span = tracer.start_as_current_span_calls[0][2]
|
||||
recording = tracer.recordings[-1]
|
||||
prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
assert json.loads(span.attributes[f"{prefix}.payload"]) == {"value": 3}
|
||||
assert json.loads(span.attributes[LightningSpanAttributes.OPERATION_OUTPUT.value]) == result
|
||||
assert _resolve_attr(recording, f"{prefix}.payload") == {"value": 3}
|
||||
assert _resolve_attr(recording, LightningSpanAttributes.OPERATION_OUTPUT.value) == result
|
||||
|
||||
|
||||
def test_operation_span_can_be_resolved_via_annotation_links(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
provider = TracerProvider()
|
||||
exporter = InMemorySpanExporter()
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
tracer = provider.get_tracer(__name__)
|
||||
tracer = OtelTracerAdapter(provider.get_tracer(__name__))
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", lambda use_active_span_processor=True: tracer)
|
||||
monkeypatch.setattr(annotation_module, "get_active_tracer", lambda: tracer)
|
||||
|
||||
@operation(conversation_id="conv-1")
|
||||
def decorated(value: int) -> int:
|
||||
@@ -373,16 +429,10 @@ def test_operation_span_can_be_resolved_via_annotation_links(monkeypatch: pytest
|
||||
|
||||
|
||||
def test_operation_honors_propagate_flag(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
tracer = DummyTracer()
|
||||
flags: List[bool] = []
|
||||
use_span = DummyUseSpan()
|
||||
def fail_get_active_tracer() -> RecordingTracer:
|
||||
raise AssertionError("get_active_tracer should not be called when propagate=False")
|
||||
|
||||
def fake_get_tracer(use_active_span_processor: bool = True) -> DummyTracer:
|
||||
flags.append(use_active_span_processor)
|
||||
return tracer
|
||||
|
||||
monkeypatch.setattr(annotation_module, "get_tracer", fake_get_tracer)
|
||||
monkeypatch.setattr(annotation_module.trace, "use_span", use_span)
|
||||
monkeypatch.setattr(annotation_module, "get_active_tracer", fail_get_active_tracer)
|
||||
|
||||
@operation(propagate=False)
|
||||
def decorated(value: int) -> int:
|
||||
@@ -390,7 +440,15 @@ def test_operation_honors_propagate_flag(monkeypatch: pytest.MonkeyPatch) -> Non
|
||||
|
||||
assert decorated(7) == 7
|
||||
|
||||
with operation(propagate=False):
|
||||
pass
|
||||
with operation(propagate=False, value=7) as op:
|
||||
with pytest.raises(RuntimeError):
|
||||
op.span()
|
||||
|
||||
assert flags == [False, False]
|
||||
assert op.span() is not None
|
||||
assert op.span().name == AGL_OPERATION
|
||||
assert op.span().attributes == {"value": 7}
|
||||
assert op.span().status.status_code == "OK"
|
||||
assert op.span().status.description is None
|
||||
assert op.span().start_time is not None
|
||||
assert op.span().end_time is not None
|
||||
assert op.span().start_time < op.span().end_time # type: ignore
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Sequence, Tuple, cast
|
||||
|
||||
@@ -22,7 +23,7 @@ from agentlightning.semconv import AGL_ANNOTATION
|
||||
from agentlightning.store.base import UNSET, LightningStore, Unset
|
||||
from agentlightning.store.memory import InMemoryLightningStore
|
||||
from agentlightning.tracer.base import Tracer
|
||||
from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, Worker
|
||||
from agentlightning.types import LLM, Hook, NamedResources, PromptTemplate, Rollout, Span, SpanCoreFields, Worker
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
@@ -73,8 +74,9 @@ def create_agent_span(
|
||||
class DummyTracer(Tracer):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._last_trace: List[ReadableSpan] = []
|
||||
self._last_trace: List[Span] = []
|
||||
self._contexts: List[Dict[str, Any]] = []
|
||||
self._sequence_id = 0
|
||||
|
||||
def init(self, *args: Any, **kwargs: Any) -> None:
|
||||
self._last_trace.clear()
|
||||
@@ -82,7 +84,7 @@ class DummyTracer(Tracer):
|
||||
def teardown(self, *args: Any, **kwargs: Any) -> None:
|
||||
self._last_trace.clear()
|
||||
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
def get_last_trace(self) -> List[Span]:
|
||||
return list(self._last_trace)
|
||||
|
||||
@asynccontextmanager
|
||||
@@ -93,7 +95,7 @@ class DummyTracer(Tracer):
|
||||
store: Optional[LightningStore] = None,
|
||||
rollout_id: Optional[str] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
) -> AsyncGenerator[List[ReadableSpan], None]:
|
||||
) -> AsyncGenerator[List[Span], None]:
|
||||
previous = self._contexts[-1] if self._contexts else None
|
||||
current = {
|
||||
"name": name,
|
||||
@@ -112,7 +114,15 @@ class DummyTracer(Tracer):
|
||||
|
||||
def record_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> ReadableSpan:
|
||||
span = create_readable_span(name, attributes)
|
||||
self._last_trace.append(span)
|
||||
rollout_id = "rollout-dummy"
|
||||
attempt_id = "attempt-dummy"
|
||||
sequence_id = self._sequence_id
|
||||
self._sequence_id += 1
|
||||
if self._contexts:
|
||||
current = self._contexts[-1]
|
||||
rollout_id = current["rollout_id"]
|
||||
attempt_id = current["attempt_id"]
|
||||
self._last_trace.append(Span.from_opentelemetry(span, rollout_id, attempt_id, sequence_id))
|
||||
return span
|
||||
|
||||
|
||||
@@ -143,7 +153,8 @@ class HeartbeatAgent(LitAgent[Dict[str, Any]]):
|
||||
async def setup_heartbeat_runner(
|
||||
*,
|
||||
heartbeat_interval: float = 0.05,
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
|
||||
heartbeat_include_gpu: bool = False,
|
||||
) -> tuple[LitAgentRunner[Any], RecordingStore]:
|
||||
"""Create a runner wired to a RecordingStore for heartbeat tests."""
|
||||
|
||||
@@ -152,6 +163,7 @@ async def setup_heartbeat_runner(
|
||||
tracer=DummyTracer(),
|
||||
heartbeat_interval=heartbeat_interval,
|
||||
heartbeat_launch_mode=heartbeat_launch_mode,
|
||||
heartbeat_include_gpu=heartbeat_include_gpu,
|
||||
)
|
||||
agent = HeartbeatAgent()
|
||||
runner.init(agent)
|
||||
@@ -309,7 +321,9 @@ async def test_step_raises_for_invalid_result_type() -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_readable_spans_return_skip_store_when_tracer_is_otel(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_readable_spans_return_skip_store_when_tracer_is_otel(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
class DummyTracerWithOtel(DummyTracer):
|
||||
pass
|
||||
|
||||
@@ -342,12 +356,55 @@ async def test_readable_spans_return_skip_store_when_tracer_is_otel(monkeypatch:
|
||||
attempted_rollout, spans
|
||||
)
|
||||
|
||||
assert result_spans == spans
|
||||
assert store.add_otel_span_calls == 0
|
||||
assert all(isinstance(span, Span) for span in result_spans)
|
||||
# Warned, but still logged
|
||||
assert [span.name for span in result_spans] == ["otel-span"]
|
||||
assert store.add_otel_span_calls == 1
|
||||
assert "Tracer is already an OpenTelemetry tracer" in caplog.text
|
||||
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_process_readable_spans_adds_each_span() -> None:
|
||||
agent = HeartbeatAgent()
|
||||
runner, store, _ = await setup_runner(agent)
|
||||
attempted_rollout = await store.start_rollout(input={"task": "post-process"}, mode="val")
|
||||
|
||||
spans = [create_readable_span("case-2-span-a"), create_readable_span("case-2-span-b")]
|
||||
|
||||
try:
|
||||
result_spans = await runner._post_process_rollout_result( # pyright: ignore[reportPrivateUsage]
|
||||
attempted_rollout, spans
|
||||
)
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
assert [span.name for span in result_spans] == ["case-2-span-a", "case-2-span-b"]
|
||||
stored_spans = await store.query_spans(attempted_rollout.rollout_id, attempted_rollout.attempt.attempt_id)
|
||||
assert [span.name for span in stored_spans] == ["case-2-span-a", "case-2-span-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_process_span_core_fields_create_spans() -> None:
|
||||
agent = HeartbeatAgent()
|
||||
runner, store, _ = await setup_runner(agent)
|
||||
attempted_rollout = await store.start_rollout(input={"task": "reward-list"}, mode="val")
|
||||
|
||||
span_core_fields = [emit_reward(0.5, propagate=False), emit_reward(-0.2, propagate=False)]
|
||||
|
||||
try:
|
||||
result_spans = await runner._post_process_rollout_result( # pyright: ignore[reportPrivateUsage]
|
||||
attempted_rollout, span_core_fields
|
||||
)
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
assert all(span.name == AGL_ANNOTATION for span in result_spans)
|
||||
stored_spans = await store.query_spans(attempted_rollout.rollout_id, attempted_rollout.attempt.attempt_id)
|
||||
assert [span.attributes.get("agentlightning.reward.0.value") for span in stored_spans] == [0.5, -0.2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_step_handles_non_llm_resource() -> None:
|
||||
class PromptAgent(LitAgent[str]):
|
||||
@@ -585,8 +642,8 @@ async def test_agent_emits_multiple_rewards() -> None:
|
||||
class RewardListAgent(LitAgent[Dict[str, Any]]):
|
||||
def validation_rollout(
|
||||
self, task: Dict[str, Any], resources: Dict[str, Any], rollout: Any
|
||||
) -> List[ReadableSpan]:
|
||||
return [emit_reward(0.2), emit_reward(0.6)]
|
||||
) -> List[SpanCoreFields]:
|
||||
return [emit_reward(0.2, propagate=False), emit_reward(0.6, propagate=False)]
|
||||
|
||||
agent = RewardListAgent()
|
||||
runner, store, _ = await setup_runner(agent)
|
||||
@@ -859,7 +916,7 @@ async def test_iter_passes_worker_id_to_dequeue(monkeypatch: pytest.MonkeyPatch)
|
||||
@pytest.mark.asyncio
|
||||
async def test_emit_heartbeat_updates_worker_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = {"cpu_pct": 42.0, "mem_pct": 10.5}
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda: snapshot)
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.1)
|
||||
worker_label = runner.get_worker_id()
|
||||
@@ -875,10 +932,87 @@ async def test_emit_heartbeat_updates_worker_snapshot(monkeypatch: pytest.Monkey
|
||||
assert worker.last_heartbeat_time is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emit_heartbeat_passes_include_gpu(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = {"cpu_pct": 11.1}
|
||||
requested_flags: List[bool] = []
|
||||
|
||||
async def immediate_to_thread(func: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.asyncio.to_thread", immediate_to_thread)
|
||||
|
||||
def fake_system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
|
||||
requested_flags.append(include_gpu)
|
||||
return snapshot
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", fake_system_snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05, heartbeat_include_gpu=True)
|
||||
worker_label = runner.get_worker_id()
|
||||
try:
|
||||
await runner._emit_heartbeat(store) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
assert requested_flags == [True]
|
||||
assert store.worker_updates == [(worker_label, snapshot)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emit_heartbeat_skips_when_snapshot_times_out(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = {"cpu_pct": 9.9}
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05)
|
||||
interval = runner._heartbeat_interval # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
async def slow_to_thread(func: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
await asyncio.sleep(interval + 0.05)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.asyncio.to_thread", slow_to_thread)
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
try:
|
||||
await runner._emit_heartbeat(store) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
assert store.worker_updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emit_heartbeat_skips_when_store_update_times_out(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = {"cpu_pct": 8.8}
|
||||
|
||||
async def immediate_to_thread(func: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.asyncio.to_thread", immediate_to_thread)
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05)
|
||||
interval = runner._heartbeat_interval # pyright: ignore[reportPrivateUsage]
|
||||
original_update_worker = store.update_worker
|
||||
|
||||
async def slow_update_worker(*args: Any, **kwargs: Any) -> Worker:
|
||||
await asyncio.sleep(interval + 0.05)
|
||||
return await original_update_worker(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(store, "update_worker", slow_update_worker)
|
||||
|
||||
try:
|
||||
await runner._emit_heartbeat(store) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
assert store.worker_updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_loop_runs_until_stopped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Test that heartbeat loop runs with the default launch mode until stopped."""
|
||||
snapshot = {"timestamp": 1234567890}
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda: snapshot)
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05)
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
@@ -897,3 +1031,300 @@ async def test_heartbeat_loop_runs_until_stopped(monkeypatch: pytest.MonkeyPatch
|
||||
assert len(store.worker_updates) == update_count
|
||||
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asyncio_heartbeat_loop_runs_until_stopped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Test that asyncio heartbeat loop runs until stopped (explicit asyncio mode test)."""
|
||||
snapshot = {"timestamp": 1234567890}
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.05, heartbeat_launch_mode="asyncio")
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
await asyncio.sleep(0.12)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
|
||||
update_count = len(store.worker_updates)
|
||||
assert update_count >= 1
|
||||
assert all(stats == snapshot for _, stats in store.worker_updates if stats is not None)
|
||||
|
||||
await asyncio.sleep(0.06)
|
||||
assert len(store.worker_updates) == update_count
|
||||
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_heartbeat_loop_runs_until_stopped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = {"timestamp": time.time()}
|
||||
include_flags: List[bool] = []
|
||||
|
||||
def fake_system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
|
||||
include_flags.append(include_gpu)
|
||||
return snapshot
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", fake_system_snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.05,
|
||||
heartbeat_launch_mode="thread",
|
||||
heartbeat_include_gpu=True,
|
||||
)
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
await asyncio.sleep(0.3)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
assert len(store.worker_updates) >= 1
|
||||
assert all(stats == snapshot for _, stats in store.worker_updates if stats is not None)
|
||||
assert include_flags and all(include_flags)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_heartbeat_handles_producer_exception(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Test that thread mode producer handles exceptions and continues running."""
|
||||
call_count = 0
|
||||
|
||||
def fake_system_snapshot(_include_gpu: bool = False) -> Dict[str, Any]:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise RuntimeError("Simulated snapshot failure")
|
||||
return {"cpu_pct": 25.0}
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", fake_system_snapshot)
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.02,
|
||||
heartbeat_launch_mode="thread",
|
||||
)
|
||||
runner._interval_jitter = 0.01 # pyright: ignore[reportPrivateUsage]
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
# Wait long enough for at least 2 cycles (2 * (interval + jitter) + buffer)
|
||||
await asyncio.sleep(0.15)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
# Verify the producer logged the exception but continued
|
||||
assert "system_snapshot failed" in caplog.text
|
||||
assert call_count >= 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_heartbeat_handles_consumer_exception(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Test that thread mode consumer handles store update exceptions and continues."""
|
||||
snapshot = {"cpu_pct": 33.0}
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", lambda include_gpu=False: snapshot)
|
||||
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.02,
|
||||
heartbeat_launch_mode="thread",
|
||||
)
|
||||
runner._interval_jitter = 0.01 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
update_count = 0
|
||||
original_update_worker = store.update_worker
|
||||
|
||||
async def failing_update_worker(*args: Any, **kwargs: Any) -> Worker:
|
||||
nonlocal update_count
|
||||
update_count += 1
|
||||
if update_count == 1:
|
||||
raise RuntimeError("Simulated store failure")
|
||||
return await original_update_worker(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(store, "update_worker", failing_update_worker)
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
# Wait long enough for at least 2 cycles
|
||||
await asyncio.sleep(0.15)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
# Verify the consumer logged the exception but continued
|
||||
assert "update failed" in caplog.text
|
||||
assert update_count >= 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_heartbeat_waits_for_first_snapshot(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Test that thread mode consumer skips update when no snapshot is available yet."""
|
||||
snapshot_ready = False
|
||||
|
||||
def fake_system_snapshot(_include_gpu: bool = False) -> Dict[str, Any]:
|
||||
nonlocal snapshot_ready
|
||||
if not snapshot_ready:
|
||||
time.sleep(0.1) # Simulate slow first snapshot
|
||||
snapshot_ready = True
|
||||
return {"cpu_pct": 42.0}
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", fake_system_snapshot)
|
||||
|
||||
caplog.set_level(logging.DEBUG)
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.05,
|
||||
heartbeat_launch_mode="thread",
|
||||
)
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
await asyncio.sleep(0.25)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
# Verify the consumer logged that no snapshot was available initially
|
||||
assert "no snapshot yet" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_heartbeat_logs_stale_snapshot_only_once(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Test that stale snapshot warning is only logged once per stale snapshot."""
|
||||
call_count = 0
|
||||
|
||||
def fake_system_snapshot(_include_gpu: bool = False) -> Dict[str, Any]:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
# Only create one snapshot, then hang
|
||||
if call_count == 1:
|
||||
return {"cpu_pct": 50.0}
|
||||
# Simulate hung producer
|
||||
time.sleep(10)
|
||||
return {"cpu_pct": 99.0}
|
||||
|
||||
monkeypatch.setattr("agentlightning.runner.agent.system_snapshot", fake_system_snapshot)
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.01,
|
||||
heartbeat_launch_mode="thread",
|
||||
)
|
||||
runner._interval_jitter = 0.01 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
# Wait long enough for the single snapshot to become stale
|
||||
# and for multiple consumer iterations to check it
|
||||
# stale_after = 0.01 + 0.01 + 1.0 = 1.02s
|
||||
await asyncio.sleep(1.3)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
# Count how many times the stale warning appears
|
||||
stale_warnings = caplog.text.count("snapshot stale")
|
||||
# Should only warn once, even though consumer checked multiple times
|
||||
assert stale_warnings == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_disabled_when_interval_zero() -> None:
|
||||
"""Test that heartbeat loop is not started when interval is 0 or negative."""
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.0)
|
||||
try:
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is None
|
||||
finally:
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_disabled_when_no_worker_id(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""Test that heartbeat loop returns None when worker_id is not set."""
|
||||
store = RecordingStore()
|
||||
runner = LitAgentRunner[Any](
|
||||
tracer=DummyTracer(),
|
||||
heartbeat_interval=0.05,
|
||||
)
|
||||
agent = HeartbeatAgent()
|
||||
runner.init(agent)
|
||||
# Note: NOT calling init_worker, so worker_id remains None
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is None
|
||||
assert "Cannot start heartbeat loop without worker_id" in caplog.text
|
||||
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emit_heartbeat_propagates_cancelled_error() -> None:
|
||||
"""Test that CancelledError is properly propagated in _emit_heartbeat."""
|
||||
runner, store = await setup_heartbeat_runner(heartbeat_interval=0.1)
|
||||
|
||||
async def cancelling_to_thread(_func: Any, *_args: Any, **_kwargs: Any) -> Any:
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
original_to_thread = asyncio.to_thread
|
||||
try:
|
||||
asyncio.to_thread = cancelling_to_thread # type: ignore
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await runner._emit_heartbeat(store) # pyright: ignore[reportPrivateUsage]
|
||||
finally:
|
||||
asyncio.to_thread = original_to_thread # type: ignore
|
||||
teardown_runner(runner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asyncio_heartbeat_continues_after_exception(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""Test that asyncio heartbeat loop continues running after exceptions."""
|
||||
call_count = 0
|
||||
|
||||
async def failing_emit_heartbeat(_self: Any, _store: Any) -> None:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise RuntimeError("Simulated heartbeat failure")
|
||||
|
||||
monkeypatch.setattr(LitAgentRunner, "_emit_heartbeat", failing_emit_heartbeat)
|
||||
|
||||
caplog.set_level(logging.ERROR)
|
||||
runner, store = await setup_heartbeat_runner(
|
||||
heartbeat_interval=0.02,
|
||||
heartbeat_launch_mode="asyncio",
|
||||
)
|
||||
runner._interval_jitter = 0.01 # pyright: ignore[reportPrivateUsage]
|
||||
stop_heartbeat = runner._start_heartbeat_loop(store) # pyright: ignore[reportPrivateUsage]
|
||||
assert stop_heartbeat is not None
|
||||
|
||||
try:
|
||||
# Wait long enough for at least 2 cycles (2 * (interval + jitter) + buffer)
|
||||
await asyncio.sleep(0.15)
|
||||
finally:
|
||||
await stop_heartbeat()
|
||||
teardown_runner(runner)
|
||||
|
||||
# Verify the loop logged the exception but continued
|
||||
assert "Heartbeat failed" in caplog.text
|
||||
assert call_count >= 2
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user