Compare commits

...

67 Commits

Author SHA1 Message Date
Yuqing dc2c5e5d8b Merge branch 'main' into dev/database-store 2025-11-12 12:02:35 +08:00
Yuge Zhang 5ae7933d41 Use unified server launcher for LiteLLM Proxy (#292) 2025-11-12 02:45:29 +08:00
yuqing 2eb207ec79 fix lint warning and update uv.lock 2025-11-11 16:53:01 +08:00
Yuge Zhang 2ab977ed18 Dashboard - build into Python package (#291) 2025-11-11 16:29:02 +08:00
Yuge Zhang 1eae9a34f0 Fix dashboard pipeline (#289) 2025-11-11 00:27:02 +08:00
Yuge Zhang 4e7748b059 Dashboard - tests and infrastructure (#288) 2025-11-10 22:30:00 +08:00
Yuge Zhang 582f67cade Python Server Launcher (#286) 2025-11-10 16:08:05 +08:00
부창규 9e23ba6b50 Rename the function properly in Spider (#285) 2025-11-10 14:57:06 +08:00
yuqing de8df23805 Merge branch 'dev/database-store' of https://github.com/microsoft/agent-lightning into dev/database-store 2025-11-07 12:05:59 +08:00
yuqing 5fc8ee0bf0 fix lint errors 2025-11-07 11:58:35 +08:00
yuqing fe2d05218f To break the big transaction inside timeout healthy checking 2025-11-07 11:05:12 +08:00
yuqing c1b7827e5f fix lint error 2025-11-06 15:54:19 +08:00
yuqing 3162ed6fb6 fix lint issue 2025-11-06 15:39:23 +08:00
Yuqing 5983eec570 Update agentlightning/store/database/retry_helper.py
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2025-11-06 15:23:59 +08:00
Yuqing c21e065d24 Update agentlightning/store/database/orm/rollout.py
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2025-11-06 15:23:29 +08:00
yuqing b6312db2a5 fix pre-commit warnings 2025-11-06 15:22:54 +08:00
yuqing 6131c1f7b5 rename to SqlLightningStore
operations with SQLAlchemy.
- Refactored existing tests to utilize the new SqlLightningStore, ensuring compatibility with the new implementation.
- Adjusted timeout behavior in tests to align with the new polling intervals.
2025-11-06 15:07:30 +08:00
yuqing 294a1cc1f6 Merge remote-tracking branch 'origin/main' into dev/database-store 2025-11-06 14:41:14 +08:00
yuqing 93814a9a9c Enhance rollout and attempt handling with type unification and new resource management features
- Updated type hints to include `Union` for rollouts and attempted rollouts across multiple files.
- Improved `query_rollouts` and `get_rollout_by_id` methods to return `AttemptedRollout` where applicable.
- Added `query_resources` method for resource management in `DatabaseLightningStore`.
- Refactored `as_attempt` and `as_rollout` methods to utilize `model_dump` for better serialization. (Serialization cost is reduced)
- Updated tests to validate status transitions and ensure correct behavior with new rollout logic.
2025-11-06 11:09:24 +08:00
Yuge Zhang 3f8a3ac0f1 Preserve interface for SQL store testing (#279) 2025-11-06 00:05:40 +08:00
yuqing 70838ee86c Merge remote-tracking branch 'origin/main' into dev/database-store 2025-11-05 21:51:28 +08:00
Yuge Zhang e0b55ab057 Fix preparing status transition on rollout when creating attempts (#278) 2025-11-05 17:45:57 +08:00
yuqing 5950276103 fix typo 2025-11-05 15:51:24 +08:00
Yuge Zhang 421f2773c7 RESTful API improvements (#275) 2025-11-05 14:12:58 +08:00
yuqing 1857e39d5e reuse test_memory.py 2025-11-05 12:05:43 +08:00
Yuge Zhang f717f9982f Fix: Port conflict in tracer tests (#271) 2025-11-05 11:10:05 +08:00
Shenghua Chen 44dbfde0b4 fix room_selector example which always run the first task (#270) 2025-11-05 10:47:07 +08:00
yuqing f7fe24a4b9 all tests passed with some FIXME 2025-11-04 22:49:31 +08:00
yuqing a477bec3e1 only corner cases for status propagation 2025-11-04 20:33:29 +08:00
Yuge Zhang 713511902d Add Tinker × Agent-lightning tuning articles to docs (#269) 2025-11-04 15:43:24 +08:00
Ni Hao 80531c9c28 fix openai_agent version for compatibility issue. (#265)
* fix openai_agent version for compatibility issue.

* gen uv.lock

---------

Co-authored-by: Hao Ni (CSI Interfusion Co Ltd) <v-nhao@microsoft.com>
2025-11-04 14:54:44 +08:00
Yuge Zhang 37daf2104f Adding VERL replacement for Tinker (#264)
* Adding VERL replacement for Tinker

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2025-11-04 12:09:36 +08:00
yuqing cbd6498e73 update error messages 2025-11-04 10:50:10 +08:00
yuqing a0b68333b8 support periodic background tasks for attempt timeout checking 2025-11-04 10:14:33 +08:00
Yuge Zhang 9afdd4570c docs: add deepwiki badge to readme (#263) 2025-11-03 08:09:06 +00:00
yuqing 96a0d58fce configurable retry added 2025-11-03 15:55:40 +08:00
yuqing b8940feffb interface implemented, but still have some issues for the periodical operations 2025-11-03 11:35:49 +08:00
Vishal V 55fbe66fe7 docs: fix typos in train-first-agent.md (#260)
Co-authored-by: Vishal <VishalV@ibm.com>
2025-11-02 16:01:46 +08:00
Yuge Zhang 3794c97c1e Store RESTful API updates (#259) 2025-11-02 16:01:32 +08:00
Yuge Zhang 848623766d Azure OpenAI Finetuning example (#256) 2025-11-01 23:44:30 +08:00
Yuge Zhang 4cd09ec900 Add contributor and maintainer guides (#239) 2025-11-01 13:08:59 +08:00
yuqing 47b0a4e493 initial version with all dbstore interface implemented, but not all tests passed 2025-10-31 23:35:30 +08:00
Zhiyuan He 3ed5e1e5b5 Fix training metrics before and after processing (#145) 2025-10-31 23:09:10 +08:00
Ni Hao 3f372ff7b3 Replace AgentOps mock server with bypassable client (#202)
---------

Co-authored-by: Hao Ni (CSI Interfusion Co Ltd) <v-nhao@microsoft.com>
2025-10-31 23:05:39 +08:00
Yuge Zhang c453c41fd2 Fix tests-full failure to checkout PR branch (#253) 2025-10-31 11:29:17 +00:00
Yuge Zhang 5c8ac61af6 Bump version to 0.2.2 (#248) 2025-10-31 18:14:43 +08:00
Zhiyuan He a02e1b91d9 Add support for verl 0.6.0 (#246) 2025-10-31 15:30:38 +08:00
Yuge Zhang 496e793f0b Tinker Integration (#245) 2025-10-30 12:44:17 +08:00
Yuge Zhang 80d306ff54 [Release] v0.2.1 (#243) 2025-10-30 08:31:45 +08:00
Yuge Zhang 5f67bfe137 Normalize Store FastAPI (#241) 2025-10-29 21:43:37 +08:00
John Eismeier f8c45b6ca8 propose fix a couple of typos and avoid emacs backup files (#237)
Signed-off-by: John E <jeis4wpi@outlook.com>
2025-10-29 09:42:58 +08:00
ddsfda99 01955aead7 Fix store port conflict handling (issue #221) (#227) 2025-10-28 14:58:41 +08:00
Yuge Zhang 0a9e3d75f2 Fix CI GPU trigger (#234) 2025-10-28 14:55:18 +08:00
Ni Hao a3b2db18fa Make the number of tasks on the server and client consistent. (#187) 2025-10-28 14:27:46 +08:00
Yuge Zhang 268bd77ce6 Refine notes and triggering conditions (#233) 2025-10-28 12:30:37 +08:00
Yuge Zhang ab6ea3c131 Augment backport implementation (#230) 2025-10-28 03:57:16 +00:00
Yuge Zhang d16538da96 Internal API update in preparation for Tinker integration (#226) 2025-10-28 11:26:49 +08:00
Yuge Zhang 8ce40a0410 Add backport support (#229) 2025-10-28 11:21:36 +08:00
Yuge Zhang a9c7dbef22 Track Actions Status in Issue Comment Responder (#228) 2025-10-28 00:01:31 +08:00
Yuge Zhang e69d24f4a8 Fix response ID security (#217) 2025-10-27 21:07:10 +08:00
Yuge Zhang 955a0cc9a3 Revert #224 (#225) 2025-10-27 20:38:12 +08:00
Yuge Zhang 3966db6d2a Post comment on workflow run ready (#224) 2025-10-27 19:46:41 +08:00
Yuge Zhang 584600d72e Fix property list too long in issue responder (#223) 2025-10-27 11:20:34 +00:00
Yuge Zhang 955524658d Update CI triggering mechanism (#222) 2025-10-27 19:10:16 +08:00
Yuge Zhang 4d5e133a06 Add vLLM blog link to resources (#215) 2025-10-26 09:57:07 +08:00
Yuge Zhang df2a159b00 Add tutorial for launching workers on separate machines (#213) 2025-10-25 14:16:27 +08:00
Yuge Zhang 675fc86727 Issue comment responder (#214) 2025-10-25 13:00:02 +08:00
208 changed files with 39446 additions and 1401 deletions
+32
View File
@@ -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
+33
View File
@@ -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
+17 -19
View File
@@ -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
+42 -24
View File
@@ -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 }}
+21 -19
View File
@@ -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
+17 -19
View File
@@ -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
+17 -19
View File
@@ -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
+309
View File
@@ -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}'.`);
}
+8
View File
@@ -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: |
+8
View File
@@ -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
+25 -19
View File
@@ -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
+58 -3
View File
@@ -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
View File
@@ -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
+50
View File
@@ -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
'
+4 -1
View File
@@ -8,6 +8,7 @@
[![Documentation](https://img.shields.io/badge/GitHub%20Pages-Documentation-blue)](https://microsoft.github.io/agent-lightning/)
[![PyPI version](https://badge.fury.io/py/agentlightning.svg)](https://badge.fury.io/py/agentlightning)
[![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/microsoft/agent-lightning)
[![Discord](https://img.shields.io/badge/Discord-Join-5865F2?logo=discord&logoColor=white)](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 -1
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
__version__ = "0.2.1"
__version__ = "0.2.2"
from .adapter import *
from .algorithm import *
+56 -18
View File
@@ -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
+1 -35
View File
@@ -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.
+4 -2
View File
@@ -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
+43
View File
@@ -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(
-30
View File
@@ -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())
+20 -2
View File
@@ -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
+3 -1
View File
@@ -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] = []
+1
View File
@@ -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
+169 -140
View File
@@ -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
View File
@@ -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.
+10 -2
View File
@@ -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
+4 -1
View File
@@ -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",
]
+35 -4
View File
@@ -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
+186
View File
@@ -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)
+101
View File
@@ -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)
+685
View File
@@ -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()
+75 -10
View File
@@ -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
+10 -1
View File
@@ -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,
+27 -62
View File
@@ -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]:
"""
+1
View File
@@ -117,6 +117,7 @@ AttemptStatus = Literal[
]
"""The status of an attempt."""
RolloutMode = Literal["train", "val", "test"]
"""Possible rollout modes."""
+6
View File
@@ -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
+999
View File
@@ -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")
+3 -8
View File
@@ -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
+9 -1
View File
@@ -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()
+107 -2
View File
@@ -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 AgentLightnings 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()
+132
View File
@@ -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
+47
View File
@@ -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;
+12
View File
@@ -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`;
+20
View File
@@ -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;
+13
View File
@@ -0,0 +1,13 @@
// Copyright (c) Microsoft. All rights reserved.
export const allModes = {
MD: {
viewport: 'md',
},
LG: {
viewport: 'lg',
},
XL: {
viewport: 'xl',
},
} as const;
+79
View File
@@ -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];
+8
View File
@@ -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]);
+5
View File
@@ -0,0 +1,5 @@
# Generated files
dist
# Theme files
theme.css
+28
View File
@@ -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"]
}
]
}
}
+27
View File
@@ -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`
+50
View File
@@ -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',
},
},
]);
+10890
View File
File diff suppressed because it is too large Load Diff
+75
View File
@@ -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"
}
}
+16
View File
@@ -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',
},
},
},
};
+13
View File
@@ -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>
+3
View File
@@ -0,0 +1,3 @@
// Copyright (c) Microsoft. All rights reserved.
import '../src/main.js';
+30
View File
@@ -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>
);
}
+41
View File
@@ -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')} />
),
};
+436
View File
@@ -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)',
},
});
+4
View File
@@ -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

+4
View File
@@ -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;
+30
View File
@@ -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,
};
},
);
+64
View File
@@ -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;
+346
View File
@@ -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;
+6
View File
@@ -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,
};
},
);
+80
View File
@@ -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;
+4
View File
@@ -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,
};
},
);
+115
View File
@@ -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');
});
});
+4
View File
@@ -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