Compare commits
67 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| dc2c5e5d8b | |||
| 5ae7933d41 | |||
| 2eb207ec79 | |||
| 2ab977ed18 | |||
| 1eae9a34f0 | |||
| 4e7748b059 | |||
| 582f67cade | |||
| 9e23ba6b50 | |||
| de8df23805 | |||
| 5fc8ee0bf0 | |||
| fe2d05218f | |||
| c1b7827e5f | |||
| 3162ed6fb6 | |||
| 5983eec570 | |||
| c21e065d24 | |||
| b6312db2a5 | |||
| 6131c1f7b5 | |||
| 294a1cc1f6 | |||
| 93814a9a9c | |||
| 3f8a3ac0f1 | |||
| 70838ee86c | |||
| e0b55ab057 | |||
| 5950276103 | |||
| 421f2773c7 | |||
| 1857e39d5e | |||
| f717f9982f | |||
| 44dbfde0b4 | |||
| f7fe24a4b9 | |||
| a477bec3e1 | |||
| 713511902d | |||
| 80531c9c28 | |||
| 37daf2104f | |||
| cbd6498e73 | |||
| a0b68333b8 | |||
| 9afdd4570c | |||
| 96a0d58fce | |||
| b8940feffb | |||
| 55fbe66fe7 | |||
| 3794c97c1e | |||
| 848623766d | |||
| 4cd09ec900 | |||
| 47b0a4e493 | |||
| 3ed5e1e5b5 | |||
| 3f372ff7b3 | |||
| c453c41fd2 | |||
| 5c8ac61af6 | |||
| a02e1b91d9 | |||
| 496e793f0b | |||
| 80d306ff54 | |||
| 5f67bfe137 | |||
| f8c45b6ca8 | |||
| 01955aead7 | |||
| 0a9e3d75f2 | |||
| a3b2db18fa | |||
| 268bd77ce6 | |||
| ab6ea3c131 | |||
| d16538da96 | |||
| 8ce40a0410 | |||
| a9c7dbef22 | |||
| e69d24f4a8 | |||
| 955a0cc9a3 | |||
| 3966db6d2a | |||
| 584600d72e | |||
| 955524658d | |||
| 4d5e133a06 | |||
| df2a159b00 | |||
| 675fc86727 |
@@ -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@v4
|
||||
- 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,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@v4
|
||||
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
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-apo, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-apo' });
|
||||
|
||||
apo:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
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
|
||||
@@ -46,7 +44,7 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-calc-x, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-calc-x' });
|
||||
|
||||
calc-x:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-calc-x' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: Calc-X (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
timeout-minutes: 90
|
||||
@@ -49,7 +47,7 @@ jobs:
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
@@ -156,16 +154,16 @@ jobs:
|
||||
|
||||
- name: Calc-X training with external store
|
||||
run: |
|
||||
set -ex
|
||||
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 &
|
||||
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
|
||||
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
|
||||
@@ -178,10 +176,30 @@ jobs:
|
||||
sleep 5
|
||||
done
|
||||
echo "train_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_external_store
|
||||
|
||||
- name: Calc-X 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 }}
|
||||
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-compat, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-compat' });
|
||||
|
||||
backward-compatibility:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
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
|
||||
@@ -47,7 +45,7 @@ jobs:
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
@@ -56,6 +54,10 @@ jobs:
|
||||
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
|
||||
if: matrix.setup-script == 'stable'
|
||||
- name: Freeze dependencies
|
||||
run: |
|
||||
set -ex
|
||||
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-spider, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-spider' });
|
||||
|
||||
spider:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
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
|
||||
@@ -49,7 +47,7 @@ jobs:
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-unsloth, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-unsloth' });
|
||||
|
||||
unsloth:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
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
|
||||
@@ -48,7 +46,7 @@ jobs:
|
||||
run: df -h
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
|
||||
@@ -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}'.`);
|
||||
}
|
||||
@@ -26,6 +26,14 @@ jobs:
|
||||
- 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: |
|
||||
|
||||
@@ -60,6 +60,14 @@ jobs:
|
||||
- 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
|
||||
|
||||
@@ -8,27 +8,25 @@ on:
|
||||
|
||||
workflow_dispatch:
|
||||
|
||||
pull_request_target:
|
||||
types: [reopened, ready_for_review]
|
||||
repository_dispatch:
|
||||
types: [ci-gpu, ci-all]
|
||||
|
||||
run-name: >-
|
||||
${{ github.event_name == 'repository_dispatch'
|
||||
&& format(
|
||||
'PR #{0} - Label {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:
|
||||
label-check:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
should-run: ${{ steps.evaluate.outputs.should-run }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Decide whether to run
|
||||
id: evaluate
|
||||
uses: actions/github-script@v7
|
||||
with:
|
||||
script: |
|
||||
const evaluateTrigger = require('./scripts/trigger_on_label.js');
|
||||
evaluateTrigger({ core, context, labelName: 'ci-gpu' });
|
||||
|
||||
tests-full:
|
||||
needs: label-check
|
||||
if: needs.label-check.outputs.should-run == 'true'
|
||||
if: >
|
||||
github.event_name != 'repository_dispatch' ||
|
||||
github.event.action == 'ci-gpu' ||
|
||||
github.event.action == 'ci-all'
|
||||
name: GPU Test with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
|
||||
|
||||
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
|
||||
@@ -48,7 +46,7 @@ jobs:
|
||||
run: nvidia-smi
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number) || github.ref }}
|
||||
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
|
||||
@@ -75,6 +73,14 @@ jobs:
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- 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: Launch LiteLLM Proxy
|
||||
run: |
|
||||
./scripts/litellm_run.sh
|
||||
|
||||
@@ -5,9 +5,9 @@ permissions:
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
branches: [ main, stable/**/* ]
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
branches: [ main, stable/**/* ]
|
||||
workflow_dispatch:
|
||||
|
||||
schedule:
|
||||
@@ -41,13 +41,15 @@ jobs:
|
||||
--group torch-cpu \
|
||||
--group torch-stable \
|
||||
--group trl \
|
||||
--group tinker \
|
||||
--group agents \
|
||||
--no-default-groups
|
||||
if: matrix.setup == 'slow'
|
||||
# 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_python_headers.py
|
||||
run: uv run --locked --no-sync scripts/check_headers.py
|
||||
- name: Run Black
|
||||
run: uv run --locked --no-sync black --check .
|
||||
- name: Run isort
|
||||
@@ -59,6 +61,28 @@ jobs:
|
||||
run: uv run --locked --no-sync pyright -p pyrightconfig.json
|
||||
if: matrix.setup == 'slow'
|
||||
|
||||
lint-js:
|
||||
name: Lint - JavaScript
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- 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
|
||||
@@ -131,8 +155,39 @@ jobs:
|
||||
name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
|
||||
path: requirements-freeze.txt
|
||||
compression-level: 0
|
||||
|
||||
- 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: Run tests
|
||||
run: |
|
||||
uv run pytest -v --durations=0 tests
|
||||
env:
|
||||
PYTEST_ADDOPTS: "--color=yes"
|
||||
|
||||
test-js:
|
||||
name: Test - JavaScript
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
- 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
|
||||
|
||||
+14
@@ -189,6 +189,9 @@ cython_debug/
|
||||
# you could uncomment the following to ignore the enitre vscode folder
|
||||
.vscode/
|
||||
|
||||
# Emacs backup files
|
||||
*~
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
@@ -204,3 +207,14 @@ cython_debug/
|
||||
|
||||
# Claude
|
||||
.claude/*.local.json
|
||||
|
||||
# Temporary and backup files
|
||||
*.tmp
|
||||
*.bak
|
||||
*.backup
|
||||
|
||||
# Dashboard generated files
|
||||
agentlightning/dashboard/**/*.css
|
||||
agentlightning/dashboard/**/*.js
|
||||
agentlightning/dashboard/**/*.html
|
||||
agentlightning/dashboard/**/*.svg
|
||||
|
||||
@@ -24,3 +24,53 @@ repos:
|
||||
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
|
||||
'
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
[](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.**
|
||||
@@ -39,6 +40,8 @@ To start using Agent-lightning, check out our [documentation](https://microsoft.
|
||||
|
||||
## ⚡ Articles
|
||||
|
||||
- 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.
|
||||
@@ -89,7 +92,7 @@ If you find Agent Lightning useful in your research or projects, please cite our
|
||||
|
||||
## ⚡ Contributing
|
||||
|
||||
This project welcomes contributions and suggestions. 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.
|
||||
This project welcomes contributions and suggestions. Start by reading the [Contributing Guide](docs/community/contributing.md) for 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.
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
__version__ = "0.2.1"
|
||||
__version__ = "0.2.2"
|
||||
|
||||
from .adapter import *
|
||||
from .algorithm import *
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
@@ -14,6 +15,8 @@ from agentlightning.types import Span, SpanNames, Triplet
|
||||
|
||||
from .base import TraceAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Transition(BaseModel):
|
||||
"""A single transition within a reinforcement learning trajectory.
|
||||
@@ -426,7 +429,16 @@ class TraceTree:
|
||||
If we don't, when we want to select the LLM completion span with agent as filter.
|
||||
We will never get the correct span underneath.
|
||||
"""
|
||||
# If the current node has only one child, recursively repair its hierarchy directly.
|
||||
# This special-case handling is needed because when a trace is manually ended
|
||||
# (via agentops.end_trace), the AgentOps provider automatically wraps all spans
|
||||
# under an extra synthetic root node (e.g., "run_one.session").
|
||||
if len(self.children) == 1:
|
||||
self.children[0].repair_hierarchy()
|
||||
return
|
||||
|
||||
nodes_to_repair = list(self.children)
|
||||
|
||||
for repair_node in nodes_to_repair:
|
||||
if len(self.children) == 1:
|
||||
# If there is only one child, we don't need to repair the hierarchy.
|
||||
@@ -505,6 +517,30 @@ class TraceTree:
|
||||
|
||||
return rewards
|
||||
|
||||
def span_to_triplet(self, span: Span, agent_name: str) -> Triplet:
|
||||
"""Convert a span to a triplet.
|
||||
|
||||
Subclass can override this method to add more fields to the triplet,
|
||||
such as chat messages and tool calls.
|
||||
"""
|
||||
prompt_token_ids = span.attributes.get("prompt_token_ids", []) # type: ignore
|
||||
response_token_ids = span.attributes.get("response_token_ids", []) # type: ignore
|
||||
response_id = span.attributes.get("gen_ai.response.id", None) # type: ignore
|
||||
|
||||
logprobs_content = span.attributes.get("logprobs.content", None) # type: ignore
|
||||
if isinstance(logprobs_content, str):
|
||||
logprobs_content = json.loads(logprobs_content)
|
||||
response: Dict[str, Any] = {"token_ids": response_token_ids, "logprobs": logprobs_content}
|
||||
else:
|
||||
response = {"token_ids": response_token_ids}
|
||||
|
||||
return Triplet(
|
||||
prompt={"token_ids": prompt_token_ids},
|
||||
response=response,
|
||||
reward=None,
|
||||
metadata=dict(response_id=response_id, agent_name=agent_name),
|
||||
)
|
||||
|
||||
def to_trajectory(
|
||||
self,
|
||||
llm_call_match: str = r"openai\.chat\.completion",
|
||||
@@ -513,6 +549,7 @@ class TraceTree:
|
||||
dedup_llm_call: bool = True,
|
||||
reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
|
||||
final_reward: Optional[float] = None,
|
||||
_skip_empty_token_spans: bool = False,
|
||||
) -> List[Triplet]:
|
||||
"""Convert the trace tree into a trajectory of [`Triplet`][agentlightning.Triplet] items.
|
||||
|
||||
@@ -537,25 +574,23 @@ class TraceTree:
|
||||
within_llm_call=False if dedup_llm_call else None,
|
||||
existing_llm_call_response_ids=set(),
|
||||
)
|
||||
id_transitions = [
|
||||
(
|
||||
llm_call.id,
|
||||
Triplet(
|
||||
prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])}, # type: ignore
|
||||
response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])}, # type: ignore
|
||||
reward=None,
|
||||
metadata=dict(
|
||||
response_id=llm_call.span.attributes.get( # type: ignore
|
||||
"gen_ai.response.id", None
|
||||
), # it works at least for OpenAI
|
||||
agent_name=agent_name,
|
||||
),
|
||||
),
|
||||
)
|
||||
for llm_call, agent_name in llm_calls
|
||||
]
|
||||
|
||||
rewards = self.match_rewards(reward_match, [call for call, _ in llm_calls])
|
||||
id_transitions: List[Tuple[str, Triplet]] = []
|
||||
# We need to filter out the LLM calls with unrecorded token IDs
|
||||
filtered_llm_calls: List[Tuple[TraceTree, str]] = []
|
||||
for llm_call, agent_name in llm_calls:
|
||||
triplet = self.span_to_triplet(llm_call.span, agent_name)
|
||||
# This is a hot-fix for Tinker+CrewAI, which has some anonymous requests outside the trained agent.
|
||||
# TODO: We might need to reconsider this.
|
||||
if _skip_empty_token_spans and (
|
||||
not triplet.prompt.get("token_ids") or not triplet.response.get("token_ids")
|
||||
):
|
||||
logger.warning(f"Skipping LLM call with unrecorded token IDs: {triplet}")
|
||||
continue
|
||||
filtered_llm_calls.append((llm_call, agent_name))
|
||||
id_transitions.append((llm_call.id, triplet))
|
||||
|
||||
rewards = self.match_rewards(reward_match, [call for call, _ in filtered_llm_calls])
|
||||
transitions = [
|
||||
transition.model_copy(update={"reward": rewards.get(id, None)}) for id, transition in id_transitions
|
||||
]
|
||||
@@ -597,12 +632,14 @@ class TracerTraceToTriplet(TraceToTripletBase):
|
||||
agent_match: Optional[str] = None,
|
||||
exclude_llm_call_in_reward: bool = True,
|
||||
reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
|
||||
_skip_empty_token_spans: bool = False,
|
||||
):
|
||||
self.repair_hierarchy = repair_hierarchy
|
||||
self.llm_call_match = llm_call_match
|
||||
self.agent_match = agent_match
|
||||
self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
|
||||
self.reward_match = reward_match
|
||||
self._skip_empty_token_spans = _skip_empty_token_spans
|
||||
|
||||
def visualize(
|
||||
self,
|
||||
@@ -654,6 +691,7 @@ class TracerTraceToTriplet(TraceToTripletBase):
|
||||
agent_match=self.agent_match,
|
||||
exclude_llm_call_in_reward=self.exclude_llm_call_in_reward,
|
||||
reward_match=self.reward_match,
|
||||
_skip_empty_token_spans=self._skip_empty_token_spans,
|
||||
)
|
||||
return trajectory
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ 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
|
||||
from agentlightning.reward import find_final_reward
|
||||
from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
|
||||
|
||||
@@ -56,41 +57,6 @@ APPLY_EDIT_PROMPT_FILES = [
|
||||
]
|
||||
|
||||
|
||||
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 = []
|
||||
|
||||
|
||||
class APO(Algorithm, Generic[T_task]):
|
||||
"""Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
|
||||
|
||||
|
||||
@@ -128,11 +128,13 @@ class Baseline(FastAlgorithm):
|
||||
# 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}")
|
||||
except ValueError:
|
||||
logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
|
||||
|
||||
async def _enqueue_rollouts(
|
||||
self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import random
|
||||
from typing import Iterator, List, Sequence, TypeVar
|
||||
|
||||
from agentlightning.types import Dataset
|
||||
|
||||
T_task = TypeVar("T_task")
|
||||
|
||||
|
||||
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 = []
|
||||
@@ -99,6 +99,8 @@ class VERL(Algorithm):
|
||||
|
||||
# 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)
|
||||
|
||||
def run(
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from typing import Iterable
|
||||
|
||||
from agentlightning.instrumentation.agentops import AgentOpsServerManager
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Start AgentOps server")
|
||||
parser.add_argument("--daemon", action="store_true", help="Run server as a daemon")
|
||||
parser.add_argument("--port", type=int, default=8002, help="Port to run the server on")
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
manager = AgentOpsServerManager(daemon=args.daemon, port=args.port)
|
||||
try:
|
||||
manager.start()
|
||||
# Wait forever
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
manager.stop()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -6,23 +6,41 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Iterable
|
||||
|
||||
from agentlightning.logging import configure_logger
|
||||
from agentlightning.store.client_server import LightningStoreServer
|
||||
from agentlightning.store.memory import InMemoryLightningStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main(argv: Iterable[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Run a LightningStore server")
|
||||
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.",
|
||||
)
|
||||
args = parser.parse_args(list(argv) if argv is not None else None)
|
||||
|
||||
configure_logger()
|
||||
|
||||
store = InMemoryLightningStore()
|
||||
server = LightningStoreServer(store, host="0.0.0.0", port=args.port)
|
||||
asyncio.run(server.run_forever())
|
||||
server = LightningStoreServer(
|
||||
store,
|
||||
host="0.0.0.0",
|
||||
port=args.port,
|
||||
cors_allow_origins=args.cors_origins,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -332,7 +332,9 @@ class DevTaskLoader(AgentLightningClient):
|
||||
if isinstance(resources, ResourcesUpdate):
|
||||
self._resources_update = resources
|
||||
else:
|
||||
self._resources_update = ResourcesUpdate(resources_id="local", resources=resources)
|
||||
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] = []
|
||||
|
||||
@@ -149,6 +149,7 @@ def emit_reward(reward: float) -> ReadableSpan:
|
||||
if not isinstance(reward, float):
|
||||
raise ValueError(f"Reward must be a number, got: {type(reward)}")
|
||||
|
||||
# TODO: This should use the tracer from current context by tracer
|
||||
tracer = get_tracer()
|
||||
span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
|
||||
# Do nothing; it's just a number
|
||||
|
||||
@@ -2,29 +2,76 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import multiprocessing
|
||||
import signal
|
||||
import socket
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
from typing import Any, Callable, no_type_check
|
||||
|
||||
import flask
|
||||
import requests
|
||||
import setproctitle
|
||||
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 opentelemetry.sdk.trace.export import SpanExportResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"instrument_agentops",
|
||||
"uninstrument_agentops",
|
||||
"agentops_local_server",
|
||||
"AgentOpsServerManager",
|
||||
]
|
||||
|
||||
# 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.
|
||||
|
||||
False (default): AgentOps exporters and clients will run in local mode
|
||||
and will not attempt to communicate with the remote AgentOps service.
|
||||
True: all exporters and clients will operate in normal mode and send data
|
||||
to the AgentOps service as expected.
|
||||
"""
|
||||
global _agentops_service_enabled
|
||||
_agentops_service_enabled = enabled
|
||||
logger.info(f"Switch set to {enabled} for exporters and clients.")
|
||||
|
||||
|
||||
def _patch_exporters():
|
||||
import agentops.client.api
|
||||
import agentops.sdk.core
|
||||
import opentelemetry.exporter.otlp.proto.http.metric_exporter
|
||||
import opentelemetry.exporter.otlp.proto.http.trace_exporter
|
||||
|
||||
agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter # type: ignore
|
||||
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = BypassableOTLPMetricExporter
|
||||
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = BypassableOTLPSpanExporter
|
||||
agentops.client.api.V3Client = BypassableV3Client
|
||||
agentops.client.api.V4Client = BypassableV4Client
|
||||
|
||||
|
||||
def _unpatch_exporters():
|
||||
import agentops.client.api
|
||||
import agentops.sdk.core
|
||||
import opentelemetry.exporter.otlp.proto.http.metric_exporter
|
||||
import opentelemetry.exporter.otlp.proto.http.trace_exporter
|
||||
|
||||
agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter # type: ignore
|
||||
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = OTLPMetricExporter
|
||||
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = OTLPSpanExporter
|
||||
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():
|
||||
@@ -40,41 +87,58 @@ def _patch_new_agentops():
|
||||
|
||||
_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) # type: ignore
|
||||
if return_value is not None and hasattr(return_value, "prompt_token_ids"): # type: ignore
|
||||
attributes["prompt_token_ids"] = list(return_value.prompt_token_ids) # type: ignore
|
||||
if return_value is not None and hasattr(return_value, "response_token_ids"): # type: ignore
|
||||
attributes["response_token_ids"] = list(return_value.response_token_ids[0]) # 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 (
|
||||
not attributes.get("response_token_ids")
|
||||
and return_value is not None
|
||||
and hasattr(return_value, "choices") # type: ignore
|
||||
and return_value.choices # type: ignore
|
||||
and isinstance(return_value.choices, list) # type: ignore
|
||||
):
|
||||
first_choice = return_value.choices[0] # type: ignore
|
||||
if hasattr(first_choice, "token_ids"): # type: ignore
|
||||
attributes["response_token_ids"] = list(first_choice.token_ids) # type: ignore
|
||||
# newer versions of OpenAI client SDK
|
||||
elif hasattr(first_choice, "provider_specific_fields") and "token_ids" in first_choice.provider_specific_fields: # type: ignore
|
||||
attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"]) # type: ignore
|
||||
|
||||
# For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
|
||||
if (
|
||||
return_value is not None
|
||||
and hasattr(return_value, "http_response") # type: ignore
|
||||
and return_value.http_response is not None # type: ignore
|
||||
and hasattr(return_value.http_response, "json") # type: ignore
|
||||
and hasattr(return_value, "choices")
|
||||
and return_value.choices
|
||||
and isinstance(return_value.choices, list)
|
||||
and len(return_value.choices) > 0
|
||||
):
|
||||
json_data = return_value.http_response.json() # type: ignore
|
||||
if isinstance(json_data, dict):
|
||||
if "prompt_token_ids" in json_data:
|
||||
attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"]) # type: ignore
|
||||
if "response_token_ids" in json_data:
|
||||
attributes["response_token_ids"] = list(json_data["response_token_ids"][0]) # type: ignore
|
||||
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
|
||||
|
||||
@@ -146,6 +210,8 @@ 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()
|
||||
@@ -165,6 +231,8 @@ def instrument_agentops():
|
||||
|
||||
def uninstrument_agentops():
|
||||
"""Uninstrument agentops to stop capturing token IDs."""
|
||||
_unpatch_exporters()
|
||||
|
||||
try:
|
||||
_unpatch_new_agentops()
|
||||
except Exception:
|
||||
@@ -175,114 +243,75 @@ def uninstrument_agentops():
|
||||
pass
|
||||
|
||||
|
||||
def agentops_local_server():
|
||||
class BypassableAuthenticatedOTLPExporter(AuthenticatedOTLPExporter):
|
||||
"""
|
||||
Returns a Flask app that can be used to test agentops integration.
|
||||
This server provides endpoints for token fetching and a catch-all endpoint.
|
||||
AuthenticatedOTLPExporter with switchable service control.
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
"""
|
||||
app = flask.Flask(__name__)
|
||||
|
||||
@app.route("/v3/auth/token", methods=["POST"])
|
||||
def fetch_token(): # type: ignore
|
||||
return {"token": "dummy", "project_id": "dummy"}
|
||||
|
||||
@app.route("/", defaults={"path": ""}, methods=["GET", "POST"])
|
||||
@app.route("/<path:path>", methods=["GET", "POST"])
|
||||
def catch_all(path: str): # type: ignore
|
||||
return {"path": path}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _run_server(**kwargs: Any): # type: ignore
|
||||
"""
|
||||
Internal function to run the Flask server.
|
||||
This is used to avoid issues with multiprocessing and Flask's reloader.
|
||||
"""
|
||||
signal.signal(signal.SIGINT, signal.SIG_IGN) # Ignore SIGINT in worker processes
|
||||
setproctitle.setproctitle(multiprocessing.current_process().name)
|
||||
app = agentops_local_server()
|
||||
app.run(**kwargs)
|
||||
|
||||
|
||||
class AgentOpsServerManager:
|
||||
"""Manages a AgentOps local server to bypass the online service of AgentOps."""
|
||||
|
||||
def __init__(self, daemon: bool = True, port: int | None = None):
|
||||
self.server_process: multiprocessing.Process | None = None
|
||||
self.server_port = port
|
||||
self.daemon = daemon
|
||||
logger.info("AgentOpsServerManager initialized.")
|
||||
|
||||
def _find_available_port(self) -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
def start(self):
|
||||
if self.server_process and self.server_process.is_alive():
|
||||
logger.warning("AgentOps server process appears to be already running.")
|
||||
return
|
||||
|
||||
if self.server_port is None:
|
||||
self.server_port = self._find_available_port()
|
||||
|
||||
logger.info(f"Starting AgentOps local server on port {self.server_port}...")
|
||||
|
||||
self.server_process = multiprocessing.Process(
|
||||
target=_run_server,
|
||||
kwargs={"host": "127.0.0.1", "port": self.server_port, "use_reloader": False, "debug": False},
|
||||
daemon=self.daemon,
|
||||
name="AgentLightning-AgentOpsServer",
|
||||
)
|
||||
self.server_process.start()
|
||||
logger.info(
|
||||
f"AgentOps local server process (PID: {self.server_process.pid}) started, targeting port {self.server_port}."
|
||||
)
|
||||
for attempt in range(20): # 10 seconds total
|
||||
time.sleep(0.5) # Brief wait for server to start up
|
||||
try:
|
||||
result = requests.get(f"http://127.0.0.1:{self.server_port}/")
|
||||
if result.status_code == 200:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.debug(f"Error checking AgentOps server: {e}")
|
||||
logger.warning(f"AgentOps still not ready after {attempt} attempts. Retrying...")
|
||||
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
|
||||
if _agentops_service_enabled:
|
||||
return super().export(*args, **kwargs)
|
||||
else:
|
||||
logger.error(f"AgentOps local server failed to start or exited prematurely.")
|
||||
return
|
||||
logger.debug("SwitchableAuthenticatedOTLPExporter is switched off, skipping export.")
|
||||
return SpanExportResult.SUCCESS
|
||||
|
||||
if not self.server_process.is_alive():
|
||||
logger.error(f"AgentOps local server failed to start or exited prematurely.")
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
if self.server_process and self.server_process.is_alive():
|
||||
return True
|
||||
return False
|
||||
class BypassableOTLPMetricExporter(OTLPMetricExporter):
|
||||
"""
|
||||
OTLPMetricExporter with switchable service control.
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
"""
|
||||
|
||||
def stop(self):
|
||||
if self.server_process is not None and self.server_process.is_alive():
|
||||
logger.info(f"Stopping AgentOps local server (PID: {self.server_process.pid})...")
|
||||
self.server_process.terminate() # Send SIGTERM
|
||||
self.server_process.join(timeout=5) # Wait for clean exit
|
||||
if self.server_process.is_alive():
|
||||
logger.warning(
|
||||
f"AgentOps server (PID: {self.server_process.pid}) did not terminate gracefully, killing..."
|
||||
)
|
||||
self.server_process.kill() # Force kill
|
||||
self.server_process.join(timeout=10) # Wait for kill
|
||||
self.server_process = None
|
||||
logger.info(f"AgentOps local server stopped.")
|
||||
def export(self, *args: Any, **kwargs: Any) -> MetricExportResult:
|
||||
if _agentops_service_enabled:
|
||||
return super().export(*args, **kwargs) # type: ignore[reportUnknownMemberType]
|
||||
else:
|
||||
logger.info("AgentOps local server was not running or already stopped.")
|
||||
logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
|
||||
return MetricExportResult.SUCCESS
|
||||
|
||||
def get_port(self) -> int | None:
|
||||
# Check liveness again in case it died since start()
|
||||
if self.is_alive() and self.server_port is not None:
|
||||
return self.server_port
|
||||
# If called after server stopped or failed, port might be stale or None
|
||||
if self.server_port is not None and (self.server_process is None or not self.server_process.is_alive()):
|
||||
logger.warning(
|
||||
f"AgentOps server port {self.server_port} is stored, but server process is not alive. Returning stored port."
|
||||
)
|
||||
return self.server_port
|
||||
|
||||
class BypassableOTLPSpanExporter(OTLPSpanExporter):
|
||||
"""
|
||||
OTLPSpanExporter with switchable service control.
|
||||
When `_agentops_service_enabled` is False, skip export and return success.
|
||||
"""
|
||||
|
||||
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
|
||||
if _agentops_service_enabled:
|
||||
return super().export(*args, **kwargs)
|
||||
else:
|
||||
logger.debug("SwitchableOTLPSpanExporter is switched off, skipping export.")
|
||||
return SpanExportResult.SUCCESS
|
||||
|
||||
|
||||
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
|
||||
|
||||
+135
-134
@@ -7,15 +7,26 @@ import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional, Sequence, TypedDict, Union, cast
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
TypedDict,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import litellm
|
||||
import opentelemetry.trace as trace_api
|
||||
import uvicorn
|
||||
import yaml
|
||||
from fastapi import Request, Response
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
@@ -26,6 +37,12 @@ from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from agentlightning.types import LLM, ProxyLLM
|
||||
from agentlightning.utils.server_launcher import (
|
||||
LaunchMode,
|
||||
PythonServerLauncher,
|
||||
PythonServerLauncherArgs,
|
||||
noop_context,
|
||||
)
|
||||
|
||||
from .store.base import LightningStore
|
||||
|
||||
@@ -498,9 +515,10 @@ class LLMProxy:
|
||||
* [`stop()`][agentlightning.LLMProxy.stop] tears down the server and removes the temp config file.
|
||||
* [`restart()`][agentlightning.LLMProxy.restart] convenience wrapper to stop then start.
|
||||
|
||||
Usage Note:
|
||||
As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
|
||||
process from the main runner (i.e., tracer from agents).
|
||||
!!! note
|
||||
|
||||
As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
|
||||
process from the main runner (i.e., tracer from agents). See `launch_mode` for how to change that.
|
||||
|
||||
!!! warning
|
||||
|
||||
@@ -512,37 +530,66 @@ class LLMProxy:
|
||||
with tracers like [`AgentOpsTracer`][agentlightning.AgentOpsTracer].
|
||||
|
||||
Args:
|
||||
port: TCP port to bind.
|
||||
port: TCP port to bind. Will bind to a random port if not provided.
|
||||
model_list: LiteLLM `model_list` entries.
|
||||
store: LightningStore used for span sequence and persistence.
|
||||
host: Publicly reachable host used in resource endpoints. Defaults to best-guess IPv4.
|
||||
host: Publicly reachable host used in resource endpoints. See `host` of `launcher_args` for more details.
|
||||
litellm_config: Extra LiteLLM proxy config merged with `model_list`.
|
||||
num_retries: Default LiteLLM retry count injected into `litellm_settings`.
|
||||
num_workers: Number of workers to run in the server. Only applicable for "mp" launch mode. Ignored if launcher_args is provided.
|
||||
When `num_workers > 1`, the server will be run using [gunicorn](https://gunicorn.org/).
|
||||
launch_mode: Launch mode for the server. Defaults to "mp". Cannot be used together with launcher_args. Ignored if launcher_args is provided.
|
||||
It's recommended to use `launch_mode="mp"` to launch the proxy, which will launch the server in a separate process.
|
||||
`launch_mode="thread"` can also be used if used in caution. It will launch the server in a separate thread.
|
||||
`launch_mode="asyncio"` launches the server in the current thread as an asyncio task.
|
||||
It is NOT recommended because it often causes hanging requests. Only use it if you know what you are doing.
|
||||
launcher_args: Arguments for the server launcher. If this is provided, host, port, and launch_mode will be ignored. Cannot be used together with port, host, and launch_mode.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
port: int,
|
||||
port: int | None = None,
|
||||
model_list: List[ModelConfig] | None = None,
|
||||
store: Optional[LightningStore] = None,
|
||||
host: str | None = None,
|
||||
litellm_config: Dict[str, Any] | None = None,
|
||||
num_retries: int = 0,
|
||||
num_workers: int = 1,
|
||||
launch_mode: LaunchMode = "mp",
|
||||
launcher_args: PythonServerLauncherArgs | None = None,
|
||||
_add_return_token_ids: bool = True,
|
||||
):
|
||||
self.store = store
|
||||
self.host = host or _get_default_ipv4_address()
|
||||
self.port = port
|
||||
|
||||
if launcher_args is not None and (
|
||||
port is not None or host is not None or launch_mode != "mp" or num_workers != 1
|
||||
):
|
||||
raise ValueError("port, host, launch_mode, and num_workers cannot be set when launcher_args is provided.")
|
||||
|
||||
self.server_launcher_args = launcher_args or PythonServerLauncherArgs(
|
||||
port=port,
|
||||
host=host,
|
||||
launch_mode=launch_mode,
|
||||
n_workers=num_workers,
|
||||
# NOTE: This /health endpoint can be slow sometimes because it actually probes the backend LLM service.
|
||||
healthcheck_url="/health",
|
||||
startup_timeout=60.0,
|
||||
)
|
||||
|
||||
if self.server_launcher_args.healthcheck_url is None:
|
||||
logger.warning("healthcheck_url is not set. LLM Proxy will not be checked for healthiness after starting.")
|
||||
|
||||
self.model_list = model_list or []
|
||||
self.litellm_config = litellm_config or {}
|
||||
|
||||
# Ensure num_retries is present inside the litellm_settings block.
|
||||
self.litellm_config.setdefault("litellm_settings", {})
|
||||
self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
|
||||
self.server_launcher = PythonServerLauncher(app, self.server_launcher_args, noop_context())
|
||||
|
||||
self._server_thread = None
|
||||
self._config_file = None
|
||||
self._uvicorn_server = None
|
||||
self._ready_event = threading.Event()
|
||||
|
||||
self._add_return_token_ids = _add_return_token_ids
|
||||
|
||||
def get_store(self) -> Optional[LightningStore]:
|
||||
"""Get the store used by the proxy.
|
||||
@@ -561,44 +608,15 @@ class LLMProxy:
|
||||
self.store = store
|
||||
|
||||
def update_model_list(self, model_list: List[ModelConfig]) -> None:
|
||||
"""Replace the in-memory model list and hot-restart if running.
|
||||
"""Replace the in-memory model list.
|
||||
|
||||
Args:
|
||||
model_list: New list of model entries.
|
||||
"""
|
||||
self.model_list = model_list
|
||||
logger.info(f"Updating LLMProxy model list to: {model_list}")
|
||||
if self.is_running():
|
||||
self.restart()
|
||||
# Do nothing if the server is not running.
|
||||
|
||||
def update_port(self, port: int) -> None:
|
||||
"""Update the port for the proxy.
|
||||
|
||||
Args:
|
||||
port: The new port to use for the proxy.
|
||||
"""
|
||||
self.port = port
|
||||
|
||||
def _wait_until_started(self, startup_timeout: float = 20.0):
|
||||
"""Block until the uvicorn server reports started or timeout.
|
||||
|
||||
Args:
|
||||
startup_timeout: Maximum seconds to wait.
|
||||
"""
|
||||
start = time.time()
|
||||
while True:
|
||||
if self._uvicorn_server is None:
|
||||
break
|
||||
if self._uvicorn_server.started:
|
||||
self._ready_event.set()
|
||||
break
|
||||
if self._uvicorn_server.should_exit:
|
||||
break
|
||||
if time.time() - start > startup_timeout:
|
||||
break
|
||||
time.sleep(0.01)
|
||||
|
||||
def initialize(self):
|
||||
"""Initialize global middleware and LiteLLM callbacks.
|
||||
|
||||
@@ -637,24 +655,17 @@ class LLMProxy:
|
||||
logger.info("Adding a new middleware to the FastAPI app.")
|
||||
app.add_middleware(RolloutAttemptMiddleware)
|
||||
|
||||
if not initialize_llm_callbacks():
|
||||
if not initialize_llm_callbacks(self._add_return_token_ids):
|
||||
# If it's not the first time to initialize the callbacks, also
|
||||
# reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
|
||||
_reset_litellm_logging_worker()
|
||||
|
||||
def start(self):
|
||||
"""Start the proxy server thread and initialize global wiring.
|
||||
@asynccontextmanager
|
||||
async def _serve_context(self) -> AsyncGenerator[None, None]:
|
||||
"""Context manager to serve the proxy server.
|
||||
|
||||
Side effects:
|
||||
|
||||
* Sets the module-level global store for middleware/exporter access.
|
||||
* Calls `initialize()` once to register middleware and callbacks.
|
||||
* Writes a temporary YAML config consumed by LiteLLM worker.
|
||||
* Launches uvicorn in a daemon thread and waits for readiness.
|
||||
See [`start`][agentlightning.LLMProxy.start] and [`stop`][agentlightning.LLMProxy.stop] for more details.
|
||||
"""
|
||||
if self.is_running():
|
||||
# Trigger restart
|
||||
self.stop()
|
||||
|
||||
if not self.store:
|
||||
raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
|
||||
@@ -675,24 +686,59 @@ class LLMProxy:
|
||||
|
||||
save_worker_config(config=self._config_file)
|
||||
|
||||
# Bind to all interfaces to allow other hosts to reach it if needed.
|
||||
self._uvicorn_server = uvicorn.Server(uvicorn.Config(app, host="0.0.0.0", port=self.port))
|
||||
|
||||
def run_server():
|
||||
# Serve uvicorn in this background thread with its own event loop.
|
||||
assert self._uvicorn_server is not None
|
||||
asyncio.run(self._uvicorn_server.serve())
|
||||
|
||||
logger.info("Starting LLMProxy server thread...")
|
||||
self._ready_event.clear()
|
||||
# FIXME: This thread should either be reused or the whole proxy should live in another process.
|
||||
# NOTE: When running the _serve_context in current process, you might encounter the following problems:
|
||||
# Problem 1: in litellm worker, <Queue at 0x70f1d028cd90 maxsize=50000> is bound to a different event loop
|
||||
# Problem 2: Proxy has conflicted opentelemetry setup with the main process.
|
||||
self._server_thread = threading.Thread(target=run_server, daemon=True)
|
||||
self._server_thread.start()
|
||||
self._wait_until_started()
|
||||
|
||||
def stop(self):
|
||||
# Ready
|
||||
logger.info("LLMProxy preparation is done. Will start the server.")
|
||||
yield
|
||||
|
||||
# Clean up
|
||||
|
||||
logger.info("LLMProxy server is cleaning up.")
|
||||
|
||||
# Remove worker config to avoid stale references.
|
||||
if self._config_file and os.path.exists(self._config_file):
|
||||
os.unlink(self._config_file)
|
||||
|
||||
logger.info("LLMProxy server finishes.")
|
||||
|
||||
async def start(self):
|
||||
"""Start the proxy server thread and initialize global wiring.
|
||||
|
||||
Side effects:
|
||||
|
||||
* Sets the module-level global store for middleware/exporter access.
|
||||
* Calls `initialize()` once to register middleware and callbacks.
|
||||
* Writes a temporary YAML config consumed by LiteLLM worker.
|
||||
* Launches uvicorn in a daemon thread and waits for readiness.
|
||||
"""
|
||||
# Refresh the serve context
|
||||
self.server_launcher.serve_context = self._serve_context()
|
||||
|
||||
if self.store is None:
|
||||
raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
|
||||
|
||||
store_capabilities = self.store.capabilities()
|
||||
if self.server_launcher.args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
|
||||
raise RuntimeError(
|
||||
"The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
|
||||
)
|
||||
elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
|
||||
raise RuntimeError(
|
||||
"The store is not thread-safe. Please use another store, or use asyncio mode to launch the server."
|
||||
)
|
||||
elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities["async_safe"]:
|
||||
raise RuntimeError("The store is not async-safe. Please use another store.")
|
||||
|
||||
logger.info(
|
||||
f"Starting LLMProxy server in {self.server_launcher.args.launch_mode} mode with store capabilities: {store_capabilities}"
|
||||
)
|
||||
|
||||
await self.server_launcher.start()
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the proxy server and clean up temporary artifacts.
|
||||
|
||||
This is a best-effort graceful shutdown with a bounded join timeout.
|
||||
@@ -701,43 +747,19 @@ class LLMProxy:
|
||||
logger.warning("LLMProxy is not running. Nothing to stop.")
|
||||
return
|
||||
|
||||
# Remove worker config to avoid stale references.
|
||||
if self._config_file and os.path.exists(self._config_file):
|
||||
os.unlink(self._config_file)
|
||||
await self.server_launcher.stop()
|
||||
|
||||
logger.info("Stopping LLMProxy server thread...")
|
||||
stop_success = True
|
||||
if self._server_thread is not None and self._uvicorn_server is not None and self._uvicorn_server.started:
|
||||
self._uvicorn_server.should_exit = True
|
||||
self._server_thread.join(timeout=10.0) # Allow time for graceful shutdown.
|
||||
if self._server_thread.is_alive():
|
||||
logger.error(
|
||||
"LLMProxy server thread is still alive after 10 seconds. Cannot kill it because it's a thread."
|
||||
)
|
||||
stop_success = False
|
||||
self._server_thread = None
|
||||
self._uvicorn_server = None
|
||||
self._config_file = None
|
||||
self._ready_event.clear()
|
||||
if not _check_port(self.host, self.port):
|
||||
logger.error(f"Port {self.port} is still in use. Stopping LLMProxy is not successful.")
|
||||
stop_success = False
|
||||
if stop_success:
|
||||
logger.info("LLMProxy server thread stopped.")
|
||||
else:
|
||||
logger.error("LLMProxy server is not stopped successfully.")
|
||||
|
||||
def restart(self, *, _port: int | None = None) -> None:
|
||||
async def restart(self, *, _port: int | None = None) -> None:
|
||||
"""Restart the proxy if running, else start it.
|
||||
|
||||
Convenience wrapper calling `stop()` followed by `start()`.
|
||||
"""
|
||||
logger.info("Restarting LLMProxy server...")
|
||||
if self.is_running():
|
||||
self.stop()
|
||||
await self.stop()
|
||||
if _port is not None:
|
||||
self.port = _port
|
||||
self.start()
|
||||
self.server_launcher_args.port = _port
|
||||
await self.start()
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Return whether the uvicorn server is active.
|
||||
@@ -745,7 +767,7 @@ class LLMProxy:
|
||||
Returns:
|
||||
bool: True if server was started and did not signal exit.
|
||||
"""
|
||||
return self._uvicorn_server is not None and self._uvicorn_server.started
|
||||
return self.server_launcher.is_running()
|
||||
|
||||
def as_resource(
|
||||
self,
|
||||
@@ -788,13 +810,13 @@ class LLMProxy:
|
||||
|
||||
if rollout_id is None and attempt_id is None:
|
||||
return ProxyLLM(
|
||||
endpoint=f"http://{self.host}:{self.port}",
|
||||
endpoint=self.server_launcher.access_endpoint,
|
||||
model=model,
|
||||
sampling_parameters=dict(sampling_parameters or {}),
|
||||
)
|
||||
elif rollout_id is not None and attempt_id is not None:
|
||||
return LLM(
|
||||
endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
|
||||
endpoint=f"{self.server_launcher.access_endpoint}/rollout/{rollout_id}/attempt/{attempt_id}",
|
||||
model=model,
|
||||
sampling_parameters=dict(sampling_parameters or {}),
|
||||
)
|
||||
@@ -827,13 +849,17 @@ def set_active_llm_proxy(proxy: LLMProxy) -> None:
|
||||
_global_llm_proxy = proxy
|
||||
|
||||
|
||||
def initialize_llm_callbacks() -> bool:
|
||||
def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
|
||||
"""Restore `litellm.callbacks` to a state that is just initialized by agent-lightning.
|
||||
|
||||
When litellm is restarted multiple times in the same process, more and more callbacks
|
||||
will be appended to `litellm.callbacks`, which may exceed the MAX_CALLBACKS limit.
|
||||
This function remembers the initial state of `litellm.callbacks` and always restore to that state.
|
||||
|
||||
Args:
|
||||
_add_return_token_ids: Whether to add the return token ids callback. Internal use only.
|
||||
Ideally the callback should automatically be enabled when the backend supports it.
|
||||
|
||||
Returns:
|
||||
Whether the callbacks are initialized for the first time.
|
||||
"""
|
||||
@@ -845,6 +871,10 @@ def initialize_llm_callbacks() -> bool:
|
||||
AddReturnTokenIds(),
|
||||
LightningOpenTelemetry(),
|
||||
]
|
||||
if _add_return_token_ids
|
||||
else [
|
||||
LightningOpenTelemetry(),
|
||||
]
|
||||
)
|
||||
_callbacks_before_litellm_start = [*litellm.callbacks] # type: ignore
|
||||
return True
|
||||
@@ -867,35 +897,6 @@ def initialize_llm_callbacks() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _get_default_ipv4_address() -> str:
|
||||
"""Determine the default outbound IPv4 address for this machine.
|
||||
|
||||
Implementation:
|
||||
Opens a UDP socket and "connects" to a public address to force route
|
||||
selection, then inspects the socket's local address. No packets are sent.
|
||||
|
||||
Returns:
|
||||
str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
|
||||
"""
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
# Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
|
||||
s.connect(("8.8.8.8", 80))
|
||||
return s.getsockname()[0]
|
||||
except Exception:
|
||||
return "127.0.0.1"
|
||||
finally:
|
||||
s.close()
|
||||
|
||||
|
||||
def _check_port(host: str, port: int) -> bool:
|
||||
"""Check if a port is available."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.settimeout(1)
|
||||
result = s.connect_ex((host, port))
|
||||
return result != 0 # True if unavailable
|
||||
|
||||
|
||||
def _check_tracer_provider() -> bool:
|
||||
"""Check if the global tracer provider is properly initialized.
|
||||
|
||||
|
||||
@@ -142,7 +142,13 @@ class ServerDataStore:
|
||||
async with self._resources_lock:
|
||||
resources = self._resource_versions.get(resources_id)
|
||||
if resources:
|
||||
return ResourcesUpdate(resources_id=resources_id, resources=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]:
|
||||
@@ -357,7 +363,9 @@ class AgentLightningServer:
|
||||
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)
|
||||
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
|
||||
|
||||
|
||||
@@ -1,14 +1,17 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .base import LightningStore
|
||||
from .base import LightningStore, LightningStoreCapabilities
|
||||
from .client_server import LightningStoreClient, LightningStoreServer
|
||||
from .database import SqlLightningStore
|
||||
from .memory import InMemoryLightningStore
|
||||
from .threading import LightningStoreThreaded
|
||||
|
||||
__all__ = [
|
||||
"LightningStore",
|
||||
"LightningStoreCapabilities",
|
||||
"LightningStoreClient",
|
||||
"LightningStoreServer",
|
||||
"InMemoryLightningStore",
|
||||
"LightningStoreThreaded",
|
||||
"SqlLightningStore",
|
||||
]
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Literal, Optional, Sequence
|
||||
from typing import Any, Dict, List, Literal, Optional, Sequence, Union, TypedDict
|
||||
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
|
||||
@@ -52,6 +52,17 @@ UNSET = _UnsetType()
|
||||
Unset = _UnsetType # Alias for convenience
|
||||
|
||||
|
||||
class LightningStoreCapabilities(TypedDict):
|
||||
"""Capability of a LightningStore implementation."""
|
||||
|
||||
thread_safe: bool
|
||||
"""Whether the store is thread-safe."""
|
||||
async_safe: bool
|
||||
"""Whether the store is async-safe."""
|
||||
zero_copy: bool
|
||||
"""Whether the store has only one copy across all threads/processes."""
|
||||
|
||||
|
||||
class LightningStore:
|
||||
"""Contract for the persistent control-plane that coordinates training rollouts.
|
||||
|
||||
@@ -74,6 +85,14 @@ class LightningStore:
|
||||
Unless stated otherwise, missing identifiers should result in a `ValueError`.
|
||||
"""
|
||||
|
||||
def capabilities(self) -> LightningStoreCapabilities:
|
||||
"""Return the capabilities of the store."""
|
||||
return LightningStoreCapabilities(
|
||||
thread_safe=False,
|
||||
async_safe=False,
|
||||
zero_copy=False,
|
||||
)
|
||||
|
||||
async def start_rollout(
|
||||
self,
|
||||
input: TaskInput,
|
||||
@@ -248,7 +267,7 @@ class LightningStore:
|
||||
|
||||
async def query_rollouts(
|
||||
self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
|
||||
) -> List[Rollout]:
|
||||
) -> List[Union[Rollout, AttemptedRollout]]:
|
||||
"""Retrieve rollouts filtered by status and/or explicit identifiers.
|
||||
|
||||
Args:
|
||||
@@ -278,7 +297,7 @@ class LightningStore:
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
|
||||
async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
|
||||
"""Fetch a rollout by identifier without mutating its state.
|
||||
|
||||
Args:
|
||||
@@ -307,6 +326,17 @@ class LightningStore:
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def query_resources(self) -> List[ResourcesUpdate]:
|
||||
"""List every stored resource snapshot in insertion order.
|
||||
|
||||
Returns:
|
||||
A chronological list of [`ResourcesUpdate`][agentlightning.ResourcesUpdate] objects.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must implement retrieval.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
|
||||
"""Return a specific named resource snapshot by identifier.
|
||||
|
||||
@@ -427,6 +457,8 @@ class LightningStore:
|
||||
This API is typically used by algorithms that maintain mutable resources (e.g., model
|
||||
checkpoints) under a stable identifier.
|
||||
|
||||
If `resources_id` does not exist, implementations should add it as a new snapshot.
|
||||
|
||||
Args:
|
||||
resources_id: Identifier of the snapshot to replace.
|
||||
resources: Updated mapping of resource names to payloads.
|
||||
@@ -436,7 +468,6 @@ class LightningStore:
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Subclasses must implement resource persistence.
|
||||
ValueError: Implementations must raise when `resources_id` does not exist.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .sqlite import SqlLightningStore
|
||||
|
||||
__all__ = [
|
||||
"SqlLightningStore",
|
||||
]
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from .attempt import AttemptInDB, SpanSeqIdInDB
|
||||
from .base import (
|
||||
AttemptStatusUpdateMessage,
|
||||
SqlAlchemyBase,
|
||||
)
|
||||
from .resources import ResourcesUpdateInDB
|
||||
from .rollout import RolloutInDB
|
||||
from .span import SpanInDB
|
||||
|
||||
__all__ = [
|
||||
"SqlAlchemyBase",
|
||||
"AttemptStatusUpdateMessage",
|
||||
"RolloutInDB",
|
||||
"AttemptInDB",
|
||||
"ResourcesUpdateInDB",
|
||||
"SpanSeqIdInDB",
|
||||
"SpanInDB",
|
||||
]
|
||||
@@ -0,0 +1,251 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import InitVar
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import JSON, Float, Integer, String, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from agentlightning.types import Attempt
|
||||
|
||||
from .base import AttemptStatusUpdateMessage, SqlAlchemyBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _generate_attempt_id() -> str:
|
||||
"""We don't need that long because attempts are limited to rollouts."""
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
|
||||
return "at-" + short_id
|
||||
|
||||
|
||||
class AttemptInDB(SqlAlchemyBase):
|
||||
__tablename__ = "attempts"
|
||||
|
||||
rollout_id: Mapped[str] = mapped_column(String, nullable=False)
|
||||
attempt_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_attempt_id)
|
||||
sequence_id: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
||||
start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
|
||||
end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
|
||||
status: Mapped[str] = mapped_column(String, default="preparing", nullable=False)
|
||||
worker_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
|
||||
last_heartbeat_time: Mapped[Optional[float]] = mapped_column(Float, nullable=False, default_factory=time.time)
|
||||
attempt_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
|
||||
|
||||
# addition columns for processing
|
||||
max_duration: Mapped[Optional[float]] = mapped_column(
|
||||
Float, nullable=True, default=None
|
||||
) # maximum duration allowed for this attempt in seconds
|
||||
max_heartbeat_interval: Mapped[Optional[float]] = mapped_column(
|
||||
Float, nullable=True, default=None
|
||||
) # maximum allowed heartbeat interval in seconds
|
||||
|
||||
version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
|
||||
__mapper_args__ = {
|
||||
"version_id_col": version_id,
|
||||
}
|
||||
|
||||
def is_unresponsive(self, current_time: float) -> bool:
|
||||
"""Check if the attempt is unresponsive based on the last heartbeat time and max_heartbeat_interval."""
|
||||
if self.max_heartbeat_interval is None:
|
||||
return False
|
||||
if self.last_heartbeat_time is None:
|
||||
return False
|
||||
return (current_time - self.last_heartbeat_time) > self.max_heartbeat_interval
|
||||
|
||||
def is_timed_out(self, current_time: float) -> bool:
|
||||
"""Check if the attempt has timed out based on the start time and max_duration."""
|
||||
if self.max_duration is None:
|
||||
return False
|
||||
return (current_time - self.start_time) > self.max_duration
|
||||
|
||||
def as_attempt(self) -> Attempt:
|
||||
return Attempt(
|
||||
**self.model_dump(
|
||||
exclude={"max_duration", "max_heartbeat_interval", "version_id"},
|
||||
mapper={"metadata": lambda obj: obj.attempt_metadata}, # type: ignore
|
||||
)
|
||||
)
|
||||
|
||||
def _validate_status_message(self, msg: Dict[str, Any]) -> None:
|
||||
"""This function validates the status update message from caller.
|
||||
Raises ValueError if the message is invalid.
|
||||
"""
|
||||
if "event" not in msg:
|
||||
raise ValueError("Status update message must contain 'event' field.")
|
||||
if "timestamp" not in msg:
|
||||
msg["timestamp"] = time.time()
|
||||
if msg["event"] not in [
|
||||
"user_update", # user update attempt status via dbstore.update_attempt()
|
||||
"span_received", # new span received
|
||||
"single_step_timeout", # single step timeout detected (from last span heartbeat)
|
||||
"overall_timeout", # overall timeout detected
|
||||
]:
|
||||
raise ValueError(f"Unsupported event type: {msg['event']}")
|
||||
if msg["event"] == "user_update" and "new_status" not in msg:
|
||||
raise ValueError("User update event must contain 'new_status' field.")
|
||||
|
||||
def get_finished_statuses(self) -> List[str]:
|
||||
"""This function returns the list of statuses that are considered finished."""
|
||||
return [
|
||||
"succeeded",
|
||||
"failed",
|
||||
"timeout",
|
||||
]
|
||||
|
||||
def update_status(self, msg: Dict[str, Any]) -> Optional[AttemptStatusUpdateMessage]:
|
||||
"""This function updates the status of the attempt based on the event.
|
||||
Args:
|
||||
msg: A dictionary containing the status update message. It must contain an "event" field, and optionally a "new_status" field.
|
||||
More details about the message format can be found in the `_validate_status_message`() method.
|
||||
current_time: The current time to use for updating timestamps. If None, uses time.time().
|
||||
Returns:
|
||||
A dictionary containing the status update message: {"event": "attempt_status_updated", "old_status": old_status, "new_status": new_status}.
|
||||
IF no meaningful status update is performed, returns None.
|
||||
Raises:
|
||||
ValueError: If the event is not recognized or the status transition is invalid.
|
||||
NotImplementedError: If the event handling is not implemented for the current status.
|
||||
RuntimeError: If the new status is not set after processing the event.
|
||||
"""
|
||||
self._validate_status_message(msg)
|
||||
event = msg["event"]
|
||||
current_time = msg.get("timestamp", time.time())
|
||||
old_status = self.status
|
||||
new_status = msg.get("new_status", None)
|
||||
|
||||
# Step 1: Determine the new status based on the event and current status
|
||||
if event == "user_update":
|
||||
if not new_status:
|
||||
raise ValueError("new_status must be provided for user_update event.")
|
||||
elif event == "span_received":
|
||||
self.last_heartbeat_time = current_time
|
||||
if old_status in ["preparing", "unresponsive", "running"]:
|
||||
new_status = "running"
|
||||
elif old_status in self.get_finished_statuses():
|
||||
logger.warning(
|
||||
f"Span received after attempt is already in status {self.status}. No status update performed."
|
||||
)
|
||||
return # no further status update needed
|
||||
else:
|
||||
raise NotImplementedError(f"Event {event} is not implemented for status {old_status}.")
|
||||
elif event == "single_step_timeout":
|
||||
if old_status in [
|
||||
"preparing",
|
||||
"running",
|
||||
]:
|
||||
new_status = "unresponsive"
|
||||
else:
|
||||
logger.warning(
|
||||
f"Single step timeout detected but attempt is in status {self.status}. No status update performed."
|
||||
)
|
||||
return # no further status update needed
|
||||
elif event == "overall_timeout":
|
||||
if old_status not in self.get_finished_statuses():
|
||||
new_status = "timeout"
|
||||
else:
|
||||
logger.warning(
|
||||
f"Overall timeout detected but attempt is in status {self.status}. No status update performed."
|
||||
)
|
||||
return # no further status update needed
|
||||
else:
|
||||
raise NotImplementedError(f"Event {event} is not implemented for status update.")
|
||||
|
||||
# Step 2: Update the status
|
||||
if not new_status:
|
||||
raise RuntimeError(
|
||||
f"new_status should not be {new_status} after processing event for {event} on status {old_status}."
|
||||
)
|
||||
if new_status == old_status:
|
||||
return # no status change
|
||||
if new_status in self.get_finished_statuses():
|
||||
# when attempt is finished, set end_time
|
||||
self.end_time = current_time
|
||||
self.status = new_status
|
||||
|
||||
# Step 3: Return the status update info for further processing
|
||||
return AttemptStatusUpdateMessage(
|
||||
attempt_id=self.attempt_id,
|
||||
rollout_id=self.rollout_id,
|
||||
timestamp=current_time,
|
||||
old_status=old_status,
|
||||
new_status=new_status,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def get_latest_attempt_for_rollout(
|
||||
cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
|
||||
) -> Optional[Attempt]:
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
result = await session.scalars(
|
||||
select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.desc()).limit(1)
|
||||
)
|
||||
attempt_obj = result.one_or_none()
|
||||
if attempt_obj is None:
|
||||
return None
|
||||
return attempt_obj.as_attempt()
|
||||
|
||||
@classmethod
|
||||
async def get_attempts_for_rollout(
|
||||
cls: type[AttemptInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
|
||||
) -> List[Attempt]:
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
result = await session.scalars(
|
||||
select(cls).where(cls.rollout_id == rollout_id).order_by(cls.sequence_id.asc())
|
||||
)
|
||||
return [attempt.as_attempt() for attempt in result.all()]
|
||||
|
||||
|
||||
class SpanSeqIdInDB(SqlAlchemyBase):
|
||||
__tablename__ = "span_sequence"
|
||||
|
||||
rollout_id: Mapped[str] = mapped_column(nullable=False, primary_key=True)
|
||||
|
||||
# FIXME InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
|
||||
# attempt_id: Mapped[str] = mapped_column(nullable=False)
|
||||
attempt_id: InitVar[str] # not mapped column, just for type hinting
|
||||
|
||||
current_sequence: Mapped[int] = mapped_column(default=1, nullable=False)
|
||||
|
||||
# Versioning for optimistic concurrency control
|
||||
version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
|
||||
__mapper_args__ = {
|
||||
"version_id_col": version_id,
|
||||
# "primary_key": [rollout_id, attempt_id],
|
||||
# "primary_key": [rollout_id],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
async def get_next_sequence_id(
|
||||
cls: type[SpanSeqIdInDB],
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
rollout_id: str,
|
||||
attempt_id: str,
|
||||
external_seq_id: Optional[int] = None,
|
||||
) -> int:
|
||||
"""Get the next sequence ID with retries to handle race conditions.
|
||||
IF external_seq_id is provided and is greater than current_sequence, set current_sequence to external_seq_id.
|
||||
"""
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
seq_obj = await session.get(cls, rollout_id)
|
||||
# seq_obj = await session.get(cls, [rollout_id, attempt_id])
|
||||
if seq_obj is None:
|
||||
raise ValueError(f"Rollout {rollout_id} not found")
|
||||
else:
|
||||
current_seq = (
|
||||
external_seq_id
|
||||
if external_seq_id is not None and external_seq_id > seq_obj.current_sequence
|
||||
else seq_obj.current_sequence
|
||||
)
|
||||
seq_obj.current_sequence = current_seq + 1
|
||||
await session.flush()
|
||||
return current_seq
|
||||
@@ -0,0 +1,186 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field, TypeAdapter, computed_field
|
||||
|
||||
# from dataclasses import asdict
|
||||
from sqlalchemy import JSON, TypeDecorator
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs
|
||||
from sqlalchemy.orm import DeclarativeBase, MappedAsDataclass
|
||||
|
||||
|
||||
class SqlAlchemyBase(AsyncAttrs, MappedAsDataclass, DeclarativeBase):
|
||||
pass
|
||||
|
||||
def model_dump(
|
||||
self,
|
||||
exclude: set[str] | None = None,
|
||||
mapper: Dict[str, Callable[["SqlAlchemyBase"], Any]] | None = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Dump the SQLAlchemy model to a dictionary.
|
||||
Args:
|
||||
exclude: set[str]
|
||||
The set of field names to exclude.
|
||||
mapper: Dict[str, Callable[[SqlAlchemyBase], Any]]
|
||||
A mapping from field names to functions that take the model instance and return the value to be used for that field.
|
||||
If the key is "*", the function should return a dictionary of additional fields to be added to the output.
|
||||
Returns:
|
||||
Dict[str, Any]: The dumped model as a dictionary.
|
||||
"""
|
||||
exclude = exclude or set()
|
||||
mapper = mapper or {}
|
||||
dic = {k: getattr(self, k) for k in self.__table__.columns.keys() if k not in exclude}
|
||||
for k, func in mapper.items():
|
||||
if k == "*":
|
||||
dic.update(func(self))
|
||||
else:
|
||||
dic[k] = func(self)
|
||||
return dic
|
||||
|
||||
|
||||
class PydanticInDB(TypeDecorator[BaseModel]):
|
||||
"""Custom SQLAlchemy type to store pydantic.BaseModel as JSON in the database.
|
||||
Attributes:
|
||||
target_type: type[BaseModel], the type of the pydantic model to be stored.
|
||||
"""
|
||||
|
||||
impl = JSON
|
||||
target_type: type[BaseModel] | None = None
|
||||
|
||||
def process_bind_param(self, value: BaseModel | None, dialect: Any) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
if self.target_type is not None:
|
||||
return TypeAdapter(self.target_type).validate_python(value).model_dump_json() # type: ignore
|
||||
return json.dumps(value)
|
||||
|
||||
def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[BaseModel]:
|
||||
if value is None:
|
||||
return None
|
||||
if self.target_type is not None:
|
||||
return TypeAdapter(self.target_type).validate_json(value) # type: ignore
|
||||
dic = json.loads(value)
|
||||
return dic # type: ignore
|
||||
|
||||
|
||||
class PydanticListInDB(TypeDecorator[list[BaseModel]]):
|
||||
"""Custom SQLAlchemy type to store List[pydantic.BaseModel] as JSON in the database.
|
||||
Attributes:
|
||||
value_type: type[BaseModel], the type of the pydantic model to be stored in the list.
|
||||
"""
|
||||
|
||||
impl = JSON
|
||||
value_type: type[BaseModel] | None = None
|
||||
|
||||
def process_bind_param(self, value: List[BaseModel] | None, dialect: Any) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
if self.value_type is not None:
|
||||
lst = [TypeAdapter(self.value_type).validate_python(v).model_dump() for v in value]
|
||||
return json.dumps(lst)
|
||||
raise ValueError("target_type must be set for PydanticListInDB")
|
||||
|
||||
def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[List[BaseModel]]:
|
||||
if value is None:
|
||||
return None
|
||||
if self.value_type is not None:
|
||||
dic = json.loads(value)
|
||||
return [TypeAdapter(self.value_type).validate_python(v) for v in dic] # type: ignore
|
||||
raise ValueError("target_type must be set for PydanticListInDB")
|
||||
|
||||
|
||||
class NamedDictBase(TypeDecorator[Dict[str, Any]]):
|
||||
"""Custom SQLAlchemy type to store Dict[str, pydantic.BaseModel] as JSON in the database.
|
||||
Attributes:
|
||||
target_alias: type[Dict[str, BaseModel]], the alias type of the dict.
|
||||
value_type: type[BaseModel], the type of the values in the dict.
|
||||
|
||||
For example, given NamedResources = Dict[str, ResourceUnion],
|
||||
we can define NamedDictBase with target_alias=NamedResources and target_type=ResourceUnion.
|
||||
"""
|
||||
|
||||
impl = JSON
|
||||
target_alias: type | None = None
|
||||
value_type: type[BaseModel] | Any = None
|
||||
|
||||
def process_bind_param(self, value: Dict[str, Any] | None, dialect: Any) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
# ignore target_alias for when dumping because Dict is not a pydantic model
|
||||
if self.value_type is not None:
|
||||
dic = {
|
||||
k: TypeAdapter(self.value_type).validate_python(v).model_dump() if isinstance(v, BaseModel) else v
|
||||
for k, v in value.items()
|
||||
}
|
||||
return json.dumps(dic)
|
||||
dic = {k: v.model_dump() if isinstance(v, BaseModel) else v for k, v in value.items()}
|
||||
return json.dumps(dic)
|
||||
|
||||
def process_result_value(self, value: Optional[str], dialect: Any) -> Optional[Dict[str, Any]]:
|
||||
if value is None:
|
||||
return None
|
||||
if self.target_alias is not None:
|
||||
return TypeAdapter(self.target_alias).validate_json(value) # type: ignore
|
||||
if self.value_type is not None:
|
||||
dic = json.loads(value)
|
||||
return {k: TypeAdapter(self.value_type).validate_python(v) for k, v in dic.items()} # type: ignore
|
||||
return json.loads(value)
|
||||
|
||||
|
||||
class DatabaseRuntimeError(Exception):
|
||||
"""Raised when a runtime error occurs during database operations.
|
||||
Particularly used when the execution of a query fails.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class RaceConditionError(Exception):
|
||||
"""Raised when a race condition is detected during database operations."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class NoRolloutToDequeueError(Exception):
|
||||
"""Raised when there is no rollout available to dequeue."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AttemptStatusUpdateMessage(BaseModel):
|
||||
attempt_id: str
|
||||
rollout_id: str
|
||||
timestamp: float = Field(default_factory=time.time)
|
||||
old_status: Optional[str] = None
|
||||
new_status: str
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def event(self) -> str:
|
||||
return "attempt_status_update"
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def is_failed(self) -> bool:
|
||||
return self.new_status in ["failed", "timeout", "unresponsive"]
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def is_succeeded(self) -> bool:
|
||||
return self.new_status == "succeeded"
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def is_finished(self) -> bool:
|
||||
return self.is_failed or self.is_succeeded
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self.new_status in ["running", "preparing"]
|
||||
@@ -0,0 +1,55 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
import uuid
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from agentlightning.types import NamedResources, ResourcesUpdate
|
||||
|
||||
from .base import NamedDictBase, SqlAlchemyBase
|
||||
|
||||
|
||||
def _generate_resources_id() -> str:
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
|
||||
return "rs-" + short_id
|
||||
|
||||
|
||||
class NamedResourcesInDB(NamedDictBase):
|
||||
"""Custom SQLAlchemy type to store NamedResources as JSON in the database."""
|
||||
|
||||
target_alias = NamedResources
|
||||
|
||||
|
||||
class ResourcesUpdateInDB(SqlAlchemyBase):
|
||||
__tablename__ = "resources"
|
||||
resources: Mapped[NamedResources] = mapped_column(
|
||||
NamedResourcesInDB, nullable=False
|
||||
) # JSON serialized, convert to NamedResources when needed
|
||||
resources_id: Mapped[str] = mapped_column(primary_key=True, default_factory=_generate_resources_id)
|
||||
create_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time)
|
||||
update_time: Mapped[float] = mapped_column(nullable=False, default_factory=time.time, onupdate=time.time)
|
||||
version: Mapped[int] = mapped_column(nullable=False, default=1)
|
||||
|
||||
__mapper_args__ = {
|
||||
"version_id_col": version,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
async def get_resources_by_id(
|
||||
cls, session_factory: async_sessionmaker[AsyncSession], resources_id: str
|
||||
) -> Optional[ResourcesUpdate]:
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
obj = await session.get(cls, resources_id)
|
||||
if obj is None:
|
||||
return None
|
||||
return obj.as_resources_update()
|
||||
|
||||
def as_resources_update(self) -> ResourcesUpdate:
|
||||
return ResourcesUpdate(**self.model_dump())
|
||||
@@ -0,0 +1,201 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
|
||||
from sqlalchemy import JSON, Float, Integer, String, and_, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from agentlightning.types import AttemptedRollout, Rollout, RolloutConfig, RolloutStatus
|
||||
|
||||
from ...base import is_finished, is_queuing
|
||||
from .attempt import AttemptInDB
|
||||
from .base import AttemptStatusUpdateMessage, PydanticInDB, SqlAlchemyBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _generate_rollout_id() -> str:
|
||||
short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
|
||||
return "ro-" + short_id
|
||||
|
||||
|
||||
class RolloutConfigInDB(PydanticInDB):
|
||||
"""Custom SQLAlchemy type to store RolloutConfig as JSON in the database."""
|
||||
|
||||
target_type = RolloutConfig
|
||||
|
||||
|
||||
class RolloutInDB(SqlAlchemyBase):
|
||||
__tablename__ = "rollouts"
|
||||
|
||||
input: Mapped[Any] = mapped_column(JSON, nullable=False)
|
||||
rollout_id: Mapped[str] = mapped_column(String, primary_key=True, default_factory=_generate_rollout_id)
|
||||
start_time: Mapped[float] = mapped_column(Float, default_factory=time.time, nullable=False)
|
||||
end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True, default=None)
|
||||
mode: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
|
||||
resources_id: Mapped[Optional[str]] = mapped_column(String, nullable=True, default=None)
|
||||
status: Mapped[RolloutStatus] = mapped_column(String, default="queuing", nullable=False)
|
||||
config: Mapped[RolloutConfig] = mapped_column(
|
||||
RolloutConfigInDB, nullable=False, default_factory=RolloutConfig
|
||||
) # JSON serialized, convert to RolloutConfig when needed
|
||||
rollout_metadata: Mapped[Optional[Dict[str, Any]]] = mapped_column(
|
||||
JSON, nullable=True, default=None
|
||||
) # JSON serialized, convert to Dict when needed
|
||||
|
||||
# Attempt-related helper methods can be added here if needed
|
||||
num_attempts: Mapped[int] = mapped_column(
|
||||
Integer, default=0, nullable=False
|
||||
) # number of attempts made for this rollout
|
||||
enqueue_time: Mapped[Optional[float]] = mapped_column(
|
||||
Float, nullable=True, default_factory=time.time
|
||||
) # time when the rollout was enqueued (for FIFO scheduling)
|
||||
latest_attempt_id: Mapped[Optional[str]] = mapped_column(
|
||||
String, nullable=True, default=None
|
||||
) # the attempt_id of the latest attempt
|
||||
|
||||
# use optimistic concurrency control
|
||||
version_id: Mapped[int] = mapped_column(Integer, nullable=False, default=1)
|
||||
__mapper_args__ = {
|
||||
"version_id_col": version_id,
|
||||
}
|
||||
|
||||
def __post_init__(self):
|
||||
if self.status not in ["queuing", "running", "succeeded", "failed", "requeuing"]:
|
||||
raise ValueError(f"Invalid rollout status: {self.status}")
|
||||
|
||||
def as_rollout(self) -> Rollout:
|
||||
return Rollout(
|
||||
**self.model_dump(
|
||||
exclude={"rollout_metadata", "num_attempts", "enqueue_time", "latest_attempt_id", "version_id"},
|
||||
mapper={
|
||||
"metadata": lambda obj: obj.rollout_metadata, # type: ignore
|
||||
"config": lambda obj: obj.config if obj.config is not None else RolloutConfig(), # type: ignore
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def _validate_status_message(self, msg: Dict[str, str]) -> None:
|
||||
"""Validate the status update message.
|
||||
Raises:
|
||||
ValueError: If the message is invalid.
|
||||
"""
|
||||
if "event" not in msg:
|
||||
raise ValueError("Status update message must contain 'event' field.")
|
||||
event = msg["event"]
|
||||
if event not in [
|
||||
"attempt_status_update", # from attempt status update
|
||||
"user_update", # from user-initiated update
|
||||
]:
|
||||
raise ValueError(f"Invalid event type in status update message: {event}")
|
||||
if event == "user_update":
|
||||
if "new_status" not in msg:
|
||||
raise ValueError("Status update message for event 'user_update' must contain 'new_status' field.")
|
||||
if event == "attempt_status_update":
|
||||
# leverage AttemptStatusUpdateMessage for validation
|
||||
pass
|
||||
|
||||
async def update_status(self, msg: Dict[str, Any] | AttemptStatusUpdateMessage) -> None:
|
||||
"""Update the rollout status based on the provided message.
|
||||
Args:
|
||||
msg (Dict[str, str]): The status update message. Refer to `_validate_status_message` for the expected format.
|
||||
current_time (Optional[float]): The current time to set end_time or enqueue_time if needed.
|
||||
"""
|
||||
if isinstance(msg, dict):
|
||||
self._validate_status_message(msg)
|
||||
event = msg["event"]
|
||||
current_time = msg.get("timestamp", time.time())
|
||||
else:
|
||||
event = msg.event
|
||||
current_time = msg.timestamp
|
||||
|
||||
old_status = self.status
|
||||
new_status = self.status # initialize new_status with old_status
|
||||
|
||||
# Step 1: Determine the new status based on the event
|
||||
if event == "user_update":
|
||||
assert isinstance(msg, dict)
|
||||
new_status = msg["new_status"]
|
||||
elif event == "attempt_status_update":
|
||||
msg = AttemptStatusUpdateMessage(**msg) if isinstance(msg, dict) else msg
|
||||
if msg.attempt_id == self.latest_attempt_id:
|
||||
new_status = msg.new_status # directly take the latest attempt status
|
||||
if msg.is_succeeded:
|
||||
new_status = "succeeded"
|
||||
elif msg.is_failed:
|
||||
# no other attempts running, decide whether to requeue or fail
|
||||
config = self.config
|
||||
if config.max_attempts > self.num_attempts and msg.new_status in config.retry_condition:
|
||||
new_status = "requeuing"
|
||||
else:
|
||||
new_status = "failed"
|
||||
# elif msg.is_running and old_status in ["failed", "requeuing"]:
|
||||
# new_status = "running"
|
||||
else:
|
||||
# ignore attempts from old attempts
|
||||
new_status = old_status
|
||||
|
||||
# Step 2: Update the status if it has changed and handle follow-up actions
|
||||
if new_status is None:
|
||||
raise RuntimeError(
|
||||
f"New status of `{old_status}` and `{self.latest_attempt_id}` could not be determined from the message {msg}."
|
||||
)
|
||||
if new_status == old_status:
|
||||
return
|
||||
self.status = cast(RolloutStatus, new_status)
|
||||
|
||||
if is_finished(self): # type: ignore
|
||||
self.end_time = current_time
|
||||
if is_queuing(self): # type: ignore
|
||||
self.enqueue_time = current_time
|
||||
# When requeuing, we do not reset latest_attempt_id or num_attempts,
|
||||
# as they should persist across requeues.
|
||||
|
||||
@classmethod
|
||||
async def get_rollout_by_id(
|
||||
cls: type[RolloutInDB], session_factory: async_sessionmaker[AsyncSession], rollout_id: str
|
||||
) -> Optional[Rollout | AttemptedRollout]:
|
||||
"""Query a specific rollout from the database."""
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = await session.get(cls, rollout_id)
|
||||
if rollout_obj is None:
|
||||
return None
|
||||
if rollout_obj.latest_attempt_id is not None:
|
||||
attempt_obj = await session.get(AttemptInDB, rollout_obj.latest_attempt_id)
|
||||
if attempt_obj is not None:
|
||||
return AttemptedRollout(
|
||||
**rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt()
|
||||
)
|
||||
return rollout_obj.as_rollout()
|
||||
|
||||
@classmethod
|
||||
async def query_rollouts(
|
||||
cls: type[RolloutInDB],
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
*,
|
||||
statuses: Optional[List[str]] = None,
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> List[RolloutInDB]:
|
||||
"""
|
||||
Query rollouts from the database with optional filters.
|
||||
"""
|
||||
async with session_factory() as session:
|
||||
async with session.begin():
|
||||
conditions: list[Any] = []
|
||||
if statuses is not None:
|
||||
conditions.append(cls.status.in_(statuses))
|
||||
if ids is not None:
|
||||
conditions.append(cls.rollout_id.in_(ids))
|
||||
query = select(cls)
|
||||
if conditions:
|
||||
query = query.where(and_(*conditions))
|
||||
result = await session.scalars(query)
|
||||
rollout_objs = result.all()
|
||||
return list(rollout_objs)
|
||||
@@ -0,0 +1,101 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import JSON, Float, Integer, String
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from agentlightning.types.tracer import (
|
||||
Attributes,
|
||||
AttributeValue,
|
||||
Event,
|
||||
Link,
|
||||
OtelResource,
|
||||
Span,
|
||||
SpanContext,
|
||||
TraceStatus,
|
||||
)
|
||||
|
||||
from .base import NamedDictBase, PydanticInDB, PydanticListInDB, SqlAlchemyBase
|
||||
|
||||
|
||||
class TraceStatusInDB(PydanticInDB):
|
||||
target_type = TraceStatus
|
||||
|
||||
|
||||
class AttributesInDB(NamedDictBase):
|
||||
target_alias = None # type: ignore
|
||||
value_type = AttributeValue
|
||||
|
||||
|
||||
class EventListInDB(PydanticListInDB):
|
||||
value_type = Event
|
||||
|
||||
|
||||
class LinkListInDB(PydanticListInDB):
|
||||
value_type = Link
|
||||
|
||||
|
||||
class SpanContextInDB(PydanticInDB):
|
||||
target_type = SpanContext
|
||||
|
||||
|
||||
class OtelResourceInDB(PydanticInDB):
|
||||
target_type = OtelResource
|
||||
|
||||
|
||||
class SpanInDB(SqlAlchemyBase):
|
||||
__tablename__ = "spans"
|
||||
|
||||
rollout_id: Mapped[str] = mapped_column(String, nullable=False) # The rollout which this span belongs to.
|
||||
attempt_id: Mapped[str] = mapped_column(String, nullable=False) # The attempt which this span belongs to.
|
||||
sequence_id: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False
|
||||
) # The ID to make spans ordered within a single attempt.
|
||||
|
||||
# Current ID (in hex, formatted via trace_api.format_*)
|
||||
trace_id: Mapped[str] = mapped_column(
|
||||
String, nullable=False
|
||||
) # one rollout can have traces coming from multiple places
|
||||
|
||||
# FIXME: span_id may be not unique across different attempts/rollouts, use (rollout_id, attempt_id, sequence_id) as the primary key instead
|
||||
span_id: Mapped[str] = mapped_column(
|
||||
String, nullable=False
|
||||
) # The span ID of the span. This ID comes from the OpenTelemetry span ID generator.
|
||||
parent_id: Mapped[Optional[str]] = mapped_column(String, nullable=True) # The parent span ID of the span.
|
||||
|
||||
# Core ReadableSpan fields
|
||||
name: Mapped[str] = mapped_column(String, nullable=False)
|
||||
status: Mapped[TraceStatus] = mapped_column(TraceStatusInDB, nullable=False)
|
||||
attributes: Mapped[Attributes] = mapped_column(AttributesInDB, nullable=False)
|
||||
events: Mapped[List[Event]] = mapped_column(EventListInDB, nullable=False)
|
||||
links: Mapped[List[Link]] = mapped_column(LinkListInDB, nullable=False)
|
||||
|
||||
# Timestamps
|
||||
start_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
|
||||
end_time: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
|
||||
|
||||
# Other parsable fields
|
||||
context: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
|
||||
parent: Mapped[Optional[SpanContext]] = mapped_column(SpanContextInDB, nullable=True)
|
||||
resource: Mapped[OtelResource] = mapped_column(OtelResourceInDB, nullable=False)
|
||||
|
||||
# extra fields can be added here as needed
|
||||
extra: Mapped[Optional[Dict[str, Any]]] = mapped_column(JSON, nullable=True, default=None)
|
||||
|
||||
__mapper_args__ = {
|
||||
"primary_key": [rollout_id, attempt_id, sequence_id],
|
||||
}
|
||||
|
||||
def as_span(self) -> Span:
|
||||
return Span(
|
||||
**self.model_dump(
|
||||
exclude={"extra"},
|
||||
mapper={"*": lambda obj: obj.extra or {}}, # type: ignore
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,316 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""This file contains a configurable async retry decorator based on exception type."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import importlib
|
||||
import logging
|
||||
import random
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Optional, Type, TypeVar
|
||||
|
||||
from tenacity import AsyncRetrying, RetryCallState, retry_if_exception
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Logging setup
|
||||
# ----------------------------------------------------------------------
|
||||
logger = logging.getLogger("async_retry")
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Type alias for async callable
|
||||
# ----------------------------------------------------------------------
|
||||
F = TypeVar("F", bound=Callable[..., Awaitable[Any]])
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Dataclass definition for retry configuration
|
||||
# ----------------------------------------------------------------------
|
||||
@dataclass
|
||||
class RetryStrategy:
|
||||
"""Configuration schema for retry behavior of a specific exception type.
|
||||
The wait time before $n$-th retry is calculated as ($n$ starts from 1):
|
||||
wait_time = wait_seconds * (backoff ** (n - 1)) * (1 + jitter * U(-1, 1))
|
||||
where U(-1, 1) is a uniform random variable between -1 and 1.
|
||||
Attributes:
|
||||
max_attempts: Maximum number of attempts before giving up. Default is 1 (no retry). None means infinite retries.
|
||||
max_retry_delay: Optional maximum delay between retries in seconds. Default is None (no limit).
|
||||
wait_seconds: Base wait time in seconds before the first retry. Default is 0.0.
|
||||
max_wait_seconds: Maximum wait time in seconds between retries. Default is None (no limit).
|
||||
backoff: Exponential backoff multiplier. Default is 1.0 (no backoff).
|
||||
jitter: Fractional (relative) jitter to apply to wait time. Default is 0.0 (no jitter).
|
||||
log: Whether to log each retry attempt. Default is False.
|
||||
"""
|
||||
|
||||
max_attempts: Optional[int] = 1
|
||||
max_retry_delay: Optional[float] = None
|
||||
wait_seconds: float = 0.0
|
||||
max_wait_seconds: Optional[float] = None
|
||||
backoff: float = 1.0
|
||||
jitter: float = 0.0
|
||||
log: bool = False
|
||||
|
||||
def asdict(self) -> Dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
def __post_init__(self):
|
||||
if self.max_attempts is not None and self.max_attempts < 1:
|
||||
raise ValueError("max_attempts must be at least 1 or None for infinite retries")
|
||||
if self.wait_seconds < 0.0:
|
||||
raise ValueError("wait_seconds must be non-negative")
|
||||
if self.backoff < 1.0:
|
||||
raise ValueError("backoff must be at least 1.0")
|
||||
if not (0.0 <= self.jitter <= 1.0):
|
||||
raise ValueError("jitter must be between 0.0 and 1.0")
|
||||
|
||||
def _get_wait_time(self, attempt_number: int) -> float:
|
||||
"""Calculate the wait time before the given attempt number."""
|
||||
base_wait = self.wait_seconds * (self.backoff ** (attempt_number - 1))
|
||||
if self.jitter > 0:
|
||||
delta = base_wait * self.jitter
|
||||
wait_time = random.uniform(base_wait - delta, base_wait + delta)
|
||||
else:
|
||||
wait_time = base_wait
|
||||
wait_time = max(wait_time, 0.0)
|
||||
if self.max_wait_seconds is not None:
|
||||
wait_time = min(wait_time, self.max_wait_seconds)
|
||||
return wait_time
|
||||
|
||||
def wait_func(self, retry_state: RetryCallState) -> float:
|
||||
"""Tenacity wait function based on the given strategy."""
|
||||
return self._get_wait_time(retry_state.attempt_number)
|
||||
|
||||
def stop_func(self, retry_state: RetryCallState) -> bool:
|
||||
"""Tenacity stop function based on the given strategy."""
|
||||
if self.max_attempts is not None:
|
||||
if retry_state.attempt_number >= self.max_attempts:
|
||||
return True
|
||||
if self.max_retry_delay is not None:
|
||||
time_since_start = retry_state.seconds_since_start
|
||||
if time_since_start is None:
|
||||
logger.warning("Cannot determine time since start for retry stop condition.")
|
||||
return False
|
||||
if time_since_start >= self.max_retry_delay:
|
||||
return True
|
||||
return False
|
||||
|
||||
async def before_sleep(self, retry_state: RetryCallState):
|
||||
"""Tenacity before_sleep callback to log retry attempts."""
|
||||
if self.log:
|
||||
exc = retry_state.outcome.exception() if retry_state.outcome else None
|
||||
next_wait = self.wait_func(retry_state)
|
||||
logger.warning(
|
||||
f"[Retry] {exc.__class__.__name__}: attempt={retry_state.attempt_number}, "
|
||||
f"next_wait={next_wait:.2f}s, message={exc}"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Exception Registry — shared, reusable, and extensible
|
||||
# ----------------------------------------------------------------------
|
||||
class ExceptionRegistry:
|
||||
"""
|
||||
Global registry for mapping string keys to Exception classes.
|
||||
Supports dynamic registration and fallback to importlib.
|
||||
"""
|
||||
|
||||
_registry: Dict[str, Type[BaseException]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str, exc_type: Type[BaseException] | None = None) -> None:
|
||||
"""Register an exception type under a given name."""
|
||||
if name in cls._registry:
|
||||
logger.warning(f"Overwriting existing exception registration for name '{name}'.")
|
||||
if exc_type is None:
|
||||
# Try to dynamically import the exception class
|
||||
try:
|
||||
module_name, class_name = name.rsplit(".", 1)
|
||||
module = importlib.import_module(module_name)
|
||||
exc_type = getattr(module, class_name)
|
||||
if exc_type is None:
|
||||
raise TypeError(f"{name} is not an Exception type.")
|
||||
except (ImportError, AttributeError, ValueError, TypeError) as e:
|
||||
raise ValueError(f"Cannot resolve exception type for name '{name}': {e}")
|
||||
cls._registry[name] = exc_type
|
||||
|
||||
@classmethod
|
||||
def all_registered(cls) -> Dict[str, Type[BaseException]]:
|
||||
"""Return the current registry mapping."""
|
||||
return dict(cls._registry)
|
||||
|
||||
@classmethod
|
||||
def clear(cls):
|
||||
"""Clear all registered exception mappings."""
|
||||
cls._registry.clear()
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Async Retry Decorator
|
||||
# ----------------------------------------------------------------------
|
||||
class AsyncTypeBasedRetry:
|
||||
"""
|
||||
A configurable async retry decorator based on exception type.
|
||||
|
||||
- Takes configuration as a Dict[str, RetryStrategy].
|
||||
- Provides `from_json()` for quick loading.
|
||||
- Uses a global ExceptionRegistry to resolve exception names.
|
||||
"""
|
||||
|
||||
def __init__(self, strategies: Dict[str, RetryStrategy], default_strategy: RetryStrategy | None = None):
|
||||
self.exception_map = self._build_exception_map(strategies)
|
||||
self.default_strategy = default_strategy or RetryStrategy()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Build exception map
|
||||
# ------------------------------------------------------------------
|
||||
def _build_exception_map(self, strategies: Dict[str, RetryStrategy]) -> Dict[Type[BaseException], RetryStrategy]:
|
||||
mapping: Dict[Type[BaseException], RetryStrategy] = {}
|
||||
all_registered = ExceptionRegistry.all_registered()
|
||||
for name, strat in strategies.items():
|
||||
if name in all_registered:
|
||||
exc_type = all_registered[name]
|
||||
else:
|
||||
raise ValueError(f"Exception type '{name}' is not registered in ExceptionRegistry.")
|
||||
mapping[exc_type] = strat
|
||||
return mapping
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Retry core logic
|
||||
# ------------------------------------------------------------------
|
||||
def get_exception(self, retry_state: RetryCallState) -> Optional[BaseException]:
|
||||
"""Get the exception from the given retry state, if any."""
|
||||
return retry_state.outcome.exception() if retry_state.outcome else None
|
||||
|
||||
def get_strategy(self, retry_state: RetryCallState) -> Optional[RetryStrategy]:
|
||||
"""Get the RetryStrategy for the exception in the given retry state.
|
||||
IF no matching exception type is found, return the default strategy.
|
||||
IF no exception is found, return None.
|
||||
"""
|
||||
exc = self.get_exception(retry_state)
|
||||
if exc is None:
|
||||
return None
|
||||
for exc_type, strat in self.exception_map.items():
|
||||
if isinstance(exc, exc_type):
|
||||
return strat
|
||||
return self.default_strategy
|
||||
|
||||
def should_retry(self, exc: BaseException) -> bool:
|
||||
return any(isinstance(exc, t) for t in self.exception_map.keys())
|
||||
|
||||
def wait_func(self, retry_state: RetryCallState) -> float:
|
||||
strat = self.get_strategy(retry_state)
|
||||
if strat is None:
|
||||
return 0.0
|
||||
return strat.wait_func(retry_state)
|
||||
|
||||
def stop_func(self, retry_state: RetryCallState) -> bool:
|
||||
strat = self.get_strategy(retry_state)
|
||||
if strat is None:
|
||||
return False
|
||||
return strat.stop_func(retry_state)
|
||||
|
||||
async def before_sleep(self, retry_state: RetryCallState):
|
||||
strat = self.get_strategy(retry_state)
|
||||
if strat is None:
|
||||
return
|
||||
await strat.before_sleep(retry_state)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Decorator entry point
|
||||
# ------------------------------------------------------------------
|
||||
def __call__(self, func: F) -> F:
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*args, **kwargs): # type: ignore
|
||||
async for attempt in AsyncRetrying(
|
||||
retry=retry_if_exception(lambda e: self.should_retry(e)),
|
||||
wait=self.wait_func,
|
||||
stop=self.stop_func,
|
||||
before_sleep=self.before_sleep,
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper # type: ignore
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# A configurable async retrier for any code block
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
class AsyncRetryBlock:
|
||||
"""
|
||||
Async retry helper for a single exception type and strategy.
|
||||
|
||||
Usage:
|
||||
async with AsyncRetryBlock(strategy):
|
||||
await some_async_function()
|
||||
"""
|
||||
|
||||
def __init__(self, strategy: RetryStrategy, **retry_kwargs): # type: ignore
|
||||
self.strategy = strategy
|
||||
self._retryer = AsyncRetrying(
|
||||
wait=self._wait_func,
|
||||
stop=self._stop_func,
|
||||
before_sleep=self._before_sleep,
|
||||
**retry_kwargs, # type: ignore
|
||||
)
|
||||
|
||||
async def run(self, coro: Callable[..., Awaitable[Any]]) -> Any:
|
||||
"""Run the given coroutine with retries according to the strategy.
|
||||
For example:
|
||||
async def my_coro():
|
||||
...
|
||||
retry_block = AsyncRetryBlock(strategy)
|
||||
result = await retry_block.run(my_coro)
|
||||
"""
|
||||
async for attempt in self._retryer:
|
||||
with attempt:
|
||||
return await coro()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core: async iterator interface
|
||||
# ------------------------------------------------------------------
|
||||
def __aiter__(self) -> AsyncIterator[Any]:
|
||||
"""Return an async iterator that yields retry attempts.
|
||||
Usage:
|
||||
async for attempt in retry_block:
|
||||
with attempt:
|
||||
await some_async_function()
|
||||
"""
|
||||
return self._retryer.__aiter__()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Context manager entry
|
||||
# ------------------------------------------------------------------
|
||||
async def __aenter__(self):
|
||||
self._aiter = self._retryer.__aiter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb): # type: ignore
|
||||
# Consume the retry iterator
|
||||
try:
|
||||
# If exception occurred, let the retryer handle it
|
||||
async for attempt in self._aiter:
|
||||
with attempt:
|
||||
if exc_val:
|
||||
raise exc_val
|
||||
except Exception:
|
||||
# Allow exception to propagate if retries exhausted
|
||||
pass
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Strategy function
|
||||
# ------------------------------------------------------------------
|
||||
def _wait_func(self, retry_state: RetryCallState) -> float:
|
||||
return self.strategy.wait_func(retry_state)
|
||||
|
||||
def _stop_func(self, retry_state: RetryCallState) -> bool:
|
||||
return self.strategy.stop_func(retry_state)
|
||||
|
||||
async def _before_sleep(self, retry_state: RetryCallState):
|
||||
await self.strategy.before_sleep(retry_state)
|
||||
@@ -0,0 +1,685 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict, List, Literal, Optional, Sequence, Union
|
||||
|
||||
from apscheduler.schedulers.background import BackgroundScheduler
|
||||
from apscheduler.triggers.interval import IntervalTrigger
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import and_, or_, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.orm.exc import StaleDataError
|
||||
from tenacity import RetryError
|
||||
|
||||
from agentlightning.types import (
|
||||
Attempt,
|
||||
AttemptedRollout,
|
||||
AttemptStatus,
|
||||
NamedResources,
|
||||
ResourcesUpdate,
|
||||
Rollout,
|
||||
RolloutConfig,
|
||||
RolloutStatus,
|
||||
Span,
|
||||
TaskInput,
|
||||
)
|
||||
|
||||
from ..base import UNSET, LightningStore, Unset, is_finished
|
||||
from .orm import (
|
||||
AttemptInDB,
|
||||
ResourcesUpdateInDB,
|
||||
RolloutInDB,
|
||||
SpanInDB,
|
||||
SpanSeqIdInDB,
|
||||
SqlAlchemyBase,
|
||||
)
|
||||
from .retry_helper import AsyncRetryBlock, AsyncTypeBasedRetry, ExceptionRegistry, RetryStrategy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# TODO add periodic cleanup of old rollouts/attempts/spans
|
||||
|
||||
ExceptionRegistry.register("sqlalchemy.orm.exc.StaleDataError")
|
||||
ExceptionRegistry.register("sqlalchemy.exc.OperationalError")
|
||||
|
||||
db_retry = AsyncTypeBasedRetry(
|
||||
{
|
||||
"sqlalchemy.exc.OperationalError": RetryStrategy(
|
||||
max_attempts=5, wait_seconds=1, backoff=1.5, jitter=0.3, log=True
|
||||
),
|
||||
"sqlalchemy.orm.exc.StaleDataError": RetryStrategy(
|
||||
max_attempts=100, wait_seconds=1e-3, backoff=1.0, jitter=0.1, log=True
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _WaitForRolloutsCompleted(Exception):
|
||||
"""Internal exception to signal that not all rollouts have completed yet."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class BackgroundTaskConfig(BaseModel):
|
||||
name: str # unique name for the task
|
||||
method: str # method name to call, currently only supports methods of SqlLightningStore
|
||||
interval: Dict[Literal["seconds", "minutes", "hours"], float] # interval for the task
|
||||
is_async: bool = True # whether the task method is async, default to True
|
||||
|
||||
|
||||
class SqlLightningStore(LightningStore):
|
||||
"""
|
||||
A LightningStore implementation that uses a database backend to store and manage rollouts and attempts.
|
||||
The database backend is expected to support asynchronous operations.
|
||||
The store uses SQLAlchemy ORM models to interact with the database
|
||||
Args:
|
||||
database_url (string):
|
||||
The database URL for connecting to the database.
|
||||
If None, will read from the 'DATABASE_URL' environment variable.
|
||||
retry_for_waiting (RetryStrategy):
|
||||
Retry strategy for polling when waiting for rollouts to complete.
|
||||
If None, a default strategy will be used.
|
||||
wait_for_nonexistent_rollout (Bool):
|
||||
If True, when waiting for rollouts, will wait for all specified rollouts to complete, including non-existing ones.
|
||||
If False, will ignore non-existing rollouts as completed. (Default: False)
|
||||
background_tasks_cfg (list[Dict[str, Any]]):
|
||||
The configuration for in-process periodic tasks, following the definition of `BackgroundTaskConfig`.
|
||||
IF not provided (None as default), the dbstore will incorporate a default set of periodic tasks as follows:
|
||||
[
|
||||
BackgroundTaskConfig(name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}),
|
||||
]
|
||||
To disable all periodic tasks, provide an empty list `[]`.
|
||||
Note:
|
||||
Explicitly use async `start()` and `stop()` methods to manage the database connection lifecycle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
database_url: Optional[str] = None,
|
||||
*,
|
||||
retry_for_waiting: Optional[dict[str, Any] | RetryStrategy] = None,
|
||||
wait_for_nonexistent_rollout: bool = False,
|
||||
background_tasks_cfg: list[Dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if database_url is None:
|
||||
database_url = os.getenv("DATABASE_URL", None)
|
||||
if database_url is None:
|
||||
raise ValueError(
|
||||
"A database URL must be provided either via the 'database_url' parameter or the 'DATABASE_URL' environment variable."
|
||||
)
|
||||
|
||||
self._engine = create_async_engine(database_url, echo=False)
|
||||
self._async_session = async_sessionmaker(self._engine, expire_on_commit=False)
|
||||
|
||||
self._latest_resources_id = None
|
||||
|
||||
# special handling for retry strategy
|
||||
retry_for_waiting = retry_for_waiting or RetryStrategy(
|
||||
max_attempts=10, # set a limit for retries if timeout is specified, otherwise will change to None later
|
||||
max_retry_delay=None, # set later
|
||||
wait_seconds=10.0, # poll every 10 seconds
|
||||
max_wait_seconds=60.0, # at most wait 60 seconds between retries
|
||||
backoff=1.0,
|
||||
jitter=0.0,
|
||||
log=True,
|
||||
)
|
||||
self.retry_for_waiting = (
|
||||
retry_for_waiting if isinstance(retry_for_waiting, RetryStrategy) else RetryStrategy(**retry_for_waiting)
|
||||
)
|
||||
self.wait_for_nonexistent_rollout = wait_for_nonexistent_rollout
|
||||
|
||||
# setup in-process periodic tasks
|
||||
if background_tasks_cfg is None:
|
||||
self.background_tasks_cfg = [
|
||||
BackgroundTaskConfig(
|
||||
name="check_attempt_timeout", method="check_attempt_timeout", interval={"seconds": 10.0}
|
||||
),
|
||||
]
|
||||
else:
|
||||
self.background_tasks_cfg = [BackgroundTaskConfig(**cfg) for cfg in background_tasks_cfg]
|
||||
self._background_scheduler = BackgroundScheduler()
|
||||
|
||||
async def start(self):
|
||||
async with self._engine.begin() as conn:
|
||||
await conn.run_sync(SqlAlchemyBase.metadata.create_all)
|
||||
for task_cfg in self.background_tasks_cfg:
|
||||
self.add_background_task(task_cfg, to_scheduler_only=True)
|
||||
self._background_scheduler.start() # type: ignore
|
||||
|
||||
async def stop(self):
|
||||
await self._engine.dispose()
|
||||
self._background_scheduler.shutdown() # type: ignore
|
||||
|
||||
def add_background_task(
|
||||
self, task_cfg: Dict[str, Any] | BackgroundTaskConfig, to_scheduler_only: bool = False
|
||||
) -> None:
|
||||
"""Add a new periodic background task to the scheduler.
|
||||
Args:
|
||||
task_cfg (Dict[str, Any] | BackgroundTaskConfig): The configuration for the background task.
|
||||
to_scheduler_only (bool): If True, only add the task to the scheduler without updating the configuration list.
|
||||
Raises:
|
||||
ValueError: If the task method is not defined in SqlLightningStore.
|
||||
"""
|
||||
config = task_cfg if isinstance(task_cfg, BackgroundTaskConfig) else BackgroundTaskConfig(**task_cfg)
|
||||
if not to_scheduler_only:
|
||||
# check existing tasks
|
||||
for existing in self.background_tasks_cfg:
|
||||
if existing.name == config.name:
|
||||
logger.warning(
|
||||
f"Background task {config.name} is already scheduled, will update its configuration."
|
||||
)
|
||||
self.background_tasks_cfg.append(config)
|
||||
delta_t = timedelta(**config.interval)
|
||||
if not hasattr(self, config.method):
|
||||
raise ValueError(f"Periodic task method {config.method} is not defined in SqlLightningStore.")
|
||||
if config.is_async:
|
||||
func = lambda: asyncio.run(getattr(self, config.method)())
|
||||
else:
|
||||
func = lambda: getattr(self, config.method)()
|
||||
|
||||
self._background_scheduler.add_job( # type: ignore
|
||||
func=func,
|
||||
trigger=IntervalTrigger(**config.interval), # type: ignore
|
||||
name=f"SqlLightningStore.{config.name}",
|
||||
replace_existing=True,
|
||||
next_run_time=datetime.now() + delta_t, # schedule the first run after the interval
|
||||
)
|
||||
|
||||
# ------------------------------------------------------
|
||||
# Public methods defined in LightningStore
|
||||
# ------------------------------------------------------
|
||||
|
||||
@db_retry
|
||||
async def start_rollout(
|
||||
self,
|
||||
input: TaskInput,
|
||||
mode: Literal["train", "val", "test"] | None = None,
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> AttemptedRollout:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = RolloutInDB(
|
||||
input=input,
|
||||
mode=mode,
|
||||
resources_id=resources_id or self._latest_resources_id,
|
||||
status="queuing",
|
||||
config=config or RolloutConfig(),
|
||||
rollout_metadata=metadata,
|
||||
)
|
||||
session.add(rollout_obj)
|
||||
attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return attempted_rollout
|
||||
|
||||
@db_retry
|
||||
async def enqueue_rollout(
|
||||
self,
|
||||
input: TaskInput,
|
||||
mode: Literal["train", "val", "test"] | None = None,
|
||||
resources_id: str | None = None,
|
||||
config: RolloutConfig | None = None,
|
||||
metadata: Dict[str, Any] | None = None,
|
||||
) -> Rollout:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = RolloutInDB(
|
||||
input=input,
|
||||
mode=mode,
|
||||
resources_id=resources_id or self._latest_resources_id,
|
||||
status="queuing",
|
||||
config=config or RolloutConfig(),
|
||||
rollout_metadata=metadata,
|
||||
)
|
||||
session.add(rollout_obj)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return rollout_obj.as_rollout()
|
||||
|
||||
@db_retry
|
||||
async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
|
||||
return await self._fifo_dequeue_rollout()
|
||||
|
||||
@db_retry
|
||||
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = await session.get(RolloutInDB, rollout_id)
|
||||
if rollout_obj is None:
|
||||
raise ValueError(f"Rollout {rollout_id} not found")
|
||||
attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return attempted_rollout
|
||||
|
||||
@db_retry
|
||||
async def add_span(self, span: Span) -> Span:
|
||||
seq_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, span.rollout_id, span.attempt_id)
|
||||
return await self._add_span(span.model_dump(), seq_id=seq_id)
|
||||
|
||||
@db_retry
|
||||
async def add_otel_span(
|
||||
self,
|
||||
rollout_id: str,
|
||||
attempt_id: str,
|
||||
readable_span: ReadableSpan,
|
||||
sequence_id: int | None = None,
|
||||
) -> Span:
|
||||
sequence_id = await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id, sequence_id)
|
||||
span = Span.from_opentelemetry(
|
||||
src=readable_span,
|
||||
rollout_id=rollout_id,
|
||||
attempt_id=attempt_id,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
return await self._add_span(span.model_dump(), seq_id=sequence_id)
|
||||
|
||||
@db_retry
|
||||
async def query_rollouts(
|
||||
self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
|
||||
) -> List[Rollout]:
|
||||
rollouts = await RolloutInDB.query_rollouts(self._async_session, statuses=status, ids=rollout_ids) # type: ignore
|
||||
attempt_ids = [r.latest_attempt_id for r in rollouts if r.latest_attempt_id is not None]
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
scalars = await session.scalars(select(AttemptInDB).where(AttemptInDB.attempt_id.in_(attempt_ids)))
|
||||
attempts = scalars.all()
|
||||
attempt_map = {a.attempt_id: a.as_attempt() for a in attempts}
|
||||
return [
|
||||
(
|
||||
AttemptedRollout(**r.as_rollout().model_dump(), attempt=attempt_map[r.latest_attempt_id])
|
||||
if r.latest_attempt_id in attempt_map
|
||||
else r.as_rollout()
|
||||
)
|
||||
for r in rollouts
|
||||
] # type: ignore
|
||||
|
||||
@db_retry
|
||||
async def query_attempts(self, rollout_id: str) -> List[Attempt]:
|
||||
return await AttemptInDB.get_attempts_for_rollout(self._async_session, rollout_id) # type: ignore
|
||||
|
||||
@db_retry
|
||||
async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
|
||||
return await RolloutInDB.get_rollout_by_id(self._async_session, rollout_id)
|
||||
|
||||
@db_retry
|
||||
async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
|
||||
return await AttemptInDB.get_latest_attempt_for_rollout(self._async_session, rollout_id)
|
||||
|
||||
@db_retry
|
||||
async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
|
||||
return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, resources_id)
|
||||
|
||||
@db_retry
|
||||
async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
|
||||
if self._latest_resources_id is None:
|
||||
return None
|
||||
return await ResourcesUpdateInDB.get_resources_by_id(self._async_session, self._latest_resources_id)
|
||||
|
||||
@db_retry
|
||||
async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
|
||||
return await SpanSeqIdInDB.get_next_sequence_id(self._async_session, rollout_id, attempt_id)
|
||||
|
||||
async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
|
||||
# implementation the timeout via tenacity retry mechanism, by a `with` context
|
||||
strategy = RetryStrategy(**self.retry_for_waiting.asdict())
|
||||
if timeout is not None:
|
||||
strategy.max_retry_delay = timeout
|
||||
if strategy.max_attempts is not None:
|
||||
strategy.wait_seconds = min(strategy.wait_seconds, timeout / (strategy.max_attempts + 1))
|
||||
else:
|
||||
strategy.max_attempts = None # infinite retries
|
||||
|
||||
non_completed_ids, non_existing_ids = set(rollout_ids), set(rollout_ids)
|
||||
completed_rollouts: Dict[str, Rollout] = {}
|
||||
if len(non_completed_ids) < len(rollout_ids):
|
||||
logger.warning("Duplicate rollout_ids found in wait_for_rollouts input. Duplicates will be ignored.")
|
||||
|
||||
try:
|
||||
async for attempt in AsyncRetryBlock(
|
||||
strategy,
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
result = await session.scalars(
|
||||
select(RolloutInDB).where(RolloutInDB.rollout_id.in_(non_completed_ids))
|
||||
)
|
||||
rollouts = [r.as_rollout() for r in result.all()]
|
||||
for r in rollouts:
|
||||
if r.rollout_id in non_existing_ids:
|
||||
non_existing_ids.discard(r.rollout_id) # found existing rollout
|
||||
if is_finished(r):
|
||||
completed_rollouts[r.rollout_id] = r
|
||||
non_completed_ids.discard(r.rollout_id)
|
||||
# check termination conditions
|
||||
if self.wait_for_nonexistent_rollout:
|
||||
if len(non_completed_ids) == 0:
|
||||
return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
|
||||
raise _WaitForRolloutsCompleted(
|
||||
f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
|
||||
)
|
||||
else:
|
||||
if len(non_completed_ids) == len(non_existing_ids):
|
||||
logger.warning(f"All remaining rollouts are non-existing: {non_existing_ids}.")
|
||||
return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
|
||||
raise _WaitForRolloutsCompleted(
|
||||
f"WaitForRolloutsCompleted: requested={len(rollout_ids)}, completed={len(completed_rollouts)}, non_existing={len(non_existing_ids)}"
|
||||
)
|
||||
|
||||
except (RetryError, _WaitForRolloutsCompleted):
|
||||
return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
|
||||
except Exception as e:
|
||||
logger.error(f"Error while waiting for rollouts: {e}")
|
||||
raise e
|
||||
|
||||
# Ensure a return value in case no rollouts are completed
|
||||
return [completed_rollouts[rid] for rid in rollout_ids if rid in completed_rollouts]
|
||||
|
||||
@db_retry
|
||||
async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
conditions: List[Any] = [SpanInDB.rollout_id == rollout_id]
|
||||
if attempt_id is not None:
|
||||
if attempt_id == "latest":
|
||||
rollout_obj = await session.get(RolloutInDB, rollout_id)
|
||||
if rollout_obj is None:
|
||||
logger.warning(f"Rollout {rollout_id} does not exist. Cannot query latest attempt spans.")
|
||||
return []
|
||||
attempt_id = rollout_obj.latest_attempt_id
|
||||
conditions.append(SpanInDB.attempt_id == attempt_id)
|
||||
query = select(SpanInDB).where(and_(*conditions)).order_by(SpanInDB.sequence_id.asc())
|
||||
result = await session.scalars(query)
|
||||
span_objs = result.all()
|
||||
return [obj.as_span() for obj in span_objs]
|
||||
|
||||
@db_retry
|
||||
async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
current_time = time.time()
|
||||
resource_obj = ResourcesUpdateInDB(
|
||||
resources=resources,
|
||||
create_time=current_time,
|
||||
update_time=current_time,
|
||||
)
|
||||
session.add(resource_obj)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
self._latest_resources_id = resource_obj.resources_id
|
||||
return resource_obj.as_resources_update()
|
||||
|
||||
@db_retry
|
||||
async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
obj = await session.get(ResourcesUpdateInDB, resources_id)
|
||||
if obj is None:
|
||||
# raise ValueError(f"Failed to update resources {resources_id}. It may not exist.")
|
||||
# FIXME InMemoryLightningStore will create the resources if not exist, but the base method require to raise error
|
||||
# HACK here stick to the behavior of InMemoryLightningStore for compatibility
|
||||
current_time = time.time()
|
||||
obj = ResourcesUpdateInDB(
|
||||
resources_id=resources_id,
|
||||
resources=resources,
|
||||
create_time=current_time,
|
||||
update_time=current_time,
|
||||
)
|
||||
session.add(obj)
|
||||
else:
|
||||
obj.resources = resources
|
||||
await session.flush()
|
||||
self._latest_resources_id = resources_id
|
||||
return obj.as_resources_update()
|
||||
|
||||
@db_retry
|
||||
async def query_resources(self) -> List[ResourcesUpdate]:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
result = await session.scalars(
|
||||
select(ResourcesUpdateInDB).order_by(ResourcesUpdateInDB.create_time.asc())
|
||||
)
|
||||
resource_objs = result.all()
|
||||
return [obj.as_resources_update() for obj in resource_objs]
|
||||
|
||||
@db_retry
|
||||
async def update_rollout(
|
||||
self,
|
||||
rollout_id: str | None,
|
||||
input: TaskInput | Unset = UNSET,
|
||||
mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
|
||||
resources_id: Optional[str] | Unset = UNSET,
|
||||
status: RolloutStatus | Unset = UNSET,
|
||||
config: RolloutConfig | Unset = UNSET,
|
||||
metadata: Optional[Dict[str, Any]] | Unset = UNSET,
|
||||
) -> Rollout:
|
||||
if rollout_id is None:
|
||||
raise ValueError("rollout_id must be provided for updating a rollout.")
|
||||
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = await session.get(RolloutInDB, rollout_id)
|
||||
if rollout_obj is None:
|
||||
raise ValueError(f"Rollout {rollout_id} not found")
|
||||
# udpate fields
|
||||
if not isinstance(input, Unset):
|
||||
rollout_obj.input = input
|
||||
if not isinstance(mode, Unset):
|
||||
rollout_obj.mode = mode
|
||||
if not isinstance(resources_id, Unset):
|
||||
rollout_obj.resources_id = resources_id
|
||||
if not isinstance(status, Unset):
|
||||
await rollout_obj.update_status(dict(event="user_update", new_status=status))
|
||||
if not isinstance(config, Unset):
|
||||
rollout_obj.config = config
|
||||
if not isinstance(metadata, Unset):
|
||||
rollout_obj.rollout_metadata = metadata
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return rollout_obj.as_rollout()
|
||||
|
||||
@db_retry
|
||||
async def update_attempt(
|
||||
self,
|
||||
rollout_id: str,
|
||||
attempt_id: str | Literal["latest"],
|
||||
status: AttemptStatus | Unset = UNSET,
|
||||
worker_id: str | Unset = UNSET,
|
||||
last_heartbeat_time: float | Unset = UNSET,
|
||||
metadata: Optional[Dict[str, Any]] | Unset = UNSET,
|
||||
) -> Attempt:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
rollout_obj = await session.get(RolloutInDB, rollout_id)
|
||||
if rollout_obj is None:
|
||||
raise ValueError(f"Rollout {rollout_id} not found")
|
||||
if attempt_id == "latest":
|
||||
if rollout_obj.latest_attempt_id is None:
|
||||
raise ValueError(f"Rollout {rollout_id} has no attempts. Cannot update latest attempt.")
|
||||
attempt_id = rollout_obj.latest_attempt_id
|
||||
if attempt_id != rollout_obj.latest_attempt_id:
|
||||
logger.warning(
|
||||
f"Updating attempt {attempt_id} which is not the latest attempt for rollout {rollout_id}. Latest is {rollout_obj.latest_attempt_id}."
|
||||
)
|
||||
attempt_obj = await session.get(AttemptInDB, attempt_id)
|
||||
if attempt_obj is None:
|
||||
raise ValueError(f"No attempts found")
|
||||
if attempt_obj.rollout_id != rollout_id:
|
||||
raise ValueError(f"Attempt {attempt_id} does not belong to rollout {rollout_id}.")
|
||||
# update fields
|
||||
if not isinstance(status, Unset):
|
||||
msg = attempt_obj.update_status(dict(event="user_update", new_status=status))
|
||||
if msg is not None:
|
||||
await rollout_obj.update_status(msg)
|
||||
if not isinstance(worker_id, Unset):
|
||||
attempt_obj.worker_id = worker_id
|
||||
if not isinstance(last_heartbeat_time, Unset):
|
||||
attempt_obj.last_heartbeat_time = last_heartbeat_time
|
||||
if not isinstance(metadata, Unset):
|
||||
attempt_obj.attempt_metadata = metadata
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return attempt_obj.as_attempt()
|
||||
|
||||
# ------------------------------------------------------
|
||||
# periodic background tasks can be added here
|
||||
# ------------------------------------------------------
|
||||
|
||||
async def check_attempt_timeout(self):
|
||||
"""Periodically check for attempts that have timed out and update their status accordingly."""
|
||||
# use update with where condition to find and update timed-out attempts
|
||||
current_time = time.time()
|
||||
|
||||
timed_out_results = await self._attempt_timeout_check(current_time)
|
||||
|
||||
# TODO run the tasks with a wrapper with asyncio semaphore to limit concurrency and handle exceptions
|
||||
tasks = [self._process_timed_out_attempt(attempt, current_time) for attempt in timed_out_results]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def _process_timed_out_attempt(self, attempt_ref: AttemptInDB, current_time: float) -> None:
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
# Step 1: Update attempt status
|
||||
attempt_obj = await session.get(
|
||||
AttemptInDB, attempt_ref.attempt_id
|
||||
) # refresh the object in the new session
|
||||
if attempt_obj is None:
|
||||
raise ValueError(f"Attempt {attempt_ref.attempt_id} not found during timeout processing")
|
||||
if attempt_obj.version_id != attempt_ref.version_id:
|
||||
# version mismatch, skip processing to avoid race conditions
|
||||
raise StaleDataError(f"Attempt {attempt_ref.attempt_id} version mismatch during timeout processing")
|
||||
msg = {}
|
||||
if attempt_obj.is_timed_out(current_time):
|
||||
msg = dict(event="overall_timeout", timestamp=current_time)
|
||||
elif attempt_obj.is_unresponsive(current_time):
|
||||
msg = dict(event="single_step_timeout", timestamp=current_time)
|
||||
else:
|
||||
raise ValueError(f"Attempt {attempt_ref.attempt_id} is not timed out during timeout processing")
|
||||
msg2rollout = attempt_obj.update_status(msg)
|
||||
if msg2rollout is None:
|
||||
return # no further update needed
|
||||
|
||||
# Step 2: Update rollouts
|
||||
rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
|
||||
if rollout_obj is None:
|
||||
raise ValueError(f"Rollout {attempt_obj.rollout_id} not found during timeout processing")
|
||||
await rollout_obj.update_status(msg2rollout)
|
||||
|
||||
# ------------------------------------------------------
|
||||
# internal helper methods can be added here
|
||||
# ------------------------------------------------------
|
||||
|
||||
async def _add_span(self, span: Dict[str, Any], seq_id: Optional[int] = None) -> Span:
|
||||
"""Add a new span to the database."""
|
||||
if seq_id is not None:
|
||||
span["sequence_id"] = seq_id
|
||||
extra_dic: Dict[str, Any] = {}
|
||||
for k in list(span.keys()):
|
||||
if k not in SpanInDB.__table__.columns.keys():
|
||||
extra_dic[k] = span.pop(k)
|
||||
span["extra"] = extra_dic if extra_dic else None
|
||||
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
# create SpanInDB object
|
||||
span_obj = SpanInDB(**span)
|
||||
session.add(span_obj)
|
||||
# update attempt's last_heartbeat_time and status
|
||||
attempt_obj = await session.get(AttemptInDB, span["attempt_id"])
|
||||
if attempt_obj is None:
|
||||
raise ValueError(f"Attempt {span['attempt_id']} not found")
|
||||
# ensure the attempt and rollout are in running status
|
||||
msg = attempt_obj.update_status(dict(event="span_received"))
|
||||
if msg is not None:
|
||||
rollout_obj = await session.get(RolloutInDB, attempt_obj.rollout_id)
|
||||
if rollout_obj is None:
|
||||
raise ValueError(f"Rollout {attempt_obj.rollout_id} not found")
|
||||
await rollout_obj.update_status(msg)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return span_obj.as_span()
|
||||
|
||||
async def _fifo_dequeue_rollout(self) -> Optional[AttemptedRollout]:
|
||||
"""Dequeue the next rollout in FIFO order (the one with the earliest enqueue_time).
|
||||
Returns the RolloutInDB object if found, else None.
|
||||
Note: This method does not update the status of the rollout. The caller should handle that.
|
||||
"""
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
# use the update...returning to atomically select the next rollout and claim it by updating its status to 'preparing'
|
||||
result = await session.scalars(
|
||||
select(RolloutInDB)
|
||||
.where(RolloutInDB.status.in_(["queuing", "requeuing"]), RolloutInDB.enqueue_time.isnot(None))
|
||||
.order_by(RolloutInDB.enqueue_time.asc())
|
||||
.limit(1)
|
||||
)
|
||||
rollout_obj = result.one_or_none()
|
||||
if rollout_obj is None:
|
||||
return None # no rollout available
|
||||
# update the status of the rollout to 'preparing' via Compare-and-Swap to avoid race
|
||||
attempted_rollout = await self._start_attempt_for_rollout(session, rollout_obj)
|
||||
await session.flush() # ensure the object is written to the DB
|
||||
return attempted_rollout
|
||||
|
||||
async def _start_attempt_for_rollout(self, session: AsyncSession, rollout_obj: RolloutInDB) -> AttemptedRollout:
|
||||
"""Create a new attempt for the given rollout and update the rollout's fields."""
|
||||
# create a new attempt for this rollout
|
||||
rollout_config = rollout_obj.config
|
||||
attempt_obj = AttemptInDB(
|
||||
rollout_id=rollout_obj.rollout_id,
|
||||
sequence_id=rollout_obj.num_attempts + 1,
|
||||
status="preparing",
|
||||
max_duration=rollout_config.timeout_seconds,
|
||||
max_heartbeat_interval=rollout_config.unresponsive_seconds,
|
||||
)
|
||||
session.add(attempt_obj)
|
||||
# pre-update the rollout_obj fields for CAS
|
||||
rollout_obj.status = attempt_obj.status # type: ignore pre-update the status in the object for CAS
|
||||
rollout_obj.enqueue_time = None # pre-update the enqueue_time in the object for CAS
|
||||
rollout_obj.num_attempts += 1 # pre-update the num_attempts in the object for CAS
|
||||
rollout_obj.latest_attempt_id = attempt_obj.attempt_id # pre-update the latest_attempt_id in the object for CAS
|
||||
|
||||
# create a sequence id tracker for each attempt
|
||||
# FIXME currently InMemoryLightningStore let all attempts under the same rollout share the same span sequence for sorting
|
||||
# create a sequence id tracker for this rollout, only if not exists
|
||||
existing = await session.get(SpanSeqIdInDB, rollout_obj.rollout_id)
|
||||
if existing is None:
|
||||
seq_obj = SpanSeqIdInDB(
|
||||
rollout_id=rollout_obj.rollout_id,
|
||||
attempt_id=attempt_obj.attempt_id,
|
||||
)
|
||||
session.add(seq_obj)
|
||||
|
||||
return AttemptedRollout(**rollout_obj.as_rollout().model_dump(), attempt=attempt_obj.as_attempt())
|
||||
|
||||
async def _attempt_timeout_check(self, now: float) -> Sequence[AttemptInDB]:
|
||||
"""Scan the table for attempts that have timed out based on the given mode, and return them for further processing.
|
||||
Returns:
|
||||
list[AttemptInDB]:
|
||||
A list of AttemptInDB objects that timed out.
|
||||
"""
|
||||
async with self._async_session() as session:
|
||||
async with session.begin():
|
||||
scalars = await session.scalars(
|
||||
select(AttemptInDB).where(
|
||||
and_(
|
||||
AttemptInDB.status.in_(["preparing", "running"]),
|
||||
or_(
|
||||
and_(
|
||||
AttemptInDB.max_duration.isnot(None),
|
||||
(now - AttemptInDB.start_time) > AttemptInDB.max_duration,
|
||||
),
|
||||
and_(
|
||||
AttemptInDB.max_heartbeat_interval.isnot(None),
|
||||
(now - AttemptInDB.last_heartbeat_time) > AttemptInDB.max_heartbeat_interval,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
return scalars.all()
|
||||
@@ -26,6 +26,7 @@ from typing import (
|
||||
Sequence,
|
||||
Set,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
@@ -45,7 +46,7 @@ from agentlightning.types import (
|
||||
TaskInput,
|
||||
)
|
||||
|
||||
from .base import UNSET, LightningStore, Unset, is_finished, is_queuing
|
||||
from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset, is_finished, is_queuing
|
||||
from .utils import healthcheck, propagate_status
|
||||
|
||||
T_callable = TypeVar("T_callable", bound=Callable[..., Any])
|
||||
@@ -242,6 +243,14 @@ class InMemoryLightningStore(LightningStore):
|
||||
# Completion tracking for wait_for_rollouts (cross-loop safe)
|
||||
self._completion_events: Dict[str, threading.Event] = {}
|
||||
|
||||
def capabilities(self) -> LightningStoreCapabilities:
|
||||
"""Return the capabilities of the store."""
|
||||
return LightningStoreCapabilities(
|
||||
thread_safe=False,
|
||||
async_safe=True,
|
||||
zero_copy=False,
|
||||
)
|
||||
|
||||
@_healthcheck_wrapper
|
||||
async def start_rollout(
|
||||
self,
|
||||
@@ -369,6 +378,9 @@ class InMemoryLightningStore(LightningStore):
|
||||
self._attempts[rollout.rollout_id] = []
|
||||
self._attempts[rollout.rollout_id].append(attempt)
|
||||
|
||||
# Sync attempt status to rollout
|
||||
await self._update_rollout_unlocked(rollout.rollout_id, status="preparing")
|
||||
|
||||
return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
|
||||
|
||||
# If not in queuing state, skip this rollout and continue
|
||||
@@ -413,6 +425,9 @@ class InMemoryLightningStore(LightningStore):
|
||||
self._attempts[rollout_id] = []
|
||||
self._attempts[rollout_id].append(attempt)
|
||||
|
||||
# Sync attempt status to rollout
|
||||
await self._update_rollout_unlocked(rollout_id, status="preparing")
|
||||
|
||||
self._completion_events.setdefault(rollout.rollout_id, threading.Event())
|
||||
|
||||
return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
|
||||
@@ -420,7 +435,7 @@ class InMemoryLightningStore(LightningStore):
|
||||
@_healthcheck_wrapper
|
||||
async def query_rollouts(
|
||||
self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
|
||||
) -> List[Rollout]:
|
||||
) -> List[Union[Rollout, AttemptedRollout]]:
|
||||
"""Retrieves rollouts filtered by their status and rollout ids.
|
||||
If no status is provided, returns all rollouts.
|
||||
|
||||
@@ -439,16 +454,40 @@ class InMemoryLightningStore(LightningStore):
|
||||
status_set = set(status)
|
||||
rollouts = [rollout for rollout in rollouts if rollout.status in status_set]
|
||||
|
||||
# Attach the latest attempt to the rollout objects
|
||||
rollouts = [self._rollout_to_attempted_rollout_unlocked(rollout) for rollout in rollouts]
|
||||
|
||||
return rollouts
|
||||
|
||||
@_healthcheck_wrapper
|
||||
async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
|
||||
async def get_rollout_by_id(self, rollout_id: str) -> Optional[Union[Rollout, AttemptedRollout]]:
|
||||
"""Retrieves a specific rollout by its ID.
|
||||
|
||||
See [`LightningStore.get_rollout_by_id()`][agentlightning.LightningStore.get_rollout_by_id] for semantics.
|
||||
|
||||
If the rollout has been attempted, the latest attempt will also be returned.
|
||||
"""
|
||||
async with self._lock:
|
||||
return self._rollouts.get(rollout_id)
|
||||
rollout = self._rollouts.get(rollout_id)
|
||||
if rollout is None:
|
||||
return None
|
||||
return self._rollout_to_attempted_rollout_unlocked(rollout)
|
||||
|
||||
def _rollout_to_attempted_rollout_unlocked(self, rollout: Rollout) -> Union[Rollout, AttemptedRollout]:
|
||||
"""Query the latest attempt for the rollout, and attach it to the rollout object.
|
||||
|
||||
If the rollout has no attempts, return the rollout object itself.
|
||||
"""
|
||||
latest_attempt = self._get_latest_attempt_unlocked(rollout.rollout_id)
|
||||
if latest_attempt is None:
|
||||
return rollout
|
||||
else:
|
||||
return AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt)
|
||||
|
||||
def _get_latest_attempt_unlocked(self, rollout_id: str) -> Optional[Attempt]:
|
||||
"""The unlocked version of `get_latest_attempt`."""
|
||||
attempts = self._attempts.get(rollout_id, [])
|
||||
return max(attempts, key=lambda a: a.sequence_id) if attempts else None
|
||||
|
||||
@_healthcheck_wrapper
|
||||
async def query_attempts(self, rollout_id: str) -> List[Attempt]:
|
||||
@@ -467,10 +506,13 @@ class InMemoryLightningStore(LightningStore):
|
||||
See [`LightningStore.get_latest_attempt()`][agentlightning.LightningStore.get_latest_attempt] for semantics.
|
||||
"""
|
||||
async with self._lock:
|
||||
attempts = self._attempts.get(rollout_id, [])
|
||||
if not attempts:
|
||||
return None
|
||||
return max(attempts, key=lambda a: a.sequence_id)
|
||||
return self._get_latest_attempt_unlocked(rollout_id)
|
||||
|
||||
@_healthcheck_wrapper
|
||||
async def query_resources(self) -> List[ResourcesUpdate]:
|
||||
"""Return every stored resource snapshot in insertion order."""
|
||||
async with self._lock:
|
||||
return list(self._resources.values())
|
||||
|
||||
@_healthcheck_wrapper
|
||||
async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
|
||||
@@ -480,7 +522,14 @@ class InMemoryLightningStore(LightningStore):
|
||||
"""
|
||||
resources_id = _generate_resources_id()
|
||||
async with self._lock:
|
||||
update = ResourcesUpdate(resources_id=resources_id, resources=resources)
|
||||
current_time = time.time()
|
||||
update = ResourcesUpdate(
|
||||
resources_id=resources_id,
|
||||
resources=resources,
|
||||
create_time=current_time,
|
||||
update_time=current_time,
|
||||
version=1,
|
||||
)
|
||||
self._resources[resources_id] = update
|
||||
self._latest_resources_id = resources_id
|
||||
return update
|
||||
@@ -493,7 +542,23 @@ class InMemoryLightningStore(LightningStore):
|
||||
See [`LightningStore.update_resources()`][agentlightning.LightningStore.update_resources] for semantics.
|
||||
"""
|
||||
async with self._lock:
|
||||
update = ResourcesUpdate(resources_id=resources_id, resources=resources)
|
||||
current_time = time.time()
|
||||
if resources_id not in self._resources:
|
||||
update = ResourcesUpdate(
|
||||
resources_id=resources_id,
|
||||
resources=resources,
|
||||
create_time=current_time,
|
||||
update_time=current_time,
|
||||
version=1,
|
||||
)
|
||||
else:
|
||||
update = self._resources[resources_id].model_copy(
|
||||
update={
|
||||
"resources": resources,
|
||||
"update_time": current_time,
|
||||
"version": self._resources[resources_id].version + 1,
|
||||
}
|
||||
)
|
||||
self._resources[resources_id] = update
|
||||
self._latest_resources_id = resources_id
|
||||
return update
|
||||
|
||||
@@ -20,7 +20,7 @@ from agentlightning.types import (
|
||||
TaskInput,
|
||||
)
|
||||
|
||||
from .base import UNSET, LightningStore, Unset
|
||||
from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
|
||||
|
||||
|
||||
class LightningStoreThreaded(LightningStore):
|
||||
@@ -35,6 +35,15 @@ class LightningStoreThreaded(LightningStore):
|
||||
self.store = store
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def capabilities(self) -> LightningStoreCapabilities:
|
||||
"""Return the capabilities of the store."""
|
||||
capabilities = self.store.capabilities()
|
||||
return {
|
||||
**capabilities,
|
||||
"async_safe": True,
|
||||
"thread_safe": True,
|
||||
}
|
||||
|
||||
async def start_rollout(
|
||||
self,
|
||||
input: TaskInput,
|
||||
|
||||
@@ -15,9 +15,9 @@ from agentops.sdk.core import TracingCore
|
||||
from agentops.sdk.processors import SpanProcessor
|
||||
from opentelemetry.instrumentation.utils import suppress_instrumentation
|
||||
from opentelemetry.sdk.trace import ReadableSpan
|
||||
from opentelemetry.trace.status import StatusCode
|
||||
|
||||
from agentlightning.instrumentation import instrument_all, uninstrument_all
|
||||
from agentlightning.instrumentation.agentops import AgentOpsServerManager
|
||||
from agentlightning.store.base import LightningStore
|
||||
|
||||
from .base import Tracer
|
||||
@@ -56,46 +56,11 @@ class AgentOpsTracer(Tracer):
|
||||
self.instrument_managed = instrument_managed
|
||||
self.daemon = daemon
|
||||
|
||||
self._agentops_server_manager = AgentOpsServerManager(self.daemon)
|
||||
self._agentops_server_port_val: Optional[int] = None
|
||||
|
||||
if not self.agentops_managed:
|
||||
logger.warning("agentops_managed=False. You are responsible for AgentOps setup.")
|
||||
if not self.instrument_managed:
|
||||
logger.warning("instrument_managed=False. You are responsible for all instrumentation.")
|
||||
|
||||
def __getstate__(self):
|
||||
state = self.__dict__.copy()
|
||||
state["_agentops_server_manager"] = None # Exclude the unpicklable server manager
|
||||
# _agentops_server_port_val (int) is inherently picklable and will be included.
|
||||
logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: Any):
|
||||
self.__dict__.update(state)
|
||||
# In child process, self._agentops_server_manager will be None.
|
||||
logger.debug(f"Setting state for unpickled Trainer (PID {os.getpid()}). _agentops_server_manager is None.")
|
||||
|
||||
def init(self, *args: Any, **kwargs: Any):
|
||||
if self.agentops_managed and self._agentops_server_manager:
|
||||
self._agentops_server_manager.start()
|
||||
self._agentops_server_port_val = self._agentops_server_manager.get_port()
|
||||
if self._agentops_server_port_val is None:
|
||||
if (
|
||||
self._agentops_server_manager.server_process is not None
|
||||
and self._agentops_server_manager.server_process.is_alive()
|
||||
):
|
||||
raise RuntimeError("AgentOps server started but port is None. Check server manager logic.")
|
||||
elif (
|
||||
self._agentops_server_port_val is None and self._agentops_server_manager.server_process is None
|
||||
): # Server failed to start
|
||||
raise RuntimeError("AgentOps server manager indicates server is not running and port is None.")
|
||||
|
||||
def teardown(self):
|
||||
if self.agentops_managed:
|
||||
self._agentops_server_manager.stop()
|
||||
logger.info("AgentOps server stopped.")
|
||||
|
||||
def instrument(self, worker_id: int):
|
||||
instrument_all()
|
||||
|
||||
@@ -111,24 +76,9 @@ class AgentOpsTracer(Tracer):
|
||||
logger.info(f"[Worker {worker_id}] Instrumentation applied.")
|
||||
|
||||
if self.agentops_managed:
|
||||
if self._agentops_server_port_val: # Use the stored, picklable port value
|
||||
base_url = f"http://localhost:{self._agentops_server_port_val}"
|
||||
env_vars_to_set = {
|
||||
"AGENTOPS_API_KEY": "dummy",
|
||||
"AGENTOPS_API_ENDPOINT": base_url,
|
||||
"AGENTOPS_APP_URL": f"{base_url}/notavailable",
|
||||
"AGENTOPS_EXPORTER_ENDPOINT": f"{base_url}/traces",
|
||||
}
|
||||
for key, value in env_vars_to_set.items():
|
||||
os.environ[key] = value
|
||||
logger.info(f"[Worker {worker_id}] Env var set: {key}={value}")
|
||||
else:
|
||||
logger.warning(
|
||||
f"[Worker {worker_id}] AgentOps managed, but local server port is not available. Client may not connect as expected."
|
||||
)
|
||||
|
||||
os.environ.setdefault("AGENTOPS_API_KEY", "dummy")
|
||||
if not agentops.get_client().initialized:
|
||||
agentops.init() # type: ignore
|
||||
agentops.init(auto_start_session=False) # type: ignore
|
||||
logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
|
||||
else:
|
||||
logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
|
||||
@@ -192,15 +142,30 @@ class AgentOpsTracer(Tracer):
|
||||
if not self._lightning_span_processor:
|
||||
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
|
||||
|
||||
if store is not None and rollout_id is not None and attempt_id is not None:
|
||||
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
|
||||
with ctx as processor:
|
||||
yield processor
|
||||
elif store is None and rollout_id is None and attempt_id is None:
|
||||
with self._lightning_span_processor:
|
||||
yield self._lightning_span_processor
|
||||
else:
|
||||
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
|
||||
kwargs: dict[str, Any] = {}
|
||||
if name is not None:
|
||||
kwargs["trace_name"] = name
|
||||
elif rollout_id is not None:
|
||||
kwargs["trace_name"] = rollout_id
|
||||
trace = agentops.start_trace(**kwargs)
|
||||
status = StatusCode.OK # type: ignore
|
||||
try:
|
||||
if store is not None and rollout_id is not None and attempt_id is not None:
|
||||
ctx = self._lightning_span_processor.with_context(
|
||||
store=store, rollout_id=rollout_id, attempt_id=attempt_id
|
||||
)
|
||||
with ctx as processor:
|
||||
yield processor
|
||||
elif store is None and rollout_id is None and attempt_id is None:
|
||||
with self._lightning_span_processor:
|
||||
yield self._lightning_span_processor
|
||||
else:
|
||||
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
|
||||
except Exception as e:
|
||||
status = StatusCode.ERROR # type: ignore
|
||||
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
|
||||
finally:
|
||||
agentops.end_trace(trace, end_state=status) # type: ignore
|
||||
|
||||
def get_last_trace(self) -> List[ReadableSpan]:
|
||||
"""
|
||||
|
||||
@@ -117,6 +117,7 @@ AttemptStatus = Literal[
|
||||
]
|
||||
"""The status of an attempt."""
|
||||
|
||||
|
||||
RolloutMode = Literal["train", "val", "test"]
|
||||
"""Possible rollout modes."""
|
||||
|
||||
|
||||
@@ -194,5 +194,11 @@ class ResourcesUpdate(BaseModel):
|
||||
|
||||
resources_id: str
|
||||
"""Identifier used to version the resources."""
|
||||
create_time: float
|
||||
"""Timestamp of the creation time of the resources."""
|
||||
update_time: float
|
||||
"""Timestamp of the last update time of the resources."""
|
||||
version: int
|
||||
"""Version of the resources."""
|
||||
resources: NamedResources
|
||||
"""Mapping of resource names to their definitions."""
|
||||
|
||||
@@ -1,3 +1 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
# TODO: Implement this
|
||||
@@ -0,0 +1,999 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import multiprocessing
|
||||
import queue
|
||||
import signal
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from contextlib import asynccontextmanager, suppress
|
||||
from dataclasses import dataclass
|
||||
from multiprocessing.process import BaseProcess
|
||||
from typing import Any, AsyncContextManager, AsyncIterator, Dict, Literal, Optional
|
||||
|
||||
import aiohttp
|
||||
import requests
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from gunicorn.app.base import BaseApplication
|
||||
from gunicorn.arbiter import Arbiter
|
||||
from portpicker import pick_unused_port
|
||||
|
||||
__all__ = ["PythonServerLauncher", "PythonServerLauncherArgs", "LaunchMode"]
|
||||
|
||||
|
||||
LaunchMode = Literal["asyncio", "thread", "mp"]
|
||||
"""The launch mode for the server."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class PythonServerLauncherArgs:
|
||||
port: Optional[int] = None
|
||||
"""The TCP port to listen on. If not provided, the server will use a random available port."""
|
||||
host: Optional[str] = None
|
||||
"""The hostname or IP address to bind the server to."""
|
||||
access_host: Optional[str] = None
|
||||
"""The hostname or IP address to advertise to the client. If not provided, the server will use the default outbound IPv4 address for this machine."""
|
||||
launch_mode: LaunchMode = "asyncio"
|
||||
"""The launch mode. `asyncio` is the default mode to runs the server in the current thread.
|
||||
`thread` runs the server in a separate thread. `mp` runs the server in a separate process."""
|
||||
n_workers: int = 1
|
||||
"""The number of workers to run in the server. Only applicable for `mp` mode.
|
||||
When `n_workers > 1`, the server will be run using Gunicorn.
|
||||
"""
|
||||
healthcheck_url: Optional[str] = None
|
||||
"""The health check URL to use.
|
||||
If not provided, the server will not be checked for healthiness after starting.
|
||||
"""
|
||||
log_level: int = logging.INFO
|
||||
"""The log level to use."""
|
||||
startup_timeout: float = 60.0
|
||||
"""The timeout to wait for the server to start up."""
|
||||
kill_unhealthy_server: bool = True
|
||||
"""Whether to kill the server if it is not healthy after startup.
|
||||
This setting is ignored when `launch_mode` is not `asyncio`.
|
||||
"""
|
||||
thread_join_timeout: float = 10.0
|
||||
"""The timeout to wait for the thread to join."""
|
||||
process_join_timeout: float = 10.0
|
||||
"""The timeout to wait for the process to join."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChildEvent:
|
||||
"""An event that occurred in a child process."""
|
||||
|
||||
kind: Literal["ready", "error"]
|
||||
"""The kind of message."""
|
||||
exc_type: Optional[str] = None
|
||||
"""The type of the exception, only used for error messages."""
|
||||
message: Optional[str] = None
|
||||
"""The message of the exception, only used for error messages."""
|
||||
traceback: Optional[str] = None
|
||||
"""The traceback of the exception, only used for error messages."""
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GunicornApp(BaseApplication):
|
||||
"""
|
||||
Programmatic Gunicorn application that:
|
||||
|
||||
- Accepts a `FastAPI` app object and option dict.
|
||||
- Uses `uvicorn_worker.UvicornWorker`.
|
||||
"""
|
||||
|
||||
def __init__(self, app: FastAPI, options: Dict[str, Any]):
|
||||
self.application = app
|
||||
self.options = options
|
||||
super().__init__() # type: ignore
|
||||
|
||||
def load_config(self):
|
||||
cfg = self.cfg
|
||||
valid_keys = cfg.settings.keys() # type: ignore
|
||||
for k, v in (self.options or {}).items():
|
||||
if k in valid_keys and v is not None:
|
||||
cfg.set(k, v) # type: ignore
|
||||
|
||||
def load(self):
|
||||
return self.application
|
||||
|
||||
|
||||
async def shutdown_uvicorn_server(server: uvicorn.Server, task: asyncio.Task[None], timeout: float = 5.0) -> None:
|
||||
"""Shutdown a uvicorn server and await the serving task."""
|
||||
logger.debug("Requesting graceful shutdown of uvicorn server.")
|
||||
server.should_exit = True
|
||||
# Give uvicorn a brief window to shut down cleanly.
|
||||
try:
|
||||
logger.debug("Waiting for graceful shutdown of uvicorn server.")
|
||||
await asyncio.wait_for(task, timeout=timeout)
|
||||
logger.debug("Graceful shutdown of uvicorn server completed.")
|
||||
except asyncio.TimeoutError:
|
||||
logger.error("Graceful shutdown of uvicorn server timed out.")
|
||||
# As a last resort, cancel; this shouldn't happen under normal circumstances.
|
||||
task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
logger.warning("Uvicorn server forced to stop.")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def noop_context() -> AsyncIterator[None]:
|
||||
"""A real async context manager that does nothing (satisfies serve_context)."""
|
||||
yield
|
||||
|
||||
|
||||
async def run_uvicorn_asyncio(
|
||||
uvicorn_server: uvicorn.Server,
|
||||
serve_context: AsyncContextManager[Any],
|
||||
timeout: float = 60.0,
|
||||
health_url: Optional[str] = None,
|
||||
wait_for_serve: bool = True,
|
||||
kill_unhealthy_server: bool = True,
|
||||
) -> asyncio.Task[None]:
|
||||
"""Run two Asyncio tasks in parallel:
|
||||
|
||||
- A watcher task that waits for the server to start up and then checks for healthiness.
|
||||
- A server task that serves the server.
|
||||
"""
|
||||
server_start_exception: Optional[BaseException] = None
|
||||
|
||||
# watcher: when server.started flips True, announce READY once
|
||||
async def _watch_server() -> None:
|
||||
start_time = time.time()
|
||||
deadline = start_time + timeout # child-side startup window
|
||||
logger.debug(f"Waiting for server to start up for {timeout:.2f} seconds...")
|
||||
# Wait for the server to start up or the deadline to be reached, or an exception to be raised.
|
||||
while time.time() < deadline and not uvicorn_server.started and server_start_exception is None:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
if not uvicorn_server.started:
|
||||
# Normally, the program will not reach this point, as the server will throw the exception itself earlier.
|
||||
raise RuntimeError(f"Server did not start up within {timeout:.2f} seconds.") from server_start_exception
|
||||
|
||||
logger.info(f"Server started up in {time.time() - start_time:.2f} seconds.")
|
||||
|
||||
# Check for health endpoint status if provided
|
||||
if health_url is not None:
|
||||
logger.info(f"Probing health endpoint {health_url}...")
|
||||
async with aiohttp.ClientSession() as session:
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
async with session.get(health_url) as resp:
|
||||
if resp.status == 200:
|
||||
logger.info(
|
||||
f"Server is healthy at {health_url} in {time.time() - start_time:.2f} seconds."
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.debug(
|
||||
f"Server is NOT healthy at {health_url} in {time.time() - start_time:.2f} seconds. Got status {resp.status}."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Error probing health endpoint {health_url}: {str(e)}")
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# If the server is not healthy, kill it if requested.
|
||||
health_failed_seconds = time.time() - start_time
|
||||
if kill_unhealthy_server:
|
||||
logger.error(
|
||||
f"Server is not healthy at {health_url} after {health_failed_seconds:.2f} seconds. Shutting down server gracefully."
|
||||
)
|
||||
uvicorn_server.should_exit = True
|
||||
await serve_task
|
||||
|
||||
raise RuntimeError(
|
||||
f"Server is not healthy at {health_url} after {health_failed_seconds:.2f} seconds. It has been killed."
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Server is not healthy at {health_url} after {health_failed_seconds:.2f} seconds. It has been left running."
|
||||
)
|
||||
|
||||
else:
|
||||
logger.info("Server does not provide a health check endpoint. Skipping health check.")
|
||||
|
||||
async def _serve_server() -> None:
|
||||
nonlocal server_start_exception
|
||||
async with serve_context:
|
||||
try:
|
||||
await uvicorn_server.serve()
|
||||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||||
# Normal shutdown path; propagate without rewrapping
|
||||
raise
|
||||
except BaseException as exc:
|
||||
server_start_exception = exc
|
||||
if wait_for_serve:
|
||||
# This probably sends out earlier than watcher exception; but either one is fine.
|
||||
raise RuntimeError("Uvicorn server failed to serve") from exc
|
||||
else:
|
||||
# If the caller is not waiting for this coroutine, we just log the error.
|
||||
# It will be handled by the watch task.
|
||||
logger.exception("Uvicorn server failed to serve. Inspect the logs for details.")
|
||||
|
||||
serve_task = asyncio.create_task(_serve_server())
|
||||
watch_task = asyncio.create_task(_watch_server())
|
||||
|
||||
if wait_for_serve:
|
||||
await asyncio.gather(watch_task, serve_task)
|
||||
else:
|
||||
# Wait for watch only, the serve task will run in the background.
|
||||
await watch_task
|
||||
return serve_task
|
||||
|
||||
|
||||
def run_uvicorn_thread(
|
||||
uvicorn_server: uvicorn.Server,
|
||||
serve_context: AsyncContextManager[Any],
|
||||
event_queue: queue.Queue[ChildEvent],
|
||||
stop_event: threading.Event,
|
||||
timeout: float = 60.0,
|
||||
health_url: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Run a uvicorn server in a thread.
|
||||
|
||||
How to stop programmatically (from the main thread):
|
||||
|
||||
uvicorn_server.should_exit = True
|
||||
|
||||
This function:
|
||||
|
||||
- starts the server and waits for startup/health (if provided),
|
||||
- then blocks until the server exits,
|
||||
- shuts down cleanly if an error happens during startup/health,
|
||||
- or if the thread is stopped by stop event.
|
||||
"""
|
||||
|
||||
async def _main() -> None:
|
||||
# Start server without waiting for full lifecycle; return once startup/health is done.
|
||||
serve_task: Optional[asyncio.Task[None]] = None
|
||||
try:
|
||||
serve_task = await run_uvicorn_asyncio(
|
||||
uvicorn_server=uvicorn_server,
|
||||
serve_context=serve_context,
|
||||
timeout=timeout,
|
||||
health_url=health_url,
|
||||
wait_for_serve=False, # return after startup watcher finishes
|
||||
kill_unhealthy_server=True, # raise if health fails within timeout
|
||||
)
|
||||
event_queue.put(ChildEvent(kind="ready"))
|
||||
except Exception as exc:
|
||||
# Startup/health failed; nothing is running in the background.
|
||||
logger.exception("Uvicorn failed to start or was unhealthy.")
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error", exc_type=type(exc).__name__, message=str(exc), traceback=traceback.format_exc()
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
logger.debug("Thread server started and ready.")
|
||||
try:
|
||||
# At this point, the server is up and serving in the same thread's loop.
|
||||
# Block here until it exits (caller can stop it via setting the stop_event).
|
||||
while not stop_event.is_set():
|
||||
await asyncio.sleep(0.1)
|
||||
except asyncio.CancelledError:
|
||||
# Shutdown the server.
|
||||
logger.warning(
|
||||
"Thread server received asyncio cancellation signal. Shutting down gracefully. This is not the recommended way to stop the server."
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("Exception during the thread event waiting loop.")
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error", exc_type=type(exc).__name__, message=str(exc), traceback=traceback.format_exc()
|
||||
)
|
||||
)
|
||||
finally:
|
||||
logger.info("Requesting graceful shutdown of uvicorn server.")
|
||||
await shutdown_uvicorn_server(uvicorn_server, serve_task)
|
||||
logger.info("Uvicorn server shut down gracefully.")
|
||||
|
||||
# Each thread needs its own event loop; use asyncio.run to manage it cleanly.
|
||||
try:
|
||||
asyncio.run(_main())
|
||||
except Exception:
|
||||
# Exceptions are already logged above; don't crash the process from a thread.
|
||||
# (Caller can inspect logs or add a queue/handler if they need to propagate.)
|
||||
logger.exception("Exception within the thread server loop. Inspect the logs for details.")
|
||||
|
||||
|
||||
def run_uvicorn_subprocess(
|
||||
uvicorn_server: uvicorn.Server,
|
||||
serve_context: AsyncContextManager[Any],
|
||||
event_queue: multiprocessing.Queue[ChildEvent],
|
||||
timeout: float = 60.0,
|
||||
health_url: Optional[str] = None,
|
||||
):
|
||||
"""Run a uvicorn server in a subprocess.
|
||||
|
||||
Behavior:
|
||||
|
||||
- Start uvicorn and wait for startup/health (if provided).
|
||||
- Post `ChildEvent(kind="ready")` once the server is up.
|
||||
- Stay alive until a termination signal (SIGTERM/SIGINT).
|
||||
- On signal, request graceful shutdown and wait for the server to exit.
|
||||
|
||||
This must be used with forked multiprocessing.Process.
|
||||
"""
|
||||
|
||||
async def _main() -> None:
|
||||
stop_event = asyncio.Event()
|
||||
|
||||
# Register signal handlers
|
||||
loop = asyncio.get_running_loop()
|
||||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||||
loop.add_signal_handler(sig, stop_event.set)
|
||||
logger.debug("Subprocess signal handlers registered.")
|
||||
|
||||
serve_task: Optional[asyncio.Task[None]] = None
|
||||
|
||||
try:
|
||||
# Start server but don't block on its full lifecycle; this returns once the watcher finishes.
|
||||
serve_task = await run_uvicorn_asyncio(
|
||||
uvicorn_server=uvicorn_server,
|
||||
serve_context=serve_context,
|
||||
timeout=timeout,
|
||||
health_url=health_url,
|
||||
wait_for_serve=False, # return after startup/health passes
|
||||
kill_unhealthy_server=True, # if unhealthy, fail fast in the child
|
||||
)
|
||||
|
||||
# Announce readiness only after watcher success.
|
||||
event_queue.put(ChildEvent(kind="ready"))
|
||||
|
||||
logger.debug("Subprocess server started and ready.")
|
||||
|
||||
# Wait until we're told to stop.
|
||||
await stop_event.wait()
|
||||
|
||||
except Exception as exc:
|
||||
# Propagate any startup/health errors to the parent.
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type=type(exc).__name__,
|
||||
message=str(exc),
|
||||
traceback=traceback.format_exc(),
|
||||
)
|
||||
)
|
||||
logger.exception("Subprocess server failed to start or was unhealthy.")
|
||||
|
||||
finally:
|
||||
# Request graceful shutdown if the server is running.
|
||||
if serve_task is not None:
|
||||
logger.info("Requesting graceful shutdown of subprocess server.")
|
||||
await shutdown_uvicorn_server(uvicorn_server, serve_task)
|
||||
logger.info("Subprocess server shut down gracefully.")
|
||||
else:
|
||||
logger.info("Subprocess server was not running. Nothing to stop.")
|
||||
|
||||
try:
|
||||
asyncio.run(_main())
|
||||
except Exception as exc:
|
||||
# If something escapes _main(), make sure the parent hears about it.
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type=type(exc).__name__,
|
||||
message=str(exc),
|
||||
traceback=traceback.format_exc(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def run_gunicorn(
|
||||
gunicorn_app: GunicornApp,
|
||||
serve_context: AsyncContextManager[Any],
|
||||
event_queue: multiprocessing.Queue[ChildEvent],
|
||||
timeout: float = 60.0,
|
||||
health_url: Optional[str] = None,
|
||||
):
|
||||
"""Run a gunicorn server in a subprocess.
|
||||
|
||||
The master arbiter will reside in a non-daemon subprocess,
|
||||
and the workers will be forked from the arbiter.
|
||||
|
||||
Behavior:
|
||||
|
||||
- Start Arbiter.run() (blocking) in this process.
|
||||
- A watchdog thread waits for workers to spawn, then (optionally) verifies a health URL.
|
||||
- On success: put `ChildEvent(kind="ready")`.
|
||||
- On failure/timeout: put `ChildEvent(kind="error")` and request a graceful shutdown.
|
||||
|
||||
`serve_context` will be applied around the `arbiter.run()` call.
|
||||
"""
|
||||
# Create the arbiter up-front so the watchdog can inspect it.
|
||||
try:
|
||||
arbiter = Arbiter(gunicorn_app)
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to initialize Gunicorn Arbiter.")
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type=type(exc).__name__,
|
||||
message=str(exc),
|
||||
traceback=traceback.format_exc(),
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
runtime_error: Optional[BaseException] = None
|
||||
|
||||
def _watchdog() -> None:
|
||||
start = time.time()
|
||||
deadline = start + timeout
|
||||
|
||||
# First, wait for arbiter.workers to get populated
|
||||
while time.time() < deadline and not arbiter.WORKERS: # type: ignore
|
||||
# If arbiter died early, abort quickly.
|
||||
if runtime_error is not None:
|
||||
logger.error("Gunicorn arbiter exited during startup. Watchdog exiting.")
|
||||
return
|
||||
time.sleep(0.1)
|
||||
|
||||
if not arbiter.WORKERS: # type: ignore
|
||||
elapsed_time = time.time() - start
|
||||
logger.error("Gunicorn workers did not start within %.2f seconds.", elapsed_time)
|
||||
if runtime_error is None:
|
||||
# Timeout case: arbiter throws no exception.
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type="RuntimeError",
|
||||
message=f"Gunicorn workers did not start within {elapsed_time:.2f} seconds.",
|
||||
traceback=None,
|
||||
)
|
||||
)
|
||||
logger.info("Halting Gunicorn arbiter.")
|
||||
# Ask arbiter to stop if it's still alive.
|
||||
# It will make the watchdog exit too.
|
||||
arbiter.signal(signal.SIGTERM, inspect.currentframe()) # type: ignore
|
||||
else:
|
||||
# Timeout case: arbiter has thrown an exception.
|
||||
logger.error("Gunicorn arbiter exited during startup. Watchdog exiting.")
|
||||
return
|
||||
|
||||
# Second, check for health endpoint status if provided
|
||||
if health_url:
|
||||
while time.time() < deadline:
|
||||
# If arbiter died early, abort.
|
||||
if runtime_error is not None:
|
||||
logger.error("Gunicorn arbiter exited during health check. Watchdog exiting.")
|
||||
return
|
||||
|
||||
# Check if the server is healthy.
|
||||
try:
|
||||
resp = requests.get(health_url, timeout=2.0)
|
||||
if resp.status_code == 200:
|
||||
logger.debug(f"Server is healthy at {health_url} in {time.time() - start:.2f} seconds.")
|
||||
# Check arbiter status again.
|
||||
if runtime_error is None:
|
||||
event_queue.put(ChildEvent(kind="ready"))
|
||||
else:
|
||||
logger.error(
|
||||
"Response status is 200 but arbiter has thrown an exception. This should not happen."
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
logger.debug(
|
||||
f"Server is still not healthy at {health_url} in {time.time() - start:.2f} seconds.",
|
||||
exc_info=True,
|
||||
)
|
||||
time.sleep(0.1)
|
||||
|
||||
# Health failed: report and shut down.
|
||||
elapsed = time.time() - start
|
||||
logger.error(
|
||||
"Server is not healthy at %s after %.2f seconds. Shutting down.",
|
||||
health_url,
|
||||
elapsed,
|
||||
)
|
||||
if runtime_error is None:
|
||||
# Arbiter throws no exception. This is a simple timeout case.
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type="RuntimeError",
|
||||
message=(
|
||||
f"Server is not healthy at {health_url} after "
|
||||
f"{elapsed:.2f} seconds. It will be killed by the watchdog."
|
||||
),
|
||||
traceback=None,
|
||||
)
|
||||
)
|
||||
logger.info("Halting Gunicorn arbiter.")
|
||||
# Ask arbiter to stop if it's still alive.
|
||||
arbiter.signal(signal.SIGTERM, inspect.currentframe()) # type: ignore
|
||||
else:
|
||||
# If arbiter has thrown an exception, report it.
|
||||
logger.error("Gunicorn arbiter exited during health check. Watchdog exiting.")
|
||||
|
||||
else:
|
||||
# No health check; workers up => ready.
|
||||
if runtime_error is None:
|
||||
event_queue.put(ChildEvent(kind="ready"))
|
||||
else:
|
||||
# If arbiter has thrown an exception, report it.
|
||||
logger.error("Gunicorn arbiter exited unexpectedly before health check. Watchdog exiting.")
|
||||
|
||||
def _watchdog_with_exception() -> None:
|
||||
try:
|
||||
_watchdog()
|
||||
except Exception as exc:
|
||||
logger.exception("Exception in watchdog thread.")
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error", exc_type=type(exc).__name__, message=str(exc), traceback=traceback.format_exc()
|
||||
)
|
||||
)
|
||||
|
||||
watchdog_thread = threading.Thread(target=_watchdog_with_exception, daemon=True)
|
||||
watchdog_thread.start()
|
||||
|
||||
async def _serve() -> None:
|
||||
nonlocal runtime_error
|
||||
try:
|
||||
async with serve_context:
|
||||
arbiter.run()
|
||||
except Exception as exc:
|
||||
runtime_error = exc
|
||||
event_queue.put(
|
||||
ChildEvent(
|
||||
kind="error",
|
||||
exc_type=type(exc).__name__,
|
||||
message=str(exc),
|
||||
traceback=traceback.format_exc(),
|
||||
)
|
||||
)
|
||||
logger.exception("Gunicorn server failed to start.")
|
||||
|
||||
try:
|
||||
asyncio.run(_serve())
|
||||
# Most exceptions should have been caught within the _serve() coroutine.
|
||||
finally:
|
||||
# Ensure watchdog doesn't try to act on a dead arbiter for long.
|
||||
watchdog_thread.join(timeout=5.0)
|
||||
|
||||
|
||||
def _get_default_ipv4_address() -> str:
|
||||
"""Determine the default outbound IPv4 address for this machine.
|
||||
|
||||
Implementation:
|
||||
Opens a UDP socket and "connects" to a public address to force route
|
||||
selection, then inspects the socket's local address. No packets are sent.
|
||||
|
||||
Returns:
|
||||
str: Best-guess IPv4 like `192.168.x.y`. Falls back to `127.0.0.1`.
|
||||
"""
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
# Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
|
||||
s.connect(("8.8.8.8", 80))
|
||||
return s.getsockname()[0]
|
||||
except Exception:
|
||||
return "127.0.0.1"
|
||||
finally:
|
||||
s.close()
|
||||
|
||||
|
||||
class PythonServerLauncher:
|
||||
"""Unified launcher for FastAPI, using uvicorn or gunicorn per mode/worker count.
|
||||
|
||||
See [`PythonServerLauncherArgs`][agentlightning.utils.server_launcher.PythonServerLauncherArgs] for configuration options.
|
||||
|
||||
Args:
|
||||
app: The FastAPI app to launch.
|
||||
args: The configuration for the server.
|
||||
serve_context: An optional context manager to apply around the server startup.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, app: FastAPI, args: PythonServerLauncherArgs, serve_context: Optional[AsyncContextManager[Any]] = None
|
||||
):
|
||||
"""Initialize the launcher with the FastAPI app, configuration, and optional serve context."""
|
||||
self.app = app
|
||||
self.args = args
|
||||
self.serve_context = serve_context
|
||||
self._host: Optional[str] = self.args.host
|
||||
self._port: Optional[int] = self.args.port
|
||||
self._access_host: Optional[str] = self.args.access_host
|
||||
|
||||
# uvicorn (in-proc asyncio)
|
||||
self._uvicorn_server: Optional[uvicorn.Server] = None
|
||||
self._uvicorn_task: Optional[asyncio.Task[None]] = None # returned by run_uvicorn_asyncio()
|
||||
|
||||
# uvicorn (thread)
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._thread_event_queue: Optional[queue.Queue[ChildEvent]] = None
|
||||
self._thread_stop_event: Optional[threading.Event] = None
|
||||
|
||||
# subprocess (uvicorn / gunicorn)
|
||||
self._proc: Optional[BaseProcess] = None
|
||||
self._mp_event_queue: Optional[multiprocessing.Queue[ChildEvent]] = None
|
||||
self._gunicorn_app: Optional[GunicornApp] = None # programmatic gunicorn wrapper
|
||||
|
||||
# is_running flag
|
||||
self._is_running: bool = False
|
||||
|
||||
@property
|
||||
def endpoint(self) -> str:
|
||||
"""Return the externally advertised host:port pair regardless of accessibility."""
|
||||
return f"http://{self._ensure_host()}:{self._ensure_port()}"
|
||||
|
||||
@property
|
||||
def access_endpoint(self) -> str:
|
||||
"""Return a loopback-friendly URL so health checks succeed even when binding to 0.0.0.0."""
|
||||
return f"http://{self._ensure_access_host()}:{self._ensure_port()}"
|
||||
|
||||
@property
|
||||
def health_url(self) -> Optional[str]:
|
||||
"""Build the absolute health-check endpoint from args, if one is configured."""
|
||||
if not self.args.healthcheck_url:
|
||||
return None
|
||||
path = self.args.healthcheck_url
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
return f"{self.access_endpoint}{path}"
|
||||
|
||||
async def start(self):
|
||||
"""Starts the server according to launch_mode and n_workers."""
|
||||
logger.info(f"Starting server {self._normalize_app_ref(self.app)}...")
|
||||
mode = self.args.launch_mode
|
||||
if mode == "mp":
|
||||
await self._start_serving_process()
|
||||
elif mode == "thread":
|
||||
await self._start_uvicorn_thread()
|
||||
elif mode == "asyncio":
|
||||
await self._start_uvicorn_asyncio()
|
||||
else:
|
||||
raise ValueError(f"Unsupported launch mode: {mode}")
|
||||
logger.info(f"Server {self._normalize_app_ref(self.app)} started at {self.endpoint}")
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the server using the inverse of whatever launch mode was used to start it."""
|
||||
logger.info(f"Stopping server {self._normalize_app_ref(self.app)}...")
|
||||
mode = self.args.launch_mode
|
||||
if mode == "mp":
|
||||
await self._stop_serving_process()
|
||||
elif mode == "thread":
|
||||
await self._stop_uvicorn_thread()
|
||||
elif mode == "asyncio":
|
||||
await self._stop_uvicorn_asyncio()
|
||||
else:
|
||||
raise ValueError(f"Unsupported launch mode: {mode}")
|
||||
logger.info(f"Server {self._normalize_app_ref(self.app)} stopped")
|
||||
|
||||
async def reload(self):
|
||||
"""Restart the server by stopping it if necessary and invoking start again."""
|
||||
if self.is_running():
|
||||
await self.stop()
|
||||
await self.start()
|
||||
|
||||
async def run_forever(self):
|
||||
"""Start the server and block the caller until it exits, respecting the configured mode."""
|
||||
mode = self.args.launch_mode
|
||||
if mode == "asyncio":
|
||||
await self._start_uvicorn_asyncio()
|
||||
try:
|
||||
if self._uvicorn_task is not None:
|
||||
# Wait for the server
|
||||
# Won't allow outer cancel to directly cancel the inner task
|
||||
await asyncio.shield(self._uvicorn_task)
|
||||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||||
logger.warning("Server received cancellation signal. Shutting down gracefully.")
|
||||
await self._stop_uvicorn_asyncio()
|
||||
raise
|
||||
|
||||
elif mode == "thread":
|
||||
await self._start_uvicorn_thread()
|
||||
try:
|
||||
# Wait for the thread to exit
|
||||
while self._thread and self._thread.is_alive():
|
||||
await asyncio.sleep(0.5)
|
||||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||||
logger.warning("Server thread received cancellation signal. Shutting down gracefully.")
|
||||
await self._stop_uvicorn_thread()
|
||||
raise
|
||||
|
||||
elif mode == "mp":
|
||||
await self._start_serving_process()
|
||||
try:
|
||||
# Wait for the process to exit
|
||||
while self._proc and self._proc.is_alive():
|
||||
await asyncio.sleep(0.5)
|
||||
except (asyncio.CancelledError, KeyboardInterrupt):
|
||||
logger.warning("Server process received cancellation signal. Shutting down gracefully.")
|
||||
await self._stop_serving_process()
|
||||
raise
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported launch mode: {mode}")
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Return True if the server has been started and not yet stopped."""
|
||||
return self._is_running
|
||||
|
||||
@staticmethod
|
||||
def _normalize_app_ref(app: FastAPI) -> str:
|
||||
module = getattr(app, "__module__", None)
|
||||
if module and module != "__main__":
|
||||
return f"{module}:app"
|
||||
return "unknown:app"
|
||||
|
||||
def _ensure_host(self) -> str:
|
||||
if self._host is None:
|
||||
logger.warning("No host provided, using 0.0.0.0.")
|
||||
self._host = "0.0.0.0"
|
||||
return self._host
|
||||
|
||||
def _ensure_port(self) -> int:
|
||||
if self._port is None:
|
||||
logger.warning("No port provided, using pick_unused_port to pick a random unused port.")
|
||||
self._port = pick_unused_port()
|
||||
return self._port
|
||||
|
||||
def _ensure_access_host(self) -> str:
|
||||
if self.args.access_host is None:
|
||||
if self._ensure_host() in ("0.0.0.0", "::"):
|
||||
# Probe host normalization for 0.0.0.0
|
||||
logger.warning("No access host provided, using default outbound IPv4 address for this machine.")
|
||||
self._access_host = _get_default_ipv4_address()
|
||||
else:
|
||||
logger.warning("No access host provided, using the host provided.")
|
||||
self._access_host = self._ensure_host()
|
||||
else:
|
||||
self._access_host = self.args.access_host
|
||||
return self._access_host
|
||||
|
||||
def _create_uvicorn_server(self) -> uvicorn.Server:
|
||||
config = uvicorn.Config(
|
||||
app=self.app,
|
||||
host=self._ensure_host(),
|
||||
port=self._ensure_port(),
|
||||
log_level=self.args.log_level,
|
||||
loop="asyncio",
|
||||
)
|
||||
return uvicorn.Server(config)
|
||||
|
||||
def _ctx(self) -> AsyncContextManager[Any]:
|
||||
# Use the provided serve_context if any; otherwise a no-op async CM
|
||||
if self.serve_context is None:
|
||||
logger.info("No serve_context provided, using noop_context.")
|
||||
return noop_context()
|
||||
return self.serve_context
|
||||
|
||||
# --- Mode 1: asyncio (in-proc) using run_uvicorn_asyncio ---
|
||||
|
||||
async def _start_uvicorn_asyncio(self):
|
||||
if self.is_running():
|
||||
raise RuntimeError("Server is already running. Stopping it first.")
|
||||
|
||||
logger.info("Starting uvicorn asyncio server...")
|
||||
self._uvicorn_server = self._create_uvicorn_server()
|
||||
# Start server; return after health passes; keep serving in background task
|
||||
self._uvicorn_task = await run_uvicorn_asyncio(
|
||||
uvicorn_server=self._uvicorn_server,
|
||||
serve_context=self._ctx(),
|
||||
timeout=self.args.startup_timeout,
|
||||
health_url=self.health_url,
|
||||
wait_for_serve=False, # return once startup/health OK
|
||||
kill_unhealthy_server=self.args.kill_unhealthy_server,
|
||||
)
|
||||
self._is_running = True
|
||||
logger.info("Uvicorn asyncio server started")
|
||||
|
||||
async def _stop_uvicorn_asyncio(self):
|
||||
# Gracefully shut down the in-proc uvicorn server task if running
|
||||
logger.info("Stopping uvicorn asyncio server...")
|
||||
if self._uvicorn_server and self._uvicorn_task:
|
||||
await shutdown_uvicorn_server(self._uvicorn_server, self._uvicorn_task)
|
||||
self._uvicorn_task = None
|
||||
self._uvicorn_server = None
|
||||
self._is_running = False
|
||||
logger.info("Uvicorn asyncio server stopped")
|
||||
|
||||
# --- Mode 2: thread (in-proc) using run_uvicorn_thread ---
|
||||
|
||||
async def _start_uvicorn_thread(self):
|
||||
if self.is_running():
|
||||
raise RuntimeError("Server is already running. Stopping it first.")
|
||||
|
||||
logger.info("Starting uvicorn thread server...")
|
||||
self._uvicorn_server = self._create_uvicorn_server()
|
||||
self._thread_event_queue = queue.Queue()
|
||||
self._thread_stop_event = threading.Event()
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=run_uvicorn_thread,
|
||||
kwargs={
|
||||
"uvicorn_server": self._uvicorn_server,
|
||||
"serve_context": self._ctx(),
|
||||
"event_queue": self._thread_event_queue,
|
||||
"stop_event": self._thread_stop_event,
|
||||
"timeout": self.args.startup_timeout,
|
||||
"health_url": self.health_url,
|
||||
},
|
||||
daemon=True,
|
||||
)
|
||||
self._thread.start()
|
||||
|
||||
# Wait for ready or error event from the thread
|
||||
timeout = self.args.startup_timeout * 2 # Allows twice the timeout for the thread to get the event
|
||||
try:
|
||||
evt: ChildEvent = await asyncio.to_thread(self._thread_event_queue.get, True, timeout)
|
||||
except queue.Empty:
|
||||
if not self._thread.is_alive():
|
||||
logger.error("Threaded server failed to start and is not alive. No error event was received.")
|
||||
return
|
||||
logger.error("Threaded server failed to start and sends no event. This should not happen.")
|
||||
await self._stop_uvicorn_thread()
|
||||
return
|
||||
|
||||
if evt.kind == "error":
|
||||
logger.error("Threaded server failed to start (%s): %s\n%s", evt.exc_type, evt.message, evt.traceback)
|
||||
await asyncio.to_thread(self._thread.join, self.args.thread_join_timeout)
|
||||
if self._thread.is_alive():
|
||||
raise RuntimeError(evt.message or "Threaded server failed to start and refused to shut down.")
|
||||
else:
|
||||
logger.info("Threaded server started successfully.")
|
||||
self._is_running = True
|
||||
|
||||
async def _stop_uvicorn_thread(self):
|
||||
logger.info("Stopping uvicorn thread server...")
|
||||
if self._thread_stop_event:
|
||||
self._thread_stop_event.set()
|
||||
if self._thread:
|
||||
await asyncio.to_thread(self._thread.join, self.args.thread_join_timeout)
|
||||
if self._thread.is_alive():
|
||||
raise RuntimeError("Threaded server refused to shut down.")
|
||||
else:
|
||||
logger.info("Uvicorn thread server was not running. Nothing to stop.")
|
||||
|
||||
self._thread = None
|
||||
self._thread_event_queue = None
|
||||
self._thread_stop_event = None
|
||||
self._uvicorn_server = None
|
||||
self._is_running = False
|
||||
logger.info("Uvicorn thread server stopped")
|
||||
|
||||
# --- Mode 3: subprocess (uvicorn / gunicorn) using run_uvicorn_subprocess or run_gunicorn ---
|
||||
|
||||
async def _start_serving_process(self):
|
||||
if self.is_running():
|
||||
raise RuntimeError("Server process is already running. Stopping it first.")
|
||||
|
||||
host = self._ensure_host()
|
||||
port = self._ensure_port()
|
||||
|
||||
try:
|
||||
ctx = multiprocessing.get_context("fork")
|
||||
except ValueError as e:
|
||||
raise RuntimeError(
|
||||
"Process launch requires 'fork' start method (Linux/macOS). "
|
||||
"On Windows, use 'thread' or 'asyncio' modes."
|
||||
) from e
|
||||
self._mp_event_queue = ctx.Queue()
|
||||
|
||||
# Gunicorn path when n_workers > 1
|
||||
if self.args.n_workers > 1:
|
||||
logger.info(f"Starting Gunicorn server...")
|
||||
options = {
|
||||
"bind": f"{host}:{port}",
|
||||
"workers": int(self.args.n_workers),
|
||||
"worker_class": "uvicorn_worker.UvicornWorker",
|
||||
"loglevel": logging.getLevelName(self.args.log_level).lower(),
|
||||
"accesslog": None,
|
||||
"errorlog": "-",
|
||||
"preload_app": True,
|
||||
"graceful_timeout": int(
|
||||
self.args.process_join_timeout / 2
|
||||
), # Allow half the timeout for graceful shutdown
|
||||
}
|
||||
self._gunicorn_app = GunicornApp(self.app, options)
|
||||
|
||||
self._proc = ctx.Process(
|
||||
target=run_gunicorn,
|
||||
kwargs={
|
||||
"gunicorn_app": self._gunicorn_app,
|
||||
"serve_context": self._ctx(),
|
||||
"event_queue": self._mp_event_queue,
|
||||
"timeout": self.args.startup_timeout,
|
||||
"health_url": self.health_url,
|
||||
},
|
||||
daemon=False,
|
||||
)
|
||||
self._proc.start()
|
||||
|
||||
else:
|
||||
# Single-worker subprocess uvicorn
|
||||
logger.info("Starting uvicorn subprocess server...")
|
||||
self._uvicorn_server = self._create_uvicorn_server()
|
||||
|
||||
self._proc = ctx.Process(
|
||||
target=run_uvicorn_subprocess,
|
||||
kwargs={
|
||||
"uvicorn_server": self._uvicorn_server,
|
||||
"serve_context": self._ctx(),
|
||||
"event_queue": self._mp_event_queue,
|
||||
"timeout": self.args.startup_timeout,
|
||||
"health_url": self.health_url,
|
||||
},
|
||||
daemon=True,
|
||||
)
|
||||
self._proc.start()
|
||||
|
||||
# Wait for ready or error event from the thread
|
||||
timeout = self.args.startup_timeout * 2 # Allows twice the timeout for the thread to get the event
|
||||
try:
|
||||
evt: ChildEvent = await asyncio.to_thread(self._mp_event_queue.get, True, timeout)
|
||||
except queue.Empty:
|
||||
if not self._proc.is_alive():
|
||||
logger.error("Server process failed to start and is not alive. No error event was received.")
|
||||
return
|
||||
logger.error("Server process failed to start and sends no event. This should not happen.")
|
||||
await self._stop_serving_process()
|
||||
return
|
||||
|
||||
if evt.kind == "error":
|
||||
logger.error(
|
||||
"Server process (%s) failed to start (%s): %s\n%s",
|
||||
"gunicorn" if self.args.n_workers > 1 else "uvicorn",
|
||||
evt.exc_type,
|
||||
evt.message,
|
||||
evt.traceback,
|
||||
)
|
||||
await asyncio.to_thread(self._proc.join, self.args.process_join_timeout)
|
||||
if self._proc.is_alive():
|
||||
raise RuntimeError(evt.message or "Server process failed to start and refused to shut down.")
|
||||
else:
|
||||
logger.info("Subprocess server started successfully.")
|
||||
self._is_running = True
|
||||
|
||||
async def _stop_serving_process(self):
|
||||
logger.info("Stopping subprocess server...")
|
||||
if self._proc is not None:
|
||||
if self._proc.is_alive():
|
||||
# Prefer graceful: SIGTERM, then wait
|
||||
try:
|
||||
self._proc.terminate()
|
||||
except Exception:
|
||||
logger.exception("Error sending SIGTERM to server process.")
|
||||
await asyncio.to_thread(self._proc.join, self.args.process_join_timeout)
|
||||
|
||||
if self._proc.is_alive():
|
||||
# Still alive, send SIGKILL
|
||||
try:
|
||||
self._proc.kill()
|
||||
except Exception:
|
||||
logger.exception("Error sending SIGKILL to server process.")
|
||||
await asyncio.to_thread(self._proc.join, 5.0) # Use a constant timeout for SIGKILL
|
||||
|
||||
if self._proc.is_alive():
|
||||
raise RuntimeError("Server process failed to shut down after SIGTERM and SIGKILL.")
|
||||
else:
|
||||
logger.info("Subprocess server was not running. Nothing to stop.")
|
||||
|
||||
if self._mp_event_queue is not None:
|
||||
self._mp_event_queue.close()
|
||||
try:
|
||||
self._mp_event_queue.join_thread()
|
||||
except Exception:
|
||||
logger.exception("Error joining event queue thread.")
|
||||
|
||||
self._proc = None
|
||||
self._mp_event_queue = None
|
||||
self._gunicorn_app = None
|
||||
self._uvicorn_server = None
|
||||
self._is_running = False
|
||||
logger.info("Subprocess server stopped")
|
||||
@@ -294,7 +294,7 @@ class AgentModeDaemon:
|
||||
self._proxy_thread.start()
|
||||
print(f"Proxy server running on port {self.proxy_port}")
|
||||
|
||||
def _update_proxy_server_v1(self):
|
||||
async def _update_proxy_server_v1(self):
|
||||
model_name = self.train_information.get("model")
|
||||
if not model_name:
|
||||
raise ValueError("Model name is not set.")
|
||||
@@ -313,12 +313,7 @@ class AgentModeDaemon:
|
||||
],
|
||||
)
|
||||
|
||||
if self.llm_proxy.is_running():
|
||||
# FIXME: Need to switch to a different port right now
|
||||
# because the forked processes carried the old fd
|
||||
self.llm_proxy.restart(_port=_find_available_port())
|
||||
else:
|
||||
self.llm_proxy.start()
|
||||
await self.llm_proxy.restart()
|
||||
|
||||
def start(self):
|
||||
"""Starts the main AgentLightningServer and the proxy server."""
|
||||
@@ -352,7 +347,7 @@ class AgentModeDaemon:
|
||||
if server_addresses != self.backend_llm_server_addresses:
|
||||
self.backend_llm_server_addresses = server_addresses
|
||||
if self.mode == "v1" and not self.llm_proxy.is_running():
|
||||
self._update_proxy_server_v1()
|
||||
await self._update_proxy_server_v1()
|
||||
self.is_train = is_train
|
||||
|
||||
# 1. Update resources on the server for clients to use
|
||||
|
||||
@@ -2,10 +2,12 @@
|
||||
|
||||
# type: ignore
|
||||
|
||||
from importlib.metadata import version
|
||||
from typing import Any
|
||||
|
||||
import hydra
|
||||
import ray
|
||||
from packaging import version as packaging_version
|
||||
from verl.trainer.main_ppo import create_rl_sampler
|
||||
from verl.trainer.ppo.reward import load_reward_manager
|
||||
|
||||
@@ -39,11 +41,17 @@ def run_ppo(
|
||||
) -> None:
|
||||
if not ray.is_initialized():
|
||||
# this is for local ray cluster
|
||||
try:
|
||||
# verl >= 0.6.0
|
||||
num_cpus = config.ray_kwargs.ray_init.num_cpus
|
||||
except AttributeError:
|
||||
# verl < 0.6.0
|
||||
num_cpus = config.ray_init.num_cpus
|
||||
ray.init(
|
||||
runtime_env={
|
||||
"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"}
|
||||
},
|
||||
num_cpus=config.ray_init.num_cpus,
|
||||
num_cpus=num_cpus,
|
||||
)
|
||||
|
||||
runner = TaskRunner.remote()
|
||||
|
||||
@@ -19,7 +19,7 @@ from verl import DataProto
|
||||
from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
|
||||
from verl.trainer.ppo.core_algos import agg_loss
|
||||
from verl.trainer.ppo.metric_utils import (
|
||||
compute_data_metrics,
|
||||
_compute_response_info,
|
||||
compute_throughout_metrics,
|
||||
compute_timing_metrics,
|
||||
)
|
||||
@@ -53,6 +53,108 @@ def _timer(name: str, timing_raw: Dict[str, float]):
|
||||
timing_raw[name] += timer.last
|
||||
|
||||
|
||||
# This function is adapted from verl.
|
||||
# We introduce a new parameter `suffix` to distinguish between metrics computed
|
||||
# before and after AgentLightning’s post-processing.
|
||||
# - "Before" refers to raw reward and advantage values.
|
||||
# - "After" refers to values computed following post-processing, which involves:
|
||||
# (1) Dropping prompts that exceed the maximum allowed length.
|
||||
# (2) Adjusting the batch size to be a multiple of the mini PPO size.
|
||||
# Different suffixes are used to label these two stages accordingly.
|
||||
def compute_data_metrics(batch: DataProto, use_critic: bool = True, suffix: str = "") -> Dict[str, Any]:
|
||||
"""
|
||||
Computes various metrics from a batch of data for PPO training.
|
||||
|
||||
This function calculates metrics related to scores, rewards, advantages, returns, values,
|
||||
and sequence lengths from a batch of data. It provides statistical information (mean, max, min)
|
||||
for each metric category.
|
||||
|
||||
Args:
|
||||
batch: A DataProto object containing batch data with token-level scores, rewards, advantages, etc.
|
||||
use_critic: Whether to include critic-specific metrics. Defaults to True.
|
||||
|
||||
Returns:
|
||||
A dictionary of metrics including:
|
||||
- critic/score/mean, max, min: Statistics about sequence scores
|
||||
- critic/rewards/mean, max, min: Statistics about sequence rewards
|
||||
- critic/advantages/mean, max, min: Statistics about advantages
|
||||
- critic/returns/mean, max, min: Statistics about returns
|
||||
- critic/values/mean, max, min: Statistics about critic values (if use_critic=True)
|
||||
- critic/vf_explained_var: Explained variance of the value function (if use_critic=True)
|
||||
- response_length/mean, max, min, clip_ratio: Statistics about response lengths
|
||||
- prompt_length/mean, max, min, clip_ratio: Statistics about prompt lengths
|
||||
"""
|
||||
sequence_score = batch.batch["token_level_scores"].sum(-1)
|
||||
sequence_reward = batch.batch["token_level_rewards"].sum(-1)
|
||||
|
||||
advantages = batch.batch["advantages"]
|
||||
returns = batch.batch["returns"]
|
||||
|
||||
max_response_length = batch.batch["responses"].shape[-1]
|
||||
|
||||
prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool()
|
||||
response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool()
|
||||
|
||||
max_prompt_length = prompt_mask.size(-1)
|
||||
|
||||
response_info = _compute_response_info(batch)
|
||||
prompt_length = response_info["prompt_length"]
|
||||
response_length = response_info["response_length"]
|
||||
|
||||
valid_adv = torch.masked_select(advantages, response_mask)
|
||||
valid_returns = torch.masked_select(returns, response_mask)
|
||||
|
||||
if use_critic:
|
||||
values = batch.batch["values"]
|
||||
valid_values = torch.masked_select(values, response_mask)
|
||||
return_diff_var = torch.var(valid_returns - valid_values)
|
||||
return_var = torch.var(valid_returns)
|
||||
|
||||
metrics = {
|
||||
# score
|
||||
"critic/score/mean" + suffix: torch.mean(sequence_score).detach().item(),
|
||||
"critic/score/max" + suffix: torch.max(sequence_score).detach().item(),
|
||||
"critic/score/min" + suffix: torch.min(sequence_score).detach().item(),
|
||||
# reward
|
||||
"critic/rewards/mean" + suffix: torch.mean(sequence_reward).detach().item(),
|
||||
"critic/rewards/max" + suffix: torch.max(sequence_reward).detach().item(),
|
||||
"critic/rewards/min" + suffix: torch.min(sequence_reward).detach().item(),
|
||||
# adv
|
||||
"critic/advantages/mean" + suffix: torch.mean(valid_adv).detach().item(),
|
||||
"critic/advantages/max" + suffix: torch.max(valid_adv).detach().item(),
|
||||
"critic/advantages/min" + suffix: torch.min(valid_adv).detach().item(),
|
||||
# returns
|
||||
"critic/returns/mean" + suffix: torch.mean(valid_returns).detach().item(),
|
||||
"critic/returns/max" + suffix: torch.max(valid_returns).detach().item(),
|
||||
"critic/returns/min" + suffix: torch.min(valid_returns).detach().item(),
|
||||
**(
|
||||
{
|
||||
# values
|
||||
"critic/values/mean" + suffix: torch.mean(valid_values).detach().item(),
|
||||
"critic/values/max" + suffix: torch.max(valid_values).detach().item(),
|
||||
"critic/values/min" + suffix: torch.min(valid_values).detach().item(),
|
||||
# vf explained var
|
||||
"critic/vf_explained_var" + suffix: (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
|
||||
}
|
||||
if use_critic
|
||||
else {}
|
||||
),
|
||||
# response length
|
||||
"response_length/mean" + suffix: torch.mean(response_length).detach().item(),
|
||||
"response_length/max" + suffix: torch.max(response_length).detach().item(),
|
||||
"response_length/min" + suffix: torch.min(response_length).detach().item(),
|
||||
"response_length/clip_ratio"
|
||||
+ suffix: torch.mean(torch.eq(response_length, max_response_length).float()).detach().item(),
|
||||
# prompt length
|
||||
"prompt_length/mean" + suffix: torch.mean(prompt_length).detach().item(),
|
||||
"prompt_length/max" + suffix: torch.max(prompt_length).detach().item(),
|
||||
"prompt_length/min" + suffix: torch.min(prompt_length).detach().item(),
|
||||
"prompt_length/clip_ratio"
|
||||
+ suffix: torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
|
||||
}
|
||||
return metrics
|
||||
|
||||
|
||||
class AgentLightningTrainer(RayPPOTrainer):
|
||||
"""
|
||||
Specialized PPO trainer for agent-based reinforcement learning.
|
||||
@@ -215,6 +317,9 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
config=self.config.algorithm,
|
||||
)
|
||||
|
||||
# Calculate the metrics before processing. Refer to the comments of function `compute_data_metrics` for details.
|
||||
metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_before_processing"))
|
||||
|
||||
# after advantages are assinged, we begin to drop (1) long prompt (2) floor to ppo minisize
|
||||
keep_indices = (~batch.batch["is_drop_mask"]).nonzero(as_tuple=True)[0]
|
||||
metrics["training/n_triplets_prompt_too_long"] = (
|
||||
@@ -274,7 +379,7 @@ class AgentLightningTrainer(RayPPOTrainer):
|
||||
)
|
||||
|
||||
# compute training metrics
|
||||
metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
|
||||
metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic, suffix="_after_processing"))
|
||||
metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
|
||||
# TODO: implement actual tflpo and theoretical tflpo
|
||||
n_gpus = self.resource_pool_manager.get_n_gpus()
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
# Logs
|
||||
logs
|
||||
*.log
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
lerna-debug.log*
|
||||
.pnpm-debug.log*
|
||||
|
||||
# Diagnostic reports (https://nodejs.org/api/report.html)
|
||||
report.[0-9]*.[0-9]*.[0-9]*.[0-9]*.json
|
||||
|
||||
# Runtime data
|
||||
pids
|
||||
*.pid
|
||||
*.seed
|
||||
*.pid.lock
|
||||
|
||||
# Directory for instrumented libs generated by jscoverage/JSCover
|
||||
lib-cov
|
||||
|
||||
# Coverage directory used by tools like istanbul
|
||||
coverage
|
||||
*.lcov
|
||||
|
||||
# nyc test coverage
|
||||
.nyc_output
|
||||
|
||||
# Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files)
|
||||
.grunt
|
||||
|
||||
# Bower dependency directory (https://bower.io/)
|
||||
bower_components
|
||||
|
||||
# node-waf configuration
|
||||
.lock-wscript
|
||||
|
||||
# Compiled binary addons (https://nodejs.org/api/addons.html)
|
||||
build/Release
|
||||
|
||||
# Dependency directories
|
||||
node_modules/
|
||||
jspm_packages/
|
||||
|
||||
# Snowpack dependency directory (https://snowpack.dev/)
|
||||
web_modules/
|
||||
|
||||
# TypeScript cache
|
||||
*.tsbuildinfo
|
||||
|
||||
# Optional npm cache directory
|
||||
.npm
|
||||
|
||||
# Optional eslint cache
|
||||
.eslintcache
|
||||
|
||||
# Optional stylelint cache
|
||||
.stylelintcache
|
||||
|
||||
# Microbundle cache
|
||||
.rpt2_cache/
|
||||
.rts2_cache_cjs/
|
||||
.rts2_cache_es/
|
||||
.rts2_cache_umd/
|
||||
|
||||
# Optional REPL history
|
||||
.node_repl_history
|
||||
|
||||
# Output of 'npm pack'
|
||||
*.tgz
|
||||
|
||||
# Yarn Integrity file
|
||||
.yarn-integrity
|
||||
|
||||
# dotenv environment variable files
|
||||
.env
|
||||
.env.development.local
|
||||
.env.test.local
|
||||
.env.production.local
|
||||
.env.local
|
||||
|
||||
# parcel-bundler cache (https://parceljs.org/)
|
||||
.cache
|
||||
.parcel-cache
|
||||
|
||||
# Next.js build output
|
||||
.next
|
||||
out
|
||||
|
||||
# Nuxt.js build / generate output
|
||||
.nuxt
|
||||
dist
|
||||
|
||||
# Gatsby files
|
||||
.cache/
|
||||
# Comment in the public line in if your project uses Gatsby and not Next.js
|
||||
# https://nextjs.org/blog/next-9-1#public-directory-support
|
||||
# public
|
||||
|
||||
# vuepress build output
|
||||
.vuepress/dist
|
||||
|
||||
# vuepress v2.x temp and cache directory
|
||||
.temp
|
||||
.cache
|
||||
|
||||
# Docusaurus cache and generated files
|
||||
.docusaurus
|
||||
|
||||
# Serverless directories
|
||||
.serverless/
|
||||
|
||||
# FuseBox cache
|
||||
.fusebox/
|
||||
|
||||
# DynamoDB Local files
|
||||
.dynamodb/
|
||||
|
||||
# TernJS port file
|
||||
.tern-port
|
||||
|
||||
# Stores VSCode versions used for testing VSCode extensions
|
||||
.vscode-test
|
||||
|
||||
# yarn v2
|
||||
.yarn/cache
|
||||
.yarn/unplugged
|
||||
.yarn/build-state.yml
|
||||
.yarn/install-state.gz
|
||||
.pnp.*
|
||||
|
||||
.DS_Store
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
/** @type {import("@ianvs/prettier-plugin-sort-imports").PrettierConfig} */
|
||||
const config = {
|
||||
printWidth: 120,
|
||||
singleQuote: true,
|
||||
tabWidth: 2,
|
||||
useTabs: false,
|
||||
semi: true,
|
||||
quoteProps: 'consistent',
|
||||
jsxSingleQuote: true,
|
||||
trailingComma: 'all',
|
||||
bracketSpacing: true,
|
||||
objectWrap: 'preserve',
|
||||
arrowParens: 'always',
|
||||
proseWrap: 'preserve',
|
||||
endOfLine: 'lf',
|
||||
plugins: ['@ianvs/prettier-plugin-sort-imports'],
|
||||
importOrder: [
|
||||
'.*styles.css$',
|
||||
'',
|
||||
'dayjs',
|
||||
'^react$',
|
||||
'^next$',
|
||||
'^next/.*$',
|
||||
'<BUILTIN_MODULES>',
|
||||
'<THIRD_PARTY_MODULES>',
|
||||
'^@mantine/(.*)$',
|
||||
'^@mantinex/(.*)$',
|
||||
'^@mantine-tests/(.*)$',
|
||||
'^@docs/(.*)$',
|
||||
'^@/.*$',
|
||||
'^../(?!.*.css$).*$',
|
||||
'^./(?!.*.css$).*$',
|
||||
'\\.css$',
|
||||
],
|
||||
overrides: [
|
||||
{
|
||||
files: '*.mdx',
|
||||
options: {
|
||||
printWidth: 120,
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
export default config;
|
||||
@@ -0,0 +1,12 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
// Centralized constants that keep Storybook fixtures deterministic so Chromatic
|
||||
// snapshots do not drift when the build environment changes.
|
||||
export const STORY_DATE_NOW_MS = 1762775145209;
|
||||
export const STORY_DATE_NOW_SECONDS = Math.floor(STORY_DATE_NOW_MS / 1000);
|
||||
|
||||
// Use a fixed origin so any code that would normally read window.location.*
|
||||
// in the app can rely on the same value from Storybook fixtures. Prefer HTTPS
|
||||
// so Chromatic (which is served over HTTPS) avoids mixed-content fetch errors.
|
||||
export const STORY_BASE_URL = 'https://storybook.agentlightning.invalid';
|
||||
export const STORY_LOCATION_HREF = `${STORY_BASE_URL}/storybook`;
|
||||
@@ -0,0 +1,20 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { StorybookConfig } from '@storybook/react-vite';
|
||||
|
||||
const config: StorybookConfig = {
|
||||
core: {
|
||||
disableWhatsNewNotifications: true,
|
||||
disableTelemetry: true,
|
||||
enableCrashReports: false,
|
||||
},
|
||||
stories: ['../src/**/*.mdx', '../src/**/*.story.@(js|jsx|ts|tsx)'],
|
||||
staticDirs: ['../static'],
|
||||
addons: ['@storybook/addon-themes', '@storybook/addon-vitest'],
|
||||
framework: {
|
||||
name: '@storybook/react-vite',
|
||||
options: {},
|
||||
},
|
||||
};
|
||||
|
||||
export default config;
|
||||
@@ -0,0 +1,13 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export const allModes = {
|
||||
MD: {
|
||||
viewport: 'md',
|
||||
},
|
||||
LG: {
|
||||
viewport: 'lg',
|
||||
},
|
||||
XL: {
|
||||
viewport: 'xl',
|
||||
},
|
||||
} as const;
|
||||
@@ -0,0 +1,79 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import '@mantine/core/styles.css';
|
||||
import 'mantine-datatable/styles.css';
|
||||
import '../src/styles/theme.css';
|
||||
import '../src/styles/app.css';
|
||||
|
||||
import { initialize, mswLoader } from 'msw-storybook-addon';
|
||||
import { ColorSchemeScript, MantineProvider } from '@mantine/core';
|
||||
import { shadcnCssVariableResolver } from '../src/cssVariableResolver';
|
||||
import { theme as mantineTheme } from '../src/theme';
|
||||
import { STORY_DATE_NOW_MS } from './constants';
|
||||
|
||||
type ColorSchemeValue = 'light' | 'dark';
|
||||
|
||||
initialize({
|
||||
onUnhandledRequest: 'bypass',
|
||||
serviceWorker: {
|
||||
url: '/mockServiceWorker.js',
|
||||
},
|
||||
});
|
||||
|
||||
const fixedDateNow = (() => {
|
||||
const patched = Date.now as typeof Date.now & { __storybookPatched?: boolean };
|
||||
if (patched.__storybookPatched) {
|
||||
return patched;
|
||||
}
|
||||
const replacement = (() => STORY_DATE_NOW_MS) as typeof Date.now & { __storybookPatched?: boolean };
|
||||
replacement.__storybookPatched = true;
|
||||
return replacement;
|
||||
})();
|
||||
|
||||
Date.now = fixedDateNow;
|
||||
|
||||
export const parameters = {
|
||||
layout: 'fullscreen',
|
||||
options: {
|
||||
showPanel: false,
|
||||
// @ts-expect-error – storybook throws build error for (a: any, b: any)
|
||||
storySort: (a, b) => a.title.localeCompare(b.title, undefined, { numeric: true }),
|
||||
},
|
||||
backgrounds: { disable: true },
|
||||
viewport: {
|
||||
options: {
|
||||
md: { name: 'md', styles: { width: '1280px', height: '800px' } },
|
||||
lg: { name: 'lg', styles: { width: '1920px', height: '1080px' } },
|
||||
xl: { name: 'xl', styles: { width: '2560px', height: '1440px' } },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
export const globalTypes = {
|
||||
theme: {
|
||||
name: 'Theme',
|
||||
description: 'Mantine color scheme',
|
||||
defaultValue: 'light',
|
||||
toolbar: {
|
||||
icon: 'mirror',
|
||||
items: [
|
||||
{ value: 'light', title: 'Light' },
|
||||
{ value: 'dark', title: 'Dark' },
|
||||
],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
export const decorators = [
|
||||
(Story: any, context: any) => {
|
||||
const scheme = (context.parameters.theme ?? context.globals.theme ?? 'light') as ColorSchemeValue;
|
||||
return (
|
||||
<MantineProvider theme={mantineTheme} cssVariablesResolver={shadcnCssVariableResolver} forceColorScheme={scheme}>
|
||||
<ColorSchemeScript />
|
||||
<Story />
|
||||
</MantineProvider>
|
||||
);
|
||||
},
|
||||
];
|
||||
|
||||
export const loaders = [mswLoader];
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { setProjectAnnotations } from '@storybook/react-vite';
|
||||
import * as projectAnnotations from './preview';
|
||||
|
||||
// This is an important step to apply the right configuration when testing your stories.
|
||||
// More info at: https://storybook.js.org/docs/api/portable-stories/portable-stories-vitest#setprojectannotations
|
||||
setProjectAnnotations([projectAnnotations]);
|
||||
@@ -0,0 +1,5 @@
|
||||
# Generated files
|
||||
dist
|
||||
|
||||
# Theme files
|
||||
theme.css
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"extends": ["stylelint-config-standard-scss"],
|
||||
"rules": {
|
||||
"custom-property-pattern": null,
|
||||
"selector-class-pattern": null,
|
||||
"scss/no-duplicate-mixins": null,
|
||||
"declaration-empty-line-before": null,
|
||||
"declaration-block-no-redundant-longhand-properties": null,
|
||||
"alpha-value-notation": null,
|
||||
"custom-property-empty-line-before": null,
|
||||
"property-no-vendor-prefix": null,
|
||||
"color-function-notation": null,
|
||||
"length-zero-no-unit": null,
|
||||
"selector-not-notation": null,
|
||||
"no-descending-specificity": null,
|
||||
"comment-empty-line-before": null,
|
||||
"scss/at-mixin-pattern": null,
|
||||
"scss/at-rule-no-unknown": null,
|
||||
"value-keyword-case": null,
|
||||
"media-feature-range-notation": null,
|
||||
"selector-pseudo-class-no-unknown": [
|
||||
true,
|
||||
{
|
||||
"ignorePseudoClasses": ["global"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
# Agent-lightning Dashboard
|
||||
|
||||
This is the dashboard for Agent-lightning. It is a web application that allows you to inspect your Agent-lightning store and debug running experiments.
|
||||
|
||||
The dashboard is built with React, Mantine UI, and Storybook.
|
||||
|
||||
## npm scripts
|
||||
|
||||
## Build and dev scripts
|
||||
|
||||
- `dev` – start development server
|
||||
- `build` – build production version of the app
|
||||
- `preview` – locally preview production build
|
||||
|
||||
### Testing scripts
|
||||
|
||||
- `eslint` - runs ESLint
|
||||
- `stylelint` - runs Stylelint
|
||||
- `prettier` - runs Prettier
|
||||
- `typecheck` - runs TypeScript typecheck
|
||||
- `vitest` – runs vitest tests
|
||||
- `chromatic` – runs chromatic tests
|
||||
|
||||
### Other scripts
|
||||
|
||||
- `storybook` – starts storybook dev server
|
||||
- `build-storybook` – build production storybook bundle to `storybook-static`
|
||||
@@ -0,0 +1,50 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
// @ts-check
|
||||
import stylistic from '@stylistic/eslint-plugin';
|
||||
import mantine from 'eslint-config-mantine';
|
||||
import { defineConfig } from 'eslint/config';
|
||||
import tseslint from 'typescript-eslint';
|
||||
|
||||
export default defineConfig([
|
||||
// These are arrays → safe to spread
|
||||
...tseslint.configs.recommended,
|
||||
stylistic.configs.customize({ semi: true }),
|
||||
|
||||
// mantine is often a single object → include as-is (or spread only if it's actually an array)
|
||||
...(Array.isArray(mantine) ? mantine : [mantine]),
|
||||
|
||||
// ignores go as their own entry
|
||||
{ ignores: ['**/*.{mjs,cjs,js,d.ts,d.mts}'] },
|
||||
|
||||
// file-specific rules
|
||||
{
|
||||
files: ['**/*.story.tsx'],
|
||||
rules: { 'no-console': 'off' },
|
||||
},
|
||||
|
||||
// project/TS settings + your custom rules
|
||||
{
|
||||
languageOptions: {
|
||||
parserOptions: {
|
||||
tsconfigRootDir: process.cwd(),
|
||||
project: ['./tsconfig.json'],
|
||||
},
|
||||
},
|
||||
rules: {
|
||||
// Disabling conflict rules with prettier
|
||||
'@stylistic/brace-style': ['error', '1tbs', { allowSingleLine: false }],
|
||||
'@stylistic/no-trailing-spaces': 'error',
|
||||
'@stylistic/no-multiple-empty-lines': ['error', { max: 2, maxEOF: 1 }],
|
||||
'@stylistic/jsx-quotes': ['error', 'prefer-single'],
|
||||
'@stylistic/multiline-ternary': 'off',
|
||||
'@stylistic/arrow-parens': ['error', 'always'],
|
||||
'@stylistic/jsx-closing-bracket-location': 'off',
|
||||
'@stylistic/operator-linebreak': 'off',
|
||||
'@stylistic/jsx-newline': 'off',
|
||||
'@stylistic/jsx-one-expression-per-line': 'off',
|
||||
'@stylistic/indent': 'off',
|
||||
'@stylistic/indent-binary-ops': 'off',
|
||||
},
|
||||
},
|
||||
]);
|
||||
Generated
+10890
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,75 @@
|
||||
{
|
||||
"name": "agent-lightning-dashboard",
|
||||
"type": "module",
|
||||
"version": "0.2.2",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "tsc && vite build",
|
||||
"preview": "vite preview",
|
||||
"typecheck": "tsc --noEmit",
|
||||
"eslint": "eslint .",
|
||||
"stylelint": "stylelint '**/*.css'",
|
||||
"prettier": "prettier --check \"**/*.{ts,tsx,mjs,cjs}\"",
|
||||
"vitest": "vitest run --project unit",
|
||||
"vitest-storybook": "vitest run --project storybook",
|
||||
"storybook": "storybook dev -p 6006",
|
||||
"build-storybook": "storybook build",
|
||||
"chromatic": "chromatic"
|
||||
},
|
||||
"dependencies": {
|
||||
"@mantine/core": "8.3.5",
|
||||
"@mantine/hooks": "8.3.5",
|
||||
"@monaco-editor/react": "^4.7.0",
|
||||
"@reduxjs/toolkit": "^2.9.2",
|
||||
"@tabler/icons-react": "^3.35.0",
|
||||
"clsx": "^2.1.1",
|
||||
"dayjs": "^1.11.18",
|
||||
"mantine-datatable": "^8.2.0",
|
||||
"react": "^19.2.0",
|
||||
"react-dom": "^19.2.0",
|
||||
"react-redux": "^9.2.0",
|
||||
"react-router-dom": "^7.9.4"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@eslint/js": "^9.37.0",
|
||||
"@ianvs/prettier-plugin-sort-imports": "^4.7.0",
|
||||
"@storybook/addon-themes": "^9.1.10",
|
||||
"@storybook/addon-vitest": "^9.1.16",
|
||||
"@storybook/react": "^9.1.10",
|
||||
"@storybook/react-vite": "^9.1.10",
|
||||
"@stylistic/eslint-plugin": "^5.5.0",
|
||||
"@testing-library/dom": "^10.4.1",
|
||||
"@testing-library/jest-dom": "^6.9.1",
|
||||
"@testing-library/react": "^16.3.0",
|
||||
"@testing-library/user-event": "^14.6.1",
|
||||
"@types/node": "^24.7.1",
|
||||
"@types/react": "^19.2.2",
|
||||
"@types/react-dom": "^19.2.1",
|
||||
"@vitejs/plugin-react": "^5.0.4",
|
||||
"chromatic": "^13.3.3",
|
||||
"eslint": "^9.37.0",
|
||||
"eslint-config-mantine": "^4.0.3",
|
||||
"eslint-plugin-jsx-a11y": "^6.10.2",
|
||||
"eslint-plugin-react": "^7.37.5",
|
||||
"identity-obj-proxy": "^3.0.0",
|
||||
"jsdom": "^27.0.0",
|
||||
"msw": "^2.11.6",
|
||||
"msw-storybook-addon": "^2.0.6",
|
||||
"postcss": "^8.5.6",
|
||||
"postcss-preset-mantine": "1.18.0",
|
||||
"postcss-simple-vars": "^7.0.1",
|
||||
"prettier": "^3.6.2",
|
||||
"prop-types": "^15.8.1",
|
||||
"storybook": "^9.1.10",
|
||||
"stylelint": "^16.25.0",
|
||||
"stylelint-config-standard-scss": "^16.0.0",
|
||||
"typescript": "^5.9.3",
|
||||
"typescript-eslint": "^8.46.0",
|
||||
"vite": "^7.1.9",
|
||||
"vite-tsconfig-paths": "^5.1.4",
|
||||
"vitest": "^4.0.0",
|
||||
"playwright": "^1.56.1",
|
||||
"@vitest/browser-playwright": "4.0.4",
|
||||
"@vitest/coverage-v8": "4.0.4"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
module.exports = {
|
||||
plugins: {
|
||||
'postcss-preset-mantine': {},
|
||||
'postcss-simple-vars': {
|
||||
variables: {
|
||||
'mantine-breakpoint-xs': '36em',
|
||||
'mantine-breakpoint-sm': '48em',
|
||||
'mantine-breakpoint-md': '62em',
|
||||
'mantine-breakpoint-lg': '75em',
|
||||
'mantine-breakpoint-xl': '88em',
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,13 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<link rel="icon" type="image/svg+xml" href="../src/favicon.svg" />
|
||||
<meta name="viewport" content="minimum-scale=1, initial-scale=1, width=device-width, user-scalable=no" />
|
||||
<title>Agent-lightning Dashboard</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/main.tsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,3 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import '../src/main.js';
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import '@mantine/core/styles.css';
|
||||
import 'mantine-datatable/styles.css';
|
||||
import './styles/theme.css';
|
||||
import './styles/app.css';
|
||||
|
||||
import { MantineProvider } from '@mantine/core';
|
||||
import { useColorScheme } from '@mantine/hooks';
|
||||
import { shadcnCssVariableResolver } from './cssVariableResolver';
|
||||
import { selectThemePreference } from './features/config/selectors';
|
||||
import { Router } from './Router';
|
||||
import { useAppSelector } from './store/hooks';
|
||||
import { shadcnTheme } from './theme';
|
||||
|
||||
export default function App() {
|
||||
const themePreference = useAppSelector(selectThemePreference);
|
||||
const systemColorScheme = useColorScheme();
|
||||
const resolvedColorScheme = themePreference === 'system' ? systemColorScheme : themePreference;
|
||||
|
||||
return (
|
||||
<MantineProvider
|
||||
theme={shadcnTheme}
|
||||
cssVariablesResolver={shadcnCssVariableResolver}
|
||||
forceColorScheme={resolvedColorScheme}
|
||||
>
|
||||
<Router />
|
||||
</MantineProvider>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createBrowserRouter, Navigate, RouterProvider } from 'react-router-dom';
|
||||
import { AppLayoutWithState } from './layouts/AppLayout';
|
||||
import { ResourcesPage } from './pages/Resources.page';
|
||||
import { RolloutsPage } from './pages/Rollouts.page';
|
||||
import { SettingsPage } from './pages/Settings.page';
|
||||
import { TracesPage } from './pages/Traces.page';
|
||||
|
||||
const router = createBrowserRouter([
|
||||
{
|
||||
path: '/',
|
||||
element: <AppLayoutWithState />,
|
||||
children: [
|
||||
{
|
||||
index: true,
|
||||
element: <Navigate to='/rollouts' replace />,
|
||||
},
|
||||
{
|
||||
path: 'rollouts',
|
||||
element: <RolloutsPage />,
|
||||
},
|
||||
{
|
||||
path: 'resources',
|
||||
element: <ResourcesPage />,
|
||||
},
|
||||
{
|
||||
path: 'traces',
|
||||
element: <TracesPage />,
|
||||
},
|
||||
{
|
||||
path: 'settings',
|
||||
element: <SettingsPage />,
|
||||
},
|
||||
],
|
||||
},
|
||||
]);
|
||||
|
||||
export function Router() {
|
||||
return <RouterProvider router={router} />;
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { Provider } from 'react-redux';
|
||||
import { initialConfigState } from '@/features/config/slice';
|
||||
import { initialRolloutsUiState } from '@/features/rollouts/slice';
|
||||
import type { AlertsState, AlertTone } from '@/features/ui/alert';
|
||||
import { initialDrawerState } from '@/features/ui/drawer/slice';
|
||||
import { createAppStore } from '@/store';
|
||||
import { STORY_BASE_URL, STORY_DATE_NOW_MS } from '../../.storybook/constants';
|
||||
import { AppAlertBanner } from './AppAlertBanner';
|
||||
|
||||
const meta: Meta<typeof AppAlertBanner> = {
|
||||
title: 'Components/AppAlertBanner',
|
||||
component: AppAlertBanner,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
};
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof AppAlertBanner>;
|
||||
|
||||
function renderWithAlert(message: string, tone: AlertTone) {
|
||||
const alertState: AlertsState = {
|
||||
alerts: [
|
||||
{
|
||||
id: 'storybook-alert',
|
||||
message,
|
||||
tone,
|
||||
isVisible: true,
|
||||
createdAt: STORY_DATE_NOW_MS,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const store = createAppStore({
|
||||
config: {
|
||||
...initialConfigState,
|
||||
baseUrl: STORY_BASE_URL,
|
||||
},
|
||||
drawer: initialDrawerState,
|
||||
rollouts: initialRolloutsUiState,
|
||||
alert: alertState,
|
||||
});
|
||||
|
||||
return (
|
||||
<Provider store={store}>
|
||||
<div style={{ padding: 24 }}>
|
||||
<AppAlertBanner />
|
||||
</div>
|
||||
</Provider>
|
||||
);
|
||||
}
|
||||
|
||||
export const InfoAlert: Story = {
|
||||
render: () => renderWithAlert('Background synchronization completed successfully.', 'info'),
|
||||
};
|
||||
|
||||
export const WarningAlert: Story = {
|
||||
render: () =>
|
||||
renderWithAlert('Rollout data may be stale. Check your network connection before continuing.', 'warning'),
|
||||
};
|
||||
|
||||
export const ErrorAlert: Story = {
|
||||
render: () =>
|
||||
renderWithAlert('Unable to reach the Agent-lightning API. Retry or adjust the backend settings.', 'error'),
|
||||
};
|
||||
|
||||
export const NoAlert: Story = {
|
||||
render: () => {
|
||||
const store = createAppStore({
|
||||
config: {
|
||||
...initialConfigState,
|
||||
baseUrl: STORY_BASE_URL,
|
||||
},
|
||||
drawer: initialDrawerState,
|
||||
rollouts: initialRolloutsUiState,
|
||||
alert: { alerts: [] },
|
||||
});
|
||||
|
||||
return (
|
||||
<Provider store={store}>
|
||||
<div style={{ padding: 24 }}>
|
||||
<AppAlertBanner />
|
||||
</div>
|
||||
</Provider>
|
||||
);
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,84 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useEffect, useState } from 'react';
|
||||
import { IconAlertCircle, IconAlertTriangle, IconInfoCircle } from '@tabler/icons-react';
|
||||
import { Notification, Portal, Transition } from '@mantine/core';
|
||||
import { hideAlert, selectHighestPriorityAlert, type AppAlert } from '@/features/ui/alert';
|
||||
import { useAppDispatch, useAppSelector } from '@/store/hooks';
|
||||
|
||||
const ALERT_META = {
|
||||
info: {
|
||||
color: 'blue',
|
||||
icon: IconInfoCircle,
|
||||
},
|
||||
warning: {
|
||||
color: 'yellow',
|
||||
icon: IconAlertTriangle,
|
||||
},
|
||||
error: {
|
||||
color: 'red',
|
||||
icon: IconAlertCircle,
|
||||
},
|
||||
} as const;
|
||||
|
||||
export function AppAlertBanner() {
|
||||
const dispatch = useAppDispatch();
|
||||
const alert = useAppSelector(selectHighestPriorityAlert);
|
||||
const [transitionAlert, setTransitionAlert] = useState<AppAlert | null>(alert);
|
||||
|
||||
useEffect(() => {
|
||||
if (alert) {
|
||||
setTransitionAlert(alert);
|
||||
}
|
||||
}, [alert]);
|
||||
|
||||
const handleClose = (id?: string) => {
|
||||
if (id) {
|
||||
dispatch(hideAlert({ id }));
|
||||
}
|
||||
};
|
||||
|
||||
const currentAlert = alert ?? transitionAlert;
|
||||
const mounted = Boolean(alert);
|
||||
|
||||
if (!currentAlert) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const meta = ALERT_META[currentAlert.tone];
|
||||
const IconComponent = meta.icon;
|
||||
|
||||
return (
|
||||
<Portal>
|
||||
<Transition
|
||||
mounted={mounted}
|
||||
transition='slide-down'
|
||||
duration={200}
|
||||
timingFunction='ease'
|
||||
onExited={() => setTransitionAlert(null)}
|
||||
>
|
||||
{(styles) => (
|
||||
<Notification
|
||||
icon={<IconComponent size={18} />}
|
||||
color={meta.color}
|
||||
variant='light'
|
||||
withCloseButton
|
||||
onClose={() => handleClose(currentAlert.id)}
|
||||
style={{
|
||||
position: 'fixed',
|
||||
top: 16,
|
||||
right: 16,
|
||||
maxWidth: 450,
|
||||
width: 'calc(100% - 32px)',
|
||||
zIndex: 2000,
|
||||
boxShadow: 'var(--mantine-shadow-md)',
|
||||
...styles,
|
||||
}}
|
||||
>
|
||||
{currentAlert.message}
|
||||
</Notification>
|
||||
)}
|
||||
</Transition>
|
||||
</Portal>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,516 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from 'react';
|
||||
import { Editor } from '@monaco-editor/react';
|
||||
import { IconCheck, IconCopy } from '@tabler/icons-react';
|
||||
import type { DataTableSortStatus } from 'mantine-datatable';
|
||||
import { createSearchParams, Link, useInRouterContext, useLocation } from 'react-router-dom';
|
||||
import {
|
||||
ActionIcon,
|
||||
Anchor,
|
||||
Badge,
|
||||
Box,
|
||||
CopyButton,
|
||||
Drawer,
|
||||
Group,
|
||||
Stack,
|
||||
Text,
|
||||
Tooltip,
|
||||
useMantineColorScheme,
|
||||
} from '@mantine/core';
|
||||
import { useGetSpansQuery } from '@/features/rollouts';
|
||||
import { closeDrawer, openDrawer, selectDrawerContent, selectDrawerIsOpen } from '@/features/ui/drawer';
|
||||
import { useAppDispatch, useAppSelector } from '@/store/hooks';
|
||||
import type { Attempt, AttemptStatus, Rollout, RolloutStatus, Span } from '@/types';
|
||||
import { formatStatusLabel } from '@/utils/format';
|
||||
import { TracesTable, type TracesTableRecord } from './TracesTable.component';
|
||||
|
||||
const ATTEMPT_STATUS_COLORS: Record<AttemptStatus, string> = {
|
||||
failed: 'red',
|
||||
preparing: 'violet',
|
||||
running: 'blue',
|
||||
succeeded: 'teal',
|
||||
timeout: 'orange',
|
||||
unresponsive: 'orange',
|
||||
};
|
||||
|
||||
const ROLLOUT_STATUS_COLORS: Record<RolloutStatus, string> = {
|
||||
cancelled: 'gray',
|
||||
failed: 'red',
|
||||
preparing: 'violet',
|
||||
queuing: 'blue',
|
||||
requeuing: 'cyan',
|
||||
running: 'blue',
|
||||
succeeded: 'teal',
|
||||
};
|
||||
|
||||
const SPAN_STATUS_COLORS: Record<Span['status']['status_code'], string> = {
|
||||
UNSET: 'gray',
|
||||
OK: 'teal',
|
||||
ERROR: 'red',
|
||||
};
|
||||
|
||||
const TRACES_SORT_FIELD_MAP: Record<string, string> = {
|
||||
name: 'name',
|
||||
traceId: 'trace_id',
|
||||
spanId: 'span_id',
|
||||
parentId: 'parent_id',
|
||||
statusCode: 'status_code',
|
||||
startTime: 'start_time',
|
||||
duration: 'duration',
|
||||
};
|
||||
|
||||
type SortDirection = 'asc' | 'desc';
|
||||
|
||||
type LocalSortState = {
|
||||
column: string;
|
||||
direction: SortDirection;
|
||||
};
|
||||
|
||||
function resolveTracesSortField(column: string): string {
|
||||
return TRACES_SORT_FIELD_MAP[column] ?? 'start_time';
|
||||
}
|
||||
|
||||
function getStatusBadgeColor(status: RolloutStatus | AttemptStatus, isAttempt: boolean) {
|
||||
if (isAttempt) {
|
||||
return ATTEMPT_STATUS_COLORS[status as AttemptStatus] ?? 'gray';
|
||||
}
|
||||
|
||||
return ROLLOUT_STATUS_COLORS[status as RolloutStatus] ?? 'gray';
|
||||
}
|
||||
|
||||
function formatJson(value: unknown) {
|
||||
try {
|
||||
return JSON.stringify(value, null, 2);
|
||||
} catch {
|
||||
return String(value);
|
||||
}
|
||||
}
|
||||
|
||||
export type AppDrawerProps = {
|
||||
opened: boolean;
|
||||
onClose: () => void;
|
||||
title?: ReactNode;
|
||||
body?: ReactNode;
|
||||
};
|
||||
|
||||
export function AppDrawer({ opened, onClose, title, body }: AppDrawerProps) {
|
||||
return (
|
||||
<Drawer
|
||||
position='right'
|
||||
size='lg'
|
||||
opened={opened}
|
||||
onClose={onClose}
|
||||
overlayProps={{ opacity: 0.5 }}
|
||||
withinPortal
|
||||
styles={{
|
||||
content: {
|
||||
display: 'flex',
|
||||
flexDirection: 'column',
|
||||
maxHeight: '100vh',
|
||||
},
|
||||
body: {
|
||||
flex: 1,
|
||||
display: 'flex',
|
||||
flexDirection: 'column',
|
||||
padding: 'var(--mantine-spacing-md)',
|
||||
minHeight: 0,
|
||||
overflow: 'hidden',
|
||||
},
|
||||
}}
|
||||
title={title}
|
||||
>
|
||||
<Stack gap='md' h='100%' style={{ flex: 1, minHeight: 0 }}>
|
||||
{body}
|
||||
</Stack>
|
||||
</Drawer>
|
||||
);
|
||||
}
|
||||
|
||||
type TraceDrawerTitleProps = {
|
||||
span: Span;
|
||||
};
|
||||
|
||||
export function TraceDrawerTitle({ span }: TraceDrawerTitleProps) {
|
||||
const spanStatusCode = span.status?.status_code ?? null;
|
||||
const spanBadgeColor = spanStatusCode ? (SPAN_STATUS_COLORS[spanStatusCode] ?? 'gray') : undefined;
|
||||
|
||||
return (
|
||||
<Stack gap={3}>
|
||||
<Group gap={6}>
|
||||
<Text fw={600}>{span.name ?? span.spanId}</Text>
|
||||
{spanStatusCode ? (
|
||||
<Badge size='sm' variant='light' color={spanBadgeColor}>
|
||||
{spanStatusCode}
|
||||
</Badge>
|
||||
) : null}
|
||||
</Group>
|
||||
<Group gap={6}>
|
||||
<Text size='sm' c='dimmed'>
|
||||
{span.spanId}
|
||||
</Text>
|
||||
<CopyButton value={span.spanId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy span ID ${span.spanId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
<Group gap='xs'>
|
||||
<Group gap={3}>
|
||||
<Text size='sm' c='dimmed' fw={500}>
|
||||
Rollout
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed'>
|
||||
{span.rolloutId}
|
||||
</Text>
|
||||
</Group>
|
||||
<Group gap={3}>
|
||||
<Text size='sm' c='dimmed' fw={500}>
|
||||
Attempt
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed'>
|
||||
{span.attemptId ?? '—'}
|
||||
</Text>
|
||||
</Group>
|
||||
</Group>
|
||||
</Stack>
|
||||
);
|
||||
}
|
||||
|
||||
type RolloutAttemptDrawerTitleProps = {
|
||||
rollout: Rollout;
|
||||
attempt: Attempt | null;
|
||||
};
|
||||
|
||||
export function RolloutAttemptDrawerTitle({ rollout, attempt }: RolloutAttemptDrawerTitleProps) {
|
||||
const rolloutId = rollout.rolloutId;
|
||||
const attemptId = attempt?.attemptId ?? null;
|
||||
const rolloutStatus = rollout.status ?? null;
|
||||
const attemptStatus = attempt?.status ?? null;
|
||||
const rolloutStatusLabel = rolloutStatus ? formatStatusLabel(rolloutStatus) : null;
|
||||
const attemptStatusLabel = attemptStatus ? formatStatusLabel(attemptStatus) : null;
|
||||
const hasStatusMismatch = rolloutStatus !== null && attemptStatus !== null && rolloutStatus !== attemptStatus;
|
||||
const rolloutBadgeColor = rolloutStatus ? getStatusBadgeColor(rolloutStatus, false) : undefined;
|
||||
const attemptBadgeColor = attemptStatus ? getStatusBadgeColor(attemptStatus, true) : undefined;
|
||||
const showRolloutBadgeInHeading = Boolean(rolloutStatusLabel && (!attemptStatus || hasStatusMismatch));
|
||||
const showAttemptBadge = Boolean(attemptStatusLabel && attemptStatus);
|
||||
|
||||
return (
|
||||
<Stack gap={3}>
|
||||
<Group gap={6}>
|
||||
<Text fw={600}>{rolloutId}</Text>
|
||||
<CopyButton value={rolloutId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy rollout ID ${rolloutId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
{showRolloutBadgeInHeading && rolloutStatusLabel ? (
|
||||
<Badge size='sm' variant='light' color={rolloutBadgeColor}>
|
||||
{rolloutStatusLabel}
|
||||
</Badge>
|
||||
) : null}
|
||||
</Group>
|
||||
<Group gap='xs'>
|
||||
{attemptId ? (
|
||||
<Group gap={3}>
|
||||
<Text size='sm' c='dimmed' fw={500}>
|
||||
Attempt
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed'>
|
||||
{attemptId}
|
||||
</Text>
|
||||
</Group>
|
||||
) : null}
|
||||
{showAttemptBadge && attemptStatusLabel ? (
|
||||
<Badge size='sm' variant='light' color={attemptBadgeColor}>
|
||||
{attemptStatusLabel}
|
||||
</Badge>
|
||||
) : null}
|
||||
{!showRolloutBadgeInHeading && !attemptStatus && rolloutStatusLabel ? (
|
||||
<Badge size='sm' variant='light' color={rolloutBadgeColor}>
|
||||
{rolloutStatusLabel}
|
||||
</Badge>
|
||||
) : null}
|
||||
</Group>
|
||||
</Stack>
|
||||
);
|
||||
}
|
||||
|
||||
type JsonEditorProps = {
|
||||
value: unknown;
|
||||
};
|
||||
|
||||
export function JsonEditor({ value }: JsonEditorProps) {
|
||||
const { colorScheme } = useMantineColorScheme();
|
||||
const editorTheme = colorScheme === 'dark' ? 'vs-dark' : 'vs-light';
|
||||
|
||||
return (
|
||||
<Box data-testid='json-editor-container' style={{ flex: 1, minHeight: 0 }}>
|
||||
<Editor
|
||||
height='100%'
|
||||
language='json'
|
||||
value={formatJson(value)}
|
||||
theme={editorTheme}
|
||||
options={{
|
||||
readOnly: true,
|
||||
domReadOnly: true,
|
||||
minimap: { enabled: false },
|
||||
automaticLayout: true,
|
||||
scrollBeyondLastLine: false,
|
||||
fontSize: 13,
|
||||
}}
|
||||
/>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
type RolloutTracesDrawerBodyProps = {
|
||||
rollout: Rollout;
|
||||
attempt: Attempt | null;
|
||||
onShowRollout: (record: TracesTableRecord) => void;
|
||||
onShowSpanDetail: (record: TracesTableRecord) => void;
|
||||
};
|
||||
|
||||
function RolloutTracesDrawerBody({ rollout, attempt, onShowRollout, onShowSpanDetail }: RolloutTracesDrawerBodyProps) {
|
||||
const [page, setPage] = useState(1);
|
||||
const [recordsPerPage, setRecordsPerPage] = useState(100);
|
||||
const [sort, setSort] = useState<LocalSortState>({
|
||||
column: 'startTime',
|
||||
direction: 'desc',
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
setPage(1);
|
||||
}, [rollout.rolloutId, attempt?.attemptId]);
|
||||
|
||||
const queryArgs = useMemo(
|
||||
() => ({
|
||||
rolloutId: rollout.rolloutId,
|
||||
attemptId: attempt?.attemptId ?? undefined,
|
||||
limit: recordsPerPage,
|
||||
offset: Math.max(0, (page - 1) * recordsPerPage),
|
||||
sortBy: resolveTracesSortField(sort.column),
|
||||
sortOrder: sort.direction,
|
||||
}),
|
||||
[rollout.rolloutId, attempt?.attemptId, recordsPerPage, page, sort],
|
||||
);
|
||||
|
||||
const { data, isFetching, isError, error, refetch } = useGetSpansQuery(queryArgs);
|
||||
const spans = data?.items ?? [];
|
||||
const totalRecords = data?.total ?? 0;
|
||||
const tracesLinkSearch = useMemo(() => {
|
||||
const params = createSearchParams({
|
||||
rolloutId: rollout.rolloutId,
|
||||
...(attempt?.attemptId ? { attemptId: attempt.attemptId } : {}),
|
||||
});
|
||||
return params.toString();
|
||||
}, [attempt?.attemptId, rollout.rolloutId]);
|
||||
const tracesLinkHref = tracesLinkSearch ? `/traces?${tracesLinkSearch}` : '/traces';
|
||||
const isWithinRouter = useInRouterContext();
|
||||
|
||||
const handleSortStatusChange = useCallback((status: DataTableSortStatus<TracesTableRecord>) => {
|
||||
setSort({
|
||||
column: status.columnAccessor as string,
|
||||
direction: status.direction,
|
||||
});
|
||||
}, []);
|
||||
|
||||
const handlePageChange = useCallback((nextPage: number) => {
|
||||
setPage(nextPage);
|
||||
}, []);
|
||||
|
||||
const handleRecordsPerPageChange = useCallback((value: number) => {
|
||||
setRecordsPerPage(value);
|
||||
setPage(1);
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<Stack gap='md' style={{ flex: 1, minHeight: 0 }}>
|
||||
<Group justify='space-between' align='center' gap='sm' wrap='nowrap'>
|
||||
<Text size='sm' style={{ flex: 1, minWidth: 0 }}>
|
||||
Showing spans for{' '}
|
||||
<Text component='span' fw={600}>
|
||||
{rollout.rolloutId}
|
||||
{attempt ? ` · Attempt ${attempt.sequenceId} (${attempt.attemptId})` : ' · Latest attempt'}
|
||||
</Text>
|
||||
</Text>
|
||||
{isWithinRouter ? (
|
||||
<Anchor
|
||||
component={Link}
|
||||
to={tracesLinkHref}
|
||||
size='sm'
|
||||
aria-label={`Open traces page for rollout ${rollout.rolloutId}${
|
||||
attempt ? ` attempt ${attempt.sequenceId}` : ''
|
||||
}`}
|
||||
>
|
||||
View full traces
|
||||
</Anchor>
|
||||
) : (
|
||||
<Anchor
|
||||
href={tracesLinkHref}
|
||||
size='sm'
|
||||
aria-label={`Open traces page for rollout ${rollout.rolloutId}${
|
||||
attempt ? ` attempt ${attempt.sequenceId}` : ''
|
||||
}`}
|
||||
>
|
||||
View full traces
|
||||
</Anchor>
|
||||
)}
|
||||
</Group>
|
||||
<Box data-testid='traces-drawer-table-container' style={{ flex: 1, minHeight: 0, overflow: 'auto' }}>
|
||||
<TracesTable
|
||||
spans={spans}
|
||||
totalRecords={totalRecords}
|
||||
isFetching={isFetching}
|
||||
isError={isError}
|
||||
error={error}
|
||||
searchTerm=''
|
||||
sort={sort}
|
||||
page={page}
|
||||
recordsPerPage={recordsPerPage}
|
||||
onSortStatusChange={handleSortStatusChange}
|
||||
onPageChange={handlePageChange}
|
||||
onRecordsPerPageChange={handleRecordsPerPageChange}
|
||||
onResetFilters={() => {}}
|
||||
onRefetch={refetch}
|
||||
onShowRollout={onShowRollout}
|
||||
onShowSpanDetail={onShowSpanDetail}
|
||||
recordsPerPageOptions={[50, 100, 200, 500]}
|
||||
/>
|
||||
</Box>
|
||||
</Stack>
|
||||
);
|
||||
}
|
||||
|
||||
export function AppDrawerContainer() {
|
||||
const dispatch = useAppDispatch();
|
||||
const isOpen = useAppSelector(selectDrawerIsOpen);
|
||||
const content = useAppSelector(selectDrawerContent);
|
||||
const isRouterAvailable = useInRouterContext();
|
||||
|
||||
const handleClose = useCallback(() => {
|
||||
dispatch(closeDrawer());
|
||||
}, [dispatch]);
|
||||
const handleNavigation = useCallback(() => {
|
||||
if (isOpen) {
|
||||
dispatch(closeDrawer());
|
||||
}
|
||||
}, [dispatch, isOpen]);
|
||||
|
||||
const derivedContent = useMemo(() => {
|
||||
if (!content) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (content.type === 'trace-detail') {
|
||||
const { span } = content;
|
||||
const title = <TraceDrawerTitle span={span} />;
|
||||
const body = <JsonEditor value={span} />;
|
||||
|
||||
return { title, body };
|
||||
}
|
||||
|
||||
const rollout = content.rollout;
|
||||
const attempt = content.attempt;
|
||||
const title = <RolloutAttemptDrawerTitle rollout={rollout} attempt={attempt} />;
|
||||
|
||||
if (content.type === 'rollout-json') {
|
||||
const jsonValue = content.isNested && content.attempt ? content.attempt : rollout;
|
||||
const body = jsonValue ? <JsonEditor value={jsonValue} /> : null;
|
||||
return { title, body };
|
||||
}
|
||||
|
||||
if (content.type === 'rollout-traces') {
|
||||
const body = (
|
||||
<RolloutTracesDrawerBody
|
||||
rollout={rollout}
|
||||
attempt={attempt}
|
||||
onShowRollout={() => {
|
||||
const attemptForRecord = attempt ?? rollout.attempt ?? null;
|
||||
dispatch(
|
||||
openDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout,
|
||||
attempt: attemptForRecord,
|
||||
isNested: content.isNested,
|
||||
}),
|
||||
);
|
||||
}}
|
||||
onShowSpanDetail={(record) => {
|
||||
const attemptForRecord = attempt ?? rollout.attempt ?? null;
|
||||
dispatch(
|
||||
openDrawer({
|
||||
type: 'trace-detail',
|
||||
span: record,
|
||||
rollout,
|
||||
attempt: attemptForRecord,
|
||||
}),
|
||||
);
|
||||
}}
|
||||
/>
|
||||
);
|
||||
|
||||
return { title, body };
|
||||
}
|
||||
|
||||
return null;
|
||||
}, [content, dispatch]);
|
||||
|
||||
if (!content || !derivedContent) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const { title, body } = derivedContent;
|
||||
|
||||
return (
|
||||
<>
|
||||
{isRouterAvailable ? <DrawerLocationWatcher onNavigation={handleNavigation} /> : null}
|
||||
<AppDrawer opened={isOpen} onClose={handleClose} title={title} body={body} />
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
type DrawerLocationWatcherProps = {
|
||||
onNavigation: () => void;
|
||||
};
|
||||
|
||||
function DrawerLocationWatcher({ onNavigation }: DrawerLocationWatcherProps) {
|
||||
const location = useLocation();
|
||||
const lastLocationKeyRef = useRef(location.key);
|
||||
|
||||
useEffect(() => {
|
||||
if (lastLocationKeyRef.current === location.key) {
|
||||
return;
|
||||
}
|
||||
lastLocationKeyRef.current = location.key;
|
||||
onNavigation();
|
||||
}, [location.key, onNavigation]);
|
||||
|
||||
return null;
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { Provider } from 'react-redux';
|
||||
import { initialConfigState } from '@/features/config/slice';
|
||||
import { initialResourcesUiState } from '@/features/resources/slice';
|
||||
import { rolloutsApi } from '@/features/rollouts';
|
||||
import { initialRolloutsUiState } from '@/features/rollouts/slice';
|
||||
import { initialTracesUiState } from '@/features/traces/slice';
|
||||
import type { DrawerContent } from '@/features/ui/drawer';
|
||||
import { createAppStore } from '@/store';
|
||||
import type { Attempt, Rollout, Span } from '@/types';
|
||||
import { STORY_BASE_URL, STORY_DATE_NOW_SECONDS } from '../../.storybook/constants';
|
||||
import { AppDrawerContainer } from './AppDrawer.component';
|
||||
|
||||
const meta = {
|
||||
title: 'Components/AppDrawer',
|
||||
component: AppDrawerContainer,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
} satisfies Meta<typeof AppDrawerContainer>;
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof AppDrawerContainer>;
|
||||
|
||||
const now = STORY_DATE_NOW_SECONDS;
|
||||
|
||||
const baseAttempt: Attempt = {
|
||||
rolloutId: 'ro-story-001',
|
||||
attemptId: 'at-story-001',
|
||||
sequenceId: 1,
|
||||
startTime: now - 3600,
|
||||
endTime: null,
|
||||
status: 'running',
|
||||
workerId: 'worker-story',
|
||||
lastHeartbeatTime: now - 42,
|
||||
metadata: { info: 'Sample metadata', runId: 'run-123' },
|
||||
};
|
||||
|
||||
const baseRollout: Rollout = {
|
||||
rolloutId: 'ro-story-001',
|
||||
input: {
|
||||
task: 'Generate daily summary',
|
||||
payload: { account: 'enterprise', date: '2024-02-19' },
|
||||
},
|
||||
startTime: now - 4000,
|
||||
endTime: null,
|
||||
mode: 'train',
|
||||
resourcesId: 'rs-story-001',
|
||||
status: 'running',
|
||||
config: { retries: 1, priority: 'high' },
|
||||
metadata: { owner: 'storybook' },
|
||||
attempt: baseAttempt,
|
||||
};
|
||||
|
||||
const noAttemptRollout: Rollout = {
|
||||
...baseRollout,
|
||||
status: 'queuing',
|
||||
attempt: null,
|
||||
};
|
||||
|
||||
const mismatchRollout: Rollout = {
|
||||
...baseRollout,
|
||||
status: 'running',
|
||||
attempt: {
|
||||
...baseAttempt,
|
||||
status: 'failed',
|
||||
endTime: now - 1200,
|
||||
metadata: { info: 'Latest attempt failed', reason: 'Timeout' },
|
||||
},
|
||||
};
|
||||
|
||||
const sampleSpan: Span = {
|
||||
rolloutId: 'ro-story-001',
|
||||
attemptId: 'at-story-001',
|
||||
sequenceId: 2,
|
||||
traceId: 'tr-story-001',
|
||||
spanId: 'sp-story-001',
|
||||
parentId: null,
|
||||
name: 'Fetch Resources',
|
||||
status: { status_code: 'OK', description: 'Completed successfully' },
|
||||
attributes: {
|
||||
'http.method': 'GET',
|
||||
'http.url': 'https://api.example.com/resources',
|
||||
'duration_ms': 120,
|
||||
},
|
||||
startTime: now - 240,
|
||||
endTime: now - 120,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
};
|
||||
|
||||
const sampleTraces: Span[] = [
|
||||
sampleSpan,
|
||||
{
|
||||
...sampleSpan,
|
||||
spanId: 'sp-story-002',
|
||||
name: 'Process Response',
|
||||
parentId: 'sp-story-001',
|
||||
sequenceId: 3,
|
||||
status: { status_code: 'ERROR', description: 'Unexpected response code' },
|
||||
attributes: {
|
||||
...sampleSpan.attributes,
|
||||
duration_ms: 240,
|
||||
},
|
||||
startTime: now - 120,
|
||||
endTime: now - 30,
|
||||
},
|
||||
];
|
||||
|
||||
function renderWithDrawer(content: DrawerContent, options?: { spans?: Span[] }) {
|
||||
const store = createAppStore({
|
||||
config: {
|
||||
...initialConfigState,
|
||||
baseUrl: STORY_BASE_URL,
|
||||
},
|
||||
rollouts: initialRolloutsUiState,
|
||||
resources: initialResourcesUiState,
|
||||
traces: initialTracesUiState,
|
||||
drawer: {
|
||||
isOpen: true,
|
||||
content,
|
||||
},
|
||||
});
|
||||
|
||||
if (content.type === 'rollout-traces' && options?.spans) {
|
||||
const defaultLimit = 100;
|
||||
const queryArgs = {
|
||||
rolloutId: content.rollout.rolloutId,
|
||||
attemptId: content.attempt?.attemptId ?? undefined,
|
||||
limit: defaultLimit,
|
||||
offset: 0,
|
||||
sortBy: 'start_time',
|
||||
sortOrder: 'desc' as const,
|
||||
};
|
||||
|
||||
store.dispatch(
|
||||
rolloutsApi.util.upsertQueryData('getSpans', queryArgs, {
|
||||
items: options.spans,
|
||||
total: options.spans.length,
|
||||
limit: defaultLimit,
|
||||
offset: 0,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<Provider store={store}>
|
||||
<AppDrawerContainer />
|
||||
</Provider>
|
||||
);
|
||||
}
|
||||
|
||||
export const RolloutJson: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout: baseRollout,
|
||||
attempt: baseRollout.attempt,
|
||||
isNested: false,
|
||||
}),
|
||||
};
|
||||
|
||||
export const NestedAttemptJson: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout: baseRollout,
|
||||
attempt: {
|
||||
...baseAttempt,
|
||||
attemptId: 'at-story-002',
|
||||
sequenceId: 2,
|
||||
status: 'failed',
|
||||
endTime: now - 1200,
|
||||
metadata: { info: 'Secondary attempt', reason: 'Timeout' },
|
||||
},
|
||||
isNested: true,
|
||||
}),
|
||||
};
|
||||
|
||||
export const RolloutTraces: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer(
|
||||
{
|
||||
type: 'rollout-traces',
|
||||
rollout: baseRollout,
|
||||
attempt: baseRollout.attempt,
|
||||
isNested: false,
|
||||
},
|
||||
{ spans: sampleTraces },
|
||||
),
|
||||
};
|
||||
|
||||
export const NoAttempt: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout: noAttemptRollout,
|
||||
attempt: null,
|
||||
isNested: false,
|
||||
}),
|
||||
};
|
||||
|
||||
export const StatusMismatch: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout: mismatchRollout,
|
||||
attempt: mismatchRollout.attempt,
|
||||
isNested: false,
|
||||
}),
|
||||
};
|
||||
|
||||
export const SpanDetail: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'trace-detail',
|
||||
span: sampleSpan,
|
||||
rollout: mismatchRollout,
|
||||
attempt: mismatchRollout.attempt,
|
||||
}),
|
||||
};
|
||||
|
||||
export const LightTheme: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'rollout-json',
|
||||
rollout: baseRollout,
|
||||
attempt: baseRollout.attempt,
|
||||
isNested: false,
|
||||
}),
|
||||
parameters: {
|
||||
theme: 'light',
|
||||
},
|
||||
};
|
||||
|
||||
export const DarkTheme: Story = {
|
||||
render: () =>
|
||||
renderWithDrawer({
|
||||
type: 'trace-detail',
|
||||
span: {
|
||||
...sampleSpan,
|
||||
spanId: 'sp-story-002',
|
||||
name: 'Process Response',
|
||||
status: { status_code: 'ERROR', description: 'Unexpected response code' },
|
||||
},
|
||||
rollout: mismatchRollout,
|
||||
attempt: mismatchRollout.attempt,
|
||||
}),
|
||||
parameters: {
|
||||
theme: 'dark',
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,322 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useCallback, useEffect, useMemo, useState, type ReactNode, type SetStateAction } from 'react';
|
||||
import { IconCheck, IconCopy, IconRefresh } from '@tabler/icons-react';
|
||||
import { DataTable, type DataTableColumn, type DataTableSortStatus } from 'mantine-datatable';
|
||||
import { ActionIcon, Box, Button, CopyButton, Group, Stack, Text, Tooltip } from '@mantine/core';
|
||||
import { useElementSize, useViewportSize } from '@mantine/hooks';
|
||||
import { getLayoutAwareWidth } from '@/layouts/helper';
|
||||
import type { Resources } from '@/types';
|
||||
import { getErrorDescriptor } from '@/utils/error';
|
||||
import { formatDateTime, safeStringify } from '@/utils/format';
|
||||
import { createResponsiveColumns, type ColumnVisibilityConfig } from '@/utils/table';
|
||||
|
||||
const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
|
||||
|
||||
const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
|
||||
resourcesId: { fixedWidth: 12, priority: 0 },
|
||||
version: { fixedWidth: 8, priority: 1 },
|
||||
createTime: { fixedWidth: 14, priority: 2 },
|
||||
updateTime: { fixedWidth: 14, priority: 2 },
|
||||
resourceCount: { fixedWidth: 8, priority: 3 },
|
||||
resourcesPreview: { minWidth: 16, priority: 4 },
|
||||
};
|
||||
|
||||
export type ResourcesTableRecord = Resources & {
|
||||
resourceCount: number;
|
||||
canExpand: boolean;
|
||||
resourcesPreview: string;
|
||||
};
|
||||
|
||||
function buildResourcesRecord(resources: Resources): ResourcesTableRecord {
|
||||
const resourceCount = Object.keys(resources.resources ?? {}).length;
|
||||
const resourcesValue =
|
||||
resources.resources === null || typeof resources.resources === 'undefined'
|
||||
? '—'
|
||||
: typeof resources.resources === 'string'
|
||||
? resources.resources
|
||||
: safeStringify(resources.resources);
|
||||
|
||||
return {
|
||||
...resources,
|
||||
resourceCount,
|
||||
canExpand: resourceCount > 0,
|
||||
resourcesPreview: resourcesValue,
|
||||
};
|
||||
}
|
||||
|
||||
type ResourcesColumnsOptions = Record<string, never>;
|
||||
|
||||
function createResourcesColumns(_options: ResourcesColumnsOptions): DataTableColumn<ResourcesTableRecord>[] {
|
||||
return [
|
||||
{
|
||||
accessor: 'resourcesId',
|
||||
title: 'Resources ID',
|
||||
sortable: true,
|
||||
render: ({ resourcesId }) => (
|
||||
<Group gap={2}>
|
||||
<Text fw={500} size='sm'>
|
||||
{resourcesId}
|
||||
</Text>
|
||||
<CopyButton value={resourcesId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy resources ID ${resourcesId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'version',
|
||||
title: 'Version',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ version }) => <Text size='sm'>{version}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'createTime',
|
||||
title: 'Created',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ createTime }) => <Text size='sm'>{formatDateTime(createTime)}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'updateTime',
|
||||
title: 'Updated',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ updateTime }) => <Text size='sm'>{formatDateTime(updateTime)}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'resourceCount',
|
||||
title: 'Count',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ resourceCount }) => <Text size='sm'>{resourceCount}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'resourcesPreview',
|
||||
title: 'Preview',
|
||||
render: ({ resourcesPreview }) => (
|
||||
<Text size='sm' ff='monospace' c='dimmed' lineClamp={1} style={{ width: '100%' }}>
|
||||
{resourcesPreview}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
type RowExpansionRenderer = (context: {
|
||||
resources: Resources;
|
||||
columns: DataTableColumn<ResourcesTableRecord>[];
|
||||
}) => ReactNode;
|
||||
|
||||
export type ResourcesTableProps = {
|
||||
resourcesList: Resources[] | undefined;
|
||||
totalRecords: number;
|
||||
isFetching: boolean;
|
||||
isError: boolean;
|
||||
error: unknown;
|
||||
searchTerm: string;
|
||||
sort: { column: string; direction: 'asc' | 'desc' };
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
onSortStatusChange: (status: DataTableSortStatus<ResourcesTableRecord>) => void;
|
||||
onPageChange: (page: number) => void;
|
||||
onRecordsPerPageChange: (value: number) => void;
|
||||
onResetFilters: () => void;
|
||||
onRefetch: () => void;
|
||||
recordsPerPageOptions?: number[];
|
||||
renderRowExpansion?: RowExpansionRenderer;
|
||||
};
|
||||
|
||||
export function ResourcesTable({
|
||||
resourcesList,
|
||||
totalRecords,
|
||||
isFetching,
|
||||
isError,
|
||||
error,
|
||||
searchTerm,
|
||||
sort,
|
||||
page,
|
||||
recordsPerPage,
|
||||
onSortStatusChange,
|
||||
onPageChange,
|
||||
onRecordsPerPageChange,
|
||||
onResetFilters,
|
||||
onRefetch,
|
||||
recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
|
||||
renderRowExpansion,
|
||||
}: ResourcesTableProps) {
|
||||
const [expandedRecordIds, setExpandedRecordIds] = useState<string[]>([]);
|
||||
const { ref: tableContainerRef, width: containerWidth } = useElementSize();
|
||||
const { width: viewportWidth } = useViewportSize();
|
||||
|
||||
const layoutAwareContainerWidth = useMemo(
|
||||
() => getLayoutAwareWidth(containerWidth, viewportWidth),
|
||||
[containerWidth, viewportWidth],
|
||||
);
|
||||
|
||||
const resourcesRecords = useMemo<ResourcesTableRecord[]>(() => {
|
||||
if (!resourcesList) {
|
||||
return [];
|
||||
}
|
||||
return resourcesList.map((resourcesItem) => buildResourcesRecord(resourcesItem));
|
||||
}, [resourcesList]);
|
||||
|
||||
const columns = useMemo(() => createResourcesColumns({}), []);
|
||||
|
||||
const responsiveColumns = useMemo(
|
||||
() => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
|
||||
[columns, layoutAwareContainerWidth],
|
||||
);
|
||||
|
||||
const totalPages = useMemo(
|
||||
() => Math.max(1, Math.ceil(Math.max(0, totalRecords) / Math.max(1, recordsPerPage))),
|
||||
[recordsPerPage, totalRecords],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (page > totalPages) {
|
||||
onPageChange(totalPages);
|
||||
}
|
||||
}, [onPageChange, page, totalPages]);
|
||||
|
||||
useEffect(() => {
|
||||
setExpandedRecordIds((current) =>
|
||||
current.filter((id) => resourcesRecords.some((record) => record.resourcesId === id && record.canExpand)),
|
||||
);
|
||||
}, [resourcesRecords]);
|
||||
|
||||
const hasActiveFilters = searchTerm.trim().length > 0;
|
||||
|
||||
const sortStatus: DataTableSortStatus<ResourcesTableRecord> = {
|
||||
columnAccessor: sort.column,
|
||||
direction: sort.direction,
|
||||
};
|
||||
|
||||
const handleSortStatusChange = useCallback(
|
||||
(status: DataTableSortStatus<ResourcesTableRecord>) => {
|
||||
onSortStatusChange(status);
|
||||
},
|
||||
[onSortStatusChange],
|
||||
);
|
||||
|
||||
const errorDescriptor = isError ? getErrorDescriptor(error) : null;
|
||||
const errorMessage = isError
|
||||
? `Resources are temporarily unavailable${errorDescriptor ? ` (${errorDescriptor})` : ''}.`
|
||||
: 'Resources are temporarily unavailable.';
|
||||
|
||||
const emptyState = (
|
||||
<Stack gap='sm' align='center' py='lg'>
|
||||
{isError ? (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
{errorMessage}
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
Use the retry button to try again, or adjust the filters to broaden the results.
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' color='gray' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Retry
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
No resources found
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
{hasActiveFilters
|
||||
? 'Try adjusting the search to see more results.'
|
||||
: 'Try refreshing to fetch the latest resources.'}
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Refresh
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
)}
|
||||
</Stack>
|
||||
);
|
||||
|
||||
return (
|
||||
<Box ref={tableContainerRef}>
|
||||
<DataTable<ResourcesTableRecord>
|
||||
classNames={{ root: 'resources-table' }}
|
||||
withTableBorder
|
||||
withColumnBorders
|
||||
highlightOnHover
|
||||
verticalAlign='center'
|
||||
minHeight={resourcesRecords.length === 0 ? 500 : undefined}
|
||||
idAccessor='resourcesId'
|
||||
records={resourcesRecords}
|
||||
columns={responsiveColumns}
|
||||
totalRecords={totalRecords}
|
||||
recordsPerPage={recordsPerPage}
|
||||
page={page}
|
||||
onPageChange={onPageChange}
|
||||
onRecordsPerPageChange={onRecordsPerPageChange}
|
||||
recordsPerPageOptions={recordsPerPageOptions}
|
||||
sortStatus={sortStatus}
|
||||
onSortStatusChange={handleSortStatusChange}
|
||||
fetching={isFetching}
|
||||
loaderSize='sm'
|
||||
emptyState={resourcesRecords.length === 0 ? emptyState : undefined}
|
||||
rowExpansion={
|
||||
renderRowExpansion
|
||||
? {
|
||||
allowMultiple: true,
|
||||
expandable: ({ record }) => record.canExpand,
|
||||
expanded: {
|
||||
recordIds: expandedRecordIds,
|
||||
onRecordIdsChange: (nextRecordIds: SetStateAction<string[]>) => {
|
||||
setExpandedRecordIds((previous) => {
|
||||
const resolved =
|
||||
typeof nextRecordIds === 'function'
|
||||
? nextRecordIds(previous)
|
||||
: ((nextRecordIds ?? []) as (string | number)[]);
|
||||
return resolved
|
||||
.map(String)
|
||||
.filter((id) =>
|
||||
resourcesRecords.some(
|
||||
(tableRecord) => tableRecord.resourcesId === id && tableRecord.canExpand,
|
||||
),
|
||||
);
|
||||
});
|
||||
},
|
||||
},
|
||||
content: ({ record }) => renderRowExpansion({ resources: record, columns: responsiveColumns }),
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useMemo, useState } from 'react';
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { IconSearch } from '@tabler/icons-react';
|
||||
import { Box, Stack, TextInput, Title } from '@mantine/core';
|
||||
import type { Resources } from '@/types';
|
||||
import { ResourcesTable } from './ResourcesTable.component';
|
||||
import { ResourcesTree } from './ResourcesTree.component';
|
||||
|
||||
const meta: Meta<typeof ResourcesTable> = {
|
||||
title: 'Components/ResourcesTable',
|
||||
component: ResourcesTable,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
};
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof ResourcesTable>;
|
||||
|
||||
const sampleResources: Resources[] = [
|
||||
{
|
||||
resourcesId: 'rs-story-001',
|
||||
version: 1,
|
||||
createTime: 1710806400,
|
||||
updateTime: 1713412800,
|
||||
resources: {
|
||||
model: {
|
||||
name: 'gpt-4',
|
||||
version: '2024-01-01',
|
||||
temperature: 0.7,
|
||||
maxTokens: 2048,
|
||||
topP: 0.9,
|
||||
},
|
||||
database: {
|
||||
host: 'db.example.com',
|
||||
port: 5432,
|
||||
name: 'production',
|
||||
pool: {
|
||||
min: 2,
|
||||
max: 10,
|
||||
idle: 30000,
|
||||
},
|
||||
},
|
||||
cache: {
|
||||
type: 'redis',
|
||||
host: 'cache.example.com',
|
||||
port: 6379,
|
||||
ttl: 3600,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
resourcesId: 'rs-story-002',
|
||||
version: 2,
|
||||
createTime: 1712217600,
|
||||
updateTime: 1714823200,
|
||||
resources: {
|
||||
model: {
|
||||
name: 'claude-3-opus',
|
||||
version: '2024-02-01',
|
||||
temperature: 0.5,
|
||||
maxTokens: 4096,
|
||||
},
|
||||
storage: {
|
||||
type: 's3',
|
||||
bucket: 'training-data',
|
||||
region: 'us-east-1',
|
||||
credentials: {
|
||||
accessKeyId: 'AKIA***',
|
||||
encrypted: true,
|
||||
},
|
||||
},
|
||||
compute: {
|
||||
instances: [
|
||||
{ id: 'i-001', type: 't3.large', zone: 'us-east-1a' },
|
||||
{ id: 'i-002', type: 't3.large', zone: 'us-east-1b' },
|
||||
],
|
||||
autoScaling: {
|
||||
min: 2,
|
||||
max: 10,
|
||||
targetCpu: 70,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
resourcesId: 'rs-story-003',
|
||||
version: 3,
|
||||
createTime: 1709251200,
|
||||
updateTime: 1711856800,
|
||||
resources: {
|
||||
model: {
|
||||
name: 'gpt-3.5-turbo',
|
||||
version: '2023-12-01',
|
||||
temperature: 0.8,
|
||||
maxTokens: 1024,
|
||||
},
|
||||
monitoring: {
|
||||
enabled: true,
|
||||
interval: 60,
|
||||
metrics: ['cpu', 'memory', 'disk', 'network'],
|
||||
alerts: {
|
||||
email: 'ops@example.com',
|
||||
slack: '#alerts',
|
||||
pagerduty: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
resourcesId: 'rs-story-004',
|
||||
version: 1,
|
||||
createTime: 1706745600,
|
||||
updateTime: 1709347200,
|
||||
resources: {
|
||||
apiKeys: {
|
||||
openai: 'sk-***',
|
||||
anthropic: 'sk-ant-***',
|
||||
replicate: 'r8-***',
|
||||
},
|
||||
rateLimits: {
|
||||
requestsPerMinute: 100,
|
||||
tokensPerDay: 1000000,
|
||||
concurrent: 5,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
resourcesId: 'rs-story-005',
|
||||
version: 1,
|
||||
createTime: 1704067200,
|
||||
updateTime: 1706668800,
|
||||
resources: {},
|
||||
},
|
||||
];
|
||||
|
||||
type WrapperProps = {
|
||||
maxWidth: number;
|
||||
resourcesList?: Resources[] | undefined;
|
||||
isFetching?: boolean;
|
||||
isError?: boolean;
|
||||
error?: unknown;
|
||||
};
|
||||
|
||||
function ResourcesTableStoryWrapper({
|
||||
maxWidth,
|
||||
resourcesList = sampleResources,
|
||||
isFetching = false,
|
||||
isError = false,
|
||||
error = null,
|
||||
}: WrapperProps) {
|
||||
const [searchTerm, setSearchTerm] = useState('');
|
||||
const [page, setPage] = useState(1);
|
||||
const [recordsPerPage, setRecordsPerPage] = useState(5);
|
||||
const [sort, setSort] = useState<{ column: string; direction: 'asc' | 'desc' }>({
|
||||
column: 'resourcesId',
|
||||
direction: 'asc',
|
||||
});
|
||||
|
||||
const baseResources = resourcesList ?? [];
|
||||
|
||||
const filteredResources = useMemo(() => {
|
||||
const normalized = searchTerm.trim().toLowerCase();
|
||||
if (normalized.length === 0) {
|
||||
return baseResources;
|
||||
}
|
||||
return baseResources.filter((resource) => resource.resourcesId.toLowerCase().includes(normalized));
|
||||
}, [baseResources, searchTerm]);
|
||||
|
||||
const sortedResources = useMemo(() => {
|
||||
const items = filteredResources.slice();
|
||||
const resolveSortValue = (resource: Resources, column: string) => {
|
||||
switch (column) {
|
||||
case 'version':
|
||||
return resource.version;
|
||||
case 'createTime':
|
||||
return resource.createTime;
|
||||
case 'updateTime':
|
||||
return resource.updateTime;
|
||||
case 'resourceCount':
|
||||
return Object.keys(resource.resources ?? {}).length;
|
||||
case 'resourcesId':
|
||||
default:
|
||||
return resource.resourcesId;
|
||||
}
|
||||
};
|
||||
|
||||
items.sort((a, b) => {
|
||||
const aValue = resolveSortValue(a, sort.column);
|
||||
const bValue = resolveSortValue(b, sort.column);
|
||||
|
||||
if (aValue === bValue) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (typeof aValue === 'number' && typeof bValue === 'number') {
|
||||
return aValue - bValue;
|
||||
}
|
||||
|
||||
return String(aValue).localeCompare(String(bValue));
|
||||
});
|
||||
|
||||
if (sort.direction === 'desc') {
|
||||
items.reverse();
|
||||
}
|
||||
|
||||
return items;
|
||||
}, [filteredResources, sort]);
|
||||
|
||||
const totalRecordsValue = sortedResources.length;
|
||||
|
||||
const pagedResources = useMemo(() => {
|
||||
const startIndex = (page - 1) * recordsPerPage;
|
||||
return sortedResources.slice(startIndex, startIndex + recordsPerPage);
|
||||
}, [page, recordsPerPage, sortedResources]);
|
||||
|
||||
return (
|
||||
<Box mx='auto' style={{ maxWidth, width: '100%', padding: 16 }}>
|
||||
<Stack gap='md'>
|
||||
<Title order={2}>Resources</Title>
|
||||
<TextInput
|
||||
placeholder='Search by Resources ID'
|
||||
value={searchTerm}
|
||||
onChange={(event) => {
|
||||
setSearchTerm(event.currentTarget.value);
|
||||
setPage(1);
|
||||
}}
|
||||
leftSection={<IconSearch size={16} />}
|
||||
data-testid='resources-search-input'
|
||||
w='100%'
|
||||
style={{ maxWidth: 360 }}
|
||||
/>
|
||||
<ResourcesTable
|
||||
resourcesList={pagedResources}
|
||||
totalRecords={totalRecordsValue}
|
||||
isFetching={isFetching}
|
||||
isError={isError}
|
||||
error={error}
|
||||
searchTerm={searchTerm}
|
||||
sort={sort}
|
||||
page={page}
|
||||
recordsPerPage={recordsPerPage}
|
||||
onSortStatusChange={(status) => {
|
||||
setSort({
|
||||
column: status.columnAccessor as string,
|
||||
direction: status.direction,
|
||||
});
|
||||
}}
|
||||
onPageChange={setPage}
|
||||
onRecordsPerPageChange={(value) => {
|
||||
setRecordsPerPage(value);
|
||||
setPage(1);
|
||||
}}
|
||||
onResetFilters={() => {
|
||||
setSearchTerm('');
|
||||
setSort({ column: 'resourcesId', direction: 'asc' });
|
||||
setPage(1);
|
||||
}}
|
||||
onRefetch={() => undefined}
|
||||
recordsPerPageOptions={[5, 10, 20]}
|
||||
renderRowExpansion={({ resources }) => <ResourcesTree resources={resources} />}
|
||||
/>
|
||||
</Stack>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
export const WideContainer: Story = {
|
||||
render: () => <ResourcesTableStoryWrapper maxWidth={1280} />,
|
||||
};
|
||||
|
||||
export const MediumContainer: Story = {
|
||||
render: () => <ResourcesTableStoryWrapper maxWidth={960} />,
|
||||
};
|
||||
|
||||
export const NarrowContainer: Story = {
|
||||
render: () => <ResourcesTableStoryWrapper maxWidth={720} />,
|
||||
};
|
||||
|
||||
export const DrawerWidth: Story = {
|
||||
render: () => <ResourcesTableStoryWrapper maxWidth={520} />,
|
||||
};
|
||||
|
||||
export const ErrorState: Story = {
|
||||
render: () => (
|
||||
<ResourcesTableStoryWrapper maxWidth={600} resourcesList={[]} isError error={new Error('Network unreachable')} />
|
||||
),
|
||||
};
|
||||
|
||||
export const EmptyResources: Story = {
|
||||
render: () => (
|
||||
<ResourcesTableStoryWrapper
|
||||
maxWidth={960}
|
||||
resourcesList={[
|
||||
{
|
||||
resourcesId: 'rs-empty-001',
|
||||
version: 1,
|
||||
createTime: 1702000000,
|
||||
updateTime: 1704600000,
|
||||
resources: {},
|
||||
},
|
||||
]}
|
||||
/>
|
||||
),
|
||||
};
|
||||
|
||||
export const LoadingState: Story = {
|
||||
render: () => <ResourcesTableStoryWrapper maxWidth={960} resourcesList={[]} isFetching />,
|
||||
};
|
||||
@@ -0,0 +1,129 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useMemo } from 'react';
|
||||
import { IconAlertCircle, IconChevronRight } from '@tabler/icons-react';
|
||||
import { Box, Group, Stack, Text, Tree, type TreeNodeData } from '@mantine/core';
|
||||
import type { Resources } from '@/types';
|
||||
import { safeStringify } from '@/utils/format';
|
||||
|
||||
function convertToTreeData(obj: any, key: string = 'root', parentPath = ''): TreeNodeData {
|
||||
const isObject = obj !== null && typeof obj === 'object' && !Array.isArray(obj);
|
||||
const isArray = Array.isArray(obj);
|
||||
const currentPath = parentPath ? `${parentPath}.${key}` : key;
|
||||
|
||||
if (isObject) {
|
||||
const children = Object.entries(obj).map(([childKey, childValue]) =>
|
||||
convertToTreeData(childValue, childKey, currentPath),
|
||||
);
|
||||
|
||||
return {
|
||||
value: currentPath,
|
||||
label: (
|
||||
<Group gap={6}>
|
||||
<Text size='sm' fw={500}>
|
||||
{key}
|
||||
</Text>
|
||||
<Text size='xs' c='dimmed'>
|
||||
(Object)
|
||||
</Text>
|
||||
</Group>
|
||||
),
|
||||
children: children.length > 0 ? children : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
if (isArray) {
|
||||
const children = obj.map((item: any, index: number) => convertToTreeData(item, `[${index}]`, currentPath));
|
||||
|
||||
return {
|
||||
value: currentPath,
|
||||
label: (
|
||||
<Group gap={6}>
|
||||
<Text size='sm' fw={500}>
|
||||
{key}
|
||||
</Text>
|
||||
<Text size='xs' c='dimmed'>
|
||||
(Array[
|
||||
{obj.length}
|
||||
])
|
||||
</Text>
|
||||
</Group>
|
||||
),
|
||||
children: children.length > 0 ? children : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
// Primitive value
|
||||
return {
|
||||
value: currentPath,
|
||||
label: (
|
||||
<Group gap={6}>
|
||||
<Text size='sm' fw={500}>
|
||||
{key}:
|
||||
</Text>
|
||||
<Text size='sm' ff='monospace' c='dimmed'>
|
||||
{safeStringify(obj)}
|
||||
</Text>
|
||||
</Group>
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
export type ResourcesTreeProps = {
|
||||
resources: Resources;
|
||||
};
|
||||
|
||||
export function ResourcesTree({ resources }: ResourcesTreeProps) {
|
||||
const resourcesDict = resources.resources ?? {};
|
||||
|
||||
const treeData = useMemo<TreeNodeData[]>(() => {
|
||||
const entries = Object.entries(resourcesDict);
|
||||
|
||||
if (entries.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
return entries.map(([key, value]) => convertToTreeData(value, key));
|
||||
}, [resourcesDict]);
|
||||
|
||||
if (treeData.length === 0) {
|
||||
return (
|
||||
<Stack gap='xs' align='center' py='md'>
|
||||
<IconAlertCircle size={24} color='gray' />
|
||||
<Text size='sm' c='dimmed'>
|
||||
No resources found
|
||||
</Text>
|
||||
</Stack>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<Box p='md' style={{ backgroundColor: 'var(--mantine-color-default-hover)' }}>
|
||||
<Tree
|
||||
data={treeData}
|
||||
levelOffset={20}
|
||||
expandOnClick
|
||||
selectOnClick
|
||||
renderNode={({ node, expanded, hasChildren, elementProps }) => (
|
||||
<Group gap={4} {...elementProps}>
|
||||
{hasChildren && (
|
||||
<Box
|
||||
style={{
|
||||
minWidth: 14,
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
justifyContent: 'center',
|
||||
transform: expanded ? 'rotate(90deg)' : 'rotate(0deg)',
|
||||
transition: 'transform 150ms ease',
|
||||
}}
|
||||
>
|
||||
<IconChevronRight size={14} />
|
||||
</Box>
|
||||
)}
|
||||
{node.label}
|
||||
</Group>
|
||||
)}
|
||||
/>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { Box, Stack, Title } from '@mantine/core';
|
||||
import type { Resources } from '@/types';
|
||||
import { ResourcesTree } from './ResourcesTree.component';
|
||||
|
||||
const meta: Meta<typeof ResourcesTree> = {
|
||||
title: 'Components/ResourcesTree',
|
||||
component: ResourcesTree,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
};
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof ResourcesTree>;
|
||||
|
||||
const simpleResources: Resources = {
|
||||
resourcesId: 'rs-simple-001',
|
||||
version: 1,
|
||||
createTime: 1704067200,
|
||||
updateTime: 1706668800,
|
||||
resources: {
|
||||
apiKey: { value: 'sk-test-key-123', type: 'secret' },
|
||||
maxRetries: { value: 3, description: 'Maximum retry attempts' },
|
||||
timeout: { value: 30000, unit: 'ms' },
|
||||
enabled: { value: true },
|
||||
},
|
||||
};
|
||||
|
||||
const nestedResources: Resources = {
|
||||
resourcesId: 'rs-nested-001',
|
||||
version: 2,
|
||||
createTime: 1709251200,
|
||||
updateTime: 1711856800,
|
||||
resources: {
|
||||
model: {
|
||||
name: 'gpt-4',
|
||||
version: '2024-01-01',
|
||||
temperature: 0.7,
|
||||
maxTokens: 2048,
|
||||
topP: 0.9,
|
||||
},
|
||||
database: {
|
||||
host: 'db.example.com',
|
||||
port: 5432,
|
||||
name: 'production',
|
||||
pool: {
|
||||
min: 2,
|
||||
max: 10,
|
||||
idle: 30000,
|
||||
},
|
||||
},
|
||||
cache: {
|
||||
type: 'redis',
|
||||
host: 'cache.example.com',
|
||||
port: 6379,
|
||||
ttl: 3600,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const arrayResources: Resources = {
|
||||
resourcesId: 'rs-array-001',
|
||||
version: 3,
|
||||
createTime: 1712217600,
|
||||
updateTime: 1714823200,
|
||||
resources: {
|
||||
compute: {
|
||||
instances: [
|
||||
{ id: 'i-001', type: 't3.large', zone: 'us-east-1a', status: 'running' },
|
||||
{ id: 'i-002', type: 't3.large', zone: 'us-east-1b', status: 'running' },
|
||||
{ id: 'i-003', type: 't3.xlarge', zone: 'us-east-1c', status: 'stopped' },
|
||||
],
|
||||
autoScaling: {
|
||||
min: 2,
|
||||
max: 10,
|
||||
targetCpu: 70,
|
||||
},
|
||||
},
|
||||
tags: ['production', 'ml-training', 'auto-scale'],
|
||||
ports: [80, 443, 8080],
|
||||
},
|
||||
};
|
||||
|
||||
const complexResources: Resources = {
|
||||
resourcesId: 'rs-complex-001',
|
||||
version: 4,
|
||||
createTime: 1706745600,
|
||||
updateTime: 1710000000,
|
||||
resources: {
|
||||
model: {
|
||||
name: 'claude-3-opus',
|
||||
version: '2024-02-01',
|
||||
temperature: 0.5,
|
||||
maxTokens: 4096,
|
||||
providers: [
|
||||
{ name: 'anthropic', priority: 1, enabled: true },
|
||||
{ name: 'aws-bedrock', priority: 2, enabled: false },
|
||||
],
|
||||
},
|
||||
storage: {
|
||||
type: 's3',
|
||||
bucket: 'training-data',
|
||||
region: 'us-east-1',
|
||||
credentials: {
|
||||
accessKeyId: 'AKIA***',
|
||||
encrypted: true,
|
||||
},
|
||||
lifecycle: {
|
||||
transitionToIA: 30,
|
||||
transitionToGlacier: 90,
|
||||
expiration: 365,
|
||||
},
|
||||
},
|
||||
monitoring: {
|
||||
enabled: true,
|
||||
interval: 60,
|
||||
metrics: ['cpu', 'memory', 'disk', 'network'],
|
||||
alerts: {
|
||||
email: 'ops@example.com',
|
||||
slack: '#alerts',
|
||||
pagerduty: true,
|
||||
thresholds: {
|
||||
cpu: { warning: 70, critical: 90 },
|
||||
memory: { warning: 80, critical: 95 },
|
||||
disk: { warning: 75, critical: 90 },
|
||||
},
|
||||
},
|
||||
},
|
||||
apiKeys: {
|
||||
openai: 'sk-***',
|
||||
anthropic: 'sk-ant-***',
|
||||
replicate: 'r8-***',
|
||||
},
|
||||
rateLimits: {
|
||||
requestsPerMinute: 100,
|
||||
tokensPerDay: 1000000,
|
||||
concurrent: 5,
|
||||
burstMultiplier: 1.5,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const emptyResources: Resources = {
|
||||
resourcesId: 'rs-empty-001',
|
||||
version: 1,
|
||||
createTime: 1702000000,
|
||||
updateTime: 1704600000,
|
||||
resources: {},
|
||||
};
|
||||
|
||||
type WrapperProps = {
|
||||
resources: Resources;
|
||||
maxWidth?: number;
|
||||
};
|
||||
|
||||
function ResourcesTreeStoryWrapper({ resources, maxWidth = 800 }: WrapperProps) {
|
||||
return (
|
||||
<Box mx='auto' style={{ maxWidth, width: '100%', padding: 16 }}>
|
||||
<Stack gap='md'>
|
||||
<Title order={2}>
|
||||
Resources:
|
||||
{resources.resourcesId}
|
||||
</Title>
|
||||
<ResourcesTree resources={resources} />
|
||||
</Stack>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
export const SimpleValues: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={simpleResources} />,
|
||||
};
|
||||
|
||||
export const NestedObjects: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={nestedResources} />,
|
||||
};
|
||||
|
||||
export const WithArrays: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={arrayResources} />,
|
||||
};
|
||||
|
||||
export const ComplexStructure: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={complexResources} maxWidth={1000} />,
|
||||
};
|
||||
|
||||
export const EmptyResources: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={emptyResources} />,
|
||||
};
|
||||
|
||||
export const NarrowContainer: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={complexResources} maxWidth={500} />,
|
||||
};
|
||||
|
||||
export const WideContainer: Story = {
|
||||
render: () => <ResourcesTreeStoryWrapper resources={complexResources} maxWidth={1400} />,
|
||||
};
|
||||
@@ -0,0 +1,787 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useCallback, useEffect, useMemo, useState, type ReactNode, type SetStateAction } from 'react';
|
||||
import {
|
||||
IconAlertCircle,
|
||||
IconCheck,
|
||||
IconCopy,
|
||||
IconFileDescription,
|
||||
IconRefresh,
|
||||
IconReload,
|
||||
IconTimeline,
|
||||
} from '@tabler/icons-react';
|
||||
import { DataTable, type DataTableColumn, type DataTableSortStatus } from 'mantine-datatable';
|
||||
import {
|
||||
ActionIcon,
|
||||
Alert,
|
||||
Badge,
|
||||
Box,
|
||||
Button,
|
||||
CopyButton,
|
||||
Group,
|
||||
MultiSelect,
|
||||
Stack,
|
||||
Text,
|
||||
Tooltip,
|
||||
} from '@mantine/core';
|
||||
import { useElementSize, useViewportSize } from '@mantine/hooks';
|
||||
import {
|
||||
type Attempt,
|
||||
type AttemptStatus,
|
||||
type Rollout,
|
||||
type RolloutMode,
|
||||
type RolloutsSortState,
|
||||
type RolloutStatus,
|
||||
} from '@/features/rollouts';
|
||||
import { getLayoutAwareWidth } from '@/layouts/helper';
|
||||
import {
|
||||
clampToNow,
|
||||
formatDateTime,
|
||||
formatDuration,
|
||||
formatRelativeTime,
|
||||
formatStatusLabel,
|
||||
safeStringify,
|
||||
toTimestamp,
|
||||
} from '@/utils/format';
|
||||
import { createResponsiveColumns, type ColumnVisibilityConfig } from '@/utils/table';
|
||||
|
||||
const ROLLOUT_STATUS_OPTIONS: RolloutStatus[] = [
|
||||
'queuing',
|
||||
'preparing',
|
||||
'running',
|
||||
'failed',
|
||||
'succeeded',
|
||||
'cancelled',
|
||||
'requeuing',
|
||||
];
|
||||
|
||||
const ATTEMPT_STATUS_COLORS: Record<AttemptStatus, string> = {
|
||||
failed: 'red',
|
||||
preparing: 'violet',
|
||||
running: 'blue',
|
||||
succeeded: 'teal',
|
||||
timeout: 'orange',
|
||||
unresponsive: 'orange',
|
||||
};
|
||||
|
||||
const ROLLOUT_STATUS_COLORS: Record<RolloutStatus, string> = {
|
||||
cancelled: 'gray',
|
||||
failed: 'red',
|
||||
preparing: 'violet',
|
||||
queuing: 'gray',
|
||||
requeuing: 'gray',
|
||||
running: 'blue',
|
||||
succeeded: 'teal',
|
||||
};
|
||||
|
||||
const ROLLOUT_MODE_OPTIONS: RolloutMode[] = ['train', 'val', 'test'];
|
||||
|
||||
const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
|
||||
|
||||
const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
|
||||
rolloutId: { fixedWidth: 12.5, priority: 0 },
|
||||
actionsPlaceholder: { fixedWidth: 6.5, priority: 0 },
|
||||
inputText: { minWidth: 14, priority: 1 },
|
||||
statusValue: { fixedWidth: 10, priority: 1 },
|
||||
startTimestamp: { fixedWidth: 12, priority: 2 },
|
||||
durationSeconds: { fixedWidth: 10, priority: 2 },
|
||||
attemptId: { fixedWidth: 12, priority: 3 },
|
||||
resourcesId: { fixedWidth: 10, priority: 3 },
|
||||
mode: { fixedWidth: 8, priority: 3 },
|
||||
lastHeartbeatTimestamp: { fixedWidth: 10, priority: 3 },
|
||||
workerId: { fixedWidth: 10, priority: 3 },
|
||||
};
|
||||
|
||||
export type RolloutTableRecord = Rollout & {
|
||||
attemptId: string | null;
|
||||
attemptSequence: number | null;
|
||||
isNested: boolean;
|
||||
canExpand: boolean;
|
||||
inputText: string;
|
||||
attemptStatus?: AttemptStatus;
|
||||
statusValue: string;
|
||||
startTimestamp: number | null;
|
||||
durationSeconds: number | null;
|
||||
lastHeartbeatTimestamp: number | null;
|
||||
workerId: string | null;
|
||||
actionsPlaceholder?: null;
|
||||
};
|
||||
|
||||
function selectHeartbeatTimestamp(attempt?: Attempt | null): number | null {
|
||||
if (!attempt || attempt.lastHeartbeatTime == null || Number.isNaN(attempt.lastHeartbeatTime)) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return attempt.lastHeartbeatTime;
|
||||
}
|
||||
|
||||
export function buildRolloutRecord(rollout: Rollout): RolloutTableRecord {
|
||||
const latestAttempt = rollout.attempt;
|
||||
const inputValue =
|
||||
rollout.input === null || typeof rollout.input === 'undefined'
|
||||
? '—'
|
||||
: typeof rollout.input === 'string'
|
||||
? rollout.input
|
||||
: safeStringify(rollout.input);
|
||||
const startTimestamp = toTimestamp(latestAttempt?.startTime ?? rollout.startTime);
|
||||
const endTimestamp = toTimestamp(latestAttempt?.endTime ?? rollout.endTime);
|
||||
const durationSeconds = clampToNow(startTimestamp, endTimestamp);
|
||||
const attemptStatus = latestAttempt?.status;
|
||||
const sequenceId = latestAttempt?.sequenceId;
|
||||
const statusValue =
|
||||
attemptStatus && attemptStatus !== rollout.status ? `${rollout.status}-${attemptStatus}` : rollout.status;
|
||||
|
||||
return {
|
||||
...rollout,
|
||||
attempt: latestAttempt ?? null,
|
||||
attemptId: latestAttempt?.attemptId ?? null,
|
||||
attemptSequence: latestAttempt?.sequenceId ?? null,
|
||||
isNested: false,
|
||||
canExpand: Boolean(sequenceId && sequenceId > 1),
|
||||
inputText: inputValue,
|
||||
attemptStatus,
|
||||
statusValue,
|
||||
startTimestamp,
|
||||
durationSeconds,
|
||||
lastHeartbeatTimestamp: rollout.attempt?.lastHeartbeatTime ?? null,
|
||||
workerId: latestAttempt?.workerId ?? null,
|
||||
actionsPlaceholder: null,
|
||||
};
|
||||
}
|
||||
|
||||
function buildAttemptRecord(rollout: Rollout, attempt: Attempt): RolloutTableRecord {
|
||||
const inputValue =
|
||||
rollout.input === null || typeof rollout.input === 'undefined'
|
||||
? '—'
|
||||
: typeof rollout.input === 'string'
|
||||
? rollout.input
|
||||
: safeStringify(rollout.input);
|
||||
const startTimestamp = toTimestamp(attempt.startTime ?? rollout.startTime);
|
||||
const endTimestamp = toTimestamp(attempt.endTime);
|
||||
const durationSeconds = clampToNow(startTimestamp, endTimestamp);
|
||||
const lastHeartbeatTimestamp = selectHeartbeatTimestamp(attempt);
|
||||
|
||||
return {
|
||||
...rollout,
|
||||
attempt,
|
||||
attemptId: attempt.attemptId,
|
||||
attemptSequence: attempt.sequenceId,
|
||||
isNested: true,
|
||||
canExpand: false,
|
||||
inputText: inputValue,
|
||||
attemptStatus: attempt.status,
|
||||
statusValue: attempt.status,
|
||||
startTimestamp,
|
||||
durationSeconds,
|
||||
lastHeartbeatTimestamp,
|
||||
workerId: attempt.workerId ?? null,
|
||||
actionsPlaceholder: null,
|
||||
};
|
||||
}
|
||||
|
||||
function getStatusBadge(status: string, kind: 'rollout' | 'attempt') {
|
||||
const color =
|
||||
kind === 'rollout'
|
||||
? (ROLLOUT_STATUS_COLORS[status as RolloutStatus] ?? 'gray')
|
||||
: (ATTEMPT_STATUS_COLORS[status as AttemptStatus] ?? 'gray');
|
||||
|
||||
return (
|
||||
<Badge size='sm' variant='light' color={color}>
|
||||
{formatStatusLabel(status)}
|
||||
</Badge>
|
||||
);
|
||||
}
|
||||
|
||||
type RolloutColumnsOptions = {
|
||||
statusFilters: RolloutStatus[];
|
||||
onStatusFilterChange: (values: RolloutStatus[]) => void;
|
||||
onStatusFilterReset: () => void;
|
||||
modeFilters: RolloutMode[];
|
||||
onModeFilterChange: (values: RolloutMode[]) => void;
|
||||
onModeFilterReset: () => void;
|
||||
onViewRawJson?: (record: RolloutTableRecord) => void;
|
||||
onViewTraces?: (record: RolloutTableRecord) => void;
|
||||
};
|
||||
|
||||
function createRolloutColumns({
|
||||
statusFilters,
|
||||
onStatusFilterChange,
|
||||
onStatusFilterReset,
|
||||
modeFilters,
|
||||
onModeFilterChange,
|
||||
onModeFilterReset,
|
||||
onViewRawJson,
|
||||
onViewTraces,
|
||||
}: RolloutColumnsOptions): DataTableColumn<RolloutTableRecord>[] {
|
||||
const statusOptions = ROLLOUT_STATUS_OPTIONS.map((status) => ({
|
||||
value: status,
|
||||
label: formatStatusLabel(status),
|
||||
}));
|
||||
const modeOptions = ROLLOUT_MODE_OPTIONS.map((mode) => ({
|
||||
value: mode,
|
||||
label: formatStatusLabel(mode),
|
||||
}));
|
||||
|
||||
return [
|
||||
{
|
||||
accessor: 'rolloutId',
|
||||
title: 'Rollout',
|
||||
sortable: true,
|
||||
render: ({ rolloutId }) => (
|
||||
<Group gap={2}>
|
||||
<Text fw={500} size='sm'>
|
||||
{rolloutId}
|
||||
</Text>
|
||||
<CopyButton value={rolloutId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy rollout ID ${rolloutId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'attemptId',
|
||||
title: 'Attempt',
|
||||
sortable: true,
|
||||
render: ({ attemptId, attemptSequence, isNested }) => (
|
||||
<Group gap={2}>
|
||||
<Text size='sm' c={attemptId ? undefined : 'dimmed'}>
|
||||
{attemptId ?? '—'}
|
||||
</Text>
|
||||
{attemptId && (
|
||||
<CopyButton value={attemptId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy attempt ID ${attemptId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
)}
|
||||
{attemptSequence && (isNested || attemptSequence > 1) && (
|
||||
<Badge leftSection={<IconReload size={12} />} pl={6} pr={6}>
|
||||
{attemptSequence}
|
||||
</Badge>
|
||||
)}
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'inputText',
|
||||
title: 'Input',
|
||||
render: ({ inputText }) => (
|
||||
<Text
|
||||
size='sm'
|
||||
ff='monospace'
|
||||
c='dimmed'
|
||||
lineClamp={1}
|
||||
title={inputText}
|
||||
style={{ width: '100%', wordBreak: 'break-all', overflow: 'hidden' }}
|
||||
>
|
||||
{inputText}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'statusValue',
|
||||
title: 'Status',
|
||||
sortable: true,
|
||||
filter: ({ close }) => (
|
||||
<Stack gap='xs'>
|
||||
<MultiSelect
|
||||
label='Status'
|
||||
description='Filter rollouts by status'
|
||||
data={statusOptions}
|
||||
value={statusFilters}
|
||||
placeholder='Select statuses...'
|
||||
searchable
|
||||
clearable
|
||||
comboboxProps={{ withinPortal: false }}
|
||||
onChange={(values) => onStatusFilterChange(values as RolloutStatus[])}
|
||||
/>
|
||||
<Button
|
||||
variant='light'
|
||||
size='xs'
|
||||
onClick={() => {
|
||||
onStatusFilterReset();
|
||||
close();
|
||||
}}
|
||||
disabled={statusFilters.length === 0}
|
||||
>
|
||||
Clear
|
||||
</Button>
|
||||
</Stack>
|
||||
),
|
||||
filtering: statusFilters.length > 0,
|
||||
render: ({ status, attemptStatus, isNested }) => {
|
||||
if (isNested) {
|
||||
return <Group gap={4}>{getStatusBadge(attemptStatus ?? 'unknown', 'attempt')}</Group>;
|
||||
}
|
||||
|
||||
if (attemptStatus && attemptStatus !== status) {
|
||||
return (
|
||||
<Group gap={4}>
|
||||
{getStatusBadge(status, 'rollout')}
|
||||
<Text size='sm' c='dimmed'>
|
||||
—
|
||||
</Text>
|
||||
{getStatusBadge(attemptStatus, 'attempt')}
|
||||
</Group>
|
||||
);
|
||||
}
|
||||
|
||||
return getStatusBadge(status, 'rollout');
|
||||
},
|
||||
},
|
||||
{
|
||||
accessor: 'resourcesId',
|
||||
title: 'Resources',
|
||||
sortable: true,
|
||||
render: ({ resourcesId }) => (
|
||||
<Text size='sm' c={resourcesId ? undefined : 'dimmed'}>
|
||||
{resourcesId ?? '—'}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'mode',
|
||||
title: 'Mode',
|
||||
sortable: true,
|
||||
filter: ({ close }) => (
|
||||
<Stack gap='xs'>
|
||||
<MultiSelect
|
||||
label='Mode'
|
||||
description='Filter rollouts by mode'
|
||||
data={modeOptions}
|
||||
value={modeFilters}
|
||||
placeholder='Select modes...'
|
||||
searchable
|
||||
clearable
|
||||
comboboxProps={{ withinPortal: false }}
|
||||
onChange={(values) => onModeFilterChange(values as RolloutMode[])}
|
||||
/>
|
||||
<Button
|
||||
variant='light'
|
||||
size='xs'
|
||||
onClick={() => {
|
||||
onModeFilterReset();
|
||||
close();
|
||||
}}
|
||||
disabled={modeFilters.length === 0}
|
||||
>
|
||||
Clear
|
||||
</Button>
|
||||
</Stack>
|
||||
),
|
||||
filtering: modeFilters.length > 0,
|
||||
render: ({ mode }) => (
|
||||
<Text size='sm' c={mode ? undefined : 'dimmed'}>
|
||||
{mode ?? '—'}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'startTimestamp',
|
||||
title: 'Start Time',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ startTimestamp }) => <Text size='sm'>{formatDateTime(startTimestamp)}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'durationSeconds',
|
||||
title: 'Duration',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ durationSeconds }) => <Text size='sm'>{formatDuration(durationSeconds)}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'lastHeartbeatTimestamp',
|
||||
title: 'Last Heartbeat',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ lastHeartbeatTimestamp, attempt, isNested }) => {
|
||||
if (!attempt && isNested) {
|
||||
return (
|
||||
<Text size='sm' c='dimmed'>
|
||||
—
|
||||
</Text>
|
||||
);
|
||||
}
|
||||
return <Text size='sm'>{formatRelativeTime(lastHeartbeatTimestamp)}</Text>;
|
||||
},
|
||||
},
|
||||
{
|
||||
accessor: 'workerId',
|
||||
title: 'Worker',
|
||||
sortable: true,
|
||||
render: ({ workerId }) => (
|
||||
<Text size='sm' c={workerId ? undefined : 'dimmed'}>
|
||||
{workerId ?? '—'}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'actionsPlaceholder',
|
||||
title: 'Actions',
|
||||
render: (record) => (
|
||||
<Group gap={4}>
|
||||
<Tooltip label='View raw JSON' withArrow disabled={!onViewRawJson}>
|
||||
<ActionIcon
|
||||
aria-label='View raw JSON'
|
||||
variant='subtle'
|
||||
color='gray'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
onViewRawJson?.(record);
|
||||
}}
|
||||
>
|
||||
<IconFileDescription size={16} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
<Tooltip label='View traces' withArrow disabled={!onViewTraces}>
|
||||
<ActionIcon
|
||||
aria-label='View traces'
|
||||
variant='subtle'
|
||||
color='gray'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
onViewTraces?.(record);
|
||||
}}
|
||||
>
|
||||
<IconTimeline size={16} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
type RowExpansionRenderer = (context: {
|
||||
rollout: Rollout;
|
||||
columns: DataTableColumn<RolloutTableRecord>[];
|
||||
}) => ReactNode;
|
||||
|
||||
export type RolloutTableProps = {
|
||||
rollouts: Rollout[] | undefined;
|
||||
totalRecords: number;
|
||||
isFetching: boolean;
|
||||
isError: boolean;
|
||||
error: unknown;
|
||||
searchTerm: string;
|
||||
statusFilters: RolloutStatus[];
|
||||
modeFilters: RolloutMode[];
|
||||
sort: RolloutsSortState;
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
onStatusFilterChange: (values: RolloutStatus[]) => void;
|
||||
onStatusFilterReset: () => void;
|
||||
onModeFilterChange: (values: RolloutMode[]) => void;
|
||||
onModeFilterReset: () => void;
|
||||
onSortStatusChange: (status: DataTableSortStatus<RolloutTableRecord>) => void;
|
||||
onPageChange: (page: number) => void;
|
||||
onRecordsPerPageChange: (value: number) => void;
|
||||
onResetFilters: () => void;
|
||||
onRefetch: () => void;
|
||||
onViewRawJson?: (record: RolloutTableRecord) => void;
|
||||
onViewTraces?: (record: RolloutTableRecord) => void;
|
||||
recordsPerPageOptions?: number[];
|
||||
renderRowExpansion?: RowExpansionRenderer;
|
||||
};
|
||||
|
||||
export function RolloutTable({
|
||||
rollouts,
|
||||
totalRecords,
|
||||
isFetching,
|
||||
isError,
|
||||
error,
|
||||
searchTerm,
|
||||
statusFilters,
|
||||
modeFilters,
|
||||
sort,
|
||||
page,
|
||||
recordsPerPage,
|
||||
onStatusFilterChange,
|
||||
onStatusFilterReset,
|
||||
onModeFilterChange,
|
||||
onModeFilterReset,
|
||||
onSortStatusChange,
|
||||
onPageChange,
|
||||
onRecordsPerPageChange,
|
||||
onResetFilters,
|
||||
onRefetch,
|
||||
onViewRawJson,
|
||||
onViewTraces,
|
||||
recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
|
||||
renderRowExpansion,
|
||||
}: RolloutTableProps) {
|
||||
const [expandedRecordIds, setExpandedRecordIds] = useState<string[]>([]);
|
||||
const { ref: tableContainerRef, width: containerWidth } = useElementSize();
|
||||
const { width: viewportWidth } = useViewportSize();
|
||||
|
||||
const layoutAwareContainerWidth = useMemo(() => {
|
||||
return getLayoutAwareWidth(containerWidth, viewportWidth);
|
||||
}, [containerWidth, viewportWidth]);
|
||||
|
||||
const rolloutRecords = useMemo<RolloutTableRecord[]>(() => {
|
||||
if (!rollouts) {
|
||||
return [];
|
||||
}
|
||||
return rollouts.map((rolloutItem) => buildRolloutRecord(rolloutItem));
|
||||
}, [rollouts]);
|
||||
|
||||
const columns = useMemo(
|
||||
() =>
|
||||
createRolloutColumns({
|
||||
statusFilters,
|
||||
onStatusFilterChange,
|
||||
onStatusFilterReset,
|
||||
modeFilters,
|
||||
onModeFilterChange,
|
||||
onModeFilterReset,
|
||||
onViewRawJson,
|
||||
onViewTraces,
|
||||
}),
|
||||
[
|
||||
statusFilters,
|
||||
onStatusFilterChange,
|
||||
onStatusFilterReset,
|
||||
modeFilters,
|
||||
onModeFilterChange,
|
||||
onModeFilterReset,
|
||||
onViewRawJson,
|
||||
onViewTraces,
|
||||
],
|
||||
);
|
||||
|
||||
const responsiveColumns = useMemo(
|
||||
() => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
|
||||
[columns, layoutAwareContainerWidth],
|
||||
);
|
||||
|
||||
const totalPages = useMemo(
|
||||
() => Math.max(1, Math.ceil(Math.max(0, totalRecords) / Math.max(1, recordsPerPage))),
|
||||
[recordsPerPage, totalRecords],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (page > totalPages) {
|
||||
onPageChange(totalPages);
|
||||
}
|
||||
}, [onPageChange, page, totalPages]);
|
||||
|
||||
useEffect(() => {
|
||||
setExpandedRecordIds((current) =>
|
||||
current.filter((id) => rolloutRecords.some((record) => record.rolloutId === id && record.canExpand)),
|
||||
);
|
||||
}, [rolloutRecords]);
|
||||
|
||||
const hasActiveFilters = searchTerm.trim().length > 0 || statusFilters.length > 0 || modeFilters.length > 0;
|
||||
|
||||
const sortStatus: DataTableSortStatus<RolloutTableRecord> = {
|
||||
columnAccessor: sort.column,
|
||||
direction: sort.direction,
|
||||
};
|
||||
|
||||
const handleSortStatusChange = useCallback(
|
||||
(status: DataTableSortStatus<RolloutTableRecord>) => {
|
||||
onSortStatusChange(status);
|
||||
},
|
||||
[onSortStatusChange],
|
||||
);
|
||||
|
||||
const errorMessage =
|
||||
isError && error && typeof error === 'object' && 'status' in (error as Record<string, unknown>)
|
||||
? `Rollouts are temporarily unavailable (status: ${String((error as Record<string, unknown>).status)}).`
|
||||
: 'Rollouts are temporarily unavailable.';
|
||||
|
||||
const emptyState = (
|
||||
<Stack gap='sm' align='center' py='lg'>
|
||||
{isError ? (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
{errorMessage}
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
Use the retry button to try again, or adjust the filters to broaden the results.
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' color='gray' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Retry
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
No rollouts found
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
{hasActiveFilters
|
||||
? 'Try adjusting the search or filters to see more results.'
|
||||
: 'Try refreshing to fetch the latest rollouts.'}
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Refresh
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
)}
|
||||
</Stack>
|
||||
);
|
||||
|
||||
return (
|
||||
<Box ref={tableContainerRef} data-testid='rollouts-table-container'>
|
||||
<DataTable<RolloutTableRecord>
|
||||
classNames={{ root: 'rollouts-table' }}
|
||||
withTableBorder
|
||||
withColumnBorders
|
||||
highlightOnHover
|
||||
verticalAlign='center'
|
||||
minHeight={rolloutRecords.length === 0 ? 500 : undefined}
|
||||
idAccessor='rolloutId'
|
||||
records={rolloutRecords}
|
||||
columns={responsiveColumns}
|
||||
totalRecords={totalRecords}
|
||||
recordsPerPage={recordsPerPage}
|
||||
page={page}
|
||||
onPageChange={onPageChange}
|
||||
onRecordsPerPageChange={onRecordsPerPageChange}
|
||||
recordsPerPageOptions={recordsPerPageOptions}
|
||||
sortStatus={sortStatus}
|
||||
onSortStatusChange={handleSortStatusChange}
|
||||
fetching={isFetching}
|
||||
loaderSize='sm'
|
||||
emptyState={rolloutRecords.length === 0 ? emptyState : undefined}
|
||||
rowExpansion={
|
||||
renderRowExpansion
|
||||
? {
|
||||
allowMultiple: true,
|
||||
expandable: ({ record }) => record.canExpand,
|
||||
expanded: {
|
||||
recordIds: expandedRecordIds,
|
||||
onRecordIdsChange: (nextRecordIds: SetStateAction<string[]>) => {
|
||||
setExpandedRecordIds((previous) => {
|
||||
const resolved =
|
||||
typeof nextRecordIds === 'function'
|
||||
? nextRecordIds(previous)
|
||||
: ((nextRecordIds ?? []) as (string | number)[]);
|
||||
return resolved
|
||||
.map(String)
|
||||
.filter((id) =>
|
||||
rolloutRecords.some((tableRecord) => tableRecord.rolloutId === id && tableRecord.canExpand),
|
||||
);
|
||||
});
|
||||
},
|
||||
},
|
||||
content: ({ record }) => renderRowExpansion({ rollout: record, columns: responsiveColumns }),
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
export type RolloutAttemptsTableProps = {
|
||||
rollout: Rollout;
|
||||
attempts: Attempt[] | undefined;
|
||||
isFetching: boolean;
|
||||
isError: boolean;
|
||||
onRetry: () => void;
|
||||
columns: DataTableColumn<RolloutTableRecord>[];
|
||||
};
|
||||
|
||||
export function RolloutAttemptsTable({
|
||||
rollout,
|
||||
attempts,
|
||||
isFetching,
|
||||
isError,
|
||||
onRetry,
|
||||
columns,
|
||||
}: RolloutAttemptsTableProps) {
|
||||
const attemptRecords = useMemo<RolloutTableRecord[]>(() => {
|
||||
if (!attempts) {
|
||||
return [];
|
||||
}
|
||||
return attempts
|
||||
.map((attempt) => buildAttemptRecord(rollout, attempt))
|
||||
.sort((a, b) => (b.attemptSequence ?? 0) - (a.attemptSequence ?? 0))
|
||||
.filter((record) => record.attemptSequence !== rollout.attempt?.sequenceId);
|
||||
}, [attempts, rollout]);
|
||||
|
||||
if (isError && !attemptRecords.length) {
|
||||
return (
|
||||
<Alert color='red' variant='light' icon={<IconAlertCircle size={16} />}>
|
||||
<Stack gap='xs'>
|
||||
<Text size='sm'>Unable to load attempts for this rollout.</Text>
|
||||
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRetry}>
|
||||
Retry
|
||||
</Button>
|
||||
</Stack>
|
||||
</Alert>
|
||||
);
|
||||
}
|
||||
|
||||
const emptyState = (
|
||||
<Stack gap='xs' align='center' py='md'>
|
||||
<Text size='sm' c='dimmed'>
|
||||
No attempts found for this rollout.
|
||||
</Text>
|
||||
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRetry}>
|
||||
Refresh
|
||||
</Button>
|
||||
</Stack>
|
||||
);
|
||||
|
||||
return (
|
||||
<DataTable<RolloutTableRecord>
|
||||
classNames={{ root: 'rollouts-table rollouts-table--nested' }}
|
||||
withColumnBorders
|
||||
noHeader
|
||||
minHeight={0}
|
||||
idAccessor='attemptId'
|
||||
verticalAlign='center'
|
||||
fetching={isFetching}
|
||||
loaderSize='sm'
|
||||
records={attemptRecords}
|
||||
columns={columns}
|
||||
emptyState={attemptRecords.length === 0 ? emptyState : undefined}
|
||||
/>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useMemo, useState } from 'react';
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { IconSearch } from '@tabler/icons-react';
|
||||
import { Box, Stack, TextInput, Title } from '@mantine/core';
|
||||
import type { RolloutsSortState } from '@/features/rollouts';
|
||||
import type { Rollout, RolloutMode, RolloutStatus } from '@/types';
|
||||
import { compareRecords } from '@/utils/table';
|
||||
import { STORY_DATE_NOW_SECONDS } from '../../.storybook/constants';
|
||||
import { buildRolloutRecord, RolloutTable, type RolloutTableRecord } from './RolloutTable.component';
|
||||
|
||||
const meta: Meta<typeof RolloutTable> = {
|
||||
title: 'Components/RolloutTable',
|
||||
component: RolloutTable,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
};
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof RolloutTable>;
|
||||
|
||||
const now = STORY_DATE_NOW_SECONDS;
|
||||
|
||||
const sampleRollouts: Rollout[] = [
|
||||
{
|
||||
rolloutId: 'ro-story-001',
|
||||
input: { task: 'Generate onboarding summary' },
|
||||
startTime: now - 3200,
|
||||
endTime: null,
|
||||
mode: 'train',
|
||||
resourcesId: 'rs-story-001',
|
||||
status: 'running',
|
||||
config: { retries: 1 },
|
||||
metadata: { owner: 'alice' },
|
||||
attempt: {
|
||||
rolloutId: 'ro-story-001',
|
||||
attemptId: 'at-story-010',
|
||||
sequenceId: 1,
|
||||
startTime: now - 3200,
|
||||
endTime: null,
|
||||
status: 'running',
|
||||
workerId: 'worker-east',
|
||||
lastHeartbeatTime: now - 45,
|
||||
metadata: { info: 'Worker is processing' },
|
||||
},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-story-002',
|
||||
input: { task: 'Classify feedback tickets' },
|
||||
startTime: now - 7200,
|
||||
endTime: now - 5400,
|
||||
mode: 'val',
|
||||
resourcesId: 'rs-story-002',
|
||||
status: 'succeeded',
|
||||
config: { retries: 2 },
|
||||
metadata: { owner: 'bob' },
|
||||
attempt: {
|
||||
rolloutId: 'ro-story-002',
|
||||
attemptId: 'at-story-011',
|
||||
sequenceId: 2,
|
||||
startTime: now - 6200,
|
||||
endTime: now - 5400,
|
||||
status: 'succeeded',
|
||||
workerId: 'worker-north',
|
||||
lastHeartbeatTime: now - 5400,
|
||||
metadata: { previousAttempt: 'at-story-010' },
|
||||
},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-story-003',
|
||||
input: { task: 'Analyze experiment results' },
|
||||
startTime: now - 10800,
|
||||
endTime: now - 9600,
|
||||
mode: 'test',
|
||||
resourcesId: 'rs-story-003',
|
||||
status: 'failed',
|
||||
config: { retries: 1 },
|
||||
metadata: { owner: 'carol' },
|
||||
attempt: {
|
||||
rolloutId: 'ro-story-003',
|
||||
attemptId: 'at-story-012',
|
||||
sequenceId: 3,
|
||||
startTime: now - 10200,
|
||||
endTime: now - 9600,
|
||||
status: 'failed',
|
||||
workerId: 'worker-west',
|
||||
lastHeartbeatTime: now - 9600,
|
||||
metadata: { reason: 'Timeout' },
|
||||
},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-story-004',
|
||||
input: { task: 'Evaluate prompt variants' },
|
||||
startTime: now - 3600,
|
||||
endTime: null,
|
||||
mode: 'train',
|
||||
resourcesId: null,
|
||||
status: 'preparing',
|
||||
config: { retries: 0 },
|
||||
metadata: { owner: 'dave' },
|
||||
attempt: null,
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-story-005',
|
||||
input: { task: 'Generate quick answers' },
|
||||
startTime: now - 1800,
|
||||
endTime: null,
|
||||
mode: 'val',
|
||||
resourcesId: 'rs-story-004',
|
||||
status: 'running',
|
||||
config: { retries: 0 },
|
||||
metadata: { owner: 'eva' },
|
||||
attempt: {
|
||||
rolloutId: 'ro-story-005',
|
||||
attemptId: 'at-story-013',
|
||||
sequenceId: 1,
|
||||
startTime: now - 1800,
|
||||
endTime: null,
|
||||
status: 'running',
|
||||
workerId: null,
|
||||
lastHeartbeatTime: now - 75,
|
||||
metadata: null,
|
||||
},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-story-006',
|
||||
input: { task: 'Compile release notes' },
|
||||
startTime: now - 9600,
|
||||
endTime: now - 9000,
|
||||
mode: null,
|
||||
resourcesId: 'rs-story-005',
|
||||
status: 'cancelled',
|
||||
config: { retries: 3 },
|
||||
metadata: null,
|
||||
attempt: {
|
||||
rolloutId: 'ro-story-006',
|
||||
attemptId: 'at-story-014',
|
||||
sequenceId: 1,
|
||||
startTime: now - 9600,
|
||||
endTime: now - 9000,
|
||||
status: 'timeout',
|
||||
workerId: 'worker-south',
|
||||
lastHeartbeatTime: now - 9000,
|
||||
metadata: { info: 'Cancelled by operator' },
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
type WrapperProps = {
|
||||
maxWidth: number;
|
||||
rollouts?: Rollout[] | undefined;
|
||||
isFetching?: boolean;
|
||||
isError?: boolean;
|
||||
error?: unknown;
|
||||
};
|
||||
|
||||
function RolloutTableStoryWrapper({
|
||||
maxWidth,
|
||||
rollouts = sampleRollouts,
|
||||
isFetching = false,
|
||||
isError = false,
|
||||
error = null,
|
||||
}: WrapperProps) {
|
||||
const [searchTerm, setSearchTerm] = useState('');
|
||||
const [statusFilters, setStatusFilters] = useState<RolloutStatus[]>([]);
|
||||
const [modeFilters, setModeFilters] = useState<RolloutMode[]>([]);
|
||||
const [page, setPage] = useState(1);
|
||||
const [recordsPerPage, setRecordsPerPage] = useState(5);
|
||||
const [sort, setSort] = useState<RolloutsSortState>({
|
||||
column: 'startTimestamp',
|
||||
direction: 'desc',
|
||||
});
|
||||
|
||||
const tableRecords = useMemo<RolloutTableRecord[]>(() => {
|
||||
if (!rollouts) {
|
||||
return [];
|
||||
}
|
||||
return rollouts.map((rolloutItem) => buildRolloutRecord(rolloutItem));
|
||||
}, [rollouts]);
|
||||
|
||||
const filteredRecords = useMemo(() => {
|
||||
const normalizedSearch = searchTerm.trim().toLowerCase();
|
||||
return tableRecords.filter((record) => {
|
||||
const matchesSearch = normalizedSearch.length === 0 || record.rolloutId.toLowerCase().includes(normalizedSearch);
|
||||
const matchesStatus = statusFilters.length === 0 || statusFilters.includes(record.status);
|
||||
const matchesMode = modeFilters.length === 0 || (record.mode !== null && modeFilters.includes(record.mode));
|
||||
return matchesSearch && matchesStatus && matchesMode;
|
||||
});
|
||||
}, [modeFilters, searchTerm, statusFilters, tableRecords]);
|
||||
|
||||
const sortedRecords = useMemo(() => {
|
||||
const sorted = filteredRecords.slice();
|
||||
if (!sorted.length) {
|
||||
return sorted;
|
||||
}
|
||||
const comparatorKey = sort.column as keyof RolloutTableRecord;
|
||||
if (!(comparatorKey in sorted[0])) {
|
||||
return sorted;
|
||||
}
|
||||
sorted.sort((a, b) => compareRecords(a, b, comparatorKey));
|
||||
if (sort.direction === 'desc') {
|
||||
sorted.reverse();
|
||||
}
|
||||
return sorted;
|
||||
}, [filteredRecords, sort]);
|
||||
|
||||
const totalRecordsValue = sortedRecords.length;
|
||||
|
||||
const pagedRecords = useMemo(() => {
|
||||
const startIndex = (page - 1) * recordsPerPage;
|
||||
const endIndex = startIndex + recordsPerPage;
|
||||
return sortedRecords.slice(startIndex, endIndex);
|
||||
}, [page, recordsPerPage, sortedRecords]);
|
||||
|
||||
const pagedRollouts = useMemo(() => pagedRecords.map((record) => record as Rollout), [pagedRecords]);
|
||||
|
||||
return (
|
||||
<Box mx='auto' style={{ maxWidth, width: '100%', padding: 16 }}>
|
||||
<Stack gap='md'>
|
||||
<Title order={2}>Rollouts</Title>
|
||||
<TextInput
|
||||
placeholder='Search by Rollout ID'
|
||||
value={searchTerm}
|
||||
onChange={(event) => setSearchTerm(event.currentTarget.value)}
|
||||
leftSection={<IconSearch size={16} />}
|
||||
data-testid='rollouts-search-input'
|
||||
w='100%'
|
||||
style={{ maxWidth: 360 }}
|
||||
/>
|
||||
<RolloutTable
|
||||
rollouts={pagedRollouts}
|
||||
totalRecords={totalRecordsValue}
|
||||
isFetching={isFetching}
|
||||
isError={isError}
|
||||
error={error}
|
||||
searchTerm={searchTerm}
|
||||
statusFilters={statusFilters}
|
||||
modeFilters={modeFilters}
|
||||
sort={sort}
|
||||
page={page}
|
||||
recordsPerPage={recordsPerPage}
|
||||
onStatusFilterChange={(values) => {
|
||||
setStatusFilters(values);
|
||||
setPage(1);
|
||||
}}
|
||||
onStatusFilterReset={() => {
|
||||
setStatusFilters([]);
|
||||
setPage(1);
|
||||
}}
|
||||
onModeFilterChange={(values) => {
|
||||
setModeFilters(values);
|
||||
setPage(1);
|
||||
}}
|
||||
onModeFilterReset={() => {
|
||||
setModeFilters([]);
|
||||
setPage(1);
|
||||
}}
|
||||
onSortStatusChange={(status) => {
|
||||
setSort({
|
||||
column: status.columnAccessor as string,
|
||||
direction: status.direction,
|
||||
});
|
||||
}}
|
||||
onPageChange={setPage}
|
||||
onRecordsPerPageChange={(value) => {
|
||||
setRecordsPerPage(value);
|
||||
setPage(1);
|
||||
}}
|
||||
onResetFilters={() => {
|
||||
setSearchTerm('');
|
||||
setStatusFilters([]);
|
||||
setModeFilters([]);
|
||||
setSort({ column: 'startTimestamp', direction: 'desc' });
|
||||
setPage(1);
|
||||
}}
|
||||
onRefetch={() => undefined}
|
||||
recordsPerPageOptions={[5, 10, 20]}
|
||||
/>
|
||||
</Stack>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
export const WideContainer: Story = {
|
||||
render: () => <RolloutTableStoryWrapper maxWidth={1280} />,
|
||||
};
|
||||
|
||||
export const MediumContainer: Story = {
|
||||
render: () => <RolloutTableStoryWrapper maxWidth={960} />,
|
||||
};
|
||||
|
||||
export const NarrowContainer: Story = {
|
||||
render: () => <RolloutTableStoryWrapper maxWidth={720} />,
|
||||
};
|
||||
|
||||
export const DrawerWidth: Story = {
|
||||
render: () => <RolloutTableStoryWrapper maxWidth={520} />,
|
||||
};
|
||||
|
||||
export const ErrorState: Story = {
|
||||
render: () => (
|
||||
<RolloutTableStoryWrapper maxWidth={600} rollouts={[]} isError error={new Error('Network unreachable')} />
|
||||
),
|
||||
};
|
||||
@@ -0,0 +1,462 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useCallback, useEffect, useMemo } from 'react';
|
||||
import {
|
||||
IconAlertCircle,
|
||||
IconCheck,
|
||||
IconCopy,
|
||||
IconFileDescription,
|
||||
IconRefresh,
|
||||
IconRouteSquare,
|
||||
} from '@tabler/icons-react';
|
||||
import { DataTable, type DataTableColumn, type DataTableSortStatus } from 'mantine-datatable';
|
||||
import { ActionIcon, Badge, Box, Button, CopyButton, Group, Stack, Text, Tooltip } from '@mantine/core';
|
||||
import { useElementSize, useViewportSize } from '@mantine/hooks';
|
||||
import { getLayoutAwareWidth } from '@/layouts/helper';
|
||||
import type { Span } from '@/types';
|
||||
import { getErrorDescriptor } from '@/utils/error';
|
||||
import { formatDateTimeWithMilliseconds, formatDuration, toTimestamp } from '@/utils/format';
|
||||
import { createResponsiveColumns, type ColumnVisibilityConfig } from '@/utils/table';
|
||||
|
||||
const DEFAULT_RECORDS_PER_PAGE_OPTIONS = [50, 100, 200, 500];
|
||||
|
||||
const COLUMN_VISIBILITY: Record<string, ColumnVisibilityConfig> = {
|
||||
name: { minWidth: 12.5, priority: 0 },
|
||||
spanId: { fixedWidth: 14, priority: 1 },
|
||||
traceId: { fixedWidth: 24, priority: 3 },
|
||||
parentId: { fixedWidth: 12, priority: 2 },
|
||||
statusCode: { fixedWidth: 8, priority: 2 },
|
||||
attributeKeys: { minWidth: 12.5, priority: 2 },
|
||||
startTime: { fixedWidth: 15, priority: 1 },
|
||||
endTime: { fixedWidth: 15, priority: 1 },
|
||||
duration: { fixedWidth: 10, priority: 3 },
|
||||
actionsPlaceholder: { fixedWidth: 6, priority: 0 },
|
||||
};
|
||||
|
||||
const STATUS_COLORS: Record<string, string> = {
|
||||
UNSET: 'gray',
|
||||
OK: 'teal',
|
||||
ERROR: 'red',
|
||||
};
|
||||
|
||||
export type TracesTableRecord = Span & {
|
||||
statusCode: string;
|
||||
attributeKeys: string;
|
||||
duration: number;
|
||||
actionsPlaceholder?: null;
|
||||
};
|
||||
|
||||
export function buildTraceRecord(span: Span): TracesTableRecord {
|
||||
const statusCode = span.status.status_code;
|
||||
const attributeKeys = Object.keys(span.attributes ?? {}).join(', ') || '';
|
||||
const startTimestamp = toTimestamp(span.startTime);
|
||||
const endTimestamp = toTimestamp(span.endTime);
|
||||
const duration = endTimestamp && startTimestamp ? endTimestamp - startTimestamp : 0;
|
||||
|
||||
return {
|
||||
...span,
|
||||
statusCode,
|
||||
attributeKeys,
|
||||
duration,
|
||||
actionsPlaceholder: null,
|
||||
};
|
||||
}
|
||||
|
||||
type TracesColumnsOptions = {
|
||||
onShowRollout?: (record: TracesTableRecord) => void;
|
||||
onShowSpanDetail?: (record: TracesTableRecord) => void;
|
||||
onParentIdClick?: (parentId: string) => void;
|
||||
spanIds: Set<string>;
|
||||
};
|
||||
|
||||
function createTracesColumns({
|
||||
onShowRollout,
|
||||
onShowSpanDetail,
|
||||
onParentIdClick,
|
||||
spanIds,
|
||||
}: TracesColumnsOptions): DataTableColumn<TracesTableRecord>[] {
|
||||
return [
|
||||
{
|
||||
accessor: 'name',
|
||||
title: 'Name',
|
||||
sortable: true,
|
||||
render: ({ name }) => (
|
||||
<Text size='sm' fw={500}>
|
||||
{name}
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'traceId',
|
||||
title: 'Trace ID',
|
||||
sortable: true,
|
||||
render: ({ traceId }) => (
|
||||
<Group gap={2}>
|
||||
<Text size='sm'>{traceId}</Text>
|
||||
<CopyButton value={traceId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy trace ID ${traceId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'spanId',
|
||||
title: 'Span ID',
|
||||
sortable: true,
|
||||
render: ({ spanId }) => (
|
||||
<Group gap={2}>
|
||||
<Text size='sm'>{spanId}</Text>
|
||||
<CopyButton value={spanId}>
|
||||
{({ copied, copy }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
|
||||
<ActionIcon
|
||||
aria-label={`Copy span ID ${spanId}`}
|
||||
variant='subtle'
|
||||
color={copied ? 'teal' : 'gray'}
|
||||
size='sm'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
copy();
|
||||
}}
|
||||
>
|
||||
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'parentId',
|
||||
title: 'Parent ID',
|
||||
sortable: true,
|
||||
render: ({ parentId }) => {
|
||||
if (!parentId) {
|
||||
return (
|
||||
<Text size='sm' c='dimmed'>
|
||||
—
|
||||
</Text>
|
||||
);
|
||||
}
|
||||
|
||||
const parentExists = spanIds.has(parentId);
|
||||
const isInteractive = parentExists && typeof onParentIdClick === 'function';
|
||||
|
||||
return (
|
||||
<Group gap={2}>
|
||||
<Text
|
||||
size='sm'
|
||||
c={parentExists ? undefined : 'red'}
|
||||
style={{ cursor: isInteractive ? 'pointer' : undefined }}
|
||||
onClick={(event) => {
|
||||
if (isInteractive) {
|
||||
event.stopPropagation();
|
||||
onParentIdClick?.(parentId);
|
||||
}
|
||||
}}
|
||||
>
|
||||
{parentId.slice(0, 8)}
|
||||
</Text>
|
||||
{!parentExists && (
|
||||
<Tooltip label='Parent span not found in table' withArrow>
|
||||
<IconAlertCircle size={14} color='red' />
|
||||
</Tooltip>
|
||||
)}
|
||||
</Group>
|
||||
);
|
||||
},
|
||||
},
|
||||
|
||||
{
|
||||
accessor: 'statusCode',
|
||||
title: 'Status',
|
||||
sortable: true,
|
||||
render: ({ statusCode }) => (
|
||||
<Badge size='sm' variant='light' color={STATUS_COLORS[statusCode] ?? 'gray'}>
|
||||
{statusCode}
|
||||
</Badge>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'attributeKeys',
|
||||
title: 'Attribute Keys',
|
||||
render: ({ attributeKeys }) =>
|
||||
attributeKeys ? (
|
||||
<Text size='sm' lineClamp={1}>
|
||||
{/* TODO: dim "." and "," and other characters are just normal text */}
|
||||
{attributeKeys}
|
||||
</Text>
|
||||
) : (
|
||||
<Text size='sm' c='dimmed'>
|
||||
—
|
||||
</Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
accessor: 'startTime',
|
||||
title: 'Start Time',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ startTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(startTime))}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'endTime',
|
||||
title: 'End Time',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ endTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(endTime))}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'duration',
|
||||
title: 'Duration',
|
||||
sortable: true,
|
||||
textAlign: 'left',
|
||||
render: ({ duration }) => <Text size='sm'>{formatDuration(duration)}</Text>,
|
||||
},
|
||||
{
|
||||
accessor: 'actionsPlaceholder',
|
||||
title: 'Actions',
|
||||
render: (record) => (
|
||||
<Group gap={2}>
|
||||
<Tooltip label='Show rollout' withArrow disabled={!onShowRollout}>
|
||||
<ActionIcon
|
||||
aria-label='Show rollout'
|
||||
variant='subtle'
|
||||
color='gray'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
onShowRollout?.(record);
|
||||
}}
|
||||
>
|
||||
<IconRouteSquare size={16} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
<Tooltip label='Show span detail' withArrow disabled={!onShowSpanDetail}>
|
||||
<ActionIcon
|
||||
aria-label='Show span detail'
|
||||
variant='subtle'
|
||||
color='gray'
|
||||
onClick={(event) => {
|
||||
event.stopPropagation();
|
||||
onShowSpanDetail?.(record);
|
||||
}}
|
||||
>
|
||||
<IconFileDescription size={16} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
</Group>
|
||||
),
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
export type TracesTableProps = {
|
||||
spans: Span[] | undefined;
|
||||
totalRecords: number;
|
||||
isFetching: boolean;
|
||||
isError: boolean;
|
||||
error: unknown;
|
||||
selectionMessage?: string;
|
||||
searchTerm: string;
|
||||
sort: { column: string; direction: 'asc' | 'desc' };
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
onSortStatusChange: (status: DataTableSortStatus<TracesTableRecord>) => void;
|
||||
onPageChange: (page: number) => void;
|
||||
onRecordsPerPageChange: (value: number) => void;
|
||||
onResetFilters: () => void;
|
||||
onRefetch: () => void;
|
||||
onShowRollout?: (record: TracesTableRecord) => void;
|
||||
onShowSpanDetail?: (record: TracesTableRecord) => void;
|
||||
onParentIdClick?: (parentId: string) => void;
|
||||
recordsPerPageOptions?: number[];
|
||||
};
|
||||
|
||||
export function TracesTable({
|
||||
spans,
|
||||
totalRecords,
|
||||
isFetching,
|
||||
isError,
|
||||
error,
|
||||
selectionMessage,
|
||||
searchTerm,
|
||||
sort,
|
||||
page,
|
||||
recordsPerPage,
|
||||
onSortStatusChange,
|
||||
onPageChange,
|
||||
onRecordsPerPageChange,
|
||||
onResetFilters,
|
||||
onRefetch,
|
||||
onShowRollout,
|
||||
onShowSpanDetail,
|
||||
onParentIdClick,
|
||||
recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
|
||||
}: TracesTableProps) {
|
||||
const { ref: tableContainerRef, width: containerWidth } = useElementSize();
|
||||
const { width: viewportWidth } = useViewportSize();
|
||||
|
||||
const traceRecords = useMemo<TracesTableRecord[]>(() => {
|
||||
if (!spans) {
|
||||
return [];
|
||||
}
|
||||
return spans.map((span) => buildTraceRecord(span));
|
||||
}, [spans]);
|
||||
|
||||
const spanIds = useMemo(() => {
|
||||
return new Set(traceRecords.map((record) => record.spanId));
|
||||
}, [traceRecords]);
|
||||
|
||||
const columns = useMemo(
|
||||
() =>
|
||||
createTracesColumns({
|
||||
onShowRollout,
|
||||
onShowSpanDetail,
|
||||
onParentIdClick,
|
||||
spanIds,
|
||||
}),
|
||||
[onShowRollout, onShowSpanDetail, onParentIdClick, spanIds],
|
||||
);
|
||||
|
||||
const layoutAwareContainerWidth = useMemo(
|
||||
() => getLayoutAwareWidth(containerWidth, viewportWidth),
|
||||
[containerWidth, viewportWidth],
|
||||
);
|
||||
|
||||
const responsiveColumns = useMemo(
|
||||
() => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
|
||||
[columns, layoutAwareContainerWidth],
|
||||
);
|
||||
|
||||
const totalPages = useMemo(
|
||||
() => Math.max(1, Math.ceil(Math.max(0, totalRecords) / Math.max(1, recordsPerPage))),
|
||||
[recordsPerPage, totalRecords],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (page > totalPages) {
|
||||
onPageChange(totalPages);
|
||||
}
|
||||
}, [onPageChange, page, totalPages]);
|
||||
|
||||
const hasActiveFilters = searchTerm.trim().length > 0;
|
||||
|
||||
const sortStatus: DataTableSortStatus<TracesTableRecord> = {
|
||||
columnAccessor: sort.column,
|
||||
direction: sort.direction,
|
||||
};
|
||||
|
||||
const handleSortStatusChange = useCallback(
|
||||
(status: DataTableSortStatus<TracesTableRecord>) => {
|
||||
onSortStatusChange(status);
|
||||
},
|
||||
[onSortStatusChange],
|
||||
);
|
||||
|
||||
const errorDescriptor = isError ? getErrorDescriptor(error) : null;
|
||||
const errorMessage = isError
|
||||
? `Traces are temporarily unavailable${errorDescriptor ? ` (${errorDescriptor})` : ''}.`
|
||||
: 'Traces are temporarily unavailable.';
|
||||
|
||||
const selectionEmptyState = selectionMessage ? (
|
||||
<Stack gap='sm' align='center' py='xl'>
|
||||
<Text fw={600} size='sm'>
|
||||
{selectionMessage}
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
Choose a rollout and attempt from the controls above to load trace results.
|
||||
</Text>
|
||||
</Stack>
|
||||
) : null;
|
||||
|
||||
const fallbackEmptyState = (
|
||||
<Stack gap='sm' align='center' py='lg'>
|
||||
{isError ? (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
{errorMessage}
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
Use the retry button to try again, or adjust the filters to broaden the results.
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' color='gray' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Retry
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Text fw={600} size='sm'>
|
||||
No traces found
|
||||
</Text>
|
||||
<Text size='sm' c='dimmed' ta='center'>
|
||||
{hasActiveFilters
|
||||
? 'Try adjusting the search to see more results.'
|
||||
: 'Try refreshing to fetch the latest traces.'}
|
||||
</Text>
|
||||
<Group gap='xs'>
|
||||
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
|
||||
Refresh
|
||||
</Button>
|
||||
{hasActiveFilters ? (
|
||||
<Button size='xs' variant='subtle' onClick={onResetFilters}>
|
||||
Clear filters
|
||||
</Button>
|
||||
) : null}
|
||||
</Group>
|
||||
</>
|
||||
)}
|
||||
</Stack>
|
||||
);
|
||||
|
||||
const emptyState = selectionEmptyState ?? fallbackEmptyState;
|
||||
|
||||
return (
|
||||
<Box ref={tableContainerRef}>
|
||||
<DataTable<TracesTableRecord>
|
||||
classNames={{ root: 'traces-table' }}
|
||||
withTableBorder
|
||||
withColumnBorders
|
||||
highlightOnHover
|
||||
verticalAlign='center'
|
||||
minHeight={traceRecords.length === 0 ? 500 : undefined}
|
||||
idAccessor='spanId'
|
||||
records={traceRecords}
|
||||
columns={responsiveColumns}
|
||||
totalRecords={totalRecords}
|
||||
recordsPerPage={recordsPerPage}
|
||||
page={page}
|
||||
onPageChange={onPageChange}
|
||||
onRecordsPerPageChange={onRecordsPerPageChange}
|
||||
recordsPerPageOptions={recordsPerPageOptions}
|
||||
sortStatus={sortStatus}
|
||||
onSortStatusChange={handleSortStatusChange}
|
||||
fetching={isFetching}
|
||||
loaderSize='sm'
|
||||
emptyState={traceRecords.length === 0 ? emptyState : undefined}
|
||||
/>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,367 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { useMemo, useState } from 'react';
|
||||
import type { Meta, StoryObj } from '@storybook/react';
|
||||
import { IconSearch } from '@tabler/icons-react';
|
||||
import { Box, Stack, TextInput, Title } from '@mantine/core';
|
||||
import type { Span } from '@/types';
|
||||
import { compareRecords } from '@/utils/table';
|
||||
import { buildTraceRecord, TracesTable, type TracesTableRecord } from './TracesTable.component';
|
||||
|
||||
const meta: Meta<typeof TracesTable> = {
|
||||
title: 'Components/TracesTable',
|
||||
component: TracesTable,
|
||||
parameters: {
|
||||
layout: 'fullscreen',
|
||||
},
|
||||
};
|
||||
|
||||
export default meta;
|
||||
|
||||
type Story = StoryObj<typeof TracesTable>;
|
||||
|
||||
const now = Math.floor(1762775145209 / 1000);
|
||||
|
||||
const sampleSpans: Span[] = [
|
||||
{
|
||||
rolloutId: 'ro-trace-001',
|
||||
attemptId: 'at-trace-001',
|
||||
sequenceId: 1,
|
||||
traceId: 'trace-abc123def456',
|
||||
spanId: 'span-root-001',
|
||||
parentId: null,
|
||||
name: 'main_task',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {
|
||||
'task.type': 'generation',
|
||||
'task.priority': 'high',
|
||||
'user.id': 'user-123',
|
||||
},
|
||||
startTime: now - 100,
|
||||
endTime: now - 10,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-001',
|
||||
attemptId: 'at-trace-001',
|
||||
sequenceId: 1,
|
||||
traceId: 'trace-abc123def456',
|
||||
spanId: 'span-child-001',
|
||||
parentId: 'span-root-001',
|
||||
name: 'llm_call',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {
|
||||
'llm.model': 'gpt-4',
|
||||
'llm.temperature': 0.7,
|
||||
'llm.max_tokens': 2048,
|
||||
},
|
||||
startTime: now - 90,
|
||||
endTime: now - 50,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-001',
|
||||
attemptId: 'at-trace-001',
|
||||
sequenceId: 1,
|
||||
traceId: 'trace-abc123def456',
|
||||
spanId: 'span-child-002',
|
||||
parentId: 'span-root-001',
|
||||
name: 'database_query',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {
|
||||
'db.system': 'postgresql',
|
||||
'db.operation': 'SELECT',
|
||||
'db.table': 'users',
|
||||
},
|
||||
startTime: now - 80,
|
||||
endTime: now - 70,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-002',
|
||||
attemptId: 'at-trace-002',
|
||||
sequenceId: 1,
|
||||
traceId: 'trace-xyz789ghi012',
|
||||
spanId: 'span-error-001',
|
||||
parentId: 'span-missing-parent',
|
||||
name: 'failed_operation',
|
||||
status: { status_code: 'ERROR', description: 'Connection timeout' },
|
||||
attributes: {
|
||||
'error.type': 'TimeoutError',
|
||||
'error.message': 'Connection timed out after 30s',
|
||||
},
|
||||
startTime: now - 150,
|
||||
endTime: now - 120,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-003',
|
||||
attemptId: 'at-trace-003',
|
||||
sequenceId: 1,
|
||||
traceId: 'trace-unset123',
|
||||
spanId: 'span-unset-001',
|
||||
parentId: null,
|
||||
name: 'pending_task',
|
||||
status: { status_code: 'UNSET', description: null },
|
||||
attributes: {},
|
||||
startTime: now - 30,
|
||||
endTime: now - 5,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-004',
|
||||
attemptId: 'at-trace-004',
|
||||
sequenceId: 2,
|
||||
traceId: 'trace-nested456',
|
||||
spanId: 'span-parent-001',
|
||||
parentId: null,
|
||||
name: 'workflow_execution',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {
|
||||
'workflow.name': 'data_processing',
|
||||
'workflow.version': '2.1.0',
|
||||
},
|
||||
startTime: now - 200,
|
||||
endTime: now - 50,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-004',
|
||||
attemptId: 'at-trace-004',
|
||||
sequenceId: 2,
|
||||
traceId: 'trace-nested456',
|
||||
spanId: 'span-child-nested-001',
|
||||
parentId: 'span-parent-001',
|
||||
name: 'step_1_validation',
|
||||
status: { status_code: 'OK', description: null },
|
||||
attributes: {
|
||||
'step.name': 'validation',
|
||||
'step.index': 1,
|
||||
},
|
||||
startTime: now - 195,
|
||||
endTime: now - 180,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
{
|
||||
rolloutId: 'ro-trace-004',
|
||||
attemptId: 'at-trace-004',
|
||||
sequenceId: 2,
|
||||
traceId: 'trace-nested456',
|
||||
spanId: 'span-child-nested-002',
|
||||
parentId: 'span-parent-001',
|
||||
name: 'step_2_processing',
|
||||
status: { status_code: 'ERROR', description: 'Validation failed' },
|
||||
attributes: {
|
||||
'step.name': 'processing',
|
||||
'step.index': 2,
|
||||
'error.type': 'ValidationError',
|
||||
},
|
||||
startTime: now - 175,
|
||||
endTime: now - 160,
|
||||
events: [],
|
||||
links: [],
|
||||
context: {},
|
||||
parent: null,
|
||||
resource: {},
|
||||
},
|
||||
];
|
||||
|
||||
type WrapperProps = {
|
||||
maxWidth: number;
|
||||
spans?: Span[] | undefined;
|
||||
isFetching?: boolean;
|
||||
isError?: boolean;
|
||||
error?: unknown;
|
||||
};
|
||||
|
||||
function TracesTableStoryWrapper({
|
||||
maxWidth,
|
||||
spans = sampleSpans,
|
||||
isFetching = false,
|
||||
isError = false,
|
||||
error = null,
|
||||
}: WrapperProps) {
|
||||
const [searchTerm, setSearchTerm] = useState('');
|
||||
const [page, setPage] = useState(1);
|
||||
const [recordsPerPage, setRecordsPerPage] = useState(10);
|
||||
const [sort, setSort] = useState<{ column: string; direction: 'asc' | 'desc' }>({
|
||||
column: 'startTime',
|
||||
direction: 'desc',
|
||||
});
|
||||
|
||||
const tableRecords = useMemo<TracesTableRecord[]>(() => {
|
||||
if (!spans) {
|
||||
return [];
|
||||
}
|
||||
return spans.map((span) => buildTraceRecord(span));
|
||||
}, [spans]);
|
||||
|
||||
const filteredRecords = useMemo(() => {
|
||||
const normalizedSearch = searchTerm.trim().toLowerCase();
|
||||
if (normalizedSearch.length === 0) {
|
||||
return tableRecords;
|
||||
}
|
||||
return tableRecords.filter(
|
||||
(record) =>
|
||||
record.traceId.toLowerCase().includes(normalizedSearch) ||
|
||||
record.spanId.toLowerCase().includes(normalizedSearch) ||
|
||||
record.name.toLowerCase().includes(normalizedSearch),
|
||||
);
|
||||
}, [searchTerm, tableRecords]);
|
||||
|
||||
const sortedRecords = useMemo(() => {
|
||||
const sorted = filteredRecords.slice();
|
||||
if (!sorted.length) {
|
||||
return sorted;
|
||||
}
|
||||
const comparatorKey = sort.column as keyof TracesTableRecord;
|
||||
if (!(comparatorKey in sorted[0])) {
|
||||
return sorted;
|
||||
}
|
||||
sorted.sort((a, b) => compareRecords(a, b, comparatorKey));
|
||||
if (sort.direction === 'desc') {
|
||||
sorted.reverse();
|
||||
}
|
||||
return sorted;
|
||||
}, [filteredRecords, sort]);
|
||||
|
||||
const totalRecordsValue = sortedRecords.length;
|
||||
|
||||
const pagedRecords = useMemo(() => {
|
||||
const startIndex = (page - 1) * recordsPerPage;
|
||||
const endIndex = startIndex + recordsPerPage;
|
||||
return sortedRecords.slice(startIndex, endIndex);
|
||||
}, [page, recordsPerPage, sortedRecords]);
|
||||
|
||||
const pagedSpans = useMemo(() => pagedRecords.map((record) => record as Span), [pagedRecords]);
|
||||
|
||||
const handleShowRollout = (record: any) => {
|
||||
console.log('Show rollout for:', record.rolloutId);
|
||||
};
|
||||
|
||||
const handleShowSpanDetail = (record: any) => {
|
||||
console.log('Show span detail for:', record.spanId, record);
|
||||
};
|
||||
|
||||
const handleParentIdClick = (parentId: string) => {
|
||||
console.log('Navigate to parent span:', parentId);
|
||||
setSearchTerm(parentId);
|
||||
};
|
||||
|
||||
return (
|
||||
<Box mx='auto' style={{ maxWidth, width: '100%', padding: 16 }}>
|
||||
<Stack gap='md'>
|
||||
<Title order={2}>Traces</Title>
|
||||
<TextInput
|
||||
placeholder='Search by Trace ID, Span ID, or Name'
|
||||
value={searchTerm}
|
||||
onChange={(event) => setSearchTerm(event.currentTarget.value)}
|
||||
leftSection={<IconSearch size={16} />}
|
||||
data-testid='traces-search-input'
|
||||
w='100%'
|
||||
style={{ maxWidth: 360 }}
|
||||
/>
|
||||
<TracesTable
|
||||
spans={pagedSpans}
|
||||
totalRecords={totalRecordsValue}
|
||||
isFetching={isFetching}
|
||||
isError={isError}
|
||||
error={error}
|
||||
searchTerm={searchTerm}
|
||||
sort={sort}
|
||||
page={page}
|
||||
recordsPerPage={recordsPerPage}
|
||||
onSortStatusChange={(status) => {
|
||||
setSort({
|
||||
column: status.columnAccessor as string,
|
||||
direction: status.direction,
|
||||
});
|
||||
}}
|
||||
onPageChange={setPage}
|
||||
onRecordsPerPageChange={(value) => {
|
||||
setRecordsPerPage(value);
|
||||
setPage(1);
|
||||
}}
|
||||
onResetFilters={() => {
|
||||
setSearchTerm('');
|
||||
setSort({ column: 'startTime', direction: 'desc' });
|
||||
setPage(1);
|
||||
}}
|
||||
onRefetch={() => undefined}
|
||||
onShowRollout={handleShowRollout}
|
||||
onShowSpanDetail={handleShowSpanDetail}
|
||||
onParentIdClick={handleParentIdClick}
|
||||
recordsPerPageOptions={[10, 20, 50]}
|
||||
/>
|
||||
</Stack>
|
||||
</Box>
|
||||
);
|
||||
}
|
||||
|
||||
export const WideContainer: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={1400} />,
|
||||
};
|
||||
|
||||
export const MediumContainer: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={960} />,
|
||||
};
|
||||
|
||||
export const NarrowContainer: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={720} />,
|
||||
};
|
||||
|
||||
export const DrawerWidth: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={520} />,
|
||||
};
|
||||
|
||||
export const ErrorState: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} isError error={new Error('Network unreachable')} />,
|
||||
};
|
||||
|
||||
export const LoadingState: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} isFetching />,
|
||||
};
|
||||
|
||||
export const EmptyState: Story = {
|
||||
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} />,
|
||||
};
|
||||
|
||||
export const WithMissingParent: Story = {
|
||||
render: () => (
|
||||
<TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.spanId === 'span-error-001')} />
|
||||
),
|
||||
};
|
||||
|
||||
export const NestedSpans: Story = {
|
||||
render: () => (
|
||||
<TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.traceId === 'trace-nested456')} />
|
||||
),
|
||||
};
|
||||
@@ -0,0 +1,436 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { alpha, CSSVariablesResolver } from '@mantine/core';
|
||||
|
||||
export const shadcnCssVariableResolver: CSSVariablesResolver = () => ({
|
||||
variables: {
|
||||
// variables that do not depend on color scheme
|
||||
'--mantine-heading-font-weight': '600',
|
||||
'--mantine-primary-color-filled-hover': alpha('var(--mantine-primary-color-filled)', 0.9),
|
||||
'--mantine-primary-color-light': 'var(--mantine-color-zinc-light)',
|
||||
'--mantine-primary-color-light-hover': 'var(--mantine-color-zinc-light-hover)',
|
||||
'--mantine-primary-color-light-color': 'var(--mantine-color-zinc-light-color)',
|
||||
},
|
||||
light: {
|
||||
// all variables that depend on light color scheme
|
||||
'--mantine-primary-color-contrast': 'var(--mantine-color-zinc-0)', // used as primary color contrast
|
||||
'--mantine-color-text': 'var(--mantine-color-secondary-9)', // used as text color
|
||||
'--mantine-color-body': 'var(--mantine-color-white)', // used as body color
|
||||
'--mantine-color-error': 'var(--mantine-color-error-10)', // used as error color
|
||||
'--mantine-color-placeholder': 'var(--mantine-color-secondary-10)', // used as placeholder color
|
||||
'--mantine-color-anchor': 'var(--mantine-color-secondary-10)', // used as anchor color
|
||||
|
||||
'--mantine-color-default': 'var(--mantine-color-secondary-0)', // used as default surface color
|
||||
'--mantine-color-default-hover': 'var(--mantine-color-secondary-1)', // used as default hover color
|
||||
'--mantine-color-default-color': 'var(--mantine-color-secondary-9)', // used as default text color
|
||||
'--mantine-color-default-border': 'var(--mantine-color-secondary-2)', // used as default border color
|
||||
'--mantine-color-dimmed': 'var(--mantine-color-secondary-10)', // used as dimmed text color
|
||||
|
||||
'--mantine-color-secondary-filled': 'var(--mantine-color-white)', // used as secondary surface color
|
||||
'--mantine-color-secondary-filled-hover': 'var(--mantine-color-secondary-1)', // used as secondary hover color
|
||||
|
||||
'--mantine-color-secondary-light': 'var(--mantine-color-secondary-1)', // used as primary light color
|
||||
'--mantine-color-secondary-light-hover': alpha('var(--mantine-color-secondary-light)', 0.8), // used as primary light hover color
|
||||
|
||||
'--mantine-color-secondary-text': 'var(--mantine-primary-color-contrast)', // can be used as secondary text color
|
||||
'--mantine-color-secondary-light-color': 'var(--mantine-color-secondary-8)', // used as primary light variant's text color
|
||||
|
||||
'--mantine-color-secondary-outline': 'var(--mantine-color-secondary-2)',
|
||||
'--mantine-color-secondary-outline-hover': 'var(--mantine-color-secondary-1)',
|
||||
|
||||
// all filled colors
|
||||
'--mantine-color-zinc-filled': 'var(--mantine-color-zinc-8)',
|
||||
'--mantine-color-zinc-filled-hover': alpha('var(--mantine-color-zinc-8)', 0.9),
|
||||
'--mantine-color-slate-filled': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-slate-filled-hover': alpha('var(--mantine-color-slate-8)', 0.9),
|
||||
'--mantine-color-gray-filled': 'var(--mantine-color-gray-8)',
|
||||
'--mantine-color-gray-filled-hover': alpha('var(--mantine-color-gray-8)', 0.9),
|
||||
'--mantine-color-neutral-filled': 'var(--mantine-color-neutral-8)',
|
||||
'--mantine-color-neutral-filled-hover': alpha('var(--mantine-color-neutral-8)', 0.9),
|
||||
'--mantine-color-stone-filled': 'var(--mantine-color-stone-8)',
|
||||
'--mantine-color-stone-filled-hover': alpha('var(--mantine-color-stone-8)', 0.9),
|
||||
'--mantine-color-red-filled': 'var(--mantine-color-red-5)',
|
||||
'--mantine-color-red-filled-hover': alpha('var(--mantine-color-red-5)', 0.9),
|
||||
'--mantine-color-rose-filled': 'var(--mantine-color-rose-5)',
|
||||
'--mantine-color-rose-filled-hover': alpha('var(--mantine-color-rose-5)', 0.9),
|
||||
'--mantine-color-orange-filled': 'var(--mantine-color-orange-5)',
|
||||
'--mantine-color-orange-filled-hover': alpha('var(--mantine-color-orange-5)', 0.9),
|
||||
'--mantine-color-amber-filled': 'var(--mantine-color-amber-5)',
|
||||
'--mantine-color-amber-filled-hover': alpha('var(--mantine-color-amber-5)', 0.9),
|
||||
'--mantine-color-yellow-filled': 'var(--mantine-color-yellow-4)',
|
||||
'--mantine-color-yellow-filled-hover': alpha('var(--mantine-color-yellow-4)', 0.9),
|
||||
'--mantine-color-lime-filled': 'var(--mantine-color-lime-5)',
|
||||
'--mantine-color-lime-filled-hover': alpha('var(--mantine-color-lime-5)', 0.9),
|
||||
'--mantine-color-green-filled': 'var(--mantine-color-green-6)',
|
||||
'--mantine-color-green-filled-hover': alpha('var(--mantine-color-green-6)', 0.9),
|
||||
'--mantine-color-emerald-filled': 'var(--mantine-color-emerald-5)',
|
||||
'--mantine-color-emerald-filled-hover': alpha('var(--mantine-color-emerald-5)', 0.9),
|
||||
'--mantine-color-teal-filled': 'var(--mantine-color-teal-5)',
|
||||
'--mantine-color-teal-filled-hover': alpha('var(--mantine-color-teal-5)', 0.9),
|
||||
'--mantine-color-cyan-filled': 'var(--mantine-color-cyan-5)',
|
||||
'--mantine-color-cyan-filled-hover': alpha('var(--mantine-color-cyan-5)', 0.9),
|
||||
'--mantine-color-sky-filled': 'var(--mantine-color-sky-5)',
|
||||
'--mantine-color-sky-filled-hover': alpha('var(--mantine-color-sky-5)', 0.9),
|
||||
'--mantine-color-blue-filled': 'var(--mantine-color-blue-6)',
|
||||
'--mantine-color-blue-filled-hover': alpha('var(--mantine-color-blue-6)', 0.9),
|
||||
'--mantine-color-indigo-filled': 'var(--mantine-color-indigo-5)',
|
||||
'--mantine-color-indigo-filled-hover': alpha('var(--mantine-color-indigo-5)', 0.9),
|
||||
'--mantine-color-violet-filled': 'var(--mantine-color-violet-5)',
|
||||
'--mantine-color-violet-filled-hover': alpha('var(--mantine-color-violet-5)', 0.9),
|
||||
'--mantine-color-purple-filled': 'var(--mantine-color-purple-5)',
|
||||
'--mantine-color-purple-filled-hover': alpha('var(--mantine-color-purple-5)', 0.9),
|
||||
'--mantine-color-fuchsia-filled': 'var(--mantine-color-fuchsia-5)',
|
||||
'--mantine-color-fuchsia-filled-hover': alpha('var(--mantine-color-fuchsia-5)', 0.9),
|
||||
'--mantine-color-pink-filled': 'var(--mantine-color-pink-5)',
|
||||
'--mantine-color-pink-filled-hover': alpha('var(--mantine-color-pink-5)', 0.9),
|
||||
|
||||
// all light colors
|
||||
'--mantine-color-zinc-light': alpha('var(--mantine-color-zinc-4)', 0.1),
|
||||
'--mantine-color-zinc-light-hover': alpha('var(--mantine-color-zinc-light)', 0.8),
|
||||
'--mantine-color-zinc-light-color': 'var(--mantine-color-zinc-6)',
|
||||
'--mantine-color-slate-light': alpha('var(--mantine-color-slate-4)', 0.1),
|
||||
'--mantine-color-slate-light-hover': alpha('var(--mantine-color-slate-light)', 0.8),
|
||||
'--mantine-color-slate-light-color': 'var(--mantine-color-slate-6)',
|
||||
'--mantine-color-gray-light': alpha('var(--mantine-color-gray-4)', 0.1),
|
||||
'--mantine-color-gray-light-hover': alpha('var(--mantine-color-gray-light)', 0.8),
|
||||
'--mantine-color-gray-light-color': 'var(--mantine-color-gray-6)',
|
||||
'--mantine-color-neutral-light': alpha('var(--mantine-color-neutral-4)', 0.1),
|
||||
'--mantine-color-neutral-light-hover': alpha('var(--mantine-color-neutral-light)', 0.8),
|
||||
'--mantine-color-neutral-light-color': 'var(--mantine-color-neutral-6)',
|
||||
'--mantine-color-stone-light': alpha('var(--mantine-color-stone-4)', 0.1),
|
||||
'--mantine-color-stone-light-hover': alpha('var(--mantine-color-stone-light)', 0.8),
|
||||
'--mantine-color-stone-light-color': 'var(--mantine-color-stone-6)',
|
||||
'--mantine-color-red-light': alpha('var(--mantine-color-red-4)', 0.1),
|
||||
'--mantine-color-red-light-hover': alpha('var(--mantine-color-red-light)', 0.8),
|
||||
'--mantine-color-red-light-color': 'var(--mantine-color-red-6)',
|
||||
'--mantine-color-rose-light': alpha('var(--mantine-color-rose-4)', 0.1),
|
||||
'--mantine-color-rose-light-hover': alpha('var(--mantine-color-rose-light)', 0.8),
|
||||
'--mantine-color-rose-light-color': 'var(--mantine-color-rose-6)',
|
||||
'--mantine-color-orange-light': alpha('var(--mantine-color-orange-4)', 0.1),
|
||||
'--mantine-color-orange-light-hover': alpha('var(--mantine-color-orange-light)', 0.8),
|
||||
'--mantine-color-orange-light-color': 'var(--mantine-color-orange-6)',
|
||||
'--mantine-color-amber-light': alpha('var(--mantine-color-amber-4)', 0.1),
|
||||
'--mantine-color-amber-light-hover': alpha('var(--mantine-color-amber-light)', 0.8),
|
||||
'--mantine-color-amber-light-color': 'var(--mantine-color-amber-6)',
|
||||
'--mantine-color-yellow-light': alpha('var(--mantine-color-yellow-4)', 0.1),
|
||||
'--mantine-color-yellow-light-hover': alpha('var(--mantine-color-yellow-light)', 0.8),
|
||||
'--mantine-color-yellow-light-color': 'var(--mantine-color-yellow-6)',
|
||||
'--mantine-color-lime-light': alpha('var(--mantine-color-lime-4)', 0.1),
|
||||
'--mantine-color-lime-light-hover': alpha('var(--mantine-color-lime-light)', 0.8),
|
||||
'--mantine-color-lime-light-color': 'var(--mantine-color-lime-6)',
|
||||
'--mantine-color-green-light': alpha('var(--mantine-color-green-4)', 0.1),
|
||||
'--mantine-color-green-light-hover': alpha('var(--mantine-color-green-light)', 0.8),
|
||||
'--mantine-color-green-light-color': 'var(--mantine-color-green-6)',
|
||||
'--mantine-color-emerald-light': alpha('var(--mantine-color-emerald-4)', 0.1),
|
||||
'--mantine-color-emerald-light-hover': alpha('var(--mantine-color-emerald-light)', 0.8),
|
||||
'--mantine-color-emerald-light-color': 'var(--mantine-color-emerald-6)',
|
||||
'--mantine-color-teal-light': alpha('var(--mantine-color-teal-4)', 0.1),
|
||||
'--mantine-color-teal-light-hover': alpha('var(--mantine-color-teal-light)', 0.8),
|
||||
'--mantine-color-teal-light-color': 'var(--mantine-color-teal-6)',
|
||||
'--mantine-color-cyan-light': alpha('var(--mantine-color-cyan-4)', 0.1),
|
||||
'--mantine-color-cyan-light-hover': alpha('var(--mantine-color-cyan-light)', 0.8),
|
||||
'--mantine-color-cyan-light-color': 'var(--mantine-color-cyan-6)',
|
||||
'--mantine-color-sky-light': alpha('var(--mantine-color-sky-4)', 0.1),
|
||||
'--mantine-color-sky-light-hover': alpha('var(--mantine-color-sky-light)', 0.8),
|
||||
'--mantine-color-sky-light-color': 'var(--mantine-color-sky-6)',
|
||||
'--mantine-color-blue-light': alpha('var(--mantine-color-blue-4)', 0.1),
|
||||
'--mantine-color-blue-light-hover': alpha('var(--mantine-color-blue-light)', 0.8),
|
||||
'--mantine-color-blue-light-color': 'var(--mantine-color-blue-6)',
|
||||
'--mantine-color-indigo-light': alpha('var(--mantine-color-indigo-4)', 0.1),
|
||||
'--mantine-color-indigo-light-hover': alpha('var(--mantine-color-indigo-light)', 0.8),
|
||||
'--mantine-color-indigo-light-color': 'var(--mantine-color-indigo-6)',
|
||||
'--mantine-color-violet-light': alpha('var(--mantine-color-violet-4)', 0.1),
|
||||
'--mantine-color-violet-light-hover': alpha('var(--mantine-color-violet-light)', 0.8),
|
||||
'--mantine-color-violet-light-color': 'var(--mantine-color-violet-6)',
|
||||
'--mantine-color-purple-light': alpha('var(--mantine-color-purple-4)', 0.1),
|
||||
'--mantine-color-purple-light-hover': alpha('var(--mantine-color-purple-light)', 0.8),
|
||||
'--mantine-color-purple-light-color': 'var(--mantine-color-purple-6)',
|
||||
'--mantine-color-fuchsia-light': alpha('var(--mantine-color-fuchsia-4)', 0.1),
|
||||
'--mantine-color-fuchsia-light-hover': alpha('var(--mantine-color-fuchsia-light)', 0.8),
|
||||
'--mantine-color-fuchsia-light-color': 'var(--mantine-color-fuchsia-6)',
|
||||
'--mantine-color-pink-light': alpha('var(--mantine-color-pink-4)', 0.1),
|
||||
'--mantine-color-pink-light-hover': alpha('var(--mantine-color-pink-light)', 0.8),
|
||||
'--mantine-color-pink-light-color': 'var(--mantine-color-pink-6)',
|
||||
|
||||
// all outline colors
|
||||
'--mantine-color-zinc-outline': 'var(--mantine-color-zinc-8)',
|
||||
'--mantine-color-zinc-outline-hover': alpha('var(--mantine-color-zinc-4)', 0.1),
|
||||
'--mantine-color-slate-outline': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-slate-outline-hover': alpha('var(--mantine-color-slate-4)', 0.1),
|
||||
'--mantine-color-gray-outline': 'var(--mantine-color-gray-8)',
|
||||
'--mantine-color-gray-outline-hover': alpha('var(--mantine-color-gray-4)', 0.1),
|
||||
'--mantine-color-neutral-outline': 'var(--mantine-color-neutral-8)',
|
||||
'--mantine-color-neutral-outline-hover': alpha('var(--mantine-color-neutral-4)', 0.1),
|
||||
'--mantine-color-stone-outline': 'var(--mantine-color-stone-8)',
|
||||
'--mantine-color-stone-outline-hover': alpha('var(--mantine-color-stone-4)', 0.1),
|
||||
'--mantine-color-red-outline': 'var(--mantine-color-red-5)',
|
||||
'--mantine-color-red-outline-hover': alpha('var(--mantine-color-red-4)', 0.1),
|
||||
'--mantine-color-rose-outline': 'var(--mantine-color-rose-5)',
|
||||
'--mantine-color-rose-outline-hover': alpha('var(--mantine-color-rose-4)', 0.1),
|
||||
'--mantine-color-orange-outline': 'var(--mantine-color-orange-5)',
|
||||
'--mantine-color-orange-outline-hover': alpha('var(--mantine-color-orange-4)', 0.1),
|
||||
'--mantine-color-amber-outline': 'var(--mantine-color-amber-5)',
|
||||
'--mantine-color-amber-outline-hover': alpha('var(--mantine-color-amber-4)', 0.1),
|
||||
'--mantine-color-yellow-outline': 'var(--mantine-color-yellow-4)',
|
||||
'--mantine-color-yellow-outline-hover': alpha('var(--mantine-color-yellow-4)', 0.1),
|
||||
'--mantine-color-lime-outline': 'var(--mantine-color-lime-5)',
|
||||
'--mantine-color-lime-outline-hover': alpha('var(--mantine-color-lime-4)', 0.1),
|
||||
'--mantine-color-green-outline': 'var(--mantine-color-green-6)',
|
||||
'--mantine-color-green-outline-hover': alpha('var(--mantine-color-green-4)', 0.1),
|
||||
'--mantine-color-emerald-outline': 'var(--mantine-color-emerald-5)',
|
||||
'--mantine-color-emerald-outline-hover': alpha('var(--mantine-color-emerald-4)', 0.1),
|
||||
'--mantine-color-teal-outline': 'var(--mantine-color-teal-5)',
|
||||
'--mantine-color-teal-outline-hover': alpha('var(--mantine-color-teal-4)', 0.1),
|
||||
'--mantine-color-cyan-outline': 'var(--mantine-color-cyan-5)',
|
||||
'--mantine-color-cyan-outline-hover': alpha('var(--mantine-color-cyan-4)', 0.1),
|
||||
'--mantine-color-sky-outline': 'var(--mantine-color-sky-5)',
|
||||
'--mantine-color-sky-outline-hover': alpha('var(--mantine-color-sky-4)', 0.1),
|
||||
'--mantine-color-blue-outline': 'var(--mantine-color-blue-6)',
|
||||
'--mantine-color-blue-outline-hover': alpha('var(--mantine-color-blue-4)', 0.1),
|
||||
'--mantine-color-indigo-outline': 'var(--mantine-color-indigo-5)',
|
||||
'--mantine-color-indigo-outline-hover': alpha('var(--mantine-color-indigo-4)', 0.1),
|
||||
'--mantine-color-violet-outline': 'var(--mantine-color-violet-5)',
|
||||
'--mantine-color-violet-outline-hover': alpha('var(--mantine-color-violet-4)', 0.1),
|
||||
'--mantine-color-purple-outline': 'var(--mantine-color-purple-5)',
|
||||
'--mantine-color-purple-outline-hover': alpha('var(--mantine-color-purple-4)', 0.1),
|
||||
'--mantine-color-fuchsia-outline': 'var(--mantine-color-fuchsia-5)',
|
||||
'--mantine-color-fuchsia-outline-hover': alpha('var(--mantine-color-fuchsia-4)', 0.1),
|
||||
'--mantine-color-pink-outline': 'var(--mantine-color-pink-5)',
|
||||
'--mantine-color-pink-outline-hover': alpha('var(--mantine-color-pink-4)', 0.1),
|
||||
|
||||
// all contrast colors
|
||||
'--mantine-color-zinc-contrast': 'var(--mantine-color-zinc-0)',
|
||||
'--mantine-color-slate-contrast': 'var(--mantine-color-slate-0)',
|
||||
'--mantine-color-gray-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-neutral-contrast': 'var(--mantine-color-neutral-0)',
|
||||
'--mantine-color-stone-contrast': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-red-contrast': 'var(--mantine-color-red-0)',
|
||||
'--mantine-color-rose-contrast': 'var(--mantine-color-rose-0)',
|
||||
'--mantine-color-orange-contrast': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-amber-contrast': 'var(--mantine-color-amber-0)',
|
||||
'--mantine-color-yellow-contrast': '#422006',
|
||||
'--mantine-color-lime-contrast': 'var(--mantine-color-lime-0)',
|
||||
'--mantine-color-green-contrast': 'var(--mantine-color-rose-0)',
|
||||
'--mantine-color-emerald-contrast': 'var(--mantine-color-emerald-0)',
|
||||
'--mantine-color-teal-contrast': 'var(--mantine-color-teal-0)',
|
||||
'--mantine-color-cyan-contrast': 'var(--mantine-color-cyan-0)',
|
||||
'--mantine-color-sky-contrast': 'var(--mantine-color-sky-0)',
|
||||
'--mantine-color-blue-contrast': 'var(--mantine-color-slate-0)',
|
||||
'--mantine-color-indigo-contrast': 'var(--mantine-color-indigo-0)',
|
||||
'--mantine-color-violet-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-purple-contrast': 'var(--mantine-color-purple-0)',
|
||||
'--mantine-color-fuchsia-contrast': 'var(--mantine-color-fuchsia-0)',
|
||||
'--mantine-color-pink-contrast': 'var(--mantine-color-pink-0)',
|
||||
},
|
||||
dark: {
|
||||
// all variables that depend on dark color scheme
|
||||
'--mantine-primary-color-contrast': 'var(--mantine-color-zinc-8)', // used as primary color contrast
|
||||
'--mantine-color-text': 'var(--mantine-color-secondary-0)', // used as text color
|
||||
'--mantine-color-body': 'var(--mantine-color-secondary-9)', // used as body color
|
||||
'--mantine-color-error': 'var(--mantine-color-error-10)', // used as error color
|
||||
'--mantine-color-placeholder': 'var(--mantine-color-secondary-4)', // used as placeholder color
|
||||
'--mantine-color-anchor': 'var(--mantine-color-secondary-4)', // used as anchor color
|
||||
|
||||
'--mantine-color-default': 'var(--mantine-color-secondary-9)', // used as default surface color
|
||||
'--mantine-color-default-hover': 'var(--mantine-color-secondary-7)', // used as default hover color
|
||||
'--mantine-color-default-color': 'var(--mantine-color-secondary-1)', // used as default text color
|
||||
'--mantine-color-default-border': 'var(--mantine-color-secondary-7)', // used as default border color
|
||||
'--mantine-color-dimmed': 'var(--mantine-color-secondary-4)', // used as dimmed text color
|
||||
|
||||
'--mantine-color-secondary-filled': 'var(--mantine-color-secondary-8)', // used as secondary surface color
|
||||
'--mantine-color-secondary-filled-hover': alpha('var(--mantine-color-secondary-filled)', 0.9), // used as secondary hover color
|
||||
|
||||
'--mantine-color-secondary-light': 'var(--mantine-color-secondary-7)', // used as primary light color
|
||||
'--mantine-color-secondary-light-hover': alpha('var(--mantine-color-secondary-light)', 0.8), // used as primary light hover color
|
||||
|
||||
'--mantine-color-secondary-text': 'var(--mantine-primary-color-contrast)', // can be used as secondary text color
|
||||
'--mantine-color-secondary-light-color': 'var(--mantine-color-secondary-0)', // used as primary light text color
|
||||
|
||||
'--mantine-color-secondary-outline': 'var(--mantine-color-secondary-7)',
|
||||
'--mantine-color-secondary-outline-hover': 'var(--mantine-color-secondary-7)',
|
||||
|
||||
// all filled colors
|
||||
'--mantine-color-zinc-filled': 'var(--mantine-color-zinc-0)',
|
||||
'--mantine-color-zinc-filled-hover': alpha('var(--mantine-color-zinc-0)', 0.9),
|
||||
'--mantine-color-slate-filled': 'var(--mantine-color-slate-0)',
|
||||
'--mantine-color-slate-filled-hover': alpha('var(--mantine-color-slate-0)', 0.9),
|
||||
'--mantine-color-gray-filled': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-gray-filled-hover': alpha('var(--mantine-color-gray-0)', 0.9),
|
||||
'--mantine-color-neutral-filled': 'var(--mantine-color-neutral-0)',
|
||||
'--mantine-color-neutral-filled-hover': alpha('var(--mantine-color-neutral-0)', 0.9),
|
||||
'--mantine-color-stone-filled': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-stone-filled-hover': alpha('var(--mantine-color-stone-0)', 0.9),
|
||||
'--mantine-color-red-filled': 'var(--mantine-color-red-5)',
|
||||
'--mantine-color-red-filled-hover': alpha('var(--mantine-color-red-5)', 0.9),
|
||||
'--mantine-color-rose-filled': 'var(--mantine-color-rose-5)',
|
||||
'--mantine-color-rose-filled-hover': alpha('var(--mantine-color-rose-5)', 0.9),
|
||||
'--mantine-color-orange-filled': 'var(--mantine-color-orange-6)',
|
||||
'--mantine-color-orange-filled-hover': alpha('var(--mantine-color-orange-6)', 0.9),
|
||||
'--mantine-color-amber-filled': 'var(--mantine-color-amber-5)',
|
||||
'--mantine-color-amber-filled-hover': alpha('var(--mantine-color-amber-5)', 0.9),
|
||||
'--mantine-color-yellow-filled': 'var(--mantine-color-yellow-4)',
|
||||
'--mantine-color-yellow-filled-hover': alpha('var(--mantine-color-yellow-4)', 0.9),
|
||||
'--mantine-color-lime-filled': 'var(--mantine-color-lime-4)',
|
||||
'--mantine-color-lime-filled-hover': alpha('var(--mantine-color-lime-4)', 0.9),
|
||||
'--mantine-color-green-filled': 'var(--mantine-color-green-5)',
|
||||
'--mantine-color-green-filled-hover': alpha('var(--mantine-color-green-5)', 0.9),
|
||||
'--mantine-color-emerald-filled': 'var(--mantine-color-emerald-5)',
|
||||
'--mantine-color-emerald-filled-hover': alpha('var(--mantine-color-emerald-5)', 0.9),
|
||||
'--mantine-color-teal-filled': 'var(--mantine-color-teal-4)',
|
||||
'--mantine-color-teal-filled-hover': alpha('var(--mantine-color-teal-4)', 0.9),
|
||||
'--mantine-color-cyan-filled': 'var(--mantine-color-cyan-4)',
|
||||
'--mantine-color-cyan-filled-hover': alpha('var(--mantine-color-cyan-4)', 0.9),
|
||||
'--mantine-color-sky-filled': 'var(--mantine-color-sky-4)',
|
||||
'--mantine-color-sky-filled-hover': alpha('var(--mantine-color-sky-4)', 0.9),
|
||||
'--mantine-color-blue-filled': 'var(--mantine-color-blue-5)',
|
||||
'--mantine-color-blue-filled-hover': alpha('var(--mantine-color-blue-5)', 0.9),
|
||||
'--mantine-color-indigo-filled': 'var(--mantine-color-indigo-6)',
|
||||
'--mantine-color-indigo-filled-hover': alpha('var(--mantine-color-indigo-6)', 0.9),
|
||||
'--mantine-color-violet-filled': 'var(--mantine-color-violet-6)',
|
||||
'--mantine-color-violet-filled-hover': alpha('var(--mantine-color-violet-6)', 0.9),
|
||||
'--mantine-color-purple-filled': 'var(--mantine-color-purple-6)',
|
||||
'--mantine-color-purple-filled-hover': alpha('var(--mantine-color-purple-6)', 0.9),
|
||||
'--mantine-color-fuchsia-filled': 'var(--mantine-color-fuchsia-7)',
|
||||
'--mantine-color-fuchsia-filled-hover': alpha('var(--mantine-color-fuchsia-7)', 0.9),
|
||||
'--mantine-color-pink-filled': 'var(--mantine-color-pink-6)',
|
||||
'--mantine-color-pink-filled-hover': alpha('var(--mantine-color-pink-6)', 0.9),
|
||||
|
||||
// all light colors
|
||||
'--mantine-color-zinc-light': alpha('var(--mantine-color-zinc-4)', 0.15),
|
||||
'--mantine-color-zinc-light-hover': alpha('var(--mantine-color-zinc-light)', 0.8),
|
||||
'--mantine-color-zinc-light-color': 'var(--mantine-color-zinc-3)',
|
||||
'--mantine-color-slate-light': alpha('var(--mantine-color-slate-4)', 0.15),
|
||||
'--mantine-color-slate-light-hover': alpha('var(--mantine-color-slate-light)', 0.8),
|
||||
'--mantine-color-slate-light-color': 'var(--mantine-color-slate-3)',
|
||||
'--mantine-color-gray-light': alpha('var(--mantine-color-gray-4)', 0.15),
|
||||
'--mantine-color-gray-light-hover': alpha('var(--mantine-color-gray-light)', 0.8),
|
||||
'--mantine-color-gray-light-color': 'var(--mantine-color-gray-3)',
|
||||
'--mantine-color-neutral-light': alpha('var(--mantine-color-neutral-4)', 0.15),
|
||||
'--mantine-color-neutral-light-hover': alpha('var(--mantine-color-neutral-light)', 0.8),
|
||||
'--mantine-color-neutral-light-color': 'var(--mantine-color-neutral-3)',
|
||||
'--mantine-color-stone-light': alpha('var(--mantine-color-stone-4)', 0.15),
|
||||
'--mantine-color-stone-light-hover': alpha('var(--mantine-color-stone-light)', 0.8),
|
||||
'--mantine-color-stone-light-color': 'var(--mantine-color-stone-3)',
|
||||
'--mantine-color-red-light': alpha('var(--mantine-color-red-4)', 0.15),
|
||||
'--mantine-color-red-light-hover': alpha('var(--mantine-color-red-light)', 0.8),
|
||||
'--mantine-color-red-light-color': 'var(--mantine-color-red-3)',
|
||||
'--mantine-color-rose-light': alpha('var(--mantine-color-rose-4)', 0.15),
|
||||
'--mantine-color-rose-light-hover': alpha('var(--mantine-color-rose-light)', 0.8),
|
||||
'--mantine-color-rose-light-color': 'var(--mantine-color-rose-3)',
|
||||
'--mantine-color-orange-light': alpha('var(--mantine-color-orange-4)', 0.15),
|
||||
'--mantine-color-orange-light-hover': alpha('var(--mantine-color-orange-light)', 0.8),
|
||||
'--mantine-color-orange-light-color': 'var(--mantine-color-orange-3)',
|
||||
'--mantine-color-amber-light': alpha('var(--mantine-color-amber-4)', 0.15),
|
||||
'--mantine-color-amber-light-hover': alpha('var(--mantine-color-amber-light)', 0.8),
|
||||
'--mantine-color-amber-light-color': 'var(--mantine-color-amber-3)',
|
||||
'--mantine-color-yellow-light': alpha('var(--mantine-color-yellow-4)', 0.15),
|
||||
'--mantine-color-yellow-light-hover': alpha('var(--mantine-color-yellow-light)', 0.8),
|
||||
'--mantine-color-yellow-light-color': 'var(--mantine-color-yellow-3)',
|
||||
'--mantine-color-lime-light': alpha('var(--mantine-color-lime-4)', 0.15),
|
||||
'--mantine-color-lime-light-hover': alpha('var(--mantine-color-lime-light)', 0.8),
|
||||
'--mantine-color-lime-light-color': 'var(--mantine-color-lime-3)',
|
||||
'--mantine-color-green-light': alpha('var(--mantine-color-green-4)', 0.15),
|
||||
'--mantine-color-green-light-hover': alpha('var(--mantine-color-green-light)', 0.8),
|
||||
'--mantine-color-green-light-color': 'var(--mantine-color-green-3)',
|
||||
'--mantine-color-emerald-light': alpha('var(--mantine-color-emerald-4)', 0.15),
|
||||
'--mantine-color-emerald-light-hover': alpha('var(--mantine-color-emerald-light)', 0.8),
|
||||
'--mantine-color-emerald-light-color': 'var(--mantine-color-emerald-3)',
|
||||
'--mantine-color-teal-light': alpha('var(--mantine-color-teal-4)', 0.15),
|
||||
'--mantine-color-teal-light-hover': alpha('var(--mantine-color-teal-light)', 0.8),
|
||||
'--mantine-color-teal-light-color': 'var(--mantine-color-teal-3)',
|
||||
'--mantine-color-cyan-light': alpha('var(--mantine-color-cyan-4)', 0.15),
|
||||
'--mantine-color-cyan-light-hover': alpha('var(--mantine-color-cyan-light)', 0.8),
|
||||
'--mantine-color-cyan-light-color': 'var(--mantine-color-cyan-3)',
|
||||
'--mantine-color-sky-light': alpha('var(--mantine-color-sky-4)', 0.15),
|
||||
'--mantine-color-sky-light-hover': alpha('var(--mantine-color-sky-light)', 0.8),
|
||||
'--mantine-color-sky-light-color': 'var(--mantine-color-sky-3)',
|
||||
'--mantine-color-blue-light': alpha('var(--mantine-color-blue-4)', 0.15),
|
||||
'--mantine-color-blue-light-hover': alpha('var(--mantine-color-blue-light)', 0.8),
|
||||
'--mantine-color-blue-light-color': 'var(--mantine-color-blue-3)',
|
||||
'--mantine-color-indigo-light': alpha('var(--mantine-color-indigo-4)', 0.15),
|
||||
'--mantine-color-indigo-light-hover': alpha('var(--mantine-color-indigo-light)', 0.8),
|
||||
'--mantine-color-indigo-light-color': 'var(--mantine-color-indigo-3)',
|
||||
'--mantine-color-violet-light': alpha('var(--mantine-color-violet-4)', 0.15),
|
||||
'--mantine-color-violet-light-hover': alpha('var(--mantine-color-violet-light)', 0.8),
|
||||
'--mantine-color-violet-light-color': 'var(--mantine-color-violet-3)',
|
||||
'--mantine-color-purple-light': alpha('var(--mantine-color-purple-4)', 0.15),
|
||||
'--mantine-color-purple-light-hover': alpha('var(--mantine-color-purple-light)', 0.8),
|
||||
'--mantine-color-purple-light-color': 'var(--mantine-color-purple-3)',
|
||||
'--mantine-color-fuchsia-light': alpha('var(--mantine-color-fuchsia-4)', 0.15),
|
||||
'--mantine-color-fuchsia-light-hover': alpha('var(--mantine-color-fuchsia-light)', 0.8),
|
||||
'--mantine-color-fuchsia-light-color': 'var(--mantine-color-fuchsia-3)',
|
||||
'--mantine-color-pink-light': alpha('var(--mantine-color-pink-4)', 0.15),
|
||||
'--mantine-color-pink-light-hover': alpha('var(--mantine-color-pink-light)', 0.8),
|
||||
'--mantine-color-pink-light-color': 'var(--mantine-color-pink-3)',
|
||||
|
||||
// all outline colors
|
||||
'--mantine-color-zinc-outline': 'var(--mantine-color-zinc-0)',
|
||||
'--mantine-color-zinc-outline-hover': alpha('var(--mantine-color-zinc-4)', 0.15),
|
||||
'--mantine-color-slate-outline': 'var(--mantine-color-slate-0)',
|
||||
'--mantine-color-slate-outline-hover': alpha('var(--mantine-color-slate-4)', 0.15),
|
||||
'--mantine-color-gray-outline': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-gray-outline-hover': alpha('var(--mantine-color-gray-4)', 0.15),
|
||||
'--mantine-color-neutral-outline': 'var(--mantine-color-neutral-0)',
|
||||
'--mantine-color-neutral-outline-hover': alpha('var(--mantine-color-neutral-4)', 0.15),
|
||||
'--mantine-color-stone-outline': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-stone-outline-hover': alpha('var(--mantine-color-stone-4)', 0.15),
|
||||
'--mantine-color-red-outline': 'var(--mantine-color-red-5)',
|
||||
'--mantine-color-red-outline-hover': alpha('var(--mantine-color-red-4)', 0.15),
|
||||
'--mantine-color-rose-outline': 'var(--mantine-color-rose-5)',
|
||||
'--mantine-color-rose-outline-hover': alpha('var(--mantine-color-rose-4)', 0.15),
|
||||
'--mantine-color-orange-outline': 'var(--mantine-color-orange-6)',
|
||||
'--mantine-color-orange-outline-hover': alpha('var(--mantine-color-orange-4)', 0.15),
|
||||
'--mantine-color-amber-outline': 'var(--mantine-color-amber-5)',
|
||||
'--mantine-color-amber-outline-hover': alpha('var(--mantine-color-amber-4)', 0.15),
|
||||
'--mantine-color-yellow-outline': 'var(--mantine-color-yellow-4)',
|
||||
'--mantine-color-yellow-outline-hover': alpha('var(--mantine-color-yellow-4)', 0.15),
|
||||
'--mantine-color-lime-outline': 'var(--mantine-color-lime-4)',
|
||||
'--mantine-color-lime-outline-hover': alpha('var(--mantine-color-lime-4)', 0.15),
|
||||
'--mantine-color-green-outline': 'var(--mantine-color-green-5)',
|
||||
'--mantine-color-green-outline-hover': alpha('var(--mantine-color-green-4)', 0.15),
|
||||
'--mantine-color-emerald-outline': 'var(--mantine-color-emerald-5)',
|
||||
'--mantine-color-emerald-outline-hover': alpha('var(--mantine-color-emerald-4)', 0.15),
|
||||
'--mantine-color-teal-outline': 'var(--mantine-color-teal-4)',
|
||||
'--mantine-color-teal-outline-hover': alpha('var(--mantine-color-teal-4)', 0.15),
|
||||
'--mantine-color-cyan-outline': 'var(--mantine-color-cyan-4)',
|
||||
'--mantine-color-cyan-outline-hover': alpha('var(--mantine-color-cyan-4)', 0.15),
|
||||
'--mantine-color-sky-outline': 'var(--mantine-color-sky-4)',
|
||||
'--mantine-color-sky-outline-hover': alpha('var(--mantine-color-sky-4)', 0.15),
|
||||
'--mantine-color-blue-outline': 'var(--mantine-color-blue-5)',
|
||||
'--mantine-color-blue-outline-hover': alpha('var(--mantine-color-blue-4)', 0.15),
|
||||
'--mantine-color-indigo-outline': 'var(--mantine-color-indigo-6)',
|
||||
'--mantine-color-indigo-outline-hover': alpha('var(--mantine-color-indigo-4)', 0.15),
|
||||
'--mantine-color-violet-outline': 'var(--mantine-color-violet-6)',
|
||||
'--mantine-color-violet-outline-hover': alpha('var(--mantine-color-violet-4)', 0.15),
|
||||
'--mantine-color-purple-outline': 'var(--mantine-color-purple-6)',
|
||||
'--mantine-color-purple-outline-hover': alpha('var(--mantine-color-purple-4)', 0.15),
|
||||
'--mantine-color-fuchsia-outline': 'var(--mantine-color-fuchsia-7)',
|
||||
'--mantine-color-fuchsia-outline-hover': alpha('var(--mantine-color-fuchsia-4)', 0.15),
|
||||
'--mantine-color-pink-outline': 'var(--mantine-color-pink-6)',
|
||||
'--mantine-color-pink-outline-hover': alpha('var(--mantine-color-pink-4)', 0.15),
|
||||
|
||||
// all contrast colors
|
||||
'--mantine-color-zinc-contrast': 'var(--mantine-color-zinc-8)',
|
||||
'--mantine-color-slate-contrast': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-gray-contrast': 'var(--mantine-color-gray-8)',
|
||||
'--mantine-color-neutral-contrast': 'var(--mantine-color-neutral-8)',
|
||||
'--mantine-color-stone-contrast': 'var(--mantine-color-stone-8)',
|
||||
'--mantine-color-red-contrast': 'var(--mantine-color-red-0)',
|
||||
'--mantine-color-rose-contrast': 'var(--mantine-color-rose-0)',
|
||||
'--mantine-color-orange-contrast': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-amber-contrast': 'var(--mantine-color-stone-8)',
|
||||
'--mantine-color-yellow-contrast': '#422006',
|
||||
'--mantine-color-lime-contrast': 'var(--mantine-color-stone-8)',
|
||||
'--mantine-color-green-contrast': 'var(--mantine-color-green-9)',
|
||||
'--mantine-color-emerald-contrast': 'var(--mantine-color-stone-0)',
|
||||
'--mantine-color-teal-contrast': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-cyan-contrast': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-sky-contrast': 'var(--mantine-color-slate-8)',
|
||||
'--mantine-color-blue-contrast': 'var(--mantine-color-slate-0)',
|
||||
'--mantine-color-indigo-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-violet-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-purple-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-fuchsia-contrast': 'var(--mantine-color-gray-0)',
|
||||
'--mantine-color-pink-contrast': 'var(--mantine-color-gray-0)',
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,4 @@
|
||||
<svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">
|
||||
<path d="M8.06935 0.740967C8.46471 0.740967 8.78513 1.06143 8.78513 1.45675C8.78513 1.73028 8.6317 1.96783 8.40619 2.08833V2.6357C8.60265 2.66378 9.18471 2.67955 9.92197 3.01465C10.8483 3.4357 11.121 3.77803 11.4378 4.2357C11.6095 4.48376 11.7205 4.74914 11.7909 5.0357H11.9009C12.273 5.0357 12.5746 5.33732 12.5746 5.70938V6.38307C12.5746 6.75514 12.273 7.05675 11.9009 7.05675H11.8192C11.7591 7.40026 11.6762 7.69328 11.6062 7.85675C11.413 8.30755 11.129 8.48833 10.9746 8.53044C11.0869 8.57254 11.4395 8.67604 11.6483 8.82517C11.943 9.0357 12.2378 9.39174 12.2378 9.75149C12.2378 10.0462 12.1957 10.4673 11.9851 10.6778C11.7074 10.9556 11.2272 11.4778 10.9746 11.6883L6.34303 15.8568L7.35356 13.4989L10.3851 9.54096H8.02724L8.86934 6.59359L8.19658 7.34393L8.19567 7.35149L5.2483 10.762H7.6904L7.10093 12.7831L5.62724 11.8989L4.91146 11.4357C4.53251 11.1831 4.32198 11.0989 4.06935 10.6778C3.91617 10.4225 3.90093 10.0462 3.90093 9.75149C3.90093 9.39174 4.19567 9.0357 4.4904 8.82517C4.69915 8.67604 4.78514 8.61465 4.99567 8.53044C4.82724 8.44623 4.7257 8.30755 4.53251 7.85675C4.46245 7.69328 4.37956 7.40026 4.31951 7.05675H4.23777C3.8657 7.05675 3.56409 6.75514 3.56409 6.38307V5.70938C3.56409 5.33732 3.8657 5.0357 4.23777 5.0357H4.34788C4.41815 4.74914 4.5292 4.48376 4.70093 4.2357C5.01778 3.77803 5.2904 3.4357 6.21672 3.01465C6.95396 2.67955 7.53602 2.66378 7.73251 2.6357V2.08833C7.50704 1.96783 7.35356 1.73028 7.35356 1.45675C7.35356 1.06143 7.67403 0.740967 8.06935 0.740967ZM6.80619 5.0357C6.50389 5.0357 6.25882 5.28077 6.25882 5.58307C6.25882 5.88538 6.50389 6.13044 6.80619 6.13044C7.1085 6.13044 7.35356 5.88538 7.35356 5.58307C7.35356 5.28077 7.1085 5.0357 6.80619 5.0357ZM9.3325 5.0357C9.03018 5.0357 8.78513 5.28077 8.78513 5.58307C8.78513 5.88538 9.03018 6.13044 9.3325 6.13044C9.63481 6.13044 9.87987 5.88538 9.87987 5.58307C9.87987 5.28077 9.63481 5.0357 9.3325 5.0357Z" fill="#F69047"/>
|
||||
<path d="M12.2279 9.63738C12.2342 9.67527 12.2378 9.71342 12.2378 9.75165C12.2378 10.0464 12.1957 10.4674 11.9851 10.678C11.7074 10.9558 11.2272 11.478 10.9746 11.6885L6.34305 15.8569L7.35357 13.499L7.9831 12.677C8.20912 12.6273 8.41543 12.5774 8.57462 12.5306C9.29041 12.3201 10.1325 11.6885 10.7641 11.099C11.2418 10.6531 11.9076 9.97136 12.2279 9.63738ZM9.62725 3.77271C10.3248 3.77271 10.8904 4.33825 10.8904 5.03586V6.80428C10.8904 7.50191 10.3248 8.06744 9.62725 8.06744H8.4483L8.86935 6.59376L8.19659 7.34409L8.19568 7.35165L7.57479 8.06744H6.59568C5.89805 8.06744 5.33252 7.50191 5.33252 6.80428V5.03586C5.33252 4.33825 5.89805 3.77271 6.59568 3.77271H9.62725ZM6.8062 5.03586C6.5039 5.03586 6.25884 5.28093 6.25884 5.58323C6.25884 5.88554 6.5039 6.1306 6.8062 6.1306C7.10851 6.1306 7.35357 5.88554 7.35357 5.58323C7.35357 5.28093 7.10851 5.03586 6.8062 5.03586ZM9.33251 5.03586C9.0302 5.03586 8.78514 5.28093 8.78514 5.58323C8.78514 5.88554 9.0302 6.1306 9.33251 6.1306C9.63483 6.1306 9.87988 5.88554 9.87988 5.58323C9.87988 5.28093 9.63483 5.03586 9.33251 5.03586Z" fill="#C45259"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 3.0 KiB |
@@ -0,0 +1,4 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export * from './slice';
|
||||
export * from './selectors';
|
||||
@@ -0,0 +1,8 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { RootState } from '../../store';
|
||||
|
||||
export const selectConfig = (state: RootState) => state.config;
|
||||
export const selectAutoRefreshMs = (state: RootState) => state.config.autoRefreshMs;
|
||||
export const selectBaseUrl = (state: RootState) => state.config.baseUrl;
|
||||
export const selectThemePreference = (state: RootState) => state.config.theme;
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSlice, PayloadAction } from '@reduxjs/toolkit';
|
||||
import type { ConfigState, ThemePreference } from '@/types';
|
||||
|
||||
export const initialConfigState: ConfigState = {
|
||||
baseUrl: typeof window !== 'undefined' ? window.location.origin : '',
|
||||
autoRefreshMs: 0,
|
||||
theme: 'system',
|
||||
};
|
||||
|
||||
const configSlice = createSlice({
|
||||
name: 'config',
|
||||
initialState: initialConfigState,
|
||||
reducers: {
|
||||
setBaseUrl(state, action: PayloadAction<string>) {
|
||||
state.baseUrl = action.payload;
|
||||
},
|
||||
setAutoRefreshMs(state, action: PayloadAction<number>) {
|
||||
state.autoRefreshMs = action.payload;
|
||||
},
|
||||
setTheme(state, action: PayloadAction<ThemePreference>) {
|
||||
state.theme = action.payload;
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
export const { setAutoRefreshMs, setBaseUrl, setTheme } = configSlice.actions;
|
||||
|
||||
export const configReducer = configSlice.reducer;
|
||||
@@ -0,0 +1,5 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export * from './slice';
|
||||
export * from './selectors';
|
||||
export { useGetResourcesQuery } from '../rollouts';
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createServerBackedStore } from '@test-utils';
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { rolloutsApi } from '@/features/rollouts';
|
||||
import type { Resources } from '@/types';
|
||||
import { selectResourcesQueryArgs } from './selectors';
|
||||
import {
|
||||
resetResourcesFilters,
|
||||
setResourcesPage,
|
||||
setResourcesRecordsPerPage,
|
||||
setResourcesSearchTerm,
|
||||
setResourcesSort,
|
||||
} from './slice';
|
||||
|
||||
const extractResourceIds = (resources: Resources[]): string[] => resources.map((resource) => resource.resourcesId);
|
||||
|
||||
describe('resources feature integration', () => {
|
||||
it('builds default query arguments from the UI state', () => {
|
||||
const store = createServerBackedStore();
|
||||
const queryArgs = selectResourcesQueryArgs(store.getState());
|
||||
|
||||
expect(queryArgs).toMatchObject({
|
||||
limit: 50,
|
||||
offset: 0,
|
||||
sortBy: 'update_time',
|
||||
sortOrder: 'desc',
|
||||
resourcesIdContains: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('fetches resources from the Python LightningStore server', async () => {
|
||||
const store = createServerBackedStore();
|
||||
const queryArgs = selectResourcesQueryArgs(store.getState());
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getResources.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.total).toBe(5);
|
||||
expect(data.items).toHaveLength(5);
|
||||
|
||||
const resourceIds = extractResourceIds(data.items);
|
||||
expect(resourceIds).toEqual(expect.arrayContaining(['rs-story-001', 'rs-story-005']));
|
||||
|
||||
const updateTimes = data.items.map((resource) => resource.updateTime);
|
||||
const sortedUpdateTimes = [...updateTimes].sort((a, b) => b - a);
|
||||
expect(updateTimes).toEqual(sortedUpdateTimes);
|
||||
|
||||
expect(data.items[0].resources).toBeDefined();
|
||||
expect(Object.keys(data.items[0].resources)).not.toHaveLength(0);
|
||||
});
|
||||
|
||||
it('paginates resource results based on UI state', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setResourcesRecordsPerPage(2));
|
||||
store.dispatch(setResourcesPage(2));
|
||||
|
||||
const queryArgs = selectResourcesQueryArgs(store.getState());
|
||||
expect(queryArgs).toMatchObject({ limit: 2, offset: 2 });
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getResources.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.items).toHaveLength(2);
|
||||
expect(data.total).toBe(5);
|
||||
expect(data.items.map((resource) => resource.resourcesId)).toEqual(['rs-story-005', 'rs-story-002']);
|
||||
});
|
||||
|
||||
it('applies search and sorting preferences', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(resetResourcesFilters());
|
||||
store.dispatch(setResourcesSearchTerm('rs-story-003'));
|
||||
store.dispatch(setResourcesSort({ column: 'version', direction: 'asc' }));
|
||||
|
||||
const queryArgs = selectResourcesQueryArgs(store.getState());
|
||||
expect(queryArgs).toMatchObject({
|
||||
limit: 50,
|
||||
offset: 0,
|
||||
sortBy: 'version',
|
||||
sortOrder: 'asc',
|
||||
resourcesIdContains: 'rs-story-003',
|
||||
});
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getResources.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.items).toHaveLength(1);
|
||||
expect(data.items[0].resourcesId).toBe('rs-story-003');
|
||||
expect(data.items[0].version).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,41 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSelector } from '@reduxjs/toolkit';
|
||||
import type { GetResourcesQueryArgs } from '@/features/rollouts';
|
||||
import type { RootState } from '@/store';
|
||||
import type { ResourcesSortState } from './slice';
|
||||
|
||||
const RESOURCES_SORT_FIELD_MAP: Record<string, string> = {
|
||||
resourcesId: 'resources_id',
|
||||
version: 'version',
|
||||
createTime: 'create_time',
|
||||
updateTime: 'update_time',
|
||||
};
|
||||
|
||||
const resolveResourcesSortField = (sort: ResourcesSortState): string =>
|
||||
RESOURCES_SORT_FIELD_MAP[sort.column] ?? 'update_time';
|
||||
|
||||
export const selectResourcesUiState = (state: RootState) => state.resources;
|
||||
|
||||
export const selectResourcesSearchTerm = (state: RootState) => selectResourcesUiState(state).searchTerm;
|
||||
export const selectResourcesPage = (state: RootState) => selectResourcesUiState(state).page;
|
||||
export const selectResourcesRecordsPerPage = (state: RootState) => selectResourcesUiState(state).recordsPerPage;
|
||||
export const selectResourcesSort = (state: RootState) => selectResourcesUiState(state).sort;
|
||||
|
||||
export const selectResourcesQueryArgs = createSelector(
|
||||
[selectResourcesSearchTerm, selectResourcesPage, selectResourcesRecordsPerPage, selectResourcesSort],
|
||||
(searchTerm, page, recordsPerPage, sort): GetResourcesQueryArgs => {
|
||||
const normalizedSearch = searchTerm.trim();
|
||||
const limit = Math.max(1, recordsPerPage);
|
||||
const offset = Math.max(0, (page - 1) * limit);
|
||||
const sortBy = resolveResourcesSortField(sort);
|
||||
|
||||
return {
|
||||
limit,
|
||||
offset,
|
||||
sortBy,
|
||||
sortOrder: sort.direction,
|
||||
resourcesIdContains: normalizedSearch.length > 0 ? normalizedSearch : undefined,
|
||||
};
|
||||
},
|
||||
);
|
||||
@@ -0,0 +1,64 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
|
||||
|
||||
export type SortDirection = 'asc' | 'desc';
|
||||
|
||||
export type ResourcesSortState = {
|
||||
column: string;
|
||||
direction: SortDirection;
|
||||
};
|
||||
|
||||
export type ResourcesUiState = {
|
||||
searchTerm: string;
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
sort: ResourcesSortState;
|
||||
};
|
||||
|
||||
export const initialResourcesUiState: ResourcesUiState = {
|
||||
searchTerm: '',
|
||||
page: 1,
|
||||
recordsPerPage: 50,
|
||||
sort: {
|
||||
column: 'updateTime',
|
||||
direction: 'desc',
|
||||
},
|
||||
};
|
||||
|
||||
const resourcesSlice = createSlice({
|
||||
name: 'resources',
|
||||
initialState: initialResourcesUiState,
|
||||
reducers: {
|
||||
setResourcesSearchTerm(state, action: PayloadAction<string>) {
|
||||
state.searchTerm = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setResourcesPage(state, action: PayloadAction<number>) {
|
||||
state.page = action.payload;
|
||||
},
|
||||
setResourcesRecordsPerPage(state, action: PayloadAction<number>) {
|
||||
state.recordsPerPage = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setResourcesSort(state, action: PayloadAction<ResourcesSortState>) {
|
||||
state.sort = action.payload;
|
||||
},
|
||||
resetResourcesFilters(state) {
|
||||
state.searchTerm = initialResourcesUiState.searchTerm;
|
||||
state.page = initialResourcesUiState.page;
|
||||
state.recordsPerPage = initialResourcesUiState.recordsPerPage;
|
||||
state.sort = initialResourcesUiState.sort;
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
export const {
|
||||
setResourcesSearchTerm,
|
||||
setResourcesPage,
|
||||
setResourcesRecordsPerPage,
|
||||
setResourcesSort,
|
||||
resetResourcesFilters,
|
||||
} = resourcesSlice.actions;
|
||||
|
||||
export const resourcesReducer = resourcesSlice.reducer;
|
||||
@@ -0,0 +1,346 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import type { BaseQueryFn } from '@reduxjs/toolkit/query';
|
||||
import { createApi, fetchBaseQuery, type FetchArgs, type FetchBaseQueryError } from '@reduxjs/toolkit/query/react';
|
||||
import type { RootState } from '@/store';
|
||||
import { camelCaseKeys } from '@/utils/format';
|
||||
import type {
|
||||
Attempt,
|
||||
PaginatedResponse,
|
||||
Resources,
|
||||
Rollout,
|
||||
RolloutMode,
|
||||
RolloutStatus,
|
||||
Span,
|
||||
Timestamp,
|
||||
} from '../../types';
|
||||
|
||||
const rawBaseQuery = fetchBaseQuery({
|
||||
baseUrl: '/',
|
||||
});
|
||||
|
||||
const buildAbsoluteUrl = (baseUrl: string, path: string) => {
|
||||
if (path.startsWith('http://') || path.startsWith('https://')) {
|
||||
return path;
|
||||
}
|
||||
|
||||
const normalizedBase = baseUrl.replace(/\/+$/, '');
|
||||
const normalizedPath = path.replace(/^\/+/, '');
|
||||
if (!normalizedBase) {
|
||||
return `/${normalizedPath}`;
|
||||
}
|
||||
return `${normalizedBase}/${normalizedPath}`;
|
||||
};
|
||||
|
||||
const normalizeHeartbeat = (
|
||||
attempt: Partial<Attempt> & { lastHeartbeatTime?: Timestamp | null; lastHeartBeatTime?: Timestamp | null },
|
||||
): Timestamp | null => {
|
||||
if (typeof attempt.lastHeartbeatTime === 'number') {
|
||||
return attempt.lastHeartbeatTime;
|
||||
}
|
||||
|
||||
if (typeof attempt.lastHeartBeatTime === 'number') {
|
||||
return attempt.lastHeartBeatTime;
|
||||
}
|
||||
|
||||
if (typeof attempt.startTime === 'number') {
|
||||
return attempt.startTime;
|
||||
}
|
||||
|
||||
return null;
|
||||
};
|
||||
|
||||
const normalizeAttempt = (value: unknown): Attempt | null => {
|
||||
if (value === null || typeof value === 'undefined') {
|
||||
return null;
|
||||
}
|
||||
|
||||
const camelized = camelCaseKeys(value) as Attempt & {
|
||||
lastHeartbeatTime?: Timestamp | null;
|
||||
lastHeartBeatTime?: Timestamp | null;
|
||||
};
|
||||
const { lastHeartbeatTime, lastHeartBeatTime, ...rest } = camelized;
|
||||
|
||||
return {
|
||||
...rest,
|
||||
lastHeartbeatTime: normalizeHeartbeat({ ...rest, lastHeartbeatTime, lastHeartBeatTime }),
|
||||
};
|
||||
};
|
||||
|
||||
const normalizeAttemptStrict = (value: unknown): Attempt => {
|
||||
const normalized = normalizeAttempt(value);
|
||||
if (!normalized) {
|
||||
throw new Error('Expected attempt payload');
|
||||
}
|
||||
return normalized;
|
||||
};
|
||||
|
||||
const normalizeRollout = (value: unknown): Rollout => {
|
||||
const camelized = camelCaseKeys(value) as Rollout & { attempt?: unknown };
|
||||
const { attempt, ...rest } = camelized;
|
||||
|
||||
return {
|
||||
...rest,
|
||||
attempt: normalizeAttempt(attempt),
|
||||
};
|
||||
};
|
||||
|
||||
const normalizeSpan = (value: unknown): Span => {
|
||||
const camelized = camelCaseKeys(value) as Span & {
|
||||
status?: {
|
||||
status_code?: Span['status']['status_code'];
|
||||
statusCode?: Span['status']['status_code'];
|
||||
description?: string | null;
|
||||
};
|
||||
};
|
||||
const rawStatus = camelized.status ?? { status_code: 'UNSET', description: null };
|
||||
const result = {
|
||||
...camelized,
|
||||
parentId: camelized.parentId ?? null,
|
||||
// The following fields does not need to be normalized to camel case
|
||||
// For example, gen_ai.xxx should not become genAi.xxx
|
||||
attributes: (value as any).attributes ?? {},
|
||||
context: (value as any).context ?? {},
|
||||
parent: (value as any).parent ?? null,
|
||||
resource: (value as any).resource ?? {},
|
||||
status: {
|
||||
status_code: rawStatus.status_code ?? rawStatus.statusCode ?? 'UNSET',
|
||||
description: rawStatus.description ?? null,
|
||||
},
|
||||
};
|
||||
return result;
|
||||
};
|
||||
|
||||
const normalizeResources = (value: unknown): Resources => {
|
||||
const camelized = camelCaseKeys(value) as Resources;
|
||||
return {
|
||||
resourcesId: camelized.resourcesId,
|
||||
version: camelized.version,
|
||||
createTime: camelized.createTime,
|
||||
updateTime: camelized.updateTime,
|
||||
resources: camelized.resources ?? {},
|
||||
};
|
||||
};
|
||||
|
||||
const normalizePaginatedResponse = <T>(value: unknown, normalizer: (item: unknown) => T): PaginatedResponse<T> => {
|
||||
if (!value || typeof value !== 'object') {
|
||||
throw new Error('Expected paginated response payload');
|
||||
}
|
||||
|
||||
const converted = value as {
|
||||
items?: unknown;
|
||||
limit?: number;
|
||||
offset?: number;
|
||||
total?: number;
|
||||
};
|
||||
|
||||
const itemsSource = Array.isArray(converted.items) ? converted.items : [];
|
||||
|
||||
return {
|
||||
items: itemsSource.map((item) => normalizer(item)),
|
||||
limit: typeof converted.limit === 'number' ? converted.limit : itemsSource.length,
|
||||
offset: typeof converted.offset === 'number' ? converted.offset : 0,
|
||||
total: typeof converted.total === 'number' ? converted.total : itemsSource.length,
|
||||
};
|
||||
};
|
||||
|
||||
const dynamicBaseQuery: BaseQueryFn<string | FetchArgs, unknown, FetchBaseQueryError> = async (
|
||||
args,
|
||||
api,
|
||||
extraOptions,
|
||||
) => {
|
||||
const state = api.getState() as RootState;
|
||||
const stateBaseUrl = state.config?.baseUrl;
|
||||
const fallbackBaseUrl = typeof window !== 'undefined' ? window.location.origin : '';
|
||||
const baseUrl = stateBaseUrl && stateBaseUrl.trim().length > 0 ? stateBaseUrl : fallbackBaseUrl;
|
||||
const preparedArgs: FetchArgs =
|
||||
typeof args === 'string'
|
||||
? { url: args }
|
||||
: {
|
||||
...args,
|
||||
url: args.url ?? '',
|
||||
};
|
||||
|
||||
const absoluteUrl = buildAbsoluteUrl(baseUrl, preparedArgs.url ?? '');
|
||||
return rawBaseQuery({ ...preparedArgs, url: absoluteUrl }, api, extraOptions);
|
||||
};
|
||||
|
||||
export type GetRolloutsQueryArgs = {
|
||||
limit: number;
|
||||
offset: number;
|
||||
sortBy?: string | null;
|
||||
sortOrder?: 'asc' | 'desc';
|
||||
statusIn?: RolloutStatus[];
|
||||
rolloutIdContains?: string | null;
|
||||
modeIn?: RolloutMode[];
|
||||
};
|
||||
|
||||
export type GetResourcesQueryArgs = {
|
||||
limit: number;
|
||||
offset: number;
|
||||
sortBy?: string | null;
|
||||
sortOrder?: 'asc' | 'desc';
|
||||
resourcesIdContains?: string | null;
|
||||
};
|
||||
|
||||
export type GetRolloutAttemptsQueryArgs = {
|
||||
rolloutId: string;
|
||||
limit?: number;
|
||||
offset?: number;
|
||||
sortBy?: string | null;
|
||||
sortOrder?: 'asc' | 'desc';
|
||||
};
|
||||
|
||||
export type GetSpansQueryArgs = {
|
||||
rolloutId: string;
|
||||
attemptId?: string | null;
|
||||
limit?: number;
|
||||
offset?: number;
|
||||
sortBy?: string | null;
|
||||
sortOrder?: 'asc' | 'desc';
|
||||
traceIdContains?: string | null;
|
||||
spanIdContains?: string | null;
|
||||
parentIdContains?: string | null;
|
||||
nameContains?: string | null;
|
||||
filterLogic?: 'and' | 'or' | null;
|
||||
};
|
||||
|
||||
export const rolloutsApi = createApi({
|
||||
reducerPath: 'rolloutsApi',
|
||||
baseQuery: dynamicBaseQuery,
|
||||
tagTypes: ['Rollout', 'Span', 'Resources'],
|
||||
endpoints: (builder) => ({
|
||||
getResources: builder.query<PaginatedResponse<Resources>, GetResourcesQueryArgs>({
|
||||
query: ({ limit, offset, sortBy, sortOrder, resourcesIdContains }) => {
|
||||
const searchParams = new URLSearchParams();
|
||||
searchParams.set('limit', String(typeof limit === 'number' ? limit : -1));
|
||||
searchParams.set('offset', String(typeof offset === 'number' ? offset : 0));
|
||||
if (sortBy) {
|
||||
searchParams.set('sort_by', sortBy);
|
||||
}
|
||||
if (sortOrder) {
|
||||
searchParams.set('sort_order', sortOrder);
|
||||
}
|
||||
if (resourcesIdContains && resourcesIdContains.trim().length > 0) {
|
||||
searchParams.set('resources_id_contains', resourcesIdContains.trim());
|
||||
}
|
||||
|
||||
const queryString = searchParams.toString();
|
||||
const url = queryString.length > 0 ? `v1/agl/resources?${queryString}` : 'v1/agl/resources';
|
||||
return { url, method: 'GET' };
|
||||
},
|
||||
transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeResources),
|
||||
providesTags: (result) =>
|
||||
result
|
||||
? [
|
||||
{ type: 'Resources' as const, id: 'LIST' },
|
||||
...result.items.map((item) => ({ type: 'Resources' as const, id: item.resourcesId })),
|
||||
]
|
||||
: [{ type: 'Resources' as const, id: 'LIST' }],
|
||||
}),
|
||||
getRollouts: builder.query<PaginatedResponse<Rollout>, GetRolloutsQueryArgs>({
|
||||
query: ({ limit, offset, sortBy, sortOrder, statusIn, rolloutIdContains, modeIn }) => {
|
||||
const searchParams = new URLSearchParams();
|
||||
searchParams.set('limit', String(typeof limit === 'number' ? limit : -1));
|
||||
searchParams.set('offset', String(typeof offset === 'number' ? offset : 0));
|
||||
if (sortBy) {
|
||||
searchParams.set('sort_by', sortBy);
|
||||
}
|
||||
if (sortOrder) {
|
||||
searchParams.set('sort_order', sortOrder);
|
||||
}
|
||||
if (statusIn && statusIn.length > 0) {
|
||||
statusIn.forEach((status) => searchParams.append('status_in', status));
|
||||
}
|
||||
if (modeIn && modeIn.length > 0) {
|
||||
modeIn.forEach((mode) => searchParams.append('mode_in', mode));
|
||||
}
|
||||
if (rolloutIdContains && rolloutIdContains.trim().length > 0) {
|
||||
searchParams.set('rollout_id_contains', rolloutIdContains.trim());
|
||||
}
|
||||
|
||||
const queryString = searchParams.toString();
|
||||
const url = queryString.length > 0 ? `v1/agl/rollouts?${queryString}` : 'v1/agl/rollouts';
|
||||
return { url, method: 'GET' };
|
||||
},
|
||||
transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeRollout),
|
||||
providesTags: (result) =>
|
||||
result
|
||||
? [
|
||||
{ type: 'Rollout' as const, id: 'LIST' },
|
||||
...result.items.map((rollout) => ({ type: 'Rollout' as const, id: rollout.rolloutId })),
|
||||
]
|
||||
: [{ type: 'Rollout' as const, id: 'LIST' }],
|
||||
}),
|
||||
getRolloutAttempts: builder.query<PaginatedResponse<Attempt>, GetRolloutAttemptsQueryArgs>({
|
||||
query: ({ rolloutId, limit = -1, offset = 0, sortBy, sortOrder }) => {
|
||||
const searchParams = new URLSearchParams();
|
||||
searchParams.set('limit', String(typeof limit === 'number' ? limit : -1));
|
||||
searchParams.set('offset', String(typeof offset === 'number' ? offset : 0));
|
||||
if (sortBy) {
|
||||
searchParams.set('sort_by', sortBy);
|
||||
}
|
||||
if (sortOrder) {
|
||||
searchParams.set('sort_order', sortOrder);
|
||||
}
|
||||
const queryString = searchParams.toString();
|
||||
const url =
|
||||
queryString.length > 0
|
||||
? `v1/agl/rollouts/${rolloutId}/attempts?${queryString}`
|
||||
: `v1/agl/rollouts/${rolloutId}/attempts`;
|
||||
return { url, method: 'GET' };
|
||||
},
|
||||
transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeAttemptStrict),
|
||||
providesTags: (_result, _error, queryArgs) => [{ type: 'Rollout', id: queryArgs.rolloutId }],
|
||||
}),
|
||||
getSpans: builder.query<PaginatedResponse<Span>, GetSpansQueryArgs>({
|
||||
query: (args) => {
|
||||
if (!args.rolloutId) {
|
||||
throw new Error('rolloutId is required to fetch spans');
|
||||
}
|
||||
const searchParams = new URLSearchParams({ rollout_id: args.rolloutId });
|
||||
if (args.attemptId) {
|
||||
searchParams.set('attempt_id', args.attemptId);
|
||||
}
|
||||
if (typeof args.limit === 'number') {
|
||||
searchParams.set('limit', String(args.limit));
|
||||
}
|
||||
if (typeof args.offset === 'number') {
|
||||
searchParams.set('offset', String(args.offset));
|
||||
}
|
||||
if (args.sortBy) {
|
||||
searchParams.set('sort_by', args.sortBy);
|
||||
}
|
||||
if (args.sortOrder) {
|
||||
searchParams.set('sort_order', args.sortOrder);
|
||||
}
|
||||
if (args.traceIdContains) {
|
||||
searchParams.set('trace_id_contains', args.traceIdContains);
|
||||
}
|
||||
if (args.spanIdContains) {
|
||||
searchParams.set('span_id_contains', args.spanIdContains);
|
||||
}
|
||||
if (args.parentIdContains) {
|
||||
searchParams.set('parent_id_contains', args.parentIdContains);
|
||||
}
|
||||
if (args.nameContains) {
|
||||
searchParams.set('name_contains', args.nameContains);
|
||||
}
|
||||
if (args.filterLogic) {
|
||||
searchParams.set('filter_logic', args.filterLogic);
|
||||
}
|
||||
return { url: `v1/agl/spans?${searchParams.toString()}`, method: 'GET' };
|
||||
},
|
||||
transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeSpan),
|
||||
providesTags: (_result, _error, args) =>
|
||||
args
|
||||
? [
|
||||
{ type: 'Span' as const, id: `${args.rolloutId}:${args.attemptId ?? 'latest'}` },
|
||||
{ type: 'Span' as const, id: 'LIST' },
|
||||
]
|
||||
: [{ type: 'Span' as const, id: 'LIST' }],
|
||||
}),
|
||||
}),
|
||||
});
|
||||
|
||||
export const { useGetResourcesQuery, useGetRolloutsQuery, useGetRolloutAttemptsQuery, useGetSpansQuery } = rolloutsApi;
|
||||
@@ -0,0 +1,6 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export * from './api';
|
||||
export * from './slice';
|
||||
export * from './selectors';
|
||||
export * from '../../types';
|
||||
@@ -0,0 +1,127 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createServerBackedStore } from '@test-utils';
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { rolloutsApi } from './api';
|
||||
import { selectRolloutsQueryArgs } from './selectors';
|
||||
import {
|
||||
resetRolloutsFilters,
|
||||
setRolloutsModeFilters,
|
||||
setRolloutsPage,
|
||||
setRolloutsRecordsPerPage,
|
||||
setRolloutsSearchTerm,
|
||||
setRolloutsSort,
|
||||
setRolloutsStatusFilters,
|
||||
} from './slice';
|
||||
|
||||
describe('rollouts feature integration', () => {
|
||||
it('builds default query arguments from the UI state', () => {
|
||||
const store = createServerBackedStore();
|
||||
const queryArgs = selectRolloutsQueryArgs(store.getState());
|
||||
|
||||
expect(queryArgs).toMatchObject({
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
sortBy: 'start_time',
|
||||
sortOrder: 'desc',
|
||||
rolloutIdContains: undefined,
|
||||
statusIn: undefined,
|
||||
modeIn: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('retrieves rollouts from the Python LightningStore server', async () => {
|
||||
const store = createServerBackedStore();
|
||||
const queryArgs = selectRolloutsQueryArgs(store.getState());
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getRollouts.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.total).toBe(6);
|
||||
expect(data.items).toHaveLength(6);
|
||||
|
||||
const rolloutIds = data.items.map((rollout) => rollout.rolloutId);
|
||||
expect(rolloutIds).toEqual(
|
||||
expect.arrayContaining(['ro-story-001', 'ro-story-002', 'ro-story-003', 'ro-story-004', 'ro-story-005']),
|
||||
);
|
||||
|
||||
const startTimes = data.items.map((rollout) => rollout.startTime);
|
||||
const sortedStartTimes = [...startTimes].sort((a, b) => b - a);
|
||||
expect(startTimes).toEqual(sortedStartTimes);
|
||||
expect(data.items[0].rolloutId).toBe('ro-story-005');
|
||||
expect(data.items[0].status).toBeDefined();
|
||||
});
|
||||
|
||||
it('includes attempts directly on rollout payloads when they exist', async () => {
|
||||
const store = createServerBackedStore();
|
||||
const queryArgs = selectRolloutsQueryArgs(store.getState());
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getRollouts.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
const rolloutWithAttempt = data.items.find((rollout) => rollout.rolloutId === 'ro-story-002');
|
||||
expect(rolloutWithAttempt).toBeDefined();
|
||||
expect(rolloutWithAttempt?.attempt).not.toBeNull();
|
||||
expect(rolloutWithAttempt?.attempt?.attemptId).toBe('at-story-022');
|
||||
|
||||
const rolloutWithoutAttempt = data.items.find((rollout) => rollout.rolloutId === 'ro-story-004');
|
||||
expect(rolloutWithoutAttempt).toBeDefined();
|
||||
expect(rolloutWithoutAttempt?.attempt).toBeNull();
|
||||
});
|
||||
|
||||
it('retrieves attempts for a rollout from the Python server', async () => {
|
||||
const store = createServerBackedStore();
|
||||
const subscription = store.dispatch(
|
||||
rolloutsApi.endpoints.getRolloutAttempts.initiate({ rolloutId: 'ro-story-002' }),
|
||||
);
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.total).toBe(2);
|
||||
expect(data.items.map((attempt) => attempt.attemptId)).toEqual(['at-story-021', 'at-story-022']);
|
||||
});
|
||||
|
||||
it('paginates rollouts with custom UI state', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setRolloutsRecordsPerPage(2));
|
||||
store.dispatch(setRolloutsPage(2));
|
||||
|
||||
const queryArgs = selectRolloutsQueryArgs(store.getState());
|
||||
expect(queryArgs).toMatchObject({ limit: 2, offset: 2 });
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getRollouts.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.items).toHaveLength(2);
|
||||
expect(data.items.map((rollout) => rollout.rolloutId)).toEqual(['ro-story-004', 'ro-story-002']);
|
||||
});
|
||||
|
||||
it('filters and sorts rollouts based on UI selections', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(resetRolloutsFilters());
|
||||
store.dispatch(setRolloutsStatusFilters(['succeeded']));
|
||||
store.dispatch(setRolloutsModeFilters(['val']));
|
||||
store.dispatch(setRolloutsSearchTerm('ro-story-002'));
|
||||
store.dispatch(setRolloutsSort({ column: 'rolloutId', direction: 'asc' }));
|
||||
|
||||
const queryArgs = selectRolloutsQueryArgs(store.getState());
|
||||
expect(queryArgs).toMatchObject({
|
||||
statusIn: ['succeeded'],
|
||||
modeIn: ['val'],
|
||||
rolloutIdContains: 'ro-story-002',
|
||||
sortBy: 'rollout_id',
|
||||
sortOrder: 'asc',
|
||||
});
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getRollouts.initiate(queryArgs));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.items).toHaveLength(1);
|
||||
expect(data.items[0].rolloutId).toBe('ro-story-002');
|
||||
expect(data.items[0].status).toBe('succeeded');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,63 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSelector } from '@reduxjs/toolkit';
|
||||
import type { RootState } from '@/store';
|
||||
import type { RolloutMode, RolloutStatus } from '@/types';
|
||||
import type { RolloutsSortState } from './slice';
|
||||
|
||||
const ROLLOUTS_SORT_FIELD_MAP: Record<string, string> = {
|
||||
rolloutId: 'rollout_id',
|
||||
attemptId: 'attempt_id',
|
||||
statusValue: 'status',
|
||||
resourcesId: 'resources_id',
|
||||
mode: 'mode',
|
||||
startTimestamp: 'start_time',
|
||||
durationSeconds: 'duration',
|
||||
lastHeartbeatTimestamp: 'last_heartbeat_time',
|
||||
workerId: 'worker_id',
|
||||
};
|
||||
|
||||
const resolveRolloutsSortField = (sort: RolloutsSortState): string =>
|
||||
ROLLOUTS_SORT_FIELD_MAP[sort.column] ?? 'start_time';
|
||||
|
||||
export const selectRolloutsUiState = (state: RootState) => state.rollouts;
|
||||
export const selectRolloutsSearchTerm = (state: RootState) => state.rollouts.searchTerm;
|
||||
export const selectRolloutsStatusFilters = (state: RootState) => state.rollouts.statusFilters;
|
||||
export const selectRolloutsModeFilters = (state: RootState) => state.rollouts.modeFilters;
|
||||
export const selectRolloutsPage = (state: RootState) => state.rollouts.page;
|
||||
export const selectRolloutsRecordsPerPage = (state: RootState) => state.rollouts.recordsPerPage;
|
||||
export const selectRolloutsSort = (state: RootState) => state.rollouts.sort;
|
||||
|
||||
export const selectRolloutsQueryArgs = createSelector(
|
||||
[
|
||||
selectRolloutsSearchTerm,
|
||||
selectRolloutsStatusFilters,
|
||||
selectRolloutsModeFilters,
|
||||
selectRolloutsPage,
|
||||
selectRolloutsRecordsPerPage,
|
||||
selectRolloutsSort,
|
||||
],
|
||||
(
|
||||
searchTerm: string,
|
||||
statusFilters: RolloutStatus[],
|
||||
modeFilters: RolloutMode[],
|
||||
page: number,
|
||||
recordsPerPage: number,
|
||||
sort: RolloutsSortState,
|
||||
) => {
|
||||
const normalizedSearch = searchTerm.trim();
|
||||
const limit = Math.max(1, recordsPerPage);
|
||||
const offset = Math.max(0, (page - 1) * limit);
|
||||
const sortBy = resolveRolloutsSortField(sort);
|
||||
|
||||
return {
|
||||
limit,
|
||||
offset,
|
||||
sortBy,
|
||||
sortOrder: sort.direction,
|
||||
statusIn: statusFilters.length > 0 ? statusFilters : undefined,
|
||||
rolloutIdContains: normalizedSearch.length > 0 ? normalizedSearch : undefined,
|
||||
modeIn: modeFilters.length > 0 ? modeFilters : undefined,
|
||||
};
|
||||
},
|
||||
);
|
||||
@@ -0,0 +1,80 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSlice, PayloadAction } from '@reduxjs/toolkit';
|
||||
import type { RolloutMode, RolloutStatus } from '../../types';
|
||||
|
||||
export type SortDirection = 'asc' | 'desc';
|
||||
|
||||
export type RolloutsSortState = {
|
||||
column: string;
|
||||
direction: SortDirection;
|
||||
};
|
||||
|
||||
export type RolloutsUiState = {
|
||||
searchTerm: string;
|
||||
statusFilters: RolloutStatus[];
|
||||
modeFilters: RolloutMode[];
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
sort: RolloutsSortState;
|
||||
};
|
||||
|
||||
export const initialRolloutsUiState: RolloutsUiState = {
|
||||
searchTerm: '',
|
||||
statusFilters: [],
|
||||
modeFilters: [],
|
||||
page: 1,
|
||||
recordsPerPage: 100,
|
||||
sort: {
|
||||
column: 'startTimestamp',
|
||||
direction: 'desc',
|
||||
},
|
||||
};
|
||||
|
||||
const rolloutsSlice = createSlice({
|
||||
name: 'rollouts',
|
||||
initialState: initialRolloutsUiState,
|
||||
reducers: {
|
||||
setRolloutsSearchTerm(state, action: PayloadAction<string>) {
|
||||
state.searchTerm = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setRolloutsStatusFilters(state, action: PayloadAction<RolloutStatus[]>) {
|
||||
state.statusFilters = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setRolloutsModeFilters(state, action: PayloadAction<RolloutMode[]>) {
|
||||
state.modeFilters = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setRolloutsPage(state, action: PayloadAction<number>) {
|
||||
state.page = action.payload;
|
||||
},
|
||||
setRolloutsRecordsPerPage(state, action: PayloadAction<number>) {
|
||||
state.recordsPerPage = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setRolloutsSort(state, action: PayloadAction<RolloutsSortState>) {
|
||||
state.sort = action.payload;
|
||||
},
|
||||
resetRolloutsFilters(state) {
|
||||
state.statusFilters = initialRolloutsUiState.statusFilters;
|
||||
state.modeFilters = initialRolloutsUiState.modeFilters;
|
||||
state.searchTerm = initialRolloutsUiState.searchTerm;
|
||||
state.page = initialRolloutsUiState.page;
|
||||
state.sort = initialRolloutsUiState.sort;
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
export const {
|
||||
setRolloutsSearchTerm,
|
||||
setRolloutsStatusFilters,
|
||||
setRolloutsModeFilters,
|
||||
setRolloutsPage,
|
||||
setRolloutsRecordsPerPage,
|
||||
setRolloutsSort,
|
||||
resetRolloutsFilters,
|
||||
} = rolloutsSlice.actions;
|
||||
|
||||
export const rolloutsReducer = rolloutsSlice.reducer;
|
||||
@@ -0,0 +1,4 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export * from './slice';
|
||||
export * from './selectors';
|
||||
@@ -0,0 +1,69 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSelector } from '@reduxjs/toolkit';
|
||||
import type { GetSpansQueryArgs } from '@/features/rollouts';
|
||||
import type { RootState } from '@/store';
|
||||
import type { TracesSortState } from './slice';
|
||||
|
||||
export const selectTracesState = (state: RootState) => state.traces;
|
||||
|
||||
export const selectTracesRolloutId = (state: RootState) => selectTracesState(state).rolloutId;
|
||||
|
||||
export const selectTracesAttemptId = (state: RootState) => selectTracesState(state).attemptId;
|
||||
|
||||
export const selectTracesSearchTerm = (state: RootState) => selectTracesState(state).searchTerm;
|
||||
|
||||
export const selectTracesPage = (state: RootState) => selectTracesState(state).page;
|
||||
|
||||
export const selectTracesRecordsPerPage = (state: RootState) => selectTracesState(state).recordsPerPage;
|
||||
|
||||
export const selectTracesSort = (state: RootState) => selectTracesState(state).sort;
|
||||
|
||||
export const selectTracesViewMode = (state: RootState) => selectTracesState(state).viewMode;
|
||||
|
||||
const TRACES_SORT_FIELD_MAP: Record<string, string> = {
|
||||
name: 'name',
|
||||
traceId: 'trace_id',
|
||||
spanId: 'span_id',
|
||||
parentId: 'parent_id',
|
||||
statusCode: 'status_code',
|
||||
startTime: 'start_time',
|
||||
duration: 'duration',
|
||||
};
|
||||
|
||||
const resolveTracesSortField = (sort: TracesSortState): string => TRACES_SORT_FIELD_MAP[sort.column] ?? 'start_time';
|
||||
|
||||
export const selectTracesQueryArgs = createSelector(
|
||||
[
|
||||
selectTracesRolloutId,
|
||||
selectTracesAttemptId,
|
||||
selectTracesSearchTerm,
|
||||
selectTracesPage,
|
||||
selectTracesRecordsPerPage,
|
||||
selectTracesSort,
|
||||
],
|
||||
(rolloutId, attemptId, searchTerm, page, recordsPerPage, sort): GetSpansQueryArgs | undefined => {
|
||||
if (!rolloutId) {
|
||||
return undefined;
|
||||
}
|
||||
const limit = Math.max(1, recordsPerPage);
|
||||
const offset = Math.max(0, (page - 1) * limit);
|
||||
const normalizedSearch = searchTerm.trim();
|
||||
const containsValue = normalizedSearch.length > 0 ? normalizedSearch : undefined;
|
||||
|
||||
const sortBy = resolveTracesSortField(sort);
|
||||
|
||||
return {
|
||||
rolloutId,
|
||||
attemptId: attemptId ?? undefined,
|
||||
limit,
|
||||
offset,
|
||||
sortBy,
|
||||
sortOrder: sort.direction,
|
||||
traceIdContains: containsValue,
|
||||
spanIdContains: containsValue,
|
||||
nameContains: containsValue,
|
||||
filterLogic: containsValue ? 'or' : undefined,
|
||||
};
|
||||
},
|
||||
);
|
||||
@@ -0,0 +1,115 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
|
||||
|
||||
export type SortDirection = 'asc' | 'desc';
|
||||
|
||||
export type TracesSortState = {
|
||||
column: string;
|
||||
direction: SortDirection;
|
||||
};
|
||||
|
||||
export type TracesViewMode = 'table' | 'waterfall' | 'tree';
|
||||
|
||||
export type TracesUiState = {
|
||||
rolloutId: string | null;
|
||||
attemptId: string | null;
|
||||
searchTerm: string;
|
||||
page: number;
|
||||
recordsPerPage: number;
|
||||
sort: TracesSortState;
|
||||
viewMode: TracesViewMode;
|
||||
};
|
||||
|
||||
export const initialTracesUiState: TracesUiState = {
|
||||
rolloutId: null,
|
||||
attemptId: null,
|
||||
searchTerm: '',
|
||||
page: 1,
|
||||
recordsPerPage: 100,
|
||||
sort: {
|
||||
column: 'startTime',
|
||||
direction: 'desc',
|
||||
},
|
||||
viewMode: 'table',
|
||||
};
|
||||
|
||||
const tracesSlice = createSlice({
|
||||
name: 'traces',
|
||||
initialState: initialTracesUiState,
|
||||
reducers: {
|
||||
setTracesRolloutId(state, action: PayloadAction<string | null>) {
|
||||
state.rolloutId = action.payload;
|
||||
state.page = 1;
|
||||
state.attemptId = null;
|
||||
},
|
||||
setTracesAttemptId(state, action: PayloadAction<string | null>) {
|
||||
state.attemptId = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setTracesSearchTerm(state, action: PayloadAction<string>) {
|
||||
state.searchTerm = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setTracesPage(state, action: PayloadAction<number>) {
|
||||
state.page = action.payload;
|
||||
},
|
||||
setTracesRecordsPerPage(state, action: PayloadAction<number>) {
|
||||
state.recordsPerPage = action.payload;
|
||||
state.page = 1;
|
||||
},
|
||||
setTracesSort(state, action: PayloadAction<TracesSortState>) {
|
||||
state.sort = action.payload;
|
||||
},
|
||||
setTracesViewMode(state, action: PayloadAction<TracesViewMode>) {
|
||||
state.viewMode = action.payload;
|
||||
},
|
||||
resetTracesFilters(state) {
|
||||
state.searchTerm = initialTracesUiState.searchTerm;
|
||||
state.page = initialTracesUiState.page;
|
||||
state.sort = initialTracesUiState.sort;
|
||||
},
|
||||
hydrateTracesStateFromQuery(
|
||||
state,
|
||||
action: PayloadAction<{ rolloutId?: string | null; attemptId?: string | null }>,
|
||||
) {
|
||||
const payload = action.payload;
|
||||
if (Object.hasOwn(payload, 'rolloutId')) {
|
||||
const nextRolloutId = payload.rolloutId ?? null;
|
||||
if (state.rolloutId !== nextRolloutId) {
|
||||
state.rolloutId = nextRolloutId;
|
||||
state.page = 1;
|
||||
state.attemptId = null;
|
||||
} else if (nextRolloutId === null) {
|
||||
state.attemptId = null;
|
||||
}
|
||||
}
|
||||
|
||||
if (Object.hasOwn(payload, 'attemptId')) {
|
||||
if (state.rolloutId === null) {
|
||||
state.attemptId = null;
|
||||
return;
|
||||
}
|
||||
const nextAttemptId = payload.attemptId ?? null;
|
||||
if (state.attemptId !== nextAttemptId) {
|
||||
state.attemptId = nextAttemptId;
|
||||
state.page = 1;
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
export const {
|
||||
setTracesRolloutId,
|
||||
setTracesAttemptId,
|
||||
setTracesSearchTerm,
|
||||
setTracesPage,
|
||||
setTracesRecordsPerPage,
|
||||
setTracesSort,
|
||||
setTracesViewMode,
|
||||
resetTracesFilters,
|
||||
hydrateTracesStateFromQuery,
|
||||
} = tracesSlice.actions;
|
||||
|
||||
export const tracesReducer = tracesSlice.reducer;
|
||||
@@ -0,0 +1,98 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createServerBackedStore } from '@test-utils';
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { rolloutsApi } from '../rollouts';
|
||||
import { selectTracesQueryArgs } from './selectors';
|
||||
import { setTracesPage, setTracesRecordsPerPage, setTracesRolloutId, setTracesSearchTerm } from './slice';
|
||||
|
||||
describe('traces feature integration', () => {
|
||||
it('requires a rollout id before building query arguments', () => {
|
||||
const store = createServerBackedStore();
|
||||
expect(selectTracesQueryArgs(store.getState())).toBeUndefined();
|
||||
});
|
||||
|
||||
it('builds query arguments when targeting a rollout', () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setTracesRolloutId('ro-story-001'));
|
||||
|
||||
const queryArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(queryArgs).toBeDefined();
|
||||
expect(queryArgs).toMatchObject({
|
||||
rolloutId: 'ro-story-001',
|
||||
limit: 100,
|
||||
offset: 0,
|
||||
sortBy: 'start_time',
|
||||
sortOrder: 'desc',
|
||||
filterLogic: undefined,
|
||||
});
|
||||
|
||||
store.dispatch(setTracesSearchTerm('span-00'));
|
||||
const filteredArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(filteredArgs?.filterLogic).toBe('or');
|
||||
});
|
||||
|
||||
it('fetches spans from the Python LightningStore server', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setTracesRolloutId('ro-story-001'));
|
||||
|
||||
const queryArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(queryArgs).toBeDefined();
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getSpans.initiate(queryArgs!));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.total).toBe(3);
|
||||
const spanIds = data.items.map((span) => span.spanId).sort();
|
||||
expect(spanIds).toEqual(['span-001-root', 'span-002-llm', 'span-003-tool']);
|
||||
expect(new Set(data.items.map((span) => span.status.status_code))).toEqual(new Set(['OK']));
|
||||
});
|
||||
|
||||
it('paginates spans based on UI state', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setTracesRolloutId('ro-story-001'));
|
||||
store.dispatch(setTracesRecordsPerPage(2));
|
||||
|
||||
const firstPageArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(firstPageArgs).toMatchObject({ limit: 2, offset: 0 });
|
||||
|
||||
const firstPageSub = store.dispatch(rolloutsApi.endpoints.getSpans.initiate(firstPageArgs!));
|
||||
const firstPage = await firstPageSub.unwrap();
|
||||
firstPageSub.unsubscribe();
|
||||
|
||||
expect(firstPage.items.map((span) => span.spanId)).toEqual(['span-003-tool', 'span-002-llm']);
|
||||
|
||||
store.dispatch(setTracesPage(2));
|
||||
const secondPageArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(secondPageArgs).toMatchObject({ limit: 2, offset: 2 });
|
||||
|
||||
const secondPageSub = store.dispatch(rolloutsApi.endpoints.getSpans.initiate(secondPageArgs!));
|
||||
const secondPage = await secondPageSub.unwrap();
|
||||
secondPageSub.unsubscribe();
|
||||
|
||||
expect(secondPage.items.map((span) => span.spanId)).toEqual(['span-001-root']);
|
||||
expect(secondPage.total).toBe(3);
|
||||
});
|
||||
|
||||
it('filters spans using the search term', async () => {
|
||||
const store = createServerBackedStore();
|
||||
store.dispatch(setTracesRolloutId('ro-story-001'));
|
||||
store.dispatch(setTracesSearchTerm('span-003'));
|
||||
|
||||
const queryArgs = selectTracesQueryArgs(store.getState());
|
||||
expect(queryArgs).toMatchObject({
|
||||
traceIdContains: 'span-003',
|
||||
spanIdContains: 'span-003',
|
||||
nameContains: 'span-003',
|
||||
filterLogic: 'or',
|
||||
});
|
||||
|
||||
const subscription = store.dispatch(rolloutsApi.endpoints.getSpans.initiate(queryArgs!));
|
||||
const data = await subscription.unwrap();
|
||||
subscription.unsubscribe();
|
||||
|
||||
expect(data.items).toHaveLength(1);
|
||||
expect(data.items[0].spanId).toBe('span-003-tool');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,4 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
export * from './slice';
|
||||
export * from './selectors';
|
||||
@@ -0,0 +1,34 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import { createSelector } from '@reduxjs/toolkit';
|
||||
import type { RootState } from '@/store';
|
||||
import type { AlertsState, AlertTone } from './slice';
|
||||
|
||||
const ALERT_PRIORITY: Record<AlertTone, number> = {
|
||||
error: 3,
|
||||
warning: 2,
|
||||
info: 1,
|
||||
};
|
||||
|
||||
const selectAlertState = (state: RootState): AlertsState => state.alert;
|
||||
|
||||
export const selectVisibleAlerts = createSelector(selectAlertState, (state) =>
|
||||
state.alerts.filter((alert) => alert.isVisible),
|
||||
);
|
||||
|
||||
export const selectHighestPriorityAlert = createSelector(selectVisibleAlerts, (alerts) => {
|
||||
if (alerts.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return alerts
|
||||
.slice()
|
||||
.sort((a, b) => {
|
||||
const priorityDiff = ALERT_PRIORITY[a.tone] - ALERT_PRIORITY[b.tone];
|
||||
if (priorityDiff !== 0) {
|
||||
return priorityDiff;
|
||||
}
|
||||
return a.createdAt - b.createdAt;
|
||||
})
|
||||
.at(-1)!;
|
||||
});
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user