chore: import upstream snapshot with attribution
Deploy Documentation / deploy (push) Has been cancelled
CPU Test / Test (Utilities, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, latest, Python 3.13) (push) Has been cancelled
Dashboard / Chromatic (push) Has been cancelled
CPU Test / Lint - fast (push) Has been cancelled
CPU Test / Lint - next (push) Has been cancelled
CPU Test / Lint - slow (push) Has been cancelled
CPU Test / Lint - JavaScript (push) Has been cancelled
CPU Test / Build documentation (push) Has been cancelled
CPU Test / Test (AgentOps, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Others, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Store, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Weave, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (AgentOps, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (LLM proxy, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Others, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Store, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Utilities, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (JavaScript) (push) Has been cancelled
Deploy Documentation / deploy (push) Has been cancelled
CPU Test / Test (Utilities, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (LLM proxy, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Others, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, latest, Python 3.13) (push) Has been cancelled
Dashboard / Chromatic (push) Has been cancelled
CPU Test / Lint - fast (push) Has been cancelled
CPU Test / Lint - next (push) Has been cancelled
CPU Test / Lint - slow (push) Has been cancelled
CPU Test / Lint - JavaScript (push) Has been cancelled
CPU Test / Build documentation (push) Has been cancelled
CPU Test / Test (AgentOps, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (LLM proxy, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Others, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Store, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (Weave, legacy, Python 3.10) (push) Has been cancelled
CPU Test / Test (AgentOps, stable, Python 3.11) (push) Has been cancelled
CPU Test / Test (Store, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Utilities, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (Weave, stable, Python 3.12) (push) Has been cancelled
CPU Test / Test (AgentOps, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (LLM proxy, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Others, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Store, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (Utilities, latest, Python 3.13) (push) Has been cancelled
CPU Test / Test (JavaScript) (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
.venv
|
||||
**/.venv
|
||||
__pycache__
|
||||
.git
|
||||
.gitignore
|
||||
**/node_modules
|
||||
dist
|
||||
build
|
||||
.env
|
||||
docker
|
||||
.pytest_cache
|
||||
.vscode
|
||||
**/*.log
|
||||
examples/**/data
|
||||
@@ -0,0 +1,32 @@
|
||||
name: Backport Merged Pull Request
|
||||
on:
|
||||
pull_request_target:
|
||||
types: [closed]
|
||||
permissions:
|
||||
contents: write
|
||||
issues: write
|
||||
pull-requests: write
|
||||
|
||||
# NOTE:
|
||||
# Microsoft requires rotating BOT_PAT every 3 months.
|
||||
# Log onto agent-lightning-bot account and rotate the PAT if needed.
|
||||
|
||||
jobs:
|
||||
backport:
|
||||
name: Backport pull request
|
||||
runs-on: ubuntu-latest
|
||||
# Don't run on closed unmerged pull requests
|
||||
if: github.event.pull_request.merged
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Create backport pull requests
|
||||
uses: korthout/backport-action@v3
|
||||
with:
|
||||
branch_name: 'backport/${pull_number}/${target_branch}'
|
||||
label_pattern: ^(stable/[^ ]+)$
|
||||
github_token: ${{ secrets.BOT_PAT }}
|
||||
add_labels: backport
|
||||
add_author_as_assignee: true
|
||||
git_committer_name: agent-lightning-bot
|
||||
# This email address is not monitored.
|
||||
git_committer_email: agl.msft@outlook.com
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - APO
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - APO
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-apo.yml', label: 'apo', variants: ['legacy', 'stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Azure
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Azure
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-azure.yml', label: 'azure', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Calc-X
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Calc-X
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-calc-x.yml', label: 'calc-x', variants: ['legacy', 'stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - ChartQA
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - ChartQA
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-chartqa.yml', label: 'chartqa', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Claude Code
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Claude Code
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-claude-code.yml', label: 'claude-code', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Compatibility
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Backward Compatibility
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-compat.yml', label: 'examples-compat', variants: ['legacy', 'stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,45 @@
|
||||
name: Badge - Examples
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Calc-X
|
||||
- Examples - Spider
|
||||
- Examples - APO
|
||||
- Examples - Unsloth
|
||||
- Examples - Tinker
|
||||
- Examples - Azure
|
||||
- Examples - Claude Code
|
||||
- Examples - RAG
|
||||
- Examples - ChartQA
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-calc-x.yml', label: 'examples-calc-x.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-spider.yml', label: 'examples-spider.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-apo.yml', label: 'examples-apo.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-tinker.yml', label: 'examples-tinker.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-azure.yml', label: 'examples-azure.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-claude-code.yml', label: 'examples-claude-code.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-rag.yml', label: 'examples-rag.stable', variants: ['stable'] },
|
||||
{ workflow: 'examples-chartqa.yml', label: 'examples-chartqa.stable', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,41 @@
|
||||
name: Badge - Latest
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Calc-X
|
||||
- Examples - Spider
|
||||
- Examples - APO
|
||||
- Examples - Unsloth
|
||||
- Examples - RAG
|
||||
- Examples - Claude Code
|
||||
- GPU Test
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-calc-x.yml', label: 'calc-x.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-spider.yml', label: 'spider.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-apo.yml', label: 'apo.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-unsloth.yml', label: 'unsloth.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-claude-code.yml', label: 'claude-code.latest', variants: ['latest'] },
|
||||
{ workflow: 'examples-rag.yml', label: 'rag.latest', variants: ['latest'] },
|
||||
{ workflow: 'tests-full.yml', label: 'tests-full.latest', variants: ['latest'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - RAG
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - RAG
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-rag.yml', label: 'rag', variants: ['legacy', 'stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Spider
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Spider
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-spider.yml', label: 'spider', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Tinker
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Tinker
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-tinker.yml', label: 'tinker', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,31 @@
|
||||
name: Badge - Unit Test
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- CPU Test
|
||||
- GPU Test
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'tests-full.yml', label: 'tests-full', variants: ['legacy', 'stable'] },
|
||||
{ workflow: 'tests.yml', label: 'tests', variants: ['legacy', 'stable', 'Lint', 'documentation', 'JavaScript'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Badge - Unsloth
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows:
|
||||
- Examples - Unsloth
|
||||
types: [completed]
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: read
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
badge:
|
||||
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/github-script@v8
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
script: |
|
||||
const badgeAggregation = require('./scripts/badge_aggregation.js');
|
||||
const dependencies = [
|
||||
{ workflow: 'examples-unsloth.yml', label: 'examples-unsloth.stable', variants: ['stable'] },
|
||||
];
|
||||
await badgeAggregation({ github, context, core, dependencies });
|
||||
@@ -0,0 +1,441 @@
|
||||
name: Benchmark
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
workflow_dispatch:
|
||||
schedule:
|
||||
# Every Monday and Thursday at 3 AM UTC+8
|
||||
- cron: '0 19 * * 0,3'
|
||||
|
||||
jobs:
|
||||
benchmark:
|
||||
name: ${{ matrix.workload.kind }} (${{ matrix.backend.id }}, ${{ matrix.workload.display }})
|
||||
runs-on: ${{ matrix.workload.runner }}
|
||||
timeout-minutes: ${{ matrix.workload.timeout }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
backend:
|
||||
- id: memory
|
||||
compose_file: compose.prometheus-memory-store.yml
|
||||
- id: mongo
|
||||
compose_file: compose.prometheus-mongo-store.yml
|
||||
workload:
|
||||
- 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: 48
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu
|
||||
timeout: 120
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
--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: 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: 96
|
||||
runner:
|
||||
- self-hosted
|
||||
- 1ES.Pool=agl-runner-cpu-high
|
||||
timeout: 60
|
||||
args: >-
|
||||
--mode batch_partial
|
||||
--total-tasks 10000
|
||||
--batch-size 1024
|
||||
--remaining-tasks 256
|
||||
--n-runners 512
|
||||
--max-rounds 20
|
||||
--sleep-seconds 1.0
|
||||
|
||||
- id: micro-worker
|
||||
display: Update worker
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: worker
|
||||
- id: micro-dequeue-empty
|
||||
display: Dequeue empty
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-empty
|
||||
- id: micro-rollout
|
||||
display: Rollout + span
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: rollout
|
||||
- id: micro-dequeue-update-attempt
|
||||
display: Dequeue + update attempt
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-update-attempt
|
||||
- id: micro-dequeue-only
|
||||
display: Dequeue only
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 30
|
||||
cli: dequeue-only
|
||||
- id: micro-metrics
|
||||
display: Multi-metric fan-out
|
||||
kind: micro
|
||||
store_workers: 8
|
||||
runner: ubuntu-latest
|
||||
timeout: 15
|
||||
cli: metrics
|
||||
env:
|
||||
PYTHONUNBUFFERED: "1"
|
||||
STORE_URL: http://localhost:4747
|
||||
STORE_API_URL: http://localhost:4747/v1/agl
|
||||
PROM_URL: http://localhost:9090
|
||||
GITHUB_ACTIONS_TIMEOUT_MINUTES: ${{ matrix.workload.timeout }}
|
||||
WORKLOAD_KIND: ${{ matrix.workload.kind }}
|
||||
WORKLOAD_ID: ${{ matrix.workload.id }}
|
||||
BACKEND_ID: ${{ matrix.backend.id }}
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
COMPOSE_FILE: ${{ matrix.backend.compose_file }}
|
||||
AGL_STORE_N_WORKERS: ${{ matrix.workload.store_workers }}
|
||||
ANALYSIS_FILE: ${{ format('analysis-{0}.log', matrix.workload.id) }}
|
||||
SUMMARY_FILE: ${{ format('summary-{0}.log', matrix.workload.id) }}
|
||||
PROM_ARCHIVE_BASENAME: ${{ format('prometheus-{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
ARTIFACT_NAME: ${{ format('{0}-{1}', matrix.workload.id, matrix.backend.id) }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- 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: Check disk space
|
||||
run: df -h
|
||||
|
||||
- name: Reset benchmark data directories
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
rm -rf data
|
||||
bash setup.sh
|
||||
|
||||
- name: Launch ${{ matrix.backend.id }} Prometheus stack
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f "$COMPOSE_FILE" down -v || true
|
||||
docker compose -f "$COMPOSE_FILE" up -d --quiet-pull
|
||||
|
||||
- name: Wait for store readiness
|
||||
run: |
|
||||
set -euo pipefail
|
||||
for attempt in {1..60}; do
|
||||
if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
|
||||
sleep 1
|
||||
curl -fsS "$STORE_API_URL/rollouts" # Warm up the scraper
|
||||
sleep 15 # Allow some time for the baseline metrics to be established
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "Store did not become ready in time" >&2
|
||||
# show logs for debugging
|
||||
cd docker && docker compose -f "$COMPOSE_FILE" logs app
|
||||
exit 1
|
||||
|
||||
- name: Prepare artifact directory
|
||||
run: mkdir -p "$ARTIFACT_DIR"
|
||||
|
||||
- name: Record workload start
|
||||
run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: (Scenario) Run ${{ matrix.workload.display }} workload
|
||||
if: ${{ matrix.workload.kind == 'scenario' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
uv run --locked --no-sync python -m tests.benchmark.benchmark_store \
|
||||
--store-url "$STORE_URL" \
|
||||
${{ matrix.workload.args }}
|
||||
|
||||
- name: (Micro) Run ${{ matrix.workload.display }}
|
||||
if: ${{ matrix.workload.kind == 'micro' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
uv run --locked --no-sync python -m tests.benchmark.micro_benchmark \
|
||||
--store-url "$STORE_URL" \
|
||||
--summary-file "$ARTIFACT_DIR/$SUMMARY_FILE" \
|
||||
"${{ matrix.workload.cli }}" | tee "$ARTIFACT_DIR/${{ matrix.workload.id }}.txt"
|
||||
|
||||
- name: Record workload end
|
||||
if: ${{ always() }}
|
||||
run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Show micro benchmark summary
|
||||
if: ${{ always() && matrix.workload.kind == 'micro' }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
summary_file="$ARTIFACT_DIR/$SUMMARY_FILE"
|
||||
if [ -f "$summary_file" ]; then
|
||||
echo "Micro benchmark summary ($WORKLOAD_ID/$BACKEND_ID):"
|
||||
cat "$summary_file"
|
||||
else
|
||||
echo "Summary file not found: $summary_file"
|
||||
fi
|
||||
|
||||
- name: Run workload analysis
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
|
||||
echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/$ANALYSIS_FILE"
|
||||
exit 1
|
||||
fi
|
||||
uv run --locked --no-sync python -m tests.benchmark.analysis \
|
||||
--prom-url "$PROM_URL" \
|
||||
--store-url "$STORE_API_URL" \
|
||||
--start "$BENCHMARK_START" \
|
||||
--end "$BENCHMARK_END" \
|
||||
| tee "$ARTIFACT_DIR/$ANALYSIS_FILE"
|
||||
|
||||
- name: Collect docker logs
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
cd docker
|
||||
readarray -t services < <(docker compose -f "$COMPOSE_FILE" config --services)
|
||||
if [ "${#services[@]}" -eq 0 ]; then
|
||||
echo "No services defined in compose file."
|
||||
exit 0
|
||||
fi
|
||||
for service in "${services[@]}"; do
|
||||
docker compose -f "$COMPOSE_FILE" logs "$service" > "../$ARTIFACT_DIR/docker-${service}-${WORKLOAD_ID}-${BACKEND_ID}.log" || true
|
||||
done
|
||||
|
||||
- name: Stop ${{ matrix.backend.id }} Prometheus stack
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f "$COMPOSE_FILE" down -v || true
|
||||
|
||||
- name: Archive Prometheus metrics
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
if [ -d docker/data/prometheus ]; then
|
||||
tar -C docker/data -czf "$ARTIFACT_DIR/${PROM_ARCHIVE_BASENAME}.tar.gz" prometheus
|
||||
fi
|
||||
|
||||
- name: Upload workload artifacts
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: ${{ env.ARTIFACT_NAME }}
|
||||
path: ${{ env.ARTIFACT_DIR }}
|
||||
if-no-files-found: error
|
||||
|
||||
collection-benchmarks:
|
||||
name: collection (${{ matrix.backend.id }}, ${{ matrix.workload.id }})
|
||||
runs-on: ${{ matrix.backend.runner }}
|
||||
timeout-minutes: 15
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
backend:
|
||||
- id: memory
|
||||
needs_mongo: false
|
||||
runner: ubuntu-latest
|
||||
- id: mongo
|
||||
needs_mongo: true
|
||||
runner: ubuntu-latest
|
||||
workload:
|
||||
- id: high-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 2048
|
||||
type: insert
|
||||
- id: medium-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 128
|
||||
type: insert
|
||||
- id: low-insert
|
||||
total_tasks: 50000
|
||||
concurrency: 4
|
||||
type: insert
|
||||
- id: high-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 2048
|
||||
type: dequeue
|
||||
- id: medium-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 128
|
||||
type: dequeue
|
||||
- id: low-dequeue
|
||||
total_tasks: 50000
|
||||
concurrency: 4
|
||||
type: dequeue
|
||||
env:
|
||||
ARTIFACT_DIR: ${{ format('artifacts/{0}-{1}', matrix.backend.id, matrix.workload.id) }}
|
||||
SUMMARY_FILE: ${{ format('artifacts/{0}-{1}/summary-{0}-{1}.jsonl', matrix.backend.id, matrix.workload.id) }}
|
||||
ARTIFACT_NAME: ${{ format('collections-{0}-{1}', matrix.backend.id, matrix.workload.id) }}
|
||||
MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --extra mongo --group core-stable --group dev
|
||||
|
||||
- name: Launch MongoDB
|
||||
if: ${{ matrix.backend.needs_mongo }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f compose.mongo.yml down -v || true
|
||||
docker compose -f compose.mongo.yml up -d --quiet-pull
|
||||
for attempt in {1..60}; do
|
||||
if docker compose -f compose.mongo.yml exec -T mongo mongosh --quiet --eval 'db.runCommand({ping:1})' >/dev/null 2>&1; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "MongoDB did not become ready in time" >&2
|
||||
docker compose -f compose.mongo.yml logs mongo
|
||||
exit 1
|
||||
|
||||
- name: Run collection benchmark
|
||||
run: |
|
||||
set -euo pipefail
|
||||
mkdir -p "$ARTIFACT_DIR"
|
||||
echo "Running collection benchmark (backend=${{ matrix.backend.id }}, workload=${{ matrix.workload.id }})"
|
||||
uv run --locked --no-sync python -m tests.benchmark.collection_benchmark \
|
||||
"${{ matrix.workload.type }}" \
|
||||
--backend "${{ matrix.backend.id }}" \
|
||||
--total-tasks "${{ matrix.workload.total_tasks }}" \
|
||||
--concurrency "${{ matrix.workload.concurrency }}" \
|
||||
--task-prefix "${{ matrix.backend.id }}-${{ matrix.workload.id }}" \
|
||||
--summary-file "$SUMMARY_FILE" \
|
||||
--mongo-uri "$MONGO_URI" \
|
||||
--mongo-database agentlightning_collection_bench
|
||||
|
||||
- name: Show collection benchmark summary
|
||||
if: ${{ always() }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ -f "$SUMMARY_FILE" ]; then
|
||||
echo "Collection benchmark summary (${{ matrix.backend.id }}):"
|
||||
cat "$SUMMARY_FILE"
|
||||
else
|
||||
echo "Summary file not found: $SUMMARY_FILE"
|
||||
fi
|
||||
|
||||
- name: Stop MongoDB
|
||||
if: ${{ always() && matrix.backend.needs_mongo }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd docker
|
||||
docker compose -f compose.mongo.yml down -v || true
|
||||
|
||||
- name: Upload collection artifacts
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: ${{ env.ARTIFACT_NAME }}
|
||||
path: ${{ env.ARTIFACT_DIR }}
|
||||
if-no-files-found: error
|
||||
@@ -0,0 +1,33 @@
|
||||
name: Dashboard
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 5 AM UTC+8
|
||||
- cron: '0 21 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
push:
|
||||
branches: [ main, stable/**/* ]
|
||||
|
||||
jobs:
|
||||
dashboard:
|
||||
name: Chromatic
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Run Chromatic
|
||||
uses: chromaui/action@v13
|
||||
with:
|
||||
projectToken: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
|
||||
workingDir: dashboard
|
||||
exitZeroOnChanges: false
|
||||
@@ -0,0 +1,64 @@
|
||||
name: Deploy Documentation
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
tags:
|
||||
- 'v*'
|
||||
workflow_dispatch:
|
||||
|
||||
concurrency:
|
||||
group: docs-deploy
|
||||
cancel-in-progress: false
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pages: write
|
||||
id-token: write
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --no-default-groups --group dev
|
||||
|
||||
- name: Configure Git
|
||||
run: |
|
||||
git config --global user.name "GitHub Actions"
|
||||
git config --global user.email "actions@github.com"
|
||||
|
||||
- name: Get version and commit
|
||||
id: version
|
||||
run: |
|
||||
if [[ $GITHUB_REF == refs/tags/* ]]; then
|
||||
VERSION=${GITHUB_REF#refs/tags/v}
|
||||
SOURCE_COMMIT=${GITHUB_SHA}
|
||||
else
|
||||
VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])")
|
||||
SOURCE_COMMIT=${GITHUB_SHA}
|
||||
fi
|
||||
echo "version=$VERSION" >> $GITHUB_OUTPUT
|
||||
echo "SOURCE_COMMIT=$SOURCE_COMMIT" >> $GITHUB_ENV
|
||||
|
||||
- name: Deploy versioned docs
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
run: |
|
||||
uv run --locked --no-sync mike deploy --push --update-aliases ${{ steps.version.outputs.version }} stable
|
||||
|
||||
- name: Deploy dev docs
|
||||
if: github.ref == 'refs/heads/main'
|
||||
run: |
|
||||
uv run --locked --no-sync mike deploy --push latest
|
||||
# Always set stable to default
|
||||
uv run --locked --no-sync mike set-default --push stable
|
||||
@@ -0,0 +1,116 @@
|
||||
name: Examples - APO
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 3 AM UTC+8
|
||||
- cron: '0 19 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-apo, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'APO - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('APO - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
apo:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-apo' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: APO (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
# This job is run on GitHub hosted runners rather than self-hosted runners because it needs no GPU.
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra apo \
|
||||
--group dev --group experiment --group agents --group core-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra apo \
|
||||
--group dev --group experiment --group agents --group core-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-apo-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: APO custom algorithm
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/apo
|
||||
uv run apo_custom_algorithm_trainer.py | tee _ci_apo.log
|
||||
# Check whether the log contains "Best prompt found:"
|
||||
grep "Best prompt found:" _ci_apo.log
|
||||
env:
|
||||
# New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
- name: APO custom algorithm debugger
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/apo
|
||||
uv run apo_debug.py --mode runner
|
||||
uv run apo_debug.py --mode hook
|
||||
uv run apo_debug.py --mode trainer
|
||||
env:
|
||||
# New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: APO built-in algorithm
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/apo
|
||||
uv run room_selector_apo.py
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
if: matrix.setup-script != 'legacy'
|
||||
@@ -0,0 +1,98 @@
|
||||
name: Examples - Azure
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 4 AM UTC+8
|
||||
- cron: '0 20 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-azure, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Azure - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Azure - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
azure:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-azure' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Azure (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
|
||||
timeout-minutes: 400
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups \
|
||||
--group dev --group experiment --group agents --group core-stable
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-azure-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Azure Login
|
||||
run: |
|
||||
az login --identity
|
||||
shell: bash
|
||||
|
||||
- name: Azure OpenAI Sanity Check
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
cd examples/azure
|
||||
python capital_agent.py
|
||||
shell: bash
|
||||
env:
|
||||
AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
id: azure_openai_sanity_check
|
||||
|
||||
- name: Azure OpenAI Supervised Fine-tuning
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
cd examples/azure
|
||||
python train_capital_agent.py --n-iterations 2 --cleanup
|
||||
shell: bash
|
||||
env:
|
||||
AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
|
||||
AZURE_OPENAI_API_VERSION: 2025-04-01-preview
|
||||
AZURE_RESOURCE_GROUP: ${{ secrets.AZURE_RESOURCE_GROUP }}
|
||||
AZURE_RESOURCE_NAME: ${{ secrets.AZURE_RESOURCE_NAME }}
|
||||
id: azure_openai_finetune
|
||||
@@ -0,0 +1,438 @@
|
||||
name: Examples - Calc-X
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 3 AM UTC+8
|
||||
- cron: '0 19 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-calc-x, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Calc-X - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Calc-X - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
calc-x-perf:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-calc-x' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Calc-X Performance (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 90
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group 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 }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-calc-x-performance-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare Calc-X dataset
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
|
||||
unzip calc-x-data.zip -d data
|
||||
rm calc-x-data.zip
|
||||
|
||||
- name: Calc-X MCP sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run tests/test_mcp_calculator.py
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
- name: Calc-X sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run legacy_calc_agent_debug.py
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
# Calc-X training suddenly works after running the sanity check.
|
||||
# And it has to be run before Spider training.
|
||||
# The client side used to hang in many of my attempts.
|
||||
# Don't ask why. Don't touch this.
|
||||
- name: Calc-X training
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
python train_calc_agent.py --val-file data/test_mini.parquet --ci
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train
|
||||
|
||||
- name: Validate Calc-X training
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
calc-x-variants:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-calc-x' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Calc-X Variants (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 90
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --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 --extra weave --extra mongo --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-calc-x-variants-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare Calc-X dataset
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
|
||||
unzip calc-x-data.zip -d data
|
||||
rm calc-x-data.zip
|
||||
|
||||
- name: Calc-X MCP sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run tests/test_mcp_calculator.py
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
- name: Calc-X sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run legacy_calc_agent_debug.py
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: Training with local model
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir data/qwen_model
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --model $(realpath data/qwen_model)
|
||||
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_local_model
|
||||
|
||||
- name: Validate training with local model
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_local_model.outputs.project_name }} ${{ steps.calc_x_train_local_model.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with LLM Proxy
|
||||
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 --llm-proxy
|
||||
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_llm_proxy
|
||||
|
||||
- name: Validate training with LLM Proxy
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_llm_proxy.outputs.project_name }} ${{ steps.calc_x_train_llm_proxy.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Setup Docker environments
|
||||
run: ./scripts/mongodb_docker_run.sh
|
||||
shell: bash
|
||||
|
||||
- name: Training with MongoDB
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --mongo-uri mongodb://localhost:27017/?replicaSet=rs0
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_mongo
|
||||
|
||||
- name: Validate training with MongoDB
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_mongo.outputs.project_name }} ${{ steps.calc_x_train_mongo.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with LoRA
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --lora
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_lora
|
||||
if: matrix.setup-script != 'legacy'
|
||||
|
||||
- name: Validate training with LoRA
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_lora.outputs.project_name }} ${{ steps.calc_x_train_lora.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
if: matrix.setup-script != 'legacy'
|
||||
|
||||
- name: Training with trajectory level aggregation
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --trajectory-level
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_trajectory_level
|
||||
|
||||
- name: Validate training with trajectory level aggregation
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_trajectory_level.outputs.project_name }} ${{ steps.calc_x_train_trajectory_level.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with Weave
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --weave
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_weave
|
||||
|
||||
- name: Validate training with Weave
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_weave.outputs.project_name }} ${{ steps.calc_x_train_weave.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with external store
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
|
||||
agl store --port 4747 &
|
||||
sleep 5
|
||||
AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=runner python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast &
|
||||
sleep 5
|
||||
AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --external-store-address http://localhost:4747 --val-file data/test_mini.parquet --ci-fast
|
||||
|
||||
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
|
||||
while pgrep -f agl; do
|
||||
echo "Waiting for agl to finish..."
|
||||
sleep 5
|
||||
done
|
||||
pkill -f train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
|
||||
while pgrep -f train_calc_agent.py; do
|
||||
echo "Waiting for train_calc_agent.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "train_calc_agent.py has finished."
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_external_store
|
||||
|
||||
- name: Validate training with external store
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_external_store.outputs.project_name }} ${{ steps.calc_x_train_external_store.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Training with role-based environment variables
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
|
||||
PYTHONUNBUFFERED=1 AGL_SERVER_HOST=127.0.0.1 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=runner python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast &
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 AGL_SERVER_HOST=0.0.0.0 AGL_SERVER_PORT=5858 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast
|
||||
|
||||
pkill -f train_calc_agent.py && echo "SIGTERM sent to train_calc_agent.py" || echo "No train_calc_agent.py process found"
|
||||
while pgrep -f train_calc_agent.py; do
|
||||
echo "Waiting for train_calc_agent.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "train_calc_agent.py has finished."
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: calc_x_train_role_based_env_var
|
||||
|
||||
- name: Validate training with role-based environment variables
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train_role_based_env_var.outputs.project_name }} ${{ steps.calc_x_train_role_based_env_var.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,168 @@
|
||||
name: Examples - ChartQA
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 6 AM UTC+8
|
||||
- cron: "0 22 * * *"
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-chartqa, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'ChartQA - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('ChartQA - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
chartqa:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-chartqa' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: ChartQA (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group image --group langchain --group vllm-0-10-2 --group torch-gpu-stable
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-chartqa-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare ChartQA dataset
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/chartqa
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1fWRt9hehg8_uV7BDWSCwKTycM60JcmGN/view?usp=sharing" -O chartqa-data.zip
|
||||
unzip chartqa-data.zip
|
||||
rm chartqa-data.zip
|
||||
shell: bash
|
||||
|
||||
- name: ChartQA sanity check with GPT
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/chartqa
|
||||
uv run python debug_chartqa_agent.py
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: Run vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
uv run --no-sync vllm serve Qwen/Qwen2-VL-2B-Instruct \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--max-model-len 4096 \
|
||||
--allowed-local-media-path "$(pwd)/data" \
|
||||
--enable-prefix-caching \
|
||||
--port 8088 &
|
||||
|
||||
VLLM_READY=0
|
||||
for i in {1..100}; do
|
||||
if curl -sSf http://localhost:8088/v1/models > /dev/null 2>&1; then
|
||||
echo "vLLM server is ready!"
|
||||
VLLM_READY=1
|
||||
break
|
||||
fi
|
||||
echo "Waiting for vLLM server to be ready... (${i})"
|
||||
sleep 5
|
||||
done
|
||||
if [[ "$VLLM_READY" != "1" ]]; then
|
||||
echo "vLLM server failed to start!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: ChartQA sanity check with vLLM
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
uv run python debug_chartqa_agent.py
|
||||
shell: bash
|
||||
env:
|
||||
USE_LLM_PROXY: "1"
|
||||
OPENAI_API_BASE: http://localhost:8088/v1
|
||||
OPENAI_MODEL: Qwen/Qwen2-VL-2B-Instruct
|
||||
|
||||
- name: Stop vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
pkill -f vllm
|
||||
for i in {1..60}; do
|
||||
if ! pgrep -f vllm; then
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: ChartQA training
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/chartqa
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_chartqa_agent.py ci
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: chartqa_train
|
||||
|
||||
- name: Validate ChartQA training
|
||||
run: |
|
||||
set -euo pipefail
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.chartqa_train.outputs.project_name }} ${{ steps.chartqa_train.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,151 @@
|
||||
name: Examples - Claude Code
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 4 AM UTC+8
|
||||
- cron: "0 20 * * *"
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-claude-code, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Claude Code - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Claude Code - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
claude-code:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-claude-code' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Claude Code (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: "3.12"
|
||||
setup-script: "stable"
|
||||
- python-version: "3.13"
|
||||
setup-script: "latest"
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups \
|
||||
--group dev --group experiment --group agents --group torch-gpu-stable
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-claude-code-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Download model
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
python -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen3-Coder-30B-A3B-Instruct')"
|
||||
|
||||
- name: Launch vLLM server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct \
|
||||
--max-model-len 131072 \
|
||||
--enable-auto-tool-choice \
|
||||
--tool-call-parser qwen3_coder \
|
||||
--port 45993 &
|
||||
|
||||
VLLM_READY=0
|
||||
for i in {1..100}; do
|
||||
if curl -sSf http://localhost:45993/v1/models > /dev/null 2>&1; then
|
||||
echo "vLLM server is ready!"
|
||||
VLLM_READY=1
|
||||
break
|
||||
fi
|
||||
echo "Waiting for vLLM server to be ready... (${i})"
|
||||
sleep 5
|
||||
done
|
||||
if [[ "$VLLM_READY" != "1" ]]; then
|
||||
echo "vLLM server failed to start!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Claude Code sanity check with vLLM models
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
cd examples/claude_code
|
||||
python claude_code_agent.py vllm --backend-model-high Qwen/Qwen3-Coder-30B-A3B-Instruct --backend-model-low Qwen/Qwen3-Coder-30B-A3B-Instruct --base-url http://localhost:45993/v1 --debug
|
||||
shell: bash
|
||||
|
||||
- name: Upload sanity check artifacts for vLLM
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: claude-code-sanity-check-vllm-${{ matrix.setup-script }}
|
||||
path: |
|
||||
examples/claude_code/data/
|
||||
examples/claude_code/logs/
|
||||
if-no-files-found: error
|
||||
|
||||
- name: Cleanup vLLM
|
||||
run: |
|
||||
set -euo pipefail
|
||||
pkill -f vllm
|
||||
for i in {1..60}; do
|
||||
if ! pgrep -f vllm; then
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
rm -rf examples/claude_code/data/
|
||||
rm -rf examples/claude_code/logs/
|
||||
|
||||
- name: Claude Code sanity check with OpenAI models
|
||||
run: |
|
||||
source .venv/bin/activate
|
||||
cd examples/claude_code
|
||||
python claude_code_agent.py openai --backend-model-high gpt-5.1-codex-mini --backend-model-low gpt-4.1-mini --debug
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
|
||||
- name: Upload sanity check artifacts for OpenAI
|
||||
if: ${{ always() }}
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: claude-code-sanity-check-openai-${{ matrix.setup-script }}
|
||||
path: |
|
||||
examples/claude_code/data/
|
||||
examples/claude_code/logs/
|
||||
if-no-files-found: error
|
||||
@@ -0,0 +1,151 @@
|
||||
name: Examples - Backward Compatibility
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 6 AM UTC+8
|
||||
- cron: '0 22 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-compat, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Backward Compatibility - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Backward Compatibility - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
backward-compatibility:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-compat' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Backward Compatibility (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra apo --extra verl \
|
||||
--group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
|
||||
- name: Override VERL (stable)
|
||||
run: |
|
||||
uv pip install verl==0.5.0 vllm==0.10.2
|
||||
if: matrix.setup-script == 'stable'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-backward-compatibility-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
- name: Prepare Calc-X dataset
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
|
||||
unzip calc-x-data.zip -d data
|
||||
rm calc-x-data.zip
|
||||
|
||||
- name: APO example (legacy client-server style)
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/apo
|
||||
uv run legacy_apo_client.py &
|
||||
sleep 3 # Wait for the client to be up
|
||||
uv run legacy_apo_server.py
|
||||
pkill -f legacy_apo_client.py && echo "SIGTERM sent to legacy_apo_client.py" || echo "No legacy_apo_client.py process found"
|
||||
while pgrep -f legacy_apo_client.py; do
|
||||
echo "Waiting for legacy_apo_client.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "legacy_apo_client.py has finished."
|
||||
sleep 10
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: Calc-X MCP sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run tests/test_mcp_calculator.py
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
- name: Calc-X sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/calc_x
|
||||
uv run legacy_calc_agent_debug.py
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
|
||||
- name: Calc-X training (legacy client-server style)
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/calc_x
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python legacy_calc_agent.py &
|
||||
bash legacy_train.sh
|
||||
pkill -f legacy_calc_agent.py && echo "SIGTERM sent to legacy_calc_agent.py" || echo "No legacy_calc_agent.py process found"
|
||||
while pgrep -f legacy_calc_agent.py; do
|
||||
echo "Waiting for legacy_calc_agent.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "legacy_calc_agent.py has finished."
|
||||
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
|
||||
|
||||
- name: Validate Calc-X training
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,179 @@
|
||||
name: Examples - RAG
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 6 AM UTC+8
|
||||
- cron: '0 22 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-rag, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'RAG - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('RAG - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
rag:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-rag' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: RAG (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group rag --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group rag --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-rag-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare RAG dataset
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/rag
|
||||
mkdir -p data
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" -O data/dataset_tiny.parquet
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" -O data/chunks_candidate_tiny.pkl
|
||||
uv run gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" -O data/index_hnsw_faiss_n32e40_tiny.index
|
||||
|
||||
- name: Run WIKI Retriever MCP Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
cd examples/rag
|
||||
uv run python wiki_retriever_mcp.py &
|
||||
for i in {1..20}; do
|
||||
sleep 5
|
||||
if nc -z localhost 8099; then
|
||||
echo "MCP server is up!"
|
||||
exit 0
|
||||
else
|
||||
echo "Waiting for MCP server to start..."
|
||||
fi
|
||||
done
|
||||
echo "MCP server failed to start within expected time."
|
||||
exit 1
|
||||
|
||||
- name: Run vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
vllm serve Qwen/Qwen2.5-1.5B-Instruct \
|
||||
--enable-auto-tool-choice \
|
||||
--tool-call-parser hermes \
|
||||
--port 8000 &
|
||||
|
||||
VLLM_READY=0
|
||||
for i in {1..100}; do
|
||||
if curl -sSf http://localhost:8000/v1/models > /dev/null 2>&1; then
|
||||
echo "vLLM server is ready!"
|
||||
VLLM_READY=1
|
||||
break
|
||||
fi
|
||||
echo "Waiting for vLLM server to be ready... (${i})"
|
||||
sleep 5
|
||||
done
|
||||
if [[ "$VLLM_READY" != "1" ]]; then
|
||||
echo "vLLM server failed to start!"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Run RAG Sanity check
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/rag
|
||||
uv run python rag_agent.py
|
||||
shell: bash
|
||||
|
||||
- name: Stop vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
pkill -f vllm
|
||||
for i in {1..60}; do
|
||||
if ! pgrep -f vllm; then
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: RAG training
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/rag
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_rag.py fast
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: rag_train
|
||||
|
||||
- name: Validate RAG training
|
||||
run: |
|
||||
set -ex
|
||||
# Allow up to 5 rollouts to fail to produce rewards
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.rag_train.outputs.project_name }} ${{ steps.rag_train.outputs.run_name }} --reward-tolerance 5
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,126 @@
|
||||
name: Examples - Spider
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 4 AM UTC+8
|
||||
- cron: '0 20 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-spider, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Spider - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Spider - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
spider:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-spider' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Spider (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
# legacy is omitted because langchain doesn't work with legacy vllm versions
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group langchain --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable)
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script == 'stable'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-spider-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Prepare Spider dataset
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/spider
|
||||
uv run gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
|
||||
unzip -q spider-data.zip -d data
|
||||
rm spider-data.zip
|
||||
|
||||
- name: Spider sanity check
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/spider
|
||||
uv run sql_agent.py
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
if: success() || failure()
|
||||
|
||||
- name: Spider training
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/spider
|
||||
../../scripts/restart_ray.sh
|
||||
sleep 5
|
||||
PYTHONUNBUFFERED=1 python train_sql_agent.py fast
|
||||
sleep 10
|
||||
shell: bash
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
id: spider_train
|
||||
|
||||
- name: Validate Spider training
|
||||
run: |
|
||||
set -ex
|
||||
uv run scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }} --reward-tolerance 5
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,170 @@
|
||||
name: Examples - Tinker
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 3 AM UTC+8
|
||||
- cron: '0 19 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-tinker, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Tinker - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Tinker - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
tinker:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-tinker' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Tinker (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-cpu]
|
||||
timeout-minutes: 150
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups \
|
||||
--group dev --group experiment --group agents --group torch-cpu --group core-stable --group tinker
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -euo pipefail
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-tinker-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
# TODO: Currently only test the client tracer implementation.
|
||||
- name: Tinker LLM sanity check (tracer text)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python -m tests.test_tinker_llm tracer-text
|
||||
shell: bash
|
||||
env:
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker LLM sanity check (tracer tool)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python -m tests.test_tinker_llm tracer-tool
|
||||
shell: bash
|
||||
env:
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker Hello
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python hello.py oneclick --ci
|
||||
shell: bash
|
||||
env:
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker Q20 Evaluate (GPT-4.1)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
mkdir -p logs
|
||||
python q20_evaluate.py --ci --model gpt-4.1 --output-file logs/q20_evaluate_gpt-4.1.jsonl
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
CREWAI_DISABLE_TELEMETRY: true
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker Q20 Evaluate (Qwen3-30B-A3B-Instruct-2507)
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python q20_evaluate.py --ci --model Qwen/Qwen3-30B-A3B-Instruct-2507 --output-file logs/q20_evaluate_qwen3-30b-a3b.jsonl
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
CREWAI_DISABLE_TELEMETRY: true
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker Q20 Training Dry Run
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
python q20_train.py dryrun --model qwen4b
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
CREWAI_DISABLE_TELEMETRY: true
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
|
||||
- name: Tinker Q20 Training
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/tinker
|
||||
agl store --port 4747 &
|
||||
sleep 5
|
||||
python q20_train.py runner --n-runners 4 &
|
||||
sleep 5
|
||||
python q20_train.py algo --model qwen4b --ci
|
||||
sleep 5
|
||||
|
||||
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
|
||||
while pgrep -f agl; do
|
||||
echo "Waiting for agl to finish..."
|
||||
sleep 5
|
||||
done
|
||||
pkill -f q20_train.py && echo "SIGTERM sent to q20_train.py" || echo "No q20_train.py process found"
|
||||
while pgrep -f q20_train.py; do
|
||||
echo "Waiting for q20_train.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "q20_train.py has finished."
|
||||
shell: bash
|
||||
env:
|
||||
OPENAI_BASE_URL: ${{ secrets.AZURE_OPENAI_ENDPOINT_SWEDEN }}
|
||||
OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY_SWEDEN }}
|
||||
CREWAI_DISABLE_TELEMETRY: true
|
||||
TINKER_API_KEY: ${{ secrets.TINKER_API_KEY }}
|
||||
@@ -0,0 +1,129 @@
|
||||
name: Examples - Unsloth
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 5 AM UTC+8
|
||||
- cron: '0 21 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-unsloth, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'Unsloth - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('Unsloth - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
unsloth:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-unsloth' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Unsloth (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
# Legacy versions are not supported for Unsloth examples.
|
||||
include:
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- name: Check disk space
|
||||
run: df -h
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies
|
||||
run: |
|
||||
uv sync --frozen --no-default-groups --extra verl \
|
||||
--group dev --group experiment --group trl --group agents --group torch-gpu-stable
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-unsloth-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Prepare Unsloth model
|
||||
run: |
|
||||
set -ex
|
||||
cd examples/unsloth
|
||||
rm -rf models
|
||||
uv run hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
|
||||
|
||||
- name: Unsloth SFT example
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/unsloth
|
||||
|
||||
agl store --port 4747 &
|
||||
sleep 5
|
||||
python sft_rollout_runners.py &
|
||||
sleep 5
|
||||
python sft_algorithm.py
|
||||
|
||||
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
|
||||
while pgrep -f agl; do
|
||||
echo "Waiting for agl to finish..."
|
||||
sleep 5
|
||||
done
|
||||
pkill -f sft_rollout_runners.py && echo "SIGTERM sent to sft_rollout_runners.py" || echo "No sft_rollout_runners.py process found"
|
||||
while pgrep -f sft_rollout_runners.py; do
|
||||
echo "Waiting for sft_rollout_runners.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
echo "sft_rollout_runners.py has finished."
|
||||
sleep 10
|
||||
|
||||
# Check models/version_2 must exist
|
||||
if [ ! -d "models/version_2" ]; then
|
||||
echo "models/version_2 does not exist"
|
||||
exit 1
|
||||
fi
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
|
||||
- name: Unsloth SFT example all-in-one
|
||||
run: |
|
||||
set -ex
|
||||
source .venv/bin/activate
|
||||
cd examples/unsloth
|
||||
rm -rf models/version_1 models/version_2
|
||||
|
||||
python sft_allinone.py
|
||||
if [ ! -d "models/version_2" ]; then
|
||||
echo "models/version_2 does not exist"
|
||||
exit 1
|
||||
fi
|
||||
env:
|
||||
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
|
||||
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
|
||||
@@ -0,0 +1,309 @@
|
||||
name: Issue Comment
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
permissions:
|
||||
pull-requests: write
|
||||
issues: write
|
||||
contents: write
|
||||
actions: read
|
||||
|
||||
jobs:
|
||||
dispatch:
|
||||
# Only run for comments on pull requests AND when the comment starts with "/ci"
|
||||
if: >
|
||||
github.event.issue.pull_request != null &&
|
||||
startsWith(github.event.comment.body, '/ci')
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
dispatched: ${{ steps.dispatch.outputs.dispatched }}
|
||||
event_types: ${{ steps.dispatch.outputs.event_types }}
|
||||
correlation_id: ${{ steps.dispatch.outputs.correlation_id }}
|
||||
trigger_comment_id: ${{ steps.dispatch.outputs.trigger_comment_id }}
|
||||
ack_comment_id: ${{ steps.ack.outputs.comment_id }}
|
||||
steps:
|
||||
- name: Guardrail — allow only members/collaborators
|
||||
id: guard
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const allowed = ['MEMBER','OWNER','COLLABORATOR'];
|
||||
const assoc = context.payload.comment.author_association;
|
||||
if (!allowed.includes(assoc)) {
|
||||
core.notice(`Ignoring /ci from ${context.payload.comment.user.login} (author_association=${assoc}).`);
|
||||
core.setOutput('skip', 'true');
|
||||
}
|
||||
|
||||
- name: Trigger repository dispatch
|
||||
id: dispatch
|
||||
if: steps.guard.outputs.skip != 'true'
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const owner = context.repo.owner;
|
||||
const repo = context.repo.repo;
|
||||
const pull_number = context.payload.issue.number;
|
||||
const comment = context.payload.comment;
|
||||
|
||||
// Fetch current PR state
|
||||
const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number });
|
||||
|
||||
// Add reaction so folks know we saw it
|
||||
try {
|
||||
await github.rest.reactions.createForIssueComment({
|
||||
owner,
|
||||
repo,
|
||||
comment_id: comment.id,
|
||||
content: 'rocket'
|
||||
});
|
||||
} catch (e) {
|
||||
core.info('Could not add reaction (likely due to permissions). Continuing.');
|
||||
}
|
||||
|
||||
const labels = (pr.labels ?? []).map(label => label.name);
|
||||
const directCiLabels = labels.filter(label => label.startsWith('ci-'));
|
||||
const hasCiAll = directCiLabels.includes('ci-all');
|
||||
const dedupe = new Set(
|
||||
directCiLabels.filter(label => label !== 'ci-all')
|
||||
);
|
||||
|
||||
if (!hasCiAll && dedupe.size === 0) {
|
||||
core.notice('No ci-* labels found on the pull request; nothing to dispatch.');
|
||||
core.setOutput('dispatched', 'false');
|
||||
core.setOutput('event_types', '');
|
||||
return;
|
||||
}
|
||||
|
||||
const correlation_id = `id-${comment.id}-${Date.now().toString(36)}`;
|
||||
|
||||
const clientPayload = {
|
||||
correlation_id,
|
||||
pull_number,
|
||||
pr_ref: `refs/pull/${pull_number}/merge`,
|
||||
pr_head_ref: pr.head.ref,
|
||||
pr_head_sha: pr.head.sha,
|
||||
pr_base_ref: pr.base.ref,
|
||||
pr_base_sha: pr.base.sha,
|
||||
trigger_comment_id: comment.id,
|
||||
trigger_comment_user: comment.user.login,
|
||||
};
|
||||
|
||||
const eventTypes = hasCiAll
|
||||
? ['ci-all']
|
||||
: Array.from(dedupe);
|
||||
for (const eventType of eventTypes) {
|
||||
await github.rest.repos.createDispatchEvent({
|
||||
owner,
|
||||
repo,
|
||||
event_type: eventType,
|
||||
client_payload: { ...clientPayload, ci_label: eventType }
|
||||
});
|
||||
core.notice(`Dispatched '${eventType}' event for PR #${pull_number}.`);
|
||||
}
|
||||
|
||||
core.setOutput('dispatched', 'true');
|
||||
core.setOutput('event_types', eventTypes.join(','));
|
||||
core.setOutput('correlation_id', correlation_id);
|
||||
core.setOutput('trigger_comment_id', String(comment.id));
|
||||
|
||||
- name: Acknowledge in thread (optional)
|
||||
if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched == 'true'
|
||||
id: ack
|
||||
uses: actions/github-script@v8
|
||||
env:
|
||||
EVENT_TYPES: ${{ steps.dispatch.outputs.event_types }}
|
||||
CORRELATION_ID: ${{ steps.dispatch.outputs.correlation_id }}
|
||||
with:
|
||||
script: |
|
||||
const eventTypes = (process.env.EVENT_TYPES || '')
|
||||
.split(',')
|
||||
.map(label => label.trim())
|
||||
.filter(Boolean);
|
||||
const formatted = eventTypes.map(label => `\`repository_dispatch:${label}\``).join(', ');
|
||||
const { owner, repo } = context.repo;
|
||||
const issue_number = context.payload.issue.number;
|
||||
const body = [
|
||||
`✅ CI trigger requested by @${context.payload.comment.user.login}.`,
|
||||
`Fired ${formatted}.`,
|
||||
'',
|
||||
`_Collecting run links for correlation \`${process.env.CORRELATION_ID}\`…_`
|
||||
].join('\n');
|
||||
const { data: comment } = await github.rest.issues.createComment({
|
||||
owner, repo, issue_number,
|
||||
body
|
||||
});
|
||||
core.setOutput('comment_id', String(comment.id));
|
||||
|
||||
- name: Notify missing ci label
|
||||
if: steps.guard.outputs.skip != 'true' && steps.dispatch.outputs.dispatched != 'true'
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
const issue_number = context.payload.issue.number;
|
||||
await github.rest.issues.createComment({
|
||||
owner,
|
||||
repo,
|
||||
issue_number,
|
||||
body: `⚠️ CI trigger ignored because the pull request has no \`ci-*\` labels (e.g. \`ci-apo\`, \`ci-calc-x\`). Add the desired labels and try \`/ci\` again.`
|
||||
});
|
||||
|
||||
watch:
|
||||
needs: dispatch
|
||||
if: needs.dispatch.outputs.dispatched == 'true'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 180
|
||||
steps:
|
||||
- name: Track dispatched runs and update comment
|
||||
uses: actions/github-script@v8
|
||||
env:
|
||||
CORRELATION_ID: ${{ needs.dispatch.outputs.correlation_id }}
|
||||
ACK_COMMENT_ID: ${{ needs.dispatch.outputs.ack_comment_id }}
|
||||
TRIGGER_COMMENT_ID: ${{ needs.dispatch.outputs.trigger_comment_id }}
|
||||
with:
|
||||
script: |
|
||||
const owner = context.repo.owner;
|
||||
const repo = context.repo.repo;
|
||||
const correlationId = process.env.CORRELATION_ID;
|
||||
if (!correlationId) {
|
||||
core.warning('No correlation id supplied; nothing to watch.');
|
||||
return;
|
||||
}
|
||||
|
||||
const ackCommentId = Number(process.env.ACK_COMMENT_ID || 0);
|
||||
if (!ackCommentId) {
|
||||
core.warning('No comment id available for updates; skipping watch.');
|
||||
return;
|
||||
}
|
||||
const triggerCommentId = Number(process.env.TRIGGER_COMMENT_ID || 0);
|
||||
if (!triggerCommentId) {
|
||||
core.warning('No trigger comment id available; skipping watch.');
|
||||
return;
|
||||
}
|
||||
|
||||
const prefix = `🚀 CI Watcher for correlation ${correlationId} triggered by comment ${triggerCommentId}`;
|
||||
core.notice(`Watching workflow runs for correlation '${correlationId}' using comment ${ackCommentId}.`);
|
||||
|
||||
function fmt(run) {
|
||||
const status = run.status;
|
||||
const conclusion = run.conclusion;
|
||||
const badge = status === 'completed'
|
||||
? (conclusion === 'success' ? '🟢' : conclusion === 'failure' ? '🔴' : '🟡')
|
||||
: (status === 'in_progress' ? '🟣' : '⚪️');
|
||||
const title = run.display_title || run.name || `run ${run.id}`;
|
||||
const statusText = status === 'completed' ? `${status}/${conclusion}` : status;
|
||||
return `- ${badge} [${title}](${run.html_url}) — \`${statusText}\``;
|
||||
}
|
||||
|
||||
const signatureOf = runs =>
|
||||
runs
|
||||
.map(run => `${run.id}:${run.status}/${run.conclusion || ''}`)
|
||||
.sort()
|
||||
.join('|');
|
||||
|
||||
const deadlineMs = Date.now() + 175 * 60 * 1000; // 175 minutes
|
||||
let found = [];
|
||||
|
||||
async function searchOnce() {
|
||||
const runs = await github.paginate(
|
||||
github.rest.actions.listWorkflowRunsForRepo,
|
||||
{ owner, repo, event: 'repository_dispatch', per_page: 100 }
|
||||
);
|
||||
const cutoff = new Date(Date.now() - 60 * 60 * 1000); // last hour
|
||||
return runs.filter(run => {
|
||||
const createdAt = new Date(run.created_at);
|
||||
const title = String(run.display_title || run.name || '');
|
||||
return createdAt >= cutoff && title.includes(correlationId);
|
||||
});
|
||||
}
|
||||
|
||||
while (Date.now() < deadlineMs) {
|
||||
found = await searchOnce();
|
||||
if (found.length > 0) {
|
||||
core.notice(`Discovered ${found.length} workflow run(s) for correlation '${correlationId}'.`);
|
||||
break;
|
||||
}
|
||||
core.notice(`No runs found yet for correlation '${correlationId}'; retrying shortly.`);
|
||||
await new Promise(res => setTimeout(res, 10000));
|
||||
}
|
||||
|
||||
if (found.length === 0) {
|
||||
core.notice(`Watcher timed out with no runs for correlation '${correlationId}'; notifying thread.`);
|
||||
await github.rest.issues.updateComment({
|
||||
owner,
|
||||
repo,
|
||||
comment_id: ackCommentId,
|
||||
body: [
|
||||
prefix,
|
||||
`⚠️ I couldn't find any workflow runs for correlation \`${correlationId}\`.`,
|
||||
`They may be delayed or misconfigured.`
|
||||
].join('\n')
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const runIds = new Set(found.map(run => run.id));
|
||||
let lastSignature = '';
|
||||
|
||||
async function refreshRuns() {
|
||||
const ids = Array.from(runIds);
|
||||
const refreshed = [];
|
||||
for (const id of ids) {
|
||||
const { data } = await github.rest.actions.getWorkflowRun({
|
||||
owner,
|
||||
repo,
|
||||
run_id: id
|
||||
});
|
||||
refreshed.push(data);
|
||||
}
|
||||
return refreshed;
|
||||
}
|
||||
|
||||
async function updateCommentIfChanged(runs, allDone) {
|
||||
const signature = signatureOf(runs);
|
||||
if (signature === lastSignature) {
|
||||
// Run statuses unchanged; skipping comment update.
|
||||
return;
|
||||
}
|
||||
lastSignature = signature;
|
||||
core.notice(`Updating comment ${ackCommentId} with ${runs.length} run status entries (allDone=${allDone}).`);
|
||||
await github.rest.issues.updateComment({
|
||||
owner,
|
||||
repo,
|
||||
comment_id: ackCommentId,
|
||||
body: [
|
||||
prefix,
|
||||
`🏃♀️ Tracking ${runs.length} workflow run(s):`,
|
||||
'',
|
||||
...runs.map(fmt),
|
||||
'',
|
||||
allDone ? '✅ All runs completed.' : '_Still running…_'
|
||||
].join('\n')
|
||||
});
|
||||
}
|
||||
|
||||
await updateCommentIfChanged(found, found.every(run => run.status === 'completed'));
|
||||
|
||||
while (Date.now() < deadlineMs) {
|
||||
const latest = await searchOnce();
|
||||
for (const run of latest) {
|
||||
if (!runIds.has(run.id)) {
|
||||
runIds.add(run.id);
|
||||
core.notice(`Detected additional run ${run.id} (${run.name || run.display_title || 'unnamed'}) for correlation '${correlationId}'.`);
|
||||
}
|
||||
}
|
||||
const current = await refreshRuns();
|
||||
const allDone = current.every(run => run.status === 'completed');
|
||||
await updateCommentIfChanged(current, allDone);
|
||||
if (allDone) {
|
||||
core.notice(`All runs for correlation '${correlationId}' completed; stopping watcher.`);
|
||||
break;
|
||||
}
|
||||
await new Promise(res => setTimeout(res, 60000));
|
||||
}
|
||||
|
||||
if (Date.now() >= deadlineMs) {
|
||||
core.warning(`Watcher hit the deadline while monitoring correlation '${correlationId}'.`);
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
# Pre-defined workflow with workflow_dispatch trigger,
|
||||
# convenient for testing and debugging.
|
||||
|
||||
name: Playground
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
playground:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v6
|
||||
- name: Run script
|
||||
run: |
|
||||
echo "Hello, world!"
|
||||
@@ -0,0 +1,59 @@
|
||||
name: PyPI Nightly Build
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run daily at 6:00 AM UTC+8
|
||||
- cron: '0 22 * * *'
|
||||
workflow_dispatch: # Allow manual trigger
|
||||
|
||||
jobs:
|
||||
publish-test-pypi:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --no-default-groups --group dev
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Get current version
|
||||
id: get_version
|
||||
run: |
|
||||
VERSION=$(grep '^version = ' pyproject.toml | sed 's/version = "\(.*\)"/\1/')
|
||||
echo "version=$VERSION" >> $GITHUB_OUTPUT
|
||||
echo "Current version: $VERSION"
|
||||
|
||||
- name: Create development version
|
||||
run: |
|
||||
# Create a dev version with timestamp
|
||||
TIMESTAMP=$(date +%Y%m%d%H%M%S)
|
||||
DEV_VERSION="${{ steps.get_version.outputs.version }}.dev$TIMESTAMP"
|
||||
echo "Creating dev version: $DEV_VERSION"
|
||||
./scripts/bump_version.sh "$DEV_VERSION"
|
||||
|
||||
- name: Build package
|
||||
run: |
|
||||
uv build
|
||||
|
||||
- name: Publish to Test PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
repository-url: https://test.pypi.org/legacy/
|
||||
@@ -0,0 +1,81 @@
|
||||
name: PyPI Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'v*' # Trigger on version tags like v1.0.0, v1.2.3, etc.
|
||||
workflow_dispatch: # Allow manual trigger
|
||||
|
||||
jobs:
|
||||
check-version:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
outputs:
|
||||
version: ${{ steps.get_version.outputs.version }}
|
||||
tag_version: ${{ steps.get_tag.outputs.tag_version }}
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v6
|
||||
|
||||
- name: Get version from pyproject.toml
|
||||
id: get_version
|
||||
run: |
|
||||
VERSION=$(grep '^version = ' pyproject.toml | sed 's/version = "\(.*\)"/\1/')
|
||||
echo "version=$VERSION" >> $GITHUB_OUTPUT
|
||||
echo "Package version: $VERSION"
|
||||
|
||||
- name: Get tag version
|
||||
id: get_tag
|
||||
run: |
|
||||
TAG_VERSION=${GITHUB_REF#refs/tags/v}
|
||||
echo "tag_version=$TAG_VERSION" >> $GITHUB_OUTPUT
|
||||
echo "Tag version: $TAG_VERSION"
|
||||
|
||||
- name: Verify version matches tag
|
||||
run: |
|
||||
if [ "${{ steps.get_version.outputs.version }}" != "${{ steps.get_tag.outputs.tag_version }}" ]; then
|
||||
echo "Error: Version in pyproject.toml (${{ steps.get_version.outputs.version }}) does not match tag (${{ steps.get_tag.outputs.tag_version }})"
|
||||
exit 1
|
||||
fi
|
||||
echo "Version check passed!"
|
||||
|
||||
publish-pypi:
|
||||
needs: check-version
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --no-default-groups --group dev
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Build package
|
||||
run: |
|
||||
uv build
|
||||
|
||||
- name: Verify package contents
|
||||
run: |
|
||||
uv run --locked --no-sync python -m tarfile -l dist/*.tar.gz
|
||||
uv run --locked --no-sync python -m zipfile -l dist/*.whl
|
||||
|
||||
- name: Publish to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -0,0 +1,363 @@
|
||||
name: GPU Test
|
||||
permissions:
|
||||
contents: read
|
||||
on:
|
||||
schedule:
|
||||
# Every day at 5 AM UTC+8
|
||||
- cron: '0 21 * * *'
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
repository_dispatch:
|
||||
types: [ci-gpu, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'GPU Test - PR #{0} - {1} - {2}',
|
||||
github.event.client_payload.pull_number,
|
||||
github.event.client_payload.ci_label,
|
||||
github.event.client_payload.correlation_id
|
||||
)
|
||||
|| format('GPU Test - {0}', github.event_name) }}
|
||||
|
||||
jobs:
|
||||
tests-full:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-gpu' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Full Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})
|
||||
|
||||
runs-on: ${{ matrix.mark.runs-on }}
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
mark:
|
||||
- id: store
|
||||
display-name: Store
|
||||
pytest-mark: 'store' # store tests should not require gpu
|
||||
runs-on: ubuntu-latest
|
||||
has-gpu: false
|
||||
# AgentOps needs to be separated because it injects tricky global state.
|
||||
- id: agentops
|
||||
display-name: AgentOps
|
||||
pytest-mark: 'agentops' # including agentops+litellm tests here
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
has-gpu: true
|
||||
# Similar for Weave.
|
||||
- id: weave
|
||||
display-name: Weave
|
||||
pytest-mark: 'weave'
|
||||
runs-on: ubuntu-latest # No GPU tests for Weave.
|
||||
has-gpu: false
|
||||
# Other tests that require GPU
|
||||
- id: gpu
|
||||
display-name: GPU required
|
||||
pytest-mark: '(gpu or llmproxy) and not agentops'
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
has-gpu: true
|
||||
# Other uncovered tests
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not weave and not gpu and not llmproxy'
|
||||
runs-on: ubuntu-latest
|
||||
has-gpu: false
|
||||
env:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
if: matrix.mark.has-gpu
|
||||
run: nvidia-smi
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.env.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
|
||||
- name: Sync dependencies (latest, gpu)
|
||||
if: matrix.env.setup-script == 'latest' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group torch-gpu-stable
|
||||
# Don't install vllm/pytorch on CPU counterparts
|
||||
- name: Sync dependencies (latest, cpu)
|
||||
if: matrix.env.setup-script == 'latest' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group core-stable
|
||||
- name: Sync dependencies (stable, gpu)
|
||||
if: matrix.env.setup-script == 'stable' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.env.setup-script }}
|
||||
- name: Sync dependencies (stable, cpu)
|
||||
if: matrix.env.setup-script == 'stable' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group langchain --group core-stable
|
||||
# Don't install langchain for legacy dependency because it has conflicts with torch.
|
||||
- name: Sync dependencies (legacy, gpu)
|
||||
if: matrix.env.setup-script == 'legacy' && matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group torch-gpu-legacy
|
||||
- name: Sync dependencies (legacy, cpu)
|
||||
if: matrix.env.setup-script == 'legacy' && !matrix.mark.has-gpu
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --extra mongo --group dev --group agents --group core-legacy
|
||||
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-tests-full-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Setup Docker environments
|
||||
run: ./scripts/mongodb_docker_run.sh
|
||||
shell: bash
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
# mongo, openai, gpu, all enabled by default
|
||||
- name: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests -m "${{ matrix.mark.pytest-mark }}${{ matrix.env.setup-script == 'legacy' && ' and not langchain' || '' }}"
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
AGL_TEST_MONGO_URI: mongodb://localhost:27017/?replicaSet=rs0
|
||||
|
||||
|
||||
minimal-examples:
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-gpu' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Minimal Examples with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
|
||||
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
steps:
|
||||
- name: Check GPU status
|
||||
run: nvidia-smi
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group langchain --group torch-gpu-stable
|
||||
if: matrix.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group langchain --group torch-gpu-${{ matrix.setup-script }}
|
||||
if: matrix.setup-script == 'stable'
|
||||
# Don't install langchain for legacy dependency because it has conflicts with torch.
|
||||
- name: Sync dependencies (legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra mongo --group dev --group agents --group torch-gpu-legacy
|
||||
if: matrix.setup-script == 'legacy'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- name: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
env:
|
||||
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
|
||||
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
|
||||
|
||||
- name: Write Traces via Otel Tracer
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python write_traces.py otel
|
||||
sleep 5
|
||||
|
||||
- name: Write Traces via AgentOps Tracer
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
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
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
agl store --port 45993 --log-level DEBUG &
|
||||
sleep 5
|
||||
python write_traces.py otel --use-client
|
||||
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
|
||||
while pgrep -f agl; do
|
||||
echo "Waiting for agl to finish..."
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: Write Traces via AgentOps Tracer with Client
|
||||
env:
|
||||
OPENAI_BASE_URL: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
agl store --port 45993 --log-level DEBUG &
|
||||
sleep 5
|
||||
python write_traces.py agentops --use-client
|
||||
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
|
||||
while pgrep -f agl; do
|
||||
echo "Waiting for agl to finish..."
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: vLLM Server
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python vllm_server.py Qwen/Qwen2.5-0.5B-Instruct
|
||||
|
||||
- name: LLM Proxy (OpenAI backend)
|
||||
env:
|
||||
OPENAI_API_BASE: http://localhost:12306/
|
||||
OPENAI_API_KEY: dummy
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
|
||||
python llm_proxy.py openai gpt-4.1-mini &
|
||||
|
||||
LLM_PROXY_READY=0
|
||||
for attempt in $(seq 1 30); do
|
||||
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
|
||||
LLM_PROXY_READY=1
|
||||
break
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
if [[ "$LLM_PROXY_READY" != "1" ]]; then
|
||||
echo "LLM proxy failed to become healthy" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python llm_proxy.py test gpt-4.1-mini
|
||||
|
||||
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
|
||||
while pgrep -f llm_proxy.py; do
|
||||
echo "Waiting for llm_proxy.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: LLM Proxy (vLLM backend)
|
||||
if: matrix.setup-script != 'legacy' # Skip if return_token_ids is not supported
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct &
|
||||
|
||||
LLM_PROXY_READY=0
|
||||
for attempt in $(seq 1 30); do
|
||||
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
|
||||
LLM_PROXY_READY=1
|
||||
break
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
if [[ "$LLM_PROXY_READY" != "1" ]]; then
|
||||
echo "LLM proxy failed to become healthy" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
|
||||
|
||||
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
|
||||
while pgrep -f llm_proxy.py; do
|
||||
echo "Waiting for llm_proxy.py to finish..."
|
||||
sleep 5
|
||||
done
|
||||
|
||||
- name: MultiMetrics backend example
|
||||
run: |
|
||||
set -euo pipefail
|
||||
source .venv/bin/activate
|
||||
cd examples/minimal
|
||||
python write_metrics.py --duration 8 --prom-port 9105 --prom-host 0.0.0.0 2>&1 | tee metrics.log &
|
||||
pid=$!
|
||||
|
||||
for attempt in $(seq 1 20); do
|
||||
if curl -sSf http://localhost:9105/metrics | grep -q minimal_requests_total; then
|
||||
echo "Metrics endpoint responding"
|
||||
wait $pid
|
||||
cat metrics.log
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
echo "Metrics endpoint did not respond"
|
||||
exit 1
|
||||
@@ -0,0 +1,238 @@
|
||||
name: CPU Test
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main, stable/**/* ]
|
||||
pull_request:
|
||||
branches: [ main, stable/**/* ]
|
||||
workflow_dispatch:
|
||||
|
||||
schedule:
|
||||
# Every day at noon and midnight
|
||||
- cron: '0 0,12 * * *'
|
||||
|
||||
jobs:
|
||||
|
||||
lint:
|
||||
strategy:
|
||||
matrix:
|
||||
setup: [fast, slow, next]
|
||||
fail-fast: false
|
||||
name: Lint - ${{ matrix.setup }}
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: '3.12'
|
||||
- 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 \
|
||||
--group torch-cpu \
|
||||
--group torch-stable \
|
||||
--group trl \
|
||||
--group tinker \
|
||||
--group agents \
|
||||
--group langchain \
|
||||
--no-default-groups
|
||||
if: matrix.setup != 'fast'
|
||||
# This pre-commit skips JavaScript on purpose.
|
||||
- name: Run pre-commit
|
||||
uses: pre-commit/action@v3.0.1
|
||||
- name: Check Python headers
|
||||
run: uv run --locked --no-sync scripts/check_headers.py
|
||||
if: matrix.setup == 'fast'
|
||||
- name: Run Black
|
||||
run: uv run --locked --no-sync black --check .
|
||||
if: matrix.setup != 'next'
|
||||
- name: Run isort
|
||||
run: uv run --locked --no-sync isort --check-only .
|
||||
if: matrix.setup != 'next'
|
||||
- name: Run pyright (fast)
|
||||
run: uv run --locked --no-sync pyright -p pyrightconfig.fast.json
|
||||
if: matrix.setup == 'fast'
|
||||
- name: Run pyright (slow)
|
||||
run: uv run --locked --no-sync pyright -p pyrightconfig.json
|
||||
if: matrix.setup != 'fast'
|
||||
|
||||
lint-js:
|
||||
name: Lint - JavaScript
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Run ESLint
|
||||
run: cd dashboard && npm run eslint
|
||||
- name: Run Prettier
|
||||
run: cd dashboard && npm run prettier
|
||||
- name: Run Stylelint
|
||||
run: cd dashboard && npm run stylelint
|
||||
- name: Run Typecheck
|
||||
run: cd dashboard && npm run typecheck
|
||||
- name: Verify build
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
docs:
|
||||
name: Build documentation
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.12'
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Sync dependencies
|
||||
run: uv sync --frozen --no-default-groups --group dev
|
||||
- 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
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: documentation-site
|
||||
path: site/
|
||||
compression-level: 6
|
||||
|
||||
test:
|
||||
strategy:
|
||||
matrix:
|
||||
mark:
|
||||
# store has many tests and is a good isolated group.
|
||||
- id: store
|
||||
display-name: Store
|
||||
pytest-mark: 'store'
|
||||
# AgentOps needs to be separated because it injects tricky global state.
|
||||
- id: agentops
|
||||
display-name: AgentOps
|
||||
pytest-mark: 'agentops'
|
||||
# Similar for Weave.
|
||||
- id: weave
|
||||
display-name: Weave
|
||||
pytest-mark: 'weave'
|
||||
# litellm proxy tests are slow
|
||||
- id: llmproxy
|
||||
display-name: LLM proxy
|
||||
pytest-mark: 'llmproxy'
|
||||
# Robustness of utilities is important. There are many tests.
|
||||
- id: utils
|
||||
display-name: Utilities
|
||||
pytest-mark: 'utils'
|
||||
# unmarked tests: adapter, execution engine, etc.
|
||||
- id: others
|
||||
display-name: Others
|
||||
pytest-mark: 'not store and not agentops and not weave and not llmproxy and not utils'
|
||||
env:
|
||||
- python-version: '3.10'
|
||||
setup-script: 'legacy'
|
||||
- python-version: '3.11'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.12'
|
||||
setup-script: 'stable'
|
||||
- python-version: '3.13'
|
||||
setup-script: 'latest'
|
||||
fail-fast: false
|
||||
|
||||
name: Test (${{ matrix.mark.display-name }}, ${{ matrix.env.setup-script }}, Python ${{ matrix.env.python-version }})
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: ${{ matrix.env.python-version }}
|
||||
- name: Upgrade dependencies (latest)
|
||||
run: uv lock --upgrade
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
- name: Sync dependencies (latest)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-stable
|
||||
if: matrix.env.setup-script == 'latest'
|
||||
- name: Sync dependencies (stable & legacy)
|
||||
run: uv sync --frozen --no-default-groups --extra apo --extra weave --group dev --group agents --group langchain --group core-${{ matrix.env.setup-script }}
|
||||
if: matrix.env.setup-script != 'latest'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
uv pip freeze | tee requirements-freeze.txt
|
||||
echo "UV_LOCKED=1" >> $GITHUB_ENV
|
||||
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
|
||||
- name: Upload dependencies artifact
|
||||
uses: actions/upload-artifact@v6
|
||||
with:
|
||||
name: dependencies-${{ matrix.mark.id }}-${{ matrix.env.python-version }}-${{ matrix.env.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Build dashboard
|
||||
run: cd dashboard && npm run build
|
||||
|
||||
- name: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests -m "not mongo and not openai and not gpu and (${{ matrix.mark.pytest-mark }})"
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
|
||||
test-js:
|
||||
name: Test (JavaScript)
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: dashboard/package-lock.json
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
python-version: '3.12'
|
||||
- name: Sync Python dependencies
|
||||
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group core-stable
|
||||
- name: Install JavaScript dependencies
|
||||
run: cd dashboard && npm ci
|
||||
- name: Run vitest
|
||||
run: cd dashboard && npm run vitest
|
||||
+220
@@ -0,0 +1,220 @@
|
||||
# Agentlightning specific files
|
||||
verl_old
|
||||
meta-llama/**
|
||||
**/debug/**/*.png
|
||||
**/debug/**/*.json
|
||||
requirements-freeze*.txt
|
||||
/playground
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
#uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||
.pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# MacOS
|
||||
.DS_Store
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
# Abstra
|
||||
# Abstra is an AI-powered process automation framework.
|
||||
# Ignore directories containing user credentials, local state, and settings.
|
||||
# Learn more at https://abstra.io/docs
|
||||
.abstra/
|
||||
|
||||
# Visual Studio Code
|
||||
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||
# you could uncomment the following to ignore the enitre vscode folder
|
||||
.vscode/
|
||||
|
||||
# Emacs backup files
|
||||
*~
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Cursor
|
||||
# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
|
||||
# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
|
||||
# refer to https://docs.cursor.com/context/ignore-files
|
||||
.cursorignore
|
||||
.cursorindexingignore
|
||||
|
||||
# Claude
|
||||
.claude/*.local.json
|
||||
|
||||
# Dashboard generated files
|
||||
agentlightning/dashboard/**/*.css
|
||||
agentlightning/dashboard/**/*.js
|
||||
agentlightning/dashboard/**/*.html
|
||||
agentlightning/dashboard/**/*.svg
|
||||
|
||||
# Docker data
|
||||
docker/data/
|
||||
@@ -0,0 +1,77 @@
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v6.0.0
|
||||
hooks:
|
||||
- id: end-of-file-fixer
|
||||
exclude: (.*store-openapi\.json$)
|
||||
- id: trailing-whitespace
|
||||
- id: check-yaml
|
||||
exclude: ^mkdocs\.yml$
|
||||
- id: check-toml
|
||||
- id: check-added-large-files
|
||||
args: ["--maxkb=1024"]
|
||||
exclude: (^uv\.lock$)|(^docs/assets/.*\.svg$)|(.*store-openapi\.json$)
|
||||
- id: check-shebang-scripts-are-executable
|
||||
- id: detect-private-key
|
||||
- repo: https://github.com/pycqa/isort
|
||||
rev: 6.0.1
|
||||
hooks:
|
||||
- id: isort
|
||||
args: ["."]
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 25.1.0
|
||||
hooks:
|
||||
- id: black
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
args: ["."]
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: prettier
|
||||
name: prettier (dashboard)
|
||||
language: system
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
entry: >
|
||||
bash -c '
|
||||
cd dashboard || exit 1
|
||||
if [ -d node_modules ]; then
|
||||
echo "✅ node_modules already exists"
|
||||
npx prettier --cache --write "**/*.{ts,tsx,mjs,cjs}"
|
||||
else
|
||||
echo "⚠️ node_modules not found — npx is not reliable. Skipping."
|
||||
fi
|
||||
'
|
||||
|
||||
- id: eslint
|
||||
name: eslint (dashboard)
|
||||
language: system
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
entry: >
|
||||
bash -c '
|
||||
cd dashboard || exit 1
|
||||
if [ -d node_modules ]; then
|
||||
echo "✅ node_modules already exists"
|
||||
npx eslint --cache --fix .
|
||||
else
|
||||
echo "⚠️ node_modules not found — npx is not reliable. Skipping."
|
||||
fi
|
||||
'
|
||||
|
||||
- id: stylelint
|
||||
name: stylelint (dashboard)
|
||||
language: system
|
||||
pass_filenames: false
|
||||
always_run: true
|
||||
entry: >
|
||||
bash -c '
|
||||
cd dashboard || exit 1
|
||||
if [ -d node_modules ]; then
|
||||
echo "✅ node_modules already exists"
|
||||
npx stylelint --cache --fix "**/*.css"
|
||||
else
|
||||
echo "⚠️ node_modules not found — npx is not reliable. Skipping."
|
||||
fi
|
||||
'
|
||||
@@ -0,0 +1 @@
|
||||
3.12
|
||||
@@ -0,0 +1,41 @@
|
||||
# Repository Guidelines
|
||||
|
||||
## Architecture Overview
|
||||
Agent Lightning runs through a continuous loop: runners and tracers emit spans, `LightningStore` (`agentlightning/store/`) keeps them synchronized, and algorithms in `agentlightning/algorithm/` consume those traces to improve behavior.
|
||||
|
||||
## Project Structure & Module Organization
|
||||
- `agentlightning/`: adapters, execution stack, training loop, tracer, reward logic, and the `agl` CLI.
|
||||
- `docs/` & `examples/`: narrative and procedural docs (assets in `docs/assets/`, navigation in `mkdocs.yml`) plus runnable workflows whose READMEs point to their companion how-to guides. `docs/how-to` covers task-focused instructions, while `docs/tutorials` explains concepts and subsystems.
|
||||
- `dashboard/`, `scripts/`, `tests/`: UI bundles, release/dataset/CI automation, and mirrored coverage of the runtime tree. Record download steps rather than committing binaries.
|
||||
|
||||
## Build, Test, and Development Commands
|
||||
- `uv sync --group dev` — provision tooling once per environment.
|
||||
- `uv run --no-sync pytest -v` — execute the full suite; add a path or `-k expr` to narrow the run.
|
||||
- `uv run --no-sync pyright` — enforce static typing parity with CI.
|
||||
- `uv run --no-sync pre-commit run --all-files --show-diff-on-failure` and `uv run --no-sync mkdocs build --strict` — keep formatting tidy and documentation valid.
|
||||
Always commit the refreshed `uv.lock` when dependencies shift, and mention optional groups (VERL, APO, GPU) in PR notes.
|
||||
|
||||
## Common Issues & Fixes
|
||||
- When `uv run` errors with `Permission denied` under `~/.cache`, override both cache locations inline: ``UV_CACHE="$(pwd)/.cache_uv" XDG_CACHE_HOME="$(pwd)/.cache_xdg" uv run --no-sync <command>``.
|
||||
|
||||
## Coding Style & Naming Conventions
|
||||
- Target `requires-python >= 3.10`, four-space indentation, 120-character lines (though docstrings may run longer), and formatter-owned diffs (Black + isort, `black` profile). Use `snake_case` for modules, functions, and variables; `PascalCase` for classes and React components; lowercase hyphenation for CLI flags, branch names, and TypeScript filenames.
|
||||
- Maintain exhaustive type hints (pyright enforces them) and prefer shared dataclasses or Pydantic models from `agentlightning.types`.
|
||||
- Author Google-style docstrings for new modules or public methods—succinct descriptions, no redundant type info, no redundant `Key features/components` bullet points. Use mkdocs styles: `[][]` syntax for cross-references and single backticks for inline code blocks.
|
||||
- Writing logs is encouraged, especially for long functions with multiple steps and try-except blocks that catch all exceptions. Use `logging.getLogger(__name__)` to get loggers. Distinguish between DEBUG, INFO, WARNING, and ERROR logs.
|
||||
|
||||
## Testing Guidelines
|
||||
- Mirror runtime directories under `tests/` and match filenames for quick traceability.
|
||||
- Parametrize pytest cases and apply markers (`openai`, `gpu`, `agentops`, `mongo`, `llmproxy`) so optional suites can be skipped via selectors like `-m "not mongo"` yet still exercised in CI.
|
||||
- Lean on fixtures, favor real stores/spans/agents over mocks, and drive coverage across the majority of branches.
|
||||
- If an imported module is missing from the environment, check whether `uv sync` has been run with the right groups. Do not make stubs for external dependencies unless necessary.
|
||||
|
||||
## Example Contributions
|
||||
- Ship each example with a README that includes smoke-test instructions so maintainers can validate quickly. The README must contain an "Included Files" section summarizing every file and its role.
|
||||
- Keep runnable example modules self-contained with a module-level docstring describing CLI usage. Document important or educational classes/functions with targeted docstrings and inline comments where clarity matters.
|
||||
- Add a CI workflow per example named `examples-<name>.yml` in `.github/workflows/`. Register it in `badge-<name>.yml`, `badge-examples.yml`, and `badge-latest.yml` when applicable so badges stay accurate.
|
||||
|
||||
## Commit & Pull Request Guidelines
|
||||
- Branch from a fresh `main` using `feature/<slug>`, `fix/<slug>`, `docs/<slug>`, or `chore/<slug>`.
|
||||
- Write imperative, scoped commits, reference issues with `Fixes #123`, and rerun pre-commit plus the relevant pytest/doc builds before pushing.
|
||||
- Use PR descriptions to summarize intent, list verification commands, call out dependency or docs-navigation updates, and link new docs/examples via `mkdocs.yml` or `examples/README.md`. Include logs for dashboard changes.
|
||||
@@ -0,0 +1,19 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
@@ -0,0 +1,77 @@
|
||||
# Responsible AI Transparency Documentation - Agent Lightning
|
||||
|
||||
## OVERVIEW
|
||||
|
||||
Agent Lightning is a flexible and extensible framework that enables seamless agent optimization for any existing agent framework. Agent optimization includes various data-driven techniques to customize the agent for better performance, including but not limited to model fine-tuning, prompt tuning, and model selection. And the agent frameworks refer to popular and easy-to-use agent developing frameworks such as OpenAI Agents SDK, Microsoft AutoGen, and LangChain.
|
||||
|
||||
### WHAT CAN AGENT LIGHTNING DO
|
||||
Agent lightning was developed to bridge the gap between agent workflow development and agent optimization, empowering developers to go beyond static, pre-trained models and unlock the full potential of adaptive, learning-based agents. Agent Lightning is a training framework which can be used for any LLMs.
|
||||
|
||||
### INTENDED USES
|
||||
Agent Lightning is best suited for agent researchers and developers. They can easily fine-tune models in existing agent frameworks with Agent Lightning. This can improve model performance on the targeted scenarios.
|
||||
|
||||
### OUT-OF-SCOPE USES
|
||||
Agent Lightning is not well-suited for users who are not familiar with agent development and machine learning concepts.
|
||||
|
||||
We do not recommend using Agent Lightning in commercial or real-world applications without further testing and development. It is being released for research purposes.
|
||||
|
||||
Agent Lightning was not designed or evaluated for all possible downstream purposes. Developers should consider its inherent limitations as they select use cases, and evaluate and mitigate for accuracy, safety, and fairness concerns specific to each intended downstream use.
|
||||
|
||||
Agent Lightning should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.
|
||||
|
||||
We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare).
|
||||
|
||||
## HOW TO GET STARTED
|
||||
To begin using Agent Lightning, here are some instructions.
|
||||
1. Install dependencies, including Python, uv, PyTorch, FlashAttention, vLLM, verl.
|
||||
2. Clone and install Agent Lightning.
|
||||
3. Convert the dataset (provided by the user) into parquet file, which contains multiple columns. Each column contains a data id, an input and an expected output.
|
||||
4. Run agent, which is developed by the user.
|
||||
5. Run the training process via “bash train.sh”
|
||||
|
||||
## EVALUATION
|
||||
Agent Lightning was evaluated on its ability to correctly complete 3 example tasks: (1) Math. The model needs to answer some math questions, and when answering one question, the model can use the calculator as its tool to help answer. (2) Text2SQL. The model is given a question related to the database, and it is required to generate a SQL which can query the database, find the information to answer the question. (3) Retrieval-Augmented Generation (RAG). The model is given a question which needs some information from Wikipedia to answer. The model is required to generate some queries to find the related information in Wikipedia, and answer the question according to retrieved documents.
|
||||
|
||||
### EVALUATION METHODS AND RESULTS
|
||||
For detailed evaluation methods and results, please refer to the latest version of our [technical report](https://arxiv.org/abs/2508.03680).
|
||||
|
||||
|
||||
## LIMITATIONS
|
||||
Agent Lightning was developed for research and experimental purposes. Further testing and validation are needed before considering its application in commercial or real-world scenarios.
|
||||
|
||||
Agent Lightning was designed and tested using the English language. Performance in other languages may vary and should be assessed by someone who is both an expert in the expected outputs and a native speaker of that language.
|
||||
|
||||
Outputs generated by AI may include factual errors, fabrication, or speculation. Users are responsible for assessing the accuracy of generated content. All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
|
||||
Agent Lightning inherits any biases, errors, or omissions produced by its base model. Developers are advised to choose an appropriate base LLM/MLLM carefully, depending on the intended use case.
|
||||
We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model.
|
||||
|
||||
## BEST PRACTICES
|
||||
Better performance can be achieved by following the instructions in how to get started section.
|
||||
|
||||
We strongly encourage users to use LLMs/MLLMs that support robust Responsible AI mitigations, such as Azure Open AI (AOAI) services. Such services continually update their safety and RAI mitigations with the latest industry standards for responsible use. For more on AOAI’s best practices when employing foundations models for scripts and applications:
|
||||
- [Blog post on responsible AI features in AOAI that were presented at Ignite 2023](https://techcommunity.microsoft.com/t5/ai-azure-ai-services-blog/announcing-new-ai-safety-amp-responsible-ai-features-in-azure/ba-p/3983686)
|
||||
- [Overview of Responsible AI practices for Azure OpenAI models](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/overview)
|
||||
- [Azure OpenAI Transparency Note](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/transparency-note)
|
||||
- [OpenAI’s Usage policies](https://openai.com/policies/usage-policies)
|
||||
- [Azure OpenAI’s Code of Conduct](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/code-of-conduct)
|
||||
|
||||
Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.
|
||||
|
||||
Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning.
|
||||
|
||||
It is the user’s responsibility to ensure that the use of Agent Lightning complies with relevant data protection regulations and organizational guidelines.
|
||||
|
||||
## LICENSE
|
||||
We use the MIT license.
|
||||
|
||||
## CONTACT
|
||||
We welcome feedback and collaboration from our audience. If you have suggestions, questions, or observe unexpected/offensive behavior in our technology, please contact us at agent-lightning@microsoft.com.
|
||||
|
||||
If the team receives reports of undesired behavior or identifies issues independently, we will update this repository with appropriate mitigations.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
*Last updated: September 6, 2025*
|
||||
*Document version: 1.0*
|
||||
@@ -0,0 +1,120 @@
|
||||
<p align="center">
|
||||
<img src="docs/assets/readme-banner.svg" alt="Agent-lightning-banner" style="width:600px"/>
|
||||
</p>
|
||||
|
||||
# Agent Lightning⚡
|
||||
|
||||
[](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
|
||||
[](https://microsoft.github.io/agent-lightning/)
|
||||
[](https://badge.fury.io/py/agentlightning)
|
||||
[](LICENSE)
|
||||
[](https://deepwiki.com/microsoft/agent-lightning)
|
||||
[](https://discord.gg/RYk7CdvDR7)
|
||||
|
||||
**The absolute trainer to light up AI agents.**
|
||||
|
||||
Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with other users and contributors.
|
||||
|
||||
## ⚡ Core Features
|
||||
|
||||
- Turn your agent into an optimizable beast with **ZERO CODE CHANGE** (almost)! 💤
|
||||
- Build with **ANY** agent framework (LangChain, OpenAI Agent SDK, AutoGen, CrewAI, Microsoft Agent Framework...); or even WITHOUT agent framework (Python OpenAI). You name it! 🤖
|
||||
- **Selectively** optimize one or more agents in a multi-agent system. 🎯
|
||||
- Embraces **Algorithms** like Reinforcement Learning, Automatic Prompt Optimization, Supervised Fine-tuning and more. 🤗
|
||||
|
||||
Read more on our [documentation website](https://microsoft.github.io/agent-lightning/).
|
||||
|
||||
<p align="center">
|
||||
<img src="docs/assets/readme-diff.svg" alt="Agent-Lightning Core Quickstart" style="width:100%"/>
|
||||
</p>
|
||||
|
||||
## ⚡ Installation
|
||||
|
||||
```bash
|
||||
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/ --pre agentlightning
|
||||
```
|
||||
|
||||
Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
|
||||
|
||||
To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).
|
||||
|
||||
## ⚡ 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.
|
||||
- 8/5/2025 [Agent Lightning: Train ANY AI Agents with Reinforcement Learning](https://arxiv.org/abs/2508.03680) arXiv paper.
|
||||
- 7/26/2025 [We discovered an approach to train any AI agent with RL, with (almost) zero code changes.](https://www.reddit.com/r/LocalLLaMA/comments/1m9m670/we_discovered_an_approach_to_train_any_ai_agent/) Reddit.
|
||||
- 6/6/2025 [Agent Lightning - Microsoft Research](https://www.microsoft.com/en-us/research/project/agent-lightning/) Project page.
|
||||
|
||||
## ⚡ Community Projects
|
||||
|
||||
- [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
|
||||
- [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
|
||||
- [Youtu-Agent](https://github.com/TencentCloudADP/Youtu-agent) — Youtu-Agent lets you build and train your agent with ease. Built with [a modified branch](https://github.com/microsoft/agent-lightning/tree/contrib/youtu-agent-lightning) of Agent Lightning, Youtu-Agent has verified up to 128 GPUs RL training on maths/code and search capabilities with steady convergence. Also check [the recipe](https://github.com/TencentCloudADP/youtu-agent/tree/rl/agl) and their blog [*Stop Wrestling with Your Agent RL: How Youtu-Agent Achieved Stable, 128-GPU Scaling Without Breaking a Sweat*](https://spotted-coconut-df8.notion.site/Stop-Wrestling-with-Your-Agent-RL-How-Youtu-Agent-Achieved-Stable-128-GPU-Scaling-Without-Breaking-2ca5e8f089ba80539a98c582b65e0233).
|
||||
|
||||
## ⚡ Architecture
|
||||
|
||||
Agent Lightning keeps the moving parts to a minimum so you can focus on your idea, not the plumbing. Your agent continues to run as usual; you can still use any agent framework you like; you drop in the lightweight `agl.emit_xxx()` helper, or let the tracer collect every prompt, tool call, and reward. Those events become structured spans that flow into the LightningStore, a central hub that keeps tasks, resources, and traces in sync.
|
||||
|
||||
On the other side of the store sits the algorithm you choose, or write yourself. The algorithm reads spans, learns from them, and posts updated resources such as refined prompt templates or new policy weights. The Trainer ties it all together: it streams datasets to runners, ferries resources between the store and the algorithm, and updates the inference engine when improvements land. You can either stop there, or simply let the same loop keep turning.
|
||||
|
||||
No rewrites, no lock-in, just a clear path from first rollout to steady improvement.
|
||||
|
||||
<p align="center">
|
||||
<img src="docs/assets/readme-architecture.svg" alt="Agent-lightning Architecture" style="width:100%"/>
|
||||
</p>
|
||||
|
||||
## ⚡ CI Status
|
||||
|
||||
| Workflow | Status |
|
||||
|----------|--------|
|
||||
| CPU Tests | [](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
|
||||
| Full Tests | [](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
|
||||
| UI Tests | [](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml) |
|
||||
| Examples Integration | [](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml) |
|
||||
| Latest Dependency Compatibility | [](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml) |
|
||||
| Legacy Examples Compatibility | [](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml) |
|
||||
|
||||
## ⚡ Citation
|
||||
|
||||
If you find Agent Lightning useful in your research or projects, please cite our paper:
|
||||
|
||||
```bibtex
|
||||
@misc{luo2025agentlightningtrainai,
|
||||
title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning},
|
||||
author={Xufang Luo and Yuge Zhang and Zhiyuan He and Zilong Wang and Siyun Zhao and Dongsheng Li and Luna K. Qiu and Yuqing Yang},
|
||||
year={2025},
|
||||
eprint={2508.03680},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.AI},
|
||||
url={https://arxiv.org/abs/2508.03680},
|
||||
}
|
||||
```
|
||||
|
||||
## ⚡ Contributing
|
||||
|
||||
This project welcomes contributions and suggestions. Start by reading the [Contributing Guide](docs/community/contributing.md) for recommended contribution points, environment setup, branching conventions, and pull request expectations. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
|
||||
|
||||
When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.
|
||||
|
||||
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
|
||||
|
||||
## ⚡ Trademarks
|
||||
|
||||
This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow [Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general). Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party's policies.
|
||||
|
||||
## ⚡ Responsible AI
|
||||
|
||||
This project has been evaluated and certified to comply with the Microsoft Responsible AI Standard. The team will continue to monitor and maintain the repository, addressing any severe issues, including potential harms, if they arise.
|
||||
|
||||
## ⚡ License
|
||||
|
||||
This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`microsoft/agent-lightning`
|
||||
- 原始仓库：https://github.com/microsoft/agent-lightning
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
<!-- BEGIN MICROSOFT SECURITY.MD V1.0.0 BLOCK -->
|
||||
|
||||
## Security
|
||||
|
||||
Microsoft takes the security of our software products and services seriously, which
|
||||
includes all source code repositories in our GitHub organizations.
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues.**
|
||||
|
||||
For security reporting information, locations, contact information, and policies,
|
||||
please review the latest guidance for Microsoft repositories at
|
||||
[https://aka.ms/SECURITY.md](https://aka.ms/SECURITY.md).
|
||||
|
||||
<!-- END MICROSOFT SECURITY.MD BLOCK -->
|
||||
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
__version__ = "0.3.1"
|
||||
|
||||
from .adapter import *
|
||||
from .algorithm import *
|
||||
from .client import AgentLightningClient, DevTaskLoader # deprecated # type: ignore
|
||||
from .config import *
|
||||
from .emitter import *
|
||||
from .env_var import *
|
||||
from .execution import *
|
||||
from .litagent import *
|
||||
from .llm_proxy import *
|
||||
from .logging import configure_logger # deprecated # type: ignore
|
||||
from .logging import setup as setup_logging # type: ignore
|
||||
from .logging import setup_module as setup_module_logging # type: ignore
|
||||
from .runner import *
|
||||
from .server import AgentLightningServer # deprecated # type: ignore
|
||||
from .store import *
|
||||
from .tracer import *
|
||||
from .trainer import *
|
||||
from .types import *
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .base import Adapter, OtelTraceAdapter, TraceAdapter
|
||||
from .messages import TraceToMessages
|
||||
from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
|
||||
|
||||
__all__ = [
|
||||
"TraceAdapter",
|
||||
"OtelTraceAdapter",
|
||||
"Adapter",
|
||||
"TraceToTripletBase",
|
||||
"TracerTraceToTriplet",
|
||||
"LlmProxyTraceToTriplet",
|
||||
"TraceToMessages",
|
||||
]
|
||||
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Generic, Sequence, TypeVar
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
from agentlightning.types import Span
|
||||
|
||||
T_from = TypeVar("T_from")
|
||||
T_to = TypeVar("T_to")
|
||||
|
||||
|
||||
class Adapter(Generic[T_from, T_to]):
|
||||
"""Base class for synchronous adapters that convert data from one format to another.
|
||||
|
||||
The class defines a minimal protocol so that adapters can be treated like callables while
|
||||
still allowing subclasses to supply the concrete transformation logic.
|
||||
|
||||
!!! note
|
||||
Subclasses must override [`adapt()`][agentlightning.Adapter.adapt] to provide
|
||||
the actual conversion.
|
||||
|
||||
Type Variables:
|
||||
|
||||
T_from: Source data type supplied to the adapter.
|
||||
|
||||
T_to: Target data type produced by the adapter.
|
||||
|
||||
Examples:
|
||||
>>> class IntToStrAdapter(Adapter[int, str]):
|
||||
... def adapt(self, source: int) -> str:
|
||||
... return str(source)
|
||||
...
|
||||
>>> adapter = IntToStrAdapter()
|
||||
>>> adapter(42)
|
||||
'42'
|
||||
"""
|
||||
|
||||
def __call__(self, source: T_from, /) -> T_to:
|
||||
"""Convert the data to the target format.
|
||||
|
||||
This method delegates to [`adapt()`][agentlightning.Adapter.adapt] so that an
|
||||
instance of [`Adapter`][agentlightning.Adapter] can be used like a standard
|
||||
function.
|
||||
|
||||
Args:
|
||||
source: Input data in the source format.
|
||||
|
||||
Returns:
|
||||
Data converted to the target format.
|
||||
"""
|
||||
return self.adapt(source)
|
||||
|
||||
def adapt(self, source: T_from, /) -> T_to:
|
||||
"""Convert the data to the target format.
|
||||
|
||||
Subclasses must override this method with the concrete transformation logic. The base
|
||||
implementation raises `NotImplementedError` to make the requirement explicit.
|
||||
|
||||
Args:
|
||||
source: Input data in the source format.
|
||||
|
||||
Returns:
|
||||
Data converted to the target format.
|
||||
"""
|
||||
raise NotImplementedError("Adapter.adapt() is not implemented")
|
||||
|
||||
|
||||
class OtelTraceAdapter(Adapter[Sequence[ReadableSpan], T_to], Generic[T_to]):
|
||||
"""Base class for adapters that convert OpenTelemetry trace spans into other formats.
|
||||
|
||||
This specialization of [`Adapter`][agentlightning.Adapter] expects a list of
|
||||
`opentelemetry.sdk.trace.ReadableSpan` instances and produces any target format, such as
|
||||
reinforcement learning trajectories, structured logs, or analytics-ready payloads.
|
||||
|
||||
Examples:
|
||||
>>> class TraceToDictAdapter(OtelTraceAdapter[dict]):
|
||||
... def adapt(self, spans: List[ReadableSpan]) -> dict:
|
||||
... return {"count": len(spans)}
|
||||
...
|
||||
>>> adapter = TraceToDictAdapter()
|
||||
>>> adapter([span1, span2])
|
||||
{'count': 2}
|
||||
"""
|
||||
|
||||
|
||||
class TraceAdapter(Adapter[Sequence[Span], T_to], Generic[T_to]):
|
||||
"""Base class for adapters that convert trace spans into other formats.
|
||||
|
||||
This class specializes [`Adapter`][agentlightning.Adapter] for working with
|
||||
[`Span`][agentlightning.Span] instances emitted by Agent Lightning instrumentation.
|
||||
Subclasses receive entire trace slices and return a format suited for the downstream consumer,
|
||||
for example reinforcement learning training data or observability metrics.
|
||||
"""
|
||||
@@ -0,0 +1,270 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING, Any, Dict, Generator, Iterable, List, Optional, Sequence, TypedDict, Union, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from agentlightning.types import Span
|
||||
|
||||
from .base import TraceAdapter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from openai.types.chat import (
|
||||
ChatCompletionFunctionToolParam,
|
||||
ChatCompletionMessageFunctionToolCallParam,
|
||||
ChatCompletionMessageParam,
|
||||
)
|
||||
|
||||
|
||||
class OpenAIMessages(TypedDict):
|
||||
"""OpenAI-style chat messages with optional tool definitions.
|
||||
|
||||
Attributes:
|
||||
messages: Ordered chat messages that describe the conversation.
|
||||
tools: Tool specifications available to the assistant, if any.
|
||||
"""
|
||||
|
||||
messages: List[ChatCompletionMessageParam]
|
||||
tools: Optional[List[ChatCompletionFunctionToolParam]]
|
||||
|
||||
|
||||
class _RawSpanInfo(TypedDict):
|
||||
"""Intermediate representation parsed from a span.
|
||||
|
||||
Attributes:
|
||||
prompt: Prompt messages reconstructed from span attributes.
|
||||
completion: Assistant completions following tool invocations.
|
||||
request: Request payload recorded in the trace.
|
||||
response: Response payload recorded in the trace.
|
||||
tools: Tool call metadata extracted from child spans.
|
||||
"""
|
||||
|
||||
prompt: List[Dict[str, Any]]
|
||||
completion: List[Dict[str, Any]]
|
||||
request: Dict[str, Any]
|
||||
response: Dict[str, Any]
|
||||
tools: List[Dict[str, Any]]
|
||||
|
||||
|
||||
def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
|
||||
"""Convert flattened trace attributes into nested structures.
|
||||
|
||||
Attributes emitted by the tracing pipeline often arrive as dotted paths (for example
|
||||
`gen_ai.prompt.0.role`). This helper groups those keys into nested dictionaries or lists so that
|
||||
downstream processing can operate on structured data.
|
||||
|
||||
Args:
|
||||
data: Flat dictionary whose keys are dotted paths.
|
||||
prefix: Top-level key (for example `gen_ai.prompt`) that determines which attributes are
|
||||
grouped.
|
||||
|
||||
Returns:
|
||||
A nested dictionary (no numeric index detected) or list (numeric indices detected) containing
|
||||
the grouped values.
|
||||
"""
|
||||
result: Union[Dict[str, Any], List[Any]] = {}
|
||||
|
||||
# Collect keys that match the prefix
|
||||
relevant = {k[len(prefix) + 1 :]: v for k, v in data.items() if k.startswith(prefix + ".")}
|
||||
|
||||
# Detect if we have numeric indices (-> list) or not (-> dict)
|
||||
indexed = any(part.split(".")[0].isdigit() for part in relevant.keys())
|
||||
|
||||
if indexed:
|
||||
# Group by index
|
||||
grouped: Dict[int, Dict[str, Any]] = defaultdict(dict)
|
||||
for k, v in relevant.items():
|
||||
parts = k.split(".")
|
||||
if not parts[0].isdigit():
|
||||
continue
|
||||
idx, rest = int(parts[0]), ".".join(parts[1:])
|
||||
grouped[idx][rest] = v
|
||||
# Recursively build
|
||||
result = []
|
||||
for i in sorted(grouped.keys()):
|
||||
result.append(group_genai_dict({f"{prefix}.{rest}": val for rest, val in grouped[i].items()}, prefix))
|
||||
else:
|
||||
# No indices: build dict
|
||||
nested: Dict[str, Any] = defaultdict(dict)
|
||||
for k, v in relevant.items():
|
||||
if "." in k:
|
||||
head, _tail = k.split(".", 1)
|
||||
nested[head][f"{prefix}.{k}"] = v
|
||||
else:
|
||||
result[k] = v
|
||||
# Recurse into nested dicts
|
||||
for head, subdict in nested.items():
|
||||
result[head] = group_genai_dict(subdict, prefix + "." + head)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
|
||||
"""Convert raw trace payloads into OpenAI-style chat messages.
|
||||
|
||||
The function consumes an iterable produced by
|
||||
[`TraceToMessages.adapt()`][agentlightning.TraceToMessages.adapt] and yields
|
||||
structures that match the OpenAI fine-tuning JSONL schema, including tool definitions.
|
||||
|
||||
Args:
|
||||
prompt_completion_list: Raw prompt/completion/tool payloads extracted from a trace.
|
||||
|
||||
Returns:
|
||||
A generator that yields [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages]
|
||||
entries compatible with the OpenAI Functions fine-tuning format.
|
||||
"""
|
||||
|
||||
# Import locally to avoid legacy OpenAI version type import errors
|
||||
from openai.types.chat import (
|
||||
ChatCompletionAssistantMessageParam,
|
||||
ChatCompletionFunctionToolParam,
|
||||
ChatCompletionMessageFunctionToolCallParam,
|
||||
ChatCompletionMessageParam,
|
||||
)
|
||||
|
||||
for pc_entry in prompt_completion_list:
|
||||
messages: List[ChatCompletionMessageParam] = []
|
||||
|
||||
# Extract messages
|
||||
for msg in pc_entry["prompt"]:
|
||||
role = msg["role"]
|
||||
|
||||
if role == "assistant" and "tool_calls" in msg:
|
||||
# Use the tool_calls directly
|
||||
# This branch is usually not used in the wild.
|
||||
tool_calls: List[ChatCompletionMessageFunctionToolCallParam] = [
|
||||
ChatCompletionMessageFunctionToolCallParam(
|
||||
id=call["id"],
|
||||
type="function",
|
||||
function={"name": call["name"], "arguments": call["arguments"]},
|
||||
)
|
||||
for call in msg["tool_calls"]
|
||||
]
|
||||
messages.append(
|
||||
ChatCompletionAssistantMessageParam(role="assistant", content=None, tool_calls=tool_calls)
|
||||
)
|
||||
else:
|
||||
# Normal user/system/tool content
|
||||
message = cast(
|
||||
ChatCompletionMessageParam,
|
||||
TypeAdapter(ChatCompletionMessageParam).validate_python(
|
||||
dict(role=role, content=msg.get("content", ""), tool_call_id=msg.get("tool_call_id", None))
|
||||
),
|
||||
)
|
||||
messages.append(message)
|
||||
|
||||
# Extract completions (assistant outputs after tool responses)
|
||||
for comp in pc_entry["completion"]:
|
||||
if comp.get("role") == "assistant":
|
||||
content = comp.get("content")
|
||||
if pc_entry["tools"]:
|
||||
tool_calls = [
|
||||
ChatCompletionMessageFunctionToolCallParam(
|
||||
id=tool["call"]["id"],
|
||||
type=tool["call"]["type"],
|
||||
function={"name": tool["name"], "arguments": tool["parameters"]},
|
||||
)
|
||||
for tool in pc_entry["tools"]
|
||||
]
|
||||
messages.append(
|
||||
ChatCompletionAssistantMessageParam(role="assistant", content=content, tool_calls=tool_calls)
|
||||
)
|
||||
else:
|
||||
messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=content))
|
||||
|
||||
# Build tools definitions (if available)
|
||||
if "functions" in pc_entry["request"]:
|
||||
tools = [
|
||||
ChatCompletionFunctionToolParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": fn["name"],
|
||||
"description": fn.get("description", ""),
|
||||
"parameters": (
|
||||
json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
|
||||
),
|
||||
},
|
||||
)
|
||||
for fn in pc_entry["request"]["functions"]
|
||||
]
|
||||
yield OpenAIMessages(messages=messages, tools=tools)
|
||||
else:
|
||||
yield OpenAIMessages(messages=messages, tools=None)
|
||||
|
||||
|
||||
class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
|
||||
"""Convert trace spans into OpenAI-compatible conversation messages.
|
||||
|
||||
The adapter reconstructs prompts, completions, tool calls, and function definitions from
|
||||
`gen_ai.*` span attributes. The resulting objects match the JSONL structure expected by the
|
||||
OpenAI fine-tuning pipeline.
|
||||
|
||||
!!! warning
|
||||
The adapter assumes all spans share a common trace and that tool call spans are direct
|
||||
children of the associated completion span.
|
||||
"""
|
||||
|
||||
def get_tool_calls(self, completion: Span, all_spans: Sequence[Span], /) -> Iterable[Dict[str, Any]]:
|
||||
"""Yield tool call payloads for a completion span.
|
||||
|
||||
Args:
|
||||
completion: The completion span whose descendants should be inspected.
|
||||
all_spans: The complete span list belonging to the trace.
|
||||
|
||||
Yields:
|
||||
Dictionaries describing tool calls with identifiers, names, and arguments.
|
||||
|
||||
Raises:
|
||||
ValueError: If a candidate tool span cannot be converted into a dictionary.
|
||||
"""
|
||||
# Get all the spans that are children of the completion span
|
||||
children = [span for span in all_spans if span.parent_id == completion.span_id]
|
||||
# Get the tool calls from the children
|
||||
for maybe_tool_call in children:
|
||||
tool_call = group_genai_dict(maybe_tool_call.attributes, "tool")
|
||||
if not isinstance(tool_call, dict):
|
||||
raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
|
||||
if tool_call:
|
||||
yield tool_call
|
||||
|
||||
def adapt(self, source: Sequence[Span], /) -> List[OpenAIMessages]:
|
||||
"""Transform trace spans into OpenAI chat payloads.
|
||||
|
||||
Args:
|
||||
source: Spans containing `gen_ai.*` attributes emitted by the tracing pipeline.
|
||||
|
||||
Returns:
|
||||
A list of [`OpenAIMessages`][agentlightning.adapter.messages.OpenAIMessages] entries that
|
||||
capture prompts, completions, tools, and metadata.
|
||||
"""
|
||||
raw_prompt_completions: List[_RawSpanInfo] = []
|
||||
|
||||
for span in source:
|
||||
attributes = {k: v for k, v in span.attributes.items()}
|
||||
|
||||
# Get all related information from the trace span
|
||||
prompt = group_genai_dict(attributes, "gen_ai.prompt") or []
|
||||
completion = group_genai_dict(attributes, "gen_ai.completion") or []
|
||||
request = group_genai_dict(attributes, "gen_ai.request") or {}
|
||||
response = group_genai_dict(attributes, "gen_ai.response") or {}
|
||||
if not isinstance(prompt, list):
|
||||
raise ValueError(f"Extracted prompt from trace is not a list: {prompt}")
|
||||
if not isinstance(completion, list):
|
||||
raise ValueError(f"Extracted completion from trace is not a list: {completion}")
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError(f"Extracted request from trace is not a dict: {request}")
|
||||
if not isinstance(response, dict):
|
||||
raise ValueError(f"Extracted response from trace is not a dict: {response}")
|
||||
if prompt or completion or request or response:
|
||||
tools = list(self.get_tool_calls(span, source)) or []
|
||||
raw_prompt_completions.append(
|
||||
_RawSpanInfo(
|
||||
prompt=prompt or [], completion=completion, request=request, response=response, tools=tools
|
||||
)
|
||||
)
|
||||
|
||||
return list(convert_to_openai_messages(raw_prompt_completions))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .base import Algorithm
|
||||
from .decorator import algo
|
||||
from .fast import Baseline, FastAlgorithm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .apo import APO as APOType
|
||||
from .verl import VERL as VERLType
|
||||
|
||||
__all__ = ["Algorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
|
||||
|
||||
# Shortcuts for usages like algo.APO(...)
|
||||
|
||||
|
||||
def APO(*args: Any, **kwargs: Any) -> APOType[Any]:
|
||||
from .apo import APO as APOImplementation
|
||||
|
||||
return APOImplementation(*args, **kwargs)
|
||||
|
||||
|
||||
def VERL(*args: Any, **kwargs: Any) -> VERLType:
|
||||
from .verl import VERL as VERLImplementation
|
||||
|
||||
return VERLImplementation(*args, **kwargs)
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .apo import APO
|
||||
|
||||
__all__ = ["APO"]
|
||||
@@ -0,0 +1,895 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""
|
||||
APO with textual gradients that read rollout spans and outputs to modify the prompt.
|
||||
|
||||
- algo: beam search with span-aware textual gradients -> apply_edit via LLM
|
||||
- rollout: same pattern as your example, but task is a dict (T_task)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Counter,
|
||||
Dict,
|
||||
Generic,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Set,
|
||||
Tuple,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
import poml
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from agentlightning.adapter.messages import TraceToMessages
|
||||
from agentlightning.algorithm.base import Algorithm
|
||||
from agentlightning.algorithm.utils import batch_iter_over_dataset, with_llm_proxy, with_store
|
||||
from agentlightning.reward import find_final_reward
|
||||
from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
|
||||
|
||||
class RolloutResultForAPO(TypedDict):
|
||||
"""This must be all JSON serializable to be processable by POML."""
|
||||
|
||||
status: RolloutStatus
|
||||
final_reward: Optional[float]
|
||||
spans: List[Dict[str, Any]]
|
||||
messages: List[Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VersionedPromptTemplate:
|
||||
version: str
|
||||
prompt_template: PromptTemplate
|
||||
score: Optional[float] = None
|
||||
|
||||
|
||||
GRADIENT_PROMPT_FILES = [
|
||||
Path(__file__).parent / "prompts" / "text_gradient_variant01.poml",
|
||||
Path(__file__).parent / "prompts" / "text_gradient_variant02.poml",
|
||||
Path(__file__).parent / "prompts" / "text_gradient_variant03.poml",
|
||||
]
|
||||
|
||||
APPLY_EDIT_PROMPT_FILES = [
|
||||
Path(__file__).parent / "prompts" / "apply_edit_variant01.poml",
|
||||
Path(__file__).parent / "prompts" / "apply_edit_variant02.poml",
|
||||
]
|
||||
|
||||
|
||||
class APO(Algorithm, Generic[T_task]):
|
||||
"""Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
|
||||
|
||||
APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
|
||||
to improve prompts through a beam search process. It evaluates prompts on rollouts,
|
||||
computes critiques based on the results, and applies edits to generate improved prompts.
|
||||
|
||||
The algorithm operates in rounds, where each round:
|
||||
|
||||
1. Samples parent prompts from the current beam
|
||||
2. Generates new prompts by computing textual gradients and applying edits
|
||||
3. Evaluates all candidates on a validation set
|
||||
4. Selects the top-k prompts for the next round
|
||||
|
||||
Based on the ideas from:
|
||||
|
||||
- [ProTeGi](https://aclanthology.org/2023.emnlp-main.494.pdf)
|
||||
- [TextGrad](https://github.com/zou-group/textgrad)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
async_openai_client: AsyncOpenAI,
|
||||
*,
|
||||
gradient_model: str = "gpt-5-mini",
|
||||
apply_edit_model: str = "gpt-4.1-mini",
|
||||
diversity_temperature: float = 1.0,
|
||||
gradient_batch_size: int = 4,
|
||||
val_batch_size: int = 16,
|
||||
beam_width: int = 4,
|
||||
branch_factor: int = 4,
|
||||
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,
|
||||
):
|
||||
"""
|
||||
Initialize the APO algorithm with configuration parameters.
|
||||
|
||||
Args:
|
||||
async_openai_client: AsyncOpenAI client for making LLM API calls.
|
||||
gradient_model: Model name for computing textual gradients (critiques).
|
||||
apply_edit_model: Model name for applying edits based on critiques.
|
||||
diversity_temperature: Temperature parameter for LLM calls to control diversity.
|
||||
gradient_batch_size: Number of rollout results to sample for gradient computation.
|
||||
val_batch_size: Number of validation examples to use for evaluation.
|
||||
beam_width: Number of top-scoring prompts to keep in the beam at each round.
|
||||
branch_factor: Number of new prompt candidates to generate from each parent prompt
|
||||
by applying textual gradient edits. This controls the expansion of the search tree.
|
||||
beam_rounds: Number of beam search rounds to perform.
|
||||
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
|
||||
self.apply_edit_model = apply_edit_model
|
||||
self.diversity_temperature = diversity_temperature
|
||||
self.gradient_batch_size = gradient_batch_size
|
||||
self.val_batch_size = val_batch_size
|
||||
self.beam_width = beam_width
|
||||
self.branch_factor = branch_factor
|
||||
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")
|
||||
self._history_best_version: Optional[str] = None
|
||||
|
||||
self._version_counter: int = 0
|
||||
|
||||
self._poml_trace = _poml_trace
|
||||
|
||||
def _create_versioned_prompt(
|
||||
self,
|
||||
prompt_template: PromptTemplate,
|
||||
*,
|
||||
score: Optional[float] = None,
|
||||
) -> VersionedPromptTemplate:
|
||||
"""
|
||||
Wrap a prompt template with a new monotonically increasing version identifier.
|
||||
"""
|
||||
version = f"v{self._version_counter}"
|
||||
self._version_counter += 1
|
||||
return VersionedPromptTemplate(version=version, prompt_template=prompt_template, score=score)
|
||||
|
||||
def _format_log_prefix(
|
||||
self,
|
||||
*,
|
||||
round_num: Optional[int] = None,
|
||||
beam_idx: Optional[int] = None,
|
||||
branch_idx: Optional[int] = None,
|
||||
prompt_version: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Construct the standardized log prefix.
|
||||
"""
|
||||
parts: List[str] = []
|
||||
if round_num is not None:
|
||||
parts.append(f"Round {round_num:02d}")
|
||||
if beam_idx is not None:
|
||||
parts.append(f"Beam {beam_idx:02d}")
|
||||
if branch_idx is not None:
|
||||
parts.append(f"Branch {branch_idx:02d}")
|
||||
if prompt_version is not None:
|
||||
parts.append(f"Prompt {prompt_version}")
|
||||
if not parts:
|
||||
return ""
|
||||
return f"[{' | '.join(parts)}]"
|
||||
|
||||
def _log(self, level: int, message: str, *, prefix: Optional[str] = None) -> None:
|
||||
"""
|
||||
Log a message with an optional standardized prefix.
|
||||
"""
|
||||
effective_prefix = prefix
|
||||
if effective_prefix:
|
||||
logger.log(level, f"{effective_prefix} {message}")
|
||||
else:
|
||||
logger.log(level, message)
|
||||
|
||||
def get_seed_prompt_template(self) -> Tuple[str, PromptTemplate]:
|
||||
"""
|
||||
Extract the initial prompt template from the algorithm's resources.
|
||||
|
||||
Returns:
|
||||
A tuple of (resource_name, prompt_template) representing the seed prompt.
|
||||
|
||||
Raises:
|
||||
ValueError: If initial_resources is not set or no PromptTemplate is found.
|
||||
"""
|
||||
initial_resources = self.get_initial_resources()
|
||||
if initial_resources is None:
|
||||
raise ValueError(
|
||||
"initial_resources are not set for APO algorithm. "
|
||||
"Use algorithm.set_initial_resources() to set initial resources or set it in Trainer()"
|
||||
)
|
||||
for name, resource in initial_resources.items():
|
||||
if isinstance(resource, PromptTemplate):
|
||||
return name, resource
|
||||
raise ValueError("No prompt template resource found in initial_resources")
|
||||
|
||||
def get_adapter(self) -> TraceToMessages:
|
||||
"""
|
||||
Get the adapter for converting spans to messages.
|
||||
|
||||
Returns:
|
||||
The TraceToMessages instance for this algorithm.
|
||||
|
||||
Raises:
|
||||
ValueError: If the adapter is not a TraceToMessages.
|
||||
"""
|
||||
adapter = super().get_adapter()
|
||||
if not isinstance(adapter, TraceToMessages):
|
||||
raise ValueError("Adapter must be a TraceToMessages for APO algorithm")
|
||||
return adapter
|
||||
|
||||
def get_best_prompt(self) -> PromptTemplate:
|
||||
"""
|
||||
Retrieve the best prompt discovered during optimization.
|
||||
|
||||
Returns:
|
||||
The prompt template with the highest validation score found so far.
|
||||
|
||||
Raises:
|
||||
ValueError: If no best prompt has been found yet (run() not called).
|
||||
"""
|
||||
if self._history_best_prompt is None:
|
||||
raise ValueError("No best prompt found")
|
||||
return self._history_best_prompt
|
||||
|
||||
async def compute_textual_gradient(
|
||||
self,
|
||||
current_prompt: VersionedPromptTemplate,
|
||||
rollout_results: List[RolloutResultForAPO],
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Compute a textual gradient (critique) for the current prompt based on rollout results.
|
||||
|
||||
This method samples rollout results, sends them to an LLM along with the current prompt,
|
||||
and generates a critique describing how the prompt could be improved.
|
||||
|
||||
Args:
|
||||
current_prompt: The prompt template to critique.
|
||||
rollout_results: List of rollout results containing spans, messages, and rewards.
|
||||
|
||||
Returns:
|
||||
A textual critique generated by the LLM, or None if generation fails.
|
||||
"""
|
||||
tg_template = random.choice(self.gradient_prompt_files)
|
||||
|
||||
if len(rollout_results) < self.gradient_batch_size:
|
||||
self._log(
|
||||
logging.WARNING,
|
||||
f"Only {len(rollout_results)} rollouts available, but {self.gradient_batch_size} are needed. Using all rollouts.",
|
||||
prefix=prefix,
|
||||
)
|
||||
sampled_rollout_results = rollout_results
|
||||
else:
|
||||
sampled_rollout_results = random.sample(rollout_results, self.gradient_batch_size)
|
||||
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Gradient will be computed with {self.gradient_model} for {len(sampled_rollout_results)} rollouts with template: {tg_template.name}",
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
tg_msg = poml.poml( # type: ignore
|
||||
tg_template,
|
||||
context={
|
||||
"experiments": sampled_rollout_results,
|
||||
"prompt_template": current_prompt.prompt_template.template,
|
||||
},
|
||||
format="openai_chat",
|
||||
)
|
||||
self._log(
|
||||
logging.DEBUG,
|
||||
f"Gradient computed with {self.gradient_model} prompt: {tg_msg}",
|
||||
prefix=prefix,
|
||||
)
|
||||
critique_response = await self.async_openai_client.chat.completions.create(
|
||||
model=self.gradient_model,
|
||||
messages=tg_msg["messages"], # type: ignore
|
||||
temperature=self.diversity_temperature,
|
||||
)
|
||||
critique_text = critique_response.choices[0].message.content
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Gradient computed with {self.gradient_model} has result: {critique_text}",
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
return critique_text
|
||||
|
||||
async def textual_gradient_and_apply_edit(
|
||||
self,
|
||||
current_prompt: VersionedPromptTemplate,
|
||||
rollout: List[RolloutResultForAPO],
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Generate an improved prompt by computing a textual gradient and applying an edit.
|
||||
|
||||
This is the main optimization step that:
|
||||
|
||||
1. Computes a critique (textual gradient) based on rollout performance
|
||||
2. Uses another LLM to apply the critique and generate an improved prompt
|
||||
|
||||
Args:
|
||||
current_prompt: The current prompt template to improve.
|
||||
rollout: List of rollout results to base the critique on.
|
||||
|
||||
Returns:
|
||||
The improved prompt text, or the original prompt if gradient computation fails.
|
||||
"""
|
||||
# 1) Critique
|
||||
critique_text = await self.compute_textual_gradient(
|
||||
current_prompt,
|
||||
rollout,
|
||||
prefix=prefix,
|
||||
)
|
||||
if not critique_text:
|
||||
self._log(
|
||||
logging.ERROR,
|
||||
"Failed to compute critique for prompt.",
|
||||
prefix=prefix,
|
||||
)
|
||||
return current_prompt.prompt_template.template
|
||||
|
||||
# 2) Apply edit
|
||||
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}",
|
||||
prefix=prefix,
|
||||
)
|
||||
ae_msg = poml.poml( # type: ignore
|
||||
ae_template,
|
||||
context={
|
||||
"prompt_template": current_prompt.prompt_template.template,
|
||||
"critique": critique_text,
|
||||
},
|
||||
format="openai_chat",
|
||||
)
|
||||
|
||||
ae_response = await self.async_openai_client.chat.completions.create(
|
||||
model=self.apply_edit_model,
|
||||
messages=ae_msg["messages"], # type: ignore
|
||||
temperature=self.diversity_temperature,
|
||||
)
|
||||
new_prompt = ae_response.choices[0].message.content
|
||||
if new_prompt:
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Edit generated by {self.apply_edit_model}: {new_prompt[:50]}...",
|
||||
prefix=prefix,
|
||||
)
|
||||
return new_prompt
|
||||
|
||||
@with_store
|
||||
async def get_rollout_results(
|
||||
self,
|
||||
store: LightningStore,
|
||||
rollout: List[Rollout],
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
) -> List[RolloutResultForAPO]:
|
||||
"""
|
||||
Convert completed rollouts to APO-compatible result format.
|
||||
|
||||
Fetches spans for each rollout, adapts them to messages, and packages them
|
||||
with rewards and status information for gradient computation.
|
||||
|
||||
Args:
|
||||
rollout: List of completed rollout metadata.
|
||||
|
||||
Returns:
|
||||
List of rollout results formatted for APO processing.
|
||||
"""
|
||||
rollout_results: List[RolloutResultForAPO] = []
|
||||
adapter = self.get_adapter()
|
||||
for r in rollout:
|
||||
spans = await store.query_spans(r.rollout_id)
|
||||
messages = adapter.adapt(spans)
|
||||
rollout_result = RolloutResultForAPO(
|
||||
status=r.status,
|
||||
final_reward=find_final_reward(spans),
|
||||
spans=[span.model_dump() for span in spans],
|
||||
messages=messages,
|
||||
)
|
||||
self._log(
|
||||
logging.DEBUG,
|
||||
f"Rollout result for {r.rollout_id}: status {rollout_result['status']} with final reward {rollout_result['final_reward']}. "
|
||||
f"{len(rollout_result['spans'])} spans and {len(rollout_result['messages'])} messages.",
|
||||
prefix=prefix,
|
||||
)
|
||||
rollout_results.append(rollout_result)
|
||||
return rollout_results
|
||||
|
||||
async def evaluate_prompt_on_batch(
|
||||
self,
|
||||
prompt: VersionedPromptTemplate,
|
||||
resource_name: str,
|
||||
dataset: Sequence[T_task],
|
||||
mode: RolloutMode,
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
) -> Tuple[List[RolloutResultForAPO], float]:
|
||||
"""
|
||||
Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.
|
||||
|
||||
This method:
|
||||
|
||||
1. Adds the prompt as a named resource to the store
|
||||
2. Enqueues rollouts for each task in the dataset
|
||||
3. Waits for rollouts to complete (with timeout)
|
||||
4. Computes and returns the average reward
|
||||
|
||||
Args:
|
||||
prompt: The prompt template string to evaluate.
|
||||
resource_name: The name to register the prompt under in the store.
|
||||
dataset: Sequence of tasks to evaluate the prompt on.
|
||||
mode: Rollout mode ("train" or "val") for logging/tracking.
|
||||
|
||||
Returns:
|
||||
A tuple of (rollout_results, average_reward) where rollout_results contains
|
||||
detailed information for each rollout and average_reward is the mean final reward.
|
||||
"""
|
||||
store = self.get_store()
|
||||
preview = prompt.prompt_template.template[:50]
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f'Evaluating prompt "{preview}..." on {len(dataset)} tasks in {mode} mode',
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
# Install prompt as named resource
|
||||
resources: NamedResources = {resource_name: prompt.prompt_template}
|
||||
resource_update = await store.update_resources(prompt.version, resources)
|
||||
|
||||
rollout_ids: List[str] = []
|
||||
for t in dataset:
|
||||
r = await store.enqueue_rollout(input=t, mode=mode, resources_id=resource_update.resources_id)
|
||||
rollout_ids.append(r.rollout_id)
|
||||
|
||||
deadline = time.time() + self.rollout_batch_timeout
|
||||
finished: List[Rollout] = []
|
||||
while time.time() < deadline:
|
||||
finished = await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)
|
||||
if len(finished) >= len(rollout_ids):
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"All {len(rollout_ids)} rollouts finished within timeout.",
|
||||
prefix=prefix,
|
||||
)
|
||||
break
|
||||
else:
|
||||
self._log(
|
||||
logging.DEBUG,
|
||||
f"Only {len(finished)} rollouts finished within timeout. Waiting for remaining {len(rollout_ids) - len(finished)} rollouts.",
|
||||
prefix=prefix,
|
||||
)
|
||||
# Sleep to avoid busy-waiting
|
||||
await asyncio.sleep(2.0)
|
||||
|
||||
rollout_results = await self.get_rollout_results(
|
||||
finished,
|
||||
prefix=prefix,
|
||||
)
|
||||
final_rewards = [rr["final_reward"] for rr in rollout_results]
|
||||
|
||||
avg = float(sum([r or 0.0 for r in final_rewards]) / max(1, len(final_rewards)))
|
||||
status_counter = Counter([rr["status"] for rr in rollout_results])
|
||||
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Evaluated {len(rollout_results)} rollouts. Statuses: {status_counter}. Rewards: {final_rewards}, average is {avg}",
|
||||
prefix=prefix,
|
||||
)
|
||||
return rollout_results, avg
|
||||
|
||||
def _initialize_beam(
|
||||
self,
|
||||
train_dataset: Optional[Dataset[T_task]],
|
||||
val_dataset: Optional[Dataset[T_task]],
|
||||
) -> Tuple[str, PromptTemplate, Iterator[Sequence[T_task]], Iterator[Sequence[T_task]]]:
|
||||
"""
|
||||
Initialize the beam search with seed prompt and dataset iterators.
|
||||
|
||||
Args:
|
||||
train_dataset: Dataset for computing gradients.
|
||||
val_dataset: Dataset for evaluating prompts.
|
||||
|
||||
Returns:
|
||||
Tuple of (resource_name, seed_prompt, grad_iterator, val_iterator).
|
||||
|
||||
Raises:
|
||||
ValueError: If either dataset is None.
|
||||
"""
|
||||
resource_name, seed_prompt = self.get_seed_prompt_template()
|
||||
|
||||
if train_dataset is None:
|
||||
raise ValueError("train_dataset is required for APO algorithm")
|
||||
if val_dataset is None:
|
||||
raise ValueError("val_dataset is required for APO algorithm")
|
||||
|
||||
grad_dataset_iterator = batch_iter_over_dataset(train_dataset, self.gradient_batch_size)
|
||||
val_dataset_iterator = batch_iter_over_dataset(val_dataset, self.val_batch_size)
|
||||
|
||||
# Initialize history tracking
|
||||
self._history_best_prompt = seed_prompt
|
||||
self._history_best_score = float("-inf")
|
||||
|
||||
return resource_name, seed_prompt, grad_dataset_iterator, val_dataset_iterator
|
||||
|
||||
def _sample_parent_prompts(
|
||||
self,
|
||||
beam: List[VersionedPromptTemplate],
|
||||
round_num: int,
|
||||
) -> List[Tuple[int, VersionedPromptTemplate]]:
|
||||
"""
|
||||
Sample parent prompts from the current beam for generating new candidates.
|
||||
|
||||
If the beam has fewer prompts than beam_width, replicates existing prompts.
|
||||
Otherwise, randomly samples beam_width prompts.
|
||||
|
||||
Args:
|
||||
beam: Current list of prompt templates in the beam.
|
||||
round_num: Current round number (for logging, 0-indexed).
|
||||
|
||||
Returns:
|
||||
List of parent prompts to generate children from.
|
||||
"""
|
||||
display_round = round_num + 1
|
||||
if len(beam) < self.beam_width:
|
||||
prefix = self._format_log_prefix(round_num=display_round)
|
||||
self._log(
|
||||
logging.WARNING,
|
||||
f"Beam width is currently {self.beam_width}, but only {len(beam)} prompts in beam. Replicating all prompts.",
|
||||
prefix=prefix,
|
||||
)
|
||||
return [(i % len(beam), beam[i % len(beam)]) for i in range(self.beam_width)]
|
||||
|
||||
selected_indices = random.sample(range(len(beam)), self.beam_width)
|
||||
return [(idx, beam[idx]) for idx in selected_indices]
|
||||
|
||||
async def _generate_candidate_prompts(
|
||||
self,
|
||||
parent_prompts: List[Tuple[int, VersionedPromptTemplate]],
|
||||
resource_name: str,
|
||||
grad_dataset_iterator: Iterator[Sequence[T_task]],
|
||||
round_num: int,
|
||||
) -> List[VersionedPromptTemplate]:
|
||||
"""
|
||||
Generate new candidate prompts from parents using textual gradients.
|
||||
|
||||
For each parent prompt, generates branch_factor new candidates by:
|
||||
|
||||
1. Evaluating the parent on a training batch
|
||||
2. Computing textual gradient
|
||||
3. Applying edit to generate improved prompt
|
||||
|
||||
Args:
|
||||
parent_prompts: List of parent prompts to generate children from.
|
||||
resource_name: Name to register prompts under in the store.
|
||||
grad_dataset_iterator: Iterator over training data batches.
|
||||
round_num: Current round number (for logging, 0-indexed).
|
||||
|
||||
Returns:
|
||||
List of newly generated prompt templates.
|
||||
"""
|
||||
display_round = round_num + 1
|
||||
round_prefix = self._format_log_prefix(round_num=display_round)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Applying {self.branch_factor} edits to each of the {len(parent_prompts)} parents based on "
|
||||
"gradients computed on training dataset",
|
||||
prefix=round_prefix,
|
||||
)
|
||||
|
||||
parent_prompts_str = [
|
||||
f"{p.version}:{p.score:.3f}" if p.score is not None else p.version for _, p in parent_prompts
|
||||
]
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Parent prompts: {', '.join(parent_prompts_str)}",
|
||||
prefix=round_prefix,
|
||||
)
|
||||
|
||||
candidates: List[VersionedPromptTemplate] = []
|
||||
used_beam_indices: Set[int] = set()
|
||||
for real_beam_idx, (beam_idx, prompt) in enumerate(parent_prompts):
|
||||
if beam_idx in used_beam_indices:
|
||||
beam_prefix = self._format_log_prefix(
|
||||
round_num=display_round,
|
||||
beam_idx=beam_idx + 1,
|
||||
prompt_version=prompt.version,
|
||||
)
|
||||
self._log(
|
||||
logging.WARNING,
|
||||
"Duplicated beam index found. Might be caused by beam_width too high. "
|
||||
+ f"The real index of this beam is {real_beam_idx + 1}.",
|
||||
prefix=beam_prefix,
|
||||
)
|
||||
else:
|
||||
used_beam_indices.add(beam_idx)
|
||||
for branch_idx in range(self.branch_factor):
|
||||
parent_prefix = self._format_log_prefix(
|
||||
round_num=display_round,
|
||||
beam_idx=beam_idx + 1,
|
||||
branch_idx=branch_idx + 1,
|
||||
prompt_version=prompt.version,
|
||||
)
|
||||
baseline_score = f"{prompt.score:.3f}" if prompt.score is not None else "N/A"
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Use parent prompt {prompt.version} as a baseline to generate a new prompt. Baseline score: {baseline_score}",
|
||||
prefix=parent_prefix,
|
||||
)
|
||||
grad_samples = next(grad_dataset_iterator)
|
||||
rollout_results, _ = await self.evaluate_prompt_on_batch(
|
||||
prompt,
|
||||
resource_name,
|
||||
grad_samples,
|
||||
mode="train",
|
||||
prefix=parent_prefix,
|
||||
)
|
||||
new_prompt = await self.textual_gradient_and_apply_edit(
|
||||
prompt,
|
||||
rollout_results,
|
||||
prefix=parent_prefix,
|
||||
)
|
||||
if not new_prompt:
|
||||
self._log(
|
||||
logging.ERROR,
|
||||
f"Failed to compute edit for prompt: {prompt.prompt_template.template}",
|
||||
prefix=parent_prefix,
|
||||
)
|
||||
continue
|
||||
new_prompt_template = PromptTemplate(template=new_prompt, engine="f-string")
|
||||
versioned_candidate = self._create_versioned_prompt(new_prompt_template)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"New prompt template created from parent {prompt.version}: {versioned_candidate.version}",
|
||||
prefix=parent_prefix,
|
||||
)
|
||||
candidate_prefix = self._format_log_prefix(
|
||||
round_num=display_round, prompt_version=versioned_candidate.version
|
||||
)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"New prompt template created from parent {prompt.version}:\n```\n{new_prompt}\n```",
|
||||
prefix=candidate_prefix,
|
||||
)
|
||||
candidates.append(versioned_candidate)
|
||||
|
||||
return candidates
|
||||
|
||||
async def _evaluate_and_select_beam(
|
||||
self,
|
||||
candidates: List[VersionedPromptTemplate],
|
||||
resource_name: str,
|
||||
val_dataset_iterator: Iterator[Sequence[T_task]],
|
||||
round_num: int,
|
||||
) -> List[VersionedPromptTemplate]:
|
||||
"""
|
||||
Evaluate all candidate prompts on validation data and select top-k for the beam.
|
||||
|
||||
Args:
|
||||
candidates: List of candidate prompts to evaluate.
|
||||
resource_name: Name to register prompts under in the store.
|
||||
val_dataset_iterator: Iterator over validation data batches.
|
||||
round_num: Current round number (for logging, 0-indexed).
|
||||
|
||||
Returns:
|
||||
List of top beam_width prompts sorted by validation score (best first).
|
||||
|
||||
Raises:
|
||||
ValueError: If no candidates remain after evaluation.
|
||||
"""
|
||||
display_round = round_num + 1
|
||||
round_prefix = self._format_log_prefix(round_num=display_round)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Evaluating {len(candidates)} candidates on validation dataset",
|
||||
prefix=round_prefix,
|
||||
)
|
||||
|
||||
val_batch = next(val_dataset_iterator)
|
||||
|
||||
for prompt in candidates:
|
||||
candidate_prefix = self._format_log_prefix(
|
||||
round_num=display_round,
|
||||
prompt_version=prompt.version,
|
||||
)
|
||||
_, score = await self.evaluate_prompt_on_batch(
|
||||
prompt,
|
||||
resource_name,
|
||||
val_batch,
|
||||
mode="val",
|
||||
prefix=candidate_prefix,
|
||||
)
|
||||
prompt.score = score
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Candidate score: {score:.3f}",
|
||||
prefix=candidate_prefix,
|
||||
)
|
||||
|
||||
# Sort by score (descending) and select top beam_width
|
||||
sorted_prompts = [p for p in sorted(candidates, key=lambda x: cast(float, x.score), reverse=True)]
|
||||
selected_prompts = sorted_prompts[: self.beam_width]
|
||||
selected_versions = [
|
||||
f"{prompt.version}:{prompt.score:.3f}" if prompt.score is not None else prompt.version
|
||||
for prompt in selected_prompts
|
||||
]
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Top {len(selected_prompts)} candidates on validation dataset: {selected_versions}",
|
||||
prefix=round_prefix,
|
||||
)
|
||||
|
||||
if len(selected_prompts) == 0:
|
||||
raise ValueError("No beam candidates any more")
|
||||
|
||||
return selected_prompts
|
||||
|
||||
async def _update_best_prompt(
|
||||
self,
|
||||
beam: List[VersionedPromptTemplate],
|
||||
resource_name: str,
|
||||
val_dataset: Dataset[T_task],
|
||||
round_num: int,
|
||||
) -> None:
|
||||
"""
|
||||
Evaluate the best prompt in the beam on the full validation set and update history.
|
||||
|
||||
Args:
|
||||
beam: Current beam of prompts (sorted, best first).
|
||||
resource_name: Name to register prompts under in the store.
|
||||
val_dataset: Full validation dataset.
|
||||
round_num: Current round number (for logging, 0-indexed).
|
||||
"""
|
||||
display_round = round_num + 1
|
||||
best_prompt = beam[0]
|
||||
prefix = self._format_log_prefix(round_num=display_round, prompt_version=best_prompt.version)
|
||||
_, best_score = await self.evaluate_prompt_on_batch(
|
||||
best_prompt,
|
||||
resource_name,
|
||||
cast(Sequence[T_task], val_dataset),
|
||||
mode="val",
|
||||
prefix=prefix,
|
||||
)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Beam leader score: {best_score:.3f}",
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
if best_score > self._history_best_score:
|
||||
prev = self._history_best_score
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Best prompt updated. New best score: {best_score:.3f} (prev: {prev:.3f})",
|
||||
prefix=prefix,
|
||||
)
|
||||
self._history_best_prompt = best_prompt.prompt_template
|
||||
self._history_best_score = best_score
|
||||
self._history_best_version = best_prompt.version
|
||||
else:
|
||||
self._log(
|
||||
logging.WARNING,
|
||||
f"Best prompt not updated. Current score: {best_score:.3f} vs. history best: {self._history_best_score:.3f})",
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
@with_llm_proxy()
|
||||
@with_store
|
||||
async def run(
|
||||
self,
|
||||
store: LightningStore, # Injected by decorator - callers should not provide this parameter
|
||||
llm_proxy: Optional[LLMProxy], # Injected by decorator - callers should not provide this parameter
|
||||
train_dataset: Optional[Dataset[T_task]] = None,
|
||||
val_dataset: Optional[Dataset[T_task]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Execute the APO algorithm to optimize prompts through beam search with textual gradients.
|
||||
|
||||
The algorithm performs iterative prompt optimization over multiple rounds:
|
||||
|
||||
- Each round: samples parent prompts, generates new candidates via textual gradients,
|
||||
evaluates all candidates on validation data, and keeps the top performers
|
||||
- Tracks the historically best prompt across all rounds
|
||||
- Uses different training data samples for each gradient computation to ensure diversity
|
||||
|
||||
Args:
|
||||
train_dataset: Dataset of tasks for computing textual gradients. Required.
|
||||
val_dataset: Dataset of tasks for evaluating and selecting prompts. Required.
|
||||
|
||||
Raises:
|
||||
ValueError: If train_dataset or val_dataset is None, or if resources are not set.
|
||||
"""
|
||||
# Initialize beam search
|
||||
resource_name, seed_prompt, grad_iterator, val_iterator = self._initialize_beam(train_dataset, val_dataset)
|
||||
|
||||
if self._poml_trace:
|
||||
poml.set_trace(trace_dir="pomltrace")
|
||||
|
||||
# Validation datasets are guaranteed to be non-None after initialization
|
||||
assert val_dataset is not None
|
||||
|
||||
# Start with seed prompt in the beam
|
||||
seed_versioned = self._create_versioned_prompt(seed_prompt)
|
||||
beam: List[VersionedPromptTemplate] = [seed_versioned]
|
||||
self._history_best_prompt = seed_prompt
|
||||
self._history_best_version = seed_versioned.version
|
||||
|
||||
# Optionally evaluate seed prompt on validation set to establish baseline
|
||||
if self.run_initial_validation:
|
||||
seed_prefix = self._format_log_prefix(round_num=0, prompt_version=seed_versioned.version)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
"Evaluating seed prompt on validation dataset before optimization...",
|
||||
prefix=seed_prefix,
|
||||
)
|
||||
_, seed_score = await self.evaluate_prompt_on_batch(
|
||||
seed_versioned,
|
||||
resource_name,
|
||||
cast(Sequence[T_task], val_dataset),
|
||||
mode="val",
|
||||
prefix=seed_prefix,
|
||||
)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Seed prompt baseline score: {seed_score:.3f}",
|
||||
prefix=seed_prefix,
|
||||
)
|
||||
self._history_best_prompt = seed_prompt
|
||||
self._history_best_score = seed_score
|
||||
self._history_best_version = seed_versioned.version
|
||||
|
||||
# Run beam search for specified number of rounds
|
||||
for rnd in range(self.beam_rounds):
|
||||
display_round = rnd + 1
|
||||
round_prefix = self._format_log_prefix(round_num=display_round)
|
||||
self._log(
|
||||
logging.INFO,
|
||||
f"Round {display_round}/{self.beam_rounds}...",
|
||||
prefix=round_prefix,
|
||||
)
|
||||
|
||||
# Sample parent prompts from current beam
|
||||
parent_prompts = self._sample_parent_prompts(beam, rnd)
|
||||
|
||||
# Generate new candidate prompts from parents
|
||||
new_candidates = await self._generate_candidate_prompts(parent_prompts, resource_name, grad_iterator, rnd)
|
||||
|
||||
# Combine existing beam with new candidates
|
||||
all_candidates = [*beam, *new_candidates]
|
||||
|
||||
# Evaluate and select top-k prompts for next beam
|
||||
beam = await self._evaluate_and_select_beam(all_candidates, resource_name, val_iterator, rnd)
|
||||
|
||||
# Update historically best prompt if improved
|
||||
await self._update_best_prompt(beam, resource_name, val_dataset, rnd)
|
||||
@@ -0,0 +1,22 @@
|
||||
<poml>
|
||||
<p>Revise the given prompt template using the critique as constraints and improvement guide.</p>
|
||||
<cp caption="Revision Rules">
|
||||
<list listStyle="decimal">
|
||||
<item>Rewrite or restructure the prompt if critique implies it.</item>
|
||||
<item>Explicitly include any requested output format, structure, or word limit, if requested by the critique.</item>
|
||||
<item>Prioritize mechanism-first phrasing: define what to do, then how to do it.</item>
|
||||
<item>Preserve placeholder variables inside curly brackets.</item>
|
||||
</list>
|
||||
</cp>
|
||||
<output-format>
|
||||
Return only the improved prompt template with placeholders intact. Do not include other explanations on how you did it, or headers and introductory texts.
|
||||
</output-format>
|
||||
<human-msg>
|
||||
<cp caption="Prompt Template">
|
||||
<text whiteSpace="pre">{{ prompt_template }}</text>
|
||||
</cp>
|
||||
<cp caption="Critique">
|
||||
<text whiteSpace="pre">{{ critique }}</text>
|
||||
</cp>
|
||||
</human-msg>
|
||||
</poml>
|
||||
@@ -0,0 +1,18 @@
|
||||
<!-- Conservative Edit Prompt -->
|
||||
|
||||
<poml>
|
||||
<p>Revise the prompt to address ONE critique point clearly and effectively. Preserve all variable names in curly-brackets.</p>
|
||||
<p>Do not address more than one critique point. Focus on the single most critical issue.</p>
|
||||
<p>Keep the new prompt close in tone, length, and structure to the original.</p>
|
||||
<output-format>
|
||||
Return only the revised full prompt. Do not include explanations, comparisons, or other text.
|
||||
</output-format>
|
||||
<human-msg>
|
||||
<cp caption="PROMPT" level="3">
|
||||
<text whiteSpace="pre">{{ prompt_template }}</text>
|
||||
</cp>
|
||||
<cp caption="CRITIQUE" level="3">
|
||||
<text whiteSpace="pre">{{ critique }}</text>
|
||||
</cp>
|
||||
</human-msg>
|
||||
</poml>
|
||||
@@ -0,0 +1,18 @@
|
||||
<poml>
|
||||
<p>You optimize a prompt template.</p>
|
||||
<cp caption="Original Prompt Template">
|
||||
<text whiteSpace="pre">{{ prompt_template }}</text>
|
||||
</cp>
|
||||
<cp caption="Experiments with Original Prompt Template">
|
||||
<cp for="experiment in experiments" caption="Experiment {{ loop.index + 1 }}">
|
||||
<p>This experiment has {{ experiment.status }}. It gets a final reward: {{ experiment.final_reward }}</p>
|
||||
<cp caption="Rollout Traces (Chat Messages, Grader Requests included)">
|
||||
<object data="{{ experiment.messages }}" />
|
||||
</cp>
|
||||
</cp>
|
||||
</cp>
|
||||
<cp caption="Your Task">
|
||||
Produce a brief critique listing specific causes for the error or ways to raise reward next time.
|
||||
Return a bullet list with concrete, testable changes (format, constraints, ordering, definitions).
|
||||
</cp>
|
||||
</poml>
|
||||
@@ -0,0 +1,16 @@
|
||||
<poml>
|
||||
<role>You are a prompt engineer.</role>
|
||||
<task>Analyze where the current prompt failed to elicit the right mechanism.</task>
|
||||
<cp caption="Current Prompt Template">
|
||||
<text whiteSpace="pre">{{ prompt_template }}</text>
|
||||
</cp>
|
||||
<cp caption="Sample Runs with Current Prompt Template">
|
||||
<p>The following are the OpenTelemetry spans collected from the sample runs with the current prompt template. They should contain both prompt, responses and rewards.</p>
|
||||
<cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }} Diagnostics">
|
||||
<object for="span in experiment.spans" data="{{ span }}" />
|
||||
</cp>
|
||||
</cp>
|
||||
<output-format>
|
||||
Write 3-5 short bullets titled 'Critique:' focusing on missing constraints, ordering, or formatting.
|
||||
</output-format>
|
||||
</poml>
|
||||
@@ -0,0 +1,107 @@
|
||||
<poml>
|
||||
|
||||
<role>You are an expert prompt engineer.</role>
|
||||
|
||||
<task>Your task is to analyze the prompt and provide a critique of the prompt. Follow the steps below to create the critique.
|
||||
|
||||
<cp caption="1. Structural Issues">
|
||||
<p>These flaws block clarity and logic. Always check them first.</p>
|
||||
|
||||
<list>
|
||||
<item><b>Missing goal</b>: The prompt never defines what success looks like. Ask: <i>Can I summarize its output goal in one line?</i></item>
|
||||
<item><b>Contradictions</b>: Two or more instructions conflict. Search for words like *never*, *always*, *except*, *but also*.</item>
|
||||
<item><b>Circular dependencies</b>: The model is told to do A before B and B before A.</item>
|
||||
<item><b>No stop condition</b>: The prompt doesn’t say when the task is done. Flag any open-ended verbs: <i>explore,</i> <i>analyze further,</i> <i>continue indefinitely.</i></item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="2. Instruction Quality">
|
||||
<p>Examine how the instructions are stated and ordered to ensure clarity and enforceability.</p>
|
||||
<list>
|
||||
<item><b>Vague verbs</b>: Avoid terms like <i>optimize,</i> <i>improve,</i> and <i>ensure.</i> Use precise, measurable instructions.</item>
|
||||
<item><b>Lack of hierarchy</b>: All rules appear equally important, making conflict resolution impossible. Clarify rule precedence.</item>
|
||||
<item><b>Mixed abstraction</b>: High-level policies are interleaved with implementation details. Keep principles separate from step-by-step actions.</item>
|
||||
<item><b>Overlapping scope</b>: Similar instructions appear in several sections with minor changes. Identify and consolidate duplicates.</item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="3. Control and Behavior">
|
||||
<p>Review boundaries on model autonomy, tool use, and communication style.</p>
|
||||
<list>
|
||||
<item><b>No tool limits</b>: Limits on tool calls, retries, or time not specified. Define boundaries for operations.</item>
|
||||
<item><b>Unclear uncertainty handling</b>: Conflicting instructions regarding clarifying uncertainties vs. never asking users. Select one behavior.</item>
|
||||
<item><b>Verbosity confusion</b>: Some parts demand detailed answers, others specify brevity. Highlight and resolve inconsistency.</item>
|
||||
<item><b>Feedback omission</b>: No plan for progress reporting or preamble during multi-step operations.</item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="4. Input and Output Specification">
|
||||
<p>Assess if required data and expected output formats are clearly defined.</p>
|
||||
<list>
|
||||
<item><b>No input defaults</b>: What should happen if a needed value is absent or invalid isn’t explained.</item>
|
||||
<item><b>Output schema missing</b>: Expected response format or sections are not spelled out.</item>
|
||||
<item><b>Format inconsistency</b>: Output style (Markdown, JSON, XML, etc.) shifts mid-prompt. Ensure format requirements are stable.</item>
|
||||
<item><b>No validation</b>: Lacks steps like <i>verify results before submitting</i> or <i>summarize at end.</i></item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="5. Scope and Safety">
|
||||
<p>Ensure prompt actions remain within safe, authorized boundaries.</p>
|
||||
<list>
|
||||
<item><b>Scope creep</b>: Open-ended statements such as <i>feel free to enhance</i> can justify unrelated changes.</item>
|
||||
<item><b>Unsafe actions</b>: Allows deletions or modifications without explicit user approval.</item>
|
||||
<item><b>No error handling</b>: What happens if a tool call fails or data is missing is not addressed.</item>
|
||||
<item><b>User authority ambiguity</b>: Model may act for multiple users or perform irreversible actions without checks.</item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="6. Efficiency and Maintainability">
|
||||
<p>Consider the prompt’s length, redundancy, and future comprehensibility.</p>
|
||||
<list>
|
||||
<item><b>Overexplained</b>: Verbose explanations where concise, numbered steps suffice.</item>
|
||||
<item><b>Redundancy</b>: Similar rules scattered in multiple aliases; centralize and summarize them.</item>
|
||||
<item><b>Hidden assumptions</b>: Implicit defaults (like timezone, language) are not stated.</item>
|
||||
<item><b>Poor auditability</b>: Lacks section markers (e.g., <code><policy></code>, <code><procedure></code>). Structure prompt for easy review.</item>
|
||||
</list>
|
||||
</cp>
|
||||
|
||||
<cp caption="7. Testing Method">
|
||||
<p>Methodical approach for reviewing a prompt:</p>
|
||||
<list>
|
||||
<item>Read the prompt fully; highlight all unclear or contradictory instructions.</item>
|
||||
<item>For each main area, answer:
|
||||
<list listStyle="decimal">
|
||||
<item>What is the intended outcome?</item>
|
||||
<item>What is the stop or completion condition?</item>
|
||||
<item>How are conflicts between rules resolved?</item>
|
||||
<item>What are the explicit limits (tools, run time, tokens)?</item>
|
||||
<item>What should the output format be?</item>
|
||||
</list>
|
||||
</item>
|
||||
<item>Rate each section: <i>clear</i>, <i>incomplete</i>, <i>contradictory</i>, or <i>redundant</i>.</item>
|
||||
<item>Summarize findings under categories: structure, control, scope, format, safety.</item>
|
||||
</list>
|
||||
<p>This method surfaces issues such as ambiguity, contradiction, missing boundaries, and output uncertainty—core failure modes in prompting identified by the GPT-5 prompting guide.</p>
|
||||
</cp>
|
||||
</task>
|
||||
|
||||
<output-format>
|
||||
Respond with a complete analysis and critique of the prompt. Be concise and direct. Less than 350 words.
|
||||
</output-format>
|
||||
|
||||
<human-msg>
|
||||
<cp caption="Prompt">
|
||||
<text whiteSpace="pre">{{ prompt_template }}</text>
|
||||
</cp>
|
||||
<cp caption="Sample Runs of the Prompts (Historical Messages and Rewards)">
|
||||
<cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }}">
|
||||
<cp caption="Overall Status">
|
||||
This run has {{ experiment.status }}. The final score is {{ experiment.final_reward }}.
|
||||
</cp>
|
||||
<cp caption="Messages">
|
||||
<object data="{{ experiment.messages }}" />
|
||||
</cp>
|
||||
</cp>
|
||||
</cp>
|
||||
</human-msg>
|
||||
</poml>
|
||||
@@ -0,0 +1,162 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import weakref
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Awaitable,
|
||||
Optional,
|
||||
Union,
|
||||
)
|
||||
|
||||
from agentlightning.adapter import TraceAdapter
|
||||
from agentlightning.client import AgentLightningClient
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Dataset, NamedResources
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.trainer import Trainer
|
||||
|
||||
|
||||
class Algorithm:
|
||||
"""Algorithm is the strategy, or tuner to train the agent."""
|
||||
|
||||
_trainer_ref: weakref.ReferenceType[Trainer] | None = None
|
||||
_llm_proxy_ref: weakref.ReferenceType["LLMProxy"] | None = None
|
||||
_store: LightningStore | None = None
|
||||
_initial_resources: NamedResources | None = None
|
||||
_adapter_ref: weakref.ReferenceType[TraceAdapter[Any]] | None = None
|
||||
|
||||
def is_async(self) -> bool:
|
||||
"""Return True if the algorithm is asynchronous."""
|
||||
return inspect.iscoroutinefunction(self.run)
|
||||
|
||||
def set_trainer(self, trainer: Trainer) -> None:
|
||||
"""
|
||||
Set the trainer for this algorithm.
|
||||
|
||||
Args:
|
||||
trainer: The Trainer instance that will handle training and validation.
|
||||
"""
|
||||
self._trainer_ref = weakref.ref(trainer)
|
||||
|
||||
def get_trainer(self) -> Trainer:
|
||||
"""
|
||||
Get the trainer for this algorithm.
|
||||
|
||||
Returns:
|
||||
The Trainer instance associated with this agent.
|
||||
"""
|
||||
if self._trainer_ref is None:
|
||||
raise ValueError("Trainer has not been set for this agent.")
|
||||
trainer = self._trainer_ref()
|
||||
if trainer is None:
|
||||
raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
|
||||
return trainer
|
||||
|
||||
def set_llm_proxy(self, llm_proxy: LLMProxy | None) -> None:
|
||||
"""
|
||||
Set the LLM proxy for this algorithm to reuse when available.
|
||||
|
||||
Args:
|
||||
llm_proxy: The LLMProxy instance configured by the trainer, if any.
|
||||
"""
|
||||
self._llm_proxy_ref = weakref.ref(llm_proxy) if llm_proxy is not None else None
|
||||
|
||||
def get_llm_proxy(self) -> Optional[LLMProxy]:
|
||||
"""
|
||||
Retrieve the configured LLM proxy instance, if one has been set.
|
||||
|
||||
Returns:
|
||||
The active LLMProxy instance or None when not configured.
|
||||
"""
|
||||
if self._llm_proxy_ref is None:
|
||||
return None
|
||||
|
||||
llm_proxy = self._llm_proxy_ref()
|
||||
if llm_proxy is None:
|
||||
raise ValueError("LLM proxy reference is no longer valid (object has been garbage collected).")
|
||||
|
||||
return llm_proxy
|
||||
|
||||
def set_adapter(self, adapter: TraceAdapter[Any]) -> None:
|
||||
"""
|
||||
Set the adapter for this algorithm to collect and convert traces.
|
||||
"""
|
||||
self._adapter_ref = weakref.ref(adapter)
|
||||
|
||||
def get_adapter(self) -> TraceAdapter[Any]:
|
||||
"""
|
||||
Retrieve the adapter for this algorithm to communicate with the runners.
|
||||
"""
|
||||
if self._adapter_ref is None:
|
||||
raise ValueError("Adapter has not been set for this algorithm.")
|
||||
adapter = self._adapter_ref()
|
||||
if adapter is None:
|
||||
raise ValueError("Adapter reference is no longer valid (object has been garbage collected).")
|
||||
return adapter
|
||||
|
||||
def set_store(self, store: LightningStore) -> None:
|
||||
"""
|
||||
Set the store for this algorithm to communicate with the runners.
|
||||
|
||||
Store is set directly instead of using weakref because its copy is meant to be
|
||||
maintained throughout the algorithm's lifecycle.
|
||||
"""
|
||||
self._store = store
|
||||
|
||||
def get_store(self) -> LightningStore:
|
||||
"""
|
||||
Retrieve the store for this algorithm to communicate with the runners.
|
||||
"""
|
||||
if self._store is None:
|
||||
raise ValueError("Store has not been set for this algorithm.")
|
||||
return self._store
|
||||
|
||||
def get_initial_resources(self) -> Optional[NamedResources]:
|
||||
"""
|
||||
Get the initial resources for this algorithm.
|
||||
"""
|
||||
return self._initial_resources
|
||||
|
||||
def set_initial_resources(self, resources: NamedResources) -> None:
|
||||
"""
|
||||
Set the initial resources for this algorithm.
|
||||
"""
|
||||
self._initial_resources = resources
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return self.run(*args, **kwargs)
|
||||
|
||||
def run(
|
||||
self,
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> Union[None, Awaitable[None]]:
|
||||
"""Subclasses should implement this method to implement the algorithm.
|
||||
|
||||
Args:
|
||||
train_dataset: The dataset to train on. Not all algorithms require a training dataset.
|
||||
val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.
|
||||
|
||||
Returns:
|
||||
Algorithm should refrain from returning anything. It should just run the algorithm.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement run().")
|
||||
|
||||
def get_client(self) -> AgentLightningClient:
|
||||
"""Get the client to communicate with the algorithm.
|
||||
|
||||
If the algorithm does not require a server-client communication, it can also create a mock client
|
||||
that never communicates with itself.
|
||||
|
||||
Deprecated and will be removed in a future version.
|
||||
|
||||
Returns:
|
||||
The AgentLightningClient instance associated with this algorithm.
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement get_client().")
|
||||
@@ -0,0 +1,264 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Awaitable,
|
||||
Dict,
|
||||
Generic,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
|
||||
from agentlightning.adapter import TraceAdapter
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Dataset, NamedResources
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
|
||||
from .base import Algorithm
|
||||
|
||||
# Algorithm function signature types
|
||||
# We've missed a lot of combinations here.
|
||||
# Let's add them in future.
|
||||
|
||||
|
||||
class AlgorithmFuncSyncFull(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
store: LightningStore,
|
||||
train_dataset: Optional[Dataset[Any]],
|
||||
val_dataset: Optional[Dataset[Any]],
|
||||
llm_proxy: Optional[LLMProxy],
|
||||
adapter: Optional[TraceAdapter[Any]],
|
||||
initial_resources: Optional[NamedResources],
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class AlgorithmFuncSyncOnlyStore(Protocol):
|
||||
def __call__(self, *, store: LightningStore) -> None: ...
|
||||
|
||||
|
||||
class AlgorithmFuncSyncOnlyDataset(Protocol):
|
||||
def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
|
||||
|
||||
|
||||
class AlgorithmFuncAsyncFull(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
store: LightningStore,
|
||||
train_dataset: Optional[Dataset[Any]],
|
||||
val_dataset: Optional[Dataset[Any]],
|
||||
llm_proxy: Optional[LLMProxy],
|
||||
adapter: Optional[TraceAdapter[Any]],
|
||||
initial_resources: Optional[NamedResources],
|
||||
) -> Awaitable[None]: ...
|
||||
|
||||
|
||||
class AlgorithmFuncAsyncOnlyStore(Protocol):
|
||||
def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
|
||||
|
||||
|
||||
class AlgorithmFuncAsyncOnlyDataset(Protocol):
|
||||
def __call__(
|
||||
self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
|
||||
) -> Awaitable[None]: ...
|
||||
|
||||
|
||||
AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
|
||||
|
||||
AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
|
||||
|
||||
|
||||
class AlgorithmFuncSyncFallback(Protocol):
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
|
||||
|
||||
|
||||
class AlgorithmFuncAsyncFallback(Protocol):
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
|
||||
|
||||
|
||||
AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
|
||||
AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
|
||||
|
||||
AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
|
||||
|
||||
|
||||
AsyncFlag = Literal[True, False]
|
||||
AF = TypeVar("AF", bound=AsyncFlag)
|
||||
|
||||
|
||||
class FunctionalAlgorithm(Algorithm, Generic[AF]):
|
||||
"""An algorithm wrapper built from a callable implementation.
|
||||
|
||||
Functional algorithms let you provide an ordinary function instead of
|
||||
subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
|
||||
the callable signature to supply optional dependencies
|
||||
such as the store, adapter, and LLM proxy.
|
||||
"""
|
||||
|
||||
@overload
|
||||
def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
|
||||
|
||||
@overload
|
||||
def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
|
||||
|
||||
def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
|
||||
"""Wrap a function that implements algorithm behaviour.
|
||||
|
||||
Args:
|
||||
algorithm_func: Sync or async callable implementing the algorithm
|
||||
contract. Arguments are detected automatically based on the
|
||||
function signature.
|
||||
"""
|
||||
super().__init__()
|
||||
self._algorithm_func = algorithm_func
|
||||
self._sig = inspect.signature(algorithm_func)
|
||||
self._is_async = inspect.iscoroutinefunction(algorithm_func)
|
||||
|
||||
# Copy function metadata to preserve type hints and other attributes
|
||||
functools.update_wrapper(self, algorithm_func) # type: ignore
|
||||
|
||||
def is_async(self) -> bool:
|
||||
return self._is_async
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self: "FunctionalAlgorithm[Literal[False]]",
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> None: ...
|
||||
|
||||
@overload
|
||||
def run(
|
||||
self: "FunctionalAlgorithm[Literal[True]]",
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> Awaitable[None]: ...
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
return self._algorithm_func(*args, **kwargs) # type: ignore
|
||||
|
||||
def run(
|
||||
self,
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> Union[None, Awaitable[None]]:
|
||||
"""Execute the wrapped function with injected dependencies.
|
||||
|
||||
Args:
|
||||
train_dataset: Optional training dataset passed through when the
|
||||
callable declares a `train_dataset` parameter.
|
||||
val_dataset: Optional validation dataset passed through when the
|
||||
callable declares a `val_dataset` parameter.
|
||||
|
||||
Returns:
|
||||
None for sync callables or an awaitable when the callable is async.
|
||||
|
||||
Raises:
|
||||
TypeError: If a dataset is provided but the function signature does
|
||||
not accept the corresponding argument.
|
||||
"""
|
||||
kwargs: Dict[str, Any] = {}
|
||||
if "store" in self._sig.parameters:
|
||||
kwargs["store"] = self.get_store()
|
||||
if "adapter" in self._sig.parameters:
|
||||
kwargs["adapter"] = self.get_adapter()
|
||||
if "llm_proxy" in self._sig.parameters:
|
||||
kwargs["llm_proxy"] = self.get_llm_proxy()
|
||||
if "initial_resources" in self._sig.parameters:
|
||||
kwargs["initial_resources"] = self.get_initial_resources()
|
||||
if "train_dataset" in self._sig.parameters:
|
||||
kwargs["train_dataset"] = train_dataset
|
||||
elif train_dataset is not None:
|
||||
raise TypeError(
|
||||
f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
|
||||
)
|
||||
if "val_dataset" in self._sig.parameters:
|
||||
kwargs["val_dataset"] = val_dataset
|
||||
elif val_dataset is not None:
|
||||
raise TypeError(
|
||||
f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
|
||||
)
|
||||
# both sync and async functions can be called with the same signature
|
||||
result = self._algorithm_func(**kwargs) # type: ignore[misc]
|
||||
if self._is_async:
|
||||
return cast(Awaitable[None], result)
|
||||
return None
|
||||
|
||||
|
||||
@overload
|
||||
def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
|
||||
|
||||
|
||||
def algo(
|
||||
func: Union[
|
||||
AlgorithmFuncSync,
|
||||
AlgorithmFuncAsync,
|
||||
AlgorithmFuncSyncFallback,
|
||||
AlgorithmFuncAsyncFallback,
|
||||
],
|
||||
) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
|
||||
"""Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].
|
||||
|
||||
The decorator inspects the callable signature to decide which dependencies
|
||||
to inject at runtime, enabling concise algorithm definitions that still
|
||||
leverage the full training runtime.
|
||||
|
||||
Args:
|
||||
func: Function implementing the algorithm logic. May be synchronous or
|
||||
asynchronous. The function can expect all of, or a subset of the following parameters:
|
||||
|
||||
- `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
|
||||
- `train_dataset`: [`Dataset`][agentlightning.Dataset],
|
||||
- `val_dataset`: [`Dataset`][agentlightning.Dataset],
|
||||
- `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
|
||||
- `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
|
||||
- `initial_resources`: [`NamedResources`][agentlightning.NamedResources],
|
||||
|
||||
If the function does not expect a parameter, the wrapper will not inject it into the call.
|
||||
Using `*args` and `**kwargs` will not work and no parameters will be injected.
|
||||
|
||||
Returns:
|
||||
FunctionalAlgorithm that proxies the callable while exposing the
|
||||
`Algorithm` interface.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
from agentlightning.algorithm.decorator import algo
|
||||
|
||||
@algo
|
||||
def batching_algorithm(*, store, train_dataset, val_dataset):
|
||||
for sample in train_dataset:
|
||||
store.enqueue_rollout(input=sample, mode="train")
|
||||
|
||||
@algo
|
||||
async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
|
||||
await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
|
||||
```
|
||||
"""
|
||||
return FunctionalAlgorithm(func)
|
||||
@@ -0,0 +1,250 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, List, Literal, Optional
|
||||
|
||||
from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
|
||||
|
||||
from .base import Algorithm
|
||||
from .utils import with_llm_proxy, with_store
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["FastAlgorithm", "Baseline"]
|
||||
|
||||
|
||||
class FastAlgorithm(Algorithm):
|
||||
"""Base class for lightweight algorithms optimised for developer workflows.
|
||||
|
||||
Fast algorithms prioritise short feedback loops so an agent developer can run
|
||||
small-scale experiments without waiting for long-running training jobs to
|
||||
finish.
|
||||
"""
|
||||
|
||||
|
||||
def _timestamp_to_iso_str(timestamp: float) -> str:
|
||||
return datetime.fromtimestamp(timestamp).isoformat()
|
||||
|
||||
|
||||
class Baseline(FastAlgorithm):
|
||||
"""Reference implementation that streams the full dataset through the rollout queue.
|
||||
|
||||
The baseline algorithm batches task submissions, waits for each rollout to
|
||||
finish, and logs every collected span and reward. It is primarily useful as
|
||||
a smoke test for the platform plumbing rather than a performant trainer.
|
||||
|
||||
The baseline algorithm will auto-start a LLM proxy if one is provided and not yet started.
|
||||
|
||||
Args:
|
||||
n_epochs: Number of dataset passes to execute for both the train and val
|
||||
splits during developer experiments.
|
||||
train_split: Fraction of the concatenated dataset to treat as training
|
||||
data. Must be strictly between 0 and 1.
|
||||
polling_interval: Interval, in seconds, to poll the store for queue
|
||||
depth and rollout completion.
|
||||
max_queue_length: Number of rollouts allowed to wait in the queue before
|
||||
throttling additional submissions.
|
||||
span_verbosity: Level of detail to include when logging span metadata.
|
||||
|
||||
Raises:
|
||||
ValueError: If `train_split` falls outside the `(0, 1)` interval.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
from agentlightning.algorithm.fast import Baseline
|
||||
|
||||
algorithm = Baseline(n_epochs=2, train_split=0.8, span_verbosity="key_values")
|
||||
trainer.fit(algorithm, train_dataset=my_train, val_dataset=my_val)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
n_epochs: int = 1,
|
||||
train_split: float = 0.5,
|
||||
polling_interval: float = 5.0,
|
||||
max_queue_length: int = 4,
|
||||
span_verbosity: Literal["keys", "key_values", "none"] = "keys",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.n_epochs = n_epochs
|
||||
self.train_split = train_split
|
||||
self.polling_interval = polling_interval
|
||||
self.max_queue_length = max_queue_length
|
||||
self.span_verbosity = span_verbosity
|
||||
if not (0.0 < self.train_split < 1.0):
|
||||
raise ValueError("train_split must be between 0 and 1.")
|
||||
|
||||
self._finished_rollout_count = 0
|
||||
|
||||
def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
|
||||
"""Format a span for logging based on the configured verbosity."""
|
||||
if self.span_verbosity == "none":
|
||||
return ""
|
||||
|
||||
prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
|
||||
elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
|
||||
|
||||
msg = (
|
||||
prefix_msg
|
||||
+ f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
|
||||
+ f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
|
||||
+ f"{elapsed} seconds. "
|
||||
)
|
||||
if self.span_verbosity == "key_values":
|
||||
msg += f"Attributes: {span.attributes}"
|
||||
else:
|
||||
msg += f"Attribute keys: {list(span.attributes.keys())}"
|
||||
return msg
|
||||
|
||||
async def _handle_rollout_finish(self, rollout: Rollout) -> None:
|
||||
"""Log attempt metadata and emit adapted traces when a rollout ends."""
|
||||
store = self.get_store()
|
||||
|
||||
rollout_id = rollout.rollout_id
|
||||
rollout_end_time = rollout.end_time or asyncio.get_event_loop().time()
|
||||
logger.info(
|
||||
f"[Rollout {rollout_id}] Finished with status {rollout.status} in {rollout_end_time - rollout.start_time:.2f} seconds."
|
||||
)
|
||||
|
||||
# Logs all the attempts and their corresponding spans
|
||||
attempts = await store.query_attempts(rollout_id)
|
||||
for attempt in attempts:
|
||||
logger.info(
|
||||
"[Rollout %s | Attempt %s] ID: %s. Status: %s. Worker: %s",
|
||||
rollout_id,
|
||||
attempt.sequence_id,
|
||||
attempt.attempt_id,
|
||||
attempt.status,
|
||||
attempt.worker_id,
|
||||
)
|
||||
spans = await store.query_spans(rollout_id=rollout_id)
|
||||
for span in spans:
|
||||
if self.span_verbosity != "none":
|
||||
logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
|
||||
|
||||
# Attempts to adapt the spans using the adapter if provided
|
||||
try:
|
||||
adapter = self.get_adapter()
|
||||
except ValueError:
|
||||
logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
|
||||
adapter = None
|
||||
if adapter is not None:
|
||||
spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
|
||||
transformed_data = adapter.adapt(spans)
|
||||
logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
|
||||
|
||||
async def _enqueue_rollouts(
|
||||
self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
|
||||
) -> None:
|
||||
"""Submit rollouts while respecting the maximum queue length."""
|
||||
store = self.get_store()
|
||||
|
||||
for index in train_indices + val_indices:
|
||||
queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
|
||||
if len(queuing_rollouts) <= 1:
|
||||
# Only enqueue a new rollout when there is at most 1 rollout in the queue.
|
||||
sample = dataset[index]
|
||||
mode = "train" if index in train_indices else "val"
|
||||
rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
|
||||
logger.info(f"[Rollout {rollout.rollout_id}] Enqueued in {mode} mode with sample: {sample}")
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
|
||||
async def _harvest_rollout_spans(self, rollout_id: str):
|
||||
"""Poll rollout status updates until completion and log transitions."""
|
||||
store = self.get_store()
|
||||
last_status: Optional[RolloutStatus] = None
|
||||
|
||||
while True:
|
||||
rollout = await store.get_rollout_by_id(rollout_id)
|
||||
if rollout is not None:
|
||||
if rollout.status in ["succeeded", "failed", "cancelled"]:
|
||||
# Rollout is finished, log all the data.
|
||||
await self._handle_rollout_finish(rollout)
|
||||
# We are done here.
|
||||
self._finished_rollout_count += 1
|
||||
logger.info(f"Finished {self._finished_rollout_count} rollouts.")
|
||||
break
|
||||
|
||||
if last_status != rollout.status:
|
||||
if last_status is not None:
|
||||
logger.info(f"[Rollout {rollout_id}] Status changed to {rollout.status}.")
|
||||
else:
|
||||
logger.info(f"[Rollout {rollout_id}] Status is initialized to {rollout.status}.")
|
||||
last_status = rollout.status
|
||||
|
||||
else:
|
||||
logger.debug(f"[Rollout {rollout_id}] Status is still {rollout.status}.")
|
||||
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
|
||||
@with_llm_proxy()
|
||||
@with_store
|
||||
async def run(
|
||||
self,
|
||||
store: LightningStore, # Injected by decorator - callers should not provide this parameter
|
||||
llm_proxy: Optional[LLMProxy], # Injected by decorator - callers should not provide this parameter
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> None:
|
||||
"""Execute the baseline loop across the provided datasets."""
|
||||
train_dataset_length = len(train_dataset) if train_dataset is not None else 0
|
||||
val_dataset_length = len(val_dataset) if val_dataset is not None else 0
|
||||
if train_dataset_length == 0 and val_dataset_length == 0:
|
||||
logger.error(
|
||||
"MockAlgorithm requires at least one dataset. Provide train_dataset or val_dataset before running."
|
||||
)
|
||||
return
|
||||
|
||||
concatenated_dataset = [train_dataset[i] for i in range(train_dataset_length) if train_dataset is not None] + [
|
||||
val_dataset[i] for i in range(val_dataset_length) if val_dataset is not None
|
||||
]
|
||||
train_indices = list(range(0, train_dataset_length))
|
||||
val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
|
||||
logger.debug(f"Train indices: {train_indices}")
|
||||
logger.debug(f"Val indices: {val_indices}")
|
||||
|
||||
# Currently we only supports a single resource update at the start.
|
||||
initial_resources = self.get_initial_resources()
|
||||
if initial_resources is not None:
|
||||
resource_update = await store.update_resources("default", initial_resources)
|
||||
resources_id = resource_update.resources_id
|
||||
logger.info(f"Initial resources set: {initial_resources}")
|
||||
else:
|
||||
logger.warning("No initial resources provided. Skip initializing resources.")
|
||||
resources_id = None
|
||||
|
||||
for epoch in range(self.n_epochs):
|
||||
harvest_tasks: List[asyncio.Task[None]] = []
|
||||
logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
|
||||
for index in train_indices + val_indices:
|
||||
logger.info(
|
||||
f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
|
||||
)
|
||||
while True:
|
||||
queuing_rollouts = await store.query_rollouts(status_in=["queuing", "requeuing"])
|
||||
if len(queuing_rollouts) <= self.max_queue_length:
|
||||
# Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
|
||||
sample = concatenated_dataset[index]
|
||||
mode = "train" if index in train_indices else "val"
|
||||
rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
|
||||
harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
|
||||
logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
|
||||
break
|
||||
else:
|
||||
# Sleep a bit and try again later.
|
||||
await asyncio.sleep(self.polling_interval)
|
||||
|
||||
# Wait for all harvest tasks to complete
|
||||
logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
|
||||
if len(harvest_tasks) > 0:
|
||||
await asyncio.gather(*harvest_tasks)
|
||||
@@ -0,0 +1,177 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
import random
|
||||
from collections.abc import Coroutine
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Concatenate,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
ParamSpec,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
overload,
|
||||
)
|
||||
|
||||
from agentlightning.types import Dataset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.llm_proxy import LLMProxy
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
from .base import Algorithm
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
T_algo = TypeVar("T_algo", bound="Algorithm")
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def batch_iter_over_dataset(dataset: Dataset[T_task], batch_size: int) -> Iterator[Sequence[T_task]]:
|
||||
"""
|
||||
Create an infinite iterator that yields batches from the dataset.
|
||||
|
||||
When batch_size >= dataset size, yields the entire shuffled dataset repeatedly.
|
||||
When batch_size < dataset size, yields batches of the specified size, reshuffling
|
||||
after each complete pass through the dataset.
|
||||
|
||||
Args:
|
||||
dataset: The dataset to iterate over.
|
||||
batch_size: The desired batch size.
|
||||
|
||||
Yields:
|
||||
Sequences of tasks from the dataset. Each task appears at most once per epoch.
|
||||
"""
|
||||
if batch_size >= len(dataset):
|
||||
while True:
|
||||
dataset_copy = [dataset[i] for i in range(len(dataset))]
|
||||
random.shuffle(dataset_copy)
|
||||
yield dataset_copy
|
||||
|
||||
else:
|
||||
current_batch: List[int] = []
|
||||
while True:
|
||||
indices = list(range(len(dataset)))
|
||||
random.shuffle(indices)
|
||||
for index in indices:
|
||||
if index in current_batch:
|
||||
continue
|
||||
current_batch.append(index)
|
||||
if len(current_batch) == batch_size:
|
||||
yield [dataset[index] for index in current_batch]
|
||||
current_batch = []
|
||||
|
||||
|
||||
def with_store(
|
||||
func: Callable[Concatenate[T_algo, LightningStore, P], Coroutine[Any, Any, R]],
|
||||
) -> Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]]:
|
||||
"""Inject the algorithm's `LightningStore` into coroutine methods.
|
||||
|
||||
The decorator calls `Algorithm.get_store()` once per invocation and passes the
|
||||
resulting store as an explicit argument to the wrapped coroutine. Decorated
|
||||
methods therefore receive the resolved store even when invoked by helper
|
||||
utilities rather than directly by the algorithm.
|
||||
|
||||
Args:
|
||||
func: The coroutine that expects `(self, store, *args, **kwargs)`.
|
||||
|
||||
Returns:
|
||||
A coroutine wrapper that automatically retrieves the store and forwards it
|
||||
to `func`.
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: T_algo, *args: P.args, **kwargs: P.kwargs) -> R:
|
||||
store = self.get_store()
|
||||
return await func(self, store, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@overload
|
||||
def with_llm_proxy(
|
||||
required: Literal[False] = False,
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[Concatenate[T_algo, Optional[LLMProxy], P], Coroutine[Any, Any, R]]],
|
||||
Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
|
||||
]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def with_llm_proxy(
|
||||
required: Literal[True],
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[Concatenate[T_algo, LLMProxy, P], Coroutine[Any, Any, R]]],
|
||||
Callable[Concatenate[T_algo, P], Coroutine[Any, Any, R]],
|
||||
]: ...
|
||||
|
||||
|
||||
def with_llm_proxy(
|
||||
required: bool = False,
|
||||
auto_start: bool = True,
|
||||
) -> Callable[
|
||||
[Callable[..., Coroutine[Any, Any, Any]]],
|
||||
Callable[..., Coroutine[Any, Any, Any]],
|
||||
]:
|
||||
"""Resolve and optionally lifecycle-manage the configured LLM proxy.
|
||||
|
||||
Args:
|
||||
required: When True, raises `ValueError` if the algorithm does not have an
|
||||
[`LLMProxy`][agentlightning.LLMProxy] set. When False, the wrapped coroutine receives
|
||||
`None` if no proxy is available.
|
||||
auto_start: When True, [`LLMProxy.start()`][agentlightning.LLMProxy.start] is invoked if the proxy is not
|
||||
already running before calling `func` and [`LLMProxy.stop()`][agentlightning.LLMProxy.stop] is
|
||||
called afterwards.
|
||||
|
||||
Returns:
|
||||
A decorator that injects the [`LLMProxy`][agentlightning.LLMProxy] (or `None`) as the first
|
||||
argument after `self` and manages automatic startup/shutdown when requested.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[..., Coroutine[Any, Any, Any]],
|
||||
) -> Callable[..., Coroutine[Any, Any, Any]]:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(self: Algorithm, *args: Any, **kwargs: Any) -> Any:
|
||||
llm_proxy = self.get_llm_proxy()
|
||||
|
||||
if required and llm_proxy is None:
|
||||
raise ValueError(
|
||||
"LLM proxy is required but not configured. Call set_llm_proxy() before using this method."
|
||||
)
|
||||
|
||||
auto_started = False
|
||||
if auto_start and llm_proxy is not None:
|
||||
if llm_proxy.is_running():
|
||||
logger.info("Proxy is already running, skipping start")
|
||||
else:
|
||||
logger.info("Starting proxy, managed by the algorithm")
|
||||
await llm_proxy.start()
|
||||
auto_started = True
|
||||
|
||||
try:
|
||||
# At type level, overloads guarantee that if `required=True`
|
||||
# then `func` expects a non-optional LLMProxy.
|
||||
return await func(self, llm_proxy, *args, **kwargs)
|
||||
finally:
|
||||
if auto_started and llm_proxy is not None:
|
||||
logger.info("Stopping proxy, managed by the algorithm")
|
||||
await llm_proxy.stop()
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .interface import VERL
|
||||
|
||||
__all__ = ["VERL"]
|
||||
@@ -0,0 +1,202 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Optional, Type
|
||||
|
||||
from hydra import compose, initialize
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from agentlightning.algorithm.base import Algorithm
|
||||
from agentlightning.client import AgentLightningClient
|
||||
from agentlightning.types import Dataset
|
||||
from agentlightning.verl.entrypoint import run_ppo # type: ignore
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.verl.daemon import AgentModeDaemon
|
||||
from agentlightning.verl.trainer import AgentLightningTrainer
|
||||
|
||||
|
||||
class VERL(Algorithm):
|
||||
"""VERL-powered algorithm that delegates training to the VERL PPO runner.
|
||||
|
||||
!!! warning
|
||||
Advanced customisation currently requires copying the VERL source and
|
||||
modifying it directly. Native hooks for overriding training behaviour
|
||||
will land in a future release.
|
||||
|
||||
Args:
|
||||
config: Dictionary mirroring the overrides passed to the VERL CLI. The
|
||||
overrides are merged with VERL's packaged defaults via Hydra before
|
||||
launching training.
|
||||
trainer_cls: Optional override for the trainer class. Experimental.
|
||||
daemon_cls: Optional override for the daemon class. Experimental.
|
||||
|
||||
!!! note "Trajectory aggregation (experimental)"
|
||||
|
||||
Trajectory-level aggregation merges an entire multi-turn rollout into a single,
|
||||
masked training sample so GPU time is spent once per trajectory rather than N times
|
||||
per turn. Enable it via:
|
||||
|
||||
```python
|
||||
config["agentlightning"]["trace_aggregator"] = {
|
||||
"level": "trajectory",
|
||||
"trajectory_max_prompt_length": 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
|
||||
|
||||
algorithm = VERL(
|
||||
config={
|
||||
"algorithm": {
|
||||
"adv_estimator": "grpo",
|
||||
"use_kl_in_reward": False,
|
||||
},
|
||||
"data": {
|
||||
"train_batch_size": 32,
|
||||
"max_prompt_length": 4096,
|
||||
"max_response_length": 2048,
|
||||
},
|
||||
"actor_rollout_ref": {
|
||||
"rollout": {
|
||||
"tensor_model_parallel_size": 1,
|
||||
"n": 4,
|
||||
"log_prob_micro_batch_size_per_gpu": 4,
|
||||
"multi_turn": {"format": "hermes"},
|
||||
"name": "vllm",
|
||||
"gpu_memory_utilization": 0.6,
|
||||
},
|
||||
"actor": {
|
||||
"ppo_mini_batch_size": 32,
|
||||
"ppo_micro_batch_size_per_gpu": 4,
|
||||
"optim": {"lr": 1e-6},
|
||||
"use_kl_loss": False,
|
||||
"kl_loss_coef": 0.0,
|
||||
"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": 8,
|
||||
"fsdp_config": {"param_offload": True},
|
||||
},
|
||||
"model": {
|
||||
"path": "Qwen/Qwen2.5-1.5B-Instruct",
|
||||
"use_remove_padding": True,
|
||||
"enable_gradient_checkpointing": True,
|
||||
},
|
||||
},
|
||||
"trainer": {
|
||||
"n_gpus_per_node": 1,
|
||||
"val_before_train": True,
|
||||
"critic_warmup": 0,
|
||||
"logger": ["console", "wandb"],
|
||||
"project_name": "AgentLightning",
|
||||
"experiment_name": "calc_x",
|
||||
"nnodes": 1,
|
||||
"save_freq": 64,
|
||||
"test_freq": 32,
|
||||
"total_epochs": 2,
|
||||
},
|
||||
}
|
||||
)
|
||||
trainer.fit(algorithm, train_dataset=my_train_dataset)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: dict[str, Any],
|
||||
trainer_cls: Optional[Type[AgentLightningTrainer]] = None,
|
||||
daemon_cls: Optional[Type[AgentModeDaemon]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Compose the base config exactly like your decorator:
|
||||
with initialize(version_base=None, config_path="pkg://agentlightning/verl"):
|
||||
base_cfg = compose(config_name="config")
|
||||
|
||||
# Merge your dict overrides
|
||||
override_conf = OmegaConf.create(config)
|
||||
# Allow adding new fields
|
||||
OmegaConf.set_struct(base_cfg, False)
|
||||
self.config = OmegaConf.merge(base_cfg, override_conf)
|
||||
self.trainer_cls = trainer_cls
|
||||
self.daemon_cls = daemon_cls
|
||||
|
||||
def run(
|
||||
self,
|
||||
train_dataset: Optional[Dataset[Any]] = None,
|
||||
val_dataset: Optional[Dataset[Any]] = None,
|
||||
) -> None:
|
||||
"""Launch the VERL PPO entrypoint with the configured runtime context.
|
||||
|
||||
Args:
|
||||
train_dataset: Optional dataset forwarded to VERL for training.
|
||||
val_dataset: Optional dataset forwarded to VERL for evaluation.
|
||||
|
||||
Raises:
|
||||
ValueError: If required dependencies such as the store, LLM proxy, or
|
||||
adapter have been garbage-collected when using the V1 execution
|
||||
mode.
|
||||
"""
|
||||
from agentlightning.verl.daemon import AgentModeDaemon
|
||||
from agentlightning.verl.trainer import AgentLightningTrainer
|
||||
|
||||
trainer_cls = self.trainer_cls or AgentLightningTrainer
|
||||
daemon_cls = self.daemon_cls or AgentModeDaemon
|
||||
try:
|
||||
store = self.get_store()
|
||||
except Exception:
|
||||
print("Store is not set. Assuming v0 execution mode.")
|
||||
run_ppo(
|
||||
self.config,
|
||||
train_dataset=train_dataset,
|
||||
val_dataset=val_dataset,
|
||||
store=None,
|
||||
llm_proxy=None,
|
||||
adapter=None,
|
||||
trainer_cls=trainer_cls,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
else:
|
||||
print("Store is set. Assuming v1 execution mode.")
|
||||
llm_proxy = self.get_llm_proxy()
|
||||
adapter = self.get_adapter()
|
||||
run_ppo(
|
||||
self.config,
|
||||
train_dataset=train_dataset,
|
||||
val_dataset=val_dataset,
|
||||
store=store,
|
||||
llm_proxy=llm_proxy,
|
||||
adapter=adapter,
|
||||
trainer_cls=trainer_cls,
|
||||
daemon_cls=daemon_cls,
|
||||
)
|
||||
|
||||
def get_client(self) -> AgentLightningClient:
|
||||
"""Create a client bound to the VERL-managed Agent Lightning server.
|
||||
|
||||
Deprecated:
|
||||
Since v0.2.
|
||||
"""
|
||||
port = self.config.agentlightning.port
|
||||
return AgentLightningClient(endpoint=f"http://localhost:{port}")
|
||||
@@ -0,0 +1,56 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Agent Lightning command line interface entry point."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib
|
||||
import sys
|
||||
from typing import Dict, Iterable, Tuple
|
||||
|
||||
_SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
|
||||
"vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
|
||||
"store": ("agentlightning.cli.store", "Run a LightningStore server."),
|
||||
"prometheus": ("agentlightning.cli.prometheus", "Serve Prometheus metrics from the multiprocess registry."),
|
||||
"agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
|
||||
}
|
||||
|
||||
_DESCRIPTION = "Agent Lightning CLI entry point.\n\nAvailable subcommands:\n" + "\n".join(
|
||||
f" {name:<10}{desc}" for name, (_, desc) in _SUBCOMMANDS.items()
|
||||
)
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
"""Dispatch to the requested Agent Lightning subcommand."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="agl",
|
||||
description=_DESCRIPTION,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("subcommand", choices=_SUBCOMMANDS.keys(), help="Subcommand to run.")
|
||||
parser.add_argument("args", nargs=argparse.REMAINDER, help=argparse.SUPPRESS)
|
||||
|
||||
parsed = parser.parse_args(list(argv) if argv is not None else None)
|
||||
module_name, _ = _SUBCOMMANDS[parsed.subcommand]
|
||||
module = importlib.import_module(module_name)
|
||||
|
||||
entry_point = getattr(module, "main", None)
|
||||
if entry_point is None:
|
||||
parser.error(f"Subcommand '{parsed.subcommand}' does not define a callable 'main'")
|
||||
|
||||
dispatch_args = parsed.args
|
||||
original_argv = sys.argv
|
||||
sys.argv = [f"{parser.prog} {parsed.subcommand}", *dispatch_args]
|
||||
try:
|
||||
result = entry_point(dispatch_args or None)
|
||||
finally:
|
||||
sys.argv = original_argv
|
||||
|
||||
if isinstance(result, int):
|
||||
return result
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,115 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Serve Prometheus metrics from the Agent Lightning multiprocess registry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
from fastapi import FastAPI
|
||||
from prometheus_client import make_asgi_app # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
from agentlightning.logging import setup as setup_logging
|
||||
from agentlightning.utils.metrics import get_prometheus_registry
|
||||
from agentlightning.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def ensure_prometheus_dir() -> str:
|
||||
"""Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
|
||||
|
||||
directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
|
||||
if directory is None:
|
||||
raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
|
||||
|
||||
Path(directory).mkdir(parents=True, exist_ok=True)
|
||||
logger.info("Serving Prometheus multiprocess metrics from %s", directory)
|
||||
return directory
|
||||
|
||||
|
||||
def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
|
||||
"""Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
|
||||
|
||||
Args:
|
||||
metrics_path: URL path to expose the Prometheus metrics endpoint on.
|
||||
|
||||
Returns:
|
||||
A FastAPI application ready to serve metrics.
|
||||
"""
|
||||
|
||||
if not metrics_path.startswith("/"):
|
||||
raise ValueError("metrics_path must start with '/'.")
|
||||
|
||||
normalized_path = metrics_path.rstrip("/")
|
||||
if normalized_path in ("", "/"):
|
||||
raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
|
||||
|
||||
app = FastAPI(title="Agent Lightning Prometheus exporter", docs_url=None, redoc_url=None)
|
||||
metrics_app = make_asgi_app(registry=get_prometheus_registry()) # pyright: ignore[reportUnknownVariableType]
|
||||
app.mount(normalized_path, metrics_app) # pyright: ignore[reportUnknownArgumentType]
|
||||
|
||||
@app.get("/health")
|
||||
async def healthcheck() -> dict[str, str]: # pyright: ignore[reportUnusedFunction]
|
||||
return {"status": "ok"}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
|
||||
parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
|
||||
parser.add_argument(
|
||||
"--metrics-path",
|
||||
default="/v1/prometheus",
|
||||
help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
||||
help="Configure the logging level for the metrics server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--access-log",
|
||||
action="store_true",
|
||||
help="Enable uvicorn access logs. Disabled by default to reduce noise.",
|
||||
)
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
setup_logging(args.log_level)
|
||||
ensure_prometheus_dir()
|
||||
|
||||
try:
|
||||
app = create_prometheus_app(args.metrics_path)
|
||||
except ValueError as exc:
|
||||
logger.error("Failed to configure prometheus app: %s", exc)
|
||||
return 1
|
||||
|
||||
launcher_args = PythonServerLauncherArgs(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level=getattr(logging, args.log_level),
|
||||
access_log=args.access_log,
|
||||
healthcheck_url="/health",
|
||||
)
|
||||
launcher = PythonServerLauncher(app, launcher_args)
|
||||
|
||||
try:
|
||||
asyncio.run(launcher.run_forever())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Received shutdown signal. Stopping Prometheus server.")
|
||||
except RuntimeError as exc:
|
||||
logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,131 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Run a LightningStore server for persistent access from multiple processes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Iterable, List
|
||||
|
||||
from agentlightning import setup_logging
|
||||
from agentlightning.store.client_server import LightningStoreServer
|
||||
from agentlightning.store.memory import InMemoryLightningStore
|
||||
from agentlightning.utils.metrics import (
|
||||
ConsoleMetricsBackend,
|
||||
MetricsBackend,
|
||||
MultiMetricsBackend,
|
||||
PrometheusMetricsBackend,
|
||||
setup_multiprocess_prometheus,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Run a LightningStore server")
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
|
||||
parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
|
||||
parser.add_argument(
|
||||
"--cors-origin",
|
||||
dest="cors_origins",
|
||||
action="append",
|
||||
help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
||||
help="Configure the logging level for the store.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tracker",
|
||||
nargs="+",
|
||||
choices=["prometheus", "console"],
|
||||
help="Enable metrics tracking. Repeat for multiple trackers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n-workers",
|
||||
default=1,
|
||||
type=int,
|
||||
help=(
|
||||
"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."
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
choices=["memory", "mongo"],
|
||||
default="memory",
|
||||
help="Backend to use for the store.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mongo-uri",
|
||||
default="mongodb://localhost:27017/?replicaSet=rs0",
|
||||
help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
|
||||
)
|
||||
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
setup_logging(args.log_level)
|
||||
|
||||
trackers: List[MetricsBackend] = []
|
||||
if args.tracker:
|
||||
if "prometheus" in args.tracker:
|
||||
logger.info("Enabling Prometheus metrics tracking.")
|
||||
if args.n_workers > 1:
|
||||
# This has to be done before prometheus_client is imported
|
||||
setup_multiprocess_prometheus()
|
||||
logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
|
||||
trackers.append(PrometheusMetricsBackend())
|
||||
|
||||
if "console" in args.tracker:
|
||||
logger.info("Enabling console metrics tracking.")
|
||||
trackers.append(ConsoleMetricsBackend())
|
||||
|
||||
if len(trackers) == 0:
|
||||
tracker: MetricsBackend | None = None
|
||||
elif len(trackers) == 1:
|
||||
tracker = trackers[0]
|
||||
else:
|
||||
tracker = MultiMetricsBackend(trackers)
|
||||
|
||||
if args.backend == "memory":
|
||||
store = InMemoryLightningStore(
|
||||
thread_safe=True, # Using thread_safe store for server
|
||||
tracker=tracker,
|
||||
)
|
||||
elif args.backend == "mongo":
|
||||
from agentlightning.store.mongo import MongoLightningStore
|
||||
|
||||
store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
|
||||
else:
|
||||
raise ValueError(f"Invalid backend: {args.backend}")
|
||||
|
||||
if args.n_workers > 1:
|
||||
logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
|
||||
launch_mode = "mp"
|
||||
else:
|
||||
logger.info("Running the server using `asyncio` launch mode.")
|
||||
launch_mode = "asyncio"
|
||||
server = LightningStoreServer(
|
||||
store,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
cors_allow_origins=args.cors_origins,
|
||||
launch_mode=launch_mode,
|
||||
tracker=tracker,
|
||||
n_workers=args.n_workers,
|
||||
)
|
||||
try:
|
||||
asyncio.run(server.run_forever())
|
||||
except RuntimeError as exc:
|
||||
logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
import sys
|
||||
|
||||
from vllm.entrypoints.cli.main import main as vllm_main
|
||||
|
||||
from agentlightning.instrumentation.vllm import instrument_vllm
|
||||
|
||||
instrument_vllm()
|
||||
if argv is not None:
|
||||
original_argv = sys.argv
|
||||
sys.argv = [original_argv[0], *list(argv)]
|
||||
try:
|
||||
vllm_main()
|
||||
finally:
|
||||
sys.argv = original_argv
|
||||
else:
|
||||
vllm_main()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,408 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Utilities for interacting with legacy Agent Lightning servers.
|
||||
|
||||
This module contains compatibility shims that speak the deprecated HTTP
|
||||
interface used by older Agent Lightning deployments. Modern code should prefer
|
||||
the store-based APIs exposed by `agentlightning.store`, but keeping these
|
||||
clients available makes it easier to migrate existing workflows incrementally.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
import urllib.parse
|
||||
import warnings
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import aiohttp
|
||||
import requests
|
||||
|
||||
from .types import NamedResources, ResourcesUpdate, RolloutLegacy, Task, TaskIfAny, TaskInput
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AgentLightningClient:
|
||||
"""Client wrapper for the legacy version-aware Agent Lightning server.
|
||||
|
||||
The client exposes synchronous and asynchronous helpers for polling tasks,
|
||||
retrieving resource bundles, and submitting rollouts. It also maintains a
|
||||
simple in-memory cache keyed by the server-provided resource identifier to
|
||||
avoid redundant network requests.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
[`AgentLightningClient`][agentlightning.client.AgentLightningClient] is part of
|
||||
the legacy client/server stack. New code should rely on the store-based APIs
|
||||
implemented in `agentlightning.store`.
|
||||
|
||||
Attributes:
|
||||
endpoint: Base URL of the Agent Lightning server.
|
||||
poll_interval: Delay in seconds between polling attempts when no task is
|
||||
available.
|
||||
timeout: Timeout in seconds applied to HTTP requests.
|
||||
task_count: Number of tasks claimed during the lifetime of this client.
|
||||
"""
|
||||
|
||||
_next_task_uri = "/task"
|
||||
_resources_uri = "/resources"
|
||||
_latest_resources_uri = "/resources/latest"
|
||||
_report_rollout_uri = "/rollout"
|
||||
|
||||
def __init__(self, endpoint: str, poll_interval: float = 5.0, timeout: float = 10.0):
|
||||
"""Initialize the client.
|
||||
|
||||
Args:
|
||||
endpoint: Root URL of the Agent Lightning server.
|
||||
poll_interval: Seconds to wait between polling attempts.
|
||||
timeout: Seconds before a request to the server is considered timed out.
|
||||
"""
|
||||
warnings.warn(
|
||||
"AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
|
||||
)
|
||||
self.endpoint = endpoint
|
||||
self.task_count = 0
|
||||
self.poll_interval = poll_interval
|
||||
self.timeout = timeout
|
||||
self._resource_cache: Dict[str, ResourcesUpdate] = {} # TODO: mechanism to evict cache
|
||||
self._default_headers = {"X-AgentLightning-Client": "true"}
|
||||
|
||||
async def _request_json_async(self, url: str) -> Optional[Dict[str, Any]]:
|
||||
"""Perform an asynchronous ``GET`` request and parse the JSON payload.
|
||||
|
||||
Args:
|
||||
url: Fully qualified URL to query.
|
||||
|
||||
Returns:
|
||||
Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
|
||||
"""
|
||||
timeout = aiohttp.ClientTimeout(total=self.timeout)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
try:
|
||||
async with session.get(url, headers=self._default_headers) as resp:
|
||||
resp.raise_for_status()
|
||||
return await resp.json()
|
||||
except Exception as e:
|
||||
logger.debug(f"Async GET request failed for {url}: {e}")
|
||||
return None
|
||||
|
||||
async def _post_json_async(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
||||
"""Perform an asynchronous ``POST`` request with a JSON body.
|
||||
|
||||
Args:
|
||||
url: Fully qualified URL that accepts the payload.
|
||||
payload: Dictionary that will be serialized and sent as JSON.
|
||||
|
||||
Returns:
|
||||
Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
|
||||
"""
|
||||
timeout = aiohttp.ClientTimeout(total=self.timeout)
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
try:
|
||||
async with session.post(url, json=payload, headers=self._default_headers) as resp:
|
||||
resp.raise_for_status()
|
||||
return await resp.json()
|
||||
except Exception as e:
|
||||
logger.debug(f"Async POST request failed for {url}: {e}")
|
||||
return None
|
||||
|
||||
async def poll_next_task_async(self) -> Optional[Task]:
|
||||
"""Poll the server asynchronously until a task becomes available.
|
||||
|
||||
Returns:
|
||||
The next [`Task`][agentlightning.Task] exposed by the server,
|
||||
or ``None`` if polling fails.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
|
||||
while True:
|
||||
response = await self._request_json_async(url)
|
||||
if response:
|
||||
task_if_any = TaskIfAny.model_validate(response)
|
||||
if task_if_any.is_available and task_if_any.task:
|
||||
self.task_count += 1
|
||||
logger.info(f"[Task {self.task_count} Received] ID: {task_if_any.task.rollout_id}")
|
||||
return task_if_any.task
|
||||
logger.debug(f"No task available yet. Retrying in {self.poll_interval} seconds...")
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
|
||||
"""Fetch a specific resource bundle by identifier.
|
||||
|
||||
Args:
|
||||
resource_id: Identifier sourced from the task metadata.
|
||||
|
||||
Returns:
|
||||
Cached or freshly downloaded
|
||||
[`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
|
||||
``None`` when the server returns an error.
|
||||
"""
|
||||
if resource_id in self._resource_cache:
|
||||
logger.debug(f"Found resources '{resource_id}' in cache.")
|
||||
return self._resource_cache[resource_id]
|
||||
|
||||
url = urllib.parse.urljoin(self.endpoint, f"{self._resources_uri}/{resource_id}")
|
||||
response = await self._request_json_async(url)
|
||||
if response:
|
||||
resources_update = ResourcesUpdate.model_validate(response)
|
||||
self._resource_cache[resource_id] = resources_update
|
||||
logger.info(f"Fetched and cached resources for ID: {resource_id}")
|
||||
return resources_update
|
||||
return None
|
||||
|
||||
async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
|
||||
"""Fetch the most recent resource bundle advertised by the server.
|
||||
|
||||
Returns:
|
||||
[`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
|
||||
newest version, or ``None`` when unavailable.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
|
||||
response = await self._request_json_async(url)
|
||||
if response:
|
||||
resources_update = ResourcesUpdate.model_validate(response)
|
||||
# Cache this result as well
|
||||
self._resource_cache[resources_update.resources_id] = resources_update
|
||||
return resources_update
|
||||
return None
|
||||
|
||||
async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
|
||||
"""Submit a completed rollout back to the server.
|
||||
|
||||
Args:
|
||||
rollout: Legacy rollout payload produced by the executor.
|
||||
|
||||
Returns:
|
||||
Parsed JSON response returned by the server, or ``None`` when the request fails.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
|
||||
payload = rollout.model_dump(mode="json")
|
||||
return await self._post_json_async(url, payload)
|
||||
|
||||
def _request_json(self, url: str) -> Optional[Dict[str, Any]]:
|
||||
"""Perform a blocking ``GET`` request and parse the JSON payload.
|
||||
|
||||
Args:
|
||||
url: Fully qualified URL to query.
|
||||
|
||||
Returns:
|
||||
Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
|
||||
"""
|
||||
try:
|
||||
response = requests.get(url, timeout=self.timeout, headers=self._default_headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.debug(f"Sync GET request failed for {url}: {e}")
|
||||
return None
|
||||
|
||||
def _post_json(self, url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
||||
"""Perform a blocking ``POST`` request with a JSON payload.
|
||||
|
||||
Args:
|
||||
url: Fully qualified URL that accepts the payload.
|
||||
payload: Dictionary that will be serialized and sent as JSON.
|
||||
|
||||
Returns:
|
||||
Parsed JSON body as a dictionary if the request succeeds; otherwise ``None``.
|
||||
"""
|
||||
try:
|
||||
response = requests.post(url, json=payload, timeout=self.timeout, headers=self._default_headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.debug(f"Sync POST request failed for {url}: {e}")
|
||||
return None
|
||||
|
||||
def poll_next_task(self) -> Optional[Task]:
|
||||
"""Poll the server synchronously until a task becomes available.
|
||||
|
||||
Returns:
|
||||
The next [`Task`][agentlightning.Task] available for execution, or
|
||||
``None`` if polling fails.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._next_task_uri)
|
||||
while True:
|
||||
response = self._request_json(url)
|
||||
if response:
|
||||
task_if_any = TaskIfAny.model_validate(response)
|
||||
if task_if_any.is_available and task_if_any.task:
|
||||
self.task_count += 1
|
||||
logger.info(f"[Task {self.task_count} Received] ID: {task_if_any.task.rollout_id}")
|
||||
return task_if_any.task
|
||||
logger.debug(f"No task available yet. Retrying in {self.poll_interval} seconds...")
|
||||
time.sleep(self.poll_interval)
|
||||
|
||||
def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
|
||||
"""Fetch a specific resource bundle by identifier.
|
||||
|
||||
Args:
|
||||
resource_id: Identifier sourced from the task metadata.
|
||||
|
||||
Returns:
|
||||
Cached or freshly downloaded
|
||||
[`ResourcesUpdate`][agentlightning.ResourcesUpdate], or
|
||||
``None`` when the server returns an error.
|
||||
"""
|
||||
if resource_id in self._resource_cache:
|
||||
logger.debug(f"Found resources '{resource_id}' in cache.")
|
||||
return self._resource_cache[resource_id]
|
||||
|
||||
url = urllib.parse.urljoin(self.endpoint, f"{self._resources_uri}/{resource_id}")
|
||||
response = self._request_json(url)
|
||||
if response:
|
||||
resources_update = ResourcesUpdate.model_validate(response)
|
||||
self._resource_cache[resource_id] = resources_update
|
||||
logger.info(f"Fetched and cached resources for ID: {resource_id}")
|
||||
return resources_update
|
||||
return None
|
||||
|
||||
def get_latest_resources(self) -> Optional[ResourcesUpdate]:
|
||||
"""Fetch the most recent resource bundle advertised by the server.
|
||||
|
||||
Returns:
|
||||
[`ResourcesUpdate`][agentlightning.ResourcesUpdate] for the
|
||||
newest version, or ``None`` when unavailable.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._latest_resources_uri)
|
||||
response = self._request_json(url)
|
||||
if response:
|
||||
resources_update = ResourcesUpdate.model_validate(response)
|
||||
self._resource_cache[resources_update.resources_id] = resources_update
|
||||
return resources_update
|
||||
return None
|
||||
|
||||
def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
|
||||
"""Submit a completed rollout back to the server.
|
||||
|
||||
Args:
|
||||
rollout: Legacy rollout payload produced by the executor.
|
||||
|
||||
Returns:
|
||||
Parsed JSON response returned by the server, or ``None`` when the request fails.
|
||||
"""
|
||||
url = urllib.parse.urljoin(self.endpoint, self._report_rollout_uri)
|
||||
payload = rollout.model_dump(mode="json")
|
||||
return self._post_json(url, payload)
|
||||
|
||||
|
||||
class DevTaskLoader(AgentLightningClient):
|
||||
"""In-memory task loader used for development and integration tests.
|
||||
|
||||
The loader mimics the behavior of the legacy HTTP server by storing tasks and
|
||||
resources locally. Polling methods simply iterate over the provided collection,
|
||||
allowing rapid iteration without provisioning any external infrastructure.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
|
||||
[`DevTaskLoader`][agentlightning.client.DevTaskLoader] is a compatibility shim.
|
||||
Prefer [`Trainer.dev`][agentlightning.Trainer.dev] for new code.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tasks: Union[List[TaskInput], List[Task]],
|
||||
resources: Union[NamedResources, ResourcesUpdate],
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize the loader with predefined tasks and resources.
|
||||
|
||||
Args:
|
||||
tasks: Sequence of task inputs or preconstructed tasks that will be served in
|
||||
order.
|
||||
resources: Static resources returned for any `resources_id` query.
|
||||
**kwargs: Additional keyword arguments forwarded to the parent client.
|
||||
|
||||
Raises:
|
||||
ValueError: If no tasks are provided or both [`Task`][agentlightning.Task]
|
||||
and [`TaskInput`][agentlightning.TaskInput] instances are mixed.
|
||||
"""
|
||||
warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
|
||||
super().__init__(endpoint="local://", **kwargs)
|
||||
self._tasks = tasks.copy()
|
||||
if len(self._tasks) == 0:
|
||||
raise ValueError("DevTaskLoader requires at least one task to be provided.")
|
||||
|
||||
# Check if tasks are mixture of TaskInput and Task
|
||||
if any(isinstance(task, Task) for task in self._tasks):
|
||||
if not all(isinstance(task, Task) for task in self._tasks):
|
||||
raise ValueError("All tasks must be either Task or TaskInput objects.")
|
||||
|
||||
self._task_index = 0
|
||||
|
||||
if isinstance(resources, ResourcesUpdate):
|
||||
self._resources_update = resources
|
||||
else:
|
||||
self._resources_update = ResourcesUpdate(
|
||||
resources_id="local", resources=resources, create_time=time.time(), update_time=time.time(), version=1
|
||||
)
|
||||
|
||||
# Store rollouts posted back to the loader for easy debugging of local runs
|
||||
self._rollouts: List[RolloutLegacy] = []
|
||||
|
||||
@property
|
||||
def rollouts(self) -> List[RolloutLegacy]:
|
||||
"""Return the rollouts posted back to the loader during development runs."""
|
||||
return self._rollouts
|
||||
|
||||
def poll_next_task(self) -> Optional[Task]:
|
||||
"""Return the next task from the local queue.
|
||||
|
||||
If [`TaskInput`][agentlightning.TaskInput] instances were provided,
|
||||
they are converted into [`Task`][agentlightning.Task] objects on the
|
||||
fly. Otherwise, the preconstructed tasks are returned in sequence.
|
||||
|
||||
Returns:
|
||||
Next task to execute.
|
||||
"""
|
||||
if self._task_index >= len(self._tasks):
|
||||
self._task_index = 0
|
||||
|
||||
task_or_input = self._tasks[self._task_index]
|
||||
|
||||
if isinstance(task_or_input, Task):
|
||||
task = task_or_input
|
||||
else:
|
||||
rollout_id = f"local_task_{self._task_index + 1:03d}"
|
||||
task = Task(
|
||||
rollout_id=rollout_id,
|
||||
input=task_or_input,
|
||||
resources_id=self._resources_update.resources_id,
|
||||
create_time=time.time(),
|
||||
)
|
||||
|
||||
self._task_index += 1
|
||||
self.task_count += 1
|
||||
logger.info(f"[Task {self.task_count} Received] Task ID: {task.rollout_id}")
|
||||
return task
|
||||
|
||||
def get_resources_by_id(self, resource_id: str) -> Optional[ResourcesUpdate]:
|
||||
logger.debug(f"DevTaskLoader checking resources for ID: {resource_id}")
|
||||
if resource_id != self._resources_update.resources_id:
|
||||
raise ValueError(
|
||||
f"Resource ID '{resource_id}' not found. Only '{self._resources_update.resources_id}' is available."
|
||||
)
|
||||
return self._resources_update
|
||||
|
||||
def get_latest_resources(self) -> Optional[ResourcesUpdate]:
|
||||
logger.debug("DevTaskLoader returning latest resources.")
|
||||
return self._resources_update
|
||||
|
||||
def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
|
||||
logger.debug(f"DevTaskLoader received rollout for task: {rollout.rollout_id}")
|
||||
self._rollouts.append(rollout)
|
||||
return {"status": "received", "rollout_id": rollout.rollout_id}
|
||||
|
||||
async def poll_next_task_async(self) -> Optional[Task]:
|
||||
return self.poll_next_task()
|
||||
|
||||
async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
|
||||
return self.get_resources_by_id(resource_id)
|
||||
|
||||
async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
|
||||
return self.get_latest_resources()
|
||||
|
||||
async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
|
||||
return self.post_rollout(rollout)
|
||||
|
||||
def __repr__(self):
|
||||
return f"DevTaskLoader(num_tasks={len(self._tasks)}, resources={self._resources_update.resources})"
|
||||
@@ -0,0 +1,348 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""
|
||||
This file is not carefully reviewed.
|
||||
It might contain unintentional bugs and issues.
|
||||
Please always review the parsed construction arguments before using them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import inspect
|
||||
import logging
|
||||
from typing import _GenericAlias # type: ignore
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
get_args,
|
||||
get_origin,
|
||||
get_type_hints,
|
||||
overload,
|
||||
)
|
||||
|
||||
CliConfigurable = Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["lightning_cli"]
|
||||
|
||||
# TypeVars for precise return type hinting with overloads
|
||||
_C = TypeVar("_C", bound=CliConfigurable)
|
||||
_C1 = TypeVar("_C1", bound=CliConfigurable)
|
||||
_C2 = TypeVar("_C2", bound=CliConfigurable)
|
||||
_C3 = TypeVar("_C3", bound=CliConfigurable)
|
||||
_C4 = TypeVar("_C4", bound=CliConfigurable)
|
||||
|
||||
|
||||
# Custom type for CLI arguments that can be string or None
|
||||
def nullable_str(value: str) -> str | None:
|
||||
"""Converts specific string values (case-insensitive) to None, otherwise returns the string."""
|
||||
if value.lower() in ["none", "null", "~", "nil"]: # Define keywords for None
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def nullable_int(value: str) -> int | None:
|
||||
"""Converts specific string values (case-insensitive) to None, otherwise returns the integer."""
|
||||
if value.lower() in ["none", "null", "~", "nil"]: # Define keywords for None
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(f"Invalid integer value: '{value}'")
|
||||
|
||||
|
||||
def nullable_float(value: str) -> float | None:
|
||||
"""Converts specific string values (case-insensitive) to None, otherwise returns the float."""
|
||||
if value.lower() in ["none", "null", "~", "nil"]: # Define keywords for None
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(f"Invalid float value: '{value}'")
|
||||
|
||||
|
||||
def _str_to_bool(v: str) -> bool:
|
||||
"""Converts common string representations of bool to Python bool (case-insensitive)."""
|
||||
if isinstance(v, bool): # type: ignore
|
||||
return v # Allow passing bools directly if used programmatically
|
||||
lowered_v = v.lower()
|
||||
if lowered_v in ("yes", "true", "t", "y", "1"):
|
||||
return True
|
||||
elif lowered_v in ("no", "false", "f", "n", "0"):
|
||||
return False
|
||||
else:
|
||||
raise argparse.ArgumentTypeError(f"Boolean value expected (e.g., 'true', 'false', 'yes', 'no'), got '{v}'")
|
||||
|
||||
|
||||
def _get_param_type_details(param_annotation: Any) -> Tuple[Any, bool, bool]:
|
||||
"""Normalize an annotation into its core type, optionality, and list status.
|
||||
|
||||
Args:
|
||||
param_annotation: The annotation to inspect.
|
||||
|
||||
Returns:
|
||||
A tuple ``(core_type, is_optional, is_list)`` describing the normalized type.
|
||||
|
||||
- For ``Optional[T]`` → ``(T, True, is_list_status_of_T)``
|
||||
- For ``List[T]`` → ``(List[T], is_optional_status_of_List, True)``
|
||||
- For ``Optional[List[T]]`` → ``(List[T], True, True)``
|
||||
"""
|
||||
is_optional = False
|
||||
is_list = False
|
||||
current_type = param_annotation
|
||||
|
||||
# Check for outer Optional
|
||||
origin = get_origin(current_type)
|
||||
if origin is Union:
|
||||
union_args = get_args(current_type)
|
||||
if len(union_args) == 2 and type(None) in union_args:
|
||||
is_optional = True
|
||||
current_type = next(arg for arg in union_args if arg is not type(None)) # Unwrap Optional
|
||||
|
||||
# Check if the (potentially unwrapped) type is a List
|
||||
origin = get_origin(current_type) # Re-check origin after potential unwrap
|
||||
if origin is list or (isinstance(current_type, _GenericAlias) and current_type.__origin__ is list):
|
||||
is_list = True
|
||||
|
||||
return current_type, is_optional, is_list
|
||||
|
||||
|
||||
def _determine_argparse_type(param_type: Any) -> Callable[[str], Any]:
|
||||
"""Determines the type for argparse based on parameter type details."""
|
||||
core_type, is_optional, _ = _get_param_type_details(param_type)
|
||||
if core_type is str and is_optional:
|
||||
return nullable_str # Special handling for Optional[str]
|
||||
elif core_type is int and is_optional:
|
||||
return nullable_int
|
||||
elif core_type is float and is_optional:
|
||||
return nullable_float
|
||||
elif core_type is bool:
|
||||
return _str_to_bool # Special handling for bool
|
||||
elif core_type in (int, float, str):
|
||||
return core_type
|
||||
return str # Default to str if no specific type is provided (including empty)
|
||||
|
||||
|
||||
def _determine_argparse_type_and_nargs(
|
||||
core_param_type: Any, is_param_list: bool # The type after unwrapping an outer Optional
|
||||
) -> Dict[str, Any]:
|
||||
"""Determines the 'type' and 'nargs' for argparse based on parameter type details."""
|
||||
kwargs: Dict[str, Any] = {}
|
||||
|
||||
if is_param_list:
|
||||
kwargs["nargs"] = "*" # Allows zero or more arguments for lists
|
||||
list_item_annotations = get_args(core_param_type) # For List[T], core_param_type is List[T]
|
||||
|
||||
if list_item_annotations and list_item_annotations[0] is not Any:
|
||||
item_ann = list_item_annotations[0]
|
||||
# Check if the list item itself is, e.g., Optional[str] or bool
|
||||
kwargs["type"] = _determine_argparse_type(item_ann)
|
||||
else:
|
||||
kwargs["type"] = str
|
||||
else: # Not a list
|
||||
kwargs["type"] = _determine_argparse_type(core_param_type)
|
||||
return kwargs
|
||||
|
||||
|
||||
def _build_help_string(cls_name: str, param_name: str, core_type: Any, is_optional: bool, is_list: bool) -> str:
|
||||
"""Constructs a descriptive help string for a CLI argument."""
|
||||
type_display_name = "Any"
|
||||
if core_type is not inspect.Parameter.empty:
|
||||
type_display_name = getattr(core_type, "__name__", str(core_type))
|
||||
|
||||
if is_list:
|
||||
list_item_args = get_args(core_type) # core_type is List[T] here
|
||||
item_name = "Any"
|
||||
if list_item_args and list_item_args[0] is not Any:
|
||||
inner_item_core_type, inner_item_optional, _ = _get_param_type_details(list_item_args[0])
|
||||
item_name = getattr(inner_item_core_type, "__name__", str(inner_item_core_type))
|
||||
if inner_item_optional: # e.g. List[Optional[str]]
|
||||
item_name = f"Optional[{item_name}]"
|
||||
type_display_name = f"List[{item_name}]"
|
||||
|
||||
full_type_display = f"Optional[{type_display_name}]" if is_optional and not is_list else type_display_name
|
||||
if is_optional and is_list: # e.g. Optional[List[str]]
|
||||
full_type_display = f"Optional[{type_display_name}]"
|
||||
|
||||
help_str = f"For {cls_name}: '{param_name}'. Inferred type: {full_type_display}."
|
||||
return help_str
|
||||
|
||||
|
||||
def _add_argument_for_parameter(
|
||||
parser: argparse.ArgumentParser,
|
||||
cls: Type[CliConfigurable],
|
||||
param_name: str,
|
||||
param_obj: inspect.Parameter,
|
||||
dest_name: str,
|
||||
resolved_param_annotation: Any = None,
|
||||
) -> None:
|
||||
"""Configures and adds a single CLI argument for an __init__ parameter."""
|
||||
if resolved_param_annotation is None:
|
||||
param_type_annotation = param_obj.annotation
|
||||
else:
|
||||
param_type_annotation = resolved_param_annotation
|
||||
|
||||
# core_type is the main type (e.g., int, str, List[str]), after unwrapping the outermost Optional.
|
||||
# is_overall_optional indicates if the parameter itself can be None (e.g. param: Optional[T] = None)
|
||||
# is_list indicates if core_type is a List.
|
||||
core_type, is_overall_optional, is_list = _get_param_type_details(param_type_annotation)
|
||||
|
||||
has_init_default = param_obj.default is not inspect.Parameter.empty
|
||||
init_default_value = param_obj.default if has_init_default else None
|
||||
|
||||
argparse_kwargs = _determine_argparse_type_and_nargs(core_type if is_list else param_type_annotation, is_list)
|
||||
|
||||
if has_init_default:
|
||||
argparse_kwargs["default"] = init_default_value
|
||||
elif is_overall_optional: # Parameter is Optional (e.g. Optional[int]) and no explicit default in __init__
|
||||
argparse_kwargs["default"] = None # So, if not provided on CLI, it becomes None.
|
||||
|
||||
argparse_kwargs["help"] = _build_help_string(cls.__name__, param_name, core_type, is_overall_optional, is_list)
|
||||
|
||||
if not has_init_default and not is_overall_optional: # Required if no __init__ default AND not Optional
|
||||
argparse_kwargs["required"] = True
|
||||
if "default" in argparse_kwargs: # Should not happen if logic is correct
|
||||
del argparse_kwargs["default"]
|
||||
|
||||
cli_arg_name = f"--{cls.__name__.lower()}.{param_name.replace('_', '-')}"
|
||||
parser.add_argument(cli_arg_name, dest=dest_name, **argparse_kwargs)
|
||||
|
||||
|
||||
def _add_arguments_for_class(
|
||||
parser: argparse.ArgumentParser,
|
||||
cls: Type[CliConfigurable],
|
||||
class_arg_configs_maps: Dict[Type[CliConfigurable], Dict[str, str]], # Maps cls to {param_name: dest_name}
|
||||
) -> None:
|
||||
"""Adds all relevant CLI arguments for a given class by processing its __init__ parameters."""
|
||||
cls_name_lower = cls.__name__.lower()
|
||||
sig = inspect.signature(cls.__init__)
|
||||
|
||||
try:
|
||||
# Resolve string annotations to actual types using get_type_hints.
|
||||
# For methods, get_type_hints automatically uses obj.__globals__ for globalns.
|
||||
resolved_hints = get_type_hints(cls.__init__)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Could not resolve type hints for {cls.__name__}.__init__ using get_type_hints: {e}. "
|
||||
f"CLI argument parsing for this class might be based on string annotations, "
|
||||
"which could be unreliable for complex types."
|
||||
)
|
||||
resolved_hints = {} # Fallback to an empty dict if resolution fails
|
||||
|
||||
if cls not in class_arg_configs_maps: # Ensure the class entry exists
|
||||
class_arg_configs_maps[cls] = {}
|
||||
|
||||
for param_name, param_obj in sig.parameters.items():
|
||||
if param_name == "self": # Skip 'self'
|
||||
continue
|
||||
|
||||
dest_name = f"{cls_name_lower}_{param_name}" # Unique destination for argparse
|
||||
class_arg_configs_maps[cls][param_name] = dest_name # Store mapping for later instantiation
|
||||
|
||||
# Use the resolved hint if available, otherwise fallback to param_obj.annotation (which might be a string)
|
||||
actual_param_annotation = resolved_hints.get(param_name, param_obj.annotation)
|
||||
_add_argument_for_parameter(parser, cls, param_name, param_obj, dest_name, actual_param_annotation)
|
||||
|
||||
|
||||
def _create_argument_parser() -> argparse.ArgumentParser:
|
||||
"""Creates and returns the main ArgumentParser with default settings."""
|
||||
return argparse.ArgumentParser(
|
||||
description="CLI configurator for application components.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter, # Automatically shows default values in help
|
||||
)
|
||||
|
||||
|
||||
def _instantiate_classes(
|
||||
parsed_args: argparse.Namespace,
|
||||
classes: Tuple[Type[CliConfigurable], ...],
|
||||
class_arg_configs_maps: Dict[Type[CliConfigurable], Dict[str, str]],
|
||||
) -> Tuple[CliConfigurable, ...]:
|
||||
"""Instantiates classes using the parsed CLI arguments and the stored mappings."""
|
||||
instances_list: List[CliConfigurable] = []
|
||||
for cls in classes:
|
||||
constructor_args: Dict[str, Any] = {}
|
||||
# Get the {__init__ param_name: argparse_dest_name} map for the current class
|
||||
param_to_dest_map = class_arg_configs_maps.get(cls, {})
|
||||
|
||||
sig = inspect.signature(cls.__init__)
|
||||
for param_name_in_sig, _ in sig.parameters.items():
|
||||
if param_name_in_sig == "self":
|
||||
continue
|
||||
|
||||
dest_name_for_arg = param_to_dest_map.get(param_name_in_sig)
|
||||
if dest_name_for_arg and hasattr(parsed_args, dest_name_for_arg):
|
||||
value = getattr(parsed_args, dest_name_for_arg)
|
||||
constructor_args[param_name_in_sig] = value
|
||||
# If an argument was required by argparse, parse_args() would have exited if missing.
|
||||
# If not required and not provided, its default value (set by argparse) is used.
|
||||
|
||||
try:
|
||||
logger.info("Instantiating %s with args: %s", cls.__name__, constructor_args)
|
||||
instances_list.append(cls(**constructor_args))
|
||||
except Exception as e:
|
||||
parsed_args_for_cls = {
|
||||
k: getattr(parsed_args, v) for k, v in param_to_dest_map.items() if hasattr(parsed_args, v)
|
||||
}
|
||||
logger.error(
|
||||
f"Error instantiating {cls.__name__} with resolved args {constructor_args}. "
|
||||
f"Parsed args for class: "
|
||||
f"{parsed_args_for_cls}. "
|
||||
f"Error: {e}"
|
||||
)
|
||||
raise
|
||||
|
||||
return tuple(instances_list)
|
||||
|
||||
|
||||
@overload
|
||||
def lightning_cli(cls1: Type[_C1]) -> _C1: ...
|
||||
@overload
|
||||
def lightning_cli(cls1: Type[_C1], cls2: Type[_C2]) -> Tuple[_C1, _C2]: ...
|
||||
@overload
|
||||
def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3]) -> Tuple[_C1, _C2, _C3]: ...
|
||||
@overload
|
||||
def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[_C4]) -> Tuple[_C1, _C2, _C3, _C4]: ...
|
||||
@overload # Fallback for more than 4 or a dynamic number of classes
|
||||
def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
|
||||
|
||||
|
||||
# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
|
||||
|
||||
|
||||
def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]: # type: ignore
|
||||
"""
|
||||
Parses command-line arguments to configure and instantiate provided CliConfigurable classes.
|
||||
|
||||
Args:
|
||||
*classes: One or more classes that inherit from CliConfigurable. Each class's
|
||||
__init__ parameters will be exposed as command-line arguments.
|
||||
|
||||
Returns:
|
||||
A tuple of instantiated objects, corresponding to the input classes in order.
|
||||
"""
|
||||
if not classes:
|
||||
return tuple() # Return an empty tuple if no classes are provided
|
||||
|
||||
parser = _create_argument_parser()
|
||||
|
||||
# This map will store {cls: {init_param_name: argparse_dest_name}}
|
||||
class_arg_configs_maps: Dict[Type[CliConfigurable], Dict[str, str]] = {}
|
||||
|
||||
for cls in classes:
|
||||
_add_arguments_for_class(parser, cls, class_arg_configs_maps)
|
||||
|
||||
parsed_args = parser.parse_args() # Uses sys.argv[1:] by default
|
||||
|
||||
# Correctly handle single class case for return type matching overloads
|
||||
instances = _instantiate_classes(parsed_args, classes, class_arg_configs_maps)
|
||||
if len(classes) == 1:
|
||||
return instances[0]
|
||||
return instances
|
||||
@@ -0,0 +1,43 @@
|
||||
# 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
|
||||
from .object import emit_object, get_object_value
|
||||
from .reward import (
|
||||
emit_reward,
|
||||
find_final_reward,
|
||||
find_reward_spans,
|
||||
get_reward_value,
|
||||
get_rewards_from_span,
|
||||
is_reward_span,
|
||||
reward,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"reward",
|
||||
"operation",
|
||||
"emit_reward",
|
||||
"get_reward_value",
|
||||
"get_rewards_from_span",
|
||||
"is_reward_span",
|
||||
"find_reward_spans",
|
||||
"find_final_reward",
|
||||
"emit_message",
|
||||
"emit_object",
|
||||
"emit_exception",
|
||||
"emit_annotation",
|
||||
"get_message_value",
|
||||
"get_object_value",
|
||||
]
|
||||
@@ -0,0 +1,374 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Helpers for emitting annotation/operation spans."""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
ContextManager,
|
||||
Dict,
|
||||
Optional,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
|
||||
from agentlightning.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import SpanCoreFields, SpanRecordingContext, TraceStatus
|
||||
from agentlightning.utils.otel import check_attributes_sanity, flatten_attributes, sanitize_attributes
|
||||
|
||||
_FnType = TypeVar("_FnType", bound=Callable[..., Any])
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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].
|
||||
|
||||
Annotation spans are used to annotate a specific event or a part of rollout.
|
||||
See [semconv][agentlightning.semconv] for conventional annotation keys in Agent-lightning.
|
||||
|
||||
If annotations contain nested dicts, they will be flattened before emitting.
|
||||
Complex objects will lead to emitting failures.
|
||||
|
||||
Args:
|
||||
annotation: Dictionary containing annotation key-value pairs.
|
||||
Representatives are rewards, tags, and metadata.
|
||||
propagate: Whether to propagate the span to tracers automatically.
|
||||
"""
|
||||
annotation_attributes = flatten_attributes(annotation, expand_leaf_lists=False)
|
||||
check_attributes_sanity(annotation_attributes)
|
||||
sanitized_attributes = sanitize_attributes(annotation_attributes)
|
||||
logger.debug("Emitting annotation span with keys %s", sanitized_attributes.keys())
|
||||
|
||||
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"),
|
||||
)
|
||||
|
||||
|
||||
class OperationContext:
|
||||
"""Context manager and decorator for tracing operations.
|
||||
|
||||
This class manages a tracer-backed span for a logical unit of work. It can be
|
||||
used either:
|
||||
|
||||
* As a decorator, in which case inputs and outputs are inferred
|
||||
automatically from the wrapped function's signature.
|
||||
* As a context manager, in which case inputs and outputs can be recorded
|
||||
explicitly via [`set_input`][agentlightning.emitter.annotation.OperationContext.set_input]
|
||||
and [`set_output`][agentlightning.emitter.annotation.OperationContext.set_output].
|
||||
|
||||
Attributes:
|
||||
name: Human-readable span name.
|
||||
initial_attributes: Attributes applied when the span is created.
|
||||
tracer: Tracer implementation used to create spans.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
|
||||
"""Initialize a new operation context.
|
||||
|
||||
Args:
|
||||
name: Human-readable name of the span.
|
||||
attributes: Initial attributes attached to the span. Values are
|
||||
JSON-serialized where necessary.
|
||||
propagate: Whether the span should be sent to active exporters.
|
||||
"""
|
||||
self.name = name
|
||||
self.initial_attributes = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
self.propagate = propagate
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot trace operation spans.")
|
||||
self.tracer = tracer
|
||||
else:
|
||||
self.tracer = DummyTracer()
|
||||
self._ctx_manager: Optional[ContextManager[SpanRecordingContext]] = None
|
||||
self._recording_context: Optional[SpanRecordingContext] = None
|
||||
self._span: Optional[SpanCoreFields] = None
|
||||
|
||||
def __enter__(self) -> "OperationContext":
|
||||
"""Enter the context manager and start a new span.
|
||||
|
||||
Returns:
|
||||
The current :class:`OperationContext` instance with an active span.
|
||||
"""
|
||||
sanitized_attrs = sanitize_attributes(self.initial_attributes)
|
||||
self._ctx_manager = self.tracer.operation_context(self.name, attributes=sanitized_attrs)
|
||||
recording_context = self._ctx_manager.__enter__()
|
||||
self._recording_context = recording_context
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[TracebackType],
|
||||
) -> None:
|
||||
"""Exit the context manager and finish the span."""
|
||||
if self._ctx_manager:
|
||||
self._ctx_manager.__exit__(exc_type, exc_val, exc_tb)
|
||||
if self._recording_context:
|
||||
self._span = self._recording_context.get_recorded_span()
|
||||
self._ctx_manager = None
|
||||
self._recording_context = None
|
||||
|
||||
def span(self) -> SpanCoreFields:
|
||||
"""Get the span that was created by this context manager."""
|
||||
if self._span is None:
|
||||
raise RuntimeError("Span is not ready yet.")
|
||||
return self._span
|
||||
|
||||
def set_input(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Record input arguments on the current span.
|
||||
|
||||
Positional arguments are stored under the `input.args.<index>` attributes,
|
||||
and keyword arguments are stored under `input.<name>` attributes.
|
||||
|
||||
This is intended for use inside a `with operation(...) as op` block.
|
||||
|
||||
Args:
|
||||
*args: Positional arguments to record.
|
||||
**kwargs: Keyword arguments to record.
|
||||
"""
|
||||
if not self._recording_context:
|
||||
raise RuntimeError("No recording context found. Cannot set input.")
|
||||
|
||||
prefix = LightningSpanAttributes.OPERATION_INPUT.value
|
||||
attributes: Dict[str, Any] = {}
|
||||
if args:
|
||||
for idx, value in enumerate(args):
|
||||
flattened = flatten_attributes({str(idx): value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{prefix}.args.{nested_key}"] = nested_value
|
||||
if kwargs:
|
||||
for key, value in kwargs.items():
|
||||
flattened = flatten_attributes({key: value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{prefix}.{nested_key}"] = nested_value
|
||||
if attributes:
|
||||
self._recording_context.record_attributes(sanitize_attributes(attributes))
|
||||
|
||||
def set_output(self, output: Any) -> None:
|
||||
"""Record the output value on the current span.
|
||||
|
||||
This is intended for use inside a `with operation(...) as op` block.
|
||||
|
||||
Args:
|
||||
output: The output value to record.
|
||||
"""
|
||||
if not self._recording_context:
|
||||
raise RuntimeError("No recording context found. Cannot set output.")
|
||||
|
||||
flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: output})
|
||||
self._recording_context.record_attributes(sanitize_attributes(flattened))
|
||||
|
||||
def __call__(self, fn: _FnType) -> _FnType:
|
||||
"""Wrap a callable so its execution is traced in a span.
|
||||
|
||||
When used as a decorator, a new span is created for each call to
|
||||
the wrapped function. The bound arguments are recorded as input
|
||||
attributes, the return value is recorded as an output attribute,
|
||||
and any exception is recorded and marks the span as an error.
|
||||
|
||||
Args:
|
||||
fn: The function or coroutine function to wrap.
|
||||
|
||||
Returns:
|
||||
The wrapped callable.
|
||||
"""
|
||||
function_name = fn.__name__
|
||||
|
||||
sig = inspect.signature(fn)
|
||||
|
||||
sanitized_init_attrs = sanitize_attributes(
|
||||
{LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
|
||||
)
|
||||
|
||||
def _record_auto_inputs(
|
||||
recording_ctx: SpanRecordingContext, args: Tuple[Any, ...], kwargs: Dict[str, Any]
|
||||
) -> None:
|
||||
"""Bind arguments to signature and log them on the span."""
|
||||
attributes: Dict[str, Any] = {}
|
||||
try:
|
||||
bound = sig.bind(*args, **kwargs)
|
||||
bound.apply_defaults()
|
||||
for name, value in bound.arguments.items():
|
||||
parameter = sig.parameters.get(name)
|
||||
if parameter and parameter.kind is inspect.Parameter.VAR_POSITIONAL:
|
||||
attr_prefix = f"{LightningSpanAttributes.OPERATION_INPUT.value}.{name}"
|
||||
for idx, item in enumerate(value):
|
||||
flattened = flatten_attributes({str(idx): item})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{attr_prefix}.{nested_key}"] = nested_value
|
||||
else:
|
||||
flattened = flatten_attributes({name: value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
|
||||
except Exception:
|
||||
if args:
|
||||
for idx, value in enumerate(args):
|
||||
flattened = flatten_attributes({str(idx): value})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.args.{nested_key}"] = (
|
||||
nested_value
|
||||
)
|
||||
if kwargs:
|
||||
flattened = flatten_attributes({"kwargs": kwargs})
|
||||
for nested_key, nested_value in flattened.items():
|
||||
attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
|
||||
if attributes:
|
||||
recording_ctx.record_attributes(sanitize_attributes(attributes))
|
||||
|
||||
def _record_auto_outputs(recording_ctx: SpanRecordingContext, result: Any) -> None:
|
||||
"""Record the output value on the span."""
|
||||
flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: result})
|
||||
recording_ctx.record_attributes(sanitize_attributes(flattened))
|
||||
|
||||
if inspect.iscoroutinefunction(fn) or (
|
||||
# For backwards compatibility.
|
||||
hasattr(asyncio, "iscoroutinefunction")
|
||||
and asyncio.iscoroutinefunction(fn) # type: ignore
|
||||
):
|
||||
|
||||
@functools.wraps(fn)
|
||||
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Async wrapper that traces the wrapped coroutine."""
|
||||
with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
|
||||
_record_auto_inputs(recording_ctx, args, kwargs)
|
||||
result = await fn(*args, **kwargs)
|
||||
_record_auto_outputs(recording_ctx, result)
|
||||
return result
|
||||
|
||||
return cast(_FnType, async_wrapper)
|
||||
|
||||
else:
|
||||
|
||||
@functools.wraps(fn)
|
||||
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
"""Sync wrapper that traces the wrapped callable."""
|
||||
with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
|
||||
_record_auto_inputs(recording_ctx, args, kwargs)
|
||||
result = fn(*args, **kwargs)
|
||||
_record_auto_outputs(recording_ctx, result)
|
||||
return result
|
||||
|
||||
return cast(_FnType, sync_wrapper)
|
||||
|
||||
|
||||
@overload
|
||||
def operation(
|
||||
fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
|
||||
) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(
|
||||
*, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
|
||||
) -> OperationContext: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(fn: _FnType, *, name: Optional[str] = None, **additional_attributes: Any) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(*, name: Optional[str] = None, **additional_attributes: Any) -> OperationContext: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(fn: _FnType, **additional_attributes: Any) -> _FnType: ...
|
||||
|
||||
|
||||
@overload
|
||||
def operation(**additional_attributes: Any) -> OperationContext: ...
|
||||
|
||||
|
||||
def operation(
|
||||
fn: Optional[_FnType] = None,
|
||||
*,
|
||||
propagate: bool = True,
|
||||
name: Optional[str] = None,
|
||||
**additional_attributes: Any,
|
||||
) -> Union[_FnType, OperationContext]:
|
||||
"""Entry point for tracking operations.
|
||||
|
||||
This helper can be used either as a decorator or as a context manager.
|
||||
The span name is fixed to [`AGL_OPERATION`][agentlightning.semconv.AGL_OPERATION];
|
||||
custom span names are not supported. Any keyword arguments are recorded as span attributes.
|
||||
|
||||
Usage as a decorator:
|
||||
|
||||
```python
|
||||
@operation
|
||||
def func(...):
|
||||
...
|
||||
|
||||
@operation(category="compute")
|
||||
def func(...):
|
||||
...
|
||||
```
|
||||
|
||||
Usage as a context manager:
|
||||
|
||||
```python
|
||||
with operation(user_id=123) as op:
|
||||
op.set_input(data=data)
|
||||
# ... do work ...
|
||||
op.set_output(result)
|
||||
```
|
||||
|
||||
Args:
|
||||
fn: When used as `@operation`, this is the wrapped function.
|
||||
When used as `operation(**attrs)`, this should be omitted (or
|
||||
left as `None`) and only keyword attributes are provided.
|
||||
propagate: Whether spans should use the active span processor. When False,
|
||||
spans will stay local and not be exported.
|
||||
name: Optional alias that populates
|
||||
[`LightningSpanAttributes.OPERATION_NAME`][agentlightning.semconv.LightningSpanAttributes.OPERATION_NAME]
|
||||
when `additional_attributes` does not already define it.
|
||||
**additional_attributes: Additional span attributes to attach at
|
||||
creation time.
|
||||
|
||||
Returns:
|
||||
Either a wrapped callable (when used as a decorator) or an
|
||||
[`OperationContext`][agentlightning.emitter.annotation.OperationContext]
|
||||
(when used as a context manager factory).
|
||||
"""
|
||||
|
||||
if name is not None:
|
||||
if LightningSpanAttributes.OPERATION_NAME.value in additional_attributes:
|
||||
raise ValueError("Cannot specify both `name` and `additional_attributes.operation_name`.")
|
||||
additional_attributes[LightningSpanAttributes.OPERATION_NAME.value] = name
|
||||
|
||||
# Case 1: Used as @operation (bare decorator or with attributes)
|
||||
if callable(fn):
|
||||
# Create context with fixed name, then immediately wrap the function
|
||||
return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn)
|
||||
|
||||
# Case 2: Used as operation(...) / with operation(...)
|
||||
# Custom span names are intentionally not supported; use AGL_OPERATION.
|
||||
if fn is not None:
|
||||
raise ValueError("Custom span names are intentionally not supported when used as a context manager.")
|
||||
return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)
|
||||
@@ -0,0 +1,54 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_EXCEPTION
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import TraceStatus
|
||||
from agentlightning.utils.otel import flatten_attributes, format_exception_attributes, sanitize_attributes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_exception(
|
||||
exception: BaseException, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True
|
||||
) -> None:
|
||||
"""Record an exception with OpenTelemetry metadata.
|
||||
|
||||
Classic OpenTelemetry records exceptions in a dedicated logging service.
|
||||
We simplify the model and use trace spans to record exceptions as well.
|
||||
|
||||
Args:
|
||||
exception: Raised exception instance to serialize into telemetry attributes.
|
||||
attributes: Additional attributes to attach to the exception span.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
!!! note
|
||||
|
||||
The helper validates its input. If a non-exception value is provided,
|
||||
a TypeError is raised to indicate a programming mistake.
|
||||
"""
|
||||
if not isinstance(exception, BaseException): # type: ignore
|
||||
raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
|
||||
span_attributes = format_exception_attributes(exception)
|
||||
|
||||
if attributes:
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
|
||||
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"),
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_MESSAGE, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import Attributes, SpanLike
|
||||
from agentlightning.utils.otel import flatten_attributes, sanitize_attributes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
|
||||
"""Emit a textual message as an OpenTelemetry span.
|
||||
|
||||
Commonly used for sending debugging and logging messages.
|
||||
|
||||
Args:
|
||||
message: Human readable message to attach as a span attribute.
|
||||
attributes: Additional attributes to attach to the message span.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
!!! note
|
||||
OpenTelemetry distinguishes between logs and spans. Emitting the message as a
|
||||
span keeps all Agent Lightning telemetry in a single data store for analysis.
|
||||
"""
|
||||
if not isinstance(message, str): # type: ignore
|
||||
raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
|
||||
|
||||
if propagate:
|
||||
tracer = get_active_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("No active tracer found. Cannot emit message span.")
|
||||
else:
|
||||
tracer = DummyTracer()
|
||||
span_attributes: Attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
|
||||
if attributes:
|
||||
flattened = flatten_attributes(attributes, expand_leaf_lists=False)
|
||||
span_attributes.update(sanitize_attributes(flattened))
|
||||
logger.debug("Emitting message span with message: %s", message)
|
||||
tracer.create_span(
|
||||
AGL_MESSAGE,
|
||||
attributes=span_attributes,
|
||||
)
|
||||
|
||||
|
||||
def get_message_value(span: SpanLike) -> Optional[str]:
|
||||
"""Extract the message string from a message span.
|
||||
|
||||
Args:
|
||||
span: Span-like object to extract the message from.
|
||||
"""
|
||||
span_attributes = span.attributes or {}
|
||||
if LightningSpanAttributes.MESSAGE_BODY.value not in span_attributes:
|
||||
return None
|
||||
message = span_attributes[LightningSpanAttributes.MESSAGE_BODY.value]
|
||||
if isinstance(message, str):
|
||||
return message
|
||||
raise TypeError(f"Message must be a string, got: {type(message)}.")
|
||||
@@ -0,0 +1,117 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from agentlightning.semconv import AGL_OBJECT, LightningSpanAttributes
|
||||
from agentlightning.tracer.base import get_active_tracer
|
||||
from agentlightning.tracer.dummy import DummyTracer
|
||||
from agentlightning.types import SpanCoreFields, SpanLike, TraceStatus
|
||||
from agentlightning.utils.otel import flatten_attributes, full_qualified_name, sanitize_attributes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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:
|
||||
object: Data structure to encode as JSON and attach to the span payload.
|
||||
attributes: Additional attributes to attach to the object span.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
!!! note
|
||||
The payload must be JSON serializable. Non-serializable objects will lead to a RuntimeError.
|
||||
"""
|
||||
span_attributes = encode_object(object)
|
||||
if 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)
|
||||
|
||||
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]:
|
||||
"""Encode an object as span attributes.
|
||||
|
||||
Args:
|
||||
object: Data structure to encode as JSON.
|
||||
"""
|
||||
span_attributes = {}
|
||||
if isinstance(object, (str, int, float, bool)):
|
||||
span_attributes = {
|
||||
LightningSpanAttributes.OBJECT_TYPE.value: type(object).__name__,
|
||||
LightningSpanAttributes.OBJECT_LITERAL.value: str(object),
|
||||
}
|
||||
elif isinstance(object, bytes):
|
||||
b64_encoded = base64.b64encode(object).decode("utf-8")
|
||||
span_attributes = {
|
||||
LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
|
||||
LightningSpanAttributes.OBJECT_LITERAL.value: b64_encoded,
|
||||
}
|
||||
else:
|
||||
try:
|
||||
serialized = json.dumps(object)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError(f"Object must be JSON serializable, got: {type(object)}.") from exc
|
||||
|
||||
span_attributes = {
|
||||
LightningSpanAttributes.OBJECT_TYPE.value: full_qualified_name(type(object)), # type: ignore
|
||||
LightningSpanAttributes.OBJECT_JSON.value: serialized,
|
||||
}
|
||||
|
||||
return span_attributes
|
||||
|
||||
|
||||
def get_object_value(span: SpanLike) -> Any:
|
||||
"""Extract the object payload from an object span.
|
||||
|
||||
Args:
|
||||
span: Span object produced by Agent Lightning emitters.
|
||||
"""
|
||||
attributes = span.attributes or {}
|
||||
if LightningSpanAttributes.OBJECT_JSON.value in attributes:
|
||||
serialized = attributes[LightningSpanAttributes.OBJECT_JSON.value]
|
||||
try:
|
||||
return json.loads(serialized) # type: ignore
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise RuntimeError("Failed to deserialize object JSON from span.") from exc
|
||||
elif LightningSpanAttributes.OBJECT_LITERAL.value in attributes:
|
||||
literal = attributes[LightningSpanAttributes.OBJECT_LITERAL.value]
|
||||
obj_type = attributes.get(LightningSpanAttributes.OBJECT_TYPE.value, "str")
|
||||
if obj_type == "str":
|
||||
return literal
|
||||
elif obj_type == "int":
|
||||
# Let it raise errors if there are any
|
||||
return int(literal) # type: ignore
|
||||
elif obj_type == "float":
|
||||
return float(literal) # type: ignore
|
||||
elif obj_type == "bool":
|
||||
return literal.lower() == "true" # type: ignore
|
||||
elif obj_type == "bytes":
|
||||
return base64.b64decode(literal.encode("utf-8")) # type: ignore
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported object type for literal deserialization: {obj_type}")
|
||||
else:
|
||||
return None
|
||||
@@ -0,0 +1,320 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Helpers for emitting reward spans and integrating with AgentOps telemetry."""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import warnings
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from agentlightning.semconv import AGL_ANNOTATION, LightningSpanAttributes, RewardPydanticModel
|
||||
from agentlightning.types import SpanCoreFields, SpanLike
|
||||
from agentlightning.utils.otel import filter_and_unflatten_attributes
|
||||
|
||||
from .annotation import emit_annotation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"reward",
|
||||
"emit_reward",
|
||||
"get_reward_value",
|
||||
"get_rewards_from_span",
|
||||
"is_reward_span",
|
||||
"find_reward_spans",
|
||||
"find_final_reward",
|
||||
]
|
||||
|
||||
|
||||
class RewardDimension(TypedDict):
|
||||
"""Type representing a single dimension in a multi-dimensional reward."""
|
||||
|
||||
name: str
|
||||
value: float
|
||||
|
||||
|
||||
class _RewardSpanData(TypedDict):
|
||||
type: Literal["reward"]
|
||||
value: Optional[float]
|
||||
|
||||
|
||||
_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
|
||||
|
||||
|
||||
def reward(fn: _FnType) -> _FnType:
|
||||
"""Decorate a reward function so its outputs are tracked as spans.
|
||||
|
||||
The decorator integrates with AgentOps when it is available and falls back to
|
||||
the built-in telemetry otherwise. Both synchronous and asynchronous functions
|
||||
are supported transparently.
|
||||
|
||||
Deprecated:
|
||||
This decorator is deprecated. Use [`emit_reward`][agentlightning.emit_reward] instead.
|
||||
|
||||
Args:
|
||||
fn: Callable that produces a numeric reward.
|
||||
|
||||
Returns:
|
||||
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:
|
||||
return {"type": "reward", "value": None}
|
||||
if not isinstance(result, (float, int)): # type: ignore
|
||||
warnings.warn(f"Reward is ignored because it is not a number: {result}")
|
||||
return {"type": "reward", "value": None}
|
||||
return {"type": "reward", "value": float(result)}
|
||||
|
||||
# Check if the function is async
|
||||
is_async = inspect.iscoroutinefunction(fn) or (
|
||||
# For backwards compatibility.
|
||||
hasattr(asyncio, "iscoroutinefunction")
|
||||
and asyncio.iscoroutinefunction(fn) # type: ignore
|
||||
)
|
||||
|
||||
if is_async:
|
||||
|
||||
async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
|
||||
if not _agentops_initialized():
|
||||
# Track the reward without AgentOps
|
||||
result = await fn(*args, **kwargs)
|
||||
emit_reward(cast(float, result))
|
||||
return result
|
||||
|
||||
result: Optional[float] = None
|
||||
|
||||
@operation
|
||||
async def agentops_reward_operation() -> _RewardSpanData:
|
||||
# The reward function we are interested in tracing
|
||||
# It takes zero inputs and return a formatted dict
|
||||
nonlocal result
|
||||
result = await fn(*args, **kwargs)
|
||||
return wrap_result(result)
|
||||
|
||||
await agentops_reward_operation()
|
||||
return result
|
||||
|
||||
return wrapper_async # type: ignore
|
||||
|
||||
else:
|
||||
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
if not _agentops_initialized():
|
||||
# Track the reward without AgentOps
|
||||
result = fn(*args, **kwargs)
|
||||
emit_reward(cast(float, result))
|
||||
return result
|
||||
|
||||
result: Optional[float] = None
|
||||
|
||||
@operation
|
||||
def agentops_reward_operation() -> _RewardSpanData:
|
||||
nonlocal result
|
||||
result = fn(*args, **kwargs)
|
||||
return wrap_result(result)
|
||||
|
||||
agentops_reward_operation()
|
||||
return result
|
||||
|
||||
return wrapper # type: ignore
|
||||
|
||||
|
||||
def emit_reward(
|
||||
reward: float | Dict[str, Any],
|
||||
*,
|
||||
primary_key: str | None = None,
|
||||
attributes: Dict[str, Any] | None = None,
|
||||
propagate: bool = True,
|
||||
) -> SpanCoreFields:
|
||||
"""Emit a reward value as an OpenTelemetry span.
|
||||
|
||||
Examples:
|
||||
Emit a single-dimensional reward:
|
||||
>>> emit_reward(1.0)
|
||||
|
||||
Emit multi-dimensional rewards:
|
||||
>>> emit_reward({"task_completion": 1.0, "efficiency": 0.8}, primary_key="task_completion")
|
||||
|
||||
Emit a reward with additional attributes (for example linking to another response span):
|
||||
>>> from agentlightning.utils.otel import make_link_attributes
|
||||
>>> emit_reward(0.5, attributes=make_link_attributes({"gen_ai.response.id": "response-123"}))
|
||||
|
||||
Or adding tags onto the reward span:
|
||||
>>> from agentlightning.utils.otel import make_tag_attributes
|
||||
>>> emit_reward(0.7, attributes=make_tag_attributes(["fast", "reliable"]))
|
||||
|
||||
Args:
|
||||
reward: Numeric reward to record. Integers and booleans are converted to
|
||||
floating point numbers for consistency.
|
||||
Use a dictionary to represent a multi-dimensional reward.
|
||||
attributes: Other optional span attributes.
|
||||
propagate: Whether to propagate the span to exporters automatically.
|
||||
|
||||
Returns:
|
||||
Span core fields capturing the recorded reward.
|
||||
"""
|
||||
logger.debug(f"Emitting reward: {reward}")
|
||||
reward_dimensions: List[RewardDimension] = []
|
||||
if isinstance(reward, dict):
|
||||
reward_dict: Dict[str, float] = {}
|
||||
for k, v in reward.items():
|
||||
if isinstance(v, (int, bool)):
|
||||
reward_dict[k] = float(v)
|
||||
elif isinstance(v, float):
|
||||
reward_dict[k] = v
|
||||
else:
|
||||
raise ValueError(f"Reward value must be a number, got: {type(v)} for key {k}")
|
||||
if primary_key is None:
|
||||
raise ValueError("When emitting a multi-dimensional reward as a dict, primary_key must be provided.")
|
||||
if primary_key not in reward_dict:
|
||||
raise ValueError(f"Primary key '{primary_key}' not found in reward dict keys: {list(reward_dict.keys())}")
|
||||
reward_dimensions.append(RewardDimension(name=primary_key, value=reward_dict[primary_key]))
|
||||
for k, v in reward_dict.items():
|
||||
if k != primary_key:
|
||||
reward_dimensions.append(RewardDimension(name=k, value=v))
|
||||
else:
|
||||
if isinstance(reward, (int, bool)):
|
||||
reward = float(reward)
|
||||
elif not isinstance(reward, float): # pyright: ignore[reportUnnecessaryIsInstance]
|
||||
raise TypeError(f"Reward must be a number, got: {type(reward)}")
|
||||
reward_dimensions.append(RewardDimension(name="primary", value=reward))
|
||||
|
||||
return emit_annotation(
|
||||
{LightningSpanAttributes.REWARD.value: reward_dimensions, **(attributes or {})}, propagate=propagate
|
||||
)
|
||||
|
||||
|
||||
def get_reward_value(span: SpanLike) -> Optional[float]:
|
||||
"""Extract the reward value from a span, if available.
|
||||
|
||||
Args:
|
||||
span: Span object produced by AgentOps or Agent Lightning emitters.
|
||||
|
||||
Returns:
|
||||
The primary reward encoded in the span or `None` when the span does not represent a reward.
|
||||
"""
|
||||
# v0.3+ emit reward format
|
||||
reward_list = get_rewards_from_span(span)
|
||||
if reward_list:
|
||||
# Reward list is ordered and the first element is the primary reward
|
||||
return reward_list[0].value
|
||||
|
||||
for key in [
|
||||
"agentops.task.output", # newer versions of agentops
|
||||
"agentops.entity.output",
|
||||
]:
|
||||
reward_dict: Dict[str, Any] | None = None
|
||||
if span.attributes:
|
||||
output = span.attributes.get(key)
|
||||
if output:
|
||||
if isinstance(output, dict):
|
||||
reward_dict = cast(Dict[str, Any], output)
|
||||
elif isinstance(output, str):
|
||||
try:
|
||||
reward_dict = cast(Dict[str, Any], json.loads(output))
|
||||
except json.JSONDecodeError:
|
||||
reward_dict = None
|
||||
|
||||
if reward_dict and reward_dict.get("type") == "reward":
|
||||
reward_value = reward_dict.get("value", None)
|
||||
if reward_value is None:
|
||||
return None
|
||||
if not isinstance(reward_value, float):
|
||||
logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
|
||||
logger.warning(
|
||||
f"Extracted reward {reward_value} from AgentOps. This format is deprecated, please migrate to using `emit_reward`."
|
||||
)
|
||||
return cast(float, reward_value)
|
||||
|
||||
# v0.2 emit reward format
|
||||
if span.name == AGL_ANNOTATION and span.attributes:
|
||||
reward_value = span.attributes.get("reward", None)
|
||||
if reward_value is None:
|
||||
return None
|
||||
if not isinstance(reward_value, float):
|
||||
logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
|
||||
logger.warning(
|
||||
f"Extracted reward {reward_value} from a legacy version of reward span. You might have inconsistent agent-lightning versions."
|
||||
)
|
||||
return cast(float, reward_value)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_rewards_from_span(span: SpanLike) -> List[RewardPydanticModel]:
|
||||
"""Extract the reward as a list from a span, if available.
|
||||
|
||||
Args:
|
||||
span: Span object produced by AgentOps or Agent Lightning emitters.
|
||||
|
||||
Returns:
|
||||
A list of reward dimensions encoded in the span or an empty list when the span does not represent a reward.
|
||||
"""
|
||||
if span.attributes and any(key.startswith(LightningSpanAttributes.REWARD.value) for key in span.attributes):
|
||||
reward_attr = filter_and_unflatten_attributes(
|
||||
cast(Any, span.attributes or {}), LightningSpanAttributes.REWARD.value
|
||||
)
|
||||
recovered_rewards = TypeAdapter(List[RewardPydanticModel]).validate_python(reward_attr)
|
||||
return recovered_rewards
|
||||
else:
|
||||
return []
|
||||
|
||||
|
||||
def is_reward_span(span: SpanLike) -> bool:
|
||||
"""Return ``True`` when the provided span encodes a reward value."""
|
||||
maybe_reward = get_reward_value(span)
|
||||
return maybe_reward is not None
|
||||
|
||||
|
||||
def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
|
||||
"""Return all reward spans in the provided sequence.
|
||||
|
||||
Args:
|
||||
spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
|
||||
|
||||
Returns:
|
||||
List of spans that could be parsed as rewards.
|
||||
"""
|
||||
return [span for span in spans if is_reward_span(span)]
|
||||
|
||||
|
||||
def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
|
||||
"""Return the last reward value present in the provided spans.
|
||||
|
||||
Args:
|
||||
spans: Sequence containing [`ReadableSpan`](https://opentelemetry.io/docs/concepts/signals/traces/) objects or mocked span-like values.
|
||||
|
||||
Returns:
|
||||
Reward value from the latest reward span, or `None` when none are found.
|
||||
"""
|
||||
for span in reversed(spans):
|
||||
reward = get_reward_value(span)
|
||||
if reward is not None:
|
||||
return reward
|
||||
return None
|
||||
@@ -0,0 +1,156 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Environment variable managements."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import overload
|
||||
|
||||
__all__ = [
|
||||
"LightningEnvVar",
|
||||
"resolve_bool_env_var",
|
||||
"resolve_int_env_var",
|
||||
"resolve_str_env_var",
|
||||
]
|
||||
|
||||
|
||||
class LightningEnvVar(Enum):
|
||||
"""Environment variables for Agent Lightning."""
|
||||
|
||||
AGL_EMITTER_DEBUG = "AGL_EMITTER_DEBUG"
|
||||
"""Enable debug logging for the emitter."""
|
||||
|
||||
AGL_MANAGED_STORE = "AGL_MANAGED_STORE"
|
||||
"""If yes, the [`ExecutionStrategy`][agentlightning.ExecutionStrategy]
|
||||
constructs LightningStore wrappers automatically. When `False` the provided
|
||||
`store` is passed directly to the bundles, allowing callers to manage
|
||||
store wrappers manually."""
|
||||
|
||||
AGL_CURRENT_ROLE = "AGL_CURRENT_ROLE"
|
||||
"""Which side(s) to run in this process. Used in
|
||||
[`ClientServerExecutionStrategy`][agentlightning.ClientServerExecutionStrategy]."""
|
||||
|
||||
AGL_SERVER_HOST = "AGL_SERVER_HOST"
|
||||
"""Interface the [`LightningStoreServer`][agentlightning.LightningStoreServer]
|
||||
binds to when running the algorithm bundle locally."""
|
||||
|
||||
AGL_SERVER_PORT = "AGL_SERVER_PORT"
|
||||
"""Port the [`LightningStoreServer`][agentlightning.LightningStoreServer] listens to."""
|
||||
|
||||
|
||||
_TRUTHY_VALUES = {"1", "true", "yes", "on"}
|
||||
_FALSY_VALUES = {"0", "false", "no", "off"}
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_bool_env_var(env_var: LightningEnvVar, override: bool, fallback: bool) -> bool: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_bool_env_var(env_var: LightningEnvVar, *, fallback: bool) -> bool: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_bool_env_var(
|
||||
env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
|
||||
) -> bool | None: ...
|
||||
|
||||
|
||||
def resolve_bool_env_var(
|
||||
env_var: LightningEnvVar, override: bool | None = None, fallback: bool | None = None
|
||||
) -> bool | None:
|
||||
"""Resolve a boolean environment variable.
|
||||
|
||||
Args:
|
||||
env_var: The environment variable to resolve.
|
||||
override: Optional override supplied by the caller.
|
||||
fallback: Default value if the environment variable is not set.
|
||||
"""
|
||||
|
||||
if override is not None:
|
||||
return override
|
||||
|
||||
env_value = os.getenv(env_var.value)
|
||||
if env_value is None:
|
||||
return fallback
|
||||
|
||||
normalized = env_value.strip().lower()
|
||||
if normalized in _TRUTHY_VALUES:
|
||||
return True
|
||||
if normalized in _FALSY_VALUES:
|
||||
return False
|
||||
|
||||
raise ValueError(f"{env_var.value} must be one of {_TRUTHY_VALUES} or {_FALSY_VALUES}")
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_int_env_var(env_var: LightningEnvVar, override: int, fallback: int) -> int: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_int_env_var(env_var: LightningEnvVar, *, fallback: int) -> int: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_int_env_var(
|
||||
env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
|
||||
) -> int | None: ...
|
||||
|
||||
|
||||
def resolve_int_env_var(
|
||||
env_var: LightningEnvVar, override: int | None = None, fallback: int | None = None
|
||||
) -> int | None:
|
||||
"""Resolve an integer environment variable.
|
||||
|
||||
Args:
|
||||
env_var: The environment variable to resolve.
|
||||
override: Optional override supplied by the caller.
|
||||
fallback: Default value if the environment variable is not set.
|
||||
"""
|
||||
if override is not None:
|
||||
return override
|
||||
|
||||
env_value = os.getenv(env_var.value)
|
||||
if env_value is None:
|
||||
return fallback
|
||||
|
||||
try:
|
||||
return int(env_value)
|
||||
except ValueError:
|
||||
raise ValueError(f"{env_var.value} must be an integer")
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_str_env_var(env_var: LightningEnvVar, override: str, fallback: str) -> str: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_str_env_var(env_var: LightningEnvVar, *, fallback: str) -> str: ...
|
||||
|
||||
|
||||
@overload
|
||||
def resolve_str_env_var(
|
||||
env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
|
||||
) -> str | None: ...
|
||||
|
||||
|
||||
def resolve_str_env_var(
|
||||
env_var: LightningEnvVar, override: str | None = None, fallback: str | None = None
|
||||
) -> str | None:
|
||||
"""Resolve a string environment variable.
|
||||
|
||||
Args:
|
||||
env_var: The environment variable to resolve.
|
||||
override: Optional override supplied by the caller.
|
||||
fallback: Default value if the environment variable is not set.
|
||||
"""
|
||||
if override is not None:
|
||||
return override
|
||||
|
||||
env_value = os.getenv(env_var.value)
|
||||
if env_value is None:
|
||||
return fallback
|
||||
|
||||
return env_value
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .base import ExecutionStrategy
|
||||
from .client_server import ClientServerExecutionStrategy
|
||||
from .events import ExecutionEvent, MultiprocessingEvent, ThreadingEvent
|
||||
from .shared_memory import SharedMemoryExecutionStrategy
|
||||
|
||||
__all__ = [
|
||||
"ExecutionStrategy",
|
||||
"ClientServerExecutionStrategy",
|
||||
"ExecutionEvent",
|
||||
"ThreadingEvent",
|
||||
"MultiprocessingEvent",
|
||||
"SharedMemoryExecutionStrategy",
|
||||
]
|
||||
@@ -0,0 +1,64 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Protocol
|
||||
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
from .events import ExecutionEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AlgorithmBundle(Protocol):
|
||||
"""Callable bundle produced by [`Trainer`][agentlightning.Trainer].
|
||||
|
||||
Execution strategies treat the returned coroutine as opaque, only providing
|
||||
the shared store instance and cooperative stop event. Bundles typically
|
||||
encapsulate algorithm setup plus adapter and LLM proxy, etc.
|
||||
"""
|
||||
|
||||
async def __call__(self, store: LightningStore, event: ExecutionEvent) -> None:
|
||||
"""Execute algorithm logic using ``store`` until completion or stop."""
|
||||
|
||||
|
||||
class RunnerBundle(Protocol):
|
||||
"""Callable bundle wrapping runner setup and the worker loop, as opposed to the
|
||||
[`AlgorithmBundle`][agentlightning.AlgorithmBundle]."""
|
||||
|
||||
async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
|
||||
"""Execute runner logic for ``worker_id`` using ``store`` and ``event``."""
|
||||
|
||||
|
||||
class ExecutionStrategy:
|
||||
"""Coordinate algorithm and runner bundles within a single process abstraction.
|
||||
|
||||
Strategies decide how many worker bundles to launch, whether to communicate
|
||||
through shared memory or an HTTP boundary, and how to react to shutdown
|
||||
signals. They intentionally avoid inspecting the bundle internals; instead,
|
||||
each bundle remains responsible for its own scheduling semantics.
|
||||
|
||||
!!! note
|
||||
Implementations must honor the [execute()][agentlightning.ExecutionStrategy.execute]
|
||||
contract by propagating `KeyboardInterrupt` and ensuring resources are
|
||||
released when an error occurs on either side of the algorithm/runner
|
||||
pair.
|
||||
"""
|
||||
|
||||
def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
|
||||
"""Run the provided bundles using the configured orchestration model.
|
||||
|
||||
Args:
|
||||
algorithm: Callable bundle responsible for algorithm execution.
|
||||
runner: Callable bundle for runner workers.
|
||||
store: Concrete [`LightningStore`][agentlightning.LightningStore]
|
||||
shared across bundles.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must provide the orchestration
|
||||
implementation.
|
||||
"""
|
||||
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,443 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
from multiprocessing.context import BaseContext
|
||||
from typing import Callable, Iterable, Literal, cast
|
||||
|
||||
from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_int_env_var, resolve_str_env_var
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
|
||||
|
||||
from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
|
||||
from .events import ExecutionEvent, MultiprocessingEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ClientServerExecutionStrategy(ExecutionStrategy):
|
||||
"""Run algorithm and runner bundles as separate processes over HTTP.
|
||||
|
||||
Execution Roles:
|
||||
|
||||
- `"algorithm"`: Start [`LightningStoreServer`][agentlightning.LightningStoreServer]
|
||||
in-process and execute the algorithm bundle against it.
|
||||
- `"runner"`: Connect to an existing server with
|
||||
[`LightningStoreClient`][agentlightning.LightningStoreClient] and run the
|
||||
runner bundle locally (spawning multiple processes when requested).
|
||||
- `"both"`: Spawn runner processes first, then execute the algorithm and
|
||||
server on the same machine. This mode orchestrates the full loop locally.
|
||||
|
||||
When `role == "both"` you may choose which side runs on the main process
|
||||
via `main_process`. The runner-on-main option is limited to
|
||||
`n_runners == 1` because each additional runner requires its own event
|
||||
loop and process.
|
||||
|
||||
!!! warning
|
||||
When `main_process == "runner"` the algorithm and HTTP server execute
|
||||
in a child process. Store mutations remain isolated inside that process,
|
||||
so the original store instance passed to
|
||||
[execute()][agentlightning.ExecutionStrategy.execute] is not updated.
|
||||
|
||||
Abort Model (four-step escalation):
|
||||
|
||||
1. Cooperative stop. Every bundle receives a shared
|
||||
[`MultiprocessingEvent`][agentlightning.MultiprocessingEvent] (`stop_evt`).
|
||||
Any failure flips the event so peers can exit cleanly. Ctrl+C on the main
|
||||
process also sets the flag.
|
||||
2. KeyboardInterrupt synthesis. Remaining subprocesses receive ``SIGINT`` to
|
||||
trigger `KeyboardInterrupt` handlers.
|
||||
3. Termination. Stubborn processes are asked to ``terminate()``
|
||||
(`SIGTERM` on POSIX).
|
||||
4. Kill. As a last resort `kill()` is invoked (`SIGKILL` on POSIX).
|
||||
|
||||
This mirrors the semantics implemented in
|
||||
[`SharedMemoryExecutionStrategy`][agentlightning.SharedMemoryExecutionStrategy]
|
||||
but adapts them to multiple processes and the HTTP client/server boundary.
|
||||
"""
|
||||
|
||||
alias: str = "cs"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
role: Literal["algorithm", "runner", "both"] | None = None,
|
||||
server_host: str | None = None,
|
||||
server_port: int | None = None,
|
||||
n_runners: int = 1,
|
||||
graceful_timeout: float = 10.0,
|
||||
terminate_timeout: float = 10.0,
|
||||
main_process: Literal["algorithm", "runner"] = "algorithm",
|
||||
managed_store: bool | None = None,
|
||||
allowed_exit_codes: Iterable[int] = (0, -15),
|
||||
) -> None:
|
||||
"""Configure the strategy.
|
||||
|
||||
Args:
|
||||
role: Which side(s) to run in this process. When omitted, the
|
||||
`AGL_CURRENT_ROLE` environment variable is used.
|
||||
server_host: Interface the HTTP server binds to when running the
|
||||
algorithm bundle locally. Defaults to `AGL_SERVER_HOST`
|
||||
or `"localhost"` if unset.
|
||||
server_port: Port for the HTTP server in "algorithm"/"both" modes.
|
||||
Defaults to `AGL_SERVER_PORT` or `4747` if unset.
|
||||
n_runners: Number of runner processes to spawn in "runner"/"both".
|
||||
graceful_timeout: How long to wait (seconds) after setting the stop
|
||||
event before escalating to signals.
|
||||
terminate_timeout: How long to wait between escalation steps beyond
|
||||
the cooperative phase (re-used for SIGINT, terminate, and kill).
|
||||
main_process: Which bundle runs on the main process when
|
||||
`role == "both"`. `"runner"` requires `n_runners == 1` and is
|
||||
primarily intended for debugging.
|
||||
managed_store: When `True` (default) the strategy constructs
|
||||
LightningStore client/server wrappers automatically. When
|
||||
`False` the provided `store` is passed directly to the
|
||||
bundles, allowing callers to manage store wrappers manually.
|
||||
allowed_exit_codes: Allowed exit codes for subprocesses.
|
||||
By default, runner can exit gracefully with code 0 or terminated
|
||||
by SIGTERM (-15).
|
||||
"""
|
||||
resolved_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE, override=role, fallback="both")
|
||||
if resolved_role not in ("algorithm", "runner", "both"):
|
||||
raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
|
||||
self.role: Literal["algorithm", "runner", "both"] = resolved_role
|
||||
self.n_runners = n_runners
|
||||
self.server_host = resolve_str_env_var(
|
||||
LightningEnvVar.AGL_SERVER_HOST, override=server_host, fallback="localhost"
|
||||
)
|
||||
self.server_port = resolve_int_env_var(LightningEnvVar.AGL_SERVER_PORT, override=server_port, fallback=4747)
|
||||
self.graceful_timeout = graceful_timeout
|
||||
self.terminate_timeout = terminate_timeout
|
||||
if main_process not in ("algorithm", "runner"):
|
||||
raise ValueError("main_process must be 'algorithm' or 'runner'")
|
||||
if main_process == "runner":
|
||||
if self.role != "both":
|
||||
raise ValueError("main_process='runner' is only supported when role='both'")
|
||||
if n_runners != 1:
|
||||
raise ValueError("main_process='runner' requires n_runners to be 1")
|
||||
self.main_process = main_process
|
||||
self.managed_store = resolve_bool_env_var(
|
||||
LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
|
||||
)
|
||||
self.allowed_exit_codes = tuple(allowed_exit_codes)
|
||||
|
||||
async def _execute_algorithm(
|
||||
self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
|
||||
) -> None:
|
||||
wrapper_store: LightningStore | None = None
|
||||
if self.managed_store:
|
||||
logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
|
||||
wrapper_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
|
||||
server_started = False
|
||||
else:
|
||||
wrapper_store = store
|
||||
server_started = False
|
||||
|
||||
try:
|
||||
if self.managed_store and isinstance(wrapper_store, LightningStoreServer):
|
||||
await wrapper_store.start()
|
||||
server_started = True
|
||||
logger.debug("Algorithm bundle starting against endpoint %s", wrapper_store.endpoint)
|
||||
await algorithm(wrapper_store, stop_evt)
|
||||
logger.debug("Algorithm bundle completed successfully")
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Algorithm received CancelledError; signaling stop event")
|
||||
stop_evt.set()
|
||||
raise
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
|
||||
stop_evt.set()
|
||||
raise
|
||||
except BaseException:
|
||||
logger.exception("Algorithm bundle crashed; signaling stop event")
|
||||
stop_evt.set()
|
||||
raise
|
||||
finally:
|
||||
if self.managed_store and isinstance(wrapper_store, LightningStoreServer) and server_started:
|
||||
try:
|
||||
await wrapper_store.stop()
|
||||
except Exception:
|
||||
logger.exception("Error stopping LightningStore server")
|
||||
else:
|
||||
logger.debug("LightningStore server shutdown completed")
|
||||
|
||||
async def _execute_runner(
|
||||
self,
|
||||
runner: RunnerBundle,
|
||||
worker_id: int,
|
||||
store: LightningStore,
|
||||
stop_evt: ExecutionEvent,
|
||||
) -> None:
|
||||
if self.managed_store:
|
||||
# If managed, we actually do not use the provided store
|
||||
client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
|
||||
else:
|
||||
client_store = store
|
||||
try:
|
||||
if self.managed_store:
|
||||
logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
|
||||
else:
|
||||
logger.debug("Runner %s executing with provided store", worker_id)
|
||||
await runner(client_store, worker_id, stop_evt)
|
||||
logger.debug("Runner %s completed successfully", worker_id)
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Runner %s received CancelledError; signaling stop event", worker_id)
|
||||
stop_evt.set()
|
||||
raise
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
|
||||
stop_evt.set()
|
||||
raise
|
||||
except BaseException:
|
||||
logger.exception("Runner %s crashed; signaling stop event", worker_id)
|
||||
stop_evt.set()
|
||||
raise
|
||||
finally:
|
||||
if self.managed_store and isinstance(client_store, LightningStoreClient):
|
||||
try:
|
||||
await client_store.close()
|
||||
except Exception:
|
||||
logger.exception("Error closing LightningStore client for runner %s", worker_id)
|
||||
else:
|
||||
logger.debug("Runner %s closed LightningStore client", worker_id)
|
||||
|
||||
def _spawn_runners(
|
||||
self,
|
||||
runner: RunnerBundle,
|
||||
store: LightningStore,
|
||||
stop_evt: ExecutionEvent,
|
||||
*,
|
||||
ctx: BaseContext,
|
||||
) -> list[multiprocessing.Process]:
|
||||
"""Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
|
||||
processes: list[multiprocessing.Process] = []
|
||||
|
||||
def _runner_sync(runner: RunnerBundle, worker_id: int, store: LightningStore, stop_evt: ExecutionEvent) -> None:
|
||||
# Runners are executed in child processes; each process owns its own
|
||||
# event loop to keep the asyncio scheduler isolated.
|
||||
try:
|
||||
asyncio.run(self._execute_runner(runner, worker_id, store, stop_evt))
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Runner (asyncio) %s received KeyboardInterrupt; exiting gracefully", worker_id)
|
||||
except BaseException as exc:
|
||||
logger.exception("Runner (asyncio) %s crashed by %s; signaling stop event", worker_id, exc)
|
||||
raise
|
||||
|
||||
for i in range(self.n_runners):
|
||||
process = cast(
|
||||
multiprocessing.Process,
|
||||
ctx.Process(target=_runner_sync, args=(runner, i, store, stop_evt), name=f"runner-{i}"), # type: ignore
|
||||
)
|
||||
process.start()
|
||||
logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
|
||||
processes.append(process)
|
||||
|
||||
return processes
|
||||
|
||||
def _spawn_algorithm_process(
|
||||
self,
|
||||
algorithm: AlgorithmBundle,
|
||||
store: LightningStore,
|
||||
stop_evt: ExecutionEvent,
|
||||
*,
|
||||
ctx: BaseContext,
|
||||
) -> multiprocessing.Process:
|
||||
"""Used when `main_process == "runner"`."""
|
||||
|
||||
def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
|
||||
try:
|
||||
asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("Algorithm (asyncio.run) received KeyboardInterrupt; exiting gracefully")
|
||||
except BaseException as exc:
|
||||
logger.exception("Algorithm (asyncio.run) crashed by %s; signaling stop event", exc)
|
||||
raise
|
||||
|
||||
process = cast(
|
||||
multiprocessing.Process,
|
||||
ctx.Process(target=_algorithm_sync, args=(algorithm, store, stop_evt), name="algorithm"), # type: ignore
|
||||
)
|
||||
process.start()
|
||||
logger.debug("Spawned algorithm process %s (pid=%s)", process.name, process.pid)
|
||||
return process
|
||||
|
||||
def _join_until_deadline(
|
||||
self,
|
||||
processes: Iterable[multiprocessing.Process],
|
||||
timeout: float,
|
||||
) -> list[multiprocessing.Process]:
|
||||
"""Join ``processes`` until ``timeout`` elapses, returning those still alive."""
|
||||
deadline = time.monotonic() + timeout
|
||||
still_alive: list[multiprocessing.Process] = []
|
||||
for process in processes:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining > 0:
|
||||
process.join(remaining)
|
||||
else:
|
||||
process.join(0)
|
||||
if process.is_alive():
|
||||
still_alive.append(process)
|
||||
return still_alive
|
||||
|
||||
def _signal_processes(
|
||||
self,
|
||||
processes: Iterable[multiprocessing.Process],
|
||||
action: Callable[[multiprocessing.Process], None],
|
||||
) -> None:
|
||||
"""Invoke ``action`` on each process while suppressing individual failures."""
|
||||
for process in processes:
|
||||
try:
|
||||
action(process)
|
||||
except Exception:
|
||||
logger.exception("Error signaling process %s (pid=%s)", process.name, process.pid)
|
||||
|
||||
def _shutdown_processes(
|
||||
self,
|
||||
processes: list[multiprocessing.Process],
|
||||
stop_evt: ExecutionEvent,
|
||||
) -> None:
|
||||
"""4-step escalation shutdown of ``processes``."""
|
||||
if not processes:
|
||||
logger.debug("No subprocesses to shutdown")
|
||||
return
|
||||
|
||||
if not stop_evt.is_set():
|
||||
logger.debug("Sending cooperative stop signal to subprocesses")
|
||||
stop_evt.set()
|
||||
else:
|
||||
logger.debug("Stop event already set; waiting for subprocesses to exit")
|
||||
|
||||
alive = self._join_until_deadline(processes, self.graceful_timeout)
|
||||
if not alive:
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"Subprocesses still alive after cooperative wait; sending SIGINT to %s",
|
||||
", ".join(p.name or str(p.pid) for p in alive),
|
||||
)
|
||||
# SIGINT is not reliable on Windows, but we do not consider such case yet.
|
||||
self._signal_processes(alive, lambda p: os.kill(cast(int, p.pid), signal.SIGINT))
|
||||
alive = self._join_until_deadline(alive, self.terminate_timeout)
|
||||
if not alive:
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"Subprocesses still alive after SIGINT wait; sending terminate() to %s",
|
||||
", ".join(p.name or str(p.pid) for p in alive),
|
||||
)
|
||||
self._signal_processes(alive, lambda p: p.terminate())
|
||||
|
||||
alive = self._join_until_deadline(alive, self.terminate_timeout)
|
||||
if not alive:
|
||||
return
|
||||
|
||||
logger.error(
|
||||
"Subprocesses still alive after terminate(); sending kill() to %s",
|
||||
", ".join(p.name or str(p.pid) for p in alive),
|
||||
)
|
||||
self._signal_processes(alive, lambda p: p.kill())
|
||||
alive = self._join_until_deadline(alive, self.terminate_timeout)
|
||||
|
||||
if alive:
|
||||
logger.error(
|
||||
"Subprocesses failed to exit even after kill(): %s", ", ".join(p.name or str(p.pid) for p in alive)
|
||||
)
|
||||
|
||||
def _check_process_exitcodes(self, processes: Iterable[multiprocessing.Process]) -> None:
|
||||
"""Raise an error if any managed process exited with a non-zero status."""
|
||||
failed = [p for p in processes if p.exitcode not in self.allowed_exit_codes + (None,)]
|
||||
if failed:
|
||||
formatted = ", ".join(f"{p.name or p.pid} (exitcode={p.exitcode})" for p in failed)
|
||||
raise RuntimeError(f"Subprocesses failed with unexpected exit codes: {formatted}")
|
||||
|
||||
def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
|
||||
logger.info(
|
||||
"Starting client-server execution with %d runner(s) [role=%s, main_process=%s]",
|
||||
self.n_runners,
|
||||
self.role,
|
||||
self.main_process,
|
||||
)
|
||||
|
||||
# Re-use the active multiprocessing context so the event and processes
|
||||
# agree on the start method (fork/spawn/forkserver).
|
||||
ctx = multiprocessing.get_context()
|
||||
stop_evt = MultiprocessingEvent(ctx=ctx)
|
||||
# Track spawned processes so we can enforce termination ordering and
|
||||
# surface non-zero exit codes back to the caller.
|
||||
processes: list[multiprocessing.Process] = []
|
||||
|
||||
exception: BaseException | None = None
|
||||
keyboard_interrupt = False
|
||||
|
||||
try:
|
||||
if self.role == "algorithm":
|
||||
logger.info("Running algorithm solely...")
|
||||
asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
|
||||
elif self.role == "runner":
|
||||
if self.n_runners == 1:
|
||||
logger.info("Running runner solely...")
|
||||
asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
|
||||
else:
|
||||
logger.info("Spawning runner processes...")
|
||||
processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
|
||||
# Wait for the processes to finish naturally.
|
||||
for process in processes:
|
||||
process.join()
|
||||
self._check_process_exitcodes(processes)
|
||||
elif self.role == "both":
|
||||
if self.main_process == "algorithm":
|
||||
logger.info("Spawning runner processes...")
|
||||
processes = self._spawn_runners(runner, store, stop_evt, ctx=ctx)
|
||||
try:
|
||||
logger.info("Running algorithm...")
|
||||
asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
|
||||
finally:
|
||||
# Always request the runner side to unwind once the
|
||||
# algorithm/server portion finishes (successfully or not).
|
||||
stop_evt.set()
|
||||
else: # main_process == "runner"
|
||||
if self.n_runners > 1:
|
||||
raise ValueError("main_process='runner' requires n_runners to be 1")
|
||||
|
||||
logger.info("Spawning algorithm process...")
|
||||
algorithm_process = self._spawn_algorithm_process(algorithm, store, stop_evt, ctx=ctx)
|
||||
processes = [algorithm_process]
|
||||
|
||||
# Run the lone runner cooperatively in-process so users can
|
||||
# attach a debugger. The algorithm + HTTP server live in
|
||||
# the background process spawned above (the provided
|
||||
# store must therefore be picklable when using spawn).
|
||||
logger.info("Running runner...")
|
||||
asyncio.run(self._execute_runner(runner, 0, store, stop_evt))
|
||||
|
||||
# Wait for the algorithm process to finish.
|
||||
algorithm_process.join()
|
||||
else:
|
||||
raise ValueError(f"Unknown role: {self.role}")
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("KeyboardInterrupt received; initiating shutdown")
|
||||
stop_evt.set()
|
||||
keyboard_interrupt = True
|
||||
except BaseException as exc:
|
||||
logger.exception("Unhandled exception in execute method")
|
||||
stop_evt.set()
|
||||
# Preserve the original exception so we can avoid masking it during
|
||||
# the cleanup phase.
|
||||
exception = exc
|
||||
raise
|
||||
finally:
|
||||
logger.info("Shutting down subprocesses")
|
||||
self._shutdown_processes(processes, stop_evt)
|
||||
if processes:
|
||||
try:
|
||||
self._check_process_exitcodes(processes)
|
||||
except RuntimeError as err:
|
||||
if exception is not None or keyboard_interrupt:
|
||||
# We already propagate/handled a different failure, so
|
||||
# emit a warning instead of raising a secondary error.
|
||||
logger.warning("Subprocesses ended abnormally during shutdown: %s", err)
|
||||
else:
|
||||
raise
|
||||
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import multiprocessing as mp
|
||||
import threading
|
||||
from multiprocessing.context import BaseContext
|
||||
from typing import Optional, Protocol
|
||||
|
||||
|
||||
class ExecutionEvent(Protocol):
|
||||
"""Protocol capturing the cooperative stop contract shared by strategies.
|
||||
|
||||
Implementations mirror the API of ``threading.Event`` and
|
||||
``multiprocessing.Event`` so the rest of the execution layer can remain
|
||||
agnostic to the underlying concurrency primitive.
|
||||
|
||||
Methods:
|
||||
|
||||
set: Signal cancellation. The call must be idempotent.
|
||||
clear: Reset the event to the unsignaled state.
|
||||
is_set: Return ``True`` when cancellation has been requested.
|
||||
wait: Block until the event is signaled or an optional timeout elapses.
|
||||
"""
|
||||
|
||||
def set(self) -> None: ...
|
||||
def clear(self) -> None: ...
|
||||
def is_set(self) -> bool: ...
|
||||
def wait(self, timeout: Optional[float] = None) -> bool: ...
|
||||
|
||||
|
||||
class ThreadingEvent:
|
||||
"""Thread-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
|
||||
|
||||
__slots__ = ("_evt",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._evt = threading.Event()
|
||||
|
||||
def set(self) -> None:
|
||||
self._evt.set()
|
||||
|
||||
def clear(self) -> None:
|
||||
self._evt.clear()
|
||||
|
||||
def is_set(self) -> bool:
|
||||
return self._evt.is_set()
|
||||
|
||||
def wait(self, timeout: Optional[float] = None) -> bool:
|
||||
return self._evt.wait(timeout)
|
||||
|
||||
|
||||
class MultiprocessingEvent:
|
||||
"""Process-safe implementation of [`ExecutionEvent`][agentlightning.ExecutionEvent]."""
|
||||
|
||||
__slots__ = ("_evt",)
|
||||
|
||||
def __init__(self, *, ctx: Optional[BaseContext] = None) -> None:
|
||||
self._evt = (ctx or mp).Event()
|
||||
|
||||
def set(self) -> None:
|
||||
self._evt.set()
|
||||
|
||||
def clear(self) -> None:
|
||||
self._evt.clear()
|
||||
|
||||
def is_set(self) -> bool:
|
||||
return self._evt.is_set()
|
||||
|
||||
def wait(self, timeout: Optional[float] = None) -> bool:
|
||||
return self._evt.wait(timeout)
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .base import ExecutionStrategy
|
||||
|
||||
|
||||
class InterProcessExecutionStrategy(ExecutionStrategy):
|
||||
"""Placeholder strategy for future inter-process primitives.
|
||||
|
||||
The class exists to reserve the `ipc` alias and make the planned
|
||||
implementation discoverable. Attempting to use it today will raise
|
||||
`NotImplementedError` once the execution contract is finalized.
|
||||
"""
|
||||
|
||||
alias: str = "ipc"
|
||||
|
||||
# TODO: to be implemented
|
||||
@@ -0,0 +1,282 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from contextlib import suppress
|
||||
from queue import SimpleQueue
|
||||
from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
|
||||
|
||||
from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.store.threading import LightningStoreThreaded
|
||||
|
||||
from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
|
||||
from .events import ExecutionEvent, ThreadingEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SharedMemoryExecutionStrategy(ExecutionStrategy):
|
||||
"""Execute bundles in a single process with cooperative worker threads.
|
||||
|
||||
Stop Model:
|
||||
|
||||
- All bundles share one [`ThreadingEvent`][agentlightning.ThreadingEvent]
|
||||
named `stop_evt`.
|
||||
- Only the main thread receives `KeyboardInterrupt`. When Ctrl+C occurs we
|
||||
set `stop_evt`.
|
||||
- Any exception raised inside a bundle sets `stop_evt` so other threads can
|
||||
unwind cooperatively.
|
||||
- Once the bundle running on the main thread exits successfully the
|
||||
treatment depends on `main_thread`:
|
||||
- `"algorithm"`: the runners are asked to stop by setting `stop_evt`.
|
||||
- `"runner"`: the algorithm keeps running until it exits naturally.
|
||||
- Background threads are marked as daemons. We join them briefly and log any
|
||||
stragglers before shutting down.
|
||||
|
||||
!!! note
|
||||
Signals other than `SIGINT` (such as `SIGTERM`) are not intercepted;
|
||||
Python's default behavior for those signals is preserved.
|
||||
"""
|
||||
|
||||
alias: str = "shm"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_runners: int = 1,
|
||||
main_thread: Literal["algorithm", "runner"] = "runner",
|
||||
join_timeout: float = 15.0,
|
||||
graceful_delay: float = 5.0,
|
||||
poll_interval: float = 0.05,
|
||||
managed_store: bool | None = None,
|
||||
) -> None:
|
||||
if main_thread not in ("algorithm", "runner"):
|
||||
raise ValueError("main_thread must be 'algorithm' or 'runner'")
|
||||
if main_thread == "runner" and n_runners != 1:
|
||||
raise ValueError(
|
||||
"When main_thread is 'runner', n_runners must be 1. "
|
||||
"Either use 'algorithm' on the main thread or set n_runners to 1."
|
||||
)
|
||||
self.n_runners = n_runners
|
||||
self.main_thread = main_thread
|
||||
self.join_timeout = join_timeout
|
||||
self.graceful_delay = graceful_delay
|
||||
self.poll_interval = poll_interval
|
||||
self.managed_store = resolve_bool_env_var(
|
||||
LightningEnvVar.AGL_MANAGED_STORE, override=managed_store, fallback=True
|
||||
)
|
||||
|
||||
async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
|
||||
"""Run `coro` until it finishes or a cooperative stop is requested.
|
||||
|
||||
Control flow:
|
||||
|
||||
1. Start the bundle coroutine as `task`.
|
||||
2. Launch a watcher that polls `stop_evt` without blocking the loop.
|
||||
3. When the stop event flips:
|
||||
a. Give the bundle `graceful_delay` seconds to finish on its own,
|
||||
because well-behaved bundles will check the event and return.
|
||||
b. Cancel the bundle task if it is still running after the grace
|
||||
period.
|
||||
4. Await both tasks and swallow `CancelledError` where appropriate.
|
||||
|
||||
This is a *backup* mechanism for bundles that might not poll the event
|
||||
frequently; cooperative shutdown (checking `stop_evt` inside the
|
||||
bundle) remains the preferred approach.
|
||||
"""
|
||||
task: asyncio.Task[Any] = asyncio.create_task(coro) # type: ignore
|
||||
task_exception: Optional[BaseException] = None
|
||||
|
||||
async def watcher() -> None:
|
||||
# Poll the threading event without blocking the event loop. Using a
|
||||
# background thread via ``asyncio.to_thread`` makes cancellation
|
||||
# difficult because ``ThreadingEvent.wait`` is not interruptible.
|
||||
# Instead we cooperatively check the flag from the loop so the
|
||||
# watcher task stays cancellable and tests don't hang when the
|
||||
# bundle finishes naturally before the stop event is set.
|
||||
while not stop_evt.is_set():
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
# Grace period: let a cooperative bundle exit on its own.
|
||||
try:
|
||||
# At this point of waiting, the main task should already see the stop event.
|
||||
await asyncio.wait_for(asyncio.shield(task), timeout=self.graceful_delay) # type: ignore
|
||||
logger.debug("Bundle finished by itself during grace period.")
|
||||
return # bundle finished by itself during grace period
|
||||
except asyncio.TimeoutError:
|
||||
# Still running after the grace window.
|
||||
pass
|
||||
except asyncio.CancelledError:
|
||||
# If someone else canceled the task already, we're done.
|
||||
logger.debug("Bundle already canceled by someone else; exiting watcher.")
|
||||
return
|
||||
|
||||
# Still running after the grace window: cancel it.
|
||||
if not task.done():
|
||||
logger.debug("Graceful delay elapsed; canceling bundle task...")
|
||||
task.cancel()
|
||||
|
||||
watcher_task = asyncio.create_task(watcher())
|
||||
result: Any = None
|
||||
|
||||
try:
|
||||
# We don't wait on FIRST_COMPLETED here, because we want the watcher
|
||||
# to be able to grant a grace window after stop_evt flips.
|
||||
await asyncio.wait(
|
||||
{task, watcher_task}, return_when=asyncio.FIRST_COMPLETED
|
||||
) # pyright: ignore[reportUnknownArgumentType]
|
||||
finally:
|
||||
# If the main task hasn't completed yet (e.g., watcher scheduled cancel),
|
||||
# finish the cancellation handshake.
|
||||
if not task.done():
|
||||
try:
|
||||
await asyncio.wait_for(task, timeout=self.graceful_delay) # second chance
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
"Bundle task did not stop after cancellation; abandoning task."
|
||||
"This thread could live until the process exits."
|
||||
)
|
||||
# We return without awaiting it. asyncio.run will still try to cancel
|
||||
# pending tasks on loop close; if the task ignores cancellation, this
|
||||
# thread may still stick. It's the best we can do in Python.
|
||||
# We don't raise an exception here, but the thread could be a zombie.
|
||||
return result
|
||||
else:
|
||||
# Task completed naturally; retrieve result.
|
||||
try:
|
||||
result = await task # type: ignore
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except BaseException as exc:
|
||||
task_exception = exc
|
||||
|
||||
watcher_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await watcher_task
|
||||
|
||||
if task_exception is not None:
|
||||
raise task_exception
|
||||
|
||||
return result # type: ignore
|
||||
|
||||
def _run_algorithm(
|
||||
self,
|
||||
algorithm: AlgorithmBundle,
|
||||
store: LightningStore,
|
||||
stop_evt: ExecutionEvent,
|
||||
thread_exceptions: Optional[SimpleQueue[BaseException]],
|
||||
) -> None:
|
||||
try:
|
||||
asyncio.run(self._run_until_completed_or_canceled(algorithm(store, stop_evt), stop_evt))
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Algorithm bundle canceled due to stop signal.")
|
||||
except BaseException as exc:
|
||||
logger.exception("Algorithm bundle crashed; signaling stop to others.")
|
||||
if thread_exceptions is not None:
|
||||
thread_exceptions.put(exc)
|
||||
stop_evt.set()
|
||||
raise
|
||||
|
||||
def _run_runner(
|
||||
self,
|
||||
runner: RunnerBundle,
|
||||
store: LightningStore,
|
||||
worker_id: int,
|
||||
stop_evt: ExecutionEvent,
|
||||
thread_exceptions: Optional[SimpleQueue[BaseException]],
|
||||
) -> None:
|
||||
try:
|
||||
asyncio.run(self._run_until_completed_or_canceled(runner(store, worker_id, stop_evt), stop_evt))
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Runner bundle (worker_id=%s) canceled due to stop signal.", worker_id)
|
||||
except BaseException as exc:
|
||||
logger.exception("Runner bundle crashed (worker_id=%s); signaling stop to others.", worker_id)
|
||||
if thread_exceptions is not None:
|
||||
thread_exceptions.put(exc)
|
||||
stop_evt.set()
|
||||
raise
|
||||
|
||||
def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
|
||||
logger.info(
|
||||
"Starting shm execution with %d runner(s); main thread runs '%s'",
|
||||
self.n_runners,
|
||||
self.main_thread,
|
||||
)
|
||||
|
||||
# Create stop event and thread-safe store.
|
||||
stop_evt = ThreadingEvent()
|
||||
if self.managed_store:
|
||||
thread_safe_store = LightningStoreThreaded(store)
|
||||
else:
|
||||
thread_safe_store = store
|
||||
|
||||
thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
|
||||
raised_from_thread: Optional[BaseException] = None
|
||||
|
||||
def make_thread(name: str, target: Callable[..., Any], args: Tuple[Any, ...]) -> threading.Thread:
|
||||
t = threading.Thread(name=name, target=target, args=args, daemon=True)
|
||||
t.start()
|
||||
return t
|
||||
|
||||
threads: List[threading.Thread] = []
|
||||
|
||||
try:
|
||||
if self.main_thread == "algorithm":
|
||||
# Start runner threads; algorithm runs on main thread.
|
||||
for i in range(self.n_runners):
|
||||
thread = make_thread(
|
||||
name=f"runner-{i}",
|
||||
target=self._run_runner,
|
||||
args=(runner, thread_safe_store, i, stop_evt, thread_exceptions),
|
||||
)
|
||||
threads.append(thread)
|
||||
|
||||
# Ctrl+C here raises KeyboardInterrupt on this stack.
|
||||
# Main thread doesn't need to collect exceptions.
|
||||
self._run_algorithm(algorithm, thread_safe_store, stop_evt, None)
|
||||
|
||||
# If algo finishes naturally, request runners to stop.
|
||||
stop_evt.set()
|
||||
|
||||
else: # main_thread == "runner"
|
||||
# Start algorithm in background; runner runs on main thread.
|
||||
thread = make_thread(
|
||||
name="algorithm",
|
||||
target=self._run_algorithm,
|
||||
args=(algorithm, thread_safe_store, stop_evt, thread_exceptions),
|
||||
)
|
||||
threads.append(thread)
|
||||
|
||||
# Ctrl+C here raises KeyboardInterrupt on this stack.
|
||||
# Main thread doesn't need to collect exceptions.
|
||||
self._run_runner(runner, thread_safe_store, 0, stop_evt, None)
|
||||
|
||||
# If runner finishes naturally, WAIT FOR ALGORITHM TO FINISH.
|
||||
thread.join()
|
||||
|
||||
if not thread_exceptions.empty():
|
||||
raised_from_thread = thread_exceptions.get()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.warning("KeyboardInterrupt received on main thread; initiating cooperative shutdown...")
|
||||
stop_evt.set()
|
||||
finally:
|
||||
# Attempt a clean join; if some threads don't comply, log and move on.
|
||||
for t in threads:
|
||||
logger.debug("Joining thread %s...", t.name)
|
||||
t.join(timeout=self.join_timeout)
|
||||
|
||||
alive = [t.name for t in threads if t.is_alive()]
|
||||
if alive:
|
||||
logger.error(
|
||||
"Threads still alive after %.1fs: %s. They are daemons; continuing shutdown.",
|
||||
self.join_timeout,
|
||||
", ".join(alive),
|
||||
)
|
||||
|
||||
if raised_from_thread is None and not thread_exceptions.empty():
|
||||
raised_from_thread = thread_exceptions.get()
|
||||
|
||||
if raised_from_thread is not None:
|
||||
raise raised_from_thread
|
||||
@@ -0,0 +1,114 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import warnings
|
||||
|
||||
AGENTOPS_INSTALLED: bool = False
|
||||
AGENTOPS_LANGCHAIN_INSTALLED: bool = False
|
||||
LITELLM_INSTALLED: bool = False
|
||||
VLLM_INSTALLED: bool = False
|
||||
WEAVE_INSTALLED: bool = False
|
||||
|
||||
try:
|
||||
from . import agentops # type: ignore
|
||||
|
||||
AGENTOPS_INSTALLED = True # type: ignore
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from . import litellm # type: ignore
|
||||
|
||||
LITELLM_INSTALLED = True # type: ignore
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# MAGIC! DO NOT TOUCH THIS!
|
||||
# vllm import will cause reward tracing function to fail and produce nothing.
|
||||
# try:
|
||||
# from . import vllm
|
||||
|
||||
# VLLM_INSTALLED = True
|
||||
# except ImportError:
|
||||
# pass
|
||||
|
||||
|
||||
try:
|
||||
from . import agentops_langchain # type: ignore
|
||||
|
||||
AGENTOPS_LANGCHAIN_INSTALLED = True # type: ignore
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def instrument_all():
|
||||
"""Instrument all the instrumentation libraries."""
|
||||
if AGENTOPS_INSTALLED:
|
||||
from .agentops import instrument_agentops
|
||||
|
||||
instrument_agentops()
|
||||
else:
|
||||
warnings.warn("agentops is not installed. It's therefore not instrumented.")
|
||||
|
||||
if LITELLM_INSTALLED:
|
||||
from .litellm import instrument_litellm
|
||||
|
||||
instrument_litellm()
|
||||
else:
|
||||
warnings.warn("litellm is not installed. It's therefore not instrumented.")
|
||||
|
||||
if VLLM_INSTALLED:
|
||||
from .vllm import instrument_vllm
|
||||
|
||||
instrument_vllm()
|
||||
else:
|
||||
warnings.warn("vllm is not installed. It's therefore not instrumented.")
|
||||
|
||||
if AGENTOPS_LANGCHAIN_INSTALLED:
|
||||
from .agentops_langchain import instrument_agentops_langchain
|
||||
|
||||
instrument_agentops_langchain()
|
||||
else:
|
||||
warnings.warn("Agentops-langchain integration is not installed. It's therefore not instrumented.")
|
||||
|
||||
|
||||
def uninstrument_all():
|
||||
"""Uninstrument all the instrumentation libraries."""
|
||||
if AGENTOPS_INSTALLED:
|
||||
try:
|
||||
from .agentops import uninstrument_agentops
|
||||
|
||||
uninstrument_agentops()
|
||||
except ImportError:
|
||||
warnings.warn("agentops is installed but uninstrument_agentops could not be imported.")
|
||||
else:
|
||||
warnings.warn("agentops is not installed. It's therefore not uninstrumented.")
|
||||
|
||||
if LITELLM_INSTALLED:
|
||||
try:
|
||||
from .litellm import uninstrument_litellm
|
||||
|
||||
uninstrument_litellm()
|
||||
except ImportError:
|
||||
warnings.warn("litellm is installed but uninstrument_litellm could not be imported.")
|
||||
else:
|
||||
warnings.warn("litellm is not installed. It's therefore not uninstrumented.")
|
||||
|
||||
if VLLM_INSTALLED:
|
||||
try:
|
||||
from .vllm import uninstrument_vllm
|
||||
|
||||
uninstrument_vllm()
|
||||
except ImportError:
|
||||
warnings.warn("vllm is installed but uninstrument_vllm could not be imported.")
|
||||
else:
|
||||
warnings.warn("vllm is not installed. It's therefore not uninstrumented.")
|
||||
|
||||
if AGENTOPS_LANGCHAIN_INSTALLED:
|
||||
try:
|
||||
from .agentops_langchain import uninstrument_agentops_langchain
|
||||
|
||||
uninstrument_agentops_langchain()
|
||||
except ImportError:
|
||||
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.")
|
||||
@@ -0,0 +1,314 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Callable, no_type_check
|
||||
|
||||
import requests
|
||||
from agentops.client.api import V3Client, V4Client
|
||||
from agentops.client.api.types import AuthTokenResponse
|
||||
from agentops.sdk.exporters import AuthenticatedOTLPExporter
|
||||
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.metrics.export import MetricExportResult
|
||||
|
||||
from agentlightning.utils.otlp import LightningStoreOTLPExporter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"instrument_agentops",
|
||||
"uninstrument_agentops",
|
||||
]
|
||||
|
||||
# Module-level storage for originals
|
||||
_original_handle_chat_attributes: Callable[..., Any] | None = None
|
||||
_original_handle_response: Callable[..., Any] | None = None
|
||||
_agentops_service_enabled = False
|
||||
|
||||
|
||||
def enable_agentops_service(enabled: bool = True) -> None:
|
||||
"""
|
||||
Enable or disable communication with the AgentOps service.
|
||||
|
||||
By default, AgentOps exporters and clients will run in local mode
|
||||
and will NOT attempt to communicate with the remote AgentOps service.
|
||||
|
||||
Args:
|
||||
enabled: If True, enable all AgentOps exporters and clients.
|
||||
All exporters and clients will operate in normal mode and send data
|
||||
to the [AgentOps service](https://www.agentops.ai).
|
||||
"""
|
||||
global _agentops_service_enabled
|
||||
_agentops_service_enabled = enabled
|
||||
logger.info(f"AgentOps service enabled is set to {enabled}.")
|
||||
|
||||
|
||||
def _patch_exporters():
|
||||
import agentops.client.api
|
||||
import agentops.sdk.core
|
||||
|
||||
agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter # type: ignore
|
||||
agentops.sdk.core.OTLPMetricExporter = BypassableOTLPMetricExporter
|
||||
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
|
||||
agentops.sdk.core.OTLPSpanExporter = BypassableOTLPSpanExporter # type: ignore
|
||||
agentops.client.api.V3Client = BypassableV3Client
|
||||
agentops.client.api.V4Client = BypassableV4Client
|
||||
|
||||
|
||||
def _unpatch_exporters():
|
||||
import agentops.client.api
|
||||
import agentops.sdk.core
|
||||
|
||||
agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter # type: ignore
|
||||
agentops.sdk.core.OTLPMetricExporter = OTLPMetricExporter
|
||||
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
|
||||
agentops.sdk.core.OTLPSpanExporter = OTLPSpanExporter # type: ignore
|
||||
agentops.client.api.V3Client = V3Client
|
||||
agentops.client.api.V4Client = V4Client
|
||||
|
||||
|
||||
def _unwrap_legacy_response(response: Any) -> Any:
|
||||
if hasattr(response, "parse") and callable(response.parse):
|
||||
return response.parse()
|
||||
return response
|
||||
|
||||
|
||||
def _patch_new_agentops():
|
||||
import agentops.instrumentation.providers.openai.stream_wrapper
|
||||
import agentops.instrumentation.providers.openai.wrappers.chat
|
||||
from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes # type: ignore
|
||||
|
||||
global _original_handle_chat_attributes
|
||||
|
||||
if _original_handle_chat_attributes is not None:
|
||||
logger.warning("AgentOps already patched. Skipping.")
|
||||
return True
|
||||
|
||||
_original_handle_chat_attributes = handle_chat_attributes # type: ignore
|
||||
|
||||
@no_type_check
|
||||
def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws): # type: ignore
|
||||
attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
|
||||
|
||||
# In some cases, response is a openai._legacy_response.LegacyAPIResponse (e.g., LiteLLM, or LangChain),
|
||||
# This is created by client.with_raw_response.create()
|
||||
return_value = _unwrap_legacy_response(return_value)
|
||||
|
||||
if (
|
||||
return_value is not None
|
||||
and hasattr(return_value, "prompt_token_ids")
|
||||
and return_value.prompt_token_ids is not None
|
||||
):
|
||||
attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
|
||||
if (
|
||||
return_value is not None
|
||||
and hasattr(return_value, "response_token_ids")
|
||||
and return_value.response_token_ids is not None
|
||||
):
|
||||
attributes["response_token_ids"] = list(return_value.response_token_ids[0])
|
||||
|
||||
# For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
|
||||
if (
|
||||
return_value is not None
|
||||
and hasattr(return_value, "choices")
|
||||
and return_value.choices
|
||||
and isinstance(return_value.choices, list)
|
||||
and len(return_value.choices) > 0
|
||||
):
|
||||
first_choice = return_value.choices[0]
|
||||
# Token IDs from "choices[0].token_ids"
|
||||
if "response_token_ids" not in attributes:
|
||||
if hasattr(first_choice, "token_ids") and first_choice.token_ids is not None:
|
||||
attributes["response_token_ids"] = list(first_choice.token_ids)
|
||||
# newer versions of OpenAI client SDK
|
||||
elif (
|
||||
hasattr(first_choice, "provider_specific_fields")
|
||||
and first_choice.provider_specific_fields.get("token_ids") is not None
|
||||
):
|
||||
attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])
|
||||
|
||||
# log probability
|
||||
# This is temporary. We need a unified convention for classifying and naming logprobs.
|
||||
if hasattr(first_choice, "logprobs") and first_choice.logprobs is not None:
|
||||
if hasattr(first_choice.logprobs, "content") and first_choice.logprobs.content is not None:
|
||||
attributes["logprobs.content"] = json.dumps(
|
||||
[logprob.model_dump() for logprob in first_choice.logprobs.content]
|
||||
)
|
||||
if hasattr(first_choice.logprobs, "refusal") and first_choice.logprobs.refusal is not None:
|
||||
attributes["logprobs.refusal"] = json.dumps(
|
||||
[logprob.model_dump() for logprob in first_choice.logprobs.refusal]
|
||||
)
|
||||
|
||||
return attributes
|
||||
|
||||
agentops.instrumentation.providers.openai.wrappers.chat.handle_chat_attributes = _handle_chat_attributes_with_tokens
|
||||
agentops.instrumentation.providers.openai.stream_wrapper.handle_chat_attributes = (
|
||||
_handle_chat_attributes_with_tokens
|
||||
)
|
||||
logger.info("Patched newer version of agentops using handle_chat_attributes")
|
||||
return True
|
||||
|
||||
|
||||
def _unpatch_new_agentops():
|
||||
import agentops.instrumentation.providers.openai.stream_wrapper
|
||||
import agentops.instrumentation.providers.openai.wrappers.chat
|
||||
|
||||
global _original_handle_chat_attributes
|
||||
if _original_handle_chat_attributes is not None:
|
||||
agentops.instrumentation.providers.openai.wrappers.chat.handle_chat_attributes = (
|
||||
_original_handle_chat_attributes
|
||||
)
|
||||
agentops.instrumentation.providers.openai.stream_wrapper.handle_chat_attributes = (
|
||||
_original_handle_chat_attributes
|
||||
)
|
||||
_original_handle_chat_attributes = None
|
||||
logger.info("Unpatched newer version of agentops using handle_chat_attributes")
|
||||
|
||||
|
||||
def _patch_old_agentops():
|
||||
import opentelemetry.instrumentation.openai.shared.chat_wrappers # type: ignore
|
||||
from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw # type: ignore
|
||||
|
||||
global _original_handle_response
|
||||
_original_handle_response = _handle_response # type: ignore
|
||||
|
||||
@dont_throw # type: ignore
|
||||
def _handle_response_with_tokens(response, span, *args, **kwargs): # type: ignore
|
||||
_original_handle_response(response, span, *args, **kwargs) # type: ignore
|
||||
if hasattr(response, "prompt_token_ids"): # type: ignore
|
||||
span.set_attribute("prompt_token_ids", list(response.prompt_token_ids)) # type: ignore
|
||||
if hasattr(response, "response_token_ids"): # type: ignore
|
||||
span.set_attribute("response_token_ids", list(response.response_token_ids[0])) # type: ignore
|
||||
|
||||
# For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
|
||||
if hasattr(response, "http_response") and hasattr(response.http_response, "json"): # type: ignore
|
||||
json_data = response.http_response.json() # type: ignore
|
||||
if isinstance(json_data, dict):
|
||||
if "prompt_token_ids" in json_data:
|
||||
span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"])) # type: ignore
|
||||
if "response_token_ids" in json_data:
|
||||
span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0])) # type: ignore
|
||||
|
||||
opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens # type: ignore
|
||||
logger.info("Patched earlier version of agentops using _handle_response")
|
||||
return True
|
||||
|
||||
|
||||
def _unpatch_old_agentops():
|
||||
import opentelemetry.instrumentation.openai.shared.chat_wrappers # type: ignore
|
||||
|
||||
global _original_handle_response
|
||||
if _original_handle_response is not None:
|
||||
opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response # type: ignore
|
||||
_original_handle_response = None
|
||||
logger.info("Unpatched earlier version of agentops using _handle_response")
|
||||
|
||||
|
||||
def instrument_agentops():
|
||||
"""
|
||||
Instrument agentops to capture token IDs.
|
||||
Automatically detects and uses the appropriate patching method based on the installed agentops version.
|
||||
"""
|
||||
_patch_exporters()
|
||||
|
||||
# Try newest version first (tested for 0.4.16)
|
||||
try:
|
||||
return _patch_new_agentops()
|
||||
except ImportError as e:
|
||||
logger.debug(f"Couldn't patch newer version of agentops: {str(e)}")
|
||||
|
||||
# Note: 0.4.15 needs another patching method, but it's too shortlived to be worth handling separately.
|
||||
|
||||
# Try older version (tested for 0.4.13)
|
||||
try:
|
||||
return _patch_old_agentops()
|
||||
except ImportError as e:
|
||||
logger.warning(f"Couldn't patch older version of agentops: {str(e)}")
|
||||
logger.error("Failed to instrument agentops - neither patching method was successful")
|
||||
return False
|
||||
|
||||
|
||||
def uninstrument_agentops():
|
||||
"""Uninstrument agentops to stop capturing token IDs."""
|
||||
_unpatch_exporters()
|
||||
|
||||
try:
|
||||
_unpatch_new_agentops()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
_unpatch_old_agentops()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, AuthenticatedOTLPExporter):
|
||||
"""
|
||||
AuthenticatedOTLPExporter with switchable service control.
|
||||
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
"""
|
||||
|
||||
def should_bypass(self) -> bool:
|
||||
return not _agentops_service_enabled
|
||||
|
||||
|
||||
class BypassableOTLPMetricExporter(OTLPMetricExporter):
|
||||
"""
|
||||
OTLPMetricExporter with switchable service control.
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
"""
|
||||
|
||||
def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
|
||||
if _agentops_service_enabled:
|
||||
return super().export(*args, **kwargs) # type: ignore[reportUnknownMemberType]
|
||||
else:
|
||||
logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
|
||||
return MetricExportResult.SUCCESS
|
||||
|
||||
|
||||
class BypassableOTLPSpanExporter(LightningStoreOTLPExporter):
|
||||
"""
|
||||
OTLPSpanExporter with switchable service control.
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
|
||||
This is used instead of BypassableAuthenticatedOTLPExporter on legacy AgentOps versions.
|
||||
"""
|
||||
|
||||
def should_bypass(self) -> bool:
|
||||
return not _agentops_service_enabled
|
||||
|
||||
|
||||
class BypassableV3Client(V3Client):
|
||||
"""
|
||||
V3Client with toggleable authentication calls.
|
||||
Returns dummy auth response when `_agentops_service_enabled` is False.
|
||||
"""
|
||||
|
||||
# Temporary synchronous override of fetch_auth_token for mock purposes.
|
||||
def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse: # type: ignore[override]
|
||||
if _agentops_service_enabled:
|
||||
return super().fetch_auth_token(*args, **kwargs) # type: ignore[override]
|
||||
else:
|
||||
logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
|
||||
return AuthTokenResponse(token="dummy", project_id="dummy")
|
||||
|
||||
|
||||
class BypassableV4Client(V4Client):
|
||||
"""
|
||||
V4Client with toggleable post requests.
|
||||
Returns dummy response when `_agentops_service_enabled` is False.
|
||||
"""
|
||||
|
||||
def post(self, *args: Any, **kwargs: Any) -> requests.Response:
|
||||
if _agentops_service_enabled:
|
||||
return super().post(*args, **kwargs)
|
||||
else:
|
||||
logger.debug("SwitchableV4Client is switched off, skipping post request.")
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response._content = b"{}"
|
||||
return response
|
||||
@@ -0,0 +1,45 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from agentops import instrumentation
|
||||
from agentops.integration.callbacks.langchain import LangchainCallbackHandler
|
||||
|
||||
original_on_chain_start = LangchainCallbackHandler.on_chain_start
|
||||
langgraph_entry = None
|
||||
|
||||
__all__ = [
|
||||
"instrument_agentops_langchain",
|
||||
"uninstrument_agentops_langchain",
|
||||
]
|
||||
|
||||
|
||||
def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
|
||||
if "name" in kwargs:
|
||||
if serialized is None: # type: ignore
|
||||
serialized = {}
|
||||
serialized = serialized.copy()
|
||||
serialized["name"] = kwargs["name"]
|
||||
if "run_id" in kwargs:
|
||||
if serialized is None: # type: ignore
|
||||
serialized = {}
|
||||
serialized = serialized.copy()
|
||||
if "id" not in serialized:
|
||||
serialized["id"] = kwargs["run_id"]
|
||||
return original_on_chain_start(self, serialized, inputs, **kwargs)
|
||||
|
||||
|
||||
def instrument_agentops_langchain():
|
||||
"""Bypass AgentOp's native support for Langchain."""
|
||||
global langgraph_entry
|
||||
langgraph_entry = instrumentation.AGENTIC_LIBRARIES.pop("langgraph", None)
|
||||
LangchainCallbackHandler.on_chain_start = on_chain_start
|
||||
|
||||
|
||||
def uninstrument_agentops_langchain():
|
||||
"""Restore AgentOp's native support for Langchain."""
|
||||
global langgraph_entry
|
||||
if langgraph_entry is not None:
|
||||
instrumentation.AGENTIC_LIBRARIES["langgraph"] = langgraph_entry
|
||||
langgraph_entry = None
|
||||
LangchainCallbackHandler.on_chain_start = original_on_chain_start
|
||||
@@ -0,0 +1,39 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""LiteLLM instrumentations.
|
||||
|
||||
It's unclear whether or not this file is useful.
|
||||
It seems that LiteLLM owns its own telemetry from their own entrance
|
||||
|
||||
[Related documentation](https://docs.litellm.ai/docs/observability/agentops_integration).
|
||||
"""
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
|
||||
__all__ = [
|
||||
"instrument_litellm",
|
||||
"uninstrument_litellm",
|
||||
]
|
||||
|
||||
original_set_attributes = OpenTelemetry.set_attributes # type: ignore
|
||||
|
||||
|
||||
def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
|
||||
original_set_attributes(self, span, kwargs, response_obj)
|
||||
# Add custom attributes
|
||||
if response_obj is not None and response_obj.get("prompt_token_ids"):
|
||||
span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
|
||||
if response_obj is not None and response_obj.get("response_token_ids"):
|
||||
span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))
|
||||
|
||||
|
||||
def instrument_litellm():
|
||||
"""Instrument litellm to capture token IDs."""
|
||||
OpenTelemetry.set_attributes = patched_set_attributes
|
||||
|
||||
|
||||
def uninstrument_litellm():
|
||||
"""Uninstrument litellm to stop capturing token IDs."""
|
||||
OpenTelemetry.set_attributes = original_set_attributes
|
||||
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import Any, List
|
||||
|
||||
import vllm.entrypoints.openai.protocol
|
||||
from vllm.entrypoints.openai.protocol import ChatCompletionResponse
|
||||
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
|
||||
|
||||
__all__ = [
|
||||
"instrument_vllm",
|
||||
"uninstrument_vllm",
|
||||
]
|
||||
|
||||
|
||||
class ChatCompletionResponsePatched(ChatCompletionResponse):
|
||||
prompt_token_ids: List[int] | None = None
|
||||
response_token_ids: List[int] | None = None
|
||||
|
||||
|
||||
original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full_generator
|
||||
|
||||
|
||||
async def chat_completion_full_generator(
|
||||
self: Any,
|
||||
request: Any,
|
||||
result_generator: Any,
|
||||
request_id: str,
|
||||
model_name: str,
|
||||
conversation: Any,
|
||||
tokenizer: Any,
|
||||
request_metadata: Any,
|
||||
) -> Any:
|
||||
prompt_token_ids: List[int] | None = None
|
||||
response_token_ids: List[List[int]] | None = None
|
||||
|
||||
async def _generate_inceptor():
|
||||
nonlocal prompt_token_ids, response_token_ids
|
||||
async for res in result_generator:
|
||||
yield res
|
||||
prompt_token_ids = res.prompt_token_ids
|
||||
response_token_ids = [output.token_ids for output in res.outputs]
|
||||
|
||||
response = await original_chat_completion_full_generator(
|
||||
self,
|
||||
request,
|
||||
_generate_inceptor(),
|
||||
request_id,
|
||||
model_name,
|
||||
conversation,
|
||||
tokenizer,
|
||||
request_metadata,
|
||||
)
|
||||
response = response.model_copy(
|
||||
update={
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"response_token_ids": response_token_ids,
|
||||
}
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def instrument_vllm():
|
||||
"""Instrument vLLM to capture token IDs generated by engine.
|
||||
|
||||
This instrumentation has been merged to upstream vLLM since v0.10.2.
|
||||
"""
|
||||
if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
|
||||
warnings.warn("vllm is already instrumented. Skip the instrumentation.")
|
||||
return
|
||||
|
||||
vllm.entrypoints.openai.protocol.ChatCompletionResponse = ChatCompletionResponsePatched
|
||||
OpenAIServingChat.chat_completion_full_generator = chat_completion_full_generator
|
||||
|
||||
|
||||
def uninstrument_vllm():
|
||||
"""Uninstrument vLLM to stop capturing token IDs generated by engine."""
|
||||
OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator
|
||||
@@ -0,0 +1,541 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import warnings
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Callable, Dict, Iterator, List
|
||||
|
||||
import weave.trace.weave_init
|
||||
from pydantic import validate_call
|
||||
from weave.trace_server import trace_server_interface as tsi
|
||||
from weave.trace_server.ids import generate_id
|
||||
from weave.trace_server_bindings.client_interface import TraceServerClientInterface
|
||||
from weave.trace_server_bindings.models import ServerInfoRes
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"instrument_weave",
|
||||
"uninstrument_weave",
|
||||
"InMemoryWeaveTraceServer",
|
||||
]
|
||||
|
||||
|
||||
class InMemoryWeaveTraceServer(TraceServerClientInterface):
|
||||
"""A minimal in-memory implementation of the TraceServerInterface.
|
||||
|
||||
It stores calls and objects in local dictionaries and returns valid Pydantic
|
||||
responses to satisfy the Weave client and FullTraceServerInterface protocol.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# Minimal storage to allow basic querying in tests
|
||||
self.calls: Dict[str, tsi.CallSchema] = {}
|
||||
self.partial_calls: Dict[str, Dict[str, Any]] = {}
|
||||
self.objs: Dict[str, Any] = {}
|
||||
self.files: Dict[str, bytes] = {}
|
||||
self.feedback: List[tsi.FeedbackCreateReq] = []
|
||||
|
||||
self._call_threading_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, *args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
|
||||
return cls()
|
||||
|
||||
def server_info(self) -> ServerInfoRes:
|
||||
return ServerInfoRes(min_required_weave_python_version="0.52.22")
|
||||
|
||||
def ensure_project_exists(self, entity: str, project: str) -> tsi.EnsureProjectExistsRes:
|
||||
return tsi.EnsureProjectExistsRes(project_name=project)
|
||||
|
||||
# --- Call API ---
|
||||
|
||||
@validate_call
|
||||
def call_start(self, req: tsi.CallStartReq) -> tsi.CallStartRes:
|
||||
# NOTE: It's not necessary that call_end must be called after call_start.
|
||||
request_content = req.start.model_dump(exclude_none=True)
|
||||
|
||||
# If id needs to be generated here, it's very likely we won't be able to find the call later.
|
||||
# This is just to make the type checker happy.
|
||||
call_id = request_content.get("id") or generate_id()
|
||||
trace_id = request_content.get("trace_id") or generate_id()
|
||||
request_content["id"] = call_id
|
||||
request_content["trace_id"] = trace_id
|
||||
|
||||
with self._call_threading_lock:
|
||||
if call_id in self.partial_calls:
|
||||
# call_end has already been called for this call.
|
||||
kwargs = {**request_content, **self.partial_calls[call_id]}
|
||||
self.calls[call_id] = tsi.CallSchema(**kwargs)
|
||||
del self.partial_calls[call_id]
|
||||
else:
|
||||
self.partial_calls[call_id] = request_content
|
||||
|
||||
return tsi.CallStartRes(id=call_id, trace_id=trace_id)
|
||||
|
||||
@validate_call
|
||||
def call_end(self, req: tsi.CallEndReq) -> tsi.CallEndRes:
|
||||
request_content = req.end.model_dump(exclude_none=True)
|
||||
call_id = req.end.id
|
||||
|
||||
with self._call_threading_lock:
|
||||
if call_id in self.partial_calls:
|
||||
# End request always override the start request content.
|
||||
kwargs = {**self.partial_calls[call_id], **request_content}
|
||||
self.calls[call_id] = tsi.CallSchema(**kwargs)
|
||||
del self.partial_calls[call_id]
|
||||
else:
|
||||
self.partial_calls[call_id] = request_content
|
||||
return tsi.CallEndRes()
|
||||
|
||||
@validate_call
|
||||
def call_start_batch(self, req: tsi.CallCreateBatchReq) -> tsi.CallCreateBatchRes:
|
||||
for item in req.batch:
|
||||
if isinstance(item, tsi.CallStartReq):
|
||||
self.call_start(item)
|
||||
elif isinstance(item, tsi.CallEndReq):
|
||||
self.call_end(item)
|
||||
return tsi.CallCreateBatchRes(res=[])
|
||||
|
||||
@validate_call
|
||||
def call_read(self, req: tsi.CallReadReq) -> tsi.CallReadRes:
|
||||
call_data = self.calls.get(req.id)
|
||||
return tsi.CallReadRes(call=call_data)
|
||||
|
||||
@validate_call
|
||||
def calls_query(self, req: tsi.CallsQueryReq) -> tsi.CallsQueryRes:
|
||||
return tsi.CallsQueryRes(calls=list(self.calls_query_stream(req)))
|
||||
|
||||
@validate_call
|
||||
def calls_query_stream(self, req: tsi.CallsQueryReq) -> Iterator[tsi.CallSchema]:
|
||||
yield from self.calls.values()
|
||||
|
||||
@validate_call
|
||||
def calls_delete(self, req: tsi.CallsDeleteReq) -> tsi.CallsDeleteRes:
|
||||
num_deleted = 0
|
||||
for call_id in req.call_ids:
|
||||
if call_id in self.calls:
|
||||
del self.calls[call_id]
|
||||
num_deleted += 1
|
||||
return tsi.CallsDeleteRes(num_deleted=num_deleted)
|
||||
|
||||
@validate_call
|
||||
def call_update(self, req: tsi.CallUpdateReq) -> tsi.CallUpdateRes:
|
||||
return tsi.CallUpdateRes()
|
||||
|
||||
@validate_call
|
||||
def calls_query_stats(self, req: tsi.CallsQueryStatsReq) -> tsi.CallsQueryStatsRes:
|
||||
return tsi.CallsQueryStatsRes(count=len(self.calls))
|
||||
|
||||
# --- Cost API ---
|
||||
|
||||
@validate_call
|
||||
def cost_create(self, req: tsi.CostCreateReq) -> tsi.CostCreateRes:
|
||||
return tsi.CostCreateRes(ids=[(generate_id(), generate_id()) for _ in req.costs])
|
||||
|
||||
@validate_call
|
||||
def cost_query(self, req: tsi.CostQueryReq) -> tsi.CostQueryRes:
|
||||
return tsi.CostQueryRes(results=[])
|
||||
|
||||
@validate_call
|
||||
def cost_purge(self, req: tsi.CostPurgeReq) -> tsi.CostPurgeRes:
|
||||
return tsi.CostPurgeRes()
|
||||
|
||||
# --- Object API (Legacy V1) ---
|
||||
|
||||
@validate_call
|
||||
def obj_create(self, req: tsi.ObjCreateReq) -> tsi.ObjCreateRes:
|
||||
digest = generate_id()
|
||||
self.objs[digest] = req.obj
|
||||
return tsi.ObjCreateRes(digest=digest)
|
||||
|
||||
@validate_call
|
||||
def obj_read(self, req: tsi.ObjReadReq) -> tsi.ObjReadRes:
|
||||
return tsi.ObjReadRes(obj=self.objs.get(req.digest, {}))
|
||||
|
||||
@validate_call
|
||||
def objs_query(self, req: tsi.ObjQueryReq) -> tsi.ObjQueryRes:
|
||||
return tsi.ObjQueryRes(objs=[])
|
||||
|
||||
@validate_call
|
||||
def obj_delete(self, req: tsi.ObjDeleteReq) -> tsi.ObjDeleteRes:
|
||||
return tsi.ObjDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Table API ---
|
||||
|
||||
@validate_call
|
||||
def table_create(self, req: tsi.TableCreateReq) -> tsi.TableCreateRes:
|
||||
return tsi.TableCreateRes(digest=generate_id(), row_digests=[])
|
||||
|
||||
@validate_call
|
||||
def table_create_from_digests(self, req: tsi.TableCreateFromDigestsReq) -> tsi.TableCreateFromDigestsRes:
|
||||
return tsi.TableCreateFromDigestsRes(digest=generate_id())
|
||||
|
||||
@validate_call
|
||||
def table_update(self, req: tsi.TableUpdateReq) -> tsi.TableUpdateRes:
|
||||
return tsi.TableUpdateRes(digest=generate_id(), updated_row_digests=[])
|
||||
|
||||
@validate_call
|
||||
def table_query(self, req: tsi.TableQueryReq) -> tsi.TableQueryRes:
|
||||
return tsi.TableQueryRes(rows=[])
|
||||
|
||||
@validate_call
|
||||
def table_query_stream(self, req: tsi.TableQueryReq) -> Iterator[tsi.TableRowSchema]:
|
||||
yield from []
|
||||
|
||||
@validate_call
|
||||
def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
|
||||
return tsi.TableQueryStatsRes(count=0)
|
||||
|
||||
@validate_call
|
||||
def table_query_stats_batch(self, req: tsi.TableQueryStatsBatchReq) -> tsi.TableQueryStatsBatchRes:
|
||||
return tsi.TableQueryStatsBatchRes(tables=[])
|
||||
|
||||
# --- Ref API ---
|
||||
|
||||
@validate_call
|
||||
def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
|
||||
return tsi.RefsReadBatchRes(vals=[])
|
||||
|
||||
# --- File API ---
|
||||
|
||||
def file_create(self, req: tsi.FileCreateReq) -> tsi.FileCreateRes:
|
||||
self.files[req.name] = req.content
|
||||
return tsi.FileCreateRes(digest=generate_id())
|
||||
|
||||
def file_content_read(self, req: tsi.FileContentReadReq) -> tsi.FileContentReadRes:
|
||||
return tsi.FileContentReadRes(content=self.files.get(req.digest, b"dummy_content"))
|
||||
|
||||
def files_stats(self, req: tsi.FilesStatsReq) -> tsi.FilesStatsRes:
|
||||
total_size = sum(len(c) for c in self.files.values())
|
||||
return tsi.FilesStatsRes(total_size_bytes=total_size)
|
||||
|
||||
# --- Feedback API ---
|
||||
|
||||
@validate_call
|
||||
def feedback_create(self, req: tsi.FeedbackCreateReq) -> tsi.FeedbackCreateRes:
|
||||
req.id = req.id or generate_id()
|
||||
self.feedback.append(req)
|
||||
return tsi.FeedbackCreateRes(
|
||||
id=req.id,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
wb_user_id="dummy_user",
|
||||
payload=req.payload,
|
||||
)
|
||||
|
||||
def feedback_create_batch(self, req: tsi.FeedbackCreateBatchReq) -> tsi.FeedbackCreateBatchRes:
|
||||
results: List[tsi.FeedbackCreateRes] = []
|
||||
for item in req.batch:
|
||||
res = self.feedback_create(item)
|
||||
results.append(res)
|
||||
return tsi.FeedbackCreateBatchRes(res=results)
|
||||
|
||||
@validate_call
|
||||
def feedback_query(self, req: tsi.FeedbackQueryReq) -> tsi.FeedbackQueryRes:
|
||||
return tsi.FeedbackQueryRes(result=[])
|
||||
|
||||
@validate_call
|
||||
def feedback_purge(self, req: tsi.FeedbackPurgeReq) -> tsi.FeedbackPurgeRes:
|
||||
self.feedback.clear()
|
||||
return tsi.FeedbackPurgeRes()
|
||||
|
||||
@validate_call
|
||||
def feedback_replace(self, req: tsi.FeedbackReplaceReq) -> tsi.FeedbackReplaceRes:
|
||||
return tsi.FeedbackReplaceRes(
|
||||
id=req.id or generate_id(),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
wb_user_id="dummy",
|
||||
payload={},
|
||||
)
|
||||
|
||||
# --- Action API ---
|
||||
|
||||
@validate_call
|
||||
def actions_execute_batch(self, req: tsi.ActionsExecuteBatchReq) -> tsi.ActionsExecuteBatchRes:
|
||||
return tsi.ActionsExecuteBatchRes()
|
||||
|
||||
# --- Execute LLM API ---
|
||||
|
||||
@validate_call
|
||||
def completions_create(self, req: tsi.CompletionsCreateReq) -> tsi.CompletionsCreateRes:
|
||||
return tsi.CompletionsCreateRes(response={"choices": [{"text": "dummy completion"}]})
|
||||
|
||||
@validate_call
|
||||
def completions_create_stream(self, req: tsi.CompletionsCreateReq) -> Iterator[dict[str, Any]]:
|
||||
yield {"choices": [{"text": "dummy "}]}
|
||||
yield {"choices": [{"text": "stream"}]}
|
||||
|
||||
# --- Execute Image Generation API ---
|
||||
|
||||
@validate_call
|
||||
def image_create(self, req: tsi.ImageGenerationCreateReq) -> tsi.ImageGenerationCreateRes:
|
||||
return tsi.ImageGenerationCreateRes(response={})
|
||||
|
||||
# --- Project Statistics API ---
|
||||
|
||||
@validate_call
|
||||
def project_stats(self, req: tsi.ProjectStatsReq) -> tsi.ProjectStatsRes:
|
||||
return tsi.ProjectStatsRes(
|
||||
trace_storage_size_bytes=0,
|
||||
objects_storage_size_bytes=0,
|
||||
tables_storage_size_bytes=0,
|
||||
files_storage_size_bytes=0,
|
||||
)
|
||||
|
||||
# --- Thread API ---
|
||||
|
||||
@validate_call
|
||||
def threads_query_stream(self, req: tsi.ThreadsQueryReq) -> Iterator[tsi.ThreadSchema]:
|
||||
yield from []
|
||||
|
||||
# --- Evaluation API (V1) ---
|
||||
|
||||
@validate_call
|
||||
def evaluate_model(self, req: tsi.EvaluateModelReq) -> tsi.EvaluateModelRes:
|
||||
return tsi.EvaluateModelRes(call_id=generate_id())
|
||||
|
||||
@validate_call
|
||||
def evaluation_status(self, req: tsi.EvaluationStatusReq) -> tsi.EvaluationStatusRes:
|
||||
return tsi.EvaluationStatusRes(status=tsi.EvaluationStatusNotFound())
|
||||
|
||||
# --- OTEL API ---
|
||||
|
||||
def otel_export(self, req: tsi.OtelExportReq) -> tsi.OtelExportRes:
|
||||
return tsi.OtelExportRes()
|
||||
|
||||
# ==========================================
|
||||
# Object Interface (V2 APIs)
|
||||
# ==========================================
|
||||
|
||||
# --- Ops ---
|
||||
def op_create(self, req: tsi.OpCreateReq) -> tsi.OpCreateRes:
|
||||
return tsi.OpCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
|
||||
|
||||
def op_read(self, req: tsi.OpReadReq) -> tsi.OpReadRes:
|
||||
return tsi.OpReadRes(op=None) # type: ignore
|
||||
|
||||
def op_list(self, req: tsi.OpListReq) -> Iterator[tsi.OpReadRes]:
|
||||
yield from []
|
||||
|
||||
def op_delete(self, req: tsi.OpDeleteReq) -> tsi.OpDeleteRes:
|
||||
return tsi.OpDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Datasets ---
|
||||
def dataset_create(self, req: tsi.DatasetCreateReq) -> tsi.DatasetCreateRes:
|
||||
return tsi.DatasetCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0)
|
||||
|
||||
def dataset_read(self, req: tsi.DatasetReadReq) -> tsi.DatasetReadRes:
|
||||
return tsi.DatasetReadRes(dataset=None) # type: ignore
|
||||
|
||||
def dataset_list(self, req: tsi.DatasetListReq) -> Iterator[tsi.DatasetReadRes]:
|
||||
yield from []
|
||||
|
||||
def dataset_delete(self, req: tsi.DatasetDeleteReq) -> tsi.DatasetDeleteRes:
|
||||
return tsi.DatasetDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Scorers ---
|
||||
def scorer_create(self, req: tsi.ScorerCreateReq) -> tsi.ScorerCreateRes:
|
||||
return tsi.ScorerCreateRes(digest=generate_id(), object_id=generate_id(), version_index=0, scorer=generate_id())
|
||||
|
||||
def scorer_read(self, req: tsi.ScorerReadReq) -> tsi.ScorerReadRes:
|
||||
return tsi.ScorerReadRes(scorer=None) # type: ignore
|
||||
|
||||
def scorer_list(self, req: tsi.ScorerListReq) -> Iterator[tsi.ScorerReadRes]:
|
||||
yield from []
|
||||
|
||||
def scorer_delete(self, req: tsi.ScorerDeleteReq) -> tsi.ScorerDeleteRes:
|
||||
return tsi.ScorerDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Evaluations (V2) ---
|
||||
def evaluation_create(self, req: tsi.EvaluationCreateReq) -> tsi.EvaluationCreateRes:
|
||||
return tsi.EvaluationCreateRes(
|
||||
digest=generate_id(), object_id=generate_id(), version_index=0, evaluation_ref=generate_id()
|
||||
)
|
||||
|
||||
def evaluation_read(self, req: tsi.EvaluationReadReq) -> tsi.EvaluationReadRes:
|
||||
return tsi.EvaluationReadRes(evaluation=None) # type: ignore
|
||||
|
||||
def evaluation_list(self, req: tsi.EvaluationListReq) -> Iterator[tsi.EvaluationReadRes]:
|
||||
yield from []
|
||||
|
||||
def evaluation_delete(self, req: tsi.EvaluationDeleteReq) -> tsi.EvaluationDeleteRes:
|
||||
return tsi.EvaluationDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Models ---
|
||||
def model_create(self, req: tsi.ModelCreateReq) -> tsi.ModelCreateRes:
|
||||
return tsi.ModelCreateRes(
|
||||
digest=generate_id(), object_id=generate_id(), version_index=0, model_ref=generate_id()
|
||||
)
|
||||
|
||||
def model_read(self, req: tsi.ModelReadReq) -> tsi.ModelReadRes:
|
||||
return tsi.ModelReadRes(model=None) # type: ignore
|
||||
|
||||
def model_list(self, req: tsi.ModelListReq) -> Iterator[tsi.ModelReadRes]:
|
||||
yield from []
|
||||
|
||||
def model_delete(self, req: tsi.ModelDeleteReq) -> tsi.ModelDeleteRes:
|
||||
return tsi.ModelDeleteRes(num_deleted=0)
|
||||
|
||||
# --- Evaluation Runs ---
|
||||
def evaluation_run_create(self, req: tsi.EvaluationRunCreateReq) -> tsi.EvaluationRunCreateRes:
|
||||
return tsi.EvaluationRunCreateRes(evaluation_run_id=generate_id())
|
||||
|
||||
def evaluation_run_read(self, req: tsi.EvaluationRunReadReq) -> tsi.EvaluationRunReadRes:
|
||||
return tsi.EvaluationRunReadRes(evaluation_run=None) # type: ignore
|
||||
|
||||
def evaluation_run_list(self, req: tsi.EvaluationRunListReq) -> Iterator[tsi.EvaluationRunReadRes]:
|
||||
yield from []
|
||||
|
||||
def evaluation_run_delete(self, req: tsi.EvaluationRunDeleteReq) -> tsi.EvaluationRunDeleteRes:
|
||||
return tsi.EvaluationRunDeleteRes(num_deleted=0)
|
||||
|
||||
def evaluation_run_finish(self, req: tsi.EvaluationRunFinishReq) -> tsi.EvaluationRunFinishRes:
|
||||
return tsi.EvaluationRunFinishRes(success=True)
|
||||
|
||||
# --- Predictions ---
|
||||
def prediction_create(self, req: tsi.PredictionCreateReq) -> tsi.PredictionCreateRes:
|
||||
return tsi.PredictionCreateRes(prediction_id=generate_id())
|
||||
|
||||
def prediction_read(self, req: tsi.PredictionReadReq) -> tsi.PredictionReadRes:
|
||||
return tsi.PredictionReadRes(prediction=None) # type: ignore
|
||||
|
||||
def prediction_list(self, req: tsi.PredictionListReq) -> Iterator[tsi.PredictionReadRes]:
|
||||
yield from []
|
||||
|
||||
def prediction_delete(self, req: tsi.PredictionDeleteReq) -> tsi.PredictionDeleteRes:
|
||||
return tsi.PredictionDeleteRes(num_deleted=0)
|
||||
|
||||
def prediction_finish(self, req: tsi.PredictionFinishReq) -> tsi.PredictionFinishRes:
|
||||
return tsi.PredictionFinishRes(success=True)
|
||||
|
||||
# --- Scores ---
|
||||
def score_create(self, req: tsi.ScoreCreateReq) -> tsi.ScoreCreateRes:
|
||||
return tsi.ScoreCreateRes(score_id=generate_id())
|
||||
|
||||
def score_read(self, req: tsi.ScoreReadReq) -> tsi.ScoreReadRes:
|
||||
return tsi.ScoreReadRes(score=None) # type: ignore
|
||||
|
||||
def score_list(self, req: tsi.ScoreListReq) -> Iterator[tsi.ScoreReadRes]:
|
||||
yield from []
|
||||
|
||||
def score_delete(self, req: tsi.ScoreDeleteReq) -> tsi.ScoreDeleteRes:
|
||||
return tsi.ScoreDeleteRes(num_deleted=0)
|
||||
|
||||
# Experimental unstable APIs
|
||||
# We don't support these APIs yet.
|
||||
def annotation_queue_create(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotation_queues_query_stream(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotation_queue_read(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotation_queue_add_calls(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotation_queues_stats(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotation_queue_items_query(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def annotator_queue_items_progress_update(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def calls_complete(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def call_start_v2(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def call_end_v2(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def call_stats(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def trace_usage(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def calls_usage(self, *args: Any, **kwargs: Any) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
# Module-level storage for originals
|
||||
_original_init_weave_get_server: Callable[..., Any] | None = None
|
||||
_original_get_entity_project_from_project_name: Callable[..., Any] | None = None
|
||||
_original_get_username: Callable[..., Any] | None = None
|
||||
|
||||
|
||||
def init_weave_get_server_factory(server: InMemoryWeaveTraceServer) -> Callable[..., Any]:
|
||||
# Bypass the usage of Weave remote server
|
||||
def init_weave_get_server(*args: Any, **kwargs: Any) -> InMemoryWeaveTraceServer:
|
||||
return server
|
||||
|
||||
return init_weave_get_server
|
||||
|
||||
|
||||
def get_entity_project_from_project_name_factory(entity_name: str) -> tuple[str, str]:
|
||||
# Bypass the usage of API
|
||||
try:
|
||||
assert _original_get_entity_project_from_project_name is not None
|
||||
if _original_get_entity_project_from_project_name is not get_entity_project_from_project_name_factory:
|
||||
return _original_get_entity_project_from_project_name(entity_name)
|
||||
else:
|
||||
warnings.warn("W&B integration might have been repeatedly/recursively instrumented.")
|
||||
return "agl", "weave"
|
||||
except weave.trace.weave_init.WeaveWandbAuthenticationException:
|
||||
# In case API is not available.
|
||||
return "agl", "weave"
|
||||
|
||||
|
||||
def get_username() -> str:
|
||||
# Bypass the usage of API
|
||||
try:
|
||||
assert _original_get_username is not None
|
||||
return _original_get_username()
|
||||
except RuntimeError:
|
||||
return "agl"
|
||||
except Exception as exc:
|
||||
warnings.warn(f"Unexpected error in get_username. Using default username. Error: {exc}")
|
||||
return "agl"
|
||||
|
||||
|
||||
def instrument_weave(server: InMemoryWeaveTraceServer):
|
||||
"""Patch the Weave/W&B integration to bypass actual network calls for testing."""
|
||||
|
||||
global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
|
||||
_original_init_weave_get_server = weave.trace.weave_init.init_weave_get_server
|
||||
_original_get_entity_project_from_project_name = weave.trace.weave_init.get_entity_project_from_project_name
|
||||
_original_get_username = weave.trace.weave_init.get_username
|
||||
weave.trace.weave_init.init_weave_get_server = init_weave_get_server_factory(server)
|
||||
weave.trace.weave_init.get_entity_project_from_project_name = get_entity_project_from_project_name_factory
|
||||
weave.trace.weave_init.get_username = get_username
|
||||
|
||||
|
||||
def uninstrument_weave():
|
||||
"""Restore the original Weave/W&B integration methods and HTTP requests."""
|
||||
global _original_init_weave_get_server, _original_get_entity_project_from_project_name, _original_get_username
|
||||
|
||||
if _original_init_weave_get_server is not None:
|
||||
weave.trace.weave_init.init_weave_get_server = _original_init_weave_get_server
|
||||
_original_init_weave_get_server = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
|
||||
if _original_get_entity_project_from_project_name is not None:
|
||||
weave.trace.weave_init.get_entity_project_from_project_name = _original_get_entity_project_from_project_name
|
||||
_original_get_entity_project_from_project_name = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
|
||||
if _original_get_username is not None:
|
||||
weave.trace.weave_init.get_username = _original_get_username
|
||||
_original_get_username = None
|
||||
else:
|
||||
raise RuntimeError("Weave/W&B integration was not instrumented.")
|
||||
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .decorator import *
|
||||
from .litagent import *
|
||||
|
||||
__all__ = [
|
||||
"LitAgent",
|
||||
"llm_rollout",
|
||||
"prompt_rollout",
|
||||
"rollout",
|
||||
]
|
||||
@@ -0,0 +1,536 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Convenience decorators for building lightweight `LitAgent` implementations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
from typing import Any, Awaitable, Callable, Dict, Protocol, TypeGuard, TypeVar, Union, overload
|
||||
|
||||
from agentlightning.types import (
|
||||
LLM,
|
||||
AttemptedRollout,
|
||||
NamedResources,
|
||||
PromptTemplate,
|
||||
ProxyLLM,
|
||||
Rollout,
|
||||
RolloutRawResult,
|
||||
)
|
||||
|
||||
from .litagent import LitAgent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
__all__ = [
|
||||
"llm_rollout",
|
||||
"prompt_rollout",
|
||||
"rollout",
|
||||
]
|
||||
|
||||
|
||||
T_contra = TypeVar("T_contra", contravariant=True)
|
||||
|
||||
|
||||
class LlmRolloutFuncSync2(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, llm: LLM) -> RolloutRawResult: ...
|
||||
|
||||
|
||||
class LlmRolloutFuncSync3(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> RolloutRawResult: ...
|
||||
|
||||
|
||||
class LlmRolloutFuncAsync2(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, llm: LLM) -> Awaitable[RolloutRawResult]: ...
|
||||
|
||||
|
||||
class LlmRolloutFuncAsync3(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> Awaitable[RolloutRawResult]: ...
|
||||
|
||||
|
||||
LlmRolloutFunc = Union[
|
||||
LlmRolloutFuncSync2[T_contra],
|
||||
LlmRolloutFuncSync3[T_contra],
|
||||
LlmRolloutFuncAsync2[T_contra],
|
||||
LlmRolloutFuncAsync3[T_contra],
|
||||
]
|
||||
|
||||
|
||||
class PromptRolloutFuncSync2(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> RolloutRawResult: ...
|
||||
|
||||
|
||||
class PromptRolloutFuncAsync2(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> Awaitable[RolloutRawResult]: ...
|
||||
|
||||
|
||||
class PromptRolloutFuncSync3(Protocol[T_contra]):
|
||||
def __call__(self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout) -> RolloutRawResult: ...
|
||||
|
||||
|
||||
class PromptRolloutFuncAsync3(Protocol[T_contra]):
|
||||
def __call__(
|
||||
self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout
|
||||
) -> Awaitable[RolloutRawResult]: ...
|
||||
|
||||
|
||||
PromptRolloutFunc = Union[
|
||||
PromptRolloutFuncSync2[T_contra],
|
||||
PromptRolloutFuncSync3[T_contra],
|
||||
PromptRolloutFuncAsync2[T_contra],
|
||||
PromptRolloutFuncAsync3[T_contra],
|
||||
]
|
||||
|
||||
|
||||
class FunctionalLitAgentFunc(Protocol[T_contra]):
|
||||
def __call__(
|
||||
self, task: T_contra, *args: Any, **kwargs: Any
|
||||
) -> Union[RolloutRawResult, Awaitable[RolloutRawResult]]: ...
|
||||
|
||||
|
||||
class FunctionalLitAgent(LitAgent[T]):
|
||||
"""Adapter that turns plain rollout functions into [`LitAgent`][agentlightning.LitAgent] instances.
|
||||
|
||||
The helper inspects the wrapped function to determine which resources to
|
||||
inject, allowing both synchronous and asynchronous callables to participate
|
||||
in the training loop without writing a dedicated subclass.
|
||||
"""
|
||||
|
||||
def __init__(self, rollout_func: FunctionalLitAgentFunc[T], *, strip_proxy: bool = True) -> None:
|
||||
"""Initialize the wrapper around a rollout function.
|
||||
|
||||
Args:
|
||||
rollout_func: Callable that implements the rollout. It may be synchronous
|
||||
or asynchronous and can optionally receive a
|
||||
[`Rollout`][agentlightning.Rollout] alongside resources such as
|
||||
`llm` or `prompt_template`.
|
||||
strip_proxy: When ``True``, convert
|
||||
[`ProxyLLM`][agentlightning.ProxyLLM] inputs into
|
||||
[`LLM`][agentlightning.LLM] instances before calling the
|
||||
rollout function. Defaults to `True`.
|
||||
"""
|
||||
super().__init__()
|
||||
self._rollout_func = rollout_func
|
||||
self._strip_proxy = strip_proxy
|
||||
self._is_async = inspect.iscoroutinefunction(rollout_func)
|
||||
self._sig = inspect.signature(rollout_func)
|
||||
|
||||
# Copy function metadata to preserve type hints and other attributes
|
||||
functools.update_wrapper(self, rollout_func) # type: ignore
|
||||
|
||||
def _accepts_rollout(self) -> bool:
|
||||
return "rollout" in self._sig.parameters
|
||||
|
||||
def _accepts_llm(self) -> bool:
|
||||
return "llm" in self._sig.parameters
|
||||
|
||||
def _accepts_prompt_template(self) -> bool:
|
||||
return "prompt_template" in self._sig.parameters
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any:
|
||||
"""Make the agent instance callable, preserving the original function behavior."""
|
||||
return self._rollout_func(*args, **kwargs) # type: ignore
|
||||
|
||||
def is_async(self) -> bool:
|
||||
return self._is_async
|
||||
|
||||
def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Execute a synchronous rollout using the wrapped function.
|
||||
|
||||
Args:
|
||||
task: Task input data.
|
||||
resources: Mapping of named resources available to the agent.
|
||||
rollout: Rollout metadata provided by the runtime.
|
||||
|
||||
Returns:
|
||||
Result produced by the wrapped rollout function.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the wrapped function is asynchronous.
|
||||
"""
|
||||
if self._is_async:
|
||||
raise RuntimeError(f"{self._rollout_func} is asynchronous. Use rollout_async instead.")
|
||||
|
||||
kwargs = self._get_kwargs(resources, rollout)
|
||||
return self._rollout_func(task, **kwargs) # type: ignore
|
||||
|
||||
async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Execute an asynchronous rollout using the wrapped function.
|
||||
|
||||
Args:
|
||||
task: Task input data.
|
||||
resources: Mapping of named resources available to the agent.
|
||||
rollout: Rollout metadata provided by the runtime.
|
||||
|
||||
Returns:
|
||||
Result produced by the wrapped rollout coroutine.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the wrapped function is synchronous.
|
||||
"""
|
||||
if not self._is_async:
|
||||
raise RuntimeError(f"{self._rollout_func} is synchronous. Use rollout instead.")
|
||||
|
||||
kwargs = self._get_kwargs(resources, rollout)
|
||||
return await self._rollout_func(task, **kwargs) # type: ignore
|
||||
|
||||
def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
|
||||
"""Prepare keyword arguments expected by the wrapped rollout function.
|
||||
|
||||
|
||||
It dynamically builds the `kwargs` dictionary by inspecting the function signature and
|
||||
including only the parameters the function accepts. This allows flexible function
|
||||
signatures that can request any combination of: rollout, llm, and/or prompt_template.
|
||||
|
||||
Args:
|
||||
resources: Mapping of named resources available for the rollout.
|
||||
rollout: Rollout metadata provided by the runtime.
|
||||
|
||||
Returns:
|
||||
Dictionary of keyword arguments to forward to the rollout function.
|
||||
"""
|
||||
|
||||
kwargs: Dict[str, Any] = {}
|
||||
if self._accepts_rollout():
|
||||
kwargs["rollout"] = rollout
|
||||
if self._accepts_llm():
|
||||
kwargs["llm"] = self._get_llm_resource(resources, rollout)
|
||||
if self._accepts_prompt_template():
|
||||
kwargs["prompt_template"] = self._get_prompt_template_resource(resources, rollout)
|
||||
|
||||
return kwargs
|
||||
|
||||
def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
|
||||
"""Retrieve the first LLM resource from the available resources.
|
||||
|
||||
Strip the ProxyLLM resource into a LLM resource if needed.
|
||||
|
||||
Args:
|
||||
resources: Mapping of named resources.
|
||||
rollout: Rollout metadata used when stripping proxy endpoints.
|
||||
|
||||
Returns:
|
||||
First [`LLM`][agentlightning.LLM] resource encountered.
|
||||
|
||||
Raises:
|
||||
ValueError: If no LLM resource is present.
|
||||
"""
|
||||
resource_found: LLM | None = None
|
||||
for name, resource in resources.items():
|
||||
if isinstance(resource, LLM):
|
||||
if resource_found is not None:
|
||||
logger.warning(f"Multiple LLM resources found in resources. Using the first one: '{name}'.")
|
||||
break
|
||||
resource_found = resource
|
||||
|
||||
if resource_found is None:
|
||||
raise ValueError("No LLM resource found in the provided resources.")
|
||||
|
||||
if self._strip_proxy:
|
||||
resource_found = self._strip_proxy_helper(resource_found, rollout)
|
||||
|
||||
return resource_found
|
||||
|
||||
def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
|
||||
"""Retrieve the first prompt template resource from the available resources.
|
||||
|
||||
Args:
|
||||
resources: Mapping of named resources.
|
||||
rollout: Rollout metadata (unused).
|
||||
|
||||
Returns:
|
||||
First [`PromptTemplate`][agentlightning.PromptTemplate] resource encountered.
|
||||
|
||||
Raises:
|
||||
ValueError: If no prompt template resource is present.
|
||||
"""
|
||||
resource_found: PromptTemplate | None = None
|
||||
for name, resource in resources.items():
|
||||
if isinstance(resource, PromptTemplate):
|
||||
if resource_found is not None:
|
||||
logger.warning(
|
||||
f"Multiple prompt template resources found in resources. Using the first one: '{name}'."
|
||||
)
|
||||
break
|
||||
resource_found = resource
|
||||
|
||||
if resource_found is None:
|
||||
raise ValueError("No prompt template resource found in the provided resources.")
|
||||
|
||||
return resource_found
|
||||
|
||||
def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
|
||||
"""Convert [`ProxyLLM`][agentlightning.ProxyLLM] instances into concrete LLMs.
|
||||
|
||||
It resolves ProxyLLM instances to their concrete LLM implementation
|
||||
by attaching the attempted rollout context. This is only used when the function
|
||||
signature accepts an `llm` parameter and strip_proxy is True.
|
||||
|
||||
Args:
|
||||
proxy_llm: Candidate LLM resource.
|
||||
rollout: Rollout metadata that provides rollout and attempt identifiers.
|
||||
|
||||
Returns:
|
||||
[`LLM`][agentlightning.LLM] with rollout context baked into the endpoint.
|
||||
|
||||
Raises:
|
||||
ValueError: If the rollout is not an
|
||||
[`AttemptedRollout`][agentlightning.AttemptedRollout].
|
||||
"""
|
||||
|
||||
if not isinstance(proxy_llm, ProxyLLM):
|
||||
# Not a ProxyLLM, nothing to strip here.
|
||||
return proxy_llm
|
||||
|
||||
# Rollout is still a Rollout here because API is not stabilized yet.
|
||||
# In practice, it must be an AttemptedRollout.
|
||||
if not isinstance(rollout, AttemptedRollout):
|
||||
raise ValueError("Rollout is not an AttemptedRollout.")
|
||||
|
||||
return proxy_llm.with_attempted_rollout(rollout)
|
||||
|
||||
|
||||
@overload
|
||||
def llm_rollout(func: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def llm_rollout(*, strip_proxy: bool = True) -> Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]: ...
|
||||
|
||||
|
||||
def llm_rollout(
|
||||
func: LlmRolloutFunc[T] | None = None, *, strip_proxy: bool = True
|
||||
) -> FunctionalLitAgent[T] | Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]:
|
||||
"""Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for LLM-based rollouts.
|
||||
|
||||
Args:
|
||||
func: Callable defining the agent's behaviour. Supported signatures include:
|
||||
|
||||
* `(task, llm) -> result`
|
||||
* `(task, llm, rollout) -> result`
|
||||
* `async (task, llm) -> result`
|
||||
* `async (task, llm, rollout) -> result`
|
||||
|
||||
strip_proxy: When `True`, convert proxy resources into concrete
|
||||
[`LLM`][agentlightning.LLM] instances before calling the
|
||||
function. Defaults to `True`.
|
||||
|
||||
Returns:
|
||||
[`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
|
||||
wraps the supplied function.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
@llm_rollout
|
||||
def my_agent(task, llm):
|
||||
return llm.endpoint
|
||||
|
||||
@llm_rollout(strip_proxy=False)
|
||||
def my_agent_no_strip(task, llm):
|
||||
return llm.model
|
||||
|
||||
result = my_agent(task, llm)
|
||||
result = my_agent.rollout(task, resources, rollout)
|
||||
```
|
||||
"""
|
||||
|
||||
def decorator(f: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]:
|
||||
_validate_llm_rollout_func(f)
|
||||
return FunctionalLitAgent(f, strip_proxy=strip_proxy)
|
||||
|
||||
if func is None:
|
||||
# Called with arguments: @llm_rollout(strip_proxy=False)
|
||||
return decorator
|
||||
else:
|
||||
# Called without arguments: @llm_rollout
|
||||
return decorator(func)
|
||||
|
||||
|
||||
def _validate_llm_rollout_func(func: Any) -> TypeGuard[LlmRolloutFunc[Any]]:
|
||||
"""Validate the function signature of an LLM rollout function.
|
||||
|
||||
Ensures the function follows the expected pattern for LLM-based rollouts:
|
||||
|
||||
- Must have at least 2 parameters
|
||||
- First parameter must be named 'task'
|
||||
- Must have a parameter named 'llm'
|
||||
- Optionally can have a 'rollout' parameter
|
||||
|
||||
Args:
|
||||
func: Function to inspect.
|
||||
|
||||
Returns:
|
||||
`True` when the signature matches the supported patterns.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function signature does not match the expected pattern.
|
||||
"""
|
||||
sig = inspect.signature(func)
|
||||
params = list(sig.parameters.keys())
|
||||
if len(params) < 2:
|
||||
raise ValueError(f"Function {func} must have at least 2 parameters.")
|
||||
if params[0] != "task":
|
||||
raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
|
||||
if "llm" not in params:
|
||||
raise ValueError(f"Function {func} must have a positional parameter called 'llm'.")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@overload
|
||||
def prompt_rollout(func: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def prompt_rollout() -> Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]: ...
|
||||
|
||||
|
||||
def prompt_rollout(
|
||||
func: PromptRolloutFunc[T] | None = None,
|
||||
) -> FunctionalLitAgent[T] | Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
|
||||
"""Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] for prompt-based rollouts.
|
||||
|
||||
This decorator is designed for agents that work with tunable prompt templates. It enables
|
||||
a workflow where algorithms manage and optimize the prompt template, while agents consume
|
||||
the template to perform rollouts. This is particularly useful for prompt optimization scenarios.
|
||||
|
||||
Args:
|
||||
func: Callable defining the agent's behavior. Supported signatures include:
|
||||
|
||||
* `(task, prompt_template) -> result`
|
||||
* `(task, prompt_template, rollout) -> result`
|
||||
* `async (task, prompt_template) -> result`
|
||||
* `async (task, prompt_template, rollout) -> result`
|
||||
|
||||
Returns:
|
||||
[`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
|
||||
wraps the supplied function.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
@prompt_rollout
|
||||
def my_agent(task, prompt_template):
|
||||
messages = prompt_template.format(task=task.input)
|
||||
return messages
|
||||
|
||||
result = my_agent(task, prompt_template)
|
||||
result = my_agent.rollout(task, resources, rollout)
|
||||
```
|
||||
"""
|
||||
|
||||
def decorator(f: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]:
|
||||
_validate_prompt_rollout_func(f)
|
||||
return FunctionalLitAgent(f)
|
||||
|
||||
if func is None:
|
||||
return decorator
|
||||
else:
|
||||
return decorator(func)
|
||||
|
||||
|
||||
def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]]:
|
||||
"""Validate the function signature of a prompt rollout function.
|
||||
|
||||
Ensures the function follows the expected pattern for prompt-template-based rollouts:
|
||||
|
||||
- Must have at least 2 parameters
|
||||
- First parameter must be named 'task'
|
||||
- Must have a parameter named 'prompt_template'
|
||||
- Optionally can have a 'rollout' parameter
|
||||
|
||||
Args:
|
||||
func: Function to inspect.
|
||||
|
||||
Returns:
|
||||
`True` when the signature matches the supported patterns.
|
||||
|
||||
Raises:
|
||||
ValueError: If the function signature does not match the expected pattern.
|
||||
"""
|
||||
sig = inspect.signature(func)
|
||||
params = list(sig.parameters.keys())
|
||||
if len(params) < 2:
|
||||
raise ValueError(f"Function {func} must have at least 2 parameters.")
|
||||
if params[0] != "task":
|
||||
raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
|
||||
if "prompt_template" not in params:
|
||||
raise ValueError(f"Function {func} must have a positional parameter called 'prompt_template'.")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., Any]]) -> FunctionalLitAgent[T]:
|
||||
"""Create a [`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] from an arbitrary rollout function.
|
||||
|
||||
This function inspects the provided callable and creates the appropriate
|
||||
agent type based on its signature. It supports both LLM-based and prompt-template-based
|
||||
agents. The returned agent instance is callable, preserving the original function's
|
||||
behavior and type hints.
|
||||
|
||||
See [`llm_rollout`][agentlightning.litagent.decorator.llm_rollout] and
|
||||
[`prompt_rollout`][agentlightning.litagent.decorator.prompt_rollout] for more details.
|
||||
|
||||
Args:
|
||||
func: Callable that implements the rollout. Supported signatures:
|
||||
|
||||
- `[async ](task, llm[, rollout])` for LLM-based agents
|
||||
- `[async ](task, prompt_template[, rollout])` for prompt-template-based agents
|
||||
|
||||
The supported output types of `func` is same as the return type of [`rollout`][agentlightning.LitAgent.rollout].
|
||||
|
||||
Returns:
|
||||
[`FunctionalLitAgent`][agentlightning.litagent.decorator.FunctionalLitAgent] that
|
||||
wraps the supplied function.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
# LLM-based agent
|
||||
@rollout
|
||||
def my_llm_agent(task, llm):
|
||||
client = OpenAI(base_url=llm.endpoint)
|
||||
response = client.chat.completions.create(
|
||||
model=llm.model,
|
||||
messages=[{"role": "user", "content": task.input}],
|
||||
)
|
||||
return response
|
||||
|
||||
# Prompt-template-based agent
|
||||
@rollout
|
||||
def my_prompt_agent(task, prompt_template):
|
||||
messages = prompt_template.format(task=task.input)
|
||||
# ... perform rollout with the formatted prompt
|
||||
return response
|
||||
|
||||
# Function is still callable with original behavior
|
||||
result = my_llm_agent(task, llm)
|
||||
|
||||
# Agent methods are also available
|
||||
result = my_llm_agent.rollout(task, resources, rollout)
|
||||
```
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the function signature doesn't match any known patterns.
|
||||
"""
|
||||
# Check if it matches the LLM rollout API pattern
|
||||
sig = inspect.signature(func)
|
||||
|
||||
try:
|
||||
if _validate_llm_rollout_func(func):
|
||||
return llm_rollout(func)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
try:
|
||||
if _validate_prompt_rollout_func(func):
|
||||
return prompt_rollout(func)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
raise NotImplementedError(
|
||||
f"Function signature {sig} does not match any known agent patterns. "
|
||||
"Expected signatures: (task, llm[, rollout]) or (task, prompt_template[, rollout]). "
|
||||
"Functions can be sync or async."
|
||||
)
|
||||
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Base abstractions for building agents that plug into Agent Lightning."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
import warnings
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
|
||||
|
||||
from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.runner import Runner
|
||||
from agentlightning.tracer import Tracer
|
||||
from agentlightning.trainer import Trainer
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
__all__ = [
|
||||
"LitAgent",
|
||||
]
|
||||
|
||||
|
||||
def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
|
||||
"""Return `True` when the rollout function uses the deprecated v0.1 signature.
|
||||
|
||||
The helper inspects the callable's signature to detect whether a `rollout_id`
|
||||
parameter is present, which indicates the legacy API.
|
||||
|
||||
Args:
|
||||
func: Function to analyze.
|
||||
|
||||
Returns:
|
||||
`True` if the callable exposes a `rollout_id` parameter.
|
||||
"""
|
||||
return "rollout_id" in inspect.signature(func).parameters
|
||||
|
||||
|
||||
class LitAgent(Generic[T]):
|
||||
"""Base class for implementing agent rollouts.
|
||||
|
||||
Subclasses override the rollout methods to process tasks while the trainer and
|
||||
runner infrastructure manages orchestration, tracing, and persistence.
|
||||
"""
|
||||
|
||||
def __init__(self, *, trained_agents: Optional[str] = None) -> None: # FIXME: str | None won't work for cli
|
||||
"""Initialize the agent instance.
|
||||
|
||||
Args:
|
||||
trained_agents: Optional identifier used by legacy tooling to mark trained
|
||||
agents.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
The `trained_agents` flag is deprecated. Configure `agent_match` in the adapter
|
||||
layer instead. See [`TracerTraceToTriplet`][agentlightning.TracerTraceToTriplet]
|
||||
for more details.
|
||||
"""
|
||||
if trained_agents is not None:
|
||||
warnings.warn(
|
||||
"`trained_agents` is deprecated. Configure `agent_match` in adapter instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self.trained_agents = trained_agents
|
||||
|
||||
self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
|
||||
self._runner_ref: weakref.ReferenceType[Runner[T]] | None = None
|
||||
|
||||
def is_async(self) -> bool:
|
||||
"""Return `True` when the agent overrides any asynchronous rollout methods.
|
||||
|
||||
Override this method for customized async detection logic.
|
||||
"""
|
||||
return (
|
||||
(
|
||||
hasattr(self, "training_rollout_async")
|
||||
and self.__class__.training_rollout_async is not LitAgent.training_rollout_async # type: ignore
|
||||
)
|
||||
or (
|
||||
hasattr(self, "validation_rollout_async")
|
||||
and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async # type: ignore
|
||||
)
|
||||
or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async) # type: ignore
|
||||
)
|
||||
|
||||
def set_trainer(self, trainer: Trainer) -> None:
|
||||
"""Attach the trainer responsible for orchestration.
|
||||
|
||||
Args:
|
||||
trainer: [`Trainer`][agentlightning.Trainer] that manages the agent.
|
||||
"""
|
||||
self._trainer_ref = weakref.ref(trainer)
|
||||
|
||||
def get_trainer(self) -> Trainer:
|
||||
"""Return the trainer associated with this agent."""
|
||||
if self._trainer_ref is None:
|
||||
raise ValueError("Trainer has not been set for this agent.")
|
||||
trainer = self._trainer_ref()
|
||||
if trainer is None:
|
||||
raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
|
||||
return trainer
|
||||
|
||||
@property
|
||||
def trainer(self) -> Trainer:
|
||||
"""Return the trainer associated with this agent."""
|
||||
return self.get_trainer()
|
||||
|
||||
def get_tracer(self) -> Tracer:
|
||||
"""Return the tracer configured for this agent."""
|
||||
if hasattr(self.runner, "tracer"):
|
||||
return self.runner.tracer # type: ignore
|
||||
else:
|
||||
return self.trainer.tracer
|
||||
|
||||
@property
|
||||
def tracer(self) -> Tracer:
|
||||
"""Return the tracer configured for this agent."""
|
||||
return self.get_tracer()
|
||||
|
||||
def set_runner(self, runner: Runner[T]) -> None:
|
||||
"""Attach the runner responsible for executing rollouts.
|
||||
|
||||
Args:
|
||||
runner: [`Runner`][agentlightning.Runner] coordinating execution.
|
||||
"""
|
||||
self._runner_ref = weakref.ref(runner)
|
||||
|
||||
def get_runner(self) -> Runner[T]:
|
||||
"""Return the runner responsible for executing rollouts."""
|
||||
if self._runner_ref is None:
|
||||
raise ValueError("Runner has not been set for this agent.")
|
||||
runner = self._runner_ref()
|
||||
if runner is None:
|
||||
raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
|
||||
return runner
|
||||
|
||||
@property
|
||||
def runner(self) -> Runner[T]:
|
||||
"""Return the runner responsible for executing rollouts."""
|
||||
return self.get_runner()
|
||||
|
||||
def on_rollout_start(self, task: Task, runner: Runner[T], tracer: Tracer) -> None:
|
||||
"""Hook invoked immediately before a rollout begins.
|
||||
|
||||
Subclasses can override this method to implement custom logic such as logging,
|
||||
metric collection, or resource setup. The default implementation is a no-op.
|
||||
|
||||
Args:
|
||||
task: [`Task`][agentlightning.Task] that will be processed.
|
||||
runner: [`Runner`][agentlightning.Runner] managing the rollout.
|
||||
tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
Override [`Hook.on_rollout_start`][agentlightning.Hook.on_rollout_start]
|
||||
instead of this method when extending agents.
|
||||
"""
|
||||
|
||||
def on_rollout_end(self, task: Task, rollout: Rollout, runner: Runner[T], tracer: Tracer) -> None:
|
||||
"""Hook invoked after a rollout completes.
|
||||
|
||||
Subclasses can override this method for cleanup or additional logging. The default
|
||||
implementation is a no-op.
|
||||
|
||||
Args:
|
||||
task: [`Task`][agentlightning.Task] that was processed.
|
||||
rollout: Resulting [`Rollout`][agentlightning.Rollout].
|
||||
runner: [`Runner`][agentlightning.Runner] managing the rollout.
|
||||
tracer: [`Tracer`][agentlightning.Tracer] associated with the runner.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
Override [`Hook.on_rollout_end`][agentlightning.Hook.on_rollout_end]
|
||||
instead of this method when extending agents.
|
||||
"""
|
||||
|
||||
def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Execute a rollout synchronously.
|
||||
|
||||
|
||||
If you don't wish to implement both training rollout and validation
|
||||
rollout separately, you can just implement `rollout` which will work for both.
|
||||
|
||||
Args:
|
||||
task: Task payload provided by the scheduler.
|
||||
resources: Mapping of named resources (for example LLMs or prompt templates).
|
||||
rollout: Rollout metadata. Avoid mutating this object directly unless a
|
||||
subclass needs to override defaults.
|
||||
|
||||
Returns:
|
||||
One of the following values:
|
||||
|
||||
* `None` when tracing is handled by the runner.
|
||||
* `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.")
|
||||
|
||||
async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Execute a rollout asynchronously.
|
||||
|
||||
Args:
|
||||
task: Task payload provided by the scheduler.
|
||||
resources: Mapping of named resources (for example LLMs or prompt templates).
|
||||
rollout: Rollout metadata. Avoid mutating this object directly unless a
|
||||
subclass needs to override defaults.
|
||||
|
||||
Returns:
|
||||
Same possible return values as
|
||||
[`rollout`][agentlightning.LitAgent.rollout].
|
||||
"""
|
||||
raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
|
||||
|
||||
def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Process a single training task synchronously.
|
||||
|
||||
By default, this method delegates to
|
||||
[`rollout`][agentlightning.LitAgent.rollout].
|
||||
"""
|
||||
return self.rollout(task, resources, rollout)
|
||||
|
||||
def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Process a single validation task synchronously.
|
||||
|
||||
Override this method when validation should differ from training. The default
|
||||
implementation delegates to
|
||||
[`training_rollout`][agentlightning.LitAgent.training_rollout].
|
||||
"""
|
||||
return self.rollout(task, resources, rollout)
|
||||
|
||||
async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Process a single training task asynchronously.
|
||||
|
||||
By default, this method delegates to
|
||||
[`rollout_async`][agentlightning.LitAgent.rollout_async].
|
||||
"""
|
||||
return await self.rollout_async(task, resources, rollout)
|
||||
|
||||
async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
|
||||
"""Process a single validation task asynchronously.
|
||||
|
||||
Override this method when validation should differ from training. The default
|
||||
implementation delegates to
|
||||
[`training_rollout_async`][agentlightning.LitAgent.training_rollout_async].
|
||||
"""
|
||||
return await self.rollout_async(task, resources, rollout)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,370 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
import warnings
|
||||
from logging.config import dictConfig
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
__all__ = ["setup", "configure_logger", "setup_module"]
|
||||
|
||||
|
||||
def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
|
||||
"""Create or reset a namespaced logger with a consistent console format.
|
||||
|
||||
This helper clears any previously attached handlers before binding a single
|
||||
`StreamHandler` that writes to standard output. The resulting logger does
|
||||
not propagate to the root logger, preventing duplicate log emission when
|
||||
applications compose multiple logging configurations.
|
||||
|
||||
!!! danger
|
||||
|
||||
This function is deprecated in favor of [`setup_logging`][agentlightning.setup_logging].
|
||||
|
||||
Args:
|
||||
level: Logging level applied both to the logger and the installed
|
||||
handler. Defaults to `logging.INFO`.
|
||||
name: Dotted path for the logger instance. Defaults to
|
||||
`"agentlightning"`.
|
||||
|
||||
Returns:
|
||||
Configured logger instance ready for immediate use.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
from agentlightning import configure_logger
|
||||
|
||||
logger = configure_logger(level=logging.INFO)
|
||||
logger.info("agent-lightning is ready!")
|
||||
```
|
||||
"""
|
||||
warnings.warn("This function is deprecated in favor of `setup_logging`.", DeprecationWarning, stacklevel=2)
|
||||
|
||||
return setup_module(level=level, name=name, console=True, color=True, propagate=False)
|
||||
|
||||
|
||||
DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s) %(message)s"
|
||||
DATE_FORMAT = "%H:%M:%S"
|
||||
|
||||
|
||||
def _to_level_value(lvl: int | str) -> int:
|
||||
if isinstance(lvl, int):
|
||||
return lvl
|
||||
val = getattr(logging, str(lvl).upper(), None)
|
||||
if val is None:
|
||||
raise ValueError(f"Invalid log level: {lvl}")
|
||||
return val
|
||||
|
||||
|
||||
def _ensure_file_handler(
|
||||
logger: logging.Logger,
|
||||
filename: str,
|
||||
*,
|
||||
level: int,
|
||||
formatter: Optional[logging.Formatter],
|
||||
) -> None:
|
||||
"""Attach a FileHandler to `logger` for `filename` if it doesn't already exist."""
|
||||
abspath = os.path.abspath(filename)
|
||||
|
||||
# Avoid duplicates
|
||||
for h in logger.handlers:
|
||||
if isinstance(h, logging.FileHandler) and getattr(h, "baseFilename", None) == abspath:
|
||||
return
|
||||
|
||||
# Ensure directory exists
|
||||
dirname = os.path.dirname(abspath)
|
||||
if dirname:
|
||||
os.makedirs(dirname, exist_ok=True)
|
||||
|
||||
fh = logging.FileHandler(abspath, encoding="utf-8")
|
||||
fh.setLevel(level)
|
||||
if formatter is not None:
|
||||
fh.setFormatter(formatter)
|
||||
else:
|
||||
fh.setFormatter(logging.Formatter(DEFAULT_FORMAT, DATE_FORMAT))
|
||||
|
||||
logger.addHandler(fh)
|
||||
|
||||
|
||||
def setup(
|
||||
level: int | str = "INFO",
|
||||
*,
|
||||
console: bool = True,
|
||||
color: bool | Dict[str, Any] = True,
|
||||
propagate: bool = False,
|
||||
disable_existing_loggers: bool = False,
|
||||
capture_warnings: bool = False,
|
||||
submodule_levels: Optional[dict[str, int | str]] = None,
|
||||
extra_handlers: Optional[list[logging.Handler]] = None,
|
||||
formatter: Optional[logging.Formatter] = None,
|
||||
apply_to: Optional[list[str]] = None,
|
||||
files: Optional[str | dict[str, str]] = None,
|
||||
) -> None:
|
||||
"""Configures logging for the `agentlightning` logger hierarchy.
|
||||
|
||||
This function provides a one-stop setup utility for configuring the
|
||||
`agentlightning` root logger and optionally its submodules or external
|
||||
loggers. It supports console logging, colored rich output, per-submodule
|
||||
log levels, and optional handler/formatter injection.
|
||||
|
||||
The setup is intentionally isolated: it does not modify the global root
|
||||
logger or loggers belonging to other libraries unless explicitly directed
|
||||
via `apply_to`.
|
||||
|
||||
Args:
|
||||
level:
|
||||
Logging level for the base `agentlightning` logger. Accepts either
|
||||
an integer (e.g., `logging.DEBUG`) or a string level name
|
||||
(e.g., `"INFO"`). Defaults to `"INFO"`.
|
||||
console:
|
||||
Whether to attach a console handler to the logger. Defaults to
|
||||
`True`.
|
||||
color:
|
||||
Enables rich-formatted output using `RichHandler` when `True`
|
||||
or a configuration dict. If `False`, a plain text formatter is
|
||||
used instead. Defaults to `True`.
|
||||
propagate:
|
||||
Whether `agentlightning` logs should propagate to ancestor
|
||||
loggers. Defaults to `False`.
|
||||
disable_existing_loggers:
|
||||
Passed to `logging.config.dictConfig`. If `True`, disables all
|
||||
existing configured loggers before applying this configuration.
|
||||
Defaults to `False`.
|
||||
capture_warnings:
|
||||
If `True`, redirects Python `warnings` emitted via the `warnings`
|
||||
module into the logging system. Defaults to `False`.
|
||||
submodule_levels:
|
||||
Mapping of submodule logger names to logging levels. If a specified
|
||||
submodule level is more verbose than the base level, a warning is emitted.
|
||||
extra_handlers:
|
||||
A list of user-provided handlers to attach to the `agentlightning` logger.
|
||||
Handlers are added idempotently; duplicates are not reattached.
|
||||
formatter:
|
||||
A formatter to apply to any handler under `agentlightning` that does not
|
||||
already have one assigned. Useful for customizing output without overwriting
|
||||
formatters on custom handlers.
|
||||
apply_to:
|
||||
A list of additional logger names to configure identically to
|
||||
`agentlightning` base logger. Their handlers are replaced with copies of the base
|
||||
handlers, and propagation is disabled to avoid duplicate log emission.
|
||||
files:
|
||||
If a string, attach a FileHandler to the base `agentlightning` logger.
|
||||
If a dict, for each `(logger_name, filename)` pair, attach a FileHandler
|
||||
directly to that logger.
|
||||
Each file handler should use the logger's effective level at creation.
|
||||
|
||||
Notes:
|
||||
* On Windows, this function forces UTF-8 mode in the console to prevent
|
||||
issues with rich output or special characters.
|
||||
* Submodule loggers can generate records below the handler's emission
|
||||
threshold. Whether such records appear depends on both the logger's
|
||||
level and the handler's level.
|
||||
* `apply_to` loggers inherit the same handlers but do not propagate
|
||||
upward, yielding isolated, consistent behavior.
|
||||
|
||||
Examples:
|
||||
Basic setup:
|
||||
|
||||
>>> setup()
|
||||
|
||||
Enabling debug mode with no color:
|
||||
|
||||
>>> setup(level="DEBUG", color=False)
|
||||
|
||||
Overriding specific submodule levels:
|
||||
|
||||
>>> setup(submodule_levels={"agentlightning.io": "DEBUG"})
|
||||
|
||||
Attaching an additional file handler:
|
||||
|
||||
>>> fh = logging.FileHandler("app.log")
|
||||
>>> setup(extra_handlers=[fh])
|
||||
"""
|
||||
# Ensure UTF-8 encoding on Windows consoles
|
||||
# Note: This change does not fully represent support for execution under the windows system.
|
||||
# It only fixes console printing issues caused by special characters.
|
||||
# TODO: More comprehensive Windows support may be needed in the future.
|
||||
if platform.system() == "Windows":
|
||||
os.environ["PYTHONUTF8"] = "1"
|
||||
|
||||
base_logger = setup_module(
|
||||
level,
|
||||
name="agentlightning",
|
||||
console=console,
|
||||
color=color,
|
||||
propagate=propagate,
|
||||
disable_existing_loggers=disable_existing_loggers,
|
||||
)
|
||||
|
||||
base_level_value = base_logger.level
|
||||
|
||||
# Apply user-provided formatter (only to handlers without one,
|
||||
# so we don't clobber custom extra_handlers)
|
||||
if formatter is not None:
|
||||
for h in base_logger.handlers:
|
||||
if h.formatter is None:
|
||||
h.setFormatter(formatter)
|
||||
|
||||
# Attach user-provided handler(s) if any, idempotently
|
||||
if extra_handlers:
|
||||
for h in extra_handlers:
|
||||
if h not in base_logger.handlers:
|
||||
base_logger.addHandler(h)
|
||||
|
||||
# Per-submodule levels
|
||||
if submodule_levels:
|
||||
for name, lvl in submodule_levels.items():
|
||||
sub_level = _to_level_value(lvl)
|
||||
|
||||
# Emit a warning if submodule level is lower (more verbose) than the global/base level
|
||||
if sub_level < base_level_value:
|
||||
base_logger.warning(
|
||||
"Submodule logger '%s' level %s (%s) is more verbose than base "
|
||||
"logger level %s (%s). Records below the base level may still be "
|
||||
"filtered out by handlers depending on their own levels.",
|
||||
name,
|
||||
lvl,
|
||||
sub_level,
|
||||
logging.getLevelName(base_level_value),
|
||||
base_level_value,
|
||||
)
|
||||
|
||||
# The logger will *create* records down to the logger's level, but a handler
|
||||
# with a higher level will still drop anything below its own threshold.
|
||||
# Effective emission is gated by both: record.level >= logger.level AND handler.level.
|
||||
logging.getLogger(name).setLevel(lvl)
|
||||
|
||||
# Attach file handlers if requested
|
||||
if files is not None:
|
||||
if isinstance(files, str):
|
||||
# Single file for the entire `agentlightning` hierarchy.
|
||||
_ensure_file_handler(
|
||||
logger=base_logger,
|
||||
filename=files,
|
||||
level=base_level_value,
|
||||
formatter=formatter,
|
||||
)
|
||||
else:
|
||||
# Per-logger files
|
||||
for logger_name, filename in files.items():
|
||||
lg = logging.getLogger(logger_name)
|
||||
# Use the logger's *effective* level at creation time
|
||||
effective_level = lg.getEffectiveLevel()
|
||||
_ensure_file_handler(
|
||||
logger=lg,
|
||||
filename=filename,
|
||||
level=effective_level,
|
||||
formatter=formatter,
|
||||
)
|
||||
|
||||
# Optionally apply the same handler setup to other loggers outside this module
|
||||
if apply_to:
|
||||
for name in apply_to:
|
||||
lg = logging.getLogger(name)
|
||||
# This removes any existing handlers so we don't duplicate output
|
||||
# and ensures these loggers share exactly the same handlers as base_logger.
|
||||
lg.handlers.clear()
|
||||
for h in base_logger.handlers:
|
||||
lg.addHandler(h)
|
||||
lg.setLevel(base_logger.level)
|
||||
# We've attached handlers directly to these loggers; if propagate
|
||||
# stayed True, records would bubble up to ancestor loggers and could be
|
||||
# emitted twice (here and on the parent/root). Setting False isolates them.
|
||||
lg.propagate = False
|
||||
|
||||
# Optionally capture warnings
|
||||
if capture_warnings:
|
||||
logging.captureWarnings(True)
|
||||
|
||||
|
||||
def setup_module(
|
||||
level: int | str = "INFO",
|
||||
*,
|
||||
name: str = "agentlightning",
|
||||
console: bool = True,
|
||||
color: bool | Dict[str, Any] = True,
|
||||
propagate: bool = False,
|
||||
disable_existing_loggers: bool = False,
|
||||
) -> logging.Logger:
|
||||
"""Initializes and returns the base logger for `agentlightning`.
|
||||
|
||||
This function constructs and applies a `dictConfig` configuration for the
|
||||
logger hierarchy rooted at `name`. It supports either rich console
|
||||
formatting (via `RichHandler`) or plain text formatting, based on the
|
||||
`color` argument.
|
||||
|
||||
Unlike [`setup_logging`][agentlightning.setup_logging], this function configures only a single logger namespace
|
||||
and does not attach extra handlers or submodule levels. It is primarily used
|
||||
internally by [`setup_logging`][agentlightning.setup_logging] but is also suitable for direct integration in
|
||||
custom logging workflows.
|
||||
"""
|
||||
root_cfg: Dict[str, Any] = {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": disable_existing_loggers,
|
||||
"loggers": {
|
||||
name: {
|
||||
"handlers": [],
|
||||
"level": level,
|
||||
"propagate": propagate,
|
||||
}
|
||||
},
|
||||
"handlers": {},
|
||||
"formatters": {},
|
||||
}
|
||||
|
||||
# Choose formatter / handler definition
|
||||
if color is not False and console:
|
||||
# Console must be true to display colored outputs
|
||||
if isinstance(color, dict):
|
||||
rich_handler_config = color
|
||||
else:
|
||||
rich_handler_config: Dict[str, Any] = {
|
||||
"rich_tracebacks": False,
|
||||
"markup": False,
|
||||
"show_time": True,
|
||||
"show_path": True,
|
||||
}
|
||||
|
||||
if not _has_width():
|
||||
# e.g., in a CI environment.
|
||||
rich_handler_config["console"] = Console(width=200)
|
||||
|
||||
root_cfg["handlers"]["console"] = {
|
||||
"class": "rich.logging.RichHandler",
|
||||
"level": level,
|
||||
**rich_handler_config,
|
||||
}
|
||||
# RichHandler manages its own style; keep formatter None
|
||||
else:
|
||||
fmt_name = "plain"
|
||||
root_cfg["formatters"][fmt_name] = {
|
||||
"format": DEFAULT_FORMAT,
|
||||
"datefmt": DATE_FORMAT,
|
||||
}
|
||||
|
||||
if console:
|
||||
root_cfg["handlers"]["console"] = {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": level,
|
||||
"formatter": fmt_name,
|
||||
}
|
||||
|
||||
# Attach selected handlers to agentlightning
|
||||
handler_names = list(root_cfg["handlers"].keys())
|
||||
root_cfg["loggers"][name]["handlers"] = handler_names
|
||||
|
||||
# Apply dictConfig (this resets the logger handlers)
|
||||
dictConfig(root_cfg)
|
||||
|
||||
return logging.getLogger(name)
|
||||
|
||||
|
||||
def _has_width() -> bool:
|
||||
"""Automatically determine whether the terminal has a width."""
|
||||
return sys.stdout.isatty()
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import warnings
|
||||
|
||||
from .emitter.reward import * # noqa: F401,F403
|
||||
|
||||
warnings.warn("agentlightning.reward is deprecated. Please use agentlightning.emitter instead.")
|
||||
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .agent import LitAgentRunner
|
||||
from .base import Runner
|
||||
from .legacy import LegacyAgentRunner
|
||||
|
||||
__all__ = [
|
||||
"Runner",
|
||||
"LegacyAgentRunner",
|
||||
"LitAgentRunner",
|
||||
]
|
||||
@@ -0,0 +1,841 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Agent runner implementation for executing agent rollouts.
|
||||
|
||||
This module provides the concrete implementation of the runner interface,
|
||||
handling the execution of agent rollouts with support for tracing, hooks,
|
||||
and distributed worker coordination.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Awaitable,
|
||||
Callable,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
from agentlightning.litagent import LitAgent
|
||||
from agentlightning.reward import emit_reward, find_final_reward
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.tracer.base import Tracer
|
||||
from agentlightning.tracer.otel import OtelTracer
|
||||
from agentlightning.types import (
|
||||
AttemptedRollout,
|
||||
Hook,
|
||||
NamedResources,
|
||||
Rollout,
|
||||
RolloutMode,
|
||||
RolloutRawResult,
|
||||
Span,
|
||||
SpanCoreFields,
|
||||
)
|
||||
from agentlightning.utils.system_snapshot import system_snapshot
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.execution.events import ExecutionEvent
|
||||
|
||||
from .base import Runner
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LitAgentRunner(Runner[T_task]):
|
||||
"""Execute [`LitAgent`][agentlightning.LitAgent] tasks with tracing support.
|
||||
|
||||
This runner manages the complete lifecycle of agent rollout execution,
|
||||
including task polling, resource management, tracing, and hooks. It supports
|
||||
both continuous iteration over tasks from the store and single-step execution.
|
||||
|
||||
Attributes:
|
||||
worker_id: Identifier for the active worker process, if any.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tracer: Tracer,
|
||||
max_rollouts: Optional[int] = None,
|
||||
poll_interval: float = 5.0,
|
||||
heartbeat_interval: float = 10.0,
|
||||
interval_jitter: float = 0.5,
|
||||
heartbeat_launch_mode: Literal["asyncio", "thread"] = "thread",
|
||||
heartbeat_include_gpu: bool = False,
|
||||
) -> None:
|
||||
"""Initialize the agent runner.
|
||||
|
||||
Args:
|
||||
tracer: [`Tracer`][agentlightning.Tracer] used for rollout spans.
|
||||
max_rollouts: Optional cap on iterations processed by
|
||||
[`iter`][agentlightning.LitAgentRunner.iter].
|
||||
poll_interval: Seconds to wait between store polls when no work is available.
|
||||
heartbeat_interval: Seconds to wait between sending heartbeats to the store.
|
||||
interval_jitter: Jitter factor for the poll interval. The actual interval will be between
|
||||
poll_interval - interval_jitter and poll_interval + interval_jitter.
|
||||
This is to avoid the overload caused by the synchronization of the runners.
|
||||
heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
|
||||
"thread" is the default and recommended mode as it prevents blocking the event loop
|
||||
under load. Use "asyncio" for simpler deployments with low worker counts.
|
||||
heartbeat_include_gpu: Whether to include GPU stats in heartbeat snapshots.
|
||||
Querying GPU stats can be slow under load, so this is disabled by default.
|
||||
"""
|
||||
super().__init__()
|
||||
self._tracer = tracer
|
||||
self._max_rollouts = max_rollouts
|
||||
self._poll_interval = poll_interval
|
||||
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
|
||||
self._agent: Optional[LitAgent[T_task]] = None
|
||||
self._hooks: Sequence[Hook] = []
|
||||
self._store: Optional[LightningStore] = None
|
||||
self.worker_id: Optional[int] = None
|
||||
|
||||
def init(self, agent: LitAgent[T_task], *, hooks: Optional[Sequence[Hook]] = None, **kwargs: Any) -> None:
|
||||
"""Initialize the runner with the agent.
|
||||
|
||||
This sets up the agent-runner relationship, registers hooks, and
|
||||
initializes the tracer.
|
||||
|
||||
Args:
|
||||
agent: [`LitAgent`][agentlightning.LitAgent] instance executed by the runner.
|
||||
hooks: Optional sequence of [`Hook`][agentlightning.Hook]
|
||||
callbacks invoked around tracing and rollout boundaries.
|
||||
**kwargs: Additional initialization arguments (currently unused).
|
||||
"""
|
||||
self._agent = agent
|
||||
self._agent.set_runner(self)
|
||||
self._hooks = [*hooks] if hooks is not None else []
|
||||
|
||||
self._tracer.init()
|
||||
|
||||
def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
|
||||
"""Initialize the runner for each worker with worker_id and store.
|
||||
|
||||
This method is called once per worker in a distributed setup to provide
|
||||
the worker with its ID and store connection.
|
||||
|
||||
Args:
|
||||
worker_id: Unique identifier for this worker process.
|
||||
store: [`LightningStore`][agentlightning.LightningStore]
|
||||
used for task coordination and persistence.
|
||||
**kwargs: Additional worker-specific initialization arguments (currently unused).
|
||||
"""
|
||||
self._store = store
|
||||
self.worker_id = worker_id
|
||||
|
||||
self._tracer.init_worker(worker_id, store)
|
||||
|
||||
def teardown(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Teardown the runner and clean up all resources.
|
||||
|
||||
This method resets all internal state including the agent, store,
|
||||
hooks, and worker ID, and calls the tracer's teardown method.
|
||||
|
||||
Args:
|
||||
*args: Additional teardown arguments (currently unused).
|
||||
**kwargs: Additional teardown keyword arguments (currently unused).
|
||||
"""
|
||||
self._agent = None
|
||||
self._store = None
|
||||
self.worker_id = None
|
||||
self._hooks = []
|
||||
|
||||
self._tracer.teardown()
|
||||
|
||||
def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
|
||||
"""Teardown the runner for a specific worker.
|
||||
|
||||
This method cleans up worker-specific resources and resets the worker ID.
|
||||
|
||||
Args:
|
||||
worker_id: Unique identifier of the worker being torn down.
|
||||
*args: Additional teardown arguments (currently unused).
|
||||
**kwargs: Additional teardown keyword arguments (currently unused).
|
||||
"""
|
||||
self.worker_id = None
|
||||
|
||||
self._tracer.teardown_worker(worker_id)
|
||||
|
||||
@property
|
||||
def tracer(self) -> Tracer:
|
||||
"""Get the tracer instance.
|
||||
|
||||
Returns:
|
||||
The Tracer instance used by this runner.
|
||||
"""
|
||||
return self._tracer
|
||||
|
||||
def get_agent(self) -> LitAgent[T_task]:
|
||||
"""Get the agent instance.
|
||||
|
||||
Returns:
|
||||
The LitAgent instance managed by this runner.
|
||||
|
||||
Raises:
|
||||
ValueError: If the agent has not been initialized via [`init`][agentlightning.LitAgentRunner.init].
|
||||
"""
|
||||
if self._agent is None:
|
||||
raise ValueError("Agent not initialized. Call init() first.")
|
||||
return self._agent
|
||||
|
||||
def get_store(self) -> LightningStore:
|
||||
"""Get the store instance.
|
||||
|
||||
Returns:
|
||||
The LightningStore instance for this worker.
|
||||
|
||||
Raises:
|
||||
ValueError: If the store has not been initialized via [`init_worker`][agentlightning.LitAgentRunner.init_worker].
|
||||
"""
|
||||
if self._store is None:
|
||||
raise ValueError("Store not initialized. Call init_worker() first.")
|
||||
return self._store
|
||||
|
||||
def get_worker_id(self) -> str:
|
||||
"""Get the formatted worker ID string.
|
||||
|
||||
Returns:
|
||||
A formatted string like "Worker-0" if initialized, or "Worker-Unknown"
|
||||
if the worker ID has not been set.
|
||||
"""
|
||||
return f"Worker-{self.worker_id}" if self.worker_id is not None else "Worker-Unknown"
|
||||
|
||||
def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
|
||||
"""Generate a standardized log prefix for the current worker.
|
||||
|
||||
This creates a consistent prefix format for log messages to identify
|
||||
which worker and rollout the message is associated with.
|
||||
|
||||
Args:
|
||||
rollout_id: Optional rollout ID to include in the prefix.
|
||||
|
||||
Returns:
|
||||
A formatted log prefix string like "[Worker 0 | Rollout xyz]",
|
||||
"[Worker 0]", "[Rollout xyz]", or "[Default Worker]".
|
||||
"""
|
||||
if self.worker_id is not None:
|
||||
if rollout_id:
|
||||
return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
|
||||
else:
|
||||
return f"[Worker {self.worker_id}]"
|
||||
if rollout_id:
|
||||
return f"[Rollout {rollout_id}]"
|
||||
return "[Default Worker]"
|
||||
|
||||
async def _trigger_hooks(
|
||||
self,
|
||||
hook_type: Literal["on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end"],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Trigger all registered hooks of a specific type.
|
||||
|
||||
This method calls the specified hook method on all registered hooks,
|
||||
catching and logging any exceptions that occur during hook execution
|
||||
to prevent them from disrupting the main execution flow.
|
||||
|
||||
Args:
|
||||
hook_type: The type of hook to trigger. Valid values are:
|
||||
"on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end".
|
||||
*args: Positional arguments to pass to the hook methods.
|
||||
**kwargs: Keyword arguments to pass to the hook methods.
|
||||
"""
|
||||
for hook in self._hooks:
|
||||
try:
|
||||
await getattr(hook, hook_type)(*args, **kwargs)
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.")
|
||||
|
||||
async def _post_process_rollout_result(
|
||||
self, rollout: AttemptedRollout, raw_result: RolloutRawResult
|
||||
) -> List[ReadableSpan] | List[Span]:
|
||||
"""Standardizes the agent's return value and report what's needed to report to the store.
|
||||
|
||||
Args:
|
||||
rollout: The rollout object for the current task.
|
||||
raw_result: The output from the agent's rollout method.
|
||||
|
||||
Returns:
|
||||
The spans that are assumed to be added to the store.
|
||||
This only serves as an estimation for logging purposes. For precise tracking, use the store directly.
|
||||
"""
|
||||
store = self.get_store()
|
||||
|
||||
trace_spans: list[Span] = []
|
||||
result_recognized: bool = False
|
||||
|
||||
# Case 0: result is None
|
||||
if raw_result is None:
|
||||
trace_spans = self._tracer.get_last_trace()
|
||||
result_recognized = True
|
||||
|
||||
# Case 1: result is a float (final reward)
|
||||
if isinstance(raw_result, (bool, int, float)):
|
||||
if isinstance(raw_result, (bool, int)):
|
||||
logger.warning(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Reward is not a number, got: {type(raw_result)}. "
|
||||
"Auto converting to float."
|
||||
)
|
||||
raw_result = float(raw_result)
|
||||
# Preserve the existing spans before another span is emitted
|
||||
trace_spans = list(self._tracer.get_last_trace())
|
||||
# This will NOT emit another span to the tracer
|
||||
reward_span_core_fields = emit_reward(raw_result, propagate=False)
|
||||
# We add it to the store manually
|
||||
sequence_id = await store.get_next_span_sequence_id(rollout.rollout_id, rollout.attempt.attempt_id)
|
||||
reward_span = Span.from_core_fields(
|
||||
reward_span_core_fields,
|
||||
rollout_id=rollout.rollout_id,
|
||||
attempt_id=rollout.attempt.attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
await store.add_span(reward_span)
|
||||
result_recognized = True
|
||||
|
||||
# Case 2-4: result is a list
|
||||
if isinstance(raw_result, list):
|
||||
# Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
|
||||
if isinstance(self._tracer, OtelTracer):
|
||||
logger.warning(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
|
||||
"The traces should have already been added to the store. "
|
||||
"Returning the traces from the rollout will result in duplicate spans."
|
||||
)
|
||||
for span in raw_result:
|
||||
added_span = await store.add_otel_span(
|
||||
rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
|
||||
)
|
||||
if added_span is not None:
|
||||
trace_spans.append(added_span)
|
||||
else:
|
||||
logger.error(
|
||||
f"{self._log_prefix(rollout.rollout_id)} Failed to add OpenTelemetry span to the store: {span}"
|
||||
)
|
||||
result_recognized = True
|
||||
|
||||
# Case 3: result is a list of Span (agentlightning spans)
|
||||
elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
|
||||
# Add the spans directly to the store
|
||||
for span in raw_result:
|
||||
await store.add_span(cast(Span, span))
|
||||
trace_spans = [cast(Span, span) for span in raw_result]
|
||||
result_recognized = True
|
||||
|
||||
# Case 4: result is a list of SpanCoreFields (agentlightning spans)
|
||||
elif len(raw_result) > 0 and all(isinstance(t, SpanCoreFields) for t in raw_result):
|
||||
# Add the spans directly to the store too, but needs to get sequence id first
|
||||
sequence_ids = await store.get_many_span_sequence_ids(
|
||||
[(rollout.rollout_id, rollout.attempt.attempt_id) for _ in range(len(raw_result))]
|
||||
)
|
||||
trace_spans = [
|
||||
Span.from_core_fields(
|
||||
cast(SpanCoreFields, span_core_fields),
|
||||
rollout_id=rollout.rollout_id,
|
||||
attempt_id=rollout.attempt.attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
for span_core_fields, sequence_id in zip(raw_result, sequence_ids, strict=True)
|
||||
]
|
||||
await store.add_many_spans(trace_spans)
|
||||
result_recognized = True
|
||||
|
||||
# Left over cases for list
|
||||
elif len(raw_result) == 0:
|
||||
logger.warning(
|
||||
f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
|
||||
"Please check your rollout implementation."
|
||||
)
|
||||
trace_spans = []
|
||||
result_recognized = True
|
||||
|
||||
else:
|
||||
types = [type(t).__name__ for t in raw_result][:10]
|
||||
raise ValueError(
|
||||
f"Invalid raw result type. It's expected to be a list of ReadableSpan or Span, "
|
||||
f"but got: {', '.join(types)}..."
|
||||
)
|
||||
|
||||
if not result_recognized:
|
||||
raise TypeError(
|
||||
f"Invalid raw result type. It's expected to be none, float, or a list of ReadableSpan or Span, "
|
||||
f"but got: {type(raw_result).__name__}..."
|
||||
)
|
||||
|
||||
return trace_spans
|
||||
|
||||
async def _emit_heartbeat(self, store: LightningStore) -> None:
|
||||
"""Send a heartbeat tick to the store.
|
||||
|
||||
Args:
|
||||
store: The lightning store to update.
|
||||
"""
|
||||
logger.debug(f"{self._log_prefix()} Preparing to emit heartbeat.")
|
||||
worker_id = self.get_worker_id()
|
||||
|
||||
try:
|
||||
snapshot = await asyncio.wait_for(
|
||||
asyncio.to_thread(system_snapshot, self._heartbeat_include_gpu),
|
||||
timeout=self._heartbeat_interval,
|
||||
)
|
||||
logger.debug(f"{self._log_prefix()} Heartbeat snapshot acquired.")
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"%s Heartbeat snapshot acquisition timed out after %.1fs, skipping.",
|
||||
self._log_prefix(),
|
||||
self._heartbeat_interval,
|
||||
)
|
||||
return
|
||||
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())
|
||||
|
||||
def _start_heartbeat_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
|
||||
"""Start a background heartbeat loop and return an async stopper."""
|
||||
|
||||
if self._heartbeat_interval <= 0:
|
||||
return None
|
||||
|
||||
if self.worker_id is None:
|
||||
logger.warning("%s Cannot start heartbeat loop without worker_id.", self._log_prefix())
|
||||
return None
|
||||
|
||||
if self._heartbeat_launch_mode == "asyncio":
|
||||
return self._start_heartbeat_asyncio_loop(store)
|
||||
if self._heartbeat_launch_mode == "thread":
|
||||
return self._start_heartbeat_thread_loop(store)
|
||||
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
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.
|
||||
|
||||
If an event is provided, the method will check it periodically (every 0.1s)
|
||||
and return early if the event is set.
|
||||
|
||||
Args:
|
||||
event: Optional [`ExecutionEvent`][agentlightning.ExecutionEvent] object that can be used to interrupt the sleep.
|
||||
If set during the sleep period, the method returns immediately.
|
||||
"""
|
||||
interval = self._poll_interval + self._random_state.uniform(-self._interval_jitter, self._interval_jitter)
|
||||
interval = max(interval, 0.01)
|
||||
if event is None:
|
||||
await asyncio.sleep(interval)
|
||||
return
|
||||
current_time = time.time()
|
||||
next_time = current_time + interval
|
||||
while time.time() < next_time:
|
||||
await asyncio.sleep(0.1)
|
||||
if event.is_set():
|
||||
return
|
||||
|
||||
async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> str:
|
||||
"""Execute a single rollout implementation.
|
||||
|
||||
This is the core method that handles the execution of a single rollout,
|
||||
including resource fetching, hook triggering, agent invocation, tracing,
|
||||
and result processing.
|
||||
|
||||
Args:
|
||||
next_rollout: The rollout to execute, containing input data, mode,
|
||||
and resources information.
|
||||
raise_on_exception: If True, exceptions during rollout execution will
|
||||
be re-raised. If False, exceptions are logged but not propagated.
|
||||
"""
|
||||
store = self.get_store()
|
||||
agent = self.get_agent()
|
||||
|
||||
rollout_id = next_rollout.rollout_id
|
||||
|
||||
resources_id = next_rollout.resources_id
|
||||
resources_update = None
|
||||
if resources_id:
|
||||
resources_update = await store.get_resources_by_id(resources_id)
|
||||
else:
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
|
||||
resources_update = await store.get_latest_resources()
|
||||
if not resources_update:
|
||||
if raise_on_exception:
|
||||
raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources")
|
||||
else:
|
||||
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
|
||||
|
||||
try:
|
||||
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
|
||||
)
|
||||
|
||||
# NOTE: This is the most costly step in the whole function
|
||||
# If the rollout method becomes unresponsive or timeouts, there is nothing we can do within the runner.
|
||||
# We might need some mechanisms in execution strategy to restart the runner. But that's a future work.
|
||||
if agent.is_async():
|
||||
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)
|
||||
|
||||
end_time = time.time()
|
||||
logger.info(
|
||||
f"{self._log_prefix(rollout_id)} Completed in "
|
||||
f"{end_time - start_time:.2f}s. Collected {len(trace_spans)} span(s). "
|
||||
f"Final reward: {last_reward}"
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
|
||||
has_exception = True
|
||||
|
||||
if raise_on_exception:
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
await self._trigger_hooks(
|
||||
hook_type="on_rollout_end", agent=agent, runner=self, rollout=next_rollout, spans=trace_spans
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
|
||||
|
||||
try:
|
||||
if has_exception:
|
||||
# possibly timed out and cancelled?
|
||||
await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="failed")
|
||||
else:
|
||||
await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="succeeded")
|
||||
except Exception:
|
||||
logger.exception(
|
||||
f"{self._log_prefix(rollout_id)} Exception during update_attempt. Giving up the update."
|
||||
)
|
||||
|
||||
return rollout_id
|
||||
|
||||
async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
|
||||
"""Run the runner, continuously iterating over tasks in the store.
|
||||
|
||||
This method polls the store for new rollouts and executes them until:
|
||||
|
||||
- The event is set (if provided)
|
||||
- The max_rollouts limit is reached (if configured)
|
||||
- No more tasks are available
|
||||
|
||||
All exceptions during rollout execution are caught and logged but not
|
||||
propagated, allowing the runner to continue processing subsequent tasks.
|
||||
|
||||
Args:
|
||||
event: Optional ExecutionEvent object to signal the runner to stop. The runner
|
||||
will check this event periodically and stop gracefully when set.
|
||||
"""
|
||||
num_tasks_processed = 0
|
||||
logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
|
||||
store = self.get_store()
|
||||
|
||||
stop_heartbeat = self._start_heartbeat_loop(store)
|
||||
|
||||
try:
|
||||
while not (event is not None and event.is_set()) and (
|
||||
self._max_rollouts is None or num_tasks_processed < self._max_rollouts
|
||||
):
|
||||
# Retrieve the next rollout
|
||||
next_rollout: Optional[Rollout] = None
|
||||
while not (event is not None and event.is_set()):
|
||||
logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
|
||||
next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
|
||||
logger.debug(f"{self._log_prefix()} Next rollout retrieved: {next_rollout}")
|
||||
if next_rollout is None:
|
||||
logger.debug(
|
||||
f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."
|
||||
)
|
||||
await self._sleep_until_next_poll(event)
|
||||
else:
|
||||
break
|
||||
|
||||
if next_rollout is None:
|
||||
return
|
||||
|
||||
# Execute the step
|
||||
await self._step_impl(next_rollout)
|
||||
|
||||
num_tasks_processed += 1
|
||||
if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
|
||||
logger.info(
|
||||
f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}"
|
||||
)
|
||||
finally:
|
||||
if stop_heartbeat is not None:
|
||||
await stop_heartbeat()
|
||||
|
||||
logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
|
||||
|
||||
async def step(
|
||||
self,
|
||||
input: T_task,
|
||||
*,
|
||||
resources: Optional[NamedResources] = None,
|
||||
mode: Optional[RolloutMode] = None,
|
||||
event: Optional[ExecutionEvent] = None,
|
||||
) -> Rollout:
|
||||
"""Execute a single task directly, bypassing the task queue.
|
||||
|
||||
This method creates a new rollout for the given input and executes it
|
||||
immediately. Unlike [`iter()`][agentlightning.LitAgentRunner.iter],
|
||||
exceptions are propagated to the caller.
|
||||
|
||||
Args:
|
||||
input: The task input to be processed by the agent.
|
||||
resources: Optional named resources to be used for this specific task.
|
||||
If provided, a new resources entry will be created in the store.
|
||||
If not provided, the latest resources from the store will be used.
|
||||
mode: Optional rollout mode ("train" or "validation"). If not provided,
|
||||
the agent's default mode will be used.
|
||||
event: Optional ExecutionEvent object to signal interruption (currently unused
|
||||
but included for interface consistency).
|
||||
|
||||
Returns:
|
||||
The completed rollout.
|
||||
|
||||
Raises:
|
||||
Exception: Any exception that occurs during rollout execution will be
|
||||
re-raised to the caller.
|
||||
"""
|
||||
store = self.get_store()
|
||||
|
||||
if resources is not None:
|
||||
resources_update = await store.add_resources(resources)
|
||||
resources_id = resources_update.resources_id
|
||||
else:
|
||||
resources_id = None
|
||||
|
||||
attempted_rollout = await self.get_store().start_rollout(
|
||||
input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
|
||||
)
|
||||
rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
|
||||
|
||||
completed_rollout = await store.get_rollout_by_id(rollout_id)
|
||||
if completed_rollout is None:
|
||||
raise RuntimeError(f"{self._log_prefix()} Failed to fetch completed rollout by id after step: {rollout_id}")
|
||||
return completed_rollout
|
||||
@@ -0,0 +1,182 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Abstract runner interface for executing agent tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from typing import TYPE_CHECKING, Any, Generic, Iterator, Optional, Sequence, TypeVar
|
||||
|
||||
from agentlightning.execution.events import ExecutionEvent
|
||||
from agentlightning.litagent import LitAgent
|
||||
from agentlightning.store.base import LightningStore
|
||||
from agentlightning.types import Hook, NamedResources, ParallelWorkerBase, Rollout, RolloutMode
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agentlightning.execution.events import ExecutionEvent
|
||||
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Runner(ParallelWorkerBase, Generic[T_task]):
|
||||
"""Abstract base class for long-running agent executors.
|
||||
|
||||
Runner implementations coordinate [`LitAgent`][agentlightning.LitAgent]
|
||||
instances, acquire work from a [`LightningStore`][agentlightning.LightningStore],
|
||||
and emit [`Rollout`][agentlightning.Rollout] objects. Subclasses decide how
|
||||
to schedule work (polling, streaming, etc.) while this base class provides a
|
||||
minimal lifecycle contract.
|
||||
"""
|
||||
|
||||
def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
|
||||
"""Prepare the runner to execute tasks for `agent`.
|
||||
|
||||
This method is called only once during the setup for all workers, not for each worker.
|
||||
|
||||
Args:
|
||||
agent: Agent instance providing task-specific logic.
|
||||
**kwargs: Optional runner-specific configuration.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must supply the initialization
|
||||
routine.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
|
||||
"""Configure worker-local state before processing tasks.
|
||||
|
||||
This method is called for **each** worker during the setup.
|
||||
|
||||
Args:
|
||||
worker_id: Unique identifier for this worker process or thread.
|
||||
store: Shared [`LightningStore`][agentlightning.LightningStore]
|
||||
backing task coordination.
|
||||
**kwargs: Optional worker-specific configuration.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must prepare per-worker resources.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def run(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Deprecated synchronous entry point.
|
||||
|
||||
Use [`iter()`][agentlightning.Runner.iter] or [`step()`][agentlightning.Runner.step] instead.
|
||||
|
||||
Raises:
|
||||
RuntimeError: Always raised to direct callers to
|
||||
[iter()][agentlightning.Runner.iter] or
|
||||
[step()][agentlightning.Runner.step].
|
||||
"""
|
||||
raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
|
||||
|
||||
def teardown(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Release resources acquired during [`init()`][agentlightning.Runner.init].
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must implement the shutdown routine.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
|
||||
"""Release per-worker resources allocated by [`init_worker()`][agentlightning.Runner.init_worker].
|
||||
|
||||
Args:
|
||||
worker_id: Identifier of the worker being torn down.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must implement the shutdown routine.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@contextmanager
|
||||
def run_context(
|
||||
self,
|
||||
*,
|
||||
agent: LitAgent[T_task],
|
||||
store: LightningStore,
|
||||
hooks: Optional[Sequence[Hook]] = None,
|
||||
worker_id: Optional[int] = None,
|
||||
) -> Iterator[Runner[T_task]]:
|
||||
"""Initialize and tear down a runner within a simple context manager.
|
||||
|
||||
The helper is primarily intended for debugging runner implementations
|
||||
outside of a full [`Trainer`][agentlightning.Trainer] stack.
|
||||
|
||||
Args:
|
||||
agent: Agent executed by this runner.
|
||||
store: Backing [`LightningStore`][agentlightning.LightningStore].
|
||||
If you don't have one, you can easily create one with
|
||||
[`InMemoryLightningStore`][agentlightning.InMemoryLightningStore].
|
||||
hooks: Optional sequence of hooks recognised by the runner.
|
||||
Not all runners support hooks.
|
||||
worker_id: Override the worker identifier used during setup. Defaults
|
||||
to `0`.
|
||||
"""
|
||||
_initialized: bool = False
|
||||
_worker_initialized: bool = False
|
||||
try:
|
||||
self.init(agent=agent, hooks=hooks)
|
||||
_initialized = True
|
||||
self.init_worker(worker_id=0, store=store)
|
||||
_worker_initialized = True
|
||||
yield self
|
||||
finally:
|
||||
try:
|
||||
if _worker_initialized:
|
||||
self.teardown_worker(worker_id=worker_id if worker_id is not None else 0)
|
||||
except Exception:
|
||||
logger.error("Error during runner worker teardown", exc_info=True)
|
||||
|
||||
try:
|
||||
if _initialized:
|
||||
self.teardown()
|
||||
except Exception:
|
||||
logger.error("Error during runner teardown", exc_info=True)
|
||||
|
||||
async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
|
||||
"""Run the runner, continuously iterating over tasks in the store.
|
||||
|
||||
This method runs in a loop, polling the store for new tasks and executing
|
||||
them until interrupted by the event or when no more tasks are available.
|
||||
|
||||
Args:
|
||||
event: Cooperative stop signal. When set, the runner should complete
|
||||
the current unit of work and exit the loop.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses provide the iteration behavior.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def step(
|
||||
self,
|
||||
input: T_task,
|
||||
*,
|
||||
resources: Optional[NamedResources] = None,
|
||||
mode: Optional[RolloutMode] = None,
|
||||
event: Optional[ExecutionEvent] = None,
|
||||
) -> Rollout:
|
||||
"""Execute a single task with the given input.
|
||||
|
||||
This method provides fine-grained control for executing individual tasks
|
||||
directly, bypassing the store's task queue.
|
||||
|
||||
Args:
|
||||
input: Task payload consumed by the agent.
|
||||
resources: Optional named resources scoped to this invocation.
|
||||
mode: Optional rollout mode such as `"train"` or `"eval"`.
|
||||
event: Cooperative stop signal for long-running tasks.
|
||||
|
||||
Returns:
|
||||
Completed rollout produced by the agent.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses provide the execution behavior.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,309 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
from agentlightning.adapter import TracerTraceToTriplet
|
||||
from agentlightning.client import AgentLightningClient
|
||||
from agentlightning.litagent import LitAgent
|
||||
from agentlightning.litagent.litagent import is_v0_1_rollout_api
|
||||
from agentlightning.tracer.base import Tracer
|
||||
from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Span, SpanLike, Triplet
|
||||
|
||||
from .base import Runner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"LegacyAgentRunner",
|
||||
]
|
||||
|
||||
|
||||
class LegacyAgentRunner(Runner[Any]):
|
||||
"""Manages the agent's execution loop and integrates with AgentOps.
|
||||
|
||||
This class orchestrates the interaction between the agent (`LitAgent`) and
|
||||
the server (`AgentLightningClient`). It handles polling for tasks, executing
|
||||
the agent's logic, and reporting results back to the server. If enabled,
|
||||
it will also automatically trace each rollout using AgentOps.
|
||||
|
||||
Attributes:
|
||||
agent: The `LitAgent` instance containing the agent's logic.
|
||||
client: The `AgentLightningClient` for server communication.
|
||||
tracer: The tracer instance for this runner/worker.
|
||||
worker_id: An optional identifier for the worker process.
|
||||
max_tasks: The maximum number of tasks to process before stopping.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: LitAgent[Any],
|
||||
client: AgentLightningClient,
|
||||
tracer: Tracer,
|
||||
triplet_exporter: TracerTraceToTriplet,
|
||||
worker_id: Optional[int] = None,
|
||||
max_tasks: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.agent = agent
|
||||
self.client = client
|
||||
self.tracer = tracer
|
||||
self.triplet_exporter = triplet_exporter
|
||||
|
||||
# Worker-specific attributes
|
||||
self.worker_id = worker_id
|
||||
self.max_tasks = max_tasks
|
||||
|
||||
# These methods are overridden by Runner, getting them back to old behavior.
|
||||
def init(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
|
||||
self.worker_id = worker_id
|
||||
|
||||
def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def teardown(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
|
||||
"""Generates a standardized log prefix for the current worker."""
|
||||
if self.worker_id is not None:
|
||||
if rollout_id:
|
||||
return f"[Worker {self.worker_id} | RolloutLegacy {rollout_id}]"
|
||||
else:
|
||||
return f"[Worker {self.worker_id}]"
|
||||
if rollout_id:
|
||||
return f"[RolloutLegacy {rollout_id}]"
|
||||
return "[Default Worker]"
|
||||
|
||||
def _to_rollout_object(
|
||||
self,
|
||||
result: RolloutRawResultLegacy,
|
||||
rollout_id: str,
|
||||
) -> RolloutLegacy:
|
||||
"""Standardizes the agent's return value into a RolloutLegacy object.
|
||||
|
||||
Args:
|
||||
result: The output from the agent's rollout method.
|
||||
rollout_id: The unique identifier for the current task.
|
||||
|
||||
Returns:
|
||||
A standardized `RolloutLegacy` object for reporting to the server.
|
||||
"""
|
||||
trace: Any = None
|
||||
final_reward: Optional[float] = None
|
||||
triplets: Optional[List[Triplet]] = None
|
||||
trace_spans: Optional[List[SpanLike]] = None
|
||||
|
||||
# Handle different types of results from the agent
|
||||
# Case 1: result is a float (final reward)
|
||||
if isinstance(result, float):
|
||||
final_reward = result
|
||||
# Case 2: result is a list of Triplets
|
||||
if isinstance(result, list) and all(isinstance(t, Triplet) for t in result):
|
||||
triplets = result # type: ignore
|
||||
# Case 3.1: result is a list of ReadableSpan (OpenTelemetry spans)
|
||||
if isinstance(result, list) and all(isinstance(t, (ReadableSpan)) for t in result):
|
||||
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
|
||||
# Case 5: result is a RolloutLegacy object
|
||||
if isinstance(result, RolloutLegacy):
|
||||
final_reward = result.final_reward
|
||||
triplets = result.triplets
|
||||
trace = result.trace
|
||||
|
||||
# If the agent has tracing enabled, use the tracer's last trace if not already set
|
||||
if self.tracer and (trace is None or trace_spans is None):
|
||||
trace_spans = self.tracer.get_last_trace() # type: ignore
|
||||
if trace_spans:
|
||||
trace = [cast(Span, span).model_dump() for span in trace_spans]
|
||||
|
||||
# Always extract triplets from the trace using TracerTraceToTriplet
|
||||
if trace_spans:
|
||||
triplets = self.triplet_exporter(trace_spans) # type: ignore
|
||||
|
||||
# If the agent has triplets, use the last one for final reward if not set
|
||||
if triplets and triplets[-1].reward is not None and final_reward is None:
|
||||
final_reward = triplets[-1].reward
|
||||
|
||||
# Create the RolloutLegacy object with standardized fields
|
||||
result_dict: Dict[str, Any] = {
|
||||
"rollout_id": rollout_id,
|
||||
}
|
||||
if final_reward is not None:
|
||||
result_dict["final_reward"] = final_reward
|
||||
if triplets is not None:
|
||||
result_dict["triplets"] = triplets
|
||||
if trace is not None:
|
||||
result_dict["trace"] = trace
|
||||
|
||||
if isinstance(result, RolloutLegacy):
|
||||
return result.model_copy(update=result_dict)
|
||||
return RolloutLegacy(**result_dict)
|
||||
|
||||
def run(self) -> bool: # type: ignore
|
||||
"""Poll the task and rollout once synchronously."""
|
||||
self.agent.set_runner(self) # Ensure the agent has a reference to this runner
|
||||
|
||||
task = self.client.poll_next_task()
|
||||
if task is None:
|
||||
logger.info(f"{self._log_prefix()} Poll returned no task. Exiting.")
|
||||
return False
|
||||
rollout_id = task.rollout_id
|
||||
|
||||
resources_id = task.resources_id
|
||||
resources_update = None
|
||||
if resources_id:
|
||||
resources_update = self.client.get_resources_by_id(resources_id)
|
||||
else:
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
|
||||
resources_update = self.client.get_latest_resources()
|
||||
if not resources_update:
|
||||
logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
|
||||
return False
|
||||
|
||||
rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task) # Default empty rollout
|
||||
|
||||
try:
|
||||
try:
|
||||
self.agent.on_rollout_start(task, self, self.tracer)
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
|
||||
|
||||
with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"): # pyright: ignore[reportPrivateUsage]
|
||||
start_time = time.time()
|
||||
rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
|
||||
# Pass the task input, not the whole task object
|
||||
if is_v0_1_rollout_api(rollout_method):
|
||||
result = cast(
|
||||
RolloutRawResultLegacy,
|
||||
rollout_method(
|
||||
task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources # type: ignore
|
||||
),
|
||||
) # type: ignore
|
||||
else:
|
||||
result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj) # type: ignore
|
||||
rollout_obj = self._to_rollout_object(result, task.rollout_id) # type: ignore
|
||||
end_time = time.time()
|
||||
logger.info(
|
||||
f"{self._log_prefix(rollout_id)} Completed in "
|
||||
f"{end_time - start_time:.2f}s. Triplet length: "
|
||||
f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
|
||||
f"Reward: {rollout_obj.final_reward}"
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
|
||||
finally:
|
||||
try:
|
||||
self.agent.on_rollout_end(task, rollout_obj, self, self.tracer) # type: ignore
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
|
||||
self.client.post_rollout(rollout_obj)
|
||||
|
||||
return True
|
||||
|
||||
def iter(self) -> int: # type: ignore
|
||||
"""Executes the synchronous polling and rollout loop."""
|
||||
num_tasks_processed = 0
|
||||
logger.info(f"{self._log_prefix()} Started sync rollouts (max: {self.max_tasks or 'unlimited'}).")
|
||||
|
||||
while self.max_tasks is None or num_tasks_processed < self.max_tasks:
|
||||
if self.run():
|
||||
num_tasks_processed += 1
|
||||
|
||||
if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
|
||||
logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self.max_tasks or 'unlimited'}")
|
||||
|
||||
logger.info(f"{self._log_prefix()} Finished sync rollouts. Processed {num_tasks_processed} tasks.")
|
||||
return num_tasks_processed
|
||||
|
||||
async def run_async(self) -> bool:
|
||||
"""Poll the task and rollout once."""
|
||||
self.agent.set_runner(self) # Ensure the agent has a reference to this runner
|
||||
|
||||
task = await self.client.poll_next_task_async()
|
||||
if task is None:
|
||||
logger.info(f"{self._log_prefix()} Poll returned no task. Exiting.")
|
||||
return False
|
||||
rollout_id = task.rollout_id
|
||||
|
||||
resources_id = task.resources_id
|
||||
resources_update = None
|
||||
if resources_id:
|
||||
resources_update = await self.client.get_resources_by_id_async(resources_id)
|
||||
else:
|
||||
logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
|
||||
resources_update = await self.client.get_latest_resources_async()
|
||||
if not resources_update:
|
||||
logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
|
||||
return False
|
||||
|
||||
rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task) # Default empty rollout
|
||||
|
||||
try:
|
||||
try:
|
||||
self.agent.on_rollout_start(task, self, self.tracer)
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
|
||||
|
||||
async with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
|
||||
start_time = time.time()
|
||||
rollout_method = (
|
||||
self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
|
||||
)
|
||||
# Pass the task input, not the whole task object
|
||||
if is_v0_1_rollout_api(rollout_method):
|
||||
result = cast(
|
||||
RolloutRawResultLegacy,
|
||||
await rollout_method(
|
||||
task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources # type: ignore
|
||||
),
|
||||
) # type: ignore
|
||||
else:
|
||||
result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj) # type: ignore
|
||||
rollout_obj = self._to_rollout_object(result, task.rollout_id) # type: ignore
|
||||
end_time = time.time()
|
||||
logger.info(
|
||||
f"{self._log_prefix(rollout_id)} Completed in "
|
||||
f"{end_time - start_time:.2f}s. Triplet length: "
|
||||
f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
|
||||
f"Reward: {rollout_obj.final_reward}"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
|
||||
finally:
|
||||
try:
|
||||
self.agent.on_rollout_end(task, rollout_obj, self, self.tracer) # type: ignore
|
||||
except Exception:
|
||||
logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
|
||||
await self.client.post_rollout_async(rollout_obj)
|
||||
|
||||
return True
|
||||
|
||||
async def iter_async(self) -> int:
|
||||
"""Executes the asynchronous polling and rollout loop."""
|
||||
num_tasks_processed = 0
|
||||
logger.info(f"{self._log_prefix()} Started async rollouts (max: {self.max_tasks or 'unlimited'}).")
|
||||
|
||||
while self.max_tasks is None or num_tasks_processed < self.max_tasks:
|
||||
if await self.run_async():
|
||||
num_tasks_processed += 1
|
||||
|
||||
if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
|
||||
logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self.max_tasks or 'unlimited'}")
|
||||
logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
|
||||
return num_tasks_processed
|
||||
@@ -0,0 +1,164 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Semantic conventions for Agent-lightning spans.
|
||||
|
||||
Conventions in this file are added on demand. We generally DO NOT add
|
||||
new semantic conventions unless it's absolutely needed for certain algorithms or scenarios.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
AGL_ANNOTATION = "agentlightning.annotation"
|
||||
"""Agent-lightning's standard span name for annotations.
|
||||
|
||||
Annotations are minimal span units for rewards, tags, and metadatas.
|
||||
They are used to "annotate" a specific event or a part of rollout.
|
||||
"""
|
||||
|
||||
AGL_MESSAGE = "agentlightning.message"
|
||||
"""Agent-lightning's standard span name for messages and logs."""
|
||||
|
||||
AGL_OBJECT = "agentlightning.object"
|
||||
"""Agent-lightning's standard span name for customized objects."""
|
||||
|
||||
AGL_EXCEPTION = "agentlightning.exception"
|
||||
"""Agent-lightning's standard span name for exceptions.
|
||||
|
||||
Used by the exception emitter to record exception details.
|
||||
"""
|
||||
|
||||
AGL_OPERATION = "agentlightning.operation"
|
||||
"""Agent-lightning's standard span name for functions.
|
||||
Wrap function or code-blocks as operations.
|
||||
"""
|
||||
|
||||
AGL_REWARD = "agentlightning.reward"
|
||||
"""Agent-lightning's standard span name for reward operations."""
|
||||
|
||||
AGL_VIRTUAL = "agentlightning.virtual"
|
||||
"""Agent-lightning's standard span name for virtual operations.
|
||||
|
||||
Mostly used in adapter when needing to represent the root or intermediate operations.
|
||||
"""
|
||||
|
||||
|
||||
class LightningResourceAttributes(Enum):
|
||||
"""Resource attribute names used in Agent-lightning spans."""
|
||||
|
||||
ROLLOUT_ID = "agentlightning.rollout_id"
|
||||
"""Resource name for rollout ID in Agent-lightning spans."""
|
||||
|
||||
ATTEMPT_ID = "agentlightning.attempt_id"
|
||||
"""Resource name for attempt ID in Agent-lightning spans."""
|
||||
|
||||
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.
|
||||
|
||||
Exception types can't be found here because they are defined in OpenTelemetry's official semantic conventions.
|
||||
"""
|
||||
|
||||
REWARD = "agentlightning.reward"
|
||||
"""Attribute prefix for rewards-related data in reward spans.
|
||||
|
||||
It should be used as a prefix. For example, "agentlightning.reward.0.value" can
|
||||
be used to track a specific metric. See [RewardAttributes][agentlightning.semconv.RewardAttributes].
|
||||
"""
|
||||
|
||||
LINK = "agentlightning.link"
|
||||
"""Attribute name for linking the current span to another span or other objects like requests/responses."""
|
||||
|
||||
TAG = "agentlightning.tag"
|
||||
"""Attribute name for tagging spans with customized strings."""
|
||||
|
||||
MESSAGE_BODY = "agentlightning.message.body"
|
||||
"""Attribute name for message text in message spans."""
|
||||
|
||||
OBJECT_TYPE = "agentlightning.object.type"
|
||||
"""Attribute name for object type (full qualified name) in object spans.
|
||||
|
||||
I think builtin types like str, int, bool, list, dict are self-explanatory and
|
||||
should also be qualified to use here.
|
||||
"""
|
||||
|
||||
OBJECT_LITERAL = "agentlightning.object.literal"
|
||||
"""Attribute name for object literal value in object spans (for str, int, bool, ...)."""
|
||||
|
||||
OBJECT_JSON = "agentlightning.object.json"
|
||||
"""Attribute name for object serialized value (JSON) in object spans."""
|
||||
|
||||
OPERATION_NAME = "agentlightning.operation.name"
|
||||
"""Attribute name for operation name in operation spans, normally the function name."""
|
||||
|
||||
OPERATION_INPUT = "agentlightning.operation.input"
|
||||
"""Attribute name for operation input in operation spans."""
|
||||
|
||||
OPERATION_OUTPUT = "agentlightning.operation.output"
|
||||
"""Attribute name for operation output in operation spans."""
|
||||
|
||||
|
||||
class RewardAttributes(Enum):
|
||||
"""Multi-dimensional reward attributes will look like:
|
||||
|
||||
```json
|
||||
{"agentlightning.reward.0.name": "efficiency", "agentlightning.reward.0.value": 0.75}
|
||||
```
|
||||
|
||||
The first reward in the reward list will automatically be the primary reward.
|
||||
If the reward list has greater than 1, it shall be a multi-dimensional case.
|
||||
"""
|
||||
|
||||
REWARD_NAME = "name"
|
||||
"""Key for each dimension in multi-dimensional reward spans."""
|
||||
|
||||
REWARD_VALUE = "value"
|
||||
"""Value for each dimension in multi-dimensional reward spans."""
|
||||
|
||||
|
||||
class RewardPydanticModel(BaseModel):
|
||||
"""A stricter implementation of RewardAttributes used in otel helpers."""
|
||||
|
||||
name: str
|
||||
"""Name of the reward dimension."""
|
||||
|
||||
value: float
|
||||
"""Value of the reward dimension."""
|
||||
|
||||
|
||||
class LinkAttributes(Enum):
|
||||
"""Standard link types used in Agent-lightning spans.
|
||||
|
||||
The link is more powerful than [OpenTelemetry link](https://opentelemetry.io/docs/specs/otel/trace/api/#link)
|
||||
in that it supports linking to a queryset of spans.
|
||||
It can even link to span object that hasn't been emitted yet.
|
||||
"""
|
||||
|
||||
KEY_MATCH = "key_match"
|
||||
"""Linking to spans with matching attribute keys.
|
||||
|
||||
`trace_id` and `span_id` are reserved and will be used to link to specific spans directly.
|
||||
|
||||
For example, it can be `gen_ai.response.id` if intended to be link to a chat completion response span.
|
||||
Or it can be `span_id` to link to a specific span by its ID.
|
||||
"""
|
||||
|
||||
VALUE_MATCH = "value_match"
|
||||
"""Linking to spans with corresponding attribute values on those keys."""
|
||||
|
||||
|
||||
class LinkPydanticModel(BaseModel):
|
||||
"""A stricter implementation of LinkAttributes used in otel helpers."""
|
||||
|
||||
key_match: str
|
||||
"""The attribute key to match on the target spans."""
|
||||
|
||||
value_match: str
|
||||
"""The attribute value to match on the target spans."""
|
||||
@@ -0,0 +1,401 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Legacy HTTP server compatible with the original Agent Lightning protocol.
|
||||
|
||||
The implementation in this module predates the modern store-powered runtime and
|
||||
is kept for backwards compatibility with older deployments. New applications
|
||||
should migrate to the store architecture where possible.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import warnings
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, Dict, List, Literal, Optional
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException, Path
|
||||
|
||||
from .types import (
|
||||
GenericResponse,
|
||||
NamedResources,
|
||||
ResourcesUpdate,
|
||||
RolloutLegacy,
|
||||
Task,
|
||||
TaskIfAny,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ServerDataStore:
|
||||
"""Async-safe container for in-memory server state.
|
||||
|
||||
The store tracks queued tasks, claimed tasks, uploaded rollouts, and the
|
||||
currently published resources. All interactions are guarded by asyncio locks
|
||||
so that the FastAPI handlers can safely run in parallel.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
[`ServerDataStore`][agentlightning.server.ServerDataStore] is part of
|
||||
the legacy client/server stack. Use [`LightningStore`][agentlightning.LightningStore] instead.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._task_queue: asyncio.Queue[Task] = asyncio.Queue()
|
||||
self._processing_tasks: Dict[str, Task] = {} # Currently processing tasks
|
||||
self._completed_rollouts: Dict[str, RolloutLegacy] = {}
|
||||
|
||||
# Store for versioned resources
|
||||
self._resource_versions: Dict[str, NamedResources] = {}
|
||||
self._latest_resources_id: Optional[str] = None
|
||||
|
||||
# Locks for thread-safe access
|
||||
self._results_lock = asyncio.Lock()
|
||||
self._resources_lock = asyncio.Lock()
|
||||
|
||||
async def add_task(
|
||||
self,
|
||||
sample: Any,
|
||||
mode: Literal["train", "val", "test"] | None = None,
|
||||
resources_id: str | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Enqueue a new task and return the generated rollout identifier.
|
||||
|
||||
Args:
|
||||
sample: Payload that describes the task input.
|
||||
mode: Phase in which the sample should be executed (`"train"`, `"val"`, or
|
||||
`"test"`).
|
||||
resources_id: Identifier of a resource bundle that the executor should
|
||||
load before running the task.
|
||||
metadata: Optional metadata forwarded to the executor.
|
||||
|
||||
Returns:
|
||||
Unique rollout identifier assigned to the task.
|
||||
"""
|
||||
rollout_id = f"rollout-{uuid.uuid4()}"
|
||||
task = Task(
|
||||
rollout_id=rollout_id,
|
||||
input=sample,
|
||||
mode=mode,
|
||||
resources_id=resources_id,
|
||||
create_time=time.time(),
|
||||
num_claims=0,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
await self._task_queue.put(task)
|
||||
logger.info(f"Task queued: {rollout_id} (mode: {mode}, resources_id: {resources_id})")
|
||||
return rollout_id
|
||||
|
||||
async def get_next_task(self) -> Optional[Task]:
|
||||
"""Retrieve the next task from the queue without blocking.
|
||||
|
||||
Returns:
|
||||
Next [`Task`][agentlightning.Task] ready to execute, or ``None``
|
||||
when the queue is empty.
|
||||
"""
|
||||
try:
|
||||
async with self._results_lock:
|
||||
task = self._task_queue.get_nowait()
|
||||
task = task.model_copy(
|
||||
update={
|
||||
"last_claim_time": time.time(),
|
||||
"num_claims": (task.num_claims or 0) + 1,
|
||||
}
|
||||
)
|
||||
self._processing_tasks[task.rollout_id] = task
|
||||
if task.num_claims == 1:
|
||||
logger.debug(f"Next task retrieved: {task.rollout_id}")
|
||||
else:
|
||||
logger.info(f"Task {task.rollout_id} re-claimed (attempt {task.num_claims})")
|
||||
return task
|
||||
except asyncio.QueueEmpty:
|
||||
return None
|
||||
|
||||
async def update_resources(self, update: ResourcesUpdate):
|
||||
"""Persist a new resource bundle and mark it as the latest version.
|
||||
|
||||
Args:
|
||||
update: Resource payload received from a client.
|
||||
"""
|
||||
# TODO: evict old resources if necessary.
|
||||
async with self._resources_lock:
|
||||
self._resource_versions[update.resources_id] = update.resources
|
||||
self._latest_resources_id = update.resources_id
|
||||
logger.info(f"Resources updated. New version '{update.resources_id}' is now latest.")
|
||||
|
||||
async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
|
||||
"""Retrieve a specific resource bundle by identifier.
|
||||
|
||||
Args:
|
||||
resources_id: Identifier that was previously published to the store.
|
||||
|
||||
Returns:
|
||||
Matching [`ResourcesUpdate`][agentlightning.ResourcesUpdate]
|
||||
instance, or ``None`` when the identifier is unknown.
|
||||
"""
|
||||
async with self._resources_lock:
|
||||
resources = self._resource_versions.get(resources_id)
|
||||
if resources:
|
||||
return ResourcesUpdate(
|
||||
resources_id=resources_id,
|
||||
resources=resources,
|
||||
create_time=time.time(),
|
||||
update_time=time.time(),
|
||||
version=1,
|
||||
)
|
||||
return None
|
||||
|
||||
async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
|
||||
"""Return the most recent resource bundle, if one exists."""
|
||||
if self._latest_resources_id:
|
||||
return await self.get_resources_by_id(self._latest_resources_id)
|
||||
return None
|
||||
|
||||
async def store_rollout(self, rollout: RolloutLegacy):
|
||||
"""Persist a completed rollout for later inspection.
|
||||
|
||||
Args:
|
||||
rollout: Rollout returned by a client.
|
||||
"""
|
||||
async with self._results_lock:
|
||||
self._processing_tasks.pop(rollout.rollout_id, None)
|
||||
self._completed_rollouts[rollout.rollout_id] = rollout
|
||||
logger.info(f"Rollout received and stored: {rollout.rollout_id}")
|
||||
|
||||
async def retrieve_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
|
||||
"""Retrieve and remove a stored rollout by identifier.
|
||||
|
||||
Args:
|
||||
rollout_id: Identifier of the rollout to fetch.
|
||||
|
||||
Returns:
|
||||
Stored [`RolloutLegacy`][agentlightning.RolloutLegacy], or ``None``
|
||||
when the identifier is unknown.
|
||||
"""
|
||||
async with self._results_lock:
|
||||
return self._completed_rollouts.pop(rollout_id, None)
|
||||
|
||||
async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
|
||||
"""Return all completed rollouts and clear the internal buffer."""
|
||||
async with self._results_lock:
|
||||
rollouts = list(self._completed_rollouts.values())
|
||||
self._completed_rollouts.clear()
|
||||
return rollouts
|
||||
|
||||
def get_processing_tasks(self) -> Dict[str, Task]:
|
||||
"""Return a copy of currently processing tasks for timeout checking."""
|
||||
return self._processing_tasks.copy()
|
||||
|
||||
async def requeue_task(self, task: Task):
|
||||
"""Requeue a task that timed out while being processed."""
|
||||
logger.warning(f"Requeuing task {task.rollout_id} after timeout (attempt {task.num_claims})")
|
||||
async with self._results_lock:
|
||||
# Remove from processing tasks
|
||||
self._processing_tasks.pop(task.rollout_id, None)
|
||||
self._task_queue.put_nowait(task)
|
||||
|
||||
|
||||
class AgentLightningServer:
|
||||
"""High-level controller for the legacy Agent Lightning FastAPI server.
|
||||
|
||||
The controller orchestrates server start-up, task queueing, resource updates,
|
||||
and retrieval of client rollouts. It is primarily used by existing systems that
|
||||
still rely on the HTTP-based workflow.
|
||||
|
||||
!!! warning "Deprecated"
|
||||
[`AgentLightningServer`][agentlightning.server.AgentLightningServer] is part of
|
||||
the legacy client/server stack. Prefer the store-based runtime for new
|
||||
integrations.
|
||||
"""
|
||||
|
||||
def __init__(self, host: str = "127.0.0.1", port: int = 8000, task_timeout_seconds: float = 300.0):
|
||||
"""Initialize the controller.
|
||||
|
||||
Args:
|
||||
host: Hostname or IP address to bind the HTTP server to.
|
||||
port: TCP port exposed by the server.
|
||||
task_timeout_seconds: Seconds before a claimed task is considered stale and
|
||||
re-queued.
|
||||
"""
|
||||
warnings.warn(
|
||||
"AgentLightningServer is deprecated. Please use LightningStoreServer instead.", DeprecationWarning
|
||||
)
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.endpoint = f"http://{host}:{port}"
|
||||
self._task_timeout_seconds = task_timeout_seconds
|
||||
|
||||
# Defer initialization and use event for cross-thread communication
|
||||
self._store: Optional[ServerDataStore] = None
|
||||
self.loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
self.startup_event = threading.Event()
|
||||
|
||||
# Create FastAPI app instance with a lifespan manager
|
||||
self._app = FastAPI(lifespan=self._lifespan)
|
||||
self._setup_routes()
|
||||
|
||||
self._uvicorn_config = uvicorn.Config(self._app, host=self.host, port=self.port, log_level="info")
|
||||
self._uvicorn_server = uvicorn.Server(self._uvicorn_config)
|
||||
|
||||
# --- ADDED: Lifespan context manager ---
|
||||
@asynccontextmanager
|
||||
async def _lifespan(self, app: FastAPI):
|
||||
"""Manage server start-up and shutdown within the event loop."""
|
||||
logger.info("Server is starting up...")
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self._store = ServerDataStore() # Initialize data store here
|
||||
self.startup_event.set() # Signal that the server is ready
|
||||
|
||||
yield
|
||||
|
||||
logger.info("Server is shutting down.")
|
||||
self._store = None
|
||||
self.startup_event.clear() # Clear the startup event
|
||||
self.loop = None
|
||||
|
||||
async def _check_and_requeue_stale_tasks(self):
|
||||
"""Check for stale tasks and requeue them when they exceed the timeout."""
|
||||
current_time = time.time()
|
||||
# Ensure store is initialized before checking
|
||||
if not self._store:
|
||||
return
|
||||
processing_tasks = self._store.get_processing_tasks()
|
||||
|
||||
for _, task in processing_tasks.items():
|
||||
if task.last_claim_time and current_time - task.last_claim_time > self._task_timeout_seconds:
|
||||
await self._store.requeue_task(task)
|
||||
logger.warning(
|
||||
f"Task {task.rollout_id} timed out after {self._task_timeout_seconds}s, requeued (attempt {task.num_claims})"
|
||||
)
|
||||
|
||||
def _setup_routes(self):
|
||||
"""Configure the FastAPI routes that make up the legacy HTTP API."""
|
||||
|
||||
@self._app.get("/task", response_model=TaskIfAny)
|
||||
async def next_task() -> TaskIfAny: # type: ignore
|
||||
"""Provide the next available task to a client."""
|
||||
await self._check_and_requeue_stale_tasks()
|
||||
|
||||
if not self._store:
|
||||
return TaskIfAny(is_available=False)
|
||||
|
||||
task = await self._store.get_next_task()
|
||||
if task:
|
||||
logger.debug(f"Serving task {task.rollout_id} to a client.")
|
||||
return TaskIfAny(is_available=True, task=task)
|
||||
else:
|
||||
logger.debug("No task available for client.")
|
||||
return TaskIfAny(is_available=False)
|
||||
|
||||
@self._app.get("/resources/latest", response_model=ResourcesUpdate)
|
||||
async def fetch_latest_resources() -> ResourcesUpdate: # type: ignore
|
||||
"""Return the most recent resource bundle published to the server."""
|
||||
if not self._store:
|
||||
raise HTTPException(status_code=503, detail="Server not fully initialized.")
|
||||
resources_update = await self._store.get_latest_resources()
|
||||
if not resources_update:
|
||||
raise HTTPException(status_code=404, detail="No resources have been set on the server.")
|
||||
logger.debug(f"Serving latest resources '{resources_update.resources_id}' to a client.")
|
||||
return resources_update
|
||||
|
||||
@self._app.get("/resources/{resource_id}", response_model=ResourcesUpdate)
|
||||
async def fetch_resources_by_id( # type: ignore
|
||||
resource_id: str = Path(..., description="The unique identifier for the resource version.")
|
||||
) -> ResourcesUpdate:
|
||||
"""Return a specific version of resources by identifier."""
|
||||
if not self._store:
|
||||
raise HTTPException(status_code=503, detail="Server not fully initialized.")
|
||||
resources_update = await self._store.get_resources_by_id(resource_id)
|
||||
if not resources_update:
|
||||
raise HTTPException(status_code=404, detail=f"Resource ID '{resource_id}' not found.")
|
||||
logger.debug(f"Serving resources for ID '{resource_id}' to a client.")
|
||||
return resources_update
|
||||
|
||||
@self._app.post("/rollout", response_model=GenericResponse)
|
||||
async def post_rollout(payload: RolloutLegacy) -> GenericResponse: # type: ignore
|
||||
"""Persist the rollout reported by a client."""
|
||||
if not self._store:
|
||||
raise HTTPException(status_code=503, detail="Server not fully initialized.")
|
||||
await self._store.store_rollout(payload)
|
||||
return GenericResponse(
|
||||
status="ok",
|
||||
message=f"Rollout {payload.rollout_id} received and stored.",
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
"""Start the FastAPI server in the background."""
|
||||
logger.info(f"Starting server at {self.endpoint}")
|
||||
asyncio.create_task(self._uvicorn_server.serve())
|
||||
await asyncio.sleep(1) # Allow time for server to start up.
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the FastAPI server and wait for a graceful shutdown."""
|
||||
if self._uvicorn_server.started:
|
||||
logger.info("Stopping server...")
|
||||
self._uvicorn_server.should_exit = True
|
||||
await asyncio.sleep(1) # Allow time for graceful shutdown.
|
||||
logger.info("Server stopped.")
|
||||
|
||||
async def run_forever(self):
|
||||
"""Run the server indefinitely until `stop()` is invoked."""
|
||||
await self._uvicorn_server.serve()
|
||||
|
||||
async def queue_task(
|
||||
self,
|
||||
sample: Any,
|
||||
mode: Literal["train", "val", "test"] | None = None,
|
||||
resources_id: str | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Add a task to the queue for a client to process."""
|
||||
if not self._store:
|
||||
raise RuntimeError("Store not initialized. The server may not be running.")
|
||||
return await self._store.add_task(sample, mode=mode, resources_id=resources_id, metadata=metadata)
|
||||
|
||||
async def update_resources(self, resources: NamedResources) -> str:
|
||||
"""Publish a new resource bundle and return its generated identifier."""
|
||||
if not self._store:
|
||||
raise RuntimeError("Store not initialized. The server may not be running.")
|
||||
resources_id = f"res-{uuid.uuid4()}"
|
||||
update = ResourcesUpdate(
|
||||
resources_id=resources_id, resources=resources, create_time=time.time(), update_time=time.time(), version=1
|
||||
)
|
||||
await self._store.update_resources(update)
|
||||
return resources_id
|
||||
|
||||
async def get_completed_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
|
||||
"""Retrieve a specific completed rollout by identifier."""
|
||||
if not self._store:
|
||||
raise RuntimeError("Store not initialized. The server may not be running.")
|
||||
return await self._store.retrieve_rollout(rollout_id)
|
||||
|
||||
async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[RolloutLegacy]:
|
||||
"""Poll for a completed rollout until it becomes available or a timeout expires.
|
||||
|
||||
Args:
|
||||
rollout_id: Identifier of the rollout to wait for.
|
||||
timeout: Maximum number of seconds to wait. ``None`` waits indefinitely.
|
||||
|
||||
Returns:
|
||||
Retrieved rollout, or ``None`` when the timeout is reached without success.
|
||||
"""
|
||||
start_time = time.time()
|
||||
while True:
|
||||
rollout = await self.get_completed_rollout(rollout_id)
|
||||
if rollout:
|
||||
return rollout
|
||||
if timeout and (time.time() - start_time) >= timeout:
|
||||
return None
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
|
||||
"""Return every completed rollout and clear the internal buffer."""
|
||||
if not self._store:
|
||||
raise RuntimeError("Store not initialized. The server may not be running.")
|
||||
return await self._store.retrieve_completed_rollouts()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user