Compare commits

...

76 Commits

Author SHA1 Message Date
Yuge Zhang e6b09e3fea update labels 2025-11-15 15:56:47 +08:00
Yuge Zhang e9317ba4a1 . 2025-11-15 15:19:46 +08:00
Yuge Zhang 8b1514a7cb . 2025-11-15 14:59:49 +08:00
Yuge Zhang 055638af1a . 2025-11-15 14:44:42 +08:00
Yuge Zhang ed62a8e3f6 fix tests 2025-11-15 14:42:10 +08:00
Yuge Zhang c08dc622ff fix agentops instrumentation tests 2025-11-15 14:13:00 +08:00
Yuge Zhang 908773d9ca fix port conflict 2025-11-15 13:31:16 +08:00
Yuge Zhang bd050860f9 fix tests 2025-11-15 13:28:21 +08:00
Yuge Zhang 2cf14121f0 add minimal examples tests 2025-11-15 13:04:16 +08:00
Yuge Zhang 2228aedf45 update documentation 2025-11-15 12:43:36 +08:00
Yuge Zhang 852d76f8bb update readme 2025-11-15 12:32:42 +08:00
Yuge Zhang d89d67c445 update minimal examples 2025-11-15 12:19:07 +08:00
Yuge Zhang 0f6f54b0fa add tests 2025-11-15 12:00:38 +08:00
Yuge Zhang 4f9bab87a4 . 2025-11-15 11:29:47 +08:00
Yuge Zhang 528020d372 add tests to llm proxy 2025-11-15 11:24:49 +08:00
Yuge Zhang 1cbd1c5f41 add tests 2025-11-15 10:43:47 +08:00
Yuge Zhang 87f6aed26b patch candidates 2025-11-15 10:22:41 +08:00
Yuge Zhang c9f0b90918 resolve comments and fix tests 2025-11-15 10:04:00 +08:00
Yuge Zhang 057d1f8d59 delete unused files 2025-11-15 01:27:56 +08:00
Yuge Zhang c3403c4b94 fix agentops and otel 2025-11-15 01:25:15 +08:00
Yuge Zhang 5173326d90 . 2025-11-15 00:29:08 +08:00
Yuge Zhang 51849445d5 move agentops logic to otel 2025-11-15 00:13:54 +08:00
Yuge Zhang b9831fea5c extract common util 2025-11-14 18:30:04 +08:00
Yuge Zhang 2e19bc6456 support llm proxy 2025-11-14 18:15:33 +08:00
Yuge Zhang bd301fd9ee agentops+ 2025-11-14 17:54:42 +08:00
Yuge Zhang cea5baf26d support agentops 2025-11-14 17:33:44 +08:00
Yuge Zhang 705b3ba98e add new examples 2025-11-14 15:52:12 +08:00
Yuge Zhang 5874453dd1 span proto to pan 2025-11-14 15:17:32 +08:00
Yuge Zhang 926d1ec7e7 Merge branch 'main' of github.com:microsoft/agent-lightning into feature/otlp 2025-11-14 14:43:47 +08:00
Yuge Zhang b3640786be . 2025-11-14 14:43:17 +08:00
Yuge Zhang e49b75b7d8 Check all matching jobs per variant (#310) 2025-11-13 17:10:50 +00:00
Yuge Zhang e9b953e19d . 2025-11-14 01:02:20 +08:00
Yuge Zhang 2824fe7e12 . 2025-11-13 23:48:01 +08:00
Yuge Zhang 5106b73999 modify store 2025-11-13 23:45:53 +08:00
Yuge Zhang 85c581a41f update examples 2025-11-13 23:27:46 +08:00
Yuge Zhang 67e19143af first examples 2025-11-13 23:12:27 +08:00
Yuge Zhang eab691b1a1 Refactor logging (#306) 2025-11-13 22:48:52 +08:00
Yuge Zhang fd6494873d Make health timeout configurable (#305) 2025-11-13 19:46:02 +08:00
Yuge Zhang 6cbfc1fee0 Fix CI Badge and make Calc-X pipeline faster (#304) 2025-11-13 18:06:18 +08:00
Yuge Zhang b986ae132a Use PythonServerLauncher in LightningStoreServer (#303) 2025-11-13 14:22:54 +08:00
Yuge Zhang f24a47969e Increase graceful timeout on CI (#302) 2025-11-13 10:15:14 +08:00
Yuge Zhang a0bc1827d9 [Release] v0.2.2 (#298) 2025-11-12 23:54:35 +08:00
Yuge Zhang f2869cea30 Fix local model support in VERL (#299) 2025-11-12 22:56:10 +08:00
Geng Zhang 77cf447717 fix stream response for anthropic and openai api (#293)
Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>
2025-11-12 21:29:02 +08:00
Yuge Zhang 790ed3efb3 View worker status on Dashboard (#296) 2025-11-12 21:27:31 +08:00
Yuge Zhang 5ae7933d41 Use unified server launcher for LiteLLM Proxy (#292) 2025-11-12 02:45:29 +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
Yuge Zhang 3f8a3ac0f1 Preserve interface for SQL store testing (#279) 2025-11-06 00:05:40 +08:00
Yuge Zhang e0b55ab057 Fix preparing status transition on rollout when creating attempts (#278) 2025-11-05 17:45:57 +08:00
Yuge Zhang 421f2773c7 RESTful API improvements (#275) 2025-11-05 14:12:58 +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
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
Yuge Zhang 9afdd4570c docs: add deepwiki badge to readme (#263) 2025-11-03 08:09:06 +00: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
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
234 changed files with 44589 additions and 1725 deletions
+29
View File
@@ -0,0 +1,29 @@
name: Badge - Compatibility
on:
workflow_run:
workflows:
- Examples - Backward Compatibility
types: [completed]
workflow_dispatch:
permissions:
actions: read
contents: read
jobs:
badge:
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/github-script@v8
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const badgeAggregation = require('./scripts/badge_aggregation.js');
const dependencies = [
{ workflow: 'examples-compat.yml', label: 'examples-compat', variants: ['legacy', 'stable'] },
];
await badgeAggregation({ github, context, core, dependencies });
+31
View File
@@ -0,0 +1,31 @@
name: Badge - Unit Test
on:
workflow_run:
workflows:
- CPU Test
- GPU Test
types: [completed]
workflow_dispatch:
permissions:
actions: read
contents: read
jobs:
badge:
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'workflow_run' && github.event.workflow_run.head_branch == 'main') }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/github-script@v8
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const badgeAggregation = require('./scripts/badge_aggregation.js');
const dependencies = [
{ workflow: 'tests-full.yml', label: 'tests-full', variants: ['legacy', 'stable'] },
{ workflow: 'tests.yml', label: 'tests', variants: ['legacy', 'stable', 'Lint', 'documentation', 'JavaScript'] },
];
await badgeAggregation({ github, context, core, dependencies });
+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
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'APO - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+113 -11
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Calc-X - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
@@ -22,12 +22,12 @@ run-name: >-
|| format('Calc-X - {0}', github.event_name) }}
jobs:
calc-x:
calc-x-perf:
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 }})
name: Calc-X Performance (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
timeout-minutes: 90
strategy:
@@ -74,7 +74,7 @@ jobs:
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-calc-x-${{ matrix.python-version }}-${{ matrix.setup-script }}
name: dependencies-calc-x-performance-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
@@ -116,13 +116,11 @@ jobs:
# Don't ask why. Don't touch this.
- name: Calc-X training
run: |
set -ex
source .venv/bin/activate
cd examples/calc_x
../../scripts/restart_ray.sh
sleep 5
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci
sleep 10
python train_calc_agent.py --val-file data/test_mini.parquet --ci
shell: bash
env:
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
@@ -137,14 +135,118 @@ jobs:
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
- name: Calc-X training LLM Proxy
calc-x-variants:
if: >
github.event_name != 'repository_dispatch' ||
github.event.action == 'ci-calc-x' ||
github.event.action == 'ci-all'
name: Calc-X Variants (Python ${{ matrix.python-version }}, ${{ matrix.setup-script }})
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
timeout-minutes: 90
strategy:
matrix:
include:
- python-version: '3.10'
setup-script: 'legacy'
- python-version: '3.12'
setup-script: 'stable'
- python-version: '3.13'
setup-script: 'latest'
fail-fast: false
steps:
- name: Check GPU status
run: nvidia-smi
- name: Check disk space
run: df -h
- uses: actions/checkout@v4
with:
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
- uses: astral-sh/setup-uv@v7
with:
enable-cache: true
python-version: ${{ matrix.python-version }}
- name: Upgrade dependencies (latest)
run: uv lock --upgrade
if: matrix.setup-script == 'latest'
- name: Sync dependencies (latest)
run: |
uv sync --frozen --no-default-groups --extra verl \
--group dev --group experiment --group agents --group torch-gpu-stable
if: matrix.setup-script == 'latest'
- name: Sync dependencies (stable & legacy)
run: |
uv sync --frozen --no-default-groups --extra verl \
--group dev --group experiment --group agents --group torch-gpu-${{ matrix.setup-script }}
if: matrix.setup-script != 'latest'
- name: Freeze dependencies
run: |
set -ex
uv pip freeze | tee requirements-freeze.txt
echo "UV_LOCKED=1" >> $GITHUB_ENV
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-calc-x-variants-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
- name: Launch LiteLLM Proxy
run: |
./scripts/litellm_run.sh
env:
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
- name: Prepare Calc-X dataset
run: |
set -ex
cd examples/calc_x
uv run gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
unzip calc-x-data.zip -d data
rm calc-x-data.zip
- name: Calc-X MCP sanity check
run: |
set -ex
cd examples/calc_x
uv run tests/test_mcp_calculator.py
env:
OPENAI_API_BASE: http://localhost:12306/
OPENAI_API_KEY: dummy
- name: Calc-X sanity check
run: |
set -ex
cd examples/calc_x
uv run legacy_calc_agent_debug.py
env:
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
- name: Training with local model
run: |
set -ex
source .venv/bin/activate
cd examples/calc_x
../../scripts/restart_ray.sh
sleep 5
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci --llm-proxy
hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir data/qwen_model
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --model $(realpath data/qwen_model)
sleep 10
shell: bash
env:
WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
id: calc_x_train_local_model
- name: Training with LLM Proxy
run: |
set -ex
source .venv/bin/activate
cd examples/calc_x
../../scripts/restart_ray.sh
sleep 5
PYTHONUNBUFFERED=1 python train_calc_agent.py --val-file data/test_mini.parquet --ci-fast --llm-proxy
sleep 10
shell: bash
env:
@@ -152,7 +254,7 @@ jobs:
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
id: calc_x_train_llm_proxy
- name: Calc-X training with external store
- name: Training with external store
run: |
set -euo pipefail
source .venv/bin/activate
@@ -182,7 +284,7 @@ jobs:
WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
id: calc_x_train_external_store
- name: Calc-X training with role-based environment variables
- name: Training with role-based environment variables
run: |
set -euo pipefail
source .venv/bin/activate
+5 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Backward Compatibility - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
@@ -54,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
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Spider - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Unsloth - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+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
+186 -4
View File
@@ -9,12 +9,12 @@ on:
workflow_dispatch:
repository_dispatch:
types: [ci-unsloth, ci-all]
types: [ci-gpu, ci-all]
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'GPU Test - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
@@ -46,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
@@ -69,10 +69,18 @@ jobs:
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
name: dependencies-tests-full-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
- uses: actions/setup-node@v6
with:
node-version: '22'
- 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
@@ -87,3 +95,177 @@ jobs:
PYTEST_ADDOPTS: "--color=yes"
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
minimal-examples:
if: >
github.event_name != 'repository_dispatch' ||
github.event.action == 'ci-gpu' ||
github.event.action == 'ci-all'
name: Minimal Examples with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
timeout-minutes: 30
strategy:
matrix:
include:
- python-version: '3.10'
setup-script: 'legacy'
- python-version: '3.12'
setup-script: 'stable'
- python-version: '3.13'
setup-script: 'latest'
fail-fast: false
steps:
- name: Check GPU status
run: nvidia-smi
- uses: actions/checkout@v4
with:
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
- uses: astral-sh/setup-uv@v7
with:
enable-cache: true
python-version: ${{ matrix.python-version }}
- name: Upgrade dependencies (latest)
run: uv lock --upgrade
if: matrix.setup-script == 'latest'
- name: Sync dependencies (latest)
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-stable
if: matrix.setup-script == 'latest'
- name: Sync dependencies (stable & legacy)
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
if: matrix.setup-script != 'latest'
- name: Freeze dependencies
run: |
set -ex
uv pip freeze | tee requirements-freeze.txt
echo "UV_LOCKED=1" >> $GITHUB_ENV
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
- name: Launch LiteLLM Proxy
run: |
./scripts/litellm_run.sh
env:
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
- name: Write Traces via Otel Tracer
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python write_traces.py otel
- name: Write Traces via AgentOps Tracer
env:
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python write_traces.py agentops
- name: Write Traces via Otel Tracer with Client
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
agl store --port 45993 --log-level DEBUG &
sleep 5
python write_traces.py otel --use-client
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
while pgrep -f agl; do
echo "Waiting for agl to finish..."
sleep 5
done
- name: Write Traces via AgentOps Tracer with Client
env:
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
agl store --port 45993 --log-level DEBUG &
sleep 5
python write_traces.py agentops --use-client
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
while pgrep -f agl; do
echo "Waiting for agl to finish..."
sleep 5
done
- name: vLLM Server
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python vllm_server.py Qwen/Qwen2.5-0.5B-Instruct
- name: LLM Proxy (OpenAI backend)
env:
OPENAI_API_BASE: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python llm_proxy.py openai gpt-4.1-mini &
LLM_PROXY_READY=0
for attempt in $(seq 1 30); do
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
LLM_PROXY_READY=1
break
fi
sleep 2
done
if [[ "$LLM_PROXY_READY" != "1" ]]; then
echo "LLM proxy failed to become healthy" >&2
exit 1
fi
python llm_proxy.py test gpt-4.1-mini
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
while pgrep -f llm_proxy.py; do
echo "Waiting for llm_proxy.py to finish..."
sleep 5
done
- name: LLM Proxy (vLLM backend)
if: matrix.setup-script != 'legacy' # Skip if return_token_ids is not supported
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct &
LLM_PROXY_READY=0
for attempt in $(seq 1 30); do
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
LLM_PROXY_READY=1
break
fi
sleep 2
done
if [[ "$LLM_PROXY_READY" != "1" ]]; then
echo "LLM proxy failed to become healthy" >&2
exit 1
fi
python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
while pgrep -f llm_proxy.py; do
echo "Waiting for llm_proxy.py to finish..."
sleep 5
done
+56 -1
View File
@@ -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
+9
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,9 @@ cython_debug/
# Claude
.claude/*.local.json
# 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
'
+13 -4
View File
@@ -4,10 +4,11 @@
# Agent Lightning⚡
[![Test](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml)
[![Unit Tests](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml)
[![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.**
@@ -33,12 +34,19 @@ Read more on our [documentation website](https://microsoft.github.io/agent-light
pip install agentlightning
```
For the latest nightly build (cutting-edge features), you can install from Test PyPI:
```bash
pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ agentlightning
```
Please refer to our [installation guide](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
To start using Agent-lightning, check out our [documentation](https://microsoft.github.io/agent-lightning/) and [examples](./examples).
## ⚡ Articles
- 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.
@@ -67,10 +75,11 @@ No rewrites, no lock-in, just a clear path from first rollout to steady improvem
| Workflow | Status |
|----------|--------|
| CPU Tests | [![tests workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests.yml) |
| GPU Tests | [![tests-full workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/tests-full.yml) |
| Full Tests | [![tests summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
| UI Tests | [![UI Tests](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/dashboard.yml) |
| Examples Integration | [![examples summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-examples.yml) |
| Latest Dependency Compatibility | [![latest summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-latest.yml) |
| Legacy Examples Compatibility | [![examples compatibility workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-compat.yml) |
| Legacy Examples Compatibility | [![compat summary workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-compat.yml) |
## ⚡ Citation
@@ -90,7 +99,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.
+4 -2
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 *
@@ -10,7 +10,9 @@ from .emitter import *
from .execution import *
from .litagent import *
from .llm_proxy import *
from .logging import *
from .logging import configure_logger # deprecated # type: ignore
from .logging import setup as setup_logging # type: ignore
from .logging import setup_module as setup_module_logging # type: ignore
from .runner import *
from .server import AgentLightningServer # deprecated # type: ignore
from .store import *
+9
View File
@@ -429,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.
+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.
+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())
+29 -4
View File
@@ -6,23 +6,48 @@ from __future__ import annotations
import argparse
import asyncio
import logging
from typing import Iterable
from agentlightning.logging import configure_logger
from agentlightning import setup_logging
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.",
)
parser.add_argument(
"--log-level",
default="INFO",
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
help="Configure the logging level for the store.",
)
args = parser.parse_args(list(argv) if argv is not None else None)
configure_logger()
setup_logging(args.log_level)
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,
launch_mode="asyncio",
)
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] = []
+2 -2
View File
@@ -67,8 +67,8 @@ class ClientServerExecutionStrategy(ExecutionStrategy):
server_host: str | None = None,
server_port: int | None = None,
n_runners: int = 1,
graceful_timeout: float = 5.0,
terminate_timeout: float = 5.0,
graceful_timeout: float = 10.0,
terminate_timeout: float = 10.0,
main_process: Literal["algorithm", "runner"] = "algorithm",
managed_store: bool | None = None,
) -> None:
+118 -119
View File
@@ -4,28 +4,76 @@ from __future__ import annotations
import json
import logging
import multiprocessing
import signal
import socket
import time
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 agentlightning.utils.otlp import LightningStoreOTLPExporter
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.
By default, AgentOps exporters and clients will run in local mode
and will NOT attempt to communicate with the remote AgentOps service.
Args:
enabled: If True, enable all AgentOps exporters and clients.
All exporters and clients will operate in normal mode and send data
to the [AgentOps service](https://www.agentops.ai).
"""
global _agentops_service_enabled
_agentops_service_enabled = enabled
logger.info(f"AgentOps service enabled is set to {enabled}.")
def _patch_exporters():
import agentops.client.api
import agentops.sdk.core
agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter # type: ignore
agentops.sdk.core.OTLPMetricExporter = BypassableOTLPMetricExporter
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
agentops.sdk.core.OTLPSpanExporter = BypassableOTLPSpanExporter # type: ignore
agentops.client.api.V3Client = BypassableV3Client
agentops.client.api.V4Client = BypassableV4Client
def _unpatch_exporters():
import agentops.client.api
import agentops.sdk.core
agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter # type: ignore
agentops.sdk.core.OTLPMetricExporter = OTLPMetricExporter
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
agentops.sdk.core.OTLPSpanExporter = OTLPSpanExporter # type: ignore
agentops.client.api.V3Client = V3Client
agentops.client.api.V4Client = V4Client
def _unwrap_legacy_response(response: Any) -> Any:
if hasattr(response, "parse") and callable(response.parse):
return response.parse()
return response
def _patch_new_agentops():
@@ -44,6 +92,11 @@ def _patch_new_agentops():
@no_type_check
def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws): # type: ignore
attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
# In some cases, response is a openai._legacy_response.LegacyAPIResponse (e.g., LiteLLM, or LangChain),
# This is created by client.with_raw_response.create()
return_value = _unwrap_legacy_response(return_value)
if (
return_value is not None
and hasattr(return_value, "prompt_token_ids")
@@ -89,20 +142,6 @@ def _patch_new_agentops():
[logprob.model_dump() for logprob in first_choice.logprobs.refusal]
)
# For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
if (
return_value is not None
and hasattr(return_value, "http_response")
and return_value.http_response is not None
and hasattr(return_value.http_response, "json")
):
json_data = return_value.http_response.json()
if isinstance(json_data, dict):
if json_data.get("prompt_token_ids") is not None:
attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
if json_data.get("response_token_ids") is not None:
attributes["response_token_ids"] = list(json_data["response_token_ids"][0])
return attributes
agentops.instrumentation.providers.openai.wrappers.chat.handle_chat_attributes = _handle_chat_attributes_with_tokens
@@ -173,6 +212,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()
@@ -192,6 +233,8 @@ def instrument_agentops():
def uninstrument_agentops():
"""Uninstrument agentops to stop capturing token IDs."""
_unpatch_exporters()
try:
_unpatch_new_agentops()
except Exception:
@@ -202,114 +245,70 @@ def uninstrument_agentops():
pass
def agentops_local_server():
class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, 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 should_bypass(self) -> bool:
return not _agentops_service_enabled
def _run_server(**kwargs: Any): # type: ignore
class BypassableOTLPMetricExporter(OTLPMetricExporter):
"""
Internal function to run the Flask server.
This is used to avoid issues with multiprocessing and Flask's reloader.
OTLPMetricExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
"""
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) -> MetricExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs) # type: ignore[reportUnknownMemberType]
else:
logger.error(f"AgentOps local server failed to start or exited prematurely.")
return
logger.debug("SwitchableOTLPMetricExporter is switched off, skipping export.")
return MetricExportResult.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 BypassableOTLPSpanExporter(LightningStoreOTLPExporter):
"""
OTLPSpanExporter 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.")
This is used instead of BypassableAuthenticatedOTLPExporter on legacy AgentOps versions.
"""
def should_bypass(self) -> bool:
return not _agentops_service_enabled
class BypassableV3Client(V3Client):
"""
V3Client with toggleable authentication calls.
Returns dummy auth response when `_agentops_service_enabled` is False.
"""
# Temporary synchronous override of fetch_auth_token for mock purposes.
def fetch_auth_token(self, *args: Any, **kwargs: Any) -> AuthTokenResponse: # type: ignore[override]
if _agentops_service_enabled:
return super().fetch_auth_token(*args, **kwargs) # type: ignore[override]
else:
logger.info("AgentOps local server was not running or already stopped.")
logger.debug("SwitchableV3Client is switched off, skipping fetch_auth_token request.")
return AuthTokenResponse(token="dummy", project_id="dummy")
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 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
File diff suppressed because it is too large Load Diff
+329 -13
View File
@@ -1,10 +1,18 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
import os
import platform
import sys
import warnings
from logging.config import dictConfig
from typing import Any, Dict, Optional
__all__ = ["configure_logger"]
from rich.console import Console
__all__ = ["setup", "configure_logger", "setup_module"]
def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
@@ -15,6 +23,10 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
not propagate to the root logger, preventing duplicate log emission when
applications compose multiple logging configurations.
!!! danger
This function is deprecated in favor of [`setup_logging`][agentlightning.setup_logging].
Args:
level: Logging level applied both to the logger and the installed
handler. Defaults to `logging.INFO`.
@@ -32,23 +44,327 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
logger.info("agent-lightning is ready!")
```
"""
warnings.warn("This function is deprecated in favor of `setup_logging`.", DeprecationWarning, stacklevel=2)
return setup_module(level=level, name=name, console=True, color=True, propagate=False)
DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s) %(message)s"
DATE_FORMAT = "%H:%M:%S"
def _to_level_value(lvl: int | str) -> int:
if isinstance(lvl, int):
return lvl
val = getattr(logging, str(lvl).upper(), None)
if val is None:
raise ValueError(f"Invalid log level: {lvl}")
return val
def _ensure_file_handler(
logger: logging.Logger,
filename: str,
*,
level: int,
formatter: Optional[logging.Formatter],
) -> None:
"""Attach a FileHandler to `logger` for `filename` if it doesn't already exist."""
abspath = os.path.abspath(filename)
# Avoid duplicates
for h in logger.handlers:
if isinstance(h, logging.FileHandler) and getattr(h, "baseFilename", None) == abspath:
return
# Ensure directory exists
dirname = os.path.dirname(abspath)
if dirname:
os.makedirs(dirname, exist_ok=True)
fh = logging.FileHandler(abspath, encoding="utf-8")
fh.setLevel(level)
if formatter is not None:
fh.setFormatter(formatter)
else:
fh.setFormatter(logging.Formatter(DEFAULT_FORMAT, DATE_FORMAT))
logger.addHandler(fh)
def setup(
level: int | str = "INFO",
*,
console: bool = True,
color: bool | Dict[str, Any] = True,
propagate: bool = False,
disable_existing_loggers: bool = False,
capture_warnings: bool = False,
submodule_levels: Optional[dict[str, int | str]] = None,
extra_handlers: Optional[list[logging.Handler]] = None,
formatter: Optional[logging.Formatter] = None,
apply_to: Optional[list[str]] = None,
files: Optional[str | dict[str, str]] = None,
) -> None:
"""Configures logging for the `agentlightning` logger hierarchy.
This function provides a one-stop setup utility for configuring the
`agentlightning` root logger and optionally its submodules or external
loggers. It supports console logging, colored rich output, per-submodule
log levels, and optional handler/formatter injection.
The setup is intentionally isolated: it does not modify the global root
logger or loggers belonging to other libraries unless explicitly directed
via `apply_to`.
Args:
level:
Logging level for the base `agentlightning` logger. Accepts either
an integer (e.g., `logging.DEBUG`) or a string level name
(e.g., `"INFO"`). Defaults to `"INFO"`.
console:
Whether to attach a console handler to the logger. Defaults to
`True`.
color:
Enables rich-formatted output using `RichHandler` when `True`
or a configuration dict. If `False`, a plain text formatter is
used instead. Defaults to `True`.
propagate:
Whether `agentlightning` logs should propagate to ancestor
loggers. Defaults to `False`.
disable_existing_loggers:
Passed to `logging.config.dictConfig`. If `True`, disables all
existing configured loggers before applying this configuration.
Defaults to `False`.
capture_warnings:
If `True`, redirects Python `warnings` emitted via the `warnings`
module into the logging system. Defaults to `False`.
submodule_levels:
Mapping of submodule logger names to logging levels. If a specified
submodule level is more verbose than the base level, a warning is emitted.
extra_handlers:
A list of user-provided handlers to attach to the `agentlightning` logger.
Handlers are added idempotently; duplicates are not reattached.
formatter:
A formatter to apply to any handler under `agentlightning` that does not
already have one assigned. Useful for customizing output without overwriting
formatters on custom handlers.
apply_to:
A list of additional logger names to configure identically to
`agentlightning` base logger. Their handlers are replaced with copies of the base
handlers, and propagation is disabled to avoid duplicate log emission.
files:
If a string, attach a FileHandler to the base `agentlightning` logger.
If a dict, for each `(logger_name, filename)` pair, attach a FileHandler
directly to that logger.
Each file handler should use the logger's effective level at creation.
Notes:
* On Windows, this function forces UTF-8 mode in the console to prevent
issues with rich output or special characters.
* Submodule loggers can generate records below the handler's emission
threshold. Whether such records appear depends on both the logger's
level and the handler's level.
* `apply_to` loggers inherit the same handlers but do not propagate
upward, yielding isolated, consistent behavior.
Examples:
Basic setup:
>>> setup()
Enabling debug mode with no color:
>>> setup(level="DEBUG", color=False)
Overriding specific submodule levels:
>>> setup(submodule_levels={"agentlightning.io": "DEBUG"})
Attaching an additional file handler:
>>> fh = logging.FileHandler("app.log")
>>> setup(extra_handlers=[fh])
"""
# Ensure UTF-8 encoding on Windows consoles
# Note: This change does not fully represent support for execution under the windown system.
# Note: This change does not fully represent support for execution under the windows system.
# It only fixes console printing issues caused by special characters.
# TODO: More comprehensive Windows support may be needed in the future.
if platform.system() == "Windows":
os.environ["PYTHONUTF8"] = "1"
logger = logging.getLogger(name)
logger.handlers.clear() # clear existing handlers
base_logger = setup_module(
level,
name="agentlightning",
console=console,
color=color,
propagate=propagate,
disable_existing_loggers=disable_existing_loggers,
)
# log to stdout
handler = logging.StreamHandler()
handler.setLevel(level)
formatter = logging.Formatter("%(asctime)s [%(levelname)s] (Process-%(process)d %(name)s) %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(level)
logger.propagate = False # prevent double logging
return logger
base_level_value = base_logger.level
# Apply user-provided formatter (only to handlers without one,
# so we don't clobber custom extra_handlers)
if formatter is not None:
for h in base_logger.handlers:
if h.formatter is None:
h.setFormatter(formatter)
# Attach user-provided handler(s) if any, idempotently
if extra_handlers:
for h in extra_handlers:
if h not in base_logger.handlers:
base_logger.addHandler(h)
# Per-submodule levels
if submodule_levels:
for name, lvl in submodule_levels.items():
sub_level = _to_level_value(lvl)
# Emit a warning if submodule level is lower (more verbose) than the global/base level
if sub_level < base_level_value:
base_logger.warning(
"Submodule logger '%s' level %s (%s) is more verbose than base "
"logger level %s (%s). Records below the base level may still be "
"filtered out by handlers depending on their own levels.",
name,
lvl,
sub_level,
logging.getLevelName(base_level_value),
base_level_value,
)
# The logger will *create* records down to the logger's level, but a handler
# with a higher level will still drop anything below its own threshold.
# Effective emission is gated by both: record.level >= logger.level AND handler.level.
logging.getLogger(name).setLevel(lvl)
# Attach file handlers if requested
if files is not None:
if isinstance(files, str):
# Single file for the entire `agentlightning` hierarchy.
_ensure_file_handler(
logger=base_logger,
filename=files,
level=base_level_value,
formatter=formatter,
)
else:
# Per-logger files
for logger_name, filename in files.items():
lg = logging.getLogger(logger_name)
# Use the logger's *effective* level at creation time
effective_level = lg.getEffectiveLevel()
_ensure_file_handler(
logger=lg,
filename=filename,
level=effective_level,
formatter=formatter,
)
# Optionally apply the same handler setup to other loggers outside this module
if apply_to:
for name in apply_to:
lg = logging.getLogger(name)
# This removes any existing handlers so we don't duplicate output
# and ensures these loggers share exactly the same handlers as base_logger.
lg.handlers.clear()
for h in base_logger.handlers:
lg.addHandler(h)
lg.setLevel(base_logger.level)
# We've attached handlers directly to these loggers; if propagate
# stayed True, records would bubble up to ancestor loggers and could be
# emitted twice (here and on the parent/root). Setting False isolates them.
lg.propagate = False
# Optionally capture warnings
if capture_warnings:
logging.captureWarnings(True)
def setup_module(
level: int | str = "INFO",
*,
name: str = "agentlightning",
console: bool = True,
color: bool | Dict[str, Any] = True,
propagate: bool = False,
disable_existing_loggers: bool = False,
) -> logging.Logger:
"""Initializes and returns the base logger for `agentlightning`.
This function constructs and applies a `dictConfig` configuration for the
logger hierarchy rooted at `name`. It supports either rich console
formatting (via `RichHandler`) or plain text formatting, based on the
`color` argument.
Unlike [`setup_logging`][agentlightning.setup_logging], this function configures only a single logger namespace
and does not attach extra handlers or submodule levels. It is primarily used
internally by [`setup_logging`][agentlightning.setup_logging] but is also suitable for direct integration in
custom logging workflows.
"""
root_cfg: Dict[str, Any] = {
"version": 1,
"disable_existing_loggers": disable_existing_loggers,
"loggers": {
name: {
"handlers": [],
"level": level,
"propagate": propagate,
}
},
"handlers": {},
"formatters": {},
}
# Choose formatter / handler definition
if color is not False and console:
# Console must be true to display colored outputs
if isinstance(color, dict):
rich_handler_config = color
else:
rich_handler_config: Dict[str, Any] = {
"rich_tracebacks": False,
"markup": False,
"show_time": True,
"show_path": True,
}
if not _has_width():
# e.g., in a CI environment.
rich_handler_config["console"] = Console(width=200)
root_cfg["handlers"]["console"] = {
"class": "rich.logging.RichHandler",
"level": level,
**rich_handler_config,
}
# RichHandler manages its own style; keep formatter None
else:
fmt_name = "plain"
root_cfg["formatters"][fmt_name] = {
"format": DEFAULT_FORMAT,
"datefmt": DATE_FORMAT,
}
if console:
root_cfg["handlers"]["console"] = {
"class": "logging.StreamHandler",
"level": level,
"formatter": fmt_name,
}
# Attach selected handlers to agentlightning
handler_names = list(root_cfg["handlers"].keys())
root_cfg["loggers"][name]["handlers"] = handler_names
# Apply dictConfig (this resets the logger handlers)
dictConfig(root_cfg)
return logging.getLogger(name)
def _has_width() -> bool:
"""Automatically determine whether the terminal has a width."""
return sys.stdout.isatty()
+134 -31
View File
@@ -11,8 +11,21 @@ from __future__ import annotations
import asyncio
import logging
import threading
import time
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
from contextlib import suppress
from typing import (
TYPE_CHECKING,
Any,
Awaitable,
Callable,
List,
Literal,
Optional,
Sequence,
TypeVar,
cast,
)
from opentelemetry.sdk.trace import ReadableSpan
@@ -30,6 +43,7 @@ from agentlightning.types import (
RolloutRawResult,
Span,
)
from agentlightning.utils.system_snapshot import system_snapshot
if TYPE_CHECKING:
from agentlightning.execution.events import ExecutionEvent
@@ -52,7 +66,14 @@ class LitAgentRunner(Runner[T_task]):
worker_id: Identifier for the active worker process, if any.
"""
def __init__(self, tracer: Tracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
def __init__(
self,
tracer: Tracer,
max_rollouts: Optional[int] = None,
poll_interval: float = 5.0,
heartbeat_interval: float = 10.0,
heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
) -> None:
"""Initialize the agent runner.
Args:
@@ -60,11 +81,16 @@ class LitAgentRunner(Runner[T_task]):
max_rollouts: Optional cap on iterations processed by
[`iter`][agentlightning.LitAgentRunner.iter].
poll_interval: Seconds to wait between store polls when no work is available.
heartbeat_interval: Seconds to wait between sending heartbeats to the store.
heartbeat_launch_mode: Launch mode for the heartbeat loop. Can be "asyncio" or "thread".
"asyncio" is the default and recommended mode. Use "thread" if you are experiencing blocking coroutines.
"""
super().__init__()
self._tracer = tracer
self._max_rollouts = max_rollouts
self._poll_interval = poll_interval
self._heartbeat_interval = heartbeat_interval
self._heartbeat_launch_mode = heartbeat_launch_mode
# Set later
self._agent: Optional[LitAgent[T_task]] = None
@@ -304,6 +330,67 @@ class LitAgentRunner(Runner[T_task]):
return trace_spans
async def _emit_heartbeat(self, store: LightningStore) -> None:
"""Send a heartbeat tick to the store."""
worker_id = self.get_worker_id()
try:
await store.update_worker(worker_id, system_snapshot())
except asyncio.CancelledError:
# bypass the exception
raise
except Exception:
logger.exception("%s Unable to update worker heartbeat.", self._log_prefix())
def _start_heartbeat_loop(self, store: LightningStore) -> Optional[Callable[[], Awaitable[None]]]:
"""Start a background heartbeat loop and return an async stopper."""
if self._heartbeat_interval <= 0:
return None
if self.worker_id is None:
logger.warning("%s Cannot start heartbeat loop without worker_id.", self._log_prefix())
return None
if self._heartbeat_launch_mode == "asyncio":
stop_event = asyncio.Event()
async def heartbeat_loop() -> None:
while not stop_event.is_set():
await self._emit_heartbeat(store)
with suppress(asyncio.TimeoutError):
await asyncio.wait_for(stop_event.wait(), timeout=self._heartbeat_interval)
task = asyncio.create_task(heartbeat_loop(), name=f"{self.get_worker_id()}-heartbeat")
async def stop() -> None:
stop_event.set()
with suppress(asyncio.CancelledError):
await task
return stop
if self._heartbeat_launch_mode == "thread":
stop_evt = threading.Event()
def thread_worker() -> None:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
while not stop_evt.is_set():
loop.run_until_complete(self._emit_heartbeat(store))
stop_evt.wait(self._heartbeat_interval)
thread = threading.Thread(target=thread_worker, name=f"{self.get_worker_id()}-heartbeat", daemon=True)
thread.start()
async def stop() -> None:
stop_evt.set()
await asyncio.to_thread(thread.join)
return stop
raise ValueError(f"Unsupported heartbeat launch mode: {self._heartbeat_launch_mode}")
async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
"""Sleep until the next poll interval, with optional event-based interruption.
@@ -450,39 +537,49 @@ class LitAgentRunner(Runner[T_task]):
logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
store = self.get_store()
while not (event is not None and event.is_set()) and (
self._max_rollouts is None or num_tasks_processed < self._max_rollouts
):
# Retrieve the next rollout
next_rollout: Optional[Rollout] = None
while not (event is not None and event.is_set()):
logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
next_rollout = await store.dequeue_rollout()
stop_heartbeat = self._start_heartbeat_loop(store)
try:
while not (event is not None and event.is_set()) and (
self._max_rollouts is None or num_tasks_processed < self._max_rollouts
):
# Retrieve the next rollout
next_rollout: Optional[Rollout] = None
while not (event is not None and event.is_set()):
logger.debug(f"{self._log_prefix()} Try to poll for next rollout.")
next_rollout = await store.dequeue_rollout(worker_id=self.get_worker_id())
if next_rollout is None:
logger.debug(
f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds."
)
await self._sleep_until_next_poll(event)
else:
break
if next_rollout is None:
logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
await self._sleep_until_next_poll(event)
else:
break
return
if next_rollout is None:
return
try:
# Claim the rollout but updating the current worker id
await store.update_attempt(
next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
)
except Exception:
# This exception could happen if the rollout is dequeued and the other end died for some reason
logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
continue
try:
# Claim the rollout but updating the current worker id
await store.update_attempt(
next_rollout.rollout_id, next_rollout.attempt.attempt_id, worker_id=self.get_worker_id()
)
except Exception:
# This exception could happen if the rollout is dequeued and the other end died for some reason
logger.exception(f"{self._log_prefix()} Exception during update_attempt, giving up the rollout.")
continue
# Execute the step
await self._step_impl(next_rollout)
# Execute the step
await self._step_impl(next_rollout)
num_tasks_processed += 1
if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
logger.info(f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}")
num_tasks_processed += 1
if num_tasks_processed % 10 == 0 or num_tasks_processed == 1:
logger.info(
f"{self._log_prefix()} Progress: {num_tasks_processed}/{self._max_rollouts or 'unlimited'}"
)
finally:
if stop_heartbeat is not None:
await stop_heartbeat()
logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
@@ -526,6 +623,12 @@ class LitAgentRunner(Runner[T_task]):
resources_id = None
attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
# Register the attempt as running by the current worker
await self.get_store().update_attempt(
attempted_rollout.rollout_id,
attempted_rollout.attempt.attempt_id,
worker_id=self.get_worker_id(),
)
rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
completed_rollout = await store.get_rollout_by_id(rollout_id)
+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
+2 -1
View File
@@ -1,12 +1,13 @@
# Copyright (c) Microsoft. All rights reserved.
from .base import LightningStore
from .base import LightningStore, LightningStoreCapabilities
from .client_server import LightningStoreClient, LightningStoreServer
from .memory import InMemoryLightningStore
from .threading import LightningStoreThreaded
__all__ = [
"LightningStore",
"LightningStoreCapabilities",
"LightningStoreClient",
"LightningStoreServer",
"InMemoryLightningStore",
+107 -2
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, TypedDict
from opentelemetry.sdk.trace import ReadableSpan
@@ -17,6 +17,7 @@ from agentlightning.types import (
RolloutStatus,
Span,
TaskInput,
Worker,
)
@@ -52,6 +53,22 @@ UNSET = _UnsetType()
Unset = _UnsetType # Alias for convenience
class LightningStoreCapabilities(TypedDict, total=False):
"""Capability of a LightningStore implementation.
All keys are optional and false by default.
"""
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."""
otlp_traces: bool
"""Whether the store supports OTLP/HTTP traces."""
class LightningStore:
"""Contract for the persistent control-plane that coordinates training rollouts.
@@ -74,6 +91,32 @@ class LightningStore:
Unless stated otherwise, missing identifiers should result in a `ValueError`.
"""
@property
def capabilities(self) -> LightningStoreCapabilities:
"""Return the capabilities of the store."""
return LightningStoreCapabilities(
thread_safe=False,
async_safe=False,
zero_copy=False,
otlp_traces=False,
)
def otlp_traces_endpoint(self) -> str:
"""Return the OTLP/HTTP traces endpoint of the store.
The traces can have rollout ID and attempt ID (and optionally sequence ID)
saved in the "resource" of the spans.
The store, if it supports OTLP, should be able to receive the traces and save them
via [`add_span`][agentlightning.LightningStore.add_span] or
[`add_otel_span`][agentlightning.LightningStore.add_otel_span].
The endpoint should be compatible with [OTLP HTTP protocol](https://opentelemetry.io/docs/specs/otlp/).
It's not necessarily compatible with OTLP gRPC protocol.
The returned endpoint will usually ends with `/v1/traces`.
"""
raise NotImplementedError()
async def start_rollout(
self,
input: TaskInput,
@@ -148,7 +191,7 @@ class LightningStore:
"""
raise NotImplementedError()
async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
"""Claim the oldest queued rollout and transition it to `preparing`.
This function do not block.
@@ -161,6 +204,8 @@ class LightningStore:
the number of attempts already registered for the rollout plus one.
* Return an [`AttemptedRollout`][agentlightning.AttemptedRollout] snapshot so the
runner knows both rollout metadata and the attempt identifier.
* Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
(e.g., `last_dequeue_time`) when `worker_id` is provided.
Returns:
The next attempt to execute, or `None` when no eligible rollouts are queued.
@@ -307,6 +352,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.
@@ -497,6 +553,12 @@ class LightningStore:
Similar to [`update_rollout()`][agentlightning.LightningStore.update_rollout],
parameters also default to the sentinel [`UNSET`][agentlightning.store.base.UNSET].
If `worker_id` is present, the worker status will be updated following the rules:
1. If attempt status is "succeeded" or "failed", the corresponding worker status will be set to "idle".
2. If attempt status is "unresponsive" or "timeout", the corresponding worker status will be set to "unknown".
3. Otherwise, the worker status will be set to "busy".
Args:
rollout_id: Identifier of the rollout whose attempt will be updated.
attempt_id: Attempt identifier or `"latest"` as a convenience.
@@ -513,3 +575,46 @@ class LightningStore:
ValueError: Implementations must raise when the rollout or attempt is unknown.
"""
raise NotImplementedError()
async def query_workers(
self,
) -> List[Worker]:
"""Query all workers in the system.
Returns:
A list of all workers.
"""
raise NotImplementedError()
async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
"""Retrieve a single worker by identifier.
Args:
worker_id: Identifier of the worker.
Returns:
The worker record if it exists, otherwise `None`.
Raises:
NotImplementedError: Subclasses must implement lookup semantics.
"""
raise NotImplementedError()
async def update_worker(
self,
worker_id: str,
heartbeat_stats: Dict[str, Any] | Unset = UNSET,
) -> Worker:
"""Record a heartbeat for `worker_id` and refresh telemetry.
Implementations must treat this API as heartbeat-only: it should snapshot
the latest stats when provided, stamp `last_heartbeat_time` with the
current wall clock, and rely on other store mutations (`dequeue_rollout`,
`update_attempt`, etc.) to drive the worker's busy/idle status,
assignment, and activity timestamps.
Args:
worker_id: Identifier of the worker to update.
heartbeat_stats: Replacement worker heartbeat statistics (non-null when provided).
"""
raise NotImplementedError()
File diff suppressed because it is too large Load Diff
+175 -15
View File
@@ -26,6 +26,7 @@ from typing import (
Sequence,
Set,
TypeVar,
Union,
cast,
)
@@ -43,9 +44,10 @@ from agentlightning.types import (
RolloutStatus,
Span,
TaskInput,
Worker,
)
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])
@@ -241,6 +243,58 @@ class InMemoryLightningStore(LightningStore):
# Completion tracking for wait_for_rollouts (cross-loop safe)
self._completion_events: Dict[str, threading.Event] = {}
# Worker tracking
self._workers: Dict[str, Worker] = {}
# Running rollouts cache, including preparing and running rollouts
self._running_rollout_ids: Set[str] = set()
def _get_or_create_worker(self, worker_id: str) -> Worker:
worker = self._workers.get(worker_id)
if worker is None:
worker = Worker(worker_id=worker_id)
self._workers[worker_id] = worker
return worker
def _sync_worker_with_attempt(self, attempt: Attempt) -> None:
worker_id = attempt.worker_id
if not worker_id:
return
worker = self._get_or_create_worker(worker_id)
now = time.time()
if attempt.status in ("succeeded", "failed"):
if worker.status != "idle":
worker.last_idle_time = now
worker.status = "idle"
worker.current_rollout_id = None
worker.current_attempt_id = None
elif attempt.status in ("timeout", "unresponsive"):
if worker.status != "unknown":
worker.last_idle_time = now
worker.status = "unknown"
worker.current_rollout_id = None
worker.current_attempt_id = None
else:
transitioned = worker.status != "busy" or worker.current_attempt_id != attempt.attempt_id
if transitioned:
worker.last_busy_time = now
worker.status = "busy"
worker.current_rollout_id = attempt.rollout_id
worker.current_attempt_id = attempt.attempt_id
Worker.model_validate(worker.model_dump())
@property
def capabilities(self) -> LightningStoreCapabilities:
"""Return the capabilities of the store."""
return LightningStoreCapabilities(
thread_safe=False,
async_safe=True,
zero_copy=False,
otlp_traces=False,
)
@_healthcheck_wrapper
async def start_rollout(
@@ -272,6 +326,7 @@ class InMemoryLightningStore(LightningStore):
config=rollout_config,
metadata=rollout_metadata,
)
self._running_rollout_ids.add(rollout.rollout_id)
# Create the initial attempt
attempt_id = _generate_attempt_id()
@@ -286,7 +341,7 @@ class InMemoryLightningStore(LightningStore):
self._attempts[rollout.rollout_id] = [attempt]
self._rollouts[rollout.rollout_id] = rollout
# Manully added rollout is not added to task queue. It's already preparing
# Manually added rollout is not added to task queue. It's already preparing
self._completion_events.setdefault(rollout.rollout_id, threading.Event())
return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
@@ -329,7 +384,7 @@ class InMemoryLightningStore(LightningStore):
return rollout
@_healthcheck_wrapper
async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
"""Retrieves the next task from the queue without blocking.
Returns `None` if the queue is empty.
@@ -338,6 +393,11 @@ class InMemoryLightningStore(LightningStore):
See [`LightningStore.dequeue_rollout()`][agentlightning.LightningStore.dequeue_rollout] for semantics.
"""
async with self._lock:
if worker_id is not None:
worker = self._get_or_create_worker(worker_id)
worker.last_dequeue_time = time.time()
worker.status = "idle"
# Keep looking until we find a rollout that's still in queuing status
# or the queue is empty
while self._task_queue:
@@ -348,6 +408,7 @@ class InMemoryLightningStore(LightningStore):
if is_queuing(rollout):
# Update status to preparing
rollout.status = "preparing"
self._running_rollout_ids.add(rollout.rollout_id)
# Create a new attempt (could be first attempt or retry)
attempt_id = _generate_attempt_id()
@@ -369,6 +430,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 +477,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)
@@ -439,16 +506,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 +558,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 +574,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 +594,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
@@ -590,6 +707,7 @@ class InMemoryLightningStore(LightningStore):
if current_attempt == latest_attempt:
if rollout.status == "preparing":
rollout.status = "running"
self._running_rollout_ids.add(rollout.rollout_id)
elif rollout.status in ["queuing", "requeuing"]:
try:
self._task_queue.remove(rollout)
@@ -598,6 +716,7 @@ class InMemoryLightningStore(LightningStore):
f"Trying to remove rollout {rollout.rollout_id} from the queue but it's not in the queue."
)
rollout.status = "running"
self._running_rollout_ids.add(rollout.rollout_id)
return span
@@ -843,6 +962,12 @@ class InMemoryLightningStore(LightningStore):
elif is_queuing(rollout) and rollout not in self._task_queue:
self._task_queue.append(rollout)
# Updating running rollouts cache
if rollout.status in ["preparing", "running"]:
self._running_rollout_ids.add(rollout.rollout_id)
else:
self._running_rollout_ids.discard(rollout.rollout_id)
# If the rollout is no longer in a queueing state, remove it from the queue.
if not isinstance(status, Unset) and not is_queuing(rollout) and rollout in self._task_queue:
try:
@@ -886,19 +1011,26 @@ class InMemoryLightningStore(LightningStore):
if not attempt:
raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
worker_sync_required = False
# Update fields if they are not UNSET
if not isinstance(worker_id, Unset):
attempt.worker_id = worker_id
worker_sync_required = worker_sync_required or bool(worker_id)
if not isinstance(status, Unset):
attempt.status = status
# Also update end_time if the status indicates completion
if status in ["failed", "succeeded"]:
attempt.end_time = time.time()
if not isinstance(worker_id, Unset):
attempt.worker_id = worker_id
worker_sync_required = worker_sync_required or bool(attempt.worker_id)
if not isinstance(last_heartbeat_time, Unset):
attempt.last_heartbeat_time = last_heartbeat_time
if not isinstance(metadata, Unset):
attempt.metadata = metadata
if worker_sync_required and attempt.worker_id:
self._sync_worker_with_attempt(attempt)
# Re-validate the attempt to ensure legality
Attempt.model_validate(attempt.model_dump())
@@ -916,12 +1048,40 @@ class InMemoryLightningStore(LightningStore):
return attempt
@_healthcheck_wrapper
async def query_workers(self) -> List[Worker]:
"""Return the current snapshot of all workers."""
async with self._lock:
return list(self._workers.values())
@_healthcheck_wrapper
async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
async with self._lock:
return self._workers.get(worker_id)
@_healthcheck_wrapper
async def update_worker(
self,
worker_id: str,
heartbeat_stats: Dict[str, Any] | Unset = UNSET,
) -> Worker:
"""Create or update a worker entry."""
async with self._lock:
worker = self._get_or_create_worker(worker_id)
if not isinstance(heartbeat_stats, Unset):
worker.heartbeat_stats = dict(heartbeat_stats)
worker.last_heartbeat_time = time.time()
Worker.model_validate(worker.model_dump())
return worker
async def _healthcheck(self) -> None:
"""Perform healthcheck against all running rollouts in the store."""
async with self._lock:
running_rollouts: List[AttemptedRollout] = []
for rollout in self._rollouts.values():
if rollout.status in ["preparing", "running"]:
for rollout_id in self._running_rollout_ids:
rollout = self._rollouts.get(rollout_id)
if rollout is not None and rollout.status in ["preparing", "running"]:
all_attempts = self._attempts.get(rollout.rollout_id, [])
if not all_attempts:
# The rollout is running but has no attempts, this should not happen
+33 -3
View File
@@ -18,9 +18,10 @@ from agentlightning.types import (
RolloutStatus,
Span,
TaskInput,
Worker,
)
from .base import UNSET, LightningStore, Unset
from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
class LightningStoreThreaded(LightningStore):
@@ -35,6 +36,16 @@ class LightningStoreThreaded(LightningStore):
self.store = store
self._lock = threading.Lock()
@property
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,
@@ -57,9 +68,9 @@ class LightningStoreThreaded(LightningStore):
with self._lock:
return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
with self._lock:
return await self.store.dequeue_rollout()
return await self.store.dequeue_rollout(worker_id=worker_id)
async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
with self._lock:
@@ -171,3 +182,22 @@ class LightningStoreThreaded(LightningStore):
last_heartbeat_time=last_heartbeat_time,
metadata=metadata,
)
async def query_workers(self) -> List[Worker]:
with self._lock:
return await self.store.query_workers()
async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
with self._lock:
return await self.store.get_worker_by_id(worker_id)
async def update_worker(
self,
worker_id: str,
heartbeat_stats: Dict[str, Any] | Unset = UNSET,
) -> Worker:
with self._lock:
return await self.store.update_worker(
worker_id=worker_id,
heartbeat_stats=heartbeat_stats,
)
+71 -223
View File
@@ -2,25 +2,23 @@
from __future__ import annotations
import asyncio
import logging
import os
import threading
from contextlib import asynccontextmanager, contextmanager
from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Iterator, List, Optional
from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator, List, Optional
import agentops
import agentops.sdk.core
import opentelemetry.trace as trace_api
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.sdk.trace import TracerProvider as TracerProviderImpl
from opentelemetry.trace import get_tracer_provider
from opentelemetry.trace.status import StatusCode
from agentlightning.instrumentation import instrument_all, uninstrument_all
from agentlightning.instrumentation.agentops import AgentOpsServerManager
from agentlightning.store.base import LightningStore
from .base import Tracer
from .otel import LightningSpanProcessor, OtelTracer
if TYPE_CHECKING:
from agentops.integration.callbacks.langchain import LangchainCallbackHandler
@@ -29,7 +27,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
class AgentOpsTracer(Tracer):
class AgentOpsTracer(OtelTracer):
"""Traces agent execution using AgentOps.
This tracer provides functionality to capture execution details using the
@@ -56,95 +54,37 @@ 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()
def uninstrument(self, worker_id: int):
uninstrument_all()
def init_worker(self, worker_id: int):
super().init_worker(worker_id)
logger.info(f"[Worker {worker_id}] Setting up tracer...") # worker_id included in process name
def _initialize_tracer_provider(self, worker_id: int):
logger.info(f"[Worker {worker_id}] Setting up AgentOps tracer...") # worker_id included in process name
if self.instrument_managed:
self.instrument(worker_id)
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.")
self._lightning_span_processor = LightningSpanProcessor()
try:
# new versions
instance = agentops.sdk.core.tracer
# TODO: The span processor cannot be deleted once added.
# This might be a problem if the tracer is entered and exited multiple times.
instance.provider.add_span_processor(self._lightning_span_processor) # type: ignore
except AttributeError:
# old versions
instance = TracingCore.get_instance() # type: ignore
instance._provider.add_span_processor(self._lightning_span_processor) # type: ignore
# TODO: The span processor cannot be deleted once added.
# This might be a problem if the tracer is entered and exited multiple times.
self._get_tracer_provider().add_span_processor(self._lightning_span_processor) # type: ignore
def teardown_worker(self, worker_id: int) -> None:
super().teardown_worker(worker_id)
@@ -161,7 +101,7 @@ class AgentOpsTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> AsyncGenerator[LightningSpanProcessor, None]:
) -> AsyncGenerator[trace_api.Tracer, None]:
"""
Starts a new tracing context. This should be used as a context manager.
@@ -172,12 +112,10 @@ class AgentOpsTracer(Tracer):
attempt_id: Optional attempt ID to add the spans to.
Yields:
The [`LightningSpanProcessor`][agentlightning.tracer.agentops.LightningSpanProcessor] instance to collect spans.
The OpenTelemetry tracer instance to collect spans.
"""
with self._trace_context_sync(
name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id
) as processor:
yield processor
with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
yield tracer
@contextmanager
def _trace_context_sync(
@@ -187,31 +125,50 @@ class AgentOpsTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> Iterator[LightningSpanProcessor]:
) -> Iterator[trace_api.Tracer]:
"""Implementation of `trace_context` for synchronous execution."""
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
tracer_provider = self._get_tracer_provider()
kwargs: dict[str, Any] = {}
if name is not None:
kwargs["trace_name"] = name
elif rollout_id is not None:
kwargs["trace_name"] = rollout_id
if store is not None and rollout_id is not None and attempt_id is not None:
if store.capabilities.get("otlp_traces", False) is True:
logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
else:
self._disable_native_otlp_exporter()
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
with ctx as processor:
yield processor
with ctx:
# AgentOps end_trace and start_trace must live inside the lightning span processor context.
# Otherwise some traces might not be recorded.
with self._agentops_trace_context(rollout_id, attempt_id, kwargs):
yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
elif store is None and rollout_id is None and attempt_id is None:
# TODO: Add tests to cover both paths
self._disable_native_otlp_exporter()
with self._lightning_span_processor:
yield self._lightning_span_processor
with self._agentops_trace_context(None, None, kwargs):
yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
def get_last_trace(self) -> List[ReadableSpan]:
"""
Retrieves the raw list of captured spans from the most recent trace.
Returns:
A list of OpenTelemetry `ReadableSpan` objects.
"""
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
return self._lightning_span_processor.spans()
@contextmanager
def _agentops_trace_context(self, rollout_id: Optional[str], attempt_id: Optional[str], kwargs: dict[str, Any]):
trace = agentops.start_trace(**kwargs)
status = StatusCode.OK # type: ignore
try:
yield
except Exception as e:
# TODO: I'm not sure whether this will catch errors in user code.
status = StatusCode.ERROR # type: ignore
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}: {e}")
finally:
agentops.end_trace(trace, end_state=status) # type: ignore
def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
"""
@@ -239,135 +196,26 @@ class AgentOpsTracer(Tracer):
get_langchain_callback_handler = get_langchain_handler # alias
def _get_tracer_provider(self) -> TracerProviderImpl:
try:
# new versions
instance = agentops.sdk.core.tracer
if instance.provider is None:
raise RuntimeError("AgentOps TracerProvider is not initialized.")
class LightningSpanProcessor(SpanProcessor):
"""Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
to a [`LightningStore`][agentlightning.LightningStore].
"""
if get_tracer_provider() is not instance.provider:
logger.error(
"Mismatch between global singleton TracerProvider and AgentOps TracerProvider. "
"AgentOps might not work properly."
)
def __init__(self):
self._spans: List[ReadableSpan] = []
if not isinstance(instance.provider, TracerProviderImpl): # type: ignore
raise RuntimeError("Unsupported TracerProvider type for AgentOps instrumentation.")
# Store related context and states
self._store: Optional[LightningStore] = None
self._rollout_id: Optional[str] = None
self._attempt_id: Optional[str] = None
self._lock = threading.Lock()
# private asyncio loop running in a daemon thread
self._loop_ready = threading.Event()
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
self._loop_thread.start()
self._loop_ready.wait() # loop is ready
def _loop_runner(self):
loop = asyncio.new_event_loop()
self._loop = loop
asyncio.set_event_loop(loop)
self._loop_ready.set()
loop.run_forever()
loop.close()
def __enter__(self):
self._last_trace = None
self._spans = []
return self
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
self._store = None
self._rollout_id = None
self._attempt_id = None
def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
# submit to the dedicated loop and wait synchronously
if self._loop is None:
raise RuntimeError("Loop is not initialized. This should not happen.")
# If already on the exporter loop thread, schedule and return immediately.
# ---------------------------------------------------------------------------
# WHY THIS CONDITIONAL EXISTS:
# In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
# (or another finalizer) while the Python garbage collector is running on the
# *same thread* that owns our exporter event loop ("otel-loop").
#
# When that happens, on_end() executes on the exporter loop thread itself.
# If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
# it would deadlock immediately — because the loop cannot both wait on and run
# the same coroutine. The Future stays pending forever and the loop stops
# processing scheduled callbacks.
#
# To avoid that self-deadlock, we detect when on_end() runs on the exporter
# loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
# instead of blocking with .result().
#
# This situation can occur because Python calls __del__ in whatever thread
# releases the last reference, which can easily be our loop thread if the
# object is dereferenced during loop._run_once().
# ---------------------------------------------------------------------------
if threading.current_thread() is self._loop_thread:
self._loop.call_soon_threadsafe(asyncio.create_task, coro) # type: ignore
return None
fut = asyncio.run_coroutine_threadsafe(coro, self._loop) # type: ignore
return fut.result(timeout=timeout) # raises on error # type: ignore
def shutdown(self) -> None:
if self._loop:
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop_thread.join(timeout=5)
self._loop = None
def force_flush(self, timeout_millis: int = 30000) -> bool:
return True
def spans(self) -> List[ReadableSpan]:
"""
Get the list of spans collected by this processor.
This is useful for debugging and testing purposes.
Returns:
List of ReadableSpan objects collected during tracing.
"""
return self._spans
def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
# simple context manager without nesting into asyncio
class _Ctx:
def __enter__(_): # type: ignore
with self._lock:
self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
self._last_trace = None
self._spans = []
return self
def __exit__(_, exc_type, exc, tb): # type: ignore
with self._lock:
self._store = self._rollout_id = self._attempt_id = None
return _Ctx()
def on_end(self, span: ReadableSpan) -> None:
"""
Process a span when it ends.
Args:
span: The span that has ended.
"""
# Skip if span is not sampled
if not span.context or not span.context.trace_flags.sampled:
return
if self._store and self._rollout_id and self._attempt_id:
try:
# Submit add_otel_span to the event loop and wait for it to complete
with suppress_instrumentation():
self._await_in_loop(
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
timeout=60.0,
)
except Exception:
# log; on_end MUST NOT raise
logger.exception(f"Error adding span to store: {span.name}")
self._spans.append(span)
self._tracer_provider = instance.provider
return self._tracer_provider
except AttributeError:
# old versions
instance = TracingCore.get_instance() # type: ignore
self._tracer_provider = instance._provider # type: ignore
return self._tracer_provider # type: ignore
+25
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
from opentelemetry.sdk.trace import ReadableSpan
@@ -138,3 +139,27 @@ class Tracer(ParallelWorkerBase):
"""
logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
return None
@contextmanager
def lifespan(self):
"""A context manager to manage the lifespan of the tracer.
This can be used to set up and tear down any necessary resources
for the tracer, useful for debugging purposes.
"""
has_init = False
has_init_worker = False
try:
self.init()
has_init = True
self.init_worker(0)
has_init_worker = True
yield
finally:
if has_init_worker:
self.teardown_worker(0)
if has_init:
self.teardown()
+255 -11
View File
@@ -2,16 +2,25 @@
from __future__ import annotations
import asyncio
import logging
import threading
from contextlib import asynccontextmanager
from typing import AsyncGenerator, List, Optional
from typing import Any, AsyncGenerator, Awaitable, List, Optional
import opentelemetry.trace as trace_api
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from agentops.sdk.core import BatchSpanProcessor
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from agentlightning.store.base import LightningStore
from agentlightning.types.tracer import SpanNames
from agentlightning.utils.otlp import LightningStoreOTLPExporter
from .agentops import LightningSpanProcessor # FIXME: This import should be from otel to agentops
from .base import Tracer
logger = logging.getLogger(__name__)
@@ -29,21 +38,31 @@ class OtelTracer(Tracer):
# This provider is only initialized when the worker is initialized.
self._tracer_provider: Optional[TracerProvider] = None
self._lightning_span_processor: Optional[LightningSpanProcessor] = None
self._simple_span_processor: Optional[SimpleSpanProcessor] = None
self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
self._initialized: bool = False
def init_worker(self, worker_id: int):
super().init_worker(worker_id)
self._initialize_tracer_provider(worker_id)
def _initialize_tracer_provider(self, worker_id: int):
logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")
if self._initialized:
logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
tracer_provider = TracerProvider()
trace_api.set_tracer_provider(tracer_provider)
self._tracer_provider = TracerProvider()
trace_api.set_tracer_provider(self._tracer_provider)
self._lightning_span_processor = LightningSpanProcessor()
tracer_provider.add_span_processor(self._lightning_span_processor)
self._tracer_provider.add_span_processor(self._lightning_span_processor)
self._otlp_span_exporter = LightningStoreOTLPExporter()
self._simple_span_processor = SimpleSpanProcessor(self._otlp_span_exporter)
self._tracer_provider.add_span_processor(self._simple_span_processor)
self._initialized = True
logger.info(f"[Worker {worker_id}] OpenTelemetry tracer provider initialized.")
def teardown_worker(self, worker_id: int):
super().teardown_worker(worker_id)
logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
@@ -57,7 +76,7 @@ class OtelTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> AsyncGenerator[LightningSpanProcessor, None]:
) -> AsyncGenerator[trace_api.Tracer, None]:
"""
Starts a new tracing context. This should be used as a context manager.
@@ -68,18 +87,24 @@ class OtelTracer(Tracer):
attempt_id: Optional attempt ID to add the spans to.
Yields:
The LightningSpanProcessor instance to collect spans.
The OpenTelemetry tracer instance to collect spans.
"""
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:
if store.capabilities.get("otlp_traces", False) is True:
logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
else:
self._disable_native_otlp_exporter()
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
with ctx as processor:
yield processor
with ctx:
yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
elif store is None and rollout_id is None and attempt_id is None:
self._disable_native_otlp_exporter()
with self._lightning_span_processor:
yield self._lightning_span_processor
yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
@@ -93,3 +118,222 @@ class OtelTracer(Tracer):
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
return self._lightning_span_processor.spans()
def _get_tracer_provider(self) -> TracerProviderImpl:
if self._tracer_provider is None:
raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
return self._tracer_provider
def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
tracer_provider = self._get_tracer_provider()
active_span_processor = tracer_provider._active_span_processor # pyright: ignore[reportPrivateUsage]
# Override the resources so that the server knows where the request comes from.
tracer_provider._resource = tracer_provider._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: rollout_id,
SpanNames.ATTEMPT_ID: attempt_id,
}
)
)
instrumented = False
candidates: List[str] = []
for processor in active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
if isinstance(processor, LightningSpanProcessor):
# We don't need the LightningSpanProcessor any more.
logger.debug("LightningSpanProcessor already present in TracerProvider, disabling it.")
processor.disable_store_submission = True
elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
# Instead, we rely on the OTLPSpanExporter to send spans to the store.
if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
instrumented = True
else:
candidates.append(
f"{processor.__class__.__name__} with {processor.span_exporter.__class__.__name__}"
)
else:
candidates.append(f"{processor.__class__.__name__}")
if not instrumented:
raise RuntimeError(
"Failed to enable native OTLP exporter: no BatchSpanProcessor or SimpleSpanProcessor with "
"LightningStoreOTLPExporter found in TracerProvider. Please try using a non-OTLP store."
"Candidates are: " + ", ".join(candidates)
)
def _disable_native_otlp_exporter(self):
tracer_provider = self._get_tracer_provider()
active_span_processor = tracer_provider._active_span_processor # pyright: ignore[reportPrivateUsage]
tracer_provider._resource = tracer_provider._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: "",
SpanNames.ATTEMPT_ID: "",
}
)
) # reset resource
for processor in active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
if isinstance(processor, LightningSpanProcessor):
# We will be in need of the LightningSpanProcessor again.
logger.debug("Enabling LightningSpanProcessor in TracerProvider.")
processor.disable_store_submission = False
class LightningSpanProcessor(SpanProcessor):
"""Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
to a [`LightningStore`][agentlightning.LightningStore].
It serves two purposes:
1. Records all the spans in a local buffer.
2. Submits the spans to the event loop to be added to the store.
"""
def __init__(self, disable_store_submission: bool = False):
self._disable_store_submission: bool = disable_store_submission
self._spans: List[ReadableSpan] = []
# Store related context and states
self._store: Optional[LightningStore] = None
self._rollout_id: Optional[str] = None
self._attempt_id: Optional[str] = None
self._lock = threading.Lock()
# private asyncio loop running in a daemon thread
self._loop_ready = threading.Event()
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread: Optional[threading.Thread] = None
@property
def disable_store_submission(self) -> bool:
"""Whether to disable submitting spans to the store."""
return self._disable_store_submission
@disable_store_submission.setter
def disable_store_submission(self, value: bool) -> None:
self._disable_store_submission = value
def _ensure_loop(self) -> None:
if self._loop_thread is None or self._loop is None:
self._loop_ready.clear()
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
self._loop_thread.start()
self._loop_ready.wait() # loop is ready
def _loop_runner(self):
loop = asyncio.new_event_loop()
self._loop = loop
asyncio.set_event_loop(loop)
self._loop_ready.set()
loop.run_forever()
loop.close()
def __enter__(self):
self._last_trace = None
self._spans = []
return self
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
self._store = None
self._rollout_id = None
self._attempt_id = None
def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
# submit to the dedicated loop and wait synchronously
self._ensure_loop()
if self._loop is None:
raise RuntimeError("Loop is not initialized. This should not happen.")
# If already on the exporter loop thread, schedule and return immediately.
# ---------------------------------------------------------------------------
# WHY THIS CONDITIONAL EXISTS:
# In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
# (or another finalizer) while the Python garbage collector is running on the
# *same thread* that owns our exporter event loop ("otel-loop").
#
# When that happens, on_end() executes on the exporter loop thread itself.
# If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
# it would deadlock immediately — because the loop cannot both wait on and run
# the same coroutine. The Future stays pending forever and the loop stops
# processing scheduled callbacks.
#
# To avoid that self-deadlock, we detect when on_end() runs on the exporter
# loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
# instead of blocking with .result().
#
# This situation can occur because Python calls __del__ in whatever thread
# releases the last reference, which can easily be our loop thread if the
# object is dereferenced during loop._run_once().
# ---------------------------------------------------------------------------
if threading.current_thread() is self._loop_thread:
self._loop.call_soon_threadsafe(asyncio.create_task, coro) # type: ignore
return None
fut = asyncio.run_coroutine_threadsafe(coro, self._loop) # type: ignore
return fut.result(timeout=timeout) # raises on error # type: ignore
def shutdown(self) -> None:
if self._loop:
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop = None
if self._loop_thread:
self._loop_thread.join(timeout=5)
def force_flush(self, timeout_millis: int = 30000) -> bool:
return True
def spans(self) -> List[ReadableSpan]:
"""
Get the list of spans collected by this processor.
This is useful for debugging and testing purposes.
Returns:
List of ReadableSpan objects collected during tracing.
"""
return self._spans
def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
# simple context manager without nesting into asyncio
class _Ctx:
def __enter__(_): # type: ignore
# Use _ instead of self to avoid shadowing the instance method.
with self._lock:
self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
self._last_trace = None
self._spans = []
return self
def __exit__(_, exc_type, exc, tb): # type: ignore
with self._lock:
self._store = self._rollout_id = self._attempt_id = None
return _Ctx()
def on_end(self, span: ReadableSpan) -> None:
"""
Process a span when it ends.
Args:
span: The span that has ended.
"""
# Skip if span is not sampled
if not span.context or not span.context.trace_flags.sampled:
return
if not self._disable_store_submission and self._store and self._rollout_id and self._attempt_id:
try:
# Submit add_otel_span to the event loop and wait for it to complete
with suppress_instrumentation():
self._ensure_loop()
self._await_in_loop(
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
timeout=60.0,
)
except Exception:
# log; on_end MUST NOT raise
logger.exception(f"Error adding span to store: {span.name}")
self._spans.append(span)
+28
View File
@@ -49,6 +49,8 @@ __all__ = [
"Attempt",
"AttemptedRollout",
"Hook",
"Worker",
"WorkerStatus",
]
T_co = TypeVar("T_co", covariant=True)
@@ -200,6 +202,32 @@ class AttemptedRollout(Rollout):
return self
WorkerStatus = Literal["idle", "busy", "unknown"]
class Worker(BaseModel):
"""Worker information. This is actually the same as Runner info."""
worker_id: str
"""The ID of the worker."""
status: WorkerStatus = "unknown"
"""The status of the worker."""
heartbeat_stats: Optional[Dict[str, Any]] = None
"""Statistics about the worker's heartbeat."""
last_heartbeat_time: Optional[float] = None
"""The last time when the worker has reported the stats."""
last_dequeue_time: Optional[float] = None
"""The last time when the worker has tried to dequeue a rollout."""
last_busy_time: Optional[float] = None
"""The last time when the worker has started an attempt and became busy."""
last_idle_time: Optional[float] = None
"""The last time when the worker has triggered the end of an attempt and became idle."""
current_rollout_id: Optional[str] = None
"""The ID of the current rollout that the worker is processing."""
current_attempt_id: Optional[str] = None
"""The ID of the current attempt that the worker is processing."""
TaskInput = Any
"""Task input type. Accepts arbitrary payloads."""
+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."""
+6
View File
@@ -411,6 +411,12 @@ class SpanNames(str, Enum):
"""The name of the exception span."""
VIRTUAL = "agentlightning.virtual"
"""The name of the virtual span. It represents derived spans without concrete operations."""
ROLLOUT_ID = "agentlightning.rollout_id"
"""The name of the rollout ID."""
ATTEMPT_ID = "agentlightning.attempt_id"
"""The name of the attempt ID."""
SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
"""The name of the span sequence ID."""
class SpanAttributeNames(str, Enum):
+1
View File
@@ -0,0 +1 @@
# Copyright (c) Microsoft. All rights reserved.
+428
View File
@@ -0,0 +1,428 @@
# Copyright (c) Microsoft. All rights reserved.
import gzip
import logging
from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
from fastapi import Request, Response
from google.protobuf import json_format
from google.rpc.status_pb2 import Status
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import (
ExportLogsServiceRequest,
ExportLogsServiceResponse,
)
from opentelemetry.proto.collector.metrics.v1.metrics_service_pb2 import (
ExportMetricsServiceRequest,
ExportMetricsServiceResponse,
)
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceRequest,
ExportTraceServiceResponse,
)
from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExportResult
from opentelemetry.util.types import AttributeValue
from agentlightning.store.base import LightningStore
from agentlightning.types.tracer import (
Attributes,
Event,
Link,
OtelResource,
Span,
SpanContext,
SpanNames,
TraceStatus,
convert_timestamp,
)
PROTOBUF_CT = "application/x-protobuf"
logger = logging.getLogger(__name__)
T_request = TypeVar("T_request", ExportLogsServiceRequest, ExportMetricsServiceRequest, ExportTraceServiceRequest)
T_response = TypeVar("T_response", ExportLogsServiceResponse, ExportMetricsServiceResponse, ExportTraceServiceResponse)
async def handle_otlp_export(
request: Request,
request_message_cls: Type[T_request],
response_message_cls: Type[T_response],
message_callback: Optional[Callable[[T_request], Awaitable[None]]],
signal_name: str,
) -> Response:
"""
Generic handler for /v1/traces, /v1/metrics, /v1/logs.
Convert the OTLP Protobuf request to a JSON-like object.
"""
content_type = request.headers.get("Content-Type", "").split(";")[0].strip()
if content_type != PROTOBUF_CT:
# For brevity we only support binary protobuf here.
return _bad_request_response(
request,
f"Unsupported Content-Type '{content_type}', expected '{PROTOBUF_CT}'",
content_type=PROTOBUF_CT,
)
raw_body = await request.body()
body = _read_body_maybe_gzip(request, raw_body)
# Empty request is allowed and should still succeed.
if not body:
req_msg = request_message_cls()
else:
req_msg = request_message_cls()
try:
req_msg.ParseFromString(body)
except Exception as exc:
return _bad_request_response(request, f"Unable to parse OTLP {signal_name} payload: {exc}")
if message_callback is not None:
await message_callback(req_msg)
# Build success response. Partial success field is left unset.
resp_msg = response_message_cls()
# Encode response in the same Content-Type as request.
if content_type == PROTOBUF_CT:
resp_bytes = resp_msg.SerializeToString()
else:
resp_bytes = json_format.MessageToJson(resp_msg).encode("utf-8")
resp_bytes, headers = _maybe_gzip_response(request, resp_bytes)
return Response(
content=resp_bytes,
media_type=content_type,
status_code=200,
headers=headers,
)
async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningStore) -> List[Span]:
"""Parse an OTLP proto payload into List[Span].
A store is needed here for generating a sequence ID for each span.
"""
output_spans: List[Span] = []
for resource_spans in request.resource_spans:
# Resource-level attributes & IDs
resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
# rollout_id, attempt_id from resource attributes when present.
rollout_id_resource = resource_attrs.get(SpanNames.ROLLOUT_ID)
attempt_id_resource = resource_attrs.get(SpanNames.ATTEMPT_ID)
# If sequence id is provided, all the spans will share the same sequence ID.
# unless otherwise overridden by span-level attributes.
sequence_id_resource = resource_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
# Each ScopeSpans contains multiple spans
for scope_spans in resource_spans.scope_spans:
for proto_span in scope_spans.spans:
trace_id_hex = _bytes_to_trace_id_hex(proto_span.trace_id)
span_id_hex = _bytes_to_span_id_hex(proto_span.span_id)
parent_id_hex = _bytes_to_span_id_hex(proto_span.parent_span_id) if proto_span.parent_span_id else None
# Status
status_code_str = _STATUS_CODE_MAP.get(proto_span.status.code, "UNSET")
status = TraceStatus(
status_code=status_code_str,
description=proto_span.status.message or None,
)
# Attributes
span_attrs = _kv_list_to_dict(proto_span.attributes)
# Context
context = SpanContext(
trace_id=trace_id_hex,
span_id=span_id_hex,
is_remote=False,
trace_state={},
)
# Try to get if span attributes contain something like rollout_id or attempt_id
# Override the resource-level attributes with the span-level attributes if present.
rollout_id_span = span_attrs.get(SpanNames.ROLLOUT_ID)
attempt_id_span = span_attrs.get(SpanNames.ATTEMPT_ID)
sequence_id_span = span_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
# Normalize to regular strings and ints
rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
attempt_id_raw = attempt_id_span if attempt_id_span is not None else attempt_id_resource
sequence_id_raw = sequence_id_span if sequence_id_span is not None else sequence_id_resource
rollout_id, attempt_id = _normalize_rollout_attempt_id(rollout_id_raw, attempt_id_raw)
sequence_id = _normalize_sequence_id(sequence_id_raw)
if rollout_id is None or attempt_id is None:
logger.warning(
"Both rollout_id and attempt_id must be present in resource attributes. "
"Spans will not be able to log to the store because of missing IDs: rollout_id=%s, attempt_id=%s, sequence_id=%s",
rollout_id,
attempt_id,
sequence_id,
)
continue
# Generate a new sequence ID if not provided
if sequence_id is None:
current_sequence_id = await store.get_next_span_sequence_id(
rollout_id=rollout_id, attempt_id=attempt_id
)
else:
current_sequence_id = sequence_id
# Build Span
span = Span(
rollout_id=rollout_id,
attempt_id=attempt_id,
sequence_id=current_sequence_id,
trace_id=trace_id_hex,
span_id=span_id_hex,
parent_id=parent_id_hex,
name=proto_span.name,
status=status,
attributes=span_attrs,
events=_events_from_proto(proto_span),
links=_links_from_proto(proto_span),
start_time=convert_timestamp(proto_span.start_time_unix_nano),
end_time=convert_timestamp(proto_span.end_time_unix_nano),
context=context,
parent=None, # OTLP only has parent_span_id; we don't have full SpanContext
resource=otel_resource,
)
output_spans.append(span)
return output_spans
class LightningStoreOTLPExporter(OTLPSpanExporter):
"""OTLP Exporter that write to a LightningStore-compatible backend.
The backend requires two special attributes on each span:
- `agentlightning.rollout_id`: The rollout ID to associate the span with.
- `agentlightning.attempt_id`: The attempt ID to associate the span with.
It can optionally use the following attribute to sequence spans:
- `agentlightning.span_sequence_id`: A decimal string representing the sequence ID of the span.
"""
_default_endpoint: Optional[str] = None
_rollout_id: Optional[str] = None
_attempt_id: Optional[str] = None
def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
"""Enable storing OTLP data to a specific LightningStore rollout/attempt."""
self._rollout_id = rollout_id
self._attempt_id = attempt_id
self._default_endpoint = self._endpoint
self._endpoint = endpoint
def disable_store_otlp(self) -> None:
"""Disable storing OTLP data to LightningStore."""
self._rollout_id = None
self._attempt_id = None
if self._default_endpoint is not None:
self._endpoint = self._default_endpoint
def should_bypass(self) -> bool:
"""Check if the exporter should bypass the default export if rollout_id and attempt_id are not set."""
return True
def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
if self._rollout_id is not None and self._attempt_id is not None:
# rollout_id and attempt_id are present in resource attributes
# It means that the server supports OTLP endpoint.
for span in spans:
# Override the resources so that the server knows where the request comes from.
span._resource = span._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: self._rollout_id,
SpanNames.ATTEMPT_ID: self._attempt_id,
}
)
)
return super().export(spans)
elif not self.should_bypass():
logger.debug("Rollout ID and Attempt ID not set; using default OTLP exporter behavior.")
return super().export(spans)
else:
logger.debug("Rollout ID and Attempt ID not set; bypassing export.")
return SpanExportResult.SUCCESS
def _read_body_maybe_gzip(request: Request, raw_body: bytes) -> bytes:
"""
Decompress body if Content-Encoding: gzip; otherwise return as is.
"""
encoding = request.headers.get("Content-Encoding", "").lower()
if encoding == "gzip":
return gzip.decompress(raw_body)
return raw_body
def _maybe_gzip_response(request: Request, payload: bytes) -> Tuple[bytes, Dict[str, str]]:
"""
If Accept-Encoding includes gzip, gzip the payload and set Content-Encoding header.
"""
ae = request.headers.get("Accept-Encoding", "")
tokens = [token.split(";")[0].strip().lower() for token in ae.split(",") if token.strip()]
headers: Dict[str, str] = {}
if "gzip" in tokens:
payload = gzip.compress(payload)
headers["Content-Encoding"] = "gzip"
return payload, headers
def _bad_request_response(request: Request, message: str, content_type: str = PROTOBUF_CT) -> Response:
"""
Build a 400 response whose body is a protobuf Status message, encoded
in the same Content-Type as the request (OTLP/HTTP requirement).
"""
status_msg = Status(message=message)
if content_type == PROTOBUF_CT:
body = status_msg.SerializeToString()
else:
# Fallback: JSON representation of Status.
body = json_format.MessageToJson(status_msg).encode("utf-8")
body, headers = _maybe_gzip_response(request, body)
return Response(
content=body,
status_code=400,
media_type=content_type,
headers=headers,
)
def _normalize_rollout_attempt_id(
rollout_id: Optional[AttributeValue], attempt_id: Optional[AttributeValue]
) -> Tuple[Optional[str], Optional[str]]:
"""Normalize a rollout or attempt ID to a string."""
rollout_id_str = str(rollout_id) if rollout_id is not None else None
attempt_id_str = str(attempt_id) if attempt_id is not None else None
return rollout_id_str, attempt_id_str
def _normalize_sequence_id(sequence_id: Optional[AttributeValue]) -> Optional[int]:
"""Normalize a sequence ID to an integer."""
if sequence_id is None:
return None
try:
sequence_id_int = int(str(sequence_id))
except (ValueError, TypeError):
logger.warning(
"Invalid sequence_id value in resource attributes: %r. Must be an integer or string representing an integer. Assuming None.",
sequence_id,
)
sequence_id_int = None
return sequence_id_int
def _any_value_to_python(value: AnyValue) -> Any:
"""Convert OTLP AnyValue -> plain Python value."""
kind = value.WhichOneof("value")
if kind is None:
return None
if kind == "string_value":
return value.string_value
if kind == "bool_value":
return value.bool_value
if kind == "int_value":
return int(value.int_value)
if kind == "double_value":
return float(value.double_value)
if kind == "array_value":
return [_any_value_to_python(v) for v in value.array_value.values]
if kind == "kvlist_value":
# Map<string, AnyValue> -> dict
return {kv.key: _any_value_to_python(kv.value) for kv in value.kvlist_value.values}
if kind == "bytes_value":
# Serialize bytes as hex string to stay JSON-friendly
return value.bytes_value.hex()
return None
def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
"""Convert repeated KeyValue -> Attributes dict."""
return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
_STATUS_CODE_MAP = {
ProtoStatus.STATUS_CODE_UNSET: "UNSET",
ProtoStatus.STATUS_CODE_OK: "OK",
ProtoStatus.STATUS_CODE_ERROR: "ERROR",
}
def _bytes_to_trace_id_hex(b: bytes) -> str:
# OTLP uses 16-byte trace IDs; format as 32-char hex
if not b:
return "0" * 32
return b.hex().rjust(32, "0")
def _bytes_to_span_id_hex(b: bytes) -> str:
# OTLP uses 8-byte span IDs; format as 16-char hex
if not b:
return "0" * 16
return b.hex().rjust(16, "0")
def _events_from_proto(span: ProtoSpan) -> List[Event]:
"""Event converter from OTLP ProtoSpan to List[Event]."""
return [
Event(
name=e.name,
attributes=_kv_list_to_dict(e.attributes),
timestamp=convert_timestamp(e.time_unix_nano),
)
for e in span.events
]
def _links_from_proto(span: ProtoSpan) -> List[Link]:
"""Link converter from OTLP ProtoSpan to List[Link]."""
links: List[Link] = []
for link in span.links:
trace_id_hex = _bytes_to_trace_id_hex(link.trace_id)
span_id_hex = _bytes_to_span_id_hex(link.span_id)
ctx = SpanContext(
trace_id=trace_id_hex,
span_id=span_id_hex,
is_remote=False,
trace_state={}, # OTLP trace_state is currently a string; you can parse if needed
)
links.append(
Link(
context=ctx,
attributes=_kv_list_to_dict(link.attributes) or None,
)
)
return links
def _resource_from_proto(resource: ProtoResource, schema_url: str = "") -> OtelResource:
return OtelResource(
attributes=_kv_list_to_dict(resource.attributes),
schema_url=schema_url or "",
)
File diff suppressed because it is too large Load Diff
+72
View File
@@ -0,0 +1,72 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import platform
import socket
from contextlib import suppress
from datetime import datetime
from typing import Any, Dict, List, cast
import psutil
from gpustat import GPUStat, GPUStatCollection
def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
# CPU
cpu = {
"cpu_name": platform.processor(),
"cpu_cores": psutil.cpu_count(logical=False),
"cpu_threads": psutil.cpu_count(logical=True),
"cpu_usage_pct": psutil.cpu_percent(0.05),
}
# Memory
vm = psutil.virtual_memory()
mem = {
"mem_used_gb": round(vm.used / (2**30), 2),
"mem_total_gb": round(vm.total / (2**30), 2),
"mem_pct": vm.percent,
}
# Disk
du = psutil.disk_usage("/")
disk = {
"disk_used_gb": round(du.used / (2**30), 2),
"disk_total_gb": round(du.total / (2**30), 2),
"disk_pct": du.percent,
}
# GPU
gpus: List[Dict[str, Any]] = []
with suppress(Exception):
for g in GPUStatCollection.new_query().gpus: # type: ignore
g = cast(GPUStat, g)
gpus.append(
{
"gpu": g.name, # type: ignore
"util_pct": g.utilization,
"mem_used_mb": g.memory_used,
"mem_total_mb": g.memory_total,
"temp_c": g.temperature,
}
)
# Network
net = psutil.net_io_counters()
netinfo = {
"bytes_sent_mb": round(net.bytes_sent / (2**20), 2),
"bytes_recv_mb": round(net.bytes_recv / (2**20), 2),
}
# OS / meta
return {
"timestamp": datetime.now().isoformat(timespec="seconds"),
"host": socket.gethostname(),
"os": platform.platform(),
**cpu,
**mem,
**disk,
**netinfo,
**({"gpus": gpus} if include_gpu else {}),
}
+5 -10
View File
@@ -18,13 +18,13 @@ from flask import Flask, Response, abort, request
from tensordict import TensorDict
from verl import DataProto
from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, configure_logger
from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, setup_logging
from agentlightning.adapter.triplet import TracerTraceToTriplet, TraceToTripletBase
from agentlightning.llm_proxy import LLMProxy, ModelConfig
from agentlightning.store.base import LightningStore
from agentlightning.types import Rollout, RolloutConfig, Task
configure_logger()
setup_logging()
__all__ = [
"AgentModeDaemon",
@@ -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()
+118 -6
View File
@@ -12,6 +12,7 @@ from typing import Dict, Tuple
import numpy as np
import torch
import verl
from codetiming import Timer
from omegaconf import OmegaConf
from tqdm import tqdm
@@ -19,7 +20,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 +54,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 +318,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 +380,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()
@@ -298,14 +404,20 @@ class AgentLightningTrainer(RayPPOTrainer):
assert self.async_rollout_mode, "If agent mode is enabled, async server must be enabled"
if self.adapter is not None and not isinstance(self.adapter, TraceToTripletBase):
raise ValueError("Adapter must be a TraceToTripletBase for currently VERL implementation.")
verl_version = verl.__version__
if verl_version == "0.5.0":
# Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
# However, it is possible that verl updates the naming and causes incompatibility.
# Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
model = "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:])
else:
# For other versions (e.g., 0.6.0), we use the full path to the model.
model = self.config.actor_rollout_ref.model.path
self.agent_mode_daemon = AgentModeDaemon(
self.config.agentlightning.port,
self.config.actor_rollout_ref.rollout.n,
train_information={
# Note (Zhiyuan): To avoid further patch into vllm async server, using the same sentence to get the naming here.
# However, it is possible that verl updates the naming and causes incompatibility.
# Reference: https://github.com/volcengine/verl/blob/5b5e09d9cc20625e436d01f69d9cc739ff681c54/verl/workers/rollout/vllm_rollout/vllm_async_server.py#L217
"model": "/".join(self.config.actor_rollout_ref.model.path.split("/")[-2:]),
"model": model,
"temperature": self.config.actor_rollout_ref.rollout.temperature,
},
tokenizer=self.tokenizer,
+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>
);
}
+46
View File
@@ -0,0 +1,46 @@
// 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';
import { WorkersPage } from './pages/Workers.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: 'runners',
element: <WorkersPage />,
},
{
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,583 @@
// 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, Worker } 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 WORKER_STATUS_COLORS: Record<Worker['status'], string> = {
busy: 'orange',
idle: 'teal',
unknown: 'gray',
};
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>
);
}
type WorkerDrawerTitleProps = {
worker: Worker;
};
function WorkerDrawerTitle({ worker }: WorkerDrawerTitleProps) {
const badgeColor = WORKER_STATUS_COLORS[worker.status] ?? 'gray';
return (
<Stack gap={3}>
<Group gap={6} align='center'>
<Text fw={600}>{worker.workerId}</Text>
<CopyButton value={worker.workerId}>
{({ copied, copy }) => (
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
<ActionIcon
aria-label={`Copy worker ID ${worker.workerId}`}
variant='subtle'
color={copied ? 'teal' : 'gray'}
size='sm'
onClick={(event) => {
event.stopPropagation();
copy();
}}
>
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
</ActionIcon>
</Tooltip>
)}
</CopyButton>
<Badge size='sm' variant='light' color={badgeColor}>
{formatStatusLabel(worker.status)}
</Badge>
</Group>
<Group gap='xl'>
<Group gap={4}>
<Text size='sm' c='dimmed' fw={500}>
Rollout
</Text>
<Text size='sm' c='dimmed'>
{worker.currentRolloutId ?? '—'}
</Text>
</Group>
<Group gap={4}>
<Text size='sm' c='dimmed' fw={500}>
Attempt
</Text>
<Text size='sm' c='dimmed'>
{worker.currentAttemptId ?? '—'}
</Text>
</Group>
</Group>
</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 === 'worker-detail') {
const { worker } = content;
const title = <WorkerDrawerTitle worker={worker} />;
const body = <JsonEditor value={worker} />;
return { title, body };
}
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,469 @@
// 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 },
sequenceId: { fixedWidth: 6, priority: 1 },
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: 'sequenceId',
title: 'Seq.',
sortable: true,
render: ({ sequenceId }) => <Text size='sm'>{sequenceId}</Text>,
},
{
accessor: 'traceId',
title: 'Trace ID',
sortable: true,
render: ({ traceId }) => (
<Group gap={2}>
<Text size='sm'>{traceId}</Text>
<CopyButton value={traceId}>
{({ copied, copy }) => (
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
<ActionIcon
aria-label={`Copy trace ID ${traceId}`}
variant='subtle'
color={copied ? 'teal' : 'gray'}
size='sm'
onClick={(event) => {
event.stopPropagation();
copy();
}}
>
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
</ActionIcon>
</Tooltip>
)}
</CopyButton>
</Group>
),
},
{
accessor: 'spanId',
title: 'Span ID',
sortable: true,
render: ({ spanId }) => (
<Group gap={2}>
<Text size='sm'>{spanId}</Text>
<CopyButton value={spanId}>
{({ copied, copy }) => (
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
<ActionIcon
aria-label={`Copy span ID ${spanId}`}
variant='subtle'
color={copied ? 'teal' : 'gray'}
size='sm'
onClick={(event) => {
event.stopPropagation();
copy();
}}
>
{copied ? <IconCheck size={14} /> : <IconCopy size={14} />}
</ActionIcon>
</Tooltip>
)}
</CopyButton>
</Group>
),
},
{
accessor: 'parentId',
title: 'Parent ID',
sortable: true,
render: ({ parentId }) => {
if (!parentId) {
return (
<Text size='sm' c='dimmed'>
</Text>
);
}
const parentExists = spanIds.has(parentId);
const isInteractive = parentExists && typeof onParentIdClick === 'function';
return (
<Group gap={2}>
<Text
size='sm'
c={parentExists ? undefined : 'red'}
style={{ cursor: isInteractive ? 'pointer' : undefined }}
onClick={(event) => {
if (isInteractive) {
event.stopPropagation();
onParentIdClick?.(parentId);
}
}}
>
{parentId.slice(0, 8)}
</Text>
{!parentExists && (
<Tooltip label='Parent span not found in table' withArrow>
<IconAlertCircle size={14} color='red' />
</Tooltip>
)}
</Group>
);
},
},
{
accessor: 'statusCode',
title: 'Status',
sortable: true,
render: ({ statusCode }) => (
<Badge size='sm' variant='light' color={STATUS_COLORS[statusCode] ?? 'gray'}>
{statusCode}
</Badge>
),
},
{
accessor: 'attributeKeys',
title: 'Attribute Keys',
render: ({ attributeKeys }) =>
attributeKeys ? (
<Text size='sm' lineClamp={1}>
{/* TODO: dim "." and "," and other characters are just normal text */}
{attributeKeys}
</Text>
) : (
<Text size='sm' c='dimmed'>
</Text>
),
},
{
accessor: 'startTime',
title: 'Start Time',
sortable: true,
textAlign: 'left',
render: ({ startTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(startTime))}</Text>,
},
{
accessor: 'endTime',
title: 'End Time',
sortable: true,
textAlign: 'left',
render: ({ endTime }) => <Text size='sm'>{formatDateTimeWithMilliseconds(toTimestamp(endTime))}</Text>,
},
{
accessor: 'duration',
title: 'Duration',
sortable: true,
textAlign: 'left',
render: ({ duration }) => <Text size='sm'>{formatDuration(duration)}</Text>,
},
{
accessor: 'actionsPlaceholder',
title: 'Actions',
render: (record) => (
<Group gap={2}>
<Tooltip label='Show rollout' withArrow disabled={!onShowRollout}>
<ActionIcon
aria-label='Show rollout'
variant='subtle'
color='gray'
onClick={(event) => {
event.stopPropagation();
onShowRollout?.(record);
}}
>
<IconRouteSquare size={16} />
</ActionIcon>
</Tooltip>
<Tooltip label='Show span detail' withArrow disabled={!onShowSpanDetail}>
<ActionIcon
aria-label='Show span detail'
variant='subtle'
color='gray'
onClick={(event) => {
event.stopPropagation();
onShowSpanDetail?.(record);
}}
>
<IconFileDescription size={16} />
</ActionIcon>
</Tooltip>
</Group>
),
},
];
}
export type TracesTableProps = {
spans: Span[] | undefined;
totalRecords: number;
isFetching: boolean;
isError: boolean;
error: unknown;
selectionMessage?: string;
searchTerm: string;
sort: { column: string; direction: 'asc' | 'desc' };
page: number;
recordsPerPage: number;
onSortStatusChange: (status: DataTableSortStatus<TracesTableRecord>) => void;
onPageChange: (page: number) => void;
onRecordsPerPageChange: (value: number) => void;
onResetFilters: () => void;
onRefetch: () => void;
onShowRollout?: (record: TracesTableRecord) => void;
onShowSpanDetail?: (record: TracesTableRecord) => void;
onParentIdClick?: (parentId: string) => void;
recordsPerPageOptions?: number[];
};
export function TracesTable({
spans,
totalRecords,
isFetching,
isError,
error,
selectionMessage,
searchTerm,
sort,
page,
recordsPerPage,
onSortStatusChange,
onPageChange,
onRecordsPerPageChange,
onResetFilters,
onRefetch,
onShowRollout,
onShowSpanDetail,
onParentIdClick,
recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
}: TracesTableProps) {
const { ref: tableContainerRef, width: containerWidth } = useElementSize();
const { width: viewportWidth } = useViewportSize();
const traceRecords = useMemo<TracesTableRecord[]>(() => {
if (!spans) {
return [];
}
return spans.map((span) => buildTraceRecord(span));
}, [spans]);
const spanIds = useMemo(() => {
return new Set(traceRecords.map((record) => record.spanId));
}, [traceRecords]);
const columns = useMemo(
() =>
createTracesColumns({
onShowRollout,
onShowSpanDetail,
onParentIdClick,
spanIds,
}),
[onShowRollout, onShowSpanDetail, onParentIdClick, spanIds],
);
const layoutAwareContainerWidth = useMemo(
() => getLayoutAwareWidth(containerWidth, viewportWidth),
[containerWidth, viewportWidth],
);
const responsiveColumns = useMemo(
() => createResponsiveColumns(columns, layoutAwareContainerWidth, COLUMN_VISIBILITY),
[columns, layoutAwareContainerWidth],
);
const totalPages = useMemo(
() => Math.max(1, Math.ceil(Math.max(0, totalRecords) / Math.max(1, recordsPerPage))),
[recordsPerPage, totalRecords],
);
useEffect(() => {
if (page > totalPages) {
onPageChange(totalPages);
}
}, [onPageChange, page, totalPages]);
const hasActiveFilters = searchTerm.trim().length > 0;
const sortStatus: DataTableSortStatus<TracesTableRecord> = {
columnAccessor: sort.column,
direction: sort.direction,
};
const handleSortStatusChange = useCallback(
(status: DataTableSortStatus<TracesTableRecord>) => {
onSortStatusChange(status);
},
[onSortStatusChange],
);
const errorDescriptor = isError ? getErrorDescriptor(error) : null;
const errorMessage = isError
? `Traces are temporarily unavailable${errorDescriptor ? ` (${errorDescriptor})` : ''}.`
: 'Traces are temporarily unavailable.';
const selectionEmptyState = selectionMessage ? (
<Stack gap='sm' align='center' py='xl'>
<Text fw={600} size='sm'>
{selectionMessage}
</Text>
<Text size='sm' c='dimmed' ta='center'>
Choose a rollout and attempt from the controls above to load trace results.
</Text>
</Stack>
) : null;
const fallbackEmptyState = (
<Stack gap='sm' align='center' py='lg'>
{isError ? (
<>
<Text fw={600} size='sm'>
{errorMessage}
</Text>
<Text size='sm' c='dimmed' ta='center'>
Use the retry button to try again, or adjust the filters to broaden the results.
</Text>
<Group gap='xs'>
<Button size='xs' variant='light' color='gray' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
Retry
</Button>
{hasActiveFilters ? (
<Button size='xs' variant='subtle' onClick={onResetFilters}>
Clear filters
</Button>
) : null}
</Group>
</>
) : (
<>
<Text fw={600} size='sm'>
No traces found
</Text>
<Text size='sm' c='dimmed' ta='center'>
{hasActiveFilters
? 'Try adjusting the search to see more results.'
: 'Try refreshing to fetch the latest traces.'}
</Text>
<Group gap='xs'>
<Button size='xs' variant='light' leftSection={<IconRefresh size={14} />} onClick={onRefetch}>
Refresh
</Button>
{hasActiveFilters ? (
<Button size='xs' variant='subtle' onClick={onResetFilters}>
Clear filters
</Button>
) : null}
</Group>
</>
)}
</Stack>
);
const emptyState = selectionEmptyState ?? fallbackEmptyState;
return (
<Box ref={tableContainerRef}>
<DataTable<TracesTableRecord>
classNames={{ root: 'traces-table' }}
withTableBorder
withColumnBorders
highlightOnHover
verticalAlign='center'
minHeight={traceRecords.length === 0 ? 500 : undefined}
idAccessor='spanId'
records={traceRecords}
columns={responsiveColumns}
totalRecords={totalRecords}
recordsPerPage={recordsPerPage}
page={page}
onPageChange={onPageChange}
onRecordsPerPageChange={onRecordsPerPageChange}
recordsPerPageOptions={recordsPerPageOptions}
sortStatus={sortStatus}
onSortStatusChange={handleSortStatusChange}
fetching={isFetching}
loaderSize='sm'
emptyState={traceRecords.length === 0 ? emptyState : undefined}
/>
</Box>
);
}
@@ -0,0 +1,367 @@
// Copyright (c) Microsoft. All rights reserved.
import { useMemo, useState } from 'react';
import type { Meta, StoryObj } from '@storybook/react';
import { IconSearch } from '@tabler/icons-react';
import { Box, Stack, TextInput, Title } from '@mantine/core';
import type { Span } from '@/types';
import { compareRecords } from '@/utils/table';
import { buildTraceRecord, TracesTable, type TracesTableRecord } from './TracesTable.component';
const meta: Meta<typeof TracesTable> = {
title: 'Components/TracesTable',
component: TracesTable,
parameters: {
layout: 'fullscreen',
},
};
export default meta;
type Story = StoryObj<typeof TracesTable>;
const now = Math.floor(1762775145209 / 1000);
const sampleSpans: Span[] = [
{
rolloutId: 'ro-trace-001',
attemptId: 'at-trace-001',
sequenceId: 1,
traceId: 'trace-abc123def456',
spanId: 'span-root-001',
parentId: null,
name: 'main_task',
status: { status_code: 'OK', description: null },
attributes: {
'task.type': 'generation',
'task.priority': 'high',
'user.id': 'user-123',
},
startTime: now - 100,
endTime: now - 10,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-001',
attemptId: 'at-trace-001',
sequenceId: 1,
traceId: 'trace-abc123def456',
spanId: 'span-child-001',
parentId: 'span-root-001',
name: 'llm_call',
status: { status_code: 'OK', description: null },
attributes: {
'llm.model': 'gpt-4',
'llm.temperature': 0.7,
'llm.max_tokens': 2048,
},
startTime: now - 90,
endTime: now - 50,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-001',
attemptId: 'at-trace-001',
sequenceId: 1,
traceId: 'trace-abc123def456',
spanId: 'span-child-002',
parentId: 'span-root-001',
name: 'database_query',
status: { status_code: 'OK', description: null },
attributes: {
'db.system': 'postgresql',
'db.operation': 'SELECT',
'db.table': 'users',
},
startTime: now - 80,
endTime: now - 70,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-002',
attemptId: 'at-trace-002',
sequenceId: 1,
traceId: 'trace-xyz789ghi012',
spanId: 'span-error-001',
parentId: 'span-missing-parent',
name: 'failed_operation',
status: { status_code: 'ERROR', description: 'Connection timeout' },
attributes: {
'error.type': 'TimeoutError',
'error.message': 'Connection timed out after 30s',
},
startTime: now - 150,
endTime: now - 120,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-003',
attemptId: 'at-trace-003',
sequenceId: 1,
traceId: 'trace-unset123',
spanId: 'span-unset-001',
parentId: null,
name: 'pending_task',
status: { status_code: 'UNSET', description: null },
attributes: {},
startTime: now - 30,
endTime: now - 5,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-004',
attemptId: 'at-trace-004',
sequenceId: 2,
traceId: 'trace-nested456',
spanId: 'span-parent-001',
parentId: null,
name: 'workflow_execution',
status: { status_code: 'OK', description: null },
attributes: {
'workflow.name': 'data_processing',
'workflow.version': '2.1.0',
},
startTime: now - 200,
endTime: now - 50,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-004',
attemptId: 'at-trace-004',
sequenceId: 2,
traceId: 'trace-nested456',
spanId: 'span-child-nested-001',
parentId: 'span-parent-001',
name: 'step_1_validation',
status: { status_code: 'OK', description: null },
attributes: {
'step.name': 'validation',
'step.index': 1,
},
startTime: now - 195,
endTime: now - 180,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
{
rolloutId: 'ro-trace-004',
attemptId: 'at-trace-004',
sequenceId: 2,
traceId: 'trace-nested456',
spanId: 'span-child-nested-002',
parentId: 'span-parent-001',
name: 'step_2_processing',
status: { status_code: 'ERROR', description: 'Validation failed' },
attributes: {
'step.name': 'processing',
'step.index': 2,
'error.type': 'ValidationError',
},
startTime: now - 175,
endTime: now - 160,
events: [],
links: [],
context: {},
parent: null,
resource: {},
},
];
type WrapperProps = {
maxWidth: number;
spans?: Span[] | undefined;
isFetching?: boolean;
isError?: boolean;
error?: unknown;
};
function TracesTableStoryWrapper({
maxWidth,
spans = sampleSpans,
isFetching = false,
isError = false,
error = null,
}: WrapperProps) {
const [searchTerm, setSearchTerm] = useState('');
const [page, setPage] = useState(1);
const [recordsPerPage, setRecordsPerPage] = useState(10);
const [sort, setSort] = useState<{ column: string; direction: 'asc' | 'desc' }>({
column: 'startTime',
direction: 'desc',
});
const tableRecords = useMemo<TracesTableRecord[]>(() => {
if (!spans) {
return [];
}
return spans.map((span) => buildTraceRecord(span));
}, [spans]);
const filteredRecords = useMemo(() => {
const normalizedSearch = searchTerm.trim().toLowerCase();
if (normalizedSearch.length === 0) {
return tableRecords;
}
return tableRecords.filter(
(record) =>
record.traceId.toLowerCase().includes(normalizedSearch) ||
record.spanId.toLowerCase().includes(normalizedSearch) ||
record.name.toLowerCase().includes(normalizedSearch),
);
}, [searchTerm, tableRecords]);
const sortedRecords = useMemo(() => {
const sorted = filteredRecords.slice();
if (!sorted.length) {
return sorted;
}
const comparatorKey = sort.column as keyof TracesTableRecord;
if (!(comparatorKey in sorted[0])) {
return sorted;
}
sorted.sort((a, b) => compareRecords(a, b, comparatorKey));
if (sort.direction === 'desc') {
sorted.reverse();
}
return sorted;
}, [filteredRecords, sort]);
const totalRecordsValue = sortedRecords.length;
const pagedRecords = useMemo(() => {
const startIndex = (page - 1) * recordsPerPage;
const endIndex = startIndex + recordsPerPage;
return sortedRecords.slice(startIndex, endIndex);
}, [page, recordsPerPage, sortedRecords]);
const pagedSpans = useMemo(() => pagedRecords.map((record) => record as Span), [pagedRecords]);
const handleShowRollout = (record: any) => {
console.log('Show rollout for:', record.rolloutId);
};
const handleShowSpanDetail = (record: any) => {
console.log('Show span detail for:', record.spanId, record);
};
const handleParentIdClick = (parentId: string) => {
console.log('Navigate to parent span:', parentId);
setSearchTerm(parentId);
};
return (
<Box mx='auto' style={{ maxWidth, width: '100%', padding: 16 }}>
<Stack gap='md'>
<Title order={2}>Traces</Title>
<TextInput
placeholder='Search by Trace ID, Span ID, or Name'
value={searchTerm}
onChange={(event) => setSearchTerm(event.currentTarget.value)}
leftSection={<IconSearch size={16} />}
data-testid='traces-search-input'
w='100%'
style={{ maxWidth: 360 }}
/>
<TracesTable
spans={pagedSpans}
totalRecords={totalRecordsValue}
isFetching={isFetching}
isError={isError}
error={error}
searchTerm={searchTerm}
sort={sort}
page={page}
recordsPerPage={recordsPerPage}
onSortStatusChange={(status) => {
setSort({
column: status.columnAccessor as string,
direction: status.direction,
});
}}
onPageChange={setPage}
onRecordsPerPageChange={(value) => {
setRecordsPerPage(value);
setPage(1);
}}
onResetFilters={() => {
setSearchTerm('');
setSort({ column: 'startTime', direction: 'desc' });
setPage(1);
}}
onRefetch={() => undefined}
onShowRollout={handleShowRollout}
onShowSpanDetail={handleShowSpanDetail}
onParentIdClick={handleParentIdClick}
recordsPerPageOptions={[10, 20, 50]}
/>
</Stack>
</Box>
);
}
export const WideContainer: Story = {
render: () => <TracesTableStoryWrapper maxWidth={1400} />,
};
export const MediumContainer: Story = {
render: () => <TracesTableStoryWrapper maxWidth={960} />,
};
export const NarrowContainer: Story = {
render: () => <TracesTableStoryWrapper maxWidth={720} />,
};
export const DrawerWidth: Story = {
render: () => <TracesTableStoryWrapper maxWidth={520} />,
};
export const ErrorState: Story = {
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} isError error={new Error('Network unreachable')} />,
};
export const LoadingState: Story = {
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} isFetching />,
};
export const EmptyState: Story = {
render: () => <TracesTableStoryWrapper maxWidth={960} spans={[]} />,
};
export const WithMissingParent: Story = {
render: () => (
<TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.spanId === 'span-error-001')} />
),
};
export const NestedSpans: Story = {
render: () => (
<TracesTableStoryWrapper maxWidth={1200} spans={sampleSpans.filter((s) => s.traceId === 'trace-nested456')} />
),
};
@@ -0,0 +1,362 @@
// Copyright (c) Microsoft. All rights reserved.
import { useCallback, useEffect, useMemo } from 'react';
import { IconCheck, IconCopy, IconInfoCircle, IconRefresh } 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 { Worker } from '@/types';
import { getErrorDescriptor } from '@/utils/error';
import { formatDateTime, formatRelativeTime, formatStatusLabel } 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> = {
workerId: { fixedWidth: 12, priority: 0 },
status: { fixedWidth: 6, priority: 1 },
currentRolloutId: { fixedWidth: 14, priority: 3 },
currentAttemptId: { fixedWidth: 14, priority: 3 },
lastHeartbeatTime: { fixedWidth: 10, priority: 2 },
lastBusyTime: { fixedWidth: 10, priority: 3 },
lastIdleTime: { fixedWidth: 10, priority: 3 },
lastDequeueTime: { fixedWidth: 10, priority: 1 },
actions: { fixedWidth: 5, priority: 0 },
};
export type WorkersTableRecord = Worker & {
timestamps: Record<
'lastHeartbeatTime' | 'lastBusyTime' | 'lastIdleTime' | 'lastDequeueTime',
{ absolute: string; relative: string }
>;
};
const buildTimestampMeta = (value: Worker['lastHeartbeatTime']) => ({
absolute: formatDateTime(value),
relative: formatRelativeTime(value),
});
function buildWorkerRecord(worker: Worker): WorkersTableRecord {
return {
...worker,
timestamps: {
lastHeartbeatTime: buildTimestampMeta(worker.lastHeartbeatTime),
lastBusyTime: buildTimestampMeta(worker.lastBusyTime),
lastIdleTime: buildTimestampMeta(worker.lastIdleTime),
lastDequeueTime: buildTimestampMeta(worker.lastDequeueTime),
},
};
}
type WorkersColumnsOptions = {
onShowDetails: (worker: Worker) => void;
};
const STATUS_COLORS: Record<Worker['status'], string> = {
busy: 'orange',
idle: 'teal',
unknown: 'gray',
};
function createWorkersColumns({ onShowDetails }: WorkersColumnsOptions): DataTableColumn<WorkersTableRecord>[] {
return [
{
accessor: 'workerId',
title: 'Runner ID',
sortable: true,
render: ({ workerId }) => (
<Group gap={2} wrap='nowrap'>
<Text fw={500} size='sm'>
{workerId}
</Text>
<CopyButton value={workerId}>
{({ copied, copy }) => (
<Tooltip label={copied ? 'Copied' : 'Copy'} withArrow>
<ActionIcon
aria-label={`Copy worker ID ${workerId}`}
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: 'status',
title: 'Status',
sortable: true,
render: ({ status }) => {
const color = STATUS_COLORS[status] ?? 'gray';
return (
<Badge size='sm' variant='light' color={color} radius='sm'>
{formatStatusLabel(status)}
</Badge>
);
},
},
{
accessor: 'currentRolloutId',
title: 'Current Rollout',
sortable: true,
render: ({ currentRolloutId }) => <Text size='sm'>{currentRolloutId ?? '—'}</Text>,
},
{
accessor: 'currentAttemptId',
title: 'Current Attempt',
sortable: true,
render: ({ currentAttemptId }) => <Text size='sm'>{currentAttemptId ?? '—'}</Text>,
},
{
accessor: 'lastHeartbeatTime',
title: 'Heartbeat',
sortable: true,
render: ({ timestamps }) => (
<Stack gap={0} justify='center'>
<Text size='sm'>{timestamps.lastHeartbeatTime.relative}</Text>
{timestamps.lastHeartbeatTime.absolute !== '—' && (
<Text size='xs' c='dimmed'>
{timestamps.lastHeartbeatTime.absolute}
</Text>
)}
</Stack>
),
},
{
accessor: 'lastBusyTime',
title: 'Last Busy',
sortable: true,
render: ({ timestamps }) => (
<Stack gap={0} justify='center'>
<Text size='sm'>{timestamps.lastBusyTime.relative}</Text>
{timestamps.lastBusyTime.absolute !== '—' && (
<Text size='xs' c='dimmed'>
{timestamps.lastBusyTime.absolute}
</Text>
)}
</Stack>
),
},
{
accessor: 'lastIdleTime',
title: 'Last Idle',
sortable: true,
render: ({ timestamps }) => (
<Stack gap={0} justify='center'>
<Text size='sm'>{timestamps.lastIdleTime.relative}</Text>
{timestamps.lastIdleTime.absolute !== '—' && (
<Text size='xs' c='dimmed'>
{timestamps.lastIdleTime.absolute}
</Text>
)}
</Stack>
),
},
{
accessor: 'lastDequeueTime',
title: 'Last Dequeue',
sortable: true,
render: ({ timestamps }) => (
<Stack gap={0} justify='center'>
<Text size='sm'>{timestamps.lastDequeueTime.relative}</Text>
{timestamps.lastDequeueTime.absolute !== '—' && (
<Text size='xs' c='dimmed'>
{timestamps.lastDequeueTime.absolute}
</Text>
)}
</Stack>
),
},
{
accessor: 'actions',
title: 'Actions',
textAlign: 'left',
render: (record) => (
<Tooltip label='Show runner detail' withArrow disabled={!onShowDetails}>
<ActionIcon
aria-label='Show runner detail'
variant='subtle'
color='gray'
onClick={(event) => {
event.stopPropagation();
onShowDetails(record);
}}
>
<IconInfoCircle size={16} />
</ActionIcon>
</Tooltip>
),
},
];
}
export type WorkersTableProps = {
workers: Worker[] | undefined;
totalRecords: number;
isFetching: boolean;
isError: boolean;
error: unknown;
searchTerm: string;
sort: { column: string; direction: 'asc' | 'desc' };
page: number;
recordsPerPage: number;
onSortStatusChange: (status: DataTableSortStatus<WorkersTableRecord>) => void;
onPageChange: (page: number) => void;
onRecordsPerPageChange: (value: number) => void;
onResetFilters: () => void;
onRefetch: () => void;
onShowDetails: (worker: Worker) => void;
recordsPerPageOptions?: number[];
};
export function WorkersTable({
workers,
totalRecords,
isFetching,
isError,
error,
searchTerm,
sort,
page,
recordsPerPage,
onSortStatusChange,
onPageChange,
onRecordsPerPageChange,
onResetFilters,
onRefetch,
onShowDetails,
recordsPerPageOptions = DEFAULT_RECORDS_PER_PAGE_OPTIONS,
}: WorkersTableProps) {
const { ref: tableContainerRef, width: containerWidth } = useElementSize();
const { width: viewportWidth } = useViewportSize();
const layoutAwareContainerWidth = useMemo(
() => getLayoutAwareWidth(containerWidth, viewportWidth),
[containerWidth, viewportWidth],
);
const workerRecords = useMemo<WorkersTableRecord[]>(() => {
if (!workers) {
return [];
}
return workers.map((worker) => buildWorkerRecord(worker));
}, [workers]);
const columns = useMemo(() => createWorkersColumns({ onShowDetails }), [onShowDetails]);
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<WorkersTableRecord> = {
columnAccessor: sort.column,
direction: sort.direction,
};
const handleSortStatusChange = useCallback(
(status: DataTableSortStatus<WorkersTableRecord>) => {
onSortStatusChange(status);
},
[onSortStatusChange],
);
const errorDescriptor = isError ? getErrorDescriptor(error) : null;
const errorMessage = isError
? `Workers are temporarily unavailable${errorDescriptor ? ` (${errorDescriptor})` : ''}.`
: 'Workers 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 search 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 workers found
</Text>
<Text size='sm' c='dimmed' ta='center'>
{hasActiveFilters
? 'Try adjusting the search to see more results.'
: 'Try refreshing to fetch the latest worker status.'}
</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<WorkersTableRecord>
classNames={{ root: 'workers-table' }}
withTableBorder
withColumnBorders
highlightOnHover
verticalAlign='center'
minHeight={workerRecords.length === 0 ? 400 : undefined}
idAccessor='workerId'
records={workerRecords}
columns={responsiveColumns}
totalRecords={totalRecords}
recordsPerPage={recordsPerPage}
page={page}
onPageChange={onPageChange}
onRecordsPerPageChange={onRecordsPerPageChange}
recordsPerPageOptions={recordsPerPageOptions}
sortStatus={sortStatus}
onSortStatusChange={handleSortStatusChange}
fetching={isFetching}
loaderSize='sm'
emptyState={workerRecords.length === 0 ? emptyState : undefined}
/>
</Box>
);
}
@@ -0,0 +1,154 @@
// 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 { Worker } from '@/types';
import { WorkersTable } from './WorkersTable.component';
const meta: Meta<typeof WorkersTable> = {
title: 'Components/WorkersTable',
component: WorkersTable,
parameters: {
layout: 'fullscreen',
},
};
export default meta;
type Story = StoryObj<typeof WorkersTable>;
const now = Math.floor(Date.now() / 1000);
const sampleWorkers: Worker[] = [
{
workerId: 'worker-east',
status: 'busy',
heartbeatStats: { queueDepth: 2, gpuUtilization: 0.82 },
lastHeartbeatTime: now - 20,
lastDequeueTime: now - 60,
lastBusyTime: now - 120,
lastIdleTime: now - 600,
currentRolloutId: 'ro-story-001',
currentAttemptId: 'at-story-010',
},
{
workerId: 'worker-west',
status: 'busy',
heartbeatStats: { queueDepth: 1 },
lastHeartbeatTime: now - 45,
lastDequeueTime: now - 300,
lastBusyTime: now - 200,
lastIdleTime: now - 4800,
currentRolloutId: 'ro-story-003',
currentAttemptId: 'at-story-033',
},
{
workerId: 'worker-north',
status: 'idle',
heartbeatStats: { queueDepth: 0 },
lastHeartbeatTime: now - 90,
lastDequeueTime: now - 3600,
lastBusyTime: now - 5400,
lastIdleTime: now - 5400,
currentRolloutId: null,
currentAttemptId: null,
},
{
workerId: 'worker-south',
status: 'idle',
heartbeatStats: null,
lastHeartbeatTime: now - 900,
lastDequeueTime: now - 7200,
lastBusyTime: now - 8600,
lastIdleTime: now - 8600,
currentRolloutId: null,
currentAttemptId: null,
},
{
workerId: 'worker-standby',
status: 'unknown',
heartbeatStats: { queueDepth: 0 },
lastHeartbeatTime: now - 15,
lastDequeueTime: now - 4000,
lastBusyTime: null,
lastIdleTime: null,
currentRolloutId: null,
currentAttemptId: null,
},
];
type WorkersTableStoryWrapperProps = {
maxWidth: number;
initialSort?: { column: string; direction: 'asc' | 'desc' };
};
function WorkersTableStoryWrapper({ maxWidth, initialSort }: WorkersTableStoryWrapperProps) {
const [searchTerm, setSearchTerm] = useState('');
const [page, setPage] = useState(1);
const [recordsPerPage, setRecordsPerPage] = useState(5);
const [sort, setSort] = useState<{ column: string; direction: 'asc' | 'desc' }>(
() => initialSort ?? { column: 'lastHeartbeatTime', direction: 'desc' },
);
const filteredWorkers = useMemo(() => {
const normalized = searchTerm.trim().toLowerCase();
if (!normalized) {
return sampleWorkers;
}
return sampleWorkers.filter((worker) => worker.workerId.toLowerCase().includes(normalized));
}, [searchTerm]);
return (
<Stack gap='md' p='lg'>
<Title order={2}>Workers ({maxWidth}px max width)</Title>
<TextInput
placeholder='Search'
leftSection={<IconSearch size={16} />}
value={searchTerm}
onChange={(event) => setSearchTerm(event.currentTarget.value)}
w='100%'
style={{ maxWidth: 360 }}
/>
<Box style={{ maxWidth }}>
<WorkersTable
workers={filteredWorkers}
totalRecords={filteredWorkers.length}
isFetching={false}
isError={false}
error={null}
searchTerm={searchTerm}
sort={sort}
page={page}
recordsPerPage={recordsPerPage}
onSortStatusChange={(status) => {
return setSort({ column: status.columnAccessor as string, direction: status.direction });
}}
onPageChange={setPage}
onRecordsPerPageChange={setRecordsPerPage}
onResetFilters={() => {
setSearchTerm('');
setPage(1);
}}
onRefetch={() => {}}
onShowDetails={() => {}}
/>
</Box>
</Stack>
);
}
export const Wide: Story = {
render: () => <WorkersTableStoryWrapper maxWidth={1600} />,
};
export const Narrow: Story = {
render: () => <WorkersTableStoryWrapper maxWidth={780} />,
};
export const SortedByCurrentRollout: Story = {
render: () => (
<WorkersTableStoryWrapper maxWidth={1200} initialSort={{ column: 'currentRolloutId', direction: 'asc' }} />
),
};
+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;
+409
View File
@@ -0,0 +1,409 @@
// 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,
Worker,
WorkerStatus,
} 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 normalizeWorker = (value: unknown): Worker => {
const camelized = camelCaseKeys(value) as Worker;
return {
workerId: camelized.workerId,
status: camelized.status,
heartbeatStats: camelized.heartbeatStats ?? null,
lastHeartbeatTime: camelized.lastHeartbeatTime ?? null,
lastDequeueTime: camelized.lastDequeueTime ?? null,
lastBusyTime: camelized.lastBusyTime ?? null,
lastIdleTime: camelized.lastIdleTime ?? null,
currentRolloutId: camelized.currentRolloutId ?? null,
currentAttemptId: camelized.currentAttemptId ?? null,
};
};
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 GetWorkersQueryArgs = {
limit: number;
offset: number;
sortBy?: string | null;
sortOrder?: 'asc' | 'desc';
workerIdContains?: string | null;
statusIn?: WorkerStatus[];
};
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', 'Worker'],
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' }],
}),
getWorkers: builder.query<PaginatedResponse<Worker>, GetWorkersQueryArgs>({
query: ({ limit, offset, sortBy, sortOrder, workerIdContains, statusIn }) => {
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 (workerIdContains && workerIdContains.trim().length > 0) {
searchParams.set('worker_id_contains', workerIdContains.trim());
}
if (statusIn && statusIn.length > 0) {
statusIn.forEach((status) => searchParams.append('status_in', status));
}
const queryString = searchParams.toString();
const url = queryString.length > 0 ? `v1/agl/workers?${queryString}` : 'v1/agl/workers';
return { url, method: 'GET' };
},
transformResponse: (response: unknown) => normalizePaginatedResponse(response, normalizeWorker),
providesTags: (result) =>
result
? [
{ type: 'Worker' as const, id: 'LIST' },
...result.items.map((worker) => ({ type: 'Worker' as const, id: worker.workerId })),
]
: [{ type: 'Worker' 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,
useGetWorkersQuery,
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)!;
});
+72
View File
@@ -0,0 +1,72 @@
// Copyright (c) Microsoft. All rights reserved.
import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
export type AlertTone = 'info' | 'warning' | 'error';
export type AppAlert = {
id: string;
message: string;
tone: AlertTone;
isVisible: boolean;
createdAt: number;
};
export type AlertsState = {
alerts: AppAlert[];
};
const initialState: AlertsState = {
alerts: [],
};
type ShowAlertPayload = {
id: string;
message: string;
tone?: AlertTone;
};
type HideAlertPayload = {
id: string;
};
const alertsSlice = createSlice({
name: 'alert',
initialState,
reducers: {
showAlert(state, action: PayloadAction<ShowAlertPayload>) {
const { id, message, tone = 'info' } = action.payload;
const existing = state.alerts.find((alert) => alert.id === id);
if (existing) {
existing.message = message;
existing.tone = tone;
existing.isVisible = true;
existing.createdAt = Date.now();
} else {
state.alerts.push({
id,
message,
tone,
isVisible: true,
createdAt: Date.now(),
});
}
},
hideAlert(state, action: PayloadAction<HideAlertPayload>) {
state.alerts = state.alerts.filter((alert) => alert.id !== action.payload.id);
},
dismissAlert(state, action: PayloadAction<HideAlertPayload>) {
const entry = state.alerts.find((alert) => alert.id === action.payload.id);
if (entry) {
entry.isVisible = false;
}
},
clearAlerts(state) {
state.alerts = [];
},
},
});
export const { showAlert, hideAlert, dismissAlert, clearAlerts } = alertsSlice.actions;
export const alertReducer = alertsSlice.reducer;

Some files were not shown because too many files have changed in this diff Show More