update

* fix: add BaseMessage import to sql_agent.py

* trigger

---------

Co-authored-by: Yuge Zhang <scottyugochang@gmail.com>

.github/workflows/docs.yml

@@ -57,4 +57,5 @@ jobs:
         if: github.ref == 'refs/heads/main'
         run: |
           mike deploy --push latest
-          mike set-default --push latest
+          # Always set stable to default
+          mike set-default --push stable

.github/workflows/examples.yml

@@ -1,4 +1,4 @@
-name: GPU Test
+name: Examples Test
 permissions:
   contents: read
 on:
@@ -10,22 +10,21 @@ on:
 
 jobs:
   examples:
-    runs-on: [self-hosted, linux, gpu]
-    timeout-minutes: 60
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 90
     strategy:
       matrix:
         setup: [stable, latest]
       fail-fast: false
-    container:
-      image: ghcr.io/microsoft/agent-lightning/base:latest
-      options: --gpus all --ipc=host --interactive --tty
     steps:
       - name: Check GPU status
         run: nvidia-smi
+      - name: Check disk space
+        run: df -h
       - uses: actions/checkout@v4
       - name: Create a virtual environment
         run: python3 -m venv .venv
-      - name: Install deps inside the container (${{ matrix.setup }})
+      - name: Install dependencies (${{ matrix.setup }})
         run: |
           . .venv/bin/activate
           ./scripts/setup_${{ matrix.setup }}_gpu.sh
@@ -42,6 +41,34 @@ jobs:
           name: dependencies-${{ matrix.setup }}
           path: requirements-freeze.txt
           compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          set -ex
+          . .venv/bin/activate
+          litellm --config scripts/litellm_ci.yaml --port 12306 &
+          sleep 10  # Wait for the proxy to be up
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
+
+      - name: Verify LiteLLM Proxy
+        run: |
+          set -ex
+          . .venv/bin/activate
+          python scripts/litellm_sanity_check.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Prepare Unsloth model
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/unsloth
+          rm -rf models
+          hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
+
       - name: Prepare Spider dataset
         run: |
           set -ex
@@ -58,6 +85,61 @@ jobs:
           gdown --fuzzy https://drive.google.com/file/d/1FQMyKLLd6hP9dw9rfZn1EZOWNvKaDsqw/view
           unzip calc-x-data.zip -d data
           rm calc-x-data.zip
+
+      # APO Examples test
+      - name: APO example (legacy)
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/apo
+          python legacy_apo_client.py &
+          sleep 3  # Wait for the client to be up
+          python legacy_apo_server.py
+          pkill -f legacy_apo_client.py && echo "SIGTERM sent to legacy_apo_client.py" || echo "No legacy_apo_client.py process found"
+          while pgrep -f legacy_apo_client.py; do
+            echo "Waiting for legacy_apo_client.py to finish..."
+            sleep 5
+          done
+          echo "legacy_apo_client.py has finished."
+          sleep 10
+        env:
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: APO example
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/apo
+          python apo.py | tee _ci_apo.log
+          # Check whether the log contains "Best prompt found:"
+          grep "Best prompt found:" _ci_apo.log
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+      - name: APO example debug sanity check
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/apo
+          python apo_debug.py --mode runner
+          python apo_debug.py --mode trainer
+        env:
+          # New versions follow OPENAI_BASE_URL instead of OPENAI_API_BASE
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: APO built-in algorithm
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/apo
+          python room_selector_apo.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: success() || failure()
+
       - name: Spider sanity check
         run: |
           set -ex
@@ -66,8 +148,9 @@ jobs:
           python sql_agent.py --trainer.n-workers 1 --trainer.dev true --trainer.max-tasks 2
         env:
           VERL_API_BASE: http://localhost:9999/
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+        if: success() || failure()
       - name: Calc-X MCP sanity check
         run: |
           set -ex
@@ -75,8 +158,8 @@ jobs:
           cd examples/calc_x
           python tests/test_mcp_calculator.py
         env:
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
       - name: Calc-X sanity check
         run: |
           set -ex
@@ -84,14 +167,14 @@ jobs:
           cd examples/calc_x
           python calc_agent_dev.py
         env:
-          OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          OPENAI_API_BASE: http://localhost:12306/
+          OPENAI_API_KEY: dummy
 
       # Calc-X training suddenly works after running the sanity check.
       # And it has to be run before Spider training.
       # The client side used to hang in many of my attempts.
       # Don't ask why. Don't touch this.
-      - name: Calc-X training
+      - name: Calc-X training v0.1
         run: |
           set -ex
           source .venv/bin/activate
@@ -109,8 +192,10 @@ jobs:
           sleep 10
         shell: bash
         env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
         id: calc_x_train
+        if: success() || failure()
 
       - name: Validate Calc-X training
         run: |
@@ -118,7 +203,40 @@ jobs:
           . .venv/bin/activate
           python scripts/validate_example_wandb.py ${{ steps.calc_x_train.outputs.project_name }} ${{ steps.calc_x_train.outputs.run_name }}
         env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+
+      - name: Calc-X training v0.2
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python calc_agent_v0_2.py
+          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_v0_2
+        if: success() || failure()
+
+      - name: Calc-X training v0.2 LLM Proxy
+        run: |
+          set -ex
+          source .venv/bin/activate
+          cd examples/calc_x
+          ../../scripts/restart_ray.sh
+          sleep 5
+          PYTHONUNBUFFERED=1 python calc_agent_v0_2_llm_proxy.py
+          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_v0_2_llm_proxy
+        if: success() || failure()
 
       - name: Spider training
         run: |
@@ -139,7 +257,8 @@ jobs:
         shell: bash
         env:
           VERL_API_BASE: http://localhost:9991/
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
         id: spider_train
         if: success() || failure()
 
@@ -149,8 +268,63 @@ jobs:
           . .venv/bin/activate
           python scripts/validate_example_wandb.py ${{ steps.spider_train.outputs.project_name }} ${{ steps.spider_train.outputs.run_name }}
         env:
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
 
+      # Unsloth Examples test
+      - name: Unsloth SFT example
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/unsloth
+
+          agl store --port 4747 &
+          sleep 5
+          python sft_rollout_runners.py &
+          sleep 5
+          python sft_algorithm.py
+
+          pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
+          while pgrep -f agl; do
+            echo "Waiting for agl to finish..."
+            sleep 5
+          done
+          pkill -f sft_rollout_runners.py && echo "SIGTERM sent to sft_rollout_runners.py" || echo "No sft_rollout_runners.py process found"
+          while pgrep -f sft_rollout_runners.py; do
+            echo "Waiting for sft_rollout_runners.py to finish..."
+            sleep 5
+          done
+          echo "sft_rollout_runners.py has finished."
+          sleep 10
+
+          # Check models/version_2 must exist
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        if: ${{ (success() || failure()) && matrix.setup == 'latest' }}
+
+      - name: Unsloth SFT example all-in-one
+        run: |
+          set -ex
+          . .venv/bin/activate
+          cd examples/unsloth
+          rm -rf models/version_1 models/version_2
+
+          python sft_allinone.py
+          if [ ! -d "models/version_2" ]; then
+            echo "models/version_2 does not exist"
+            exit 1
+          fi
+        env:
+          WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
+          WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
+        if: matrix.setup == 'latest'
+
+      # Cleanup
       - name: Cleanup
         run: ./scripts/cleanup.sh
         if: success() || failure()

.github/workflows/tests-full.yml

@@ -0,0 +1,69 @@
+name: GPU Test
+permissions:
+  contents: read
+on:
+  schedule:
+    # Every day at 5 AM UTC+8
+    - cron: '0 21 * * *'
+
+  workflow_dispatch:
+
+jobs:
+  tests-full:
+    runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
+    timeout-minutes: 30
+    strategy:
+      matrix:
+        setup: [stable, latest]
+      fail-fast: false
+    steps:
+      - name: Check GPU status
+        run: nvidia-smi
+      - uses: actions/checkout@v4
+      - name: Create a virtual environment
+        run: python3 -m venv .venv
+      - name: Install dependencies (${{ matrix.setup }})
+        run: |
+          . .venv/bin/activate
+          ./scripts/setup_${{ matrix.setup }}_gpu.sh
+      - name: Freeze dependencies
+        run: |
+          . .venv/bin/activate
+          which python
+          which pip
+          which uvx
+          pip list | tee requirements-freeze.txt
+      - name: Upload dependencies artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: dependencies-${{ matrix.setup }}
+          path: requirements-freeze.txt
+          compression-level: 0
+
+      - name: Launch LiteLLM Proxy
+        run: |
+          set -ex
+          . .venv/bin/activate
+          litellm --config scripts/litellm_ci.yaml --port 12306 &
+          sleep 10  # Wait for the proxy to be up
+        env:
+          AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
+          AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
+      - name: Verify LiteLLM Proxy
+        run: |
+          set -ex
+          . .venv/bin/activate
+          python scripts/litellm_sanity_check.py
+        env:
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy
+
+      - name: Run tests
+        run: |
+          set -ex
+          . .venv/bin/activate
+          pytest -v --durations=0 tests
+        env:
+          PYTEST_ADDOPTS: "--color=yes"
+          OPENAI_BASE_URL: http://localhost:12306/
+          OPENAI_API_KEY: dummy

.github/workflows/tests.yml

@@ -16,8 +16,8 @@ on:
 
 jobs:
 
-  lint:
-    name: Lint with Black
+  lint-fast:
+    name: Lint - Fast
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
@@ -29,9 +29,36 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -e .[dev]
-      - name: Run Black
+      - name: Run pre-commit
+        uses: pre-commit/action@v3.0.1
+      - name: Check Python headers
         run: |
-          black --check --diff --line-length=120 .
+          python scripts/check_python_headers.py
+      - name: Run Black
+        run: black --check .
+      - name: Run isort
+        run: isort --check-only .
+      - name: Run pyright
+        run: pyright -p pyrightconfig.fast.json
+
+  lint-slow:
+    name: Lint - Slow
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Install dependencies
+        run: |
+          ./scripts/setup_type_checking.sh
+      - name: Run Black
+        run: black --check .
+      - name: Run isort
+        run: isort --check-only .
+      - name: Run pyright
+        run: pyright -p pyrightconfig.json
 
   docs:
     name: Build documentation
@@ -94,6 +121,6 @@ jobs:
           compression-level: 0
       - name: Run tests
         run: |
-          pytest -v tests
+          pytest -v --durations=0 tests
         env:
           PYTEST_ADDOPTS: "--color=yes"

.gitignore

@@ -183,9 +183,9 @@ cython_debug/
 .abstra/
 
 # Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
 #  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
 #  you could uncomment the following to ignore the enitre vscode folder
 .vscode/
 
@@ -201,3 +201,6 @@ cython_debug/
 #  refer to https://docs.cursor.com/context/ignore-files
 .cursorignore
 .cursorindexingignore
+
+# Claude
+.claude/*.local.json

.pre-commit-config.yaml

@@ -1,8 +1,24 @@
 repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-yaml
+        exclude: ^mkdocs\.yml$
+      - id: check-toml
+      - id: check-added-large-files
+      - id: check-shebang-scripts-are-executable
+      - id: detect-private-key
+  - repo: https://github.com/pycqa/isort
+    rev: 6.0.1
+    hooks:
+      - id: isort
+        args: ["."]
   - repo: https://github.com/psf/black
     rev: 25.1.0
     hooks:
-    - id: black
-      pass_filenames: false
-      always_run: true
-      args: ["--line-length=120", "."]
+      - id: black
+        pass_filenames: false
+        always_run: true
+        args: ["."]

LICENSE

@@ -16,4 +16,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
+THE SOFTWARE.

RAI_README.md

@@ -0,0 +1,77 @@
+# Responsible AI Transparency Documentation - Agent Lightning
+
+## OVERVIEW
+
+Agent Lightning is a flexible and extensible framework that enables seamless agent optimization for any existing agent framework. Agent optimization includes various data-driven techniques to customize the agent for better performance, including but not limited to model fine-tuning, prompt tuning, and model selection. And the agent frameworks refer to popular and easy-to-use agent developing frameworks such as OpenAI Agents SDK, Microsoft AutoGen, and LangChain.
+
+### WHAT CAN AGENT LIGHTNING DO
+Agent lightning was developed to bridge the gap between agent workflow development and agent optimization, empowering developers to go beyond static, pre-trained models and unlock the full potential of adaptive, learning-based agents. Agent Lightning is a training framework which can be used for any LLMs.
+
+### INTENDED USES
+Agent Lightning is best suited for agent researchers and developers. They can easily fine-tune models in existing agent frameworks with Agent Lightning. This can improve model performance on the targeted scenarios.
+
+### OUT-OF-SCOPE USES
+Agent Lightning is not well-suited for users who are not familiar with agent development and machine learning concepts.
+
+We do not recommend using Agent Lightning in commercial or real-world applications without further testing and development. It is being released for research purposes.
+
+Agent Lightning was not designed or evaluated for all possible downstream purposes. Developers should consider its inherent limitations as they select use cases, and evaluate and mitigate for accuracy, safety, and fairness concerns specific to each intended downstream use.
+
+Agent Lightning should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.
+
+We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare).
+
+## HOW TO GET STARTED
+To begin using Agent Lightning, here are some instructions.
+1.	Install dependencies, including Python, uv, PyTorch, FlashAttention, vLLM, verl.
+2.	Clone and install Agent Lightning.
+3.	Convert the dataset (provided by the user) into parquet file, which contains multiple columns. Each column contains a data id, an input and an expected output.
+4.	Run agent, which is developed by the user.
+5.	Run the training process via “bash train.sh”
+
+## EVALUATION
+Agent Lightning was evaluated on its ability to correctly complete 3 example tasks: (1) Math. The model needs to answer some math questions, and when answering one question, the model can use the calculator as its tool to help answer. (2) Text2SQL. The model is given a question related to the database, and it is required to generate a SQL which can query the database, find the information to answer the question. (3) Retrieval-Augmented Generation (RAG). The model is given a question which needs some information from Wikipedia to answer. The model is required to generate some queries to find the related information in Wikipedia, and answer the question according to retrieved documents.
+
+### EVALUATION METHODS AND RESULTS
+For detailed evaluation methods and results, please refer to the latest version of our  [technical report](https://arxiv.org/abs/2508.03680).
+
+
+## LIMITATIONS
+Agent Lightning was developed for research and experimental purposes. Further testing and validation are needed before considering its application in commercial or real-world scenarios.
+
+Agent Lightning was designed and tested using the English language. Performance in other languages may vary and should be assessed by someone who is both an expert in the expected outputs and a native speaker of that language.
+
+Outputs generated by AI may include factual errors, fabrication, or speculation. Users are responsible for assessing the accuracy of generated content. All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
+Agent Lightning inherits any biases, errors, or omissions produced by its base model. Developers are advised to choose an appropriate base LLM/MLLM carefully, depending on the intended use case.
+We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model.
+
+## BEST PRACTICES
+Better performance can be achieved by following the instructions in how to get started section.
+
+We strongly encourage users to use LLMs/MLLMs that support robust Responsible AI mitigations, such as Azure Open AI (AOAI) services. Such services continually update their safety and RAI mitigations with the latest industry standards for responsible use. For more on AOAI’s best practices when employing foundations models for scripts and applications:
+- [Blog post on responsible AI features in AOAI that were presented at Ignite 2023](https://techcommunity.microsoft.com/t5/ai-azure-ai-services-blog/announcing-new-ai-safety-amp-responsible-ai-features-in-azure/ba-p/3983686)
+- [Overview of Responsible AI practices for Azure OpenAI models](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/overview)
+- [Azure OpenAI Transparency Note](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/transparency-note)
+- [OpenAI’s Usage policies](https://openai.com/policies/usage-policies)
+- [Azure OpenAI’s Code of Conduct](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/code-of-conduct)
+
+Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.
+
+Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning.
+
+It is the user’s responsibility to ensure that the use of Agent Lightning complies with relevant data protection regulations and organizational guidelines.
+
+## LICENSE
+We use the MIT license.
+
+## CONTACT
+We welcome feedback and collaboration from our audience. If you have suggestions, questions, or observe unexpected/offensive behavior in our technology, please contact us at agent-lightning@microsoft.com.
+
+If the team receives reports of undesired behavior or identifies issues independently, we will update this repository with appropriate mitigations.
+
+
+
+---
+
+*Last updated: September 6, 2025*
+*Document version: 1.0*

README.md

@@ -1,4 +1,6 @@
-![Agent-lightning-banner](docs/assets/readme-banner.png)
+<div style="text-align:center; margin-bottom:20px;">
+  <img src="docs/assets/readme-banner.png" alt="Agent-lightning-banner" style="max-width:600px"/>
+</div>
 
 # Agent Lightning⚡
 
@@ -28,6 +30,11 @@ Join our [Discord community](https://discord.gg/RYk7CdvDR7) to connect with othe
 - 7/26/2025 [We discovered an approach to train any AI agent with RL, with (almost) zero code changes.](https://www.reddit.com/r/LocalLLaMA/comments/1m9m670/we_discovered_an_approach_to_train_any_ai_agent/) Reddit.
 - 6/6/2025 [Agent Lightning - Microsoft Research](https://www.microsoft.com/en-us/research/project/agent-lightning/) Project page.
 
+## ⚡ Community Projects
+
+- [DeepWerewolf](https://github.com/af-74413592/DeepWerewolf) — A case study of agent RL training for the Chinese Werewolf game built with AgentScope and Agent Lightning.
+- [AgentFlow](https://agentflow.stanford.edu/) — A modular multi-agent framework that combines planner, executor, verifier, and generator agents with the Flow-GRPO algorithm to tackle long-horizon, sparse-reward tasks.
+
 ## ⚡ Installation
 
 First, let's get your environment set up. We'll be using `/path/to/agentlightning` to refer to the directory containing this README file.
@@ -143,13 +150,13 @@ If you find Agent Lightning useful in your research or projects, please cite our
 
 ```bibtex
 @misc{luo2025agentlightningtrainai,
-      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning}, 
+      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning},
       author={Xufang Luo and Yuge Zhang and Zhiyuan He and Zilong Wang and Siyun Zhao and Dongsheng Li and Luna K. Qiu and Yuqing Yang},
       year={2025},
       eprint={2508.03680},
       archivePrefix={arXiv},
       primaryClass={cs.AI},
-      url={https://arxiv.org/abs/2508.03680}, 
+      url={https://arxiv.org/abs/2508.03680},
 }
 ```
 

SECURITY.md

@@ -11,4 +11,4 @@ For security reporting information, locations, contact information, and policies
 please review the latest guidance for Microsoft repositories at
 [https://aka.ms/SECURITY.md](https://aka.ms/SECURITY.md).
 
-<!-- END MICROSOFT SECURITY.MD BLOCK -->
+<!-- END MICROSOFT SECURITY.MD BLOCK -->

a.py

@@ -0,0 +1,22 @@
+import asyncio
+
+
+async def a():
+    print("a")
+    b()
+    print("finish")
+
+
+def b():
+    print("b")
+    loop = asyncio.get_running_loop()
+    fut = asyncio.run_coroutine_threadsafe(c(), loop)
+    fut.result(timeout=5.0)
+
+
+async def c():
+    print("c")
+    await asyncio.sleep(0.1)
+
+
+asyncio.run(a())

agentlightning/__init__.py

@@ -1,10 +1,19 @@
-__version__ = "0.1.2"
+# Copyright (c) Microsoft. All rights reserved.
 
-from .client import AgentLightningClient, DevTaskLoader
-from .config import lightning_cli
-from .litagent import LitAgent
-from .logging import configure_logger
-from .reward import reward
-from .server import AgentLightningServer
-from .trainer import Trainer
+__version__ = "0.2.0"
+
+from .adapter import *
+from .algorithm import *
+from .client import AgentLightningClient, DevTaskLoader  # deprecated  # type: ignore
+from .config import *
+from .emitter import *
+from .execution import *
+from .litagent import *
+from .llm_proxy import *
+from .logging import *
+from .runner import *
+from .server import AgentLightningServer  # deprecated  # type: ignore
+from .store import *
+from .tracer import *
+from .trainer import *
 from .types import *

agentlightning/adapter/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import Adapter, TraceAdapter
+from .messages import TraceToMessages
+from .triplet import LlmProxyTraceToTriplet, TracerTraceToTriplet, TraceToTripletBase
+
+__all__ = [
+    "TraceAdapter",
+    "Adapter",
+    "TraceToTripletBase",
+    "TracerTraceToTriplet",
+    "LlmProxyTraceToTriplet",
+    "TraceToMessages",
+]

agentlightning/adapter/base.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Generic, List, TypeVar
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import Span
+
+T_from = TypeVar("T_from")
+T_to = TypeVar("T_to")
+
+
+class Adapter(Generic[T_from, T_to]):
+    """Base class for synchronous adapters that convert data from one format to another.
+
+    This class defines a simple protocol for transformation:
+
+    - The `__call__` method makes adapters callable, so they can be used like functions.
+    - Subclasses must implement the `adapt` method to define the actual conversion logic.
+
+    Type parameters:
+
+    - T_from: The source data type (input).
+    - T_to: The target data type (output).
+
+    Example:
+
+        >>> class IntToStrAdapter(Adapter[int, str]):
+        ...     def adapt(self, source: int) -> str:
+        ...         return str(source)
+        ...
+        >>> adapter = IntToStrAdapter()
+        >>> adapter(42)
+        '42'
+    """
+
+    def __call__(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        This method delegates to `adapt` and allows the adapter
+        to be invoked as a function.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        return self.adapt(source)
+
+    def adapt(self, source: T_from, /) -> T_to:
+        """Convert the data to the target format.
+
+        Subclasses should override this method with the concrete
+        transformation logic.
+
+        Args:
+            source: Input data in the source format.
+
+        Returns:
+            Data converted to the target format.
+        """
+        raise NotImplementedError("Adapter.adapt() is not implemented")
+
+
+class OtelTraceAdapter(Adapter[List[ReadableSpan], T_to], Generic[T_to]):
+    """Base class for adapters that convert OpenTelemetry trace spans into other formats.
+
+    This class specializes `Adapter` for working with OpenTelemetry `ReadableSpan`
+    objects. It expects a list of spans as input and produces a custom target format
+    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+
+    Subclasses should override `adapt` to define the desired conversion.
+
+    Type parameters:
+        T_to: The target data type that spans should be converted into.
+
+    Example:
+        >>> class TraceToDictAdapter(TraceAdapter[dict]):
+        ...     def adapt(self, spans: List[ReadableSpan]) -> dict:
+        ...         return {"count": len(spans)}
+        ...
+        >>> adapter = TraceToDictAdapter()
+        >>> adapter([span1, span2])
+        {'count': 2}
+    """
+
+
+class TraceAdapter(Adapter[List[Span], T_to], Generic[T_to]):
+    """Base class for adapters that convert trace spans into other formats.
+
+    This class specializes `Adapter` for working with trace spans. It expects a list of
+    Agent-lightning spans as input and produces a custom target format
+    (e.g., reinforcement learning training data, SFT datasets, logs, metrics).
+    """

agentlightning/adapter/messages.py

@@ -0,0 +1,217 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+from collections import defaultdict
+from typing import Any, Dict, Generator, Iterable, List, Optional, TypedDict, Union, cast
+
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionFunctionToolParam,
+    ChatCompletionMessageFunctionToolCallParam,
+    ChatCompletionMessageParam,
+)
+from pydantic import TypeAdapter
+
+from agentlightning.types import Span
+
+from .base import TraceAdapter
+
+
+class OpenAIMessages(TypedDict):
+    messages: List[ChatCompletionMessageParam]
+    tools: Optional[List[ChatCompletionFunctionToolParam]]
+
+
+class _RawSpanInfo(TypedDict):
+    prompt: List[Dict[str, Any]]
+    completion: List[Dict[str, Any]]
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    tools: List[Dict[str, Any]]
+
+
+def group_genai_dict(data: Dict[str, Any], prefix: str) -> Union[Dict[str, Any], List[Any]]:
+    """
+    Convert a flat dict with keys like 'gen_ai.prompt.0.role'
+    into structured nested dicts or lists under the given prefix.
+
+    Args:
+        data: Flat dictionary (keys are dotted paths).
+        prefix: Top-level key to extract (e.g., 'gen_ai.prompt').
+
+    Returns:
+        A nested dict (if no index detected) or list (if indexed).
+    """
+    result: Union[Dict[str, Any], List[Any]] = {}
+
+    # Collect keys that match the prefix
+    relevant = {k[len(prefix) + 1 :]: v for k, v in data.items() if k.startswith(prefix + ".")}
+
+    # Detect if we have numeric indices (-> list) or not (-> dict)
+    indexed = any(part.split(".")[0].isdigit() for part in relevant.keys())
+
+    if indexed:
+        # Group by index
+        grouped: Dict[int, Dict[str, Any]] = defaultdict(dict)
+        for k, v in relevant.items():
+            parts = k.split(".")
+            if not parts[0].isdigit():
+                continue
+            idx, rest = int(parts[0]), ".".join(parts[1:])
+            grouped[idx][rest] = v
+        # Recursively build
+        result = []
+        for i in sorted(grouped.keys()):
+            result.append(group_genai_dict({f"{prefix}.{rest}": val for rest, val in grouped[i].items()}, prefix))
+    else:
+        # No indices: build dict
+        nested: Dict[str, Any] = defaultdict(dict)
+        for k, v in relevant.items():
+            if "." in k:
+                head, _tail = k.split(".", 1)
+                nested[head][f"{prefix}.{k}"] = v
+            else:
+                result[k] = v
+        # Recurse into nested dicts
+        for head, subdict in nested.items():
+            result[head] = group_genai_dict(subdict, prefix + "." + head)
+
+    return result
+
+
+def convert_to_openai_messages(prompt_completion_list: List[_RawSpanInfo]) -> Generator[OpenAIMessages, None, None]:
+    """
+    Convert raw tool call traces + prompt/completion list
+    into OpenAI fine-tuning JSONL format (tool calling style).
+
+    https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning-functions
+    """
+    for pc_entry in prompt_completion_list:
+        messages: List[ChatCompletionMessageParam] = []
+
+        # Extract messages
+        for msg in pc_entry["prompt"]:
+            role = msg["role"]
+
+            if role == "assistant" and "tool_calls" in msg:
+                # Use the tool_calls directly
+                # This branch is usually not used in the wild.
+                tool_calls: List[ChatCompletionMessageFunctionToolCallParam] = [
+                    ChatCompletionMessageFunctionToolCallParam(
+                        id=call["id"],
+                        type="function",
+                        function={"name": call["name"], "arguments": call["arguments"]},
+                    )
+                    for call in msg["tool_calls"]
+                ]
+                messages.append(
+                    ChatCompletionAssistantMessageParam(role="assistant", content=None, tool_calls=tool_calls)
+                )
+            else:
+                # Normal user/system/tool content
+                message = cast(
+                    ChatCompletionMessageParam,
+                    TypeAdapter(ChatCompletionMessageParam).validate_python(
+                        dict(role=role, content=msg.get("content", ""), tool_call_id=msg.get("tool_call_id", None))
+                    ),
+                )
+                messages.append(message)
+
+        # Extract completions (assistant outputs after tool responses)
+        for comp in pc_entry["completion"]:
+            if comp.get("role") == "assistant":
+                content = comp.get("content")
+                if pc_entry["tools"]:
+                    tool_calls = [
+                        ChatCompletionMessageFunctionToolCallParam(
+                            id=tool["call"]["id"],
+                            type=tool["call"]["type"],
+                            function={"name": tool["name"], "arguments": tool["parameters"]},
+                        )
+                        for tool in pc_entry["tools"]
+                    ]
+                    messages.append(
+                        ChatCompletionAssistantMessageParam(role="assistant", content=content, tool_calls=tool_calls)
+                    )
+                else:
+                    messages.append(ChatCompletionAssistantMessageParam(role="assistant", content=content))
+
+        # Build tools definitions (if available)
+        if "functions" in pc_entry["request"]:
+            tools = [
+                ChatCompletionFunctionToolParam(
+                    type="function",
+                    function={
+                        "name": fn["name"],
+                        "description": fn.get("description", ""),
+                        "parameters": (
+                            json.loads(fn["parameters"]) if isinstance(fn["parameters"], str) else fn["parameters"]
+                        ),
+                    },
+                )
+                for fn in pc_entry["request"]["functions"]
+            ]
+            yield OpenAIMessages(messages=messages, tools=tools)
+        else:
+            yield OpenAIMessages(messages=messages, tools=None)
+
+
+class TraceToMessages(TraceAdapter[List[OpenAIMessages]]):
+    """
+    Adapter that converts OpenTelemetry trace spans into OpenAI-compatible message format.
+
+    This adapter processes trace spans containing LLM conversation data and transforms them
+    into structured OpenAI message format suitable for fine-tuning or analysis. It extracts
+    prompts, completions, tool calls, and function definitions from trace attributes and
+    reconstructs the conversation flow.
+
+    The adapter handles:
+    - Converting flat trace attributes into structured message objects
+    - Extracting and matching tool calls with their corresponding requests
+    - Building proper OpenAI ChatCompletionMessage objects with roles, content, and tool calls
+    - Generating function definitions for tools used in conversations
+    """
+
+    def get_tool_calls(self, completion: Span, all_spans: List[Span], /) -> Iterable[Dict[str, Any]]:
+        """Find tool calls in the trace. Returns a dict with the tool call id, name, and arguments.
+
+        The spans that are direct children of the completion span are the tool calls.
+        """
+        # Get all the spans that are children of the completion span
+        children = [span for span in all_spans if span.parent_id == completion.span_id]
+        # Get the tool calls from the children
+        for maybe_tool_call in children:
+            tool_call = group_genai_dict(maybe_tool_call.attributes, "tool")
+            if not isinstance(tool_call, dict):
+                raise ValueError(f"Extracted tool call from trace is not a dict: {tool_call}")
+            if tool_call:
+                yield tool_call
+
+    def adapt(self, source: List[Span], /) -> List[OpenAIMessages]:
+        raw_prompt_completions: List[_RawSpanInfo] = []
+
+        for span in source:
+            attributes = {k: v for k, v in span.attributes.items()}
+
+            # Get all related information from the trace span
+            prompt = group_genai_dict(attributes, "gen_ai.prompt") or []
+            completion = group_genai_dict(attributes, "gen_ai.completion") or []
+            request = group_genai_dict(attributes, "gen_ai.request") or {}
+            response = group_genai_dict(attributes, "gen_ai.response") or {}
+            if not isinstance(prompt, list):
+                raise ValueError(f"Extracted prompt from trace is not a list: {prompt}")
+            if not isinstance(completion, list):
+                raise ValueError(f"Extracted completion from trace is not a list: {completion}")
+            if not isinstance(request, dict):
+                raise ValueError(f"Extracted request from trace is not a dict: {request}")
+            if not isinstance(response, dict):
+                raise ValueError(f"Extracted response from trace is not a dict: {response}")
+            if prompt or completion or request or response:
+                tools = list(self.get_tool_calls(span, source)) or []
+                raw_prompt_completions.append(
+                    _RawSpanInfo(
+                        prompt=prompt or [], completion=completion, request=request, response=response, tools=tools
+                    )
+                )
+
+        return list(convert_to_openai_messages(raw_prompt_completions))

agentlightning/adapter/triplet.py

@@ -1,12 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
 import json
 import re
 from enum import Enum
-from typing import List, Dict, Tuple, Optional, Any
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
-from pydantic import BaseModel
-from opentelemetry import trace as trace_api
 from opentelemetry.sdk.trace import ReadableSpan
-from agentlightning.types import Triplet
+from pydantic import BaseModel
+
+from agentlightning.types import SpanNames, Triplet
+from agentlightning.types.tracer import Span
+
+from .base import TraceAdapter
 
 
 class Transition(BaseModel):
@@ -45,7 +52,7 @@ class TraceTree:
     def __init__(
         self,
         id: str,
-        span: ReadableSpan,
+        span: Span,
         children: Optional[List["TraceTree"]] = None,
     ):
         self.id = id
@@ -82,7 +89,7 @@ class TraceTree:
 
         dot = graphviz.Digraph(comment="Trace Tree")
 
-        should_visit_cache = {}
+        should_visit_cache: Dict[str, bool] = {}
 
         def should_visit(node: "TraceTree") -> bool:
             if node.id in should_visit_cache:
@@ -108,14 +115,14 @@ class TraceTree:
             vis_name = node.id[:8] + " (" + node.span.name + ")"
             if agent_name is not None:
                 vis_name += " [" + agent_name + "]"
-            dot.node(node.id, vis_name)
+            dot.node(node.id, vis_name)  # type: ignore
             for child in node.children:
                 if visit(child):
-                    dot.edge(node.id, child.id)
+                    dot.edge(node.id, child.id)  # type: ignore
             return True
 
         visit(self)
-        dot.render(filename, format="png", cleanup=True)
+        dot.render(filename, format="png", cleanup=True)  # type: ignore
 
     def names_tuple(self) -> Tuple[str, List[Any]]:
         """Return the span name, and a list of children.
@@ -126,7 +133,7 @@ class TraceTree:
         agent_name = self.agent_name()
         if agent_name is not None:
             name += " [" + agent_name + "]"
-        children_names = []
+        children_names: List[Tuple[str, List[Any]]] = []
         for child in self.children:
             child_name, child_children = child.names_tuple()
             children_names.append((child_name, child_children))
@@ -142,14 +149,18 @@ class TraceTree:
         return spans
 
     def to_json(self) -> dict[str, Any]:
+        if isinstance(self.span, ReadableSpan):
+            span_data = json.loads(self.span.to_json())
+        else:
+            span_data = self.span.model_dump()
         return {
             "id": self.id,
-            "span": self.span.to_json(),
+            "span": span_data,
             "children": [child.to_json() for child in self.children],
         }
 
     @classmethod
-    def from_spans(cls, spans: List[ReadableSpan]) -> "TraceTree":
+    def from_spans(cls, spans: List[Span]) -> "TraceTree":
         """
         Create a TraceTree from a list of spans.
         All spans without parents found will be considered as candidate root spans.
@@ -160,17 +171,18 @@ class TraceTree:
             raise ValueError("No spans provided to create TraceTree.")
 
         # Process trace items in topological order
-        id_to_span = {span.get_span_context().span_id: span for span in spans}
+        id_to_span = {span.span_id: span for span in spans}
 
-        forward_graph: dict[int, list[int]] = {}
-        root_ids: list[int] = []
+        forward_graph: dict[str, list[str]] = {}
+        root_ids: list[str] = []
         for span in spans:
-            if span.parent is None:
-                root_ids.append(span.get_span_context().span_id)
+            span_id = span.span_id
+            if span.parent_id is None:
+                root_ids.append(span.span_id)
             else:
-                if span.parent.span_id not in forward_graph:
-                    forward_graph[span.parent.span_id] = []
-                forward_graph[span.parent.span_id].append(span.get_span_context().span_id)
+                if span.parent_id not in forward_graph:
+                    forward_graph[span.parent_id] = []
+                forward_graph[span.parent_id].append(span_id)
 
         # Diff between span with data and forward_graph keys
         # Sometimes the top-level session span is lost.
@@ -178,7 +190,7 @@ class TraceTree:
         for unfound_root in unfound_roots:
             root_ids.append(unfound_root)
 
-        def visit(node_id):
+        def visit(node_id: str) -> "TraceTree":
             children: list[TraceTree] = []
             if node_id in forward_graph:
                 for child_id in forward_graph[node_id]:
@@ -186,22 +198,21 @@ class TraceTree:
 
             if node_id not in id_to_span:
                 assert len(children) > 0
-                virtual_span = ReadableSpan(
-                    context=trace_api.SpanContext(
-                        trace_id=children[0].span.get_span_context().trace_id,
-                        span_id=node_id,
-                        is_remote=False,
-                    ),
-                    name="virtual-node",
-                    kind=trace_api.SpanKind.INTERNAL,
+                virtual_span = Span.from_attributes(
+                    rollout_id=children[0].span.rollout_id,
+                    attempt_id=children[0].span.attempt_id,
+                    sequence_id=children[0].span.sequence_id,
+                    trace_id=children[0].span.trace_id,
+                    span_id=node_id,
+                    parent_id=None,
                     attributes={},
-                    start_time=min(child.start_time for child in children),
-                    end_time=max(child.end_time for child in children),
+                    start_time=min(child.start_time for child in children if child.start_time is not None),
+                    end_time=max(child.end_time for child in children if child.end_time is not None),
                 )
-                return cls(trace_api.format_span_id(node_id), virtual_span, children=children)
+                return cls(node_id, virtual_span, children=children)
             else:
                 return cls(
-                    trace_api.format_span_id(node_id),
+                    node_id,
                     id_to_span[node_id],
                     children=children,
                 )
@@ -211,14 +222,14 @@ class TraceTree:
             root_spans = [visit(root_id) for root_id in root_ids]
             virtual_root = TraceTree(
                 id="virtual-root",
-                span=ReadableSpan(
-                    context=trace_api.SpanContext(
-                        trace_id=root_spans[0].span.get_span_context().trace_id,
-                        span_id=0,
-                        is_remote=False,
-                    ),
+                span=Span.from_attributes(
+                    rollout_id=root_spans[0].span.rollout_id,
+                    attempt_id=root_spans[0].span.attempt_id,
+                    sequence_id=root_spans[0].span.sequence_id,
+                    trace_id=root_spans[0].span.trace_id,
+                    span_id=None,  # Generate one
+                    parent_id=None,
                     name="virtual-root",
-                    kind=trace_api.SpanKind.INTERNAL,
                     attributes={},
                     start_time=root_spans[0].start_time,
                     end_time=root_spans[-1].end_time,
@@ -236,26 +247,34 @@ class TraceTree:
     def agent_name(self) -> Optional[str]:
         """Return the name of agent span. Return the agent or None (not an agent at all).
         Extend this function to support more agent frameworks."""
+        attributes = self.span.attributes
+        if attributes is None:  # type: ignore
+            return None
 
         # Case 1: OpenAI Agent SDK
-        agent_name = self.span.attributes.get("agent.name")
+        agent_name = cast(Optional[str], attributes.get("agent.name"))
         if agent_name is not None:
             return agent_name
 
         # Case 2: Agentops decorator @agent
-        is_agent = self.span.attributes.get("agentops.span.kind") == "agent"
+        is_agent = attributes.get("agentops.span.kind") == "agent"
         if is_agent:
-            agent_name = self.span.attributes.get("operation.name")
+            agent_name = cast(Optional[str], attributes.get("operation.name"))
             if agent_name is not None:
                 return agent_name
 
         # Case 3: Autogen team
-        agent_name = self.span.attributes.get("recipient_agent_type")
+        agent_name = cast(Optional[str], attributes.get("recipient_agent_type"))
         if agent_name is not None:
             return agent_name
 
         # Case 4: LangGraph
-        agent_name = self.span.attributes.get("langchain.chain.type")
+        agent_name = cast(Optional[str], attributes.get("langchain.chain.type"))
+        if agent_name is not None:
+            return agent_name
+
+        # Case 5: agent-framework
+        agent_name = cast(Optional[str], attributes.get("executor.id"))
         if agent_name is not None:
             return agent_name
 
@@ -264,7 +283,7 @@ class TraceTree:
             "agentops.task.output",  # newer versions of agentops
             "agentops.entity.output",
         ]:
-            output = self.span.attributes.get(key)
+            output = self.span.attributes.get(key)  # type: ignore
             if output:
                 if isinstance(output, dict):
                     return output
@@ -273,11 +292,15 @@ class TraceTree:
                         return json.loads(output)
                     except json.JSONDecodeError:
                         return {}
+
+        # Latest emit reward format
+        if self.span.name == SpanNames.REWARD.value and self.span.attributes:
+            return {"type": "reward", "value": self.span.attributes.get("reward", None)}
         return {}
 
     def is_reward_span(self) -> bool:
         maybe_reward = self.maybe_reward_dict()
-        return maybe_reward and maybe_reward.get("type") == "reward"
+        return maybe_reward and maybe_reward.get("type") == "reward"  # type: ignore
 
     def find_llm_calls(
         self,
@@ -307,7 +330,7 @@ class TraceTree:
             is_llm_call = False
         if is_llm_call:
             # Check the response id
-            response_id = self.span.attributes.get("gen_ai.response.id")
+            response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id")  # type: ignore
             if response_id is None and within_llm_call is True:
                 is_llm_call = False
             if (
@@ -318,7 +341,7 @@ class TraceTree:
                 is_llm_call = False
 
             if is_llm_call:
-                llm_calls.append((self, within_matching_subtree))
+                llm_calls.append((self, within_matching_subtree))  # type: ignore
                 existing_llm_call_response_ids = existing_llm_call_response_ids or set()
                 if response_id is not None:
                     existing_llm_call_response_ids.add(response_id)
@@ -375,10 +398,10 @@ class TraceTree:
                     continue
                 if node is self:
                     continue
-                if node.start_time <= repair_node.start_time and node.end_time >= repair_node.end_time:
-                    duration_delta = node.end_time - repair_node.end_time + repair_node.start_time - node.start_time
+                if node.start_time <= repair_node.start_time and node.end_time >= repair_node.end_time:  # type: ignore
+                    duration_delta = node.end_time - repair_node.end_time + repair_node.start_time - node.start_time  # type: ignore
                     if duration_delta > 0 and duration_delta < closest_duration:
-                        closest_duration = duration_delta
+                        closest_duration = duration_delta  # type: ignore
                         closest_parent = node
 
             # Repair the hierarchy
@@ -392,18 +415,18 @@ class TraceTree:
         rewards: dict[str, Optional[float]] = {}
 
         if reward_match == RewardMatchPolicy.FIRST_OCCURRENCE:
-            time_sorted: List[TraceTree] = sorted(self.traverse(), key=lambda x: x.start_time)
-            assign_to: List[Tuple[str, int]] = []
+            time_sorted: List[TraceTree] = cast(List[TraceTree], sorted(self.traverse(), key=lambda x: x.start_time))  # type: ignore
+            assign_to: List[Tuple[str, int]] = []  # type: ignore
             for item in time_sorted:
                 if item.id in llm_call_ids:
-                    assign_to.append((item.id, item.end_time))
+                    assign_to.append((item.id, item.end_time))  # type: ignore
 
                 # get reward
                 agentops_output = item.maybe_reward_dict()
                 if agentops_output and agentops_output.get("type") == "reward":
                     for assign_to_id, assign_to_end_time in reversed(assign_to):
                         # This reward happens before the end of the LLM call.
-                        if assign_to_end_time > item.start_time:
+                        if assign_to_end_time > item.start_time:  # type: ignore
                             continue
                         # Ok, we found someone to assign to
                         if assign_to_id in rewards:
@@ -417,12 +440,12 @@ class TraceTree:
                 assign_to: List[Tuple[str, int]] = []
                 for child in item.children:
                     if child.id in llm_call_ids:
-                        assign_to.append(child.id)
+                        assign_to.append(child.id)  # type: ignore
 
                     agentops_output = item.maybe_reward_dict()
                     if agentops_output and agentops_output.get("type") == "reward":
                         for assign_to_id, assign_to_end_time in reversed(assign_to):
-                            if assign_to_end_time > item.start_time:
+                            if assign_to_end_time > item.start_time:  # type: ignore
                                 # This reward happens before the end of the LLM call.
                                 continue
                             if assign_to_id in rewards:
@@ -468,11 +491,11 @@ class TraceTree:
             (
                 llm_call.id,
                 Triplet(
-                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},
-                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},
+                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},  # type: ignore
+                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},  # type: ignore
                     reward=None,
                     metadata=dict(
-                        response_id=llm_call.span.attributes.get(
+                        response_id=llm_call.span.attributes.get(  # type: ignore
                             "gen_ai.response.id", None
                         ),  # it works at least for OpenAI
                         agent_name=agent_name,
@@ -498,9 +521,15 @@ class TraceTree:
         )
 
 
-class TripletExporter:
+class TraceToTripletBase(TraceAdapter[List[Triplet]]):
     """
-    A class to export triplet data from OpenTelemetry spans.
+    Base class for trace triplet adapters.
+    """
+
+
+class TracerTraceToTriplet(TraceToTripletBase):
+    """
+    An adapter to convert OpenTelemetry spans to triplet data.
 
     Attributes:
         repair_hierarchy: When `repair_hierarchy` is set to True, the trace will be repaired with the time information.
@@ -525,9 +554,41 @@ class TripletExporter:
         self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
         self.reward_match = reward_match
 
-    def export(self, spans: List[ReadableSpan]) -> List[Triplet]:
+    def visualize(
+        self,
+        source: Union[List[Span], List[ReadableSpan]],
+        /,
+        filename: str = "trace_tree",
+        interested_span_match: str | None = None,
+    ) -> TraceTree:
+        """
+        Visualize the trace tree.
+
+        Args:
+            source (List[Span]): The list of OpenTelemetry spans to visualize.
+            filename (str): The base filename for the output visualization (default: "trace_tree").
+            interested_span_match (str | None): Optional regular expression pattern to highlight or focus on specific spans in the visualization.
+
+        Returns:
+            TraceTree: The constructed trace tree object.
+        """
+        source_normalized = [
+            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
+            for span in source
+        ]
+        trace_tree = TraceTree.from_spans(source_normalized)
+        if self.repair_hierarchy:
+            trace_tree.repair_hierarchy()
+        trace_tree.visualize(filename, interested_span_match=interested_span_match)
+        return trace_tree
+
+    def adapt(self, source: Union[List[Span], List[ReadableSpan]], /) -> List[Triplet]:  # type: ignore
         """Convert OpenTelemetry spans to a list of Triplet objects."""
-        trace_tree = TraceTree.from_spans(spans)
+        source_normalized = [
+            Span.from_opentelemetry(span, "dummy", "dummy", 0) if isinstance(span, ReadableSpan) else span
+            for span in source
+        ]
+        trace_tree = TraceTree.from_spans(source_normalized)
         if self.repair_hierarchy:
             trace_tree.repair_hierarchy()
         trajectory = trace_tree.to_trajectory(
@@ -537,3 +598,190 @@ class TripletExporter:
             reward_match=self.reward_match,
         )
         return trajectory
+
+
+class LlmProxyTraceToTriplet(TraceToTripletBase):
+    """
+    Converting telemetry data emitted by the LLM Proxy to triplet data.
+    This adapter is very experimental. Should only be used when the TracerTraceToTriplet does not work at all.
+
+    IMPORTANT: Do NOT rely on timestamps here. Proxy spans can be emitted from different
+    machines with unsynchronized clocks. We therefore treat `sequence_id` as the only
+    reliable ordering primitive and perform "first occurrence" reward matching using
+    sequence order only.
+
+    Strategy:
+
+    1) Sort spans by (sequence_id, start_time).
+    2) Extract LLM calls that expose prompt/response token IDs from either:
+       - litellm_request         (sometimes only metadata, ignore if no token ids)
+       - raw_gen_ai_request      (llm.hosted_vllm.* stringified fields)
+    3) Extract rewards from spans whose attributes contain an AgentOps-style
+       reward payload or explicit REWARD span.
+    4) For each reward with sequence R, assign it to the most recent *unmatched* LLM call
+       with sequence < R. Ignore timestamps completely.
+    """
+
+    def _literal_eval_maybe(self, v: Any) -> Any:
+        import ast
+
+        if isinstance(v, str):
+            try:
+                return ast.literal_eval(v)
+            except Exception:
+                return v
+        return v
+
+    def _extract_tokens_from_raw(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int]]:
+        """Extract token ids from raw_gen_ai_request attributes.
+
+        - llm.hosted_vllm.prompt_token_ids: string -> List[int]
+        - llm.hosted_vllm.response_token_ids: string -> List[List[int]] -> take first
+        - llm.hosted_vllm.choices: string -> [{'token_ids': [...]}] -> take first
+        """
+        prompt_ids: List[int] = []
+        resp_ids: List[int] = []
+
+        # prompt
+        p = attrs.get("llm.hosted_vllm.prompt_token_ids")
+        p = self._literal_eval_maybe(p)
+        if isinstance(p, list) and all(isinstance(x, int) for x in p):  # type: ignore
+            prompt_ids = cast(List[int], p)
+
+        # response preferred path
+        r = attrs.get("llm.hosted_vllm.response_token_ids")
+        r = self._literal_eval_maybe(r)
+        if isinstance(r, list) and len(r) > 0 and isinstance(r[0], list):  # type: ignore
+            first = cast(List[Any], r[0])
+            if all(isinstance(x, int) for x in first):
+                resp_ids = cast(List[int], first)
+
+        # fallback via choices
+        if not resp_ids:
+            choices = attrs.get("llm.hosted_vllm.choices")
+            choices = self._literal_eval_maybe(choices)
+            if isinstance(choices, list) and choices:
+                cand = cast(Any, choices[0])
+                if isinstance(cand, dict):
+                    tids = cast(Dict[str, Any], cand).get("token_ids")
+                    if isinstance(tids, list) and all(isinstance(x, int) for x in tids):  # type: ignore
+                        resp_ids = cast(List[int], tids)
+
+        return prompt_ids, resp_ids
+
+    def _extract_tokens_from_openai(self, attrs: Dict[str, Any]) -> Tuple[List[int], List[int]]:
+        prompt_ids = cast(Any, attrs.get("prompt_token_ids") or [])
+        resp_ids = cast(Any, attrs.get("response_token_ids") or [])
+        prompt_ids = self._literal_eval_maybe(prompt_ids)
+        resp_ids = self._literal_eval_maybe(resp_ids)
+        if not (isinstance(prompt_ids, list) and all(isinstance(x, int) for x in prompt_ids)):  # type: ignore
+            prompt_ids = []
+        if not (isinstance(resp_ids, list) and all(isinstance(x, int) for x in resp_ids)):  # type: ignore
+            resp_ids = []
+        return cast(List[int], prompt_ids), cast(List[int], resp_ids)
+
+    def _maybe_reward_value(self, span: Span) -> Optional[float]:
+        """
+        Parse reward from typical AgentOps payload or explicit REWARD span.
+        """
+        attrs = span.attributes or {}
+
+        # AgentOps new/old keys
+        for k in ("agentops.task.output", "agentops.entity.output"):
+            v = attrs.get(k)
+            v = self._literal_eval_maybe(v)
+            if isinstance(v, dict) and cast(Dict[str, Any], v).get("type") == "reward":
+                rv = cast(Dict[str, Any], v).get("value", None)
+                if rv is None or isinstance(rv, (int, float)):
+                    return None if rv is None else float(rv)
+
+        # Explicit reward span
+        if span.name == SpanNames.REWARD.value:
+            rv = attrs.get("reward", None)
+            if rv is None or isinstance(rv, (int, float)):
+                return None if rv is None else float(rv)
+
+        return None
+
+    def _request_id_from_attrs(self, attrs: Dict[str, Any]) -> Optional[str]:
+        # Prefer OpenAI-like id if present, else proxy raw id.
+        rid = attrs.get("gen_ai.response.id") or attrs.get("llm.hosted_vllm.id")
+        return str(rid) if isinstance(rid, str) and rid else None
+
+    def adapt(self, source: List[Span], /) -> List[Triplet]:  # type: ignore
+        # 1) Sort deterministically by (sequence_id, start_time).
+        spans = sorted(
+            source,
+            key=lambda s: (s.sequence_id, s.start_time),
+        )
+
+        # 2) Collect LLM calls with token IDs.
+        llm_items: List[Dict[str, Any]] = []
+        seen_request_ids: set[str] = set()
+        for s in spans:
+            attrs = s.attributes or {}
+            prompt_ids: List[int] = []
+            resp_ids: List[int] = []
+
+            if s.name == "raw_gen_ai_request":
+                prompt_ids, resp_ids = self._extract_tokens_from_raw(attrs)
+            elif s.name == "litellm_request":
+                # Some proxies never include token ids here. Ignore unless present.
+                prompt_ids, resp_ids = self._extract_tokens_from_openai(attrs)
+
+            if prompt_ids and resp_ids:
+                rid = self._request_id_from_attrs(attrs)
+                if rid:
+                    # Duplicated request ID. This request is already handled.
+                    if rid in seen_request_ids:
+                        continue
+                    seen_request_ids.add(rid)
+                llm_items.append(
+                    dict(
+                        span=s,
+                        seq=s.sequence_id,
+                        response_ids=resp_ids,
+                        prompt_ids=prompt_ids,
+                        request_id=rid,
+                    )
+                )
+
+        # Order LLM items by sequence only.
+        llm_items.sort(key=lambda x: x["seq"])
+
+        # Collect rewards by sequence only.
+        rewards: List[Tuple[int, Optional[float]]] = []
+        for s in spans:
+            val = self._maybe_reward_value(s)
+            if val is not None:
+                rewards.append((s.sequence_id, val))
+
+        # First-occurrence matching by sequence_id only:
+        # For reward at sequence R, assign to the most recent unmatched LLM with seq < R.
+        assigned: Dict[str, Optional[float]] = {}
+        for r_seq, r_val in sorted(rewards, key=lambda x: x[0]):
+            for item in reversed(llm_items):
+                sid = item["span"].span_id
+                if sid in assigned:
+                    continue
+                if item["seq"] < r_seq:
+                    assigned[sid] = r_val
+                    break
+
+        # Build triplets in LLM sequence order.
+        triplets: List[Triplet] = []
+        for item in llm_items:
+            s = item["span"]
+            triplets.append(
+                Triplet(
+                    prompt={"token_ids": item["prompt_ids"]},
+                    response={"token_ids": item["response_ids"]},
+                    reward=assigned.get(s.span_id, None),
+                    metadata=dict(
+                        # This is called response_id to align with the other adapters.
+                        response_id=item["request_id"],
+                    ),
+                )
+            )
+
+        return triplets

agentlightning/algorithm/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import BaseAlgorithm
+from .decorator import algo
+from .fast import Baseline, FastAlgorithm
+
+if TYPE_CHECKING:
+    from .apo import APO as APOType
+    from .verl import VERL as VERLType
+
+__all__ = ["BaseAlgorithm", "algo", "FastAlgorithm", "Baseline", "APO", "VERL"]
+
+# Shortcuts for usages like algo.APO(...)
+
+
+def APO(*args: Any, **kwargs: Any) -> APOType[Any]:
+    from .apo import APO as APOImplementation
+
+    return APOImplementation(*args, **kwargs)
+
+
+def VERL(*args: Any, **kwargs: Any) -> VERLType:
+    from .verl import VERL as VERLImplementation
+
+    return VERLImplementation(*args, **kwargs)

agentlightning/algorithm/apo/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .apo import APO
+
+__all__ = ["APO"]

agentlightning/algorithm/apo/apo.py

@@ -0,0 +1,891 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+APO with textual gradients that read rollout spans and outputs to modify the prompt.
+
+- algo: beam search with span-aware textual gradients -> apply_edit via LLM
+- rollout: same pattern as your example, but task is a dict (T_task)
+"""
+
+import asyncio
+import logging
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Counter, Dict, Generic, Iterator, List, Optional, Sequence, Set, Tuple, TypedDict, TypeVar, cast
+
+import poml
+from openai import AsyncOpenAI
+
+from agentlightning.adapter.messages import TraceToMessages
+from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.reward import find_final_reward
+from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, RolloutStatus
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+
+class RolloutResultForAPO(TypedDict):
+    """This must be all JSON serializable to be processable by POML."""
+
+    status: RolloutStatus
+    final_reward: Optional[float]
+    spans: List[Dict[str, Any]]
+    messages: List[Any]
+
+
+@dataclass
+class VersionedPromptTemplate:
+    version: str
+    prompt_template: PromptTemplate
+    score: Optional[float] = None
+
+
+GRADIENT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "text_gradient_variant01.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant02.poml",
+    Path(__file__).parent / "prompts" / "text_gradient_variant03.poml",
+]
+
+APPLY_EDIT_PROMPT_FILES = [
+    Path(__file__).parent / "prompts" / "apply_edit_variant01.poml",
+    Path(__file__).parent / "prompts" / "apply_edit_variant02.poml",
+]
+
+
+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(BaseAlgorithm, Generic[T_task]):
+    """Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search.
+
+    APO is an iterative prompt optimization algorithm that uses LLM-generated textual gradients
+    to improve prompts through a beam search process. It evaluates prompts on rollouts,
+    computes critiques based on the results, and applies edits to generate improved prompts.
+
+    The algorithm operates in rounds, where each round:
+    1. Samples parent prompts from the current beam
+    2. Generates new prompts by computing textual gradients and applying edits
+    3. Evaluates all candidates on a validation set
+    4. Selects the top-k prompts for the next round
+
+    Based on the ideas from:
+    - ProTeGi: https://aclanthology.org/2023.emnlp-main.494.pdf
+    - TextGrad: https://github.com/zou-group/textgrad
+    """
+
+    def __init__(
+        self,
+        async_openai_client: AsyncOpenAI,
+        *,
+        gradient_model: str = "gpt-5-mini",
+        apply_edit_model: str = "gpt-4.1-mini",
+        diversity_temperature: float = 1.0,
+        gradient_batch_size: int = 4,
+        val_batch_size: int = 16,
+        beam_width: int = 4,
+        branch_factor: int = 4,
+        beam_rounds: int = 3,
+        rollout_batch_timeout: float = 3600.0,
+        run_initial_validation: bool = True,
+        # Internal flags for debugging
+        _poml_trace: bool = False,
+    ):
+        """
+        Initialize the APO algorithm with configuration parameters.
+
+        Args:
+            async_openai_client: AsyncOpenAI client for making LLM API calls.
+            gradient_model: Model name for computing textual gradients (critiques).
+            apply_edit_model: Model name for applying edits based on critiques.
+            diversity_temperature: Temperature parameter for LLM calls to control diversity.
+            gradient_batch_size: Number of rollout results to sample for gradient computation.
+            val_batch_size: Number of validation examples to use for evaluation.
+            beam_width: Number of top-scoring prompts to keep in the beam at each round.
+            branch_factor: Number of new prompt candidates to generate from each parent prompt
+                by applying textual gradient edits. This controls the expansion of the search tree.
+            beam_rounds: Number of beam search rounds to perform.
+            rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
+            run_initial_validation: If True, runs validation on the seed prompt before starting
+                optimization to establish a baseline score. Defaults to True.
+        """
+        self.async_openai_client = async_openai_client
+        self.gradient_model = gradient_model
+        self.apply_edit_model = apply_edit_model
+        self.diversity_temperature = diversity_temperature
+        self.gradient_batch_size = gradient_batch_size
+        self.val_batch_size = val_batch_size
+        self.beam_width = beam_width
+        self.branch_factor = branch_factor
+        self.beam_rounds = beam_rounds
+        self.rollout_batch_timeout = rollout_batch_timeout
+        self.run_initial_validation = run_initial_validation
+
+        self._history_best_prompt: Optional[PromptTemplate] = None
+        self._history_best_score: float = float("-inf")
+        self._history_best_version: Optional[str] = None
+
+        self._version_counter: int = 0
+
+        self._poml_trace = _poml_trace
+
+    def _create_versioned_prompt(
+        self,
+        prompt_template: PromptTemplate,
+        *,
+        score: Optional[float] = None,
+    ) -> VersionedPromptTemplate:
+        """
+        Wrap a prompt template with a new monotonically increasing version identifier.
+        """
+        version = f"v{self._version_counter}"
+        self._version_counter += 1
+        return VersionedPromptTemplate(version=version, prompt_template=prompt_template, score=score)
+
+    def _format_log_prefix(
+        self,
+        *,
+        round_num: Optional[int] = None,
+        beam_idx: Optional[int] = None,
+        branch_idx: Optional[int] = None,
+        prompt_version: Optional[str] = None,
+    ) -> str:
+        """
+        Construct the standardized log prefix.
+        """
+        parts: List[str] = []
+        if round_num is not None:
+            parts.append(f"Round {round_num:02d}")
+        if beam_idx is not None:
+            parts.append(f"Beam {beam_idx:02d}")
+        if branch_idx is not None:
+            parts.append(f"Branch {branch_idx:02d}")
+        if prompt_version is not None:
+            parts.append(f"Prompt {prompt_version}")
+        if not parts:
+            return ""
+        return f"[{' | '.join(parts)}]"
+
+    def _log(self, level: int, message: str, *, prefix: Optional[str] = None) -> None:
+        """
+        Log a message with an optional standardized prefix.
+        """
+        effective_prefix = prefix
+        if effective_prefix:
+            logger.log(level, f"{effective_prefix} {message}")
+        else:
+            logger.log(level, message)
+
+    def get_seed_prompt_template(self) -> Tuple[str, PromptTemplate]:
+        """
+        Extract the initial prompt template from the algorithm's resources.
+
+        Returns:
+            A tuple of (resource_name, prompt_template) representing the seed prompt.
+
+        Raises:
+            ValueError: If initial_resources is not set or no PromptTemplate is found.
+        """
+        initial_resources = self.get_initial_resources()
+        if initial_resources is None:
+            raise ValueError(
+                "initial_resources are not set for APO algorithm. "
+                "Use algorithm.set_initial_resources() to set initial resources or set it in Trainer()"
+            )
+        for name, resource in initial_resources.items():
+            if isinstance(resource, PromptTemplate):
+                return name, resource
+        raise ValueError("No prompt template resource found in initial_resources")
+
+    def get_adapter(self) -> TraceToMessages:
+        """
+        Get the adapter for converting spans to messages.
+
+        Returns:
+            The TraceToMessages instance for this algorithm.
+
+        Raises:
+            ValueError: If the adapter is not a TraceToMessages.
+        """
+        adapter = super().get_adapter()
+        if not isinstance(adapter, TraceToMessages):
+            raise ValueError("Adapter must be a TraceToMessages for APO algorithm")
+        return adapter
+
+    def get_best_prompt(self) -> PromptTemplate:
+        """
+        Retrieve the best prompt discovered during optimization.
+
+        Returns:
+            The prompt template with the highest validation score found so far.
+
+        Raises:
+            ValueError: If no best prompt has been found yet (run() not called).
+        """
+        if self._history_best_prompt is None:
+            raise ValueError("No best prompt found")
+        return self._history_best_prompt
+
+    async def compute_textual_gradient(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout_results: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Compute a textual gradient (critique) for the current prompt based on rollout results.
+
+        This method samples rollout results, sends them to an LLM along with the current prompt,
+        and generates a critique describing how the prompt could be improved.
+
+        Args:
+            current_prompt: The prompt template to critique.
+            rollout_results: List of rollout results containing spans, messages, and rewards.
+
+        Returns:
+            A textual critique generated by the LLM, or None if generation fails.
+        """
+        tg_template = random.choice(GRADIENT_PROMPT_FILES)
+
+        if len(rollout_results) < self.gradient_batch_size:
+            self._log(
+                logging.WARNING,
+                f"Only {len(rollout_results)} rollouts available, but {self.gradient_batch_size} are needed. Using all rollouts.",
+                prefix=prefix,
+            )
+            sampled_rollout_results = rollout_results
+        else:
+            sampled_rollout_results = random.sample(rollout_results, self.gradient_batch_size)
+
+        self._log(
+            logging.INFO,
+            f"Gradient will be computed with {self.gradient_model} for {len(sampled_rollout_results)} rollouts with template: {tg_template.name}",
+            prefix=prefix,
+        )
+
+        tg_msg = poml.poml(  # type: ignore
+            tg_template,
+            context={
+                "experiments": sampled_rollout_results,
+                "prompt_template": current_prompt.prompt_template.template,
+            },
+            format="openai_chat",
+        )
+        self._log(
+            logging.DEBUG,
+            f"Gradient computed with {self.gradient_model} prompt: {tg_msg}",
+            prefix=prefix,
+        )
+        critique_response = await self.async_openai_client.chat.completions.create(
+            model=self.gradient_model,
+            messages=tg_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        critique_text = critique_response.choices[0].message.content
+        self._log(
+            logging.INFO,
+            f"Gradient computed with {self.gradient_model} has result: {critique_text}",
+            prefix=prefix,
+        )
+
+        return critique_text
+
+    async def textual_gradient_and_apply_edit(
+        self,
+        current_prompt: VersionedPromptTemplate,
+        rollout: List[RolloutResultForAPO],
+        *,
+        prefix: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Generate an improved prompt by computing a textual gradient and applying an edit.
+
+        This is the main optimization step that:
+        1. Computes a critique (textual gradient) based on rollout performance
+        2. Uses another LLM to apply the critique and generate an improved prompt
+
+        Args:
+            current_prompt: The current prompt template to improve.
+            rollout: List of rollout results to base the critique on.
+
+        Returns:
+            The improved prompt text, or the original prompt if gradient computation fails.
+        """
+        # 1) Critique
+        critique_text = await self.compute_textual_gradient(
+            current_prompt,
+            rollout,
+            prefix=prefix,
+        )
+        if not critique_text:
+            self._log(
+                logging.ERROR,
+                "Failed to compute critique for prompt.",
+                prefix=prefix,
+            )
+            return current_prompt.prompt_template.template
+
+        # 2) Apply edit
+        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        self._log(
+            logging.INFO,
+            f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",
+            prefix=prefix,
+        )
+        ae_msg = poml.poml(  # type: ignore
+            ae_template,
+            context={
+                "prompt_template": current_prompt.prompt_template.template,
+                "critique": critique_text,
+            },
+            format="openai_chat",
+        )
+
+        ae_response = await self.async_openai_client.chat.completions.create(
+            model=self.apply_edit_model,
+            messages=ae_msg["messages"],  # type: ignore
+            temperature=self.diversity_temperature,
+        )
+        new_prompt = ae_response.choices[0].message.content
+        if new_prompt:
+            self._log(
+                logging.INFO,
+                f"Edit generated by {self.apply_edit_model}: {new_prompt[:50]}...",
+                prefix=prefix,
+            )
+        return new_prompt
+
+    async def get_rollout_results(
+        self,
+        rollout: List[Rollout],
+        *,
+        prefix: Optional[str] = None,
+    ) -> List[RolloutResultForAPO]:
+        """
+        Convert completed rollouts to APO-compatible result format.
+
+        Fetches spans for each rollout, adapts them to messages, and packages them
+        with rewards and status information for gradient computation.
+
+        Args:
+            rollout: List of completed rollout metadata.
+
+        Returns:
+            List of rollout results formatted for APO processing.
+        """
+        rollout_results: List[RolloutResultForAPO] = []
+        store = self.get_store()
+        adapter = self.get_adapter()
+        for r in rollout:
+            spans = await store.query_spans(r.rollout_id)
+            messages = adapter.adapt(spans)
+            rollout_result = RolloutResultForAPO(
+                status=r.status,
+                final_reward=find_final_reward(spans),
+                spans=[span.model_dump() for span in spans],
+                messages=messages,
+            )
+            self._log(
+                logging.DEBUG,
+                f"Rollout result for {r.rollout_id}: status {rollout_result['status']} with final reward {rollout_result['final_reward']}. "
+                f"{len(rollout_result['spans'])} spans and {len(rollout_result['messages'])} messages.",
+                prefix=prefix,
+            )
+            rollout_results.append(rollout_result)
+        return rollout_results
+
+    async def evaluate_prompt_on_batch(
+        self,
+        prompt: VersionedPromptTemplate,
+        resource_name: str,
+        dataset: Sequence[T_task],
+        mode: RolloutMode,
+        *,
+        prefix: Optional[str] = None,
+    ) -> Tuple[List[RolloutResultForAPO], float]:
+        """
+        Evaluate a prompt on a batch of tasks by running rollouts and computing average reward.
+
+        This method:
+        1. Adds the prompt as a named resource to the store
+        2. Enqueues rollouts for each task in the dataset
+        3. Waits for rollouts to complete (with timeout)
+        4. Computes and returns the average reward
+
+        Args:
+            prompt: The prompt template string to evaluate.
+            resource_name: The name to register the prompt under in the store.
+            dataset: Sequence of tasks to evaluate the prompt on.
+            mode: Rollout mode ("train" or "val") for logging/tracking.
+
+        Returns:
+            A tuple of (rollout_results, average_reward) where rollout_results contains
+            detailed information for each rollout and average_reward is the mean final reward.
+        """
+        store = self.get_store()
+        preview = prompt.prompt_template.template[:50]
+        self._log(
+            logging.INFO,
+            f'Evaluating prompt "{preview}..." on {len(dataset)} tasks in {mode} mode',
+            prefix=prefix,
+        )
+
+        # Install prompt as named resource
+        resources: NamedResources = {resource_name: prompt.prompt_template}
+        resource_update = await store.update_resources(prompt.version, resources)
+
+        rollout_ids: List[str] = []
+        for t in dataset:
+            r = await store.enqueue_rollout(input=t, mode=mode, resources_id=resource_update.resources_id)
+            rollout_ids.append(r.rollout_id)
+
+        deadline = time.time() + self.rollout_batch_timeout
+        finished: List[Rollout] = []
+        while time.time() < deadline:
+            finished = await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)
+            if len(finished) >= len(rollout_ids):
+                self._log(
+                    logging.INFO,
+                    f"All {len(rollout_ids)} rollouts finished within timeout.",
+                    prefix=prefix,
+                )
+                break
+            else:
+                self._log(
+                    logging.DEBUG,
+                    f"Only {len(finished)} rollouts finished within timeout. Waiting for remaining {len(rollout_ids) - len(finished)} rollouts.",
+                    prefix=prefix,
+                )
+                # Sleep to avoid busy-waiting
+                await asyncio.sleep(2.0)
+
+        rollout_results = await self.get_rollout_results(
+            finished,
+            prefix=prefix,
+        )
+        final_rewards = [rr["final_reward"] for rr in rollout_results]
+
+        avg = float(sum([r or 0.0 for r in final_rewards]) / max(1, len(final_rewards)))
+        status_counter = Counter([rr["status"] for rr in rollout_results])
+
+        self._log(
+            logging.INFO,
+            f"Evaluated {len(rollout_results)} rollouts. Statuses: {status_counter}. Rewards: {final_rewards}, average is {avg}",
+            prefix=prefix,
+        )
+        return rollout_results, avg
+
+    def _initialize_beam(
+        self,
+        train_dataset: Optional[Dataset[T_task]],
+        val_dataset: Optional[Dataset[T_task]],
+    ) -> Tuple[str, PromptTemplate, Iterator[Sequence[T_task]], Iterator[Sequence[T_task]]]:
+        """
+        Initialize the beam search with seed prompt and dataset iterators.
+
+        Args:
+            train_dataset: Dataset for computing gradients.
+            val_dataset: Dataset for evaluating prompts.
+
+        Returns:
+            Tuple of (resource_name, seed_prompt, grad_iterator, val_iterator).
+
+        Raises:
+            ValueError: If either dataset is None.
+        """
+        resource_name, seed_prompt = self.get_seed_prompt_template()
+
+        if train_dataset is None:
+            raise ValueError("train_dataset is required for APO algorithm")
+        if val_dataset is None:
+            raise ValueError("val_dataset is required for APO algorithm")
+
+        grad_dataset_iterator = batch_iter_over_dataset(train_dataset, self.gradient_batch_size)
+        val_dataset_iterator = batch_iter_over_dataset(val_dataset, self.val_batch_size)
+
+        # Initialize history tracking
+        self._history_best_prompt = seed_prompt
+        self._history_best_score = float("-inf")
+
+        return resource_name, seed_prompt, grad_dataset_iterator, val_dataset_iterator
+
+    def _sample_parent_prompts(
+        self,
+        beam: List[VersionedPromptTemplate],
+        round_num: int,
+    ) -> List[Tuple[int, VersionedPromptTemplate]]:
+        """
+        Sample parent prompts from the current beam for generating new candidates.
+
+        If the beam has fewer prompts than beam_width, replicates existing prompts.
+        Otherwise, randomly samples beam_width prompts.
+
+        Args:
+            beam: Current list of prompt templates in the beam.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of parent prompts to generate children from.
+        """
+        display_round = round_num + 1
+        if len(beam) < self.beam_width:
+            prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.WARNING,
+                f"Beam width is currently {self.beam_width}, but only {len(beam)} prompts in beam. Replicating all prompts.",
+                prefix=prefix,
+            )
+            return [(i % len(beam), beam[i % len(beam)]) for i in range(self.beam_width)]
+
+        selected_indices = random.sample(range(len(beam)), self.beam_width)
+        return [(idx, beam[idx]) for idx in selected_indices]
+
+    async def _generate_candidate_prompts(
+        self,
+        parent_prompts: List[Tuple[int, VersionedPromptTemplate]],
+        resource_name: str,
+        grad_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Generate new candidate prompts from parents using textual gradients.
+
+        For each parent prompt, generates branch_factor new candidates by:
+        1. Evaluating the parent on a training batch
+        2. Computing textual gradient
+        3. Applying edit to generate improved prompt
+
+        Args:
+            parent_prompts: List of parent prompts to generate children from.
+            resource_name: Name to register prompts under in the store.
+            grad_dataset_iterator: Iterator over training data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of newly generated prompt templates.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Applying {self.branch_factor} edits to each of the {len(parent_prompts)} parents based on "
+            "gradients computed on training dataset",
+            prefix=round_prefix,
+        )
+
+        parent_prompts_str = [
+            f"{p.version}:{p.score:.3f}" if p.score is not None else p.version for _, p in parent_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Parent prompts: {', '.join(parent_prompts_str)}",
+            prefix=round_prefix,
+        )
+
+        candidates: List[VersionedPromptTemplate] = []
+        used_beam_indices: Set[int] = set()
+        for real_beam_idx, (beam_idx, prompt) in enumerate(parent_prompts):
+            if beam_idx in used_beam_indices:
+                beam_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                self._log(
+                    logging.WARNING,
+                    "Duplicated beam index found. Might be caused by beam_width too high. "
+                    + f"The real index of this beam is {real_beam_idx + 1}.",
+                    prefix=beam_prefix,
+                )
+            else:
+                used_beam_indices.add(beam_idx)
+            for branch_idx in range(self.branch_factor):
+                parent_prefix = self._format_log_prefix(
+                    round_num=display_round,
+                    beam_idx=beam_idx + 1,
+                    branch_idx=branch_idx + 1,
+                    prompt_version=prompt.version,
+                )
+                baseline_score = f"{prompt.score:.3f}" if prompt.score is not None else "N/A"
+                self._log(
+                    logging.INFO,
+                    f"Use parent prompt {prompt.version} as a baseline to generate a new prompt. Baseline score: {baseline_score}",
+                    prefix=parent_prefix,
+                )
+                grad_samples = next(grad_dataset_iterator)
+                rollout_results, _ = await self.evaluate_prompt_on_batch(
+                    prompt,
+                    resource_name,
+                    grad_samples,
+                    mode="train",
+                    prefix=parent_prefix,
+                )
+                new_prompt = await self.textual_gradient_and_apply_edit(
+                    prompt,
+                    rollout_results,
+                    prefix=parent_prefix,
+                )
+                if not new_prompt:
+                    self._log(
+                        logging.ERROR,
+                        f"Failed to compute edit for prompt: {prompt.prompt_template.template}",
+                        prefix=parent_prefix,
+                    )
+                    continue
+                new_prompt_template = PromptTemplate(template=new_prompt, engine="f-string")
+                versioned_candidate = self._create_versioned_prompt(new_prompt_template)
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}: {versioned_candidate.version}",
+                    prefix=parent_prefix,
+                )
+                candidate_prefix = self._format_log_prefix(
+                    round_num=display_round, prompt_version=versioned_candidate.version
+                )
+                self._log(
+                    logging.INFO,
+                    f"New prompt template created from parent {prompt.version}:\n```\n{new_prompt}\n```",
+                    prefix=candidate_prefix,
+                )
+                candidates.append(versioned_candidate)
+
+        return candidates
+
+    async def _evaluate_and_select_beam(
+        self,
+        candidates: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset_iterator: Iterator[Sequence[T_task]],
+        round_num: int,
+    ) -> List[VersionedPromptTemplate]:
+        """
+        Evaluate all candidate prompts on validation data and select top-k for the beam.
+
+        Args:
+            candidates: List of candidate prompts to evaluate.
+            resource_name: Name to register prompts under in the store.
+            val_dataset_iterator: Iterator over validation data batches.
+            round_num: Current round number (for logging, 0-indexed).
+
+        Returns:
+            List of top beam_width prompts sorted by validation score (best first).
+
+        Raises:
+            ValueError: If no candidates remain after evaluation.
+        """
+        display_round = round_num + 1
+        round_prefix = self._format_log_prefix(round_num=display_round)
+        self._log(
+            logging.INFO,
+            f"Evaluating {len(candidates)} candidates on validation dataset",
+            prefix=round_prefix,
+        )
+
+        val_batch = next(val_dataset_iterator)
+
+        for prompt in candidates:
+            candidate_prefix = self._format_log_prefix(
+                round_num=display_round,
+                prompt_version=prompt.version,
+            )
+            _, score = await self.evaluate_prompt_on_batch(
+                prompt,
+                resource_name,
+                val_batch,
+                mode="val",
+                prefix=candidate_prefix,
+            )
+            prompt.score = score
+            self._log(
+                logging.INFO,
+                f"Candidate score: {score:.3f}",
+                prefix=candidate_prefix,
+            )
+
+        # Sort by score (descending) and select top beam_width
+        sorted_prompts = [p for p in sorted(candidates, key=lambda x: cast(float, x.score), reverse=True)]
+        selected_prompts = sorted_prompts[: self.beam_width]
+        selected_versions = [
+            f"{prompt.version}:{prompt.score:.3f}" if prompt.score is not None else prompt.version
+            for prompt in selected_prompts
+        ]
+        self._log(
+            logging.INFO,
+            f"Top {len(selected_prompts)} candidates on validation dataset: {selected_versions}",
+            prefix=round_prefix,
+        )
+
+        if len(selected_prompts) == 0:
+            raise ValueError("No beam candidates any more")
+
+        return selected_prompts
+
+    async def _update_best_prompt(
+        self,
+        beam: List[VersionedPromptTemplate],
+        resource_name: str,
+        val_dataset: Dataset[T_task],
+        round_num: int,
+    ) -> None:
+        """
+        Evaluate the best prompt in the beam on the full validation set and update history.
+
+        Args:
+            beam: Current beam of prompts (sorted, best first).
+            resource_name: Name to register prompts under in the store.
+            val_dataset: Full validation dataset.
+            round_num: Current round number (for logging, 0-indexed).
+        """
+        display_round = round_num + 1
+        best_prompt = beam[0]
+        prefix = self._format_log_prefix(round_num=display_round, prompt_version=best_prompt.version)
+        _, best_score = await self.evaluate_prompt_on_batch(
+            best_prompt,
+            resource_name,
+            cast(Sequence[T_task], val_dataset),
+            mode="val",
+            prefix=prefix,
+        )
+        self._log(
+            logging.INFO,
+            f"Beam leader score: {best_score:.3f}",
+            prefix=prefix,
+        )
+
+        if best_score > self._history_best_score:
+            prev = self._history_best_score
+            self._log(
+                logging.INFO,
+                f"Best prompt updated. New best score: {best_score:.3f} (prev: {prev:.3f})",
+                prefix=prefix,
+            )
+            self._history_best_prompt = best_prompt.prompt_template
+            self._history_best_score = best_score
+            self._history_best_version = best_prompt.version
+        else:
+            self._log(
+                logging.WARNING,
+                f"Best prompt not updated. Current score: {best_score:.3f} vs. history best: {self._history_best_score:.3f})",
+                prefix=prefix,
+            )
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[T_task]] = None,
+        val_dataset: Optional[Dataset[T_task]] = None,
+    ) -> None:
+        """
+        Execute the APO algorithm to optimize prompts through beam search with textual gradients.
+
+        The algorithm performs iterative prompt optimization over multiple rounds:
+        - Each round: samples parent prompts, generates new candidates via textual gradients,
+          evaluates all candidates on validation data, and keeps the top performers
+        - Tracks the historically best prompt across all rounds
+        - Uses different training data samples for each gradient computation to ensure diversity
+
+        Args:
+            train_dataset: Dataset of tasks for computing textual gradients. Required.
+            val_dataset: Dataset of tasks for evaluating and selecting prompts. Required.
+
+        Raises:
+            ValueError: If train_dataset or val_dataset is None, or if resources are not set.
+        """
+        # Initialize beam search
+        resource_name, seed_prompt, grad_iterator, val_iterator = self._initialize_beam(train_dataset, val_dataset)
+
+        if self._poml_trace:
+            poml.set_trace(trace_dir="pomltrace")
+
+        # Validation datasets are guaranteed to be non-None after initialization
+        assert val_dataset is not None
+
+        # Start with seed prompt in the beam
+        seed_versioned = self._create_versioned_prompt(seed_prompt)
+        beam: List[VersionedPromptTemplate] = [seed_versioned]
+        self._history_best_prompt = seed_prompt
+        self._history_best_version = seed_versioned.version
+
+        # Optionally evaluate seed prompt on validation set to establish baseline
+        if self.run_initial_validation:
+            seed_prefix = self._format_log_prefix(round_num=0, prompt_version=seed_versioned.version)
+            self._log(
+                logging.INFO,
+                "Evaluating seed prompt on validation dataset before optimization...",
+                prefix=seed_prefix,
+            )
+            _, seed_score = await self.evaluate_prompt_on_batch(
+                seed_versioned,
+                resource_name,
+                cast(Sequence[T_task], val_dataset),
+                mode="val",
+                prefix=seed_prefix,
+            )
+            self._log(
+                logging.INFO,
+                f"Seed prompt baseline score: {seed_score:.3f}",
+                prefix=seed_prefix,
+            )
+            self._history_best_prompt = seed_prompt
+            self._history_best_score = seed_score
+            self._history_best_version = seed_versioned.version
+
+        # Run beam search for specified number of rounds
+        for rnd in range(self.beam_rounds):
+            display_round = rnd + 1
+            round_prefix = self._format_log_prefix(round_num=display_round)
+            self._log(
+                logging.INFO,
+                f"Round {display_round}/{self.beam_rounds}...",
+                prefix=round_prefix,
+            )
+
+            # Sample parent prompts from current beam
+            parent_prompts = self._sample_parent_prompts(beam, rnd)
+
+            # Generate new candidate prompts from parents
+            new_candidates = await self._generate_candidate_prompts(parent_prompts, resource_name, grad_iterator, rnd)
+
+            # Combine existing beam with new candidates
+            all_candidates = [*beam, *new_candidates]
+
+            # Evaluate and select top-k prompts for next beam
+            beam = await self._evaluate_and_select_beam(all_candidates, resource_name, val_iterator, rnd)
+
+            # Update historically best prompt if improved
+            await self._update_best_prompt(beam, resource_name, val_dataset, rnd)

agentlightning/algorithm/apo/prompts/apply_edit_variant01.poml

@@ -0,0 +1,22 @@
+<poml>
+  <p>Revise the given prompt template using the critique as constraints and improvement guide.</p>
+  <cp caption="Revision Rules">
+    <list listStyle="decimal">
+      <item>Rewrite or restructure the prompt if critique implies it.</item>
+      <item>Explicitly include any requested output format, structure, or word limit, if requested by the critique.</item>
+      <item>Prioritize mechanism-first phrasing: define what to do, then how to do it.</item>
+      <item>Preserve placeholder variables inside curly brackets.</item>
+    </list>
+  </cp>
+  <output-format>
+    Return only the improved prompt template with placeholders intact. Do not include other explanations on how you did it, or headers and introductory texts.
+  </output-format>
+  <human-msg>
+    <cp caption="Prompt Template">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Critique">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/apo/prompts/apply_edit_variant02.poml

@@ -0,0 +1,18 @@
+<!-- Conservative Edit Prompt -->
+
+<poml>
+  <p>Revise the prompt to address ONE critique point clearly and effectively. Preserve all variable names in curly-brackets.</p>
+  <p>Do not address more than one critique point. Focus on the single most critical issue.</p>
+  <p>Keep the new prompt close in tone, length, and structure to the original.</p>
+  <output-format>
+    Return only the revised full prompt. Do not include explanations, comparisons, or other text.
+  </output-format>
+  <human-msg>
+    <cp caption="PROMPT" level="3">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="CRITIQUE" level="3">
+      <text whiteSpace="pre">{{ critique }}</text>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant01.poml

@@ -0,0 +1,18 @@
+<poml>
+  <p>You optimize a prompt template.</p>
+  <cp caption="Original Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Experiments with Original Prompt Template">
+    <cp for="experiment in experiments" caption="Experiment {{ loop.index + 1 }}">
+      <p>This experiment has {{ experiment.status }}. It gets a final reward: {{ experiment.final_reward }}</p>
+      <cp caption="Rollout Traces (Chat Messages, Grader Requests included)">
+        <object data="{{ experiment.messages }}" />
+      </cp>
+    </cp>
+  </cp>
+  <cp caption="Your Task">
+    Produce a brief critique listing specific causes for the error or ways to raise reward next time.
+    Return a bullet list with concrete, testable changes (format, constraints, ordering, definitions).
+  </cp>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant02.poml

@@ -0,0 +1,16 @@
+<poml>
+  <role>You are a prompt engineer.</role>
+  <task>Analyze where the current prompt failed to elicit the right mechanism.</task>
+  <cp caption="Current Prompt Template">
+    <text whiteSpace="pre">{{ prompt_template }}</text>
+  </cp>
+  <cp caption="Sample Runs with Current Prompt Template">
+    <p>The following are the OpenTelemetry spans collected from the sample runs with the current prompt template. They should contain both prompt, responses and rewards.</p>
+    <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }} Diagnostics">
+      <object for="span in experiment.spans" data="{{ span }}" />
+    </cp>
+  </cp>
+  <output-format>
+    Write 3-5 short bullets titled 'Critique:' focusing on missing constraints, ordering, or formatting.
+  </output-format>
+</poml>

agentlightning/algorithm/apo/prompts/text_gradient_variant03.poml

@@ -0,0 +1,107 @@
+<poml>
+
+  <role>You are an expert prompt engineer.</role>
+
+  <task>Your task is to analyze the prompt and provide a critique of the prompt. Follow the steps below to create the critique.
+
+  <cp caption="1. Structural Issues">
+    <p>These flaws block clarity and logic. Always check them first.</p>
+
+    <list>
+      <item><b>Missing goal</b>: The prompt never defines what success looks like. Ask: <i>Can I summarize its output goal in one line?</i></item>
+      <item><b>Contradictions</b>: Two or more instructions conflict. Search for words like *never*, *always*, *except*, *but also*.</item>
+      <item><b>Circular dependencies</b>: The model is told to do A before B and B before A.</item>
+      <item><b>No stop condition</b>: The prompt doesn’t say when the task is done. Flag any open-ended verbs: <i>explore,</i> <i>analyze further,</i> <i>continue indefinitely.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="2. Instruction Quality">
+    <p>Examine how the instructions are stated and ordered to ensure clarity and enforceability.</p>
+    <list>
+      <item><b>Vague verbs</b>: Avoid terms like <i>optimize,</i> <i>improve,</i> and <i>ensure.</i> Use precise, measurable instructions.</item>
+      <item><b>Lack of hierarchy</b>: All rules appear equally important, making conflict resolution impossible. Clarify rule precedence.</item>
+      <item><b>Mixed abstraction</b>: High-level policies are interleaved with implementation details. Keep principles separate from step-by-step actions.</item>
+      <item><b>Overlapping scope</b>: Similar instructions appear in several sections with minor changes. Identify and consolidate duplicates.</item>
+    </list>
+  </cp>
+
+  <cp caption="3. Control and Behavior">
+    <p>Review boundaries on model autonomy, tool use, and communication style.</p>
+    <list>
+      <item><b>No tool limits</b>: Limits on tool calls, retries, or time not specified. Define boundaries for operations.</item>
+      <item><b>Unclear uncertainty handling</b>: Conflicting instructions regarding clarifying uncertainties vs. never asking users. Select one behavior.</item>
+      <item><b>Verbosity confusion</b>: Some parts demand detailed answers, others specify brevity. Highlight and resolve inconsistency.</item>
+      <item><b>Feedback omission</b>: No plan for progress reporting or preamble during multi-step operations.</item>
+    </list>
+  </cp>
+
+  <cp caption="4. Input and Output Specification">
+    <p>Assess if required data and expected output formats are clearly defined.</p>
+    <list>
+      <item><b>No input defaults</b>: What should happen if a needed value is absent or invalid isn’t explained.</item>
+      <item><b>Output schema missing</b>: Expected response format or sections are not spelled out.</item>
+      <item><b>Format inconsistency</b>: Output style (Markdown, JSON, XML, etc.) shifts mid-prompt. Ensure format requirements are stable.</item>
+      <item><b>No validation</b>: Lacks steps like <i>verify results before submitting</i> or <i>summarize at end.</i></item>
+    </list>
+  </cp>
+
+  <cp caption="5. Scope and Safety">
+    <p>Ensure prompt actions remain within safe, authorized boundaries.</p>
+    <list>
+      <item><b>Scope creep</b>: Open-ended statements such as <i>feel free to enhance</i> can justify unrelated changes.</item>
+      <item><b>Unsafe actions</b>: Allows deletions or modifications without explicit user approval.</item>
+      <item><b>No error handling</b>: What happens if a tool call fails or data is missing is not addressed.</item>
+      <item><b>User authority ambiguity</b>: Model may act for multiple users or perform irreversible actions without checks.</item>
+    </list>
+  </cp>
+
+  <cp caption="6. Efficiency and Maintainability">
+    <p>Consider the prompt’s length, redundancy, and future comprehensibility.</p>
+    <list>
+      <item><b>Overexplained</b>: Verbose explanations where concise, numbered steps suffice.</item>
+      <item><b>Redundancy</b>: Similar rules scattered in multiple aliases; centralize and summarize them.</item>
+      <item><b>Hidden assumptions</b>: Implicit defaults (like timezone, language) are not stated.</item>
+      <item><b>Poor auditability</b>: Lacks section markers (e.g., <code>&lt;policy&gt;</code>, <code>&lt;procedure&gt;</code>). Structure prompt for easy review.</item>
+    </list>
+  </cp>
+
+  <cp caption="7. Testing Method">
+    <p>Methodical approach for reviewing a prompt:</p>
+    <list>
+      <item>Read the prompt fully; highlight all unclear or contradictory instructions.</item>
+      <item>For each main area, answer:
+      <list listStyle="decimal">
+        <item>What is the intended outcome?</item>
+        <item>What is the stop or completion condition?</item>
+        <item>How are conflicts between rules resolved?</item>
+        <item>What are the explicit limits (tools, run time, tokens)?</item>
+        <item>What should the output format be?</item>
+      </list>
+      </item>
+      <item>Rate each section: <i>clear</i>, <i>incomplete</i>, <i>contradictory</i>, or <i>redundant</i>.</item>
+      <item>Summarize findings under categories: structure, control, scope, format, safety.</item>
+    </list>
+    <p>This method surfaces issues such as ambiguity, contradiction, missing boundaries, and output uncertainty—core failure modes in prompting identified by the GPT-5 prompting guide.</p>
+  </cp>
+  </task>
+
+  <output-format>
+    Respond with a complete analysis and critique of the prompt. Be concise and direct. Less than 350 words.
+  </output-format>
+
+  <human-msg>
+    <cp caption="Prompt">
+      <text whiteSpace="pre">{{ prompt_template }}</text>
+    </cp>
+    <cp caption="Sample Runs of the Prompts (Historical Messages and Rewards)">
+      <cp for="experiment in experiments" caption="Sample Run #{{ loop.index + 1 }}">
+        <cp caption="Overall Status">
+          This run has {{ experiment.status }}. The final score is {{ experiment.final_reward }}.
+        </cp>
+        <cp caption="Messages">
+          <object data="{{ experiment.messages }}" />
+        </cp>
+      </cp>
+    </cp>
+  </human-msg>
+</poml>

agentlightning/algorithm/base.py

@@ -0,0 +1,162 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import weakref
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Optional,
+    Union,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.client import AgentLightningClient
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+    from agentlightning.trainer import Trainer
+
+
+class BaseAlgorithm:
+    """Algorithm is the strategy, or tuner to train the agent."""
+
+    _trainer_ref: weakref.ReferenceType[Trainer] | None = None
+    _llm_proxy_ref: weakref.ReferenceType["LLMProxy"] | None = None
+    _store: LightningStore | None = None
+    _initial_resources: NamedResources | None = None
+    _adapter_ref: weakref.ReferenceType[TraceAdapter[Any]] | None = None
+
+    def is_async(self) -> bool:
+        """Return True if the algorithm is asynchronous."""
+        return inspect.iscoroutinefunction(self.run)
+
+    def set_trainer(self, trainer: Trainer) -> None:
+        """
+        Set the trainer for this algorithm.
+
+        Args:
+            trainer: The Trainer instance that will handle training and validation.
+        """
+        self._trainer_ref = weakref.ref(trainer)
+
+    def get_trainer(self) -> Trainer:
+        """
+        Get the trainer for this algorithm.
+
+        Returns:
+            The Trainer instance associated with this agent.
+        """
+        if self._trainer_ref is None:
+            raise ValueError("Trainer has not been set for this agent.")
+        trainer = self._trainer_ref()
+        if trainer is None:
+            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
+        return trainer
+
+    def set_llm_proxy(self, llm_proxy: LLMProxy | None) -> None:
+        """
+        Set the LLM proxy for this algorithm to reuse when available.
+
+        Args:
+            llm_proxy: The LLMProxy instance configured by the trainer, if any.
+        """
+        self._llm_proxy_ref = weakref.ref(llm_proxy) if llm_proxy is not None else None
+
+    def get_llm_proxy(self) -> Optional[LLMProxy]:
+        """
+        Retrieve the configured LLM proxy instance, if one has been set.
+
+        Returns:
+            The active LLMProxy instance or None when not configured.
+        """
+        if self._llm_proxy_ref is None:
+            return None
+
+        llm_proxy = self._llm_proxy_ref()
+        if llm_proxy is None:
+            raise ValueError("LLM proxy reference is no longer valid (object has been garbage collected).")
+
+        return llm_proxy
+
+    def set_adapter(self, adapter: TraceAdapter[Any]) -> None:
+        """
+        Set the adapter for this algorithm to collect and convert traces.
+        """
+        self._adapter_ref = weakref.ref(adapter)
+
+    def get_adapter(self) -> TraceAdapter[Any]:
+        """
+        Retrieve the adapter for this algorithm to communicate with the runners.
+        """
+        if self._adapter_ref is None:
+            raise ValueError("Adapter has not been set for this algorithm.")
+        adapter = self._adapter_ref()
+        if adapter is None:
+            raise ValueError("Adapter reference is no longer valid (object has been garbage collected).")
+        return adapter
+
+    def set_store(self, store: LightningStore) -> None:
+        """
+        Set the store for this algorithm to communicate with the runners.
+
+        Store is set directly instead of using weakref because its copy is meant to be
+        maintained throughout the algorithm's lifecycle.
+        """
+        self._store = store
+
+    def get_store(self) -> LightningStore:
+        """
+        Retrieve the store for this algorithm to communicate with the runners.
+        """
+        if self._store is None:
+            raise ValueError("Store has not been set for this algorithm.")
+        return self._store
+
+    def get_initial_resources(self) -> Optional[NamedResources]:
+        """
+        Get the initial resources for this algorithm.
+        """
+        return self._initial_resources
+
+    def set_initial_resources(self, resources: NamedResources) -> None:
+        """
+        Set the initial resources for this algorithm.
+        """
+        self._initial_resources = resources
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self.run(*args, **kwargs)
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Subclasses should implement this method to implement the algorithm.
+
+        Args:
+            train_dataset: The dataset to train on. Not all algorithms require a training dataset.
+            val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.
+
+        Returns:
+            Algorithm should refrain from returning anything. It should just run the algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement run().")
+
+    def get_client(self) -> AgentLightningClient:
+        """Get the client to communicate with the algorithm.
+
+        If the algorithm does not require a server-client communication, it can also create a mock client
+        that never communicates with itself.
+
+        Deprecated and will be removed in a future version.
+
+        Returns:
+            The AgentLightningClient instance associated with this algorithm.
+        """
+        raise NotImplementedError("Subclasses must implement get_client().")

agentlightning/algorithm/decorator.py

@@ -0,0 +1,256 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Dict,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset, NamedResources
+
+if TYPE_CHECKING:
+    from agentlightning.llm_proxy import LLMProxy
+
+from .base import BaseAlgorithm
+
+# Algorithm function signature types
+# We've missed a lot of combinations here.
+# Let's add them in future.
+
+
+class AlgorithmFuncSyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> None: ...
+
+
+class AlgorithmFuncSyncOnlyDataset(Protocol):
+    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...
+
+
+class AlgorithmFuncAsyncFull(Protocol):
+    def __call__(
+        self,
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+        llm_proxy: Optional[LLMProxy],
+        adapter: Optional[TraceAdapter[Any]],
+        initial_resources: Optional[NamedResources],
+    ) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyStore(Protocol):
+    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...
+
+
+class AlgorithmFuncAsyncOnlyDataset(Protocol):
+    def __call__(
+        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
+    ) -> Awaitable[None]: ...
+
+
+AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]
+
+AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]
+
+
+class AlgorithmFuncSyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+class AlgorithmFuncAsyncFallback(Protocol):
+    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...
+
+
+AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
+AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]
+
+AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]
+
+
+AsyncFlag = Literal[True, False]
+AF = TypeVar("AF", bound=AsyncFlag)
+
+
+class FunctionalAlgorithm(BaseAlgorithm, Generic[AF]):
+    """A BaseAlgorithm that wraps a function-based algorithm implementation.
+
+    This class allows users to define algorithm behavior using a simple function
+    that takes train_dataset and val_dataset parameters, rather than implementing
+    a full BaseAlgorithm subclass.
+    """
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...
+
+    @overload
+    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...
+
+    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
+        """
+        Initialize the FunctionalAlgorithm with an algorithm function.
+
+        Args:
+            algorithm_func: A function that defines the algorithm's behavior.
+                           Can be sync or async with signature:
+                           (train_dataset, val_dataset) -> None
+        """
+        super().__init__()
+        self._algorithm_func = algorithm_func
+        self._sig = inspect.signature(algorithm_func)
+        self._is_async = inspect.iscoroutinefunction(algorithm_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, algorithm_func)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[False]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None: ...
+
+    @overload
+    def run(
+        self: "FunctionalAlgorithm[Literal[True]]",
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Awaitable[None]: ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._algorithm_func(*args, **kwargs)  # type: ignore
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> Union[None, Awaitable[None]]:
+        """Execute the algorithm using the wrapped function.
+
+        Args:
+            train_dataset: The dataset to train on.
+            val_dataset: The dataset to validate on.
+
+        Returns:
+            None or Awaitable[None] if the function is async.
+        """
+        kwargs: Dict[str, Any] = {}
+        if "store" in self._sig.parameters:
+            kwargs["store"] = self.get_store()
+        if "adapter" in self._sig.parameters:
+            kwargs["adapter"] = self.get_adapter()
+        if "llm_proxy" in self._sig.parameters:
+            kwargs["llm_proxy"] = self.get_llm_proxy()
+        if "initial_resources" in self._sig.parameters:
+            kwargs["initial_resources"] = self.get_initial_resources()
+        if "train_dataset" in self._sig.parameters:
+            kwargs["train_dataset"] = train_dataset
+        elif train_dataset is not None:
+            raise TypeError(
+                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        if "val_dataset" in self._sig.parameters:
+            kwargs["val_dataset"] = val_dataset
+        elif val_dataset is not None:
+            raise TypeError(
+                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
+            )
+        # both sync and async functions can be called with the same signature
+        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
+        if self._is_async:
+            return cast(Awaitable[None], result)
+        return None
+
+
+@overload
+def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...
+
+
+@overload
+def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...
+
+
+def algo(
+    func: Union[
+        AlgorithmFuncSync,
+        AlgorithmFuncAsync,
+        AlgorithmFuncSyncFallback,
+        AlgorithmFuncAsyncFallback,
+    ],
+) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
+    """Create a BaseAlgorithm from a function.
+
+    This decorator allows you to define an algorithm using a simple function
+    instead of creating a full BaseAlgorithm subclass. The returned FunctionalAlgorithm
+    instance is callable, preserving the original function's behavior.
+
+    Args:
+        func: A function that defines the algorithm's behavior with signature:
+              (train_dataset, val_dataset) -> None
+              Can be sync or async.
+
+    Returns:
+        A callable FunctionalAlgorithm instance that preserves the original function's
+        type hints and behavior while providing all algorithm functionality.
+
+    Example:
+        @algo
+        def my_algorithm(train_dataset, val_dataset):
+            # Algorithm logic here
+            for task in train_dataset:
+                # Process training tasks
+                pass
+
+        @algo
+        async def my_async_algorithm(train_dataset, val_dataset):
+            # Async algorithm logic here
+            async for task in train_dataset:
+                # Process training tasks asynchronously
+                pass
+
+        # Function is still callable with original behavior
+        my_algorithm(train_data, val_data)
+
+        # Algorithm methods are also available
+        my_algorithm.run(train_data, val_data)
+    """
+    return FunctionalAlgorithm(func)

agentlightning/algorithm/fast.py

@@ -0,0 +1,208 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from datetime import datetime
+from typing import Any, List, Literal, Optional
+
+from agentlightning.llm_proxy import ModelConfig
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
+
+from .base import BaseAlgorithm
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["FastAlgorithm", "Baseline"]
+
+
+class FastAlgorithm(BaseAlgorithm):
+    """Algorithm that can run fast and qualify for dev mode.
+
+    Fast algorithms enable agent developers to quickly iterate on agent development
+    without waiting for a long training to complete.
+    """
+
+
+def _timestamp_to_iso_str(timestamp: float) -> str:
+    return datetime.fromtimestamp(timestamp).isoformat()
+
+
+class Baseline(FastAlgorithm):
+    """A dummy implementation of algorithm interface that puts all dataset into the queue, and waits for all rollouts to complete.
+
+    Logs all collected spans and rewards.
+
+    Args:
+        model_list: Optional list of models to load into the llm proxy.
+            If both model_list and llm_proxy is provided, llm_proxy will be launched.
+            Not implemented yet.
+        n_epochs: Number of epochs to run through the dev dataset.
+        train_split: Fraction of dev dataset to use for training vs validation. Must be between 0 and 1.
+        polling_interval: Time interval (in seconds) to poll the store for queue length and for completed rollouts.
+        max_queue_length: Maximum number of rollouts to keep in the queue at any time.
+    """
+
+    def __init__(
+        self,
+        *,
+        model_list: Optional[List[ModelConfig]] = None,
+        n_epochs: int = 1,
+        train_split: float = 0.5,
+        polling_interval: float = 5.0,
+        max_queue_length: int = 4,
+        span_verbosity: Literal["keys", "key_values", "none"] = "keys",
+    ) -> None:
+        super().__init__()
+        self.n_epochs = n_epochs
+        self.train_split = train_split
+        self.polling_interval = polling_interval
+        self.max_queue_length = max_queue_length
+        self.span_verbosity = span_verbosity
+        if not (0.0 < self.train_split < 1.0):
+            raise ValueError("train_split must be between 0 and 1.")
+
+        self._finished_rollout_count = 0
+
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        if self.span_verbosity == "none":
+            return ""
+
+        prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
+        elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
+
+        msg = (
+            prefix_msg
+            + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
+            + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
+            + f"{elapsed} seconds. "
+        )
+        if self.span_verbosity == "key_values":
+            msg += f"Attributes: {span.attributes}"
+        else:
+            msg += f"Attribute keys: {list(span.attributes.keys())}"
+        return msg
+
+    async def _handle_rollout_finish(self, rollout: Rollout) -> None:
+        store = self.get_store()
+
+        rollout_id = rollout.rollout_id
+        rollout_end_time = rollout.end_time or asyncio.get_event_loop().time()
+        logger.info(
+            f"[Rollout {rollout_id}] Finished with status {rollout.status} in {rollout_end_time - rollout.start_time:.2f} seconds."
+        )
+
+        # Logs all the attempts and their corresponding spans
+        attempts = await store.query_attempts(rollout_id)
+        for attempt in attempts:
+            logger.info(
+                f"[Rollout {rollout_id} | Attempt {attempt.sequence_id}] ID: {attempt.attempt_id}. Status: {attempt.status}. Worker: {attempt.worker_id}"
+            )
+            spans = await store.query_spans(rollout_id=rollout_id)
+            for span in spans:
+                if self.span_verbosity != "none":
+                    logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
+
+        # Attempts to adapt the spans using the adapter if provided
+        try:
+            adapter = self.get_adapter()
+            spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
+            transformed_data = adapter.adapt(spans)
+            logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
+        except ValueError:
+            logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
+
+    async def _enqueue_rollouts(
+        self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
+    ) -> None:
+        store = self.get_store()
+
+        for index in train_indices + val_indices:
+            queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+            if len(queuing_rollouts) <= 1:
+                # Only enqueue a new rollout when there is at most 1 rollout in the queue.
+                sample = dataset[index]
+                mode = "train" if index in train_indices else "val"
+                rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                logger.info(f"[Rollout {rollout.rollout_id}] Enqueued in {mode} mode with sample: {sample}")
+            await asyncio.sleep(self.polling_interval)
+
+    async def _harvest_rollout_spans(self, rollout_id: str):
+        store = self.get_store()
+        last_status: Optional[RolloutStatus] = None
+
+        while True:
+            rollout = await store.get_rollout_by_id(rollout_id)
+            if rollout is not None:
+                if rollout.status in ["succeeded", "failed", "cancelled"]:
+                    # Rollout is finished, log all the data.
+                    await self._handle_rollout_finish(rollout)
+                    # We are done here.
+                    self._finished_rollout_count += 1
+                    logger.info(f"Finished {self._finished_rollout_count} rollouts.")
+                    break
+
+                if last_status != rollout.status:
+                    if last_status is not None:
+                        logger.info(f"[Rollout {rollout_id}] Status changed to {rollout.status}.")
+                    else:
+                        logger.info(f"[Rollout {rollout_id}] Status is initialized to {rollout.status}.")
+                    last_status = rollout.status
+
+                else:
+                    logger.debug(f"[Rollout {rollout_id}] Status is still {rollout.status}.")
+
+            await asyncio.sleep(self.polling_interval)
+
+    async def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        train_dataset_length = len(train_dataset) if train_dataset is not None else 0
+        val_dataset_length = len(val_dataset) if val_dataset is not None else 0
+        if train_dataset_length == 0 and val_dataset_length == 0:
+            logger.error(
+                "MockAlgorithm requires at least a train_dataset or val_dataset to run. No train_dataset or val_dataset is provided. Exiting."
+            )
+            return
+
+        concatenated_dataset = [train_dataset[i] for i in range(train_dataset_length) if train_dataset is not None] + [
+            val_dataset[i] for i in range(val_dataset_length) if val_dataset is not None
+        ]
+        train_indices = list(range(0, train_dataset_length))
+        val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
+
+        store = self.get_store()
+
+        # Currently we only supports a single resource update at the start.
+        initial_resources = self.get_initial_resources()
+        if initial_resources is not None:
+            resource_update = await store.update_resources("default", initial_resources)
+            resources_id = resource_update.resources_id
+            logger.info(f"Initial resources set: {initial_resources}")
+        else:
+            logger.warning("No initial resources provided. Skip initializing resources.")
+            resources_id = None
+
+        for epoch in range(self.n_epochs):
+            harvest_tasks: List[asyncio.Task[None]] = []
+            logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
+            for index in train_indices + val_indices:
+                queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                if len(queuing_rollouts) <= self.max_queue_length:
+                    # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
+                    sample = concatenated_dataset[index]
+                    mode = "train" if index in train_indices else "val"
+                    rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                    harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
+                    logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
+                else:
+                    # Sleep a bit and try again later.
+                    await asyncio.sleep(self.polling_interval)
+
+            # Wait for all harvest tasks to complete
+            print(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
+            if len(harvest_tasks) > 0:
+                await asyncio.gather(*harvest_tasks)

agentlightning/algorithm/verl/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .interface import VERL
+
+__all__ = ["VERL"]

agentlightning/algorithm/verl/interface.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, Optional
+
+from hydra import compose, initialize
+from omegaconf import OmegaConf
+
+from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.types import Dataset
+from agentlightning.verl.entrypoint import run_ppo  # type: ignore
+
+
+class VERL(BaseAlgorithm):
+    """Algorithm leveraging VERL as the backend framework.
+
+    **Note on Customization:**
+
+    At present, we recommend copying the source code from VERL and modifying it as needed to suit your requirements.
+    Native support for customizing training logic will be provided in future releases.
+
+    Args:
+        config: The VERL configuration, matching what is typically provided when running VERL via the command line.
+            This config will be merged with VERL's base configuration and processed by Hydra.
+    """
+
+    def __init__(self, config: dict[str, Any]):
+        super().__init__()
+
+        # Compose the base config exactly like your decorator:
+        with initialize(version_base=None, config_path="pkg://agentlightning/verl"):
+            base_cfg = compose(config_name="config")
+
+        # Merge your dict overrides
+        override_conf = OmegaConf.create(config)
+        self.config = OmegaConf.merge(base_cfg, override_conf)
+
+    def run(
+        self,
+        train_dataset: Optional[Dataset[Any]] = None,
+        val_dataset: Optional[Dataset[Any]] = None,
+    ) -> None:
+        try:
+            store = self.get_store()
+        except Exception:
+            print("Store is not set. Assuming v0 execution mode.")
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=None,
+                llm_proxy=None,
+                adapter=None,
+            )
+        else:
+            print("Store is set. Assuming v1 execution mode.")
+            llm_proxy = self.get_llm_proxy()
+            adapter = self.get_adapter()
+            run_ppo(
+                self.config,
+                train_dataset=train_dataset,
+                val_dataset=val_dataset,
+                store=store,
+                llm_proxy=llm_proxy,
+                adapter=adapter,
+            )
+
+    def get_client(self) -> AgentLightningClient:
+        port = self.config.agentlightning.port
+        return AgentLightningClient(endpoint=f"http://localhost:{port}")

agentlightning/cli/__init__.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent Lightning command line interface entry point."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from typing import Dict, Iterable, Tuple
+
+_SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
+    "vllm": ("agentlightning.cli.vllm", "Run the vLLM CLI with Agent Lightning instrumentation."),
+    "store": ("agentlightning.cli.store", "Run a LightningStore server."),
+    "agentops": ("agentlightning.cli.agentops_server", "Start the AgentOps server manager."),
+}
+
+_DESCRIPTION = "Agent Lightning CLI entry point.\n\nAvailable subcommands:\n" + "\n".join(
+    f"  {name:<10}{desc}" for name, (_, desc) in _SUBCOMMANDS.items()
+)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    """Dispatch to the requested Agent Lightning subcommand."""
+    parser = argparse.ArgumentParser(
+        prog="agl",
+        description=_DESCRIPTION,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("subcommand", choices=_SUBCOMMANDS.keys(), help="Subcommand to run.")
+    parser.add_argument("args", nargs=argparse.REMAINDER, help=argparse.SUPPRESS)
+
+    parsed = parser.parse_args(list(argv) if argv is not None else None)
+    module_name, _ = _SUBCOMMANDS[parsed.subcommand]
+    module = importlib.import_module(module_name)
+
+    entry_point = getattr(module, "main", None)
+    if entry_point is None:
+        parser.error(f"Subcommand '{parsed.subcommand}' does not define a callable 'main'")
+
+    dispatch_args = parsed.args
+    original_argv = sys.argv
+    sys.argv = [f"{parser.prog} {parsed.subcommand}", *dispatch_args]
+    try:
+        result = entry_point(dispatch_args or None)
+    finally:
+        sys.argv = original_argv
+
+    if isinstance(result, int):
+        return result
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agentlightning/cli/agentops_server.py

@@ -1,14 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import argparse
 import time
+from typing import Iterable
+
 from agentlightning.instrumentation.agentops import AgentOpsServerManager
 
 
-if __name__ == "__main__":
-    import argparse
-
+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()
+    args = parser.parse_args(list(argv) if argv is not None else None)
 
     manager = AgentOpsServerManager(daemon=args.daemon, port=args.port)
     try:
@@ -18,3 +23,8 @@ if __name__ == "__main__":
             time.sleep(1)
     except KeyboardInterrupt:
         manager.stop()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agentlightning/cli/store.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Run a LightningStore server for persistent access from multiple processes."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+from typing import Iterable
+
+from agentlightning.logging import configure_logger
+from agentlightning.store.client_server import LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+
+
+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")
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    configure_logger()
+
+    store = InMemoryLightningStore()
+    server = LightningStoreServer(store, host="0.0.0.0", port=args.port)
+    asyncio.run(server.run_forever())
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agentlightning/cli/vllm.py

@@ -1,10 +1,29 @@
-from typing import List
+# Copyright (c) Microsoft. All rights reserved.
 
-from vllm.entrypoints.cli.main import main
+from __future__ import annotations
 
-from agentlightning.instrumentation.vllm import instrument_vllm
+from typing import Iterable
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    import sys
+
+    from vllm.entrypoints.cli.main import main as vllm_main
+
+    from agentlightning.instrumentation.vllm import instrument_vllm
+
+    instrument_vllm()
+    if argv is not None:
+        original_argv = sys.argv
+        sys.argv = [original_argv[0], *list(argv)]
+        try:
+            vllm_main()
+        finally:
+            sys.argv = original_argv
+    else:
+        vllm_main()
+    return 0
 
 
 if __name__ == "__main__":
-    instrument_vllm()
-    main()
+    raise SystemExit(main())

agentlightning/client.py

@@ -1,14 +1,18 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Legacy client for interacting with a legacy Agent Lightning server."""
+
 import asyncio
 import logging
 import time
 import urllib.parse
-from typing import Any, Dict, Optional, List, Union
+import warnings
+from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
 import requests
 
-from .types import Rollout, Task, TaskInput, TaskIfAny, ResourcesUpdate, NamedResources
-
+from .types import NamedResources, ResourcesUpdate, RolloutLegacy, Task, TaskIfAny, TaskInput
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +40,9 @@ class AgentLightningClient:
             poll_interval: The interval in seconds to wait between polling for new tasks.
             timeout: The timeout in seconds for HTTP requests.
         """
+        warnings.warn(
+            "AgentLightningClient is deprecated. Please use LightningStoreClient instead.", DeprecationWarning
+        )
         self.endpoint = endpoint
         self.task_count = 0
         self.poll_interval = poll_interval
@@ -82,7 +89,7 @@ class AgentLightningClient:
                 logger.debug(f"Async POST request failed for {url}: {e}")
                 return None
 
-    async def poll_next_task_async(self) -> Task:
+    async def poll_next_task_async(self) -> Optional[Task]:
         """Polls the server asynchronously for the next task until one is available.
 
         Returns:
@@ -137,7 +144,7 @@ class AgentLightningClient:
             return resources_update
         return None
 
-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         """Posts a completed rollout to the server asynchronously.
 
         Args:
@@ -185,7 +192,7 @@ class AgentLightningClient:
             logger.debug(f"Sync POST request failed for {url}: {e}")
             return None
 
-    def poll_next_task(self) -> Task:
+    def poll_next_task(self) -> Optional[Task]:
         """Polls the server synchronously for the next task until one is available.
 
         Returns:
@@ -239,7 +246,7 @@ class AgentLightningClient:
             return resources_update
         return None
 
-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         """Posts a completed rollout to the server synchronously.
 
         Args:
@@ -277,6 +284,7 @@ class DevTaskLoader(AgentLightningClient):
             resources: Either NamedResources or ResourcesUpdate object.
             **kwargs: Additional arguments passed to the parent AgentLightningClient.
         """
+        warnings.warn("DevTaskLoader is deprecated. Please use Trainer.dev instead.", DeprecationWarning)
         super().__init__(endpoint="local://", **kwargs)
         self._tasks = tasks.copy()
         if len(self._tasks) == 0:
@@ -295,14 +303,14 @@ class DevTaskLoader(AgentLightningClient):
             self._resources_update = ResourcesUpdate(resources_id="local", resources=resources)
 
         # Store rollouts posted back to the loader for easy debugging of local runs
-        self._rollouts: List[Rollout] = []
+        self._rollouts: List[RolloutLegacy] = []
 
     @property
-    def rollouts(self) -> List[Rollout]:
+    def rollouts(self) -> List[RolloutLegacy]:
         """Return rollouts that have been posted back to the loader."""
         return self._rollouts
 
-    def poll_next_task(self) -> Task:
+    def poll_next_task(self) -> Optional[Task]:
         """Returns the next task from the local queue.
 
         If tasks are TaskInput objects, assembles them into Task objects.
@@ -344,12 +352,12 @@ class DevTaskLoader(AgentLightningClient):
         logger.debug("DevTaskLoader returning latest resources.")
         return self._resources_update
 
-    def post_rollout(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    def post_rollout(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         logger.debug(f"DevTaskLoader received rollout for task: {rollout.rollout_id}")
         self._rollouts.append(rollout)
         return {"status": "received", "rollout_id": rollout.rollout_id}
 
-    async def poll_next_task_async(self) -> Task:
+    async def poll_next_task_async(self) -> Optional[Task]:
         return self.poll_next_task()
 
     async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:
@@ -358,7 +366,7 @@ class DevTaskLoader(AgentLightningClient):
     async def get_latest_resources_async(self) -> Optional[ResourcesUpdate]:
         return self.get_latest_resources()
 
-    async def post_rollout_async(self, rollout: Rollout) -> Optional[Dict[str, Any]]:
+    async def post_rollout_async(self, rollout: RolloutLegacy) -> Optional[Dict[str, Any]]:
         return self.post_rollout(rollout)
 
     def __repr__(self):

agentlightning/config.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 """
 This file is not carefully reviewed.
 It might contain unintentional bugs and issues.
@@ -9,26 +11,28 @@ from __future__ import annotations
 import argparse
 import inspect
 import logging
+from typing import _GenericAlias  # type: ignore
 from typing import (
     Any,
+    Callable,
+    Dict,
     List,
+    Tuple,
     Type,
     TypeVar,
     Union,
-    _GenericAlias,  # type: ignore
-    get_origin,
     get_args,
-    Tuple,
-    Callable,
-    overload,
-    Dict,
+    get_origin,
     get_type_hints,
+    overload,
 )
 
 CliConfigurable = Any
 
 logger = logging.getLogger(__name__)
 
+__all__ = ["lightning_cli"]
+
 # TypeVars for precise return type hinting with overloads
 _C = TypeVar("_C", bound=CliConfigurable)
 _C1 = TypeVar("_C1", bound=CliConfigurable)
@@ -67,8 +71,8 @@ def nullable_float(value: str) -> float | None:
 
 def _str_to_bool(v: str) -> bool:
     """Converts common string representations of bool to Python bool (case-insensitive)."""
-    if isinstance(v, bool):  # Allow passing bools directly if used programmatically
-        return v
+    if isinstance(v, bool):  # type: ignore
+        return v  # Allow passing bools directly if used programmatically
     lowered_v = v.lower()
     if lowered_v in ("yes", "true", "t", "y", "1"):
         return True
@@ -305,7 +309,10 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
 
 
-def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:
+# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
+
+
+def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
     """
     Parses command-line arguments to configure and instantiate provided CliConfigurable classes.
 

agentlightning/emitter/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .exception import emit_exception
+from .message import emit_message
+from .object import emit_object
+from .reward import (
+    emit_reward,
+    find_final_reward,
+    find_reward_spans,
+    get_reward_value,
+    is_reward_span,
+    reward,
+)
+
+__all__ = [
+    "reward",
+    "emit_reward",
+    "get_reward_value",
+    "is_reward_span",
+    "find_reward_spans",
+    "find_final_reward",
+    "emit_message",
+    "emit_object",
+    "emit_exception",
+]

agentlightning/emitter/exception.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import traceback
+
+from opentelemetry.semconv.attributes import exception_attributes
+
+from agentlightning.types import SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_exception(exception: BaseException) -> None:
+    """Emit an exception as a span."""
+    if not isinstance(exception, BaseException):  # type: ignore
+        logger.error(f"Expected an BaseException instance, got: {type(exception)}. Skip emit_exception.")
+        return
+
+    tracer = get_tracer()
+    stacktrace = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__))
+    attributes = {
+        exception_attributes.EXCEPTION_TYPE: type(exception).__name__,
+        exception_attributes.EXCEPTION_MESSAGE: str(exception),
+        exception_attributes.EXCEPTION_ESCAPED: True,
+    }
+    if stacktrace.strip():
+        attributes[exception_attributes.EXCEPTION_STACKTRACE] = stacktrace
+
+    span = tracer.start_span(
+        SpanNames.EXCEPTION.value,
+        attributes=attributes,
+    )
+    logger.debug("Emitting exception span for %s", type(exception).__name__)
+    with span:
+        span.record_exception(exception)
+        # We don't set the status of the span here. They have other semantics.

agentlightning/emitter/message.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_message(message: str) -> None:
+    """Emit a string message as a span.
+
+    OpenTelemetry has a dedicated design of logs by design, but we can also use spans to emit messages.
+    So that it can all be unified in the data store and analyzed together.
+    """
+    if not isinstance(message, str):  # type: ignore
+        logger.error(f"Message must be a string, got: {type(message)}. Skip emit_message.")
+        return
+
+    tracer = get_tracer()
+    span = tracer.start_span(
+        SpanNames.MESSAGE.value,
+        attributes={SpanAttributeNames.MESSAGE.value: message},
+    )
+    logger.debug("Emitting message span with message: %s", message)
+    with span:
+        pass

agentlightning/emitter/object.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+from typing import Any
+
+from agentlightning.types import SpanAttributeNames, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+def emit_object(object: Any) -> None:
+    """Emit any object as a span. Make sure the object is JSON serializable."""
+    try:
+        serialized = json.dumps(object)
+    except (TypeError, ValueError):
+        logger.error(f"Object must be JSON serializable, got: {type(object)}. Skip emit_object.")
+        return
+
+    tracer = get_tracer()
+    span = tracer.start_span(
+        SpanNames.OBJECT.value,
+        attributes={SpanAttributeNames.OBJECT.value: serialized},
+    )
+    logger.debug("Emitting object span with payload size %d characters", len(serialized))
+    with span:
+        pass

agentlightning/emitter/reward.py

@@ -0,0 +1,215 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import inspect
+import json
+import logging
+import warnings
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    TypedDict,
+    TypeVar,
+    cast,
+)
+
+import agentops
+from agentops.sdk.decorators import operation
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import SpanLike, SpanNames
+
+from .utils import get_tracer
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "reward",
+    "emit_reward",
+    "get_reward_value",
+    "is_reward_span",
+    "find_reward_spans",
+    "find_final_reward",
+]
+
+
+class RewardSpanData(TypedDict):
+    type: Literal["reward"]
+    value: Optional[float]
+
+
+FnType = TypeVar("FnType", bound=Callable[..., Any])
+
+
+def _agentops_initialized() -> bool:
+    """Check if AgentOps is initialized in the current context."""
+    return agentops.get_client().initialized
+
+
+def reward(fn: FnType) -> FnType:
+    """
+    A decorator to wrap a function that computes rewards.
+    It will automatically handle the input and output of the function.
+    """
+
+    def wrap_result(result: Optional[float]) -> RewardSpanData:
+        """
+        Wrap the result of the function in a dict.
+        """
+        if result is None:
+            return {"type": "reward", "value": None}
+        if not isinstance(result, (float, int)):  # type: ignore
+            warnings.warn(f"Reward is ignored because it is not a number: {result}")
+            return {"type": "reward", "value": None}
+        return {"type": "reward", "value": float(result)}
+
+    # Check if the function is async
+    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
+
+    if is_async:
+
+        async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = await fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            async def agentops_reward_operation() -> RewardSpanData:
+                # The reward function we are interested in tracing
+                # It takes zero inputs and return a formatted dict
+                nonlocal result
+                result = await fn(*args, **kwargs)
+                return wrap_result(result)
+
+            await agentops_reward_operation()
+            return result
+
+        return wrapper_async  # type: ignore
+
+    else:
+
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if not _agentops_initialized():
+                # Track the reward without AgentOps
+                result = fn(*args, **kwargs)
+                emit_reward(cast(float, result))
+                return result
+
+            result: Optional[float] = None
+
+            @operation
+            def agentops_reward_operation() -> RewardSpanData:
+                nonlocal result
+                result = fn(*args, **kwargs)
+                return wrap_result(result)
+
+            agentops_reward_operation()
+            return result
+
+        return wrapper  # type: ignore
+
+
+def emit_reward(reward: float) -> ReadableSpan:
+    """
+    Record a new reward as a new span.
+    """
+    logger.debug(f"Emitting reward: {reward}")
+    if isinstance(reward, (int, bool)):
+        reward = float(reward)
+    if not isinstance(reward, float):
+        raise ValueError(f"Reward must be a number, got: {type(reward)}")
+
+    tracer = get_tracer()
+    span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
+    # Do nothing; it's just a number
+    with span:
+        pass
+    if not isinstance(span, ReadableSpan):
+        raise ValueError(f"Span is not a ReadableSpan: {span}")
+    return span
+
+
+def get_reward_value(span: SpanLike) -> Optional[float]:
+    """
+    Get the reward value from a span.
+    """
+    for key in [
+        "agentops.task.output",  # newer versions of agentops
+        "agentops.entity.output",
+    ]:
+        reward_dict: Dict[str, Any] | None = None
+        if span.attributes:
+            output = span.attributes.get(key)
+            if output:
+                if isinstance(output, dict):
+                    reward_dict = cast(Dict[str, Any], output)
+                elif isinstance(output, str):
+                    try:
+                        reward_dict = cast(Dict[str, Any], json.loads(output))
+                    except json.JSONDecodeError:
+                        reward_dict = None
+
+        if reward_dict and reward_dict.get("type") == "reward":
+            reward_value = reward_dict.get("value", None)
+            if reward_value is None:
+                return None
+            if not isinstance(reward_value, float):
+                logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+            return cast(float, reward_value)
+
+    # Latest emit reward format
+    if span.name == SpanNames.REWARD.value and span.attributes:
+        reward_value = span.attributes.get("reward", None)
+        if reward_value is None:
+            return None
+        if not isinstance(reward_value, float):
+            logger.error(f"Reward is not a number, got: {type(reward_value)}. This may cause undefined behaviors.")
+        return cast(float, reward_value)
+    return None
+
+
+def is_reward_span(span: SpanLike) -> bool:
+    """
+    Check if a span is a reward span.
+    """
+    maybe_reward = get_reward_value(span)
+    return maybe_reward is not None
+
+
+def find_reward_spans(spans: Sequence[SpanLike]) -> List[SpanLike]:
+    """
+    Find all reward spans in the given list of spans.
+
+    Args:
+        spans: A list of spans (either ReadableSpan or Span).
+
+    Returns:
+        A list of spans whose name matches the reward span name.
+    """
+    return [span for span in spans if is_reward_span(span)]
+
+
+def find_final_reward(spans: Sequence[SpanLike]) -> Optional[float]:
+    """
+    Get the last reward value from a list of spans.
+
+    Args:
+        spans: A list of spans (either ReadableSpan or Span).
+
+    Returns:
+        The reward value from the last reward span, or None if not found.
+    """
+    for span in reversed(spans):
+        reward = get_reward_value(span)
+        if reward is not None:
+            return reward
+    return None

agentlightning/emitter/utils.py

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Common utilities for the emitter module."""
+
+import opentelemetry.trace as trace_api
+from opentelemetry.trace import get_tracer_provider
+
+
+def get_tracer() -> trace_api.Tracer:
+    """Return the tracer used for AgentLightning spans.
+
+    Raises:
+        RuntimeError: If the tracer is not initialized.
+
+    Returns:
+        The AgentLightning tracer instance.
+    """
+    if hasattr(trace_api, "_TRACER_PROVIDER") and trace_api._TRACER_PROVIDER is None:  # type: ignore[attr-defined]
+        raise RuntimeError("Tracer is not initialized. Cannot emit a meaningful span.")
+
+    tracer_provider = get_tracer_provider()
+    return tracer_provider.get_tracer("agentlightning")

agentlightning/execution/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+from .client_server import ClientServerExecutionStrategy
+from .events import ExecutionEvent, MultiprocessingEvent, ThreadingEvent
+from .shared_memory import SharedMemoryExecutionStrategy
+
+__all__ = [
+    "ExecutionStrategy",
+    "ClientServerExecutionStrategy",
+    "ExecutionEvent",
+    "ThreadingEvent",
+    "MultiprocessingEvent",
+    "SharedMemoryExecutionStrategy",
+]

agentlightning/execution/base.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Protocol
+
+from agentlightning.store.base import LightningStore
+
+from .events import ExecutionEvent
+
+logger = logging.getLogger(__name__)
+
+
+class AlgorithmBundle(Protocol):
+    async def __call__(self, store: LightningStore, event: ExecutionEvent) -> None:
+        """Initalization and execution logic."""
+
+
+class RunnerBundle(Protocol):
+    async def __call__(self, store: LightningStore, worker_id: int, event: ExecutionEvent) -> None:
+        """Initalization and execution logic."""
+
+
+class ExecutionStrategy:
+    """When trainer has created the executable of algorithm and runner in two bundles,
+    the execution strategy defines how to run them together, and how many parallel runners to run.
+
+    The store is the centric place for the two bundles to communicate.
+
+    The algorithm and runner's behavior (whether runner should perform one step or run forever,
+    whether the algo would send out the tasks or not) are defined inside the bundle,
+    and does not belong to the execution strategy.
+
+    The execute should support Ctrl+C to exit gracefully.
+    """
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        raise NotImplementedError()

agentlightning/execution/client_server.py

@@ -0,0 +1,405 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import os
+import signal
+import time
+from multiprocessing.context import BaseContext
+from typing import Callable, Iterable, Literal, cast
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .events import ExecutionEvent, MultiprocessingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class ClientServerExecutionStrategy(ExecutionStrategy):
+    """Run algorithm (server) and runners (clients) as separate processes over HTTP.
+
+    **Execution Roles:**
+
+    - "algorithm": Start the HTTP server (`LightningStoreServer`) in-process and run the
+      algorithm bundle against it.
+    - "runner": Connect to an already running server via `LightningStoreClient` and
+      execute runner bundles (optionally in multiple processes).
+    - "both": Spawn the runner processes first, then launch the algorithm/server
+      bundle on the main process. This mode orchestrates the full loop locally.
+
+    When role == "both", you may choose which side runs on the main process via
+    `main_process` (debug helper). Running the runner bundle on the main process
+    is only supported with `n_runners == 1`.
+
+    Important: When `main_process == "runner"`, the algorithm runs in a subprocess
+    with the LightningStore server. This means any state modifications made during
+    execution remain in that subprocess and are NOT reflected in the original store
+    object passed to `execute()`. The main process runner accesses the store only
+    through the HTTP client interface.
+
+    **Abort / Stop Model (four-step escalation):**
+
+    1. Cooperative stop:
+       A shared :class:`~agentlightning.execution.events.MultiprocessingEvent`
+       (`stop_evt`) is passed to *all* bundles. Bundles should check it to exit.
+       Any crash (algorithm or runner) sets `stop_evt` so the other side can
+       stop cooperatively. Ctrl+C on the main process also flips the event.
+    2. KeyboardInterrupt synth:
+       Remaining subprocesses receive `SIGINT` to trigger `KeyboardInterrupt`
+       handlers.
+    3. Termination:
+       Stubborn subprocesses get `terminate()` (SIGTERM on POSIX).
+    4. Kill:
+        As a last resort we call `kill()` (SIGKILL on POSIX).
+
+    Notes:
+        This mirrors the semantics implemented in :mod:`shared_memory`, but adapted
+        to multiple processes and the HTTP client/server boundary.
+    """
+
+    alias: str = "cs"
+
+    def __init__(
+        self,
+        role: Literal["algorithm", "runner", "both"] | None = None,
+        server_host: str | None = None,
+        server_port: int | None = None,
+        n_runners: int = 1,
+        graceful_timeout: float = 5.0,
+        terminate_timeout: float = 5.0,
+        main_process: Literal["algorithm", "runner"] = "algorithm",
+    ) -> None:
+        """Configure the strategy.
+
+        Args:
+            role: Which side(s) to run in this process. When omitted, the
+                :envvar:`AGL_CURRENT_ROLE` environment variable is used.
+            server_host: Interface the HTTP server binds to when running the
+                algorithm bundle locally. Defaults to :envvar:`AGL_SERVER_HOST`
+                or ``"localhost"`` if unset.
+            server_port: Port for the HTTP server in "algorithm"/"both" modes.
+                Defaults to :envvar:`AGL_SERVER_PORT` or ``4747`` if unset.
+            n_runners: Number of runner processes to spawn in "runner"/"both".
+            graceful_timeout: How long to wait (seconds) after setting the stop
+                event before escalating to signals.
+            terminate_timeout: How long to wait between escalation steps beyond
+                the cooperative phase (re-used for SIGINT, terminate, and kill).
+            main_process: Which bundle runs on the main process when
+                `role == "both"`. `"runner"` requires `n_runners == 1` and is
+                primarily intended for debugging.
+        """
+        if role is None:
+            role_env = os.getenv("AGL_CURRENT_ROLE")
+            if role_env is None:
+                raise ValueError("role must be provided via argument or AGL_CURRENT_ROLE env var")
+            if role_env not in ("algorithm", "runner", "both"):
+                raise ValueError("role must be one of 'algorithm', 'runner', or 'both'")
+            role = role_env
+
+        if server_host is None:
+            server_host = os.getenv("AGL_SERVER_HOST", "localhost")
+
+        if server_port is None:
+            server_port_env = os.getenv("AGL_SERVER_PORT")
+            if server_port_env is None:
+                server_port = 4747
+            else:
+                try:
+                    server_port = int(server_port_env)
+                except ValueError as exc:
+                    raise ValueError("AGL_SERVER_PORT must be an integer") from exc
+
+        self.role = role
+        self.n_runners = n_runners
+        self.server_host = server_host
+        self.server_port = server_port
+        self.graceful_timeout = graceful_timeout
+        self.terminate_timeout = terminate_timeout
+        if main_process not in ("algorithm", "runner"):
+            raise ValueError("main_process must be 'algorithm' or 'runner'")
+        if main_process == "runner":
+            if role != "both":
+                raise ValueError("main_process='runner' is only supported when role='both'")
+            if n_runners != 1:
+                raise ValueError("main_process='runner' requires n_runners to be 1")
+        self.main_process = main_process
+
+    async def _execute_algorithm(
+        self, algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent
+    ) -> None:
+        logger.info("Starting LightningStore server on %s:%s", self.server_host, self.server_port)
+        server_store = LightningStoreServer(store, host=self.server_host, port=self.server_port)
+        server_started = False
+
+        try:
+            await server_store.start()
+            server_started = True
+            logger.debug("Algorithm bundle starting against endpoint %s", server_store.endpoint)
+            await algorithm(server_store, stop_evt)
+            logger.debug("Algorithm bundle completed successfully")
+        except KeyboardInterrupt:
+            logger.warning("Algorithm received KeyboardInterrupt; signaling stop event")
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Algorithm bundle crashed; signaling stop event")
+            stop_evt.set()
+            raise
+        finally:
+            if server_started:
+                try:
+                    await server_store.stop()
+                except Exception:
+                    logger.exception("Error stopping LightningStore server")
+                else:
+                    logger.debug("LightningStore server shutdown completed")
+
+    async def _execute_runner(self, runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
+        client_store = LightningStoreClient(f"http://{self.server_host}:{self.server_port}")
+        try:
+            logger.debug("Runner %s connecting to server at %s:%s", worker_id, self.server_host, self.server_port)
+            await runner(client_store, worker_id, stop_evt)
+            logger.debug("Runner %s completed successfully", worker_id)
+        except KeyboardInterrupt:
+            logger.warning("Runner %s received KeyboardInterrupt; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        except BaseException:
+            logger.exception("Runner %s crashed; signaling stop event", worker_id)
+            stop_evt.set()
+            raise
+        finally:
+            try:
+                await client_store.close()
+            except Exception:
+                logger.exception("Error closing LightningStore client for runner %s", worker_id)
+            else:
+                logger.debug("Runner %s closed LightningStore client", worker_id)
+
+    def _spawn_runners(
+        self,
+        runner: RunnerBundle,
+        stop_evt: ExecutionEvent,
+        *,
+        ctx: BaseContext,
+    ) -> list[multiprocessing.Process]:
+        """Used when `role == "runner"` or `role == "both"` and `n_runners > 1`."""
+        processes: list[multiprocessing.Process] = []
+
+        def _runner_sync(runner: RunnerBundle, worker_id: int, stop_evt: ExecutionEvent) -> None:
+            # Runners are executed in child processes; each process owns its own
+            # event loop to keep the asyncio scheduler isolated.
+            asyncio.run(self._execute_runner(runner, worker_id, stop_evt))
+
+        for i in range(self.n_runners):
+            process = cast(
+                multiprocessing.Process,
+                ctx.Process(target=_runner_sync, args=(runner, i, stop_evt), name=f"runner-{i}"),  # type: ignore
+            )
+            process.start()
+            logger.debug("Spawned runner process %s (pid=%s)", process.name, process.pid)
+            processes.append(process)
+
+        return processes
+
+    def _spawn_algorithm_process(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+        *,
+        ctx: BaseContext,
+    ) -> multiprocessing.Process:
+        """Used when `main_process == "runner"`."""
+
+        def _algorithm_sync(algorithm: AlgorithmBundle, store: LightningStore, stop_evt: ExecutionEvent) -> None:
+            asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+
+        process = cast(
+            multiprocessing.Process,
+            ctx.Process(target=_algorithm_sync, args=(algorithm, store, stop_evt), name="algorithm"),  # type: ignore
+        )
+        process.start()
+        logger.debug("Spawned algorithm process %s (pid=%s)", process.name, process.pid)
+        return process
+
+    def _join_until_deadline(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        timeout: float,
+    ) -> list[multiprocessing.Process]:
+        """Join ``processes`` until ``timeout`` elapses, returning those still alive."""
+        deadline = time.monotonic() + timeout
+        still_alive: list[multiprocessing.Process] = []
+        for process in processes:
+            remaining = deadline - time.monotonic()
+            if remaining > 0:
+                process.join(remaining)
+            else:
+                process.join(0)
+            if process.is_alive():
+                still_alive.append(process)
+        return still_alive
+
+    def _signal_processes(
+        self,
+        processes: Iterable[multiprocessing.Process],
+        action: Callable[[multiprocessing.Process], None],
+    ) -> None:
+        """Invoke ``action`` on each process while suppressing individual failures."""
+        for process in processes:
+            try:
+                action(process)
+            except Exception:
+                logger.exception("Error signaling process %s (pid=%s)", process.name, process.pid)
+
+    def _shutdown_processes(
+        self,
+        processes: list[multiprocessing.Process],
+        stop_evt: ExecutionEvent,
+    ) -> None:
+        """4-step escalation shutdown of ``processes``."""
+        if not processes:
+            logger.debug("No subprocesses to shutdown")
+            return
+
+        if not stop_evt.is_set():
+            logger.debug("Sending cooperative stop signal to subprocesses")
+            stop_evt.set()
+        else:
+            logger.debug("Stop event already set; waiting for subprocesses to exit")
+
+        alive = self._join_until_deadline(processes, self.graceful_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after cooperative wait; sending SIGINT to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        # SIGINT is not reliable on Windows, but we do not consider such case yet.
+        self._signal_processes(alive, lambda p: os.kill(cast(int, p.pid), signal.SIGINT))
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.warning(
+            "Subprocesses still alive after SIGINT wait; sending terminate() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.terminate())
+
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+        if not alive:
+            return
+
+        logger.error(
+            "Subprocesses still alive after terminate(); sending kill() to %s",
+            ", ".join(p.name or str(p.pid) for p in alive),
+        )
+        self._signal_processes(alive, lambda p: p.kill())
+        alive = self._join_until_deadline(alive, self.terminate_timeout)
+
+        if alive:
+            logger.error(
+                "Subprocesses failed to exit even after kill(): %s", ", ".join(p.name or str(p.pid) for p in alive)
+            )
+
+    def _check_process_exitcodes(self, processes: Iterable[multiprocessing.Process]) -> None:
+        """Raise an error if any managed process exited with a non-zero status."""
+        failed = [p for p in processes if p.exitcode not in (0, None)]
+        if failed:
+            formatted = ", ".join(f"{p.name or p.pid} (exitcode={p.exitcode})" for p in failed)
+            raise RuntimeError(f"Subprocesses failed: {formatted}")
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting client-server execution with %d runner(s) [role=%s, main_process=%s]",
+            self.n_runners,
+            self.role,
+            self.main_process,
+        )
+
+        # Re-use the active multiprocessing context so the event and processes
+        # agree on the start method (fork/spawn/forkserver).
+        ctx = multiprocessing.get_context()
+        stop_evt = MultiprocessingEvent(ctx=ctx)
+        # Track spawned processes so we can enforce termination ordering and
+        # surface non-zero exit codes back to the caller.
+        processes: list[multiprocessing.Process] = []
+
+        exception: BaseException | None = None
+        keyboard_interrupt = False
+
+        try:
+            if self.role == "algorithm":
+                logger.info("Running algorithm solely...")
+                asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+            elif self.role == "runner":
+                if self.n_runners == 1:
+                    logger.info("Running runner solely...")
+                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+                else:
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    # Wait for the processes to finish naturally.
+                    for process in processes:
+                        process.join()
+                    self._check_process_exitcodes(processes)
+            elif self.role == "both":
+                if self.main_process == "algorithm":
+                    logger.info("Spawning runner processes...")
+                    processes = self._spawn_runners(runner, stop_evt, ctx=ctx)
+                    try:
+                        logger.info("Running algorithm...")
+                        asyncio.run(self._execute_algorithm(algorithm, store, stop_evt))
+                    finally:
+                        # Always request the runner side to unwind once the
+                        # algorithm/server portion finishes (successfully or not).
+                        stop_evt.set()
+                else:  # main_process == "runner"
+                    if self.n_runners > 1:
+                        raise ValueError("main_process='runner' requires n_runners to be 1")
+
+                    logger.info("Spawning algorithm process...")
+                    algorithm_process = self._spawn_algorithm_process(algorithm, store, stop_evt, ctx=ctx)
+                    processes = [algorithm_process]
+
+                    # Run the lone runner cooperatively in-process so users can
+                    # attach a debugger. The algorithm + HTTP server live in
+                    # the background process spawned above (the provided
+                    # store must therefore be picklable when using spawn).
+                    logger.info("Running runner...")
+                    asyncio.run(self._execute_runner(runner, 0, stop_evt))
+
+                    # Wait for the algorithm process to finish.
+                    algorithm_process.join()
+            else:
+                raise ValueError(f"Unknown role: {self.role}")
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received; initiating shutdown")
+            stop_evt.set()
+            keyboard_interrupt = True
+        except BaseException as exc:
+            logger.exception("Unhandled exception in execute method")
+            stop_evt.set()
+            # Preserve the original exception so we can avoid masking it during
+            # the cleanup phase.
+            exception = exc
+            raise
+        finally:
+            logger.info("Shutting down subprocesses")
+            self._shutdown_processes(processes, stop_evt)
+            if processes:
+                try:
+                    self._check_process_exitcodes(processes)
+                except RuntimeError as err:
+                    if exception is not None or keyboard_interrupt:
+                        # We already propagate/handled a different failure, so
+                        # emit a warning instead of raising a secondary error.
+                        logger.warning("Subprocesses ended abnormally during shutdown: %s", err)
+                    else:
+                        raise

agentlightning/execution/events.py

@@ -0,0 +1,75 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import multiprocessing as mp
+import threading
+from multiprocessing.context import BaseContext
+from typing import Optional, Protocol
+
+
+class ExecutionEvent(Protocol):
+    """
+    A minimal protocol similar to threading.Event.
+
+    Methods:
+        set(): Signal event like a cancellation (idempotent).
+        clear(): Reset to the non-set state.
+        is_set() -> bool: True if event has been signaled.
+        wait(timeout: Optional[float] = None) -> bool:
+            Block until event is set or timeout. Returns True if event has signaled.
+    """
+
+    def set(self) -> None: ...
+    def clear(self) -> None: ...
+    def is_set(self) -> bool: ...
+    def wait(self, timeout: Optional[float] = None) -> bool: ...
+
+
+class ThreadingEvent:
+    """
+    An Event implementation using threading.Event.
+
+    Provides a thread-safe event object for signaling between threads.
+    """
+
+    __slots__ = ("_evt",)
+
+    def __init__(self) -> None:
+        self._evt = threading.Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)
+
+
+class MultiprocessingEvent:
+    """
+    An Event implementation using multiprocessing.Event.
+
+    Provides a process-safe event object for signaling between processes.
+    Optionally accepts a multiprocessing context for custom process start methods.
+    """
+
+    __slots__ = ("_evt",)
+
+    def __init__(self, *, ctx: Optional[BaseContext] = None) -> None:
+        self._evt = (ctx or mp).Event()
+
+    def set(self) -> None:
+        self._evt.set()
+
+    def clear(self) -> None:
+        self._evt.clear()
+
+    def is_set(self) -> bool:
+        return self._evt.is_set()
+
+    def wait(self, timeout: Optional[float] = None) -> bool:
+        return self._evt.wait(timeout)

agentlightning/execution/inter_process.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import ExecutionStrategy
+
+
+class InterProcessExecutionStrategy(ExecutionStrategy):
+
+    alias: str = "ipc"
+
+    # TODO: to be implemented

agentlightning/execution/shared_memory.py

@@ -0,0 +1,264 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import threading
+from contextlib import suppress
+from queue import SimpleQueue
+from typing import Any, Awaitable, Callable, List, Literal, Optional, Tuple
+
+from agentlightning.store.base import LightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+
+from .base import AlgorithmBundle, ExecutionStrategy, RunnerBundle
+from .events import ExecutionEvent, ThreadingEvent
+
+logger = logging.getLogger(__name__)
+
+
+class SharedMemoryExecutionStrategy(ExecutionStrategy):
+    """Run algorithm and runners in a single process with threads sharing memory.
+
+    Termination & abort model:
+
+    - One shared ThreadingEvent (`stop_evt`) is passed to *all* bundles.
+    - The main thread (only) receives KeyboardInterrupt on Ctrl+C; we set `stop_evt` there.
+    - If any bundle raises, we set `stop_evt` from that thread to stop the rest.
+    - After the main-thread bundle finishes normally:
+      - If main_thread is "algorithm", we also set `stop_evt` to stop the runners.
+      - If main_thread is "runner", we do not set `stop_evt` to stop the algorithm.
+        We instead wait for the algorithm to finish naturally.
+    - Background threads are daemons; we join briefly and log any stragglers.
+
+    Notes: Signals other than SIGINT (e.g., SIGTERM) are not intercepted; we respect
+    Python's default behavior for them.
+    """
+
+    alias: str = "shm"
+
+    def __init__(
+        self,
+        n_runners: int = 1,
+        main_thread: Literal["algorithm", "runner"] = "runner",
+        join_timeout: float = 15.0,
+        graceful_delay: float = 5.0,
+        poll_interval: float = 0.05,
+    ) -> None:
+        if main_thread not in ("algorithm", "runner"):
+            raise ValueError("main_thread must be 'algorithm' or 'runner'")
+        if main_thread == "runner" and n_runners != 1:
+            raise ValueError("When main_thread is 'runner', n_runners must be 1")
+        self.n_runners = n_runners
+        self.main_thread = main_thread
+        self.join_timeout = join_timeout
+        self.graceful_delay = graceful_delay
+        self.poll_interval = poll_interval
+
+    async def _run_until_completed_or_canceled(self, coro: Awaitable[Any], stop_evt: ExecutionEvent) -> Any:
+        """Run `coro` until it finishes or a cooperative stop is requested.
+
+        Control flow:
+          1) Start the bundle coroutine as `task`.
+          2) Start a watcher task that waits for `stop_evt` *without blocking* the loop
+             by periodically polling the threading event.
+          3) When the stop event flips:
+               a) Give the bundle *graceful_delay* seconds to finish on its own,
+                  because well-behaved bundles will check the event and return.
+               b) If still running after the grace period, cancel the bundle task.
+          4) Ensure both tasks are awaited; swallow `CancelledError` where appropriate.
+
+        This is a *backup* mechanism for bundles that might not poll the event
+        frequently; cooperative shutdown (checking `stop_evt` yourself) is still preferred.
+        """
+        task: asyncio.Task[Any] = asyncio.create_task(coro)  # type: ignore
+        task_exception: Optional[BaseException] = None
+
+        async def watcher() -> None:
+            # Poll the threading event without blocking the event loop. Using a
+            # background thread via ``asyncio.to_thread`` makes cancellation
+            # difficult because ``ThreadingEvent.wait`` is not interruptible.
+            # Instead we cooperatively check the flag from the loop so the
+            # watcher task stays cancellable and tests don't hang when the
+            # bundle finishes naturally before the stop event is set.
+            while not stop_evt.is_set():
+                await asyncio.sleep(self.poll_interval)
+
+            # Grace period: let a cooperative bundle exit on its own.
+            try:
+                # At this point of waiting, the main task should already see the stop event.
+                await asyncio.wait_for(asyncio.shield(task), timeout=self.graceful_delay)  # type: ignore
+                logger.debug("Bundle finished by itself during grace period.")
+                return  # bundle finished by itself during grace period
+            except asyncio.TimeoutError:
+                # Still running after the grace window.
+                pass
+            except asyncio.CancelledError:
+                # If someone else canceled the task already, we're done.
+                logger.debug("Bundle already canceled by someone else; exiting watcher.")
+                return
+
+            # Still running after the grace window: cancel it.
+            if not task.done():
+                logger.debug("Graceful delay elapsed; canceling bundle task...")
+                task.cancel()
+
+        watcher_task = asyncio.create_task(watcher())
+        result: Any = None
+
+        try:
+            # We don't wait on FIRST_COMPLETED here, because we want the watcher
+            # to be able to grant a grace window after stop_evt flips.
+            await asyncio.wait(
+                {task, watcher_task}, return_when=asyncio.FIRST_COMPLETED
+            )  # pyright: ignore[reportUnknownArgumentType]
+        finally:
+            # If the main task hasn't completed yet (e.g., watcher scheduled cancel),
+            # finish the cancellation handshake.
+            if not task.done():
+                try:
+                    await asyncio.wait_for(task, timeout=self.graceful_delay)  # second chance
+                except asyncio.TimeoutError:
+                    logger.error(
+                        "Bundle task did not stop after cancellation; abandoning task."
+                        "This thread could live until the process exits."
+                    )
+                    # We return without awaiting it. asyncio.run will still try to cancel
+                    # pending tasks on loop close; if the task ignores cancellation, this
+                    # thread may still stick. It's the best we can do in Python.
+                    # We don't raise an exception here, but the thread could be a zombie.
+                    return result
+            else:
+                # Task completed naturally; retrieve result.
+                try:
+                    result = await task  # type: ignore
+                except asyncio.CancelledError:
+                    pass
+                except BaseException as exc:
+                    task_exception = exc
+
+            watcher_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await watcher_task
+
+        if task_exception is not None:
+            raise task_exception
+
+        return result  # type: ignore
+
+    def _run_algorithm(
+        self,
+        algorithm: AlgorithmBundle,
+        store: LightningStore,
+        stop_evt: ExecutionEvent,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(algorithm(store, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Algorithm bundle canceled due to stop signal.")
+        except BaseException as exc:
+            logger.exception("Algorithm bundle crashed; signaling stop to others.")
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def _run_runner(
+        self,
+        runner: RunnerBundle,
+        store: LightningStore,
+        worker_id: int,
+        stop_evt: ExecutionEvent,
+        thread_exceptions: Optional[SimpleQueue[BaseException]],
+    ) -> None:
+        try:
+            asyncio.run(self._run_until_completed_or_canceled(runner(store, worker_id, stop_evt), stop_evt))
+        except asyncio.CancelledError:
+            logger.info("Runner bundle (worker_id=%s) canceled due to stop signal.", worker_id)
+        except BaseException as exc:
+            logger.exception("Runner bundle crashed (worker_id=%s); signaling stop to others.", worker_id)
+            if thread_exceptions is not None:
+                thread_exceptions.put(exc)
+            stop_evt.set()
+            raise
+
+    def execute(self, algorithm: AlgorithmBundle, runner: RunnerBundle, store: LightningStore) -> None:
+        logger.info(
+            "Starting shm execution with %d runner(s); main thread runs '%s'",
+            self.n_runners,
+            self.main_thread,
+        )
+
+        # Create stop event and thread-safe store.
+        stop_evt = ThreadingEvent()
+        thread_safe_store = LightningStoreThreaded(store)
+
+        thread_exceptions: SimpleQueue[BaseException] = SimpleQueue()
+        raised_from_thread: Optional[BaseException] = None
+
+        def make_thread(name: str, target: Callable[..., Any], args: Tuple[Any, ...]) -> threading.Thread:
+            t = threading.Thread(name=name, target=target, args=args, daemon=True)
+            t.start()
+            return t
+
+        threads: List[threading.Thread] = []
+
+        try:
+            if self.main_thread == "algorithm":
+                # Start runner threads; algorithm runs on main thread.
+                for i in range(self.n_runners):
+                    thread = make_thread(
+                        name=f"runner-{i}",
+                        target=self._run_runner,
+                        args=(runner, thread_safe_store, i, stop_evt, thread_exceptions),
+                    )
+                    threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_algorithm(algorithm, thread_safe_store, stop_evt, None)
+
+                # If algo finishes naturally, request runners to stop.
+                stop_evt.set()
+
+            else:  # main_thread == "runner"
+                # Start algorithm in background; runner runs on main thread.
+                thread = make_thread(
+                    name="algorithm",
+                    target=self._run_algorithm,
+                    args=(algorithm, thread_safe_store, stop_evt, thread_exceptions),
+                )
+                threads.append(thread)
+
+                # Ctrl+C here raises KeyboardInterrupt on this stack.
+                # Main thread doesn't need to collect exceptions.
+                self._run_runner(runner, thread_safe_store, 0, stop_evt, None)
+
+                # If runner finishes naturally, WAIT FOR ALGORITHM TO FINISH.
+                thread.join()
+
+            if not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        except KeyboardInterrupt:
+            logger.warning("KeyboardInterrupt received on main thread; initiating cooperative shutdown...")
+            stop_evt.set()
+        finally:
+            # Attempt a clean join; if some threads don't comply, log and move on.
+            for t in threads:
+                logger.debug("Joining thread %s...", t.name)
+                t.join(timeout=self.join_timeout)
+
+            alive = [t.name for t in threads if t.is_alive()]
+            if alive:
+                logger.error(
+                    "Threads still alive after %.1fs: %s. They are daemons; continuing shutdown.",
+                    self.join_timeout,
+                    ", ".join(alive),
+                )
+
+            if raised_from_thread is None and not thread_exceptions.empty():
+                raised_from_thread = thread_exceptions.get()
+
+        if raised_from_thread is not None:
+            raise raised_from_thread

agentlightning/instrumentation/__init__.py

@@ -1,21 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 import warnings
 
-AGENTOPS_INSTALLED = False
-AGENTOPS_LANGCHAIN_INSTALLED = False
-LITELLM_INSTALLED = False
-VLLM_INSTALLED = False
+AGENTOPS_INSTALLED: bool = False
+AGENTOPS_LANGCHAIN_INSTALLED: bool = False
+LITELLM_INSTALLED: bool = False
+VLLM_INSTALLED: bool = False
 
 try:
-    from . import agentops
+    from . import agentops  # type: ignore
 
-    AGENTOPS_INSTALLED = True
+    AGENTOPS_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 
 try:
-    from . import litellm
+    from . import litellm  # type: ignore
 
-    LITELLM_INSTALLED = True
+    LITELLM_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 
@@ -30,14 +32,15 @@ except ImportError:
 
 
 try:
-    from . import agentops_langchain
+    from . import agentops_langchain  # type: ignore
 
-    AGENTOPS_LANGCHAIN_INSTALLED = True
+    AGENTOPS_LANGCHAIN_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 
 
 def instrument_all():
+    """Instrument all the instrumentation libraries."""
     if AGENTOPS_INSTALLED:
         from .agentops import instrument_agentops
 
@@ -68,6 +71,7 @@ def instrument_all():
 
 
 def uninstrument_all():
+    """Uninstrument all the instrumentation libraries."""
     if AGENTOPS_INSTALLED:
         try:
             from .agentops import uninstrument_agentops

agentlightning/instrumentation/agentops.py

@@ -1,23 +1,35 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
 import logging
 import multiprocessing
 import signal
 import socket
 import time
+from typing import Any, Callable
 
 import flask
 import setproctitle
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "instrument_agentops",
+    "uninstrument_agentops",
+    "agentops_local_server",
+    "AgentOpsServerManager",
+]
+
 # Module-level storage for originals
-_original_handle_chat_attributes = None
-_original_handle_response = None
+_original_handle_chat_attributes: Callable[..., Any] | None = None
+_original_handle_response: Callable[..., Any] | None = None
 
 
 def _patch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
     import agentops.instrumentation.providers.openai.stream_wrapper
-    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes
+    import agentops.instrumentation.providers.openai.wrappers.chat
+    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes  # type: ignore
 
     global _original_handle_chat_attributes
 
@@ -25,23 +37,43 @@ def _patch_new_agentops():
         logger.warning("AgentOps already patched. Skipping.")
         return True
 
-    _original_handle_chat_attributes = handle_chat_attributes
+    _original_handle_chat_attributes = handle_chat_attributes  # type: ignore
 
-    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):
-        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
-        if hasattr(return_value, "prompt_token_ids"):
-            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
-        if hasattr(return_value, "response_token_ids"):
-            attributes["response_token_ids"] = list(return_value.response_token_ids[0])
+    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
+        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)  # type: ignore
+        if return_value is not None and hasattr(return_value, "prompt_token_ids"):  # type: ignore
+            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)  # type: ignore
+        if return_value is not None and hasattr(return_value, "response_token_ids"):  # type: ignore
+            attributes["response_token_ids"] = list(return_value.response_token_ids[0])  # type: ignore
+
+        # For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
+        if (
+            not attributes.get("response_token_ids")
+            and return_value is not None
+            and hasattr(return_value, "choices")  # type: ignore
+            and return_value.choices  # type: ignore
+            and isinstance(return_value.choices, list)  # type: ignore
+        ):
+            first_choice = return_value.choices[0]  # type: ignore
+            if hasattr(first_choice, "token_ids"):  # type: ignore
+                attributes["response_token_ids"] = list(first_choice.token_ids)  # type: ignore
+            # newer versions of OpenAI client SDK
+            elif hasattr(first_choice, "provider_specific_fields") and "token_ids" in first_choice.provider_specific_fields:  # type: ignore
+                attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])  # type: ignore
 
         # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(return_value, "http_response") and hasattr(return_value.http_response, "json"):
-            json_data = return_value.http_response.json()
+        if (
+            return_value is not None
+            and hasattr(return_value, "http_response")  # type: ignore
+            and return_value.http_response is not None  # type: ignore
+            and hasattr(return_value.http_response, "json")  # type: ignore
+        ):
+            json_data = return_value.http_response.json()  # type: ignore
             if isinstance(json_data, dict):
                 if "prompt_token_ids" in json_data:
-                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
+                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])  # type: ignore
                 if "response_token_ids" in json_data:
-                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])
+                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])  # type: ignore
 
         return attributes
 
@@ -54,8 +86,8 @@ def _patch_new_agentops():
 
 
 def _unpatch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
     import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat
 
     global _original_handle_chat_attributes
     if _original_handle_chat_attributes is not None:
@@ -70,40 +102,40 @@ def _unpatch_new_agentops():
 
 
 def _patch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
-    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
+    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw  # type: ignore
 
     global _original_handle_response
-    _original_handle_response = _handle_response
+    _original_handle_response = _handle_response  # type: ignore
 
-    @dont_throw
-    def _handle_response_with_tokens(response, span, *args, **kwargs):
-        _original_handle_response(response, span, *args, **kwargs)
-        if hasattr(response, "prompt_token_ids"):
-            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))
-        if hasattr(response, "response_token_ids"):
-            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))
+    @dont_throw  # type: ignore
+    def _handle_response_with_tokens(response, span, *args, **kwargs):  # type: ignore
+        _original_handle_response(response, span, *args, **kwargs)  # type: ignore
+        if hasattr(response, "prompt_token_ids"):  # type: ignore
+            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))  # type: ignore
+        if hasattr(response, "response_token_ids"):  # type: ignore
+            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))  # type: ignore
 
         # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):
-            json_data = response.http_response.json()
+        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):  # type: ignore
+            json_data = response.http_response.json()  # type: ignore
             if isinstance(json_data, dict):
                 if "prompt_token_ids" in json_data:
-                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))
+                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))  # type: ignore
                 if "response_token_ids" in json_data:
-                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))
+                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))  # type: ignore
 
-    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens
+    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens  # type: ignore
     logger.info("Patched earlier version of agentops using _handle_response")
     return True
 
 
 def _unpatch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
 
     global _original_handle_response
     if _original_handle_response is not None:
-        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response
+        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response  # type: ignore
         _original_handle_response = None
         logger.info("Unpatched earlier version of agentops using _handle_response")
 
@@ -131,6 +163,7 @@ def instrument_agentops():
 
 
 def uninstrument_agentops():
+    """Uninstrument agentops to stop capturing token IDs."""
     try:
         _unpatch_new_agentops()
     except Exception:
@@ -149,18 +182,18 @@ def agentops_local_server():
     app = flask.Flask(__name__)
 
     @app.route("/v3/auth/token", methods=["POST"])
-    def fetch_token():
+    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):
+    def catch_all(path: str):  # type: ignore
         return {"path": path}
 
     return app
 
 
-def _run_server(**kwargs):
+def _run_server(**kwargs: Any):  # type: ignore
     """
     Internal function to run the Flask server.
     This is used to avoid issues with multiprocessing and Flask's reloader.
@@ -172,6 +205,8 @@ def _run_server(**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
@@ -213,7 +248,7 @@ class AgentOpsServerManager:
         return False
 
     def stop(self):
-        if self.is_alive():
+        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

agentlightning/instrumentation/agentops_langchain.py

@@ -1,20 +1,27 @@
-from typing import Dict, Any
-from agentops.integration.callbacks.langchain import LangchainCallbackHandler
-from agentops import instrumentation
+# Copyright (c) Microsoft. All rights reserved.
 
+from typing import Any, Dict
+
+from agentops import instrumentation
+from agentops.integration.callbacks.langchain import LangchainCallbackHandler
 
 original_on_chain_start = LangchainCallbackHandler.on_chain_start
 langgraph_entry = None
 
+__all__ = [
+    "instrument_agentops_langchain",
+    "uninstrument_agentops_langchain",
+]
 
-def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
+
+def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
     if "name" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
             serialized = {}
         serialized = serialized.copy()
         serialized["name"] = kwargs["name"]
     if "run_id" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
             serialized = {}
         serialized = serialized.copy()
         if "id" not in serialized:
@@ -23,12 +30,14 @@ def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **k
 
 
 def instrument_agentops_langchain():
+    """Bypass AgentOp's native support for Langchain."""
     global langgraph_entry
     langgraph_entry = instrumentation.AGENTIC_LIBRARIES.pop("langgraph", None)
     LangchainCallbackHandler.on_chain_start = on_chain_start
 
 
 def uninstrument_agentops_langchain():
+    """Restore AgentOp's native support for Langchain."""
     global langgraph_entry
     if langgraph_entry is not None:
         instrumentation.AGENTIC_LIBRARIES["langgraph"] = langgraph_entry

agentlightning/instrumentation/litellm.py

@@ -1,26 +1,38 @@
-from typing import Optional, Any
+# Copyright (c) Microsoft. All rights reserved.
+
+"""LiteLLM instrumentations.
+
+It's unclear whether or not this file is useful.
+It seems that LiteLLM owns its own telemetry from their own entrance
+https://docs.litellm.ai/docs/observability/agentops_integration
+"""
+
+from typing import Any, Optional
 
 from litellm.integrations.opentelemetry import OpenTelemetry
 
-# It's unclear whether or not this file is useful
-# It seems that LiteLLM owns its own telemetry from their own entrance
-# https://docs.litellm.ai/docs/observability/agentops_integration
+__all__ = [
+    "instrument_litellm",
+    "uninstrument_litellm",
+]
 
-original_set_attributes = OpenTelemetry.set_attributes
+original_set_attributes = OpenTelemetry.set_attributes  # type: ignore
 
 
-def patched_set_attributes(self, span: Any, kwargs, response_obj: Optional[Any]):
+def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
     original_set_attributes(self, span, kwargs, response_obj)
     # Add custom attributes
-    if response_obj.get("prompt_token_ids"):
+    if response_obj is not None and response_obj.get("prompt_token_ids"):
         span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
-    if response_obj.get("response_token_ids"):
+    if response_obj is not None and response_obj.get("response_token_ids"):
         span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))
 
 
 def instrument_litellm():
+    """Instrument litellm to capture token IDs."""
     OpenTelemetry.set_attributes = patched_set_attributes
 
 
 def uninstrument_litellm():
+    """Uninstrument litellm to stop capturing token IDs."""
     OpenTelemetry.set_attributes = original_set_attributes

agentlightning/instrumentation/verl_chat_scheduler.py

@@ -1,148 +0,0 @@
-# type: ignore
-
-# https://github.com/volcengine/verl/blob/bd94bd61fe4193e56f2845dc794004afbef7f818/examples/ppo_trainer/naive_chat_scheduler.py
-# This file is part of VERL example. It should be included in the VERL package but it's not currently.
-
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import asyncio
-from typing import Any, Dict, List
-
-import torch
-from openai.types.chat.chat_completion import ChatCompletion
-from tensordict import TensorDict
-
-from verl.protocol import DataProto
-from verl.workers.rollout.async_server import ChatCompletionScheduler
-
-
-class NaiveChatCompletionScheduler(ChatCompletionScheduler):
-    """
-    A very naive implementation of ChatCompletionScheduler for demo purpose,
-    only do single-turn chat completion.
-    """
-
-    async def generate_sequences(self, batch: DataProto, **sampling_params) -> DataProto:
-        kwargs = dict(
-            n=self.config.n,
-            max_completion_tokens=self.config.response_length,
-            temperature=self.config.temperature,
-            top_p=self.config.top_p,
-        )
-
-        do_sample = batch.meta_info.get("do_sample", True)
-        is_validate = batch.meta_info.get("validate", False)
-        if not do_sample or is_validate:
-            kwargs["n"] = 1
-            kwargs["temperature"] = 0
-
-        kwargs.update(sampling_params)
-        print(f"[NaiveChatCompletionScheduler] generate_sequences sampling params: {kwargs}")
-
-        async def callback(completions: ChatCompletion, info: Dict[str, Any], exception: Exception):
-            assert exception is None, f"exception: {exception}"
-            conversation, batch_conversations, batch_index = (
-                info["conversation"],
-                info["batch_conversations"],
-                info["batch_index"],
-            )
-
-            conversations = []
-            for choice in completions.choices:
-                chat = conversation.copy()
-                chat.append({"role": choice.message.role, "content": choice.message.content})
-                conversations.append(chat)
-            batch_conversations[batch_index] = conversations
-
-            # NOTE: we can call tools and resubmit chat completions here.
-            # call_tools(completions, info)
-            # await self.submit_chat_completions(callback2, ...)
-
-        # TODO: we may need to control max concurrent requests here, or it will harm prefix cache hit rate.
-        tasks, batch_conversations = [], [None] * len(batch)
-        for batch_index, conversation in enumerate(batch.non_tensor_batch["raw_prompt"]):
-            # raw_prompt: [{"role": "user", "content": ""}, ["role": "assistant", "content"], ...]
-            tasks.append(
-                asyncio.create_task(
-                    self.submit_chat_completions(
-                        callback=callback,
-                        callback_additional_info={
-                            "batch_conversations": batch_conversations,
-                            "batch_index": batch_index,
-                            "conversation": list(conversation),
-                        },
-                        model=self.model_name,
-                        messages=conversation.tolist(),
-                        **kwargs,
-                    )
-                )
-            )
-        await asyncio.gather(*tasks)
-        print("[NaiveChatCompletionScheduler] generate_sequences done")
-
-        return self._postprocess(batch, batch_conversations, kwargs["n"])
-
-    def _postprocess(
-        self, batch: DataProto, batch_conversations: List[List[List[Dict[str, str]]]], n: int
-    ) -> DataProto:
-        # NOTE: consistent with batch version of generate_sequences in vllm_rollout_spmd.py
-        # prompts: left pad
-        # responses: right pad
-        # input_ids: prompt + response
-        # attention_mask: [0,0,0,0,1,1,1,1, | 1,1,1,0,0,0,0,0]
-        # position_ids:   [0,0,0,0,0,1,2,3, | 4,5,6,7,8,9,10,11]
-
-        # prompts: [prompt] from input dataset
-        prompts = [
-            self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=False)
-            for prompt in batch.non_tensor_batch["raw_prompt"]
-        ]
-
-        # flatten batch_conversations if n > 1
-        assert len(batch_conversations) == len(prompts)
-        batch_conversations = [conversation for conversations in batch_conversations for conversation in conversations]
-        assert len(batch_conversations) == len(prompts) * n
-
-        # sequences: [prompt + response]
-        sequences = [
-            self.tokenizer.apply_chat_template(conversation, add_generation_prompt=False, tokenize=False)
-            for conversation in batch_conversations
-        ]
-
-        # responses: [response]
-        # TODO: mask out tools calling tokens?
-        responses = [sequence[len(prompts[i // n]) :] for i, sequence in enumerate(sequences)]
-
-        prompts = self.tokenizer(prompts, return_tensors="pt", padding="longest", padding_side="left")
-        responses = self.tokenizer(responses, return_tensors="pt", padding="longest", padding_side="right")
-        if n > 1:
-            prompts["input_ids"] = prompts["input_ids"].repeat_interleave(n, dim=0)
-            prompts["attention_mask"] = prompts["attention_mask"].repeat_interleave(n, dim=0)
-
-        input_ids = torch.cat([prompts["input_ids"], responses["input_ids"]], dim=1)
-        attention_mask = torch.cat([prompts["attention_mask"], responses["attention_mask"]], dim=1)
-        position_ids = (attention_mask.cumsum(dim=1) - 1) * attention_mask
-
-        batch = TensorDict(
-            {
-                "prompts": prompts["input_ids"],
-                "responses": responses["input_ids"],
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-                "position_ids": position_ids,
-            },
-            batch_size=len(input_ids),
-        )
-
-        return DataProto(batch=batch)

agentlightning/instrumentation/vllm.py

@@ -1,12 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 from __future__ import annotations
 
 import warnings
-from typing import List
+from typing import Any, List
 
-from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 import vllm.entrypoints.openai.protocol
+from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 
+__all__ = [
+    "instrument_vllm",
+    "uninstrument_vllm",
+]
+
 
 class ChatCompletionResponsePatched(ChatCompletionResponse):
     prompt_token_ids: List[int] | None = None
@@ -17,15 +24,15 @@ original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full
 
 
 async def chat_completion_full_generator(
-    self,
-    request,
-    result_generator,
+    self: Any,
+    request: Any,
+    result_generator: Any,
     request_id: str,
     model_name: str,
-    conversation,
-    tokenizer,
-    request_metadata,
-):
+    conversation: Any,
+    tokenizer: Any,
+    request_metadata: Any,
+) -> Any:
     prompt_token_ids: List[int] | None = None
     response_token_ids: List[List[int]] | None = None
 
@@ -57,6 +64,10 @@ async def chat_completion_full_generator(
 
 
 def instrument_vllm():
+    """Instrument vLLM to capture token IDs generated by engine.
+
+    This instrumentation has been merged to upstream vLLM since v0.10.2.
+    """
     if vllm.entrypoints.openai.protocol.ChatCompletionResponse is ChatCompletionResponsePatched:
         warnings.warn("vllm is already instrumented. Skip the instrumentation.")
         return
@@ -66,4 +77,5 @@ def instrument_vllm():
 
 
 def uninstrument_vllm():
+    """Uninstrument vLLM to stop capturing token IDs generated by engine."""
     OpenAIServingChat.chat_completion_full_generator = original_chat_completion_full_generator

agentlightning/litagent.py

@@ -1,204 +0,0 @@
-from __future__ import annotations
-
-import logging
-import weakref
-from typing import Any, List, Dict, Union, Optional, TYPE_CHECKING
-
-from .types import NamedResources, Rollout, Task, TaskInput, Triplet, RolloutRawResult
-
-if TYPE_CHECKING:
-    from .trainer import Trainer
-    from .runner import AgentRunner
-    from .tracer import BaseTracer
-
-
-logger = logging.getLogger(__name__)
-
-
-class LitAgent:
-    """Base class for the training and validation logic of an agent.
-
-    Developers should subclass this class and implement the rollout methods
-    to define the agent's behavior for a single task. The agent's logic
-    is completely decoupled from the server communication and training
-    infrastructure.
-    """
-
-    def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
-        """
-        Initialize the LitAgent.
-
-        Args:
-            trained_agents: Optional string representing the trained agents.
-                            This can be used to track which agents have been trained by this instance.
-        """
-        self.trained_agents = trained_agents
-        self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
-        self._runner_ref: weakref.ReferenceType[AgentRunner] | None = None
-
-    def set_trainer(self, trainer: Trainer) -> None:
-        """
-        Set the trainer for this agent.
-
-        Args:
-            trainer: The Trainer instance that will handle training and validation.
-        """
-        self._trainer_ref = weakref.ref(trainer)
-
-    @property
-    def trainer(self) -> Trainer:
-        """
-        Get the trainer for this agent.
-
-        Returns:
-            The Trainer instance associated with this agent.
-        """
-        if self._trainer_ref is None:
-            raise ValueError("Trainer has not been set for this agent.")
-        trainer = self._trainer_ref()
-        if trainer is None:
-            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
-        return trainer
-
-    @property
-    def tracer(self) -> BaseTracer:
-        """
-        Get the tracer for this agent.
-
-        Returns:
-            The BaseTracer instance associated with this agent.
-        """
-        return self.trainer.tracer
-
-    def set_runner(self, runner: AgentRunner) -> None:
-        """
-        Set the runner for this agent.
-
-        Args:
-            runner: The AgentRunner instance that will handle the execution of rollouts.
-        """
-        self._runner_ref = weakref.ref(runner)
-
-    @property
-    def runner(self) -> AgentRunner:
-        """
-        Get the runner for this agent.
-
-        Returns:
-            The AgentRunner instance associated with this agent.
-        """
-        if self._runner_ref is None:
-            raise ValueError("Runner has not been set for this agent.")
-        runner = self._runner_ref()
-        if runner is None:
-            raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
-        return runner
-
-    def on_rollout_start(self, task: Task, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called immediately before a rollout begins.
-
-        Args:
-            task: The :class:`Task` object that will be processed.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method to implement custom logic such as
-        logging, metric collection, or resource setup. By default, this is a
-        no-op.
-        """
-
-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called after a rollout completes.
-
-        Args:
-            task: The :class:`Task` object that was processed.
-            rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method for cleanup or additional
-        logging. By default, this is a no-op.
-        """
-
-    def training_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
-        """Defines the agent's behavior for a single training task.
-
-        This method should contain the logic for how the agent processes an
-        input, uses the provided resources (like LLMs or prompts), and
-        produces a result.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            rollout_id: A unique identifier for the rollout, used for tracking
-                        and reporting purposes.
-            resources: A dictionary of named resources (e.g., LLMs, prompt
-                       templates) for the agent to use.
-
-        Returns:
-            The result of the rollout, which can be one of:
-            - None. The tracing should be handled by the agent runner.
-            - A float representing the final reward.
-            - A list of `Triplet` objects for detailed, step-by-step feedback.
-            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
-            - A list of dictionaries for any trace spans.
-            - A complete `Rollout` object for full control over reporting.
-        """
-        raise NotImplementedError("Subclasses must implement the `training_rollout` method.")
-
-    def validation_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
-        """Defines the agent's behavior for a single validation task.
-
-        By default, this method redirects to `training_rollout`. Override it
-        if the agent should behave differently during validation.
-
-        Args:
-            task: The task object received from the server, containing the
-                  input data and metadata.
-            rollout_id: A unique identifier for the validation rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the validation rollout. See `training_rollout` for
-            possible return types.
-        """
-        return self.training_rollout(task, rollout_id, resources)
-
-    async def training_rollout_async(
-        self, task: TaskInput, rollout_id: str, resources: NamedResources
-    ) -> RolloutRawResult:
-        """Asynchronous version of `training_rollout`.
-
-        This method should be implemented by agents that perform asynchronous
-        operations (e.g., non-blocking I/O, concurrent API calls).
-
-        Args:
-            task: The task object received from the server.
-            rollout_id: A unique identifier for the training rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the asynchronous training rollout.
-        """
-        raise NotImplementedError("Async agents must implement the `training_rollout_async` method.")
-
-    async def validation_rollout_async(
-        self, task: TaskInput, rollout_id: str, resources: NamedResources
-    ) -> RolloutRawResult:
-        """Asynchronous version of `validation_rollout`.
-
-        By default, this method redirects to `training_rollout_async`.
-        Override it for different asynchronous validation behavior.
-
-        Args:
-            task: The task object received from the server.
-            rollout_id: A unique identifier for the validation rollout,
-                        used for tracking and reporting purposes.
-            resources: A dictionary of named resources for the agent to use.
-
-        Returns:
-            The result of the asynchronous validation rollout.
-        """
-        return await self.training_rollout_async(task, rollout_id, resources)

agentlightning/litagent/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .decorator import *
+from .litagent import *
+
+__all__ = [
+    "LitAgent",
+    "llm_rollout",
+    "prompt_rollout",
+    "rollout",
+]

agentlightning/litagent/decorator.py

@@ -0,0 +1,521 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from typing import Any, Awaitable, Callable, Dict, Protocol, TypeGuard, TypeVar, Union, overload
+
+from agentlightning.types import (
+    LLM,
+    AttemptedRollout,
+    NamedResources,
+    PromptTemplate,
+    ProxyLLM,
+    Rollout,
+    RolloutRawResult,
+)
+
+from .litagent import LitAgent
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+__all__ = [
+    "llm_rollout",
+    "prompt_rollout",
+    "rollout",
+]
+
+
+T_contra = TypeVar("T_contra", contravariant=True)
+
+
+class LlmRolloutFuncSync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM) -> RolloutRawResult: ...
+
+
+class LlmRolloutFuncSync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> RolloutRawResult: ...
+
+
+class LlmRolloutFuncAsync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM) -> Awaitable[RolloutRawResult]: ...
+
+
+class LlmRolloutFuncAsync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, llm: LLM, rollout: Rollout) -> Awaitable[RolloutRawResult]: ...
+
+
+LlmRolloutFunc = Union[
+    LlmRolloutFuncSync2[T_contra],
+    LlmRolloutFuncSync3[T_contra],
+    LlmRolloutFuncAsync2[T_contra],
+    LlmRolloutFuncAsync3[T_contra],
+]
+
+
+class PromptRolloutFuncSync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> RolloutRawResult: ...
+
+
+class PromptRolloutFuncAsync2(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate) -> Awaitable[RolloutRawResult]: ...
+
+
+class PromptRolloutFuncSync3(Protocol[T_contra]):
+    def __call__(self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout) -> RolloutRawResult: ...
+
+
+class PromptRolloutFuncAsync3(Protocol[T_contra]):
+    def __call__(
+        self, task: T_contra, prompt_template: PromptTemplate, rollout: Rollout
+    ) -> Awaitable[RolloutRawResult]: ...
+
+
+PromptRolloutFunc = Union[
+    PromptRolloutFuncSync2[T_contra],
+    PromptRolloutFuncSync3[T_contra],
+    PromptRolloutFuncAsync2[T_contra],
+    PromptRolloutFuncAsync3[T_contra],
+]
+
+
+class FunctionalLitAgentFunc(Protocol[T_contra]):
+    def __call__(
+        self, task: T_contra, *args: Any, **kwargs: Any
+    ) -> Union[RolloutRawResult, Awaitable[RolloutRawResult]]: ...
+
+
+class FunctionalLitAgent(LitAgent[T]):
+    """A specialized LitAgent that wraps a function-based rollout that accepts
+    dynamically a task input and a configured resource (LLM / prompt template / ...).
+
+    This class allows users to define agent behavior using a simple function
+    that takes task input and a resource, rather than implementing a full
+    LitAgent subclass.
+    """
+
+    def __init__(self, rollout_func: FunctionalLitAgentFunc[T], *, strip_proxy: bool = True) -> None:
+        """
+        Initialize the FunctionalLitAgent with a functional rollout function.
+
+        Args:
+            rollout_func: A function that defines the agent's behavior.
+                          Can be sync or async, and can optionally accept a Rollout parameter.
+                          The function signature determines which resources are injected (llm, prompt_template, etc.).
+            strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource when the function accepts an llm parameter.
+                         Defaults to True.
+        """
+        super().__init__()
+        self._rollout_func = rollout_func
+        self._strip_proxy = strip_proxy
+        self._is_async = inspect.iscoroutinefunction(rollout_func)
+        self._sig = inspect.signature(rollout_func)
+
+        # Copy function metadata to preserve type hints and other attributes
+        functools.update_wrapper(self, rollout_func)  # type: ignore
+
+    def _accepts_rollout(self) -> bool:
+        return "rollout" in self._sig.parameters
+
+    def _accepts_llm(self) -> bool:
+        return "llm" in self._sig.parameters
+
+    def _accepts_prompt_template(self) -> bool:
+        return "prompt_template" in self._sig.parameters
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Make the agent instance callable, preserving the original function behavior."""
+        return self._rollout_func(*args, **kwargs)  # type: ignore
+
+    def is_async(self) -> bool:
+        return self._is_async
+
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute a synchronous rollout using the wrapped function.
+
+        Args:
+            task: The task input data.
+            resources: Dictionary of named resources including LLMs.
+            rollout: The rollout object with metadata.
+
+        Returns:
+            The result from the wrapped rollout function.
+        """
+        if self._is_async:
+            raise RuntimeError(f"{self._rollout_func} is asynchronous. Use rollout_async instead.")
+
+        kwargs = self._get_kwargs(resources, rollout)
+        return self._rollout_func(task, **kwargs)  # type: ignore
+
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Execute an asynchronous rollout using the wrapped function.
+
+        Args:
+            task: The task input data.
+            resources: Dictionary of named resources including LLMs.
+            rollout: The rollout object with metadata.
+
+        Returns:
+            The result from the wrapped rollout function.
+        """
+        if not self._is_async:
+            raise RuntimeError(f"{self._rollout_func} is synchronous. Use rollout instead.")
+
+        kwargs = self._get_kwargs(resources, rollout)
+        return await self._rollout_func(task, **kwargs)  # type: ignore
+
+    def _get_kwargs(self, resources: NamedResources, rollout: Rollout) -> Dict[str, Any]:
+        """Extract the kwargs needed for the rollout function based on its signature.
+
+        Dynamically builds the kwargs dictionary by inspecting the function signature and
+        including only the parameters the function accepts. This allows flexible function
+        signatures that can request any combination of: rollout, llm, and/or prompt_template.
+
+        Args:
+            resources: Dictionary of named resources available for the rollout.
+            rollout: The rollout object with metadata.
+
+        Returns:
+            A dictionary of kwargs to pass to the rollout function.
+        """
+
+        kwargs: Dict[str, Any] = {}
+        if self._accepts_rollout():
+            kwargs["rollout"] = rollout
+        if self._accepts_llm():
+            kwargs["llm"] = self._get_llm_resource(resources, rollout)
+        if self._accepts_prompt_template():
+            kwargs["prompt_template"] = self._get_prompt_template_resource(resources, rollout)
+
+        return kwargs
+
+    def _get_llm_resource(self, resources: NamedResources, rollout: Rollout) -> LLM:
+        """Extract the first LLM resource from the resources dictionary.
+
+        Strip the ProxyLLM resource into a LLM resource if needed.
+
+        Args:
+            resources: Dictionary of named resources.
+            rollout: The rollout object with metadata.
+
+        Returns:
+            The first LLM resource found.
+
+        Raises:
+            ValueError: If no LLM resource is found.
+        """
+        resource_found: LLM | None = None
+        for name, resource in resources.items():
+            if isinstance(resource, LLM):
+                if resource_found is not None:
+                    logger.warning(f"Multiple LLM resources found in resources. Using the first one: '{name}'.")
+                    break
+                resource_found = resource
+
+        if resource_found is None:
+            raise ValueError("No LLM resource found in the provided resources.")
+
+        if self._strip_proxy:
+            resource_found = self._strip_proxy_helper(resource_found, rollout)
+
+        return resource_found
+
+    def _get_prompt_template_resource(self, resources: NamedResources, rollout: Rollout) -> PromptTemplate:
+        """Extract the first PromptTemplate resource from the resources dictionary.
+
+        Args:
+            resources: Dictionary of named resources.
+            rollout: The rollout object with metadata. Not used in this method.
+
+        Returns:
+            The first PromptTemplate resource found.
+
+        Raises:
+            ValueError: If no PromptTemplate resource is found.
+        """
+        resource_found: PromptTemplate | None = None
+        for name, resource in resources.items():
+            if isinstance(resource, PromptTemplate):
+                if resource_found is not None:
+                    logger.warning(
+                        f"Multiple prompt template resources found in resources. Using the first one: '{name}'."
+                    )
+                    break
+                resource_found = resource
+
+        if resource_found is None:
+            raise ValueError("No prompt template resource found in the provided resources.")
+
+        return resource_found
+
+    def _strip_proxy_helper(self, proxy_llm: LLM, rollout: Rollout) -> LLM:
+        """Strip the ProxyLLM resource into a concrete LLM resource.
+
+        This method resolves ProxyLLM instances to their concrete LLM implementation
+        by attaching the attempted rollout context. This is only used when the function
+        signature accepts an 'llm' parameter and strip_proxy is True.
+
+        Args:
+            proxy_llm: The LLM resource, which may be a ProxyLLM.
+            rollout: The rollout object with metadata.
+
+        Returns:
+            The concrete LLM resource.
+
+        Raises:
+            ValueError: If the rollout is not an AttemptedRollout (required for stripping ProxyLLM).
+        """
+
+        if not isinstance(proxy_llm, ProxyLLM):
+            # Not a ProxyLLM, nothing to strip here.
+            return proxy_llm
+
+        # Rollout is still a Rollout here because API is not stabilized yet.
+        # In practice, it must be an AttemptedRollout.
+        if not isinstance(rollout, AttemptedRollout):
+            raise ValueError("Rollout is not an AttemptedRollout.")
+
+        return proxy_llm.with_attempted_rollout(rollout)
+
+
+@overload
+def llm_rollout(func: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
+
+
+@overload
+def llm_rollout(*, strip_proxy: bool = True) -> Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]: ...
+
+
+def llm_rollout(
+    func: LlmRolloutFunc[T] | None = None, *, strip_proxy: bool = True
+) -> FunctionalLitAgent[T] | Callable[[LlmRolloutFunc[T]], FunctionalLitAgent[T]]:
+    """Create a FunctionalLitAgent from a function that takes (task, llm[, rollout]).
+
+    This decorator allows you to define an agent using a simple function
+    instead of creating a full LitAgent subclass. The returned FunctionalLitAgent
+    instance is callable, preserving the original function's behavior.
+
+    Args:
+        func: A function that defines the agent's behavior. Can be:
+              - sync: (task, llm) -> result
+              - sync with rollout: (task, llm, rollout) -> result
+              - async: async (task, llm) -> result
+              - async with rollout: async (task, llm, rollout) -> result
+        strip_proxy: Whether to strip the ProxyLLM resource into a LLM resource.
+                     Defaults to True.
+
+    Returns:
+        A callable FunctionalLitAgent instance that preserves the original function's
+        type hints and behavior while providing all agent functionality.
+
+    Example:
+        @llm_rollout
+        def my_agent(task, llm):
+            # Agent logic here
+            return response
+
+        @llm_rollout(strip_proxy=False)
+        def my_agent_no_strip(task, llm):
+            # Agent logic here
+            return response
+
+        # Function is still callable with original behavior
+        result = my_agent(task, llm)
+
+        # Agent methods are also available
+        result = my_agent.rollout(task, resources, rollout)
+    """
+
+    def decorator(f: LlmRolloutFunc[T]) -> FunctionalLitAgent[T]:
+        _validate_llm_rollout_func(f)
+        return FunctionalLitAgent(f, strip_proxy=strip_proxy)
+
+    if func is None:
+        # Called with arguments: @llm_rollout(strip_proxy=False)
+        return decorator
+    else:
+        # Called without arguments: @llm_rollout
+        return decorator(func)
+
+
+def _validate_llm_rollout_func(func: Any) -> TypeGuard[LlmRolloutFunc[Any]]:
+    """Validate the function signature of a LLM rollout function.
+
+    Ensures the function follows the expected pattern for LLM-based rollouts:
+    - Must have at least 2 parameters
+    - First parameter must be named 'task'
+    - Must have a parameter named 'llm'
+    - Optionally can have a 'rollout' parameter
+
+    Args:
+        func: The function to validate.
+
+    Returns:
+        True if the function signature is valid.
+
+    Raises:
+        ValueError: If the function signature does not match the expected pattern.
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.keys())
+    if len(params) < 2:
+        raise ValueError(f"Function {func} must have at least 2 parameters.")
+    if params[0] != "task":
+        raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
+    if "llm" not in params:
+        raise ValueError(f"Function {func} must have a positional parameter called 'llm'.")
+
+    return True
+
+
+@overload
+def prompt_rollout(func: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]: ...
+
+
+@overload
+def prompt_rollout() -> Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]: ...
+
+
+def prompt_rollout(
+    func: PromptRolloutFunc[T] | None = None,
+) -> FunctionalLitAgent[T] | Callable[[PromptRolloutFunc[T]], FunctionalLitAgent[T]]:
+    """Create a FunctionalLitAgent from a function that takes (task, prompt_template[, rollout]).
+
+    This decorator is designed for agents that work with tunable prompt templates. It enables
+    a workflow where algorithms manage and optimize the prompt template, while agents consume
+    the template to perform rollouts. This is particularly useful for prompt optimization scenarios.
+
+    Args:
+        func: A function that defines the agent's behavior. Can be:
+              - sync: (task, prompt_template) -> result
+              - sync with rollout: (task, prompt_template, rollout) -> result
+              - async: async (task, prompt_template) -> result
+              - async with rollout: async (task, prompt_template, rollout) -> result
+
+    Returns:
+        A callable FunctionalLitAgent instance that preserves the original function's
+        type hints and behavior while providing all agent functionality.
+
+    Example:
+        @prompt_rollout
+        def my_agent(task, prompt_template):
+            # Use the prompt template to generate a response
+            messages = prompt_template.format(task=task.input)
+            # ... perform rollout with the formatted prompt
+            return response
+
+        # Function is still callable with original behavior
+        result = my_agent(task, prompt_template)
+
+        # Agent methods are also available
+        result = my_agent.rollout(task, resources, rollout)
+    """
+
+    def decorator(f: PromptRolloutFunc[T]) -> FunctionalLitAgent[T]:
+        _validate_prompt_rollout_func(f)
+        return FunctionalLitAgent(f)
+
+    if func is None:
+        return decorator
+    else:
+        return decorator(func)
+
+
+def _validate_prompt_rollout_func(func: Any) -> TypeGuard[PromptRolloutFunc[Any]]:
+    """Validate the function signature of a prompt rollout function.
+
+    Ensures the function follows the expected pattern for prompt-template-based rollouts:
+    - Must have at least 2 parameters
+    - First parameter must be named 'task'
+    - Must have a parameter named 'prompt_template'
+    - Optionally can have a 'rollout' parameter
+
+    Args:
+        func: The function to validate.
+
+    Returns:
+        True if the function signature is valid.
+
+    Raises:
+        ValueError: If the function signature does not match the expected pattern.
+    """
+    sig = inspect.signature(func)
+    params = list(sig.parameters.keys())
+    if len(params) < 2:
+        raise ValueError(f"Function {func} must have at least 2 parameters.")
+    if params[0] != "task":
+        raise ValueError(f"Function {func} must be a positional parameter called 'task'.")
+    if "prompt_template" not in params:
+        raise ValueError(f"Function {func} must have a positional parameter called 'prompt_template'.")
+
+    return True
+
+
+def rollout(func: Union[LlmRolloutFunc[T], PromptRolloutFunc[T], Callable[..., Any]]) -> FunctionalLitAgent[T]:
+    """Create a LitAgent from a function, automatically detecting the appropriate type.
+
+    This function inspects the provided callable and creates the appropriate
+    agent type based on its signature. It supports both LLM-based and prompt-template-based
+    agents. The returned agent instance is callable, preserving the original function's
+    behavior and type hints.
+
+    Args:
+        func: A function that defines the agent's behavior. Supported signatures:
+              - (task, llm[, rollout]) for LLM-based agents
+              - (task, prompt_template[, rollout]) for prompt-template-based agents
+
+    Returns:
+        A callable FunctionalLitAgent instance that preserves the original function's
+        type hints and behavior while providing all agent functionality.
+
+    Example:
+        # LLM-based agent
+        @rollout
+        def my_llm_agent(task, llm):
+            client = OpenAI(base_url=llm.endpoint)
+            response = client.chat.completions.create(
+                model=llm.model,
+                messages=[{"role": "user", "content": task.input}],
+            )
+            return response
+
+        # Prompt-template-based agent
+        @rollout
+        def my_prompt_agent(task, prompt_template):
+            messages = prompt_template.format(task=task.input)
+            # ... perform rollout with the formatted prompt
+            return response
+
+        # Function is still callable with original behavior
+        result = my_llm_agent(task, llm)
+
+        # Agent methods are also available
+        result = my_llm_agent.rollout(task, resources, rollout)
+
+    Raises:
+        NotImplementedError: If the function signature doesn't match any known patterns.
+    """
+    # Check if it matches the LLM rollout API pattern
+    sig = inspect.signature(func)
+
+    try:
+        if _validate_llm_rollout_func(func):
+            return llm_rollout(func)
+    except ValueError:
+        pass
+
+    try:
+        if _validate_prompt_rollout_func(func):
+            return prompt_rollout(func)
+    except ValueError:
+        pass
+
+    raise NotImplementedError(
+        f"Function signature {sig} does not match any known agent patterns. "
+        "Expected signatures: (task, llm[, rollout]) or (task, prompt_template[, rollout]). "
+        "Functions can be sync or async."
+    )

agentlightning/litagent/litagent.py

@@ -0,0 +1,311 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import logging
+import warnings
+import weakref
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
+
+from agentlightning.types import NamedResources, Rollout, RolloutRawResult, Task
+
+if TYPE_CHECKING:
+    from agentlightning.runner import BaseRunner
+    from agentlightning.tracer import BaseTracer
+    from agentlightning.trainer import Trainer
+
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+__all__ = [
+    "LitAgent",
+]
+
+
+def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
+    """Check if the rollout API is v0.1.
+    Inspect the function signature to see if it has a rollout_id parameter.
+
+    Args:
+        func: The function to check.
+    """
+    return "rollout_id" in inspect.signature(func).parameters
+
+
+class LitAgent(Generic[T]):
+    """Base class for the training and validation logic of an agent.
+
+    Developers should subclass this class and implement the rollout methods
+    to define the agent's behavior for a single task. The agent's logic
+    is completely decoupled from the server communication and training
+    infrastructure.
+    """
+
+    def __init__(self, *, trained_agents: Optional[str] = None) -> None:  # FIXME: str | None won't work for cli
+        """
+        Initialize the LitAgent.
+
+        Args:
+            trained_agents: Optional string representing the trained agents.
+                            This can be used to track which agents have been trained by this instance.
+                            Deprecated. Configure `agent_match` in adapter instead.
+        """
+        if trained_agents is not None:
+            warnings.warn(
+                "`trained_agents` is deprecated. Configure `agent_match` in adapter instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        self.trained_agents = trained_agents
+
+        self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
+        self._runner_ref: weakref.ReferenceType[BaseRunner[T]] | None = None
+
+    def is_async(self) -> bool:
+        """
+        Check if the agent implements asynchronous rollout methods.
+        Override this property for customized async detection logic.
+
+        Returns:
+            True if the agent has custom async rollout methods, False otherwise.
+        """
+        return (
+            (
+                hasattr(self, "training_rollout_async")
+                and self.__class__.training_rollout_async is not LitAgent.training_rollout_async  # type: ignore
+            )
+            or (
+                hasattr(self, "validation_rollout_async")
+                and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async  # type: ignore
+            )
+            or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async)  # type: ignore
+        )
+
+    def set_trainer(self, trainer: Trainer) -> None:
+        """
+        Set the trainer for this agent.
+
+        Args:
+            trainer: The Trainer instance that will handle training and validation.
+        """
+        self._trainer_ref = weakref.ref(trainer)
+
+    def get_trainer(self) -> Trainer:
+        """
+        Get the trainer for this agent.
+
+        Returns:
+            The Trainer instance associated with this agent.
+        """
+        if self._trainer_ref is None:
+            raise ValueError("Trainer has not been set for this agent.")
+        trainer = self._trainer_ref()
+        if trainer is None:
+            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
+        return trainer
+
+    @property
+    def trainer(self) -> Trainer:
+        """Convenient shortcut of self.get_trainer()."""
+        return self.get_trainer()
+
+    def get_tracer(self) -> BaseTracer:
+        """
+        Get the tracer for this agent.
+
+        Returns:
+            The BaseTracer instance associated with this agent.
+        """
+        if hasattr(self.runner, "tracer"):
+            return self.runner.tracer  # type: ignore
+        else:
+            return self.trainer.tracer
+
+    @property
+    def tracer(self) -> BaseTracer:
+        """Convenient shortcut of self.get_tracer()."""
+        return self.get_tracer()
+
+    def set_runner(self, runner: BaseRunner[T]) -> None:
+        """
+        Set the runner for this agent.
+
+        Args:
+            runner: The runner instance that will handle the execution of rollouts.
+        """
+        self._runner_ref = weakref.ref(runner)
+
+    def get_runner(self) -> BaseRunner[T]:
+        """
+        Get the runner for this agent.
+
+        Returns:
+            The runner instance associated with this agent.
+        """
+        if self._runner_ref is None:
+            raise ValueError("Runner has not been set for this agent.")
+        runner = self._runner_ref()
+        if runner is None:
+            raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
+        return runner
+
+    @property
+    def runner(self) -> BaseRunner[T]:
+        """Convenient shortcut of self.get_runner()."""
+        return self.get_runner()
+
+    def on_rollout_start(self, task: Task, runner: BaseRunner[T], tracer: BaseTracer) -> None:
+        """Hook called immediately before a rollout begins.
+
+        Deprecated in favor of `on_rollout_start` in the `Hook` interface.
+
+        Args:
+            task: The :class:`Task` object that will be processed.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The tracer instance associated with the runner.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: BaseRunner[T], tracer: BaseTracer) -> None:
+        """Hook called after a rollout completes.
+
+        Deprecated in favor of `on_rollout_end` in the `Hook` interface.
+
+        Args:
+            task: The :class:`Task` object that was processed.
+            rollout: The resulting :class:`Rollout` object.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The tracer instance associated with the runner.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """
+
+    def rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Main entry point for executing a rollout.
+
+        This method determines whether to call the synchronous or
+        asynchronous rollout method based on the agent's implementation.
+
+        If you don't wish to implement both training rollout and validation
+        rollout separately, you can just implement `rollout` which will work for both.
+
+        Args:
+            task: The task object received from the server, containing the
+                  input data and metadata.
+            resources: A dictionary of named resources (e.g., LLMs, prompt
+                       templates) for the agent to use.
+            rollout: The full rollout object, please avoid from directly modifying it.
+                     Most agents should only use `task` and `resources`. Use `rollout`
+                     only if you need to access metadata like `rollout_id`.
+
+        Returns:
+            The result of the rollout, which can be one of:
+            - None. The tracing should be handled by the agent runner.
+            - A float representing the final reward.
+            - A list of `Triplet` objects for detailed, step-by-step feedback.
+            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
+            - A list of dictionaries for any trace spans.
+            - A complete `Rollout` object for full control over reporting.
+        """
+        raise NotImplementedError("Agents must implement the `rollout` method.")
+
+    async def rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Asynchronous version of the main rollout method.
+
+        This method determines whether to call the synchronous or
+        asynchronous rollout method based on the agent's implementation.
+
+        Args:
+            task: The task object received from the server, containing the
+                  input data and metadata.
+            resources: A dictionary of named resources (e.g., LLMs, prompt
+                       templates) for the agent to use.
+            rollout: The full rollout object, please avoid from directly modifying it.
+                     Most agents should only use `task` and `resources`. Use `rollout`
+                     only if you need to access metadata like `rollout_id`.
+
+        Returns:
+            The result of the rollout, which can be one of:
+            - None. The tracing should be handled by the agent runner.
+            - A float representing the final reward.
+            - A list of `Triplet` objects for detailed, step-by-step feedback.
+            - A list of `ReadableSpan` objects for OpenTelemetry tracing.
+            - A list of dictionaries for any trace spans.
+            - A complete `Rollout` object for full control over reporting.
+        """
+        raise NotImplementedError("Agents must implement the `rollout_async` method for async operations.")
+
+    def training_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Defines the agent's behavior for a single training task.
+
+        This method should contain the logic for how the agent processes an
+        input, uses the provided resources (like LLMs or prompts), and
+        produces a result.
+
+        Args:
+            task: The task object received from the server, containing the
+                  input data and metadata.
+            resources: A dictionary of named resources (e.g., LLMs, prompt
+                       templates) for the agent to use.
+            rollout: The full rollout object, please avoid from directly modifying it.
+        """
+        return self.rollout(task, resources, rollout)
+
+    def validation_rollout(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Defines the agent's behavior for a single validation task.
+
+        By default, this method redirects to `training_rollout`. Override it
+        if the agent should behave differently during validation.
+
+        Args:
+            task: The task object received from the server, containing the
+                  input data and metadata.
+            resources: A dictionary of named resources for the agent to use.
+            rollout: The full rollout object, avoid from modifying it.
+
+        Returns:
+            The result of the validation rollout. See `rollout` for
+            possible return types.
+        """
+        return self.rollout(task, resources, rollout)
+
+    async def training_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Asynchronous version of `training_rollout`.
+
+        This method should be implemented by agents that perform asynchronous
+        operations (e.g., non-blocking I/O, concurrent API calls).
+
+        Args:
+            task: The task object received from the server.
+            resources: A dictionary of named resources for the agent to use.
+            rollout: The full rollout object, avoid from modifying it.
+
+        Returns:
+            The result of the asynchronous training rollout. See `rollout` for
+            possible return types.
+        """
+        return await self.rollout_async(task, resources, rollout)
+
+    async def validation_rollout_async(self, task: T, resources: NamedResources, rollout: Rollout) -> RolloutRawResult:
+        """Asynchronous version of `validation_rollout`.
+
+        By default, this method redirects to `training_rollout_async`.
+        Override it for different asynchronous validation behavior.
+
+        Args:
+            task: The task object received from the server.
+            resources: A dictionary of named resources for the agent to use.
+            rollout: The full rollout object, avoid from modifying it.
+
+        Returns:
+            The result of the asynchronous validation rollout. See `rollout` for
+            possible return types.
+        """
+        return await self.rollout_async(task, resources, rollout)

agentlightning/llm_proxy.py

@@ -0,0 +1,765 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import logging
+import os
+import re
+import socket
+import tempfile
+import threading
+import time
+from typing import Any, Awaitable, Callable, Dict, Iterable, List, Optional, Sequence, TypedDict, Union, cast
+
+import litellm
+import opentelemetry.trace as trace_api
+import uvicorn
+import yaml
+from fastapi import Request, Response
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from litellm.proxy.proxy_server import app, save_worker_config  # pyright: ignore[reportUnknownVariableType]
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from agentlightning.types import LLM, ProxyLLM
+
+from .store.base import LightningStore
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "LLMProxy",
+]
+
+
+class ModelConfig(TypedDict):
+    """LiteLLM model registration entry.
+
+    This mirrors the items in LiteLLM's ``model_list`` section.
+
+    Attributes:
+        model_name: Logical model name exposed by the proxy.
+        litellm_params: Parameters passed to LiteLLM for this model
+            (e.g., backend model id, api_base, additional options).
+    """  # Google style kept concise.
+
+    model_name: str
+    litellm_params: Dict[str, Any]
+
+
+def _get_pre_call_data(args: Any, kwargs: Any) -> Dict[str, Any]:
+    """Extract LiteLLM request payload from hook args.
+
+    The LiteLLM logger hooks receive ``(*args, **kwargs)`` whose third positional
+    argument or ``data=`` kwarg contains the request payload.
+
+    Args:
+        args: Positional arguments from the hook.
+        kwargs: Keyword arguments from the hook.
+
+    Returns:
+        The request payload dict.
+
+    Raises:
+        ValueError: If the payload cannot be located or is not a dict.
+    """
+    if kwargs.get("data"):
+        data = kwargs["data"]
+    elif len(args) >= 3:
+        data = args[2]
+    else:
+        raise ValueError(f"Unable to get request data from args or kwargs: {args}, {kwargs}")
+    if not isinstance(data, dict):
+        raise ValueError(f"Request data is not a dictionary: {data}")
+    return cast(Dict[str, Any], data)
+
+
+# We need global state because litellm is based on a global app.
+# Repeatedly initializing the app with different stores will cause errors.
+_initialized: bool = False
+_global_store: LightningStore | None = None
+
+
+def get_global_store() -> LightningStore:
+    """Return the globally registered LightningStore.
+
+    Used by components that are initialized without an explicit store
+    (e.g., exporter created inside OpenTelemetry).
+
+    Returns:
+        LightningStore: The active global store.
+
+    Raises:
+        ValueError: If the global store has not been set by ``LLMProxy.start()``.
+    """
+    if _global_store is None:
+        raise ValueError("Global store is not initialized. Please start a LLMProxy first.")
+    return _global_store
+
+
+def initialize() -> None:
+    """Initialize global middleware and LiteLLM callbacks once.
+
+    Idempotent. Installs:
+
+    * A FastAPI middleware that rewrites /rollout/{rid}/attempt/{aid}/... paths,
+      injects rollout/attempt/sequence headers, and forwards downstream.
+    * LiteLLM callbacks for token ids and OpenTelemetry export.
+
+    This function does not start any server. It only wires global hooks.
+    """
+    global _initialized
+    if _initialized:
+        return
+
+    # Add middleware here because it relies on the global store.
+    @app.middleware("http")
+    async def rollout_attempt_middleware(  # pyright: ignore[reportUnusedFunction]
+        request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        # Decode rollout and attempt from the URL prefix. Example:
+        #   /rollout/r123/attempt/a456/v1/chat/completions
+        # becomes
+        #   /v1/chat/completions
+        # while adding request-scoped headers for trace attribution.
+        path = request.url.path
+
+        match = re.match(r"^/rollout/([^/]+)/attempt/([^/]+)(/.*)?$", path)
+        if match:
+            rollout_id = match.group(1)
+            attempt_id = match.group(2)
+            new_path = match.group(3) if match.group(3) is not None else "/"
+
+            # Rewrite the ASGI scope path so downstream sees a clean OpenAI path.
+            request.scope["path"] = new_path
+            request.scope["raw_path"] = new_path.encode()
+
+            # Allocate a monotonic sequence id per (rollout, attempt).
+            sequence_id = await get_global_store().get_next_span_sequence_id(rollout_id, attempt_id)
+
+            # Inject headers so downstream components and exporters can retrieve them.
+            request.scope["headers"] = list(request.scope["headers"]) + [
+                (b"x-rollout-id", rollout_id.encode()),
+                (b"x-attempt-id", attempt_id.encode()),
+                (b"x-sequence-id", str(sequence_id).encode()),
+            ]
+
+        response = await call_next(request)
+        return response
+
+    # Register callbacks once on the global LiteLLM callback list.
+    litellm.callbacks.extend(  # pyright: ignore[reportUnknownMemberType]
+        [
+            AddReturnTokenIds(),
+            LightningOpenTelemetry(),
+        ]
+    )
+
+    _initialized = True
+
+
+class AddReturnTokenIds(CustomLogger):
+    """LiteLLM logger hook to request token ids from vLLM.
+
+    This mutates the outgoing request payload to include ``return_token_ids=True``
+    for backends that support token id return (e.g., vLLM).
+
+    See:
+        https://github.com/vllm-project/vllm/pull/22587
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload.
+
+        Args:
+            args: Positional args from LiteLLM.
+            kwargs: Keyword args from LiteLLM.
+
+        Returns:
+            Either an updated payload dict or an Exception to short-circuit.
+        """
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure token ids are requested from the backend when supported.
+        return {**data, "return_token_ids": True}
+
+
+class LightningSpanExporter(SpanExporter):
+    """Buffered OTEL span exporter with subtree flushing and training-store sink.
+
+    Design:
+
+    * Spans are buffered until a root span's entire subtree is available.
+    * A private event loop on a daemon thread runs async flush logic.
+    * Rollout/attempt/sequence metadata is reconstructed by merging headers
+      from any span within a subtree.
+
+    Thread-safety:
+
+    * Buffer access is protected by a re-entrant lock.
+    * Export is synchronous to the caller yet schedules an async flush on the
+      internal loop, then waits for completion.
+
+    Args:
+        store: Optional explicit LightningStore. If None, uses ``get_global_store()``.
+    """
+
+    def __init__(self, store: Optional[LightningStore] = None):
+        self._store = store
+        self._buffer: List[ReadableSpan] = []
+        self._lock: Optional[threading.RLock] = None
+
+        # Single dedicated event loop running in a daemon thread.
+        # This decouples OTEL SDK threads from our async store I/O.
+        # Deferred creation until first use.
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+        """Lazily initialize the event loop and thread on first use.
+
+        Returns:
+            asyncio.AbstractEventLoop: The initialized event loop.
+        """
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
+            self._loop_thread.start()
+        return self._loop
+
+    def _ensure_lock(self) -> threading.RLock:
+        """Lazily initialize the lock on first use.
+
+        Returns:
+            threading.RLock: The initialized lock.
+        """
+        if self._lock is None:
+            self._lock = threading.RLock()
+        return self._lock
+
+    def _get_store(self) -> LightningStore:
+        """Return the LightningStore to use.
+
+        Returns:
+            LightningStore: Explicit store if provided, else the global store.
+
+        Raises:
+            ValueError: If no global store is configured and no explicit store was given.
+        """
+        if self._store is None:
+            return get_global_store()
+        return self._store
+
+    def _run_loop(self) -> None:
+        """Run the private asyncio loop forever on the exporter thread."""
+        assert self._loop is not None, "Loop should be initialized before thread starts"
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def shutdown(self) -> None:
+        """Shut down the exporter event loop.
+
+        Safe to call at process exit.
+
+        """
+        if self._loop is None:
+            return
+
+        try:
+
+            def _stop():
+                assert self._loop is not None
+                self._loop.stop()
+
+            self._loop.call_soon_threadsafe(_stop)
+            if self._loop_thread is not None:
+                self._loop_thread.join(timeout=2.0)
+            self._loop.close()
+        except Exception:
+            logger.exception("Error during exporter shutdown")
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans via buffered subtree flush.
+
+        Appends spans to the internal buffer, then triggers an async flush on the
+        private event loop. Blocks until that flush completes.
+
+        Args:
+            spans: Sequence of spans to export.
+
+        Returns:
+            SpanExportResult: SUCCESS on flush success, else FAILURE.
+        """
+        # Buffer append under lock to protect against concurrent exporters.
+        with self._ensure_lock():
+            for span in spans:
+                self._buffer.append(span)
+
+        # Run the async flush on our private loop, synchronously from caller's POV.
+        async def _locked_flush():
+            # Take the lock inside the coroutine to serialize with other flushes.
+            with self._ensure_lock():
+                return await self._maybe_flush()
+
+        try:
+            loop = self._ensure_loop()
+            fut = asyncio.run_coroutine_threadsafe(_locked_flush(), loop)
+            fut.result()  # Bubble up any exceptions from the coroutine.
+        except Exception as e:
+            logger.exception("Export flush failed: %s", e)
+            return SpanExportResult.FAILURE
+
+        return SpanExportResult.SUCCESS
+
+    async def _maybe_flush(self):
+        """Flush ready subtrees from the buffer.
+
+        Strategy:
+            We consider a subtree "ready" if we can identify a root span. We
+            then take that root and all its descendants out of the buffer and
+            try to reconstruct rollout/attempt/sequence headers by merging any
+            span's ``metadata.requester_custom_headers`` within the subtree.
+
+        Required headers:
+            ``x-rollout-id`` (str), ``x-attempt-id`` (str), ``x-sequence-id`` (str of int)
+
+        Raises:
+            None directly. Logs and skips malformed spans.
+
+        """
+        # Iterate over current roots. Each iteration pops a whole subtree.
+        for root_span_id in self._get_root_span_ids():
+            subtree_spans = self._pop_subtrees(root_span_id)
+            if not subtree_spans:
+                continue
+
+            # Merge all custom headers found in the subtree.
+            headers_merged: Dict[str, Any] = {}
+
+            for span in subtree_spans:
+                if span.attributes is None:
+                    continue
+                headers_str = span.attributes.get("metadata.requester_custom_headers")
+                if headers_str is None:
+                    continue
+                if not isinstance(headers_str, str):
+                    logger.error(
+                        f"metadata.requester_custom_headers is not stored as a string: {headers_str}. Skipping the span."
+                    )
+                    continue
+                try:
+                    # Use literal_eval to parse the stringified dict safely.
+                    headers = ast.literal_eval(headers_str)
+                except Exception as e:
+                    logger.error(
+                        f"Failed to parse metadata.requester_custom_headers: {headers_str}, error: {e}. Skipping the span."
+                    )
+                    continue
+                if not isinstance(headers, dict):
+                    logger.error(
+                        f"metadata.requester_custom_headers is not parsed as a dict: {headers}. Skipping the span."
+                    )
+                    continue
+                headers_merged.update(cast(Dict[str, Any], headers))
+
+            if not headers_merged:
+                logger.warning(f"No headers found in {len(subtree_spans)} subtree spans. Cannot log to store.")
+                continue
+
+            # Validate and normalize required header fields.
+            rollout_id = headers_merged.get("x-rollout-id")
+            attempt_id = headers_merged.get("x-attempt-id")
+            sequence_id = headers_merged.get("x-sequence-id")
+            if not rollout_id or not attempt_id or not sequence_id or not sequence_id.isdigit():
+                logger.warning(
+                    f"Missing or invalid rollout_id, attempt_id, or sequence_id in headers: {headers_merged}. Cannot log to store."
+                )
+                continue
+            if not isinstance(rollout_id, str) or not isinstance(attempt_id, str):
+                logger.warning(
+                    f"rollout_id or attempt_id is not a string: {rollout_id}, {attempt_id}. Cannot log to store."
+                )
+                continue
+            sequence_id_decimal = int(sequence_id)
+
+            # Persist each span in the subtree with the resolved identifiers.
+            for span in subtree_spans:
+                await self._get_store().add_otel_span(
+                    rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id_decimal, readable_span=span
+                )
+
+    def _get_root_span_ids(self) -> Iterable[int]:
+        """Yield span_ids for root spans currently in the buffer.
+
+        A root span is defined as one with ``parent is None``.
+
+        Yields:
+            int: Span id for each root span found.
+        """
+        for span in self._buffer:
+            if span.parent is None:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    yield span_context.span_id
+
+    def _get_subtrees(self, root_span_id: int) -> Iterable[int]:
+        """Yield span_ids in the subtree rooted at ``root_span_id``.
+
+        Depth-first traversal over the current buffer.
+
+        Args:
+            root_span_id: The span id of the root.
+
+        Yields:
+            int: Span ids including the root and all descendants found.
+        """
+        # Yield the root span id first.
+        yield root_span_id
+        for span in self._buffer:
+            # Check whether the span's parent is the root_span_id.
+            if span.parent is not None and span.parent.span_id == root_span_id:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    # Recursively get child spans.
+                    yield from self._get_subtrees(span_context.span_id)
+
+    def _pop_subtrees(self, root_span_id: int) -> List[ReadableSpan]:
+        """Remove and return the subtree for a particular root from the buffer.
+
+        Args:
+            root_span_id: Root span id identifying the subtree.
+
+        Returns:
+            list[ReadableSpan]: Spans that were part of the subtree. Order follows buffer order.
+        """
+        subtree_span_ids = set(self._get_subtrees(root_span_id))
+        subtree_spans: List[ReadableSpan] = []
+        new_buffer: List[ReadableSpan] = []
+        for span in self._buffer:
+            span_context = span.get_span_context()
+            if span_context is not None and span_context.span_id in subtree_span_ids:
+                subtree_spans.append(span)
+            else:
+                new_buffer.append(span)
+        # Replace buffer with remaining spans to avoid re-processing.
+        self._buffer = new_buffer
+        return subtree_spans
+
+
+class LightningOpenTelemetry(OpenTelemetry):
+    """OpenTelemetry integration that exports spans to the Lightning store.
+
+    Responsibilities:
+
+    * Ensures each request is annotated with a per-attempt sequence id so spans
+      are ordered deterministically even with clock skew across nodes.
+    * Uses ``LightningSpanExporter`` to persist spans for analytics and training.
+
+    Args:
+        store: Optional explicit LightningStore for the exporter.
+    """
+
+    def __init__(self, store: LightningStore | None = None):
+        config = OpenTelemetryConfig(exporter=LightningSpanExporter(store))
+
+        # Check for tracer initialization
+        if (
+            hasattr(trace_api, "_TRACER_PROVIDER")
+            and trace_api._TRACER_PROVIDER is not None  # pyright: ignore[reportPrivateUsage]
+        ):
+            logger.error("Tracer is already initialized. OpenTelemetry may not work as expected.")
+
+        super().__init__(config=config)  # pyright: ignore[reportUnknownMemberType]
+
+
+class LLMProxy:
+    """Host a LiteLLM OpenAI-compatible proxy bound to a LightningStore.
+
+    The proxy:
+
+    * Serves an OpenAI-compatible API via uvicorn.
+    * Adds rollout/attempt routing and headers via middleware.
+    * Registers OTEL export and token-id callbacks.
+    * Writes a LiteLLM worker config file with ``model_list`` and settings.
+
+    Lifecycle:
+
+    * ``start()`` writes config, starts uvicorn server in a thread, and waits until ready.
+    * ``stop()`` tears down the server and removes the temp config file.
+    * ``restart()`` convenience wrapper to stop then start.
+
+    Usage Note:
+    As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
+    process from the main runner (i.e., tracer from agents).
+
+    Args:
+        port: TCP port to bind.
+        model_list: LiteLLM ``model_list`` entries.
+        store: LightningStore used for span sequence and persistence.
+        host: Publicly reachable host used in resource endpoints. Defaults to best-guess IPv4.
+        litellm_config: Extra LiteLLM proxy config merged with ``model_list``.
+        num_retries: Default LiteLLM retry count injected into ``litellm_settings``.
+    """
+
+    def __init__(
+        self,
+        port: int,
+        model_list: List[ModelConfig] | None = None,
+        store: Optional[LightningStore] = None,
+        host: str | None = None,
+        litellm_config: Dict[str, Any] | None = None,
+        num_retries: int = 0,
+    ):
+        self.store = store
+        self.host = host or _get_default_ipv4_address()
+        self.port = port
+        self.model_list = model_list or []
+        self.litellm_config = litellm_config or {}
+
+        # Ensure num_retries is present inside the litellm_settings block.
+        self.litellm_config.setdefault("litellm_settings", {})
+        self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
+
+        self._server_thread = None
+        self._config_file = None
+        self._uvicorn_server = None
+        self._ready_event = threading.Event()
+
+    def set_store(self, store: LightningStore) -> None:
+        """Set the store for the proxy.
+
+        Args:
+            store: The store to use for the proxy.
+        """
+        self.store = store
+
+    def update_model_list(self, model_list: List[ModelConfig]) -> None:
+        """Replace the in-memory model list and hot-restart if running.
+
+        Args:
+            model_list: New list of model entries.
+        """
+        self.model_list = model_list
+        logger.info(f"Updating LLMProxy model list to: {model_list}")
+        if self.is_running():
+            self.restart()
+        # Do nothing if the server is not running.
+
+    def update_port(self, port: int) -> None:
+        """Update the port for the proxy.
+
+        Args:
+            port: The new port to use for the proxy.
+        """
+        self.port = port
+
+    def _wait_until_started(self, startup_timeout: float = 20.0):
+        """Block until the uvicorn server reports started or timeout.
+
+        Args:
+            startup_timeout: Maximum seconds to wait.
+        """
+        start = time.time()
+        while True:
+            if self._uvicorn_server is None:
+                break
+            if self._uvicorn_server.started:
+                self._ready_event.set()
+                break
+            if self._uvicorn_server.should_exit:
+                break
+            if time.time() - start > startup_timeout:
+                break
+            time.sleep(0.01)
+
+    def start(self):
+        """Start the proxy server thread and initialize global wiring.
+
+        Side effects:
+
+        * Sets the module-level global store for middleware/exporter access.
+        * Calls ``initialize()`` once to register middleware and callbacks.
+        * Writes a temporary YAML config consumed by LiteLLM worker.
+        * Launches uvicorn in a daemon thread and waits for readiness.
+        """
+        if self.is_running():
+            # Trigger restart
+            self.stop()
+
+        if not self.store:
+            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
+
+        global _global_store
+
+        _global_store = self.store
+
+        # Initialize global middleware and callbacks once.
+        initialize()
+
+        # Persist a temp worker config for LiteLLM and point the proxy at it.
+        self._config_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False).name
+        with open(self._config_file, "w") as fp:
+            yaml.safe_dump(
+                {
+                    "model_list": self.model_list,
+                    **self.litellm_config,
+                },
+                fp,
+            )
+
+        save_worker_config(config=self._config_file)
+
+        # Bind to all interfaces to allow other hosts to reach it if needed.
+        self._uvicorn_server = uvicorn.Server(uvicorn.Config(app, host="0.0.0.0", port=self.port))
+
+        def run_server():
+            # Serve uvicorn in this background thread with its own event loop.
+            assert self._uvicorn_server is not None
+            asyncio.run(self._uvicorn_server.serve())
+
+        logger.info("Starting LLMProxy server thread...")
+        self._ready_event.clear()
+        # FIXME: This thread should either be reused or the whole proxy should live in another process.
+        # Problem 1: in litellm worker, <Queue at 0x70f1d028cd90 maxsize=50000> is bound to a different event loop
+        # Problem 2: Proxy has conflicted opentelemetry setup with the main process.
+        self._server_thread = threading.Thread(target=run_server, daemon=True)
+        self._server_thread.start()
+        self._wait_until_started()
+
+    def stop(self):
+        """Stop the proxy server and clean up temporary artifacts.
+
+        This is a best-effort graceful shutdown with a bounded join timeout.
+        """
+        if not self.is_running():
+            logger.warning("LLMProxy is not running. Nothing to stop.")
+            return
+
+        # Remove worker config to avoid stale references.
+        if self._config_file and os.path.exists(self._config_file):
+            os.unlink(self._config_file)
+
+        logger.info("Stopping LLMProxy server thread...")
+        stop_success = True
+        if self._server_thread is not None and self._uvicorn_server is not None and self._uvicorn_server.started:
+            self._uvicorn_server.should_exit = True
+            self._server_thread.join(timeout=10.0)  # Allow time for graceful shutdown.
+            if self._server_thread.is_alive():
+                logger.error(
+                    "LLMProxy server thread is still alive after 10 seconds. Cannot kill it because it's a thread."
+                )
+                stop_success = False
+            self._server_thread = None
+            self._uvicorn_server = None
+            self._config_file = None
+            self._ready_event.clear()
+            if not _check_port(self.host, self.port):
+                logger.error(f"Port {self.port} is still in use. Stopping LLMProxy is not successful.")
+                stop_success = False
+        if stop_success:
+            logger.info("LLMProxy server thread stopped.")
+        else:
+            logger.error("LLMProxy server is not stopped successfully.")
+
+    def restart(self, *, _port: int | None = None) -> None:
+        """Restart the proxy if running, else start it.
+
+        Convenience wrapper calling ``stop()`` followed by ``start()``.
+        """
+        logger.info("Restarting LLMProxy server...")
+        if self.is_running():
+            self.stop()
+        if _port is not None:
+            self.port = _port
+        self.start()
+
+    def is_running(self) -> bool:
+        """Return whether the uvicorn server is active.
+
+        Returns:
+            bool: True if server was started and did not signal exit.
+        """
+        return self._uvicorn_server is not None and self._uvicorn_server.started
+
+    def as_resource(
+        self,
+        rollout_id: str | None = None,
+        attempt_id: str | None = None,
+        model: str | None = None,
+        sampling_parameters: Dict[str, Any] | None = None,
+    ) -> LLM:
+        """Create an ``LLM`` resource pointing at this proxy with rollout context.
+
+        The returned endpoint is:
+            ``http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}``
+
+        Args:
+            rollout_id: Rollout identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            attempt_id: Attempt identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            model: Logical model name to use. If omitted and exactly one model
+                is configured, that model is used.
+            sampling_parameters: Optional default sampling parameters.
+
+        Returns:
+            LLM: Configured resource ready for OpenAI-compatible calls.
+
+        Raises:
+            ValueError: If ``model`` is omitted and zero or multiple models are configured.
+        """
+        if model is None:
+            if len(self.model_list) == 1:
+                model = self.model_list[0]["model_name"]
+            else:
+                raise ValueError(
+                    f"Multiple or zero models found in model_list: {self.model_list}. Please specify the model."
+                )
+
+        if rollout_id is None and attempt_id is None:
+            return ProxyLLM(
+                endpoint=f"http://{self.host}:{self.port}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        elif rollout_id is not None and attempt_id is not None:
+            return LLM(
+                endpoint=f"http://{self.host}:{self.port}/rollout/{rollout_id}/attempt/{attempt_id}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        else:
+            raise ValueError("Either rollout_id and attempt_id must be provided, or neither.")
+
+
+def _get_default_ipv4_address() -> str:
+    """Determine the default outbound IPv4 address for this machine.
+
+    Implementation:
+        Opens a UDP socket and "connects" to a public address to force route
+        selection, then inspects the socket's local address. No packets are sent.
+
+    Returns:
+        str: Best-guess IPv4 like ``192.168.x.y``. Falls back to ``127.0.0.1``.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # Doesn't actually contact 8.8.8.8; just forces the OS to pick a route.
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except Exception:
+        return "127.0.0.1"
+    finally:
+        s.close()
+
+
+def _check_port(host: str, port: int) -> bool:
+    """Check if a port is available."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.settimeout(1)
+        result = s.connect_ex((host, port))
+        return result != 0  # True if unavailable

agentlightning/logging.py

@@ -1,5 +1,9 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 import logging
 
+__all__ = ["configure_logger"]
+
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
     logger = logging.getLogger(name)

agentlightning/reward.py

@@ -1,66 +1,7 @@
-import asyncio
-import inspect
+# Copyright (c) Microsoft. All rights reserved.
+
 import warnings
-from typing import TypedDict, Optional
 
-from agentops.sdk.decorators import operation
+from .emitter.reward import *  # noqa: F401,F403
 
-
-class RewardSpanData(TypedDict):
-    type: "reward"
-    value: Optional[float]
-
-
-def reward(fn: callable) -> callable:
-    """
-    A decorator to wrap a function that computes rewards.
-    It will automatically handle the input and output of the function.
-    """
-
-    def wrap_result(result: Optional[float]) -> RewardSpanData:
-        """
-        Wrap the result of the function in a dict.
-        """
-        if result is None:
-            return {"type": "reward", "value": None}
-        if not isinstance(result, (float, int)):
-            warnings.warn(f"Reward is ignored because it is not a number: {result}")
-            return {"type": "reward", "value": None}
-        return {"type": "reward", "value": float(result)}
-
-    # Check if the function is async
-    is_async = asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn)
-
-    if is_async:
-
-        async def wrapper_async(*args, **kwargs):
-            result: Optional[float] = None
-
-            @operation
-            async def agentops_reward_operation() -> RewardSpanData:
-                # The reward function we are interested in tracing
-                # It takes zero inputs and return a formatted dict
-                nonlocal result
-                result = await fn(*args, **kwargs)
-                return wrap_result(result)
-
-            await agentops_reward_operation()
-            return result
-
-        return wrapper_async
-
-    else:
-
-        def wrapper(*args, **kwargs):
-            result: Optional[float] = None
-
-            @operation
-            def agentops_reward_operation() -> RewardSpanData:
-                nonlocal result
-                result = fn(*args, **kwargs)
-                return wrap_result(result)
-
-            agentops_reward_operation()
-            return result
-
-        return wrapper
+warnings.warn("agentlightning.reward is deprecated. Please use agentlightning.emitter instead.")

agentlightning/runner/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .agent import LitAgentRunner
+from .base import BaseRunner
+from .legacy import LegacyAgentRunner
+
+__all__ = [
+    "BaseRunner",
+    "LegacyAgentRunner",
+    "LitAgentRunner",
+]

agentlightning/runner/agent.py

@@ -0,0 +1,533 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent runner implementation for executing agent rollouts.
+
+This module provides the concrete implementation of the runner interface,
+handling the execution of agent rollouts with support for tracing, hooks,
+and distributed worker coordination.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, TypeVar, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.litagent import LitAgent
+from agentlightning.reward import emit_reward, find_final_reward
+from agentlightning.store.base import LightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import (
+    AttemptedRollout,
+    Hook,
+    NamedResources,
+    Rollout,
+    RolloutMode,
+    RolloutRawResult,
+    Span,
+)
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import ExecutionEvent
+
+from .base import BaseRunner
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class LitAgentRunner(BaseRunner[T_task]):
+    """Runner implementation for executing agent tasks with distributed support.
+
+    This runner manages the complete lifecycle of agent rollout execution,
+    including task polling, resource management, tracing, and hooks. It supports
+    both continuous iteration over tasks from the store and single-step execution.
+
+    Attributes:
+        worker_id: The unique identifier for this worker process.
+    """
+
+    def __init__(self, tracer: BaseTracer, max_rollouts: Optional[int] = None, poll_interval: float = 5.0) -> None:
+        """Initialize the agent runner.
+
+        Args:
+            tracer: The tracer instance for recording execution traces and spans.
+            max_rollouts: Maximum number of tasks to process in iter() mode. If None,
+                the runner will continue indefinitely until interrupted.
+            poll_interval: Time in seconds to wait between polling attempts when
+                no tasks are available in the store.
+        """
+        super().__init__()
+        self._tracer = tracer
+        self._max_rollouts = max_rollouts
+        self._poll_interval = poll_interval
+
+        # Set later
+        self._agent: Optional[LitAgent[T_task]] = None
+        self._hooks: Sequence[Hook] = []
+        self._store: Optional[LightningStore] = None
+        self.worker_id: Optional[int] = None
+
+    def init(self, agent: LitAgent[T_task], *, hooks: Optional[Sequence[Hook]] = None, **kwargs: Any) -> None:
+        """Initialize the runner with the agent.
+
+        This sets up the agent-runner relationship, registers hooks, and
+        initializes the tracer.
+
+        Args:
+            agent: The LitAgent instance to be managed by this runner.
+            hooks: Optional sequence of Hook objects to be called at various
+                lifecycle stages (on_trace_start, on_trace_end, on_rollout_start,
+                on_rollout_end).
+            **kwargs: Additional initialization arguments (currently unused).
+        """
+        self._agent = agent
+        self._agent.set_runner(self)
+        self._hooks = [*hooks] if hooks is not None else []
+
+        self._tracer.init()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Initialize the runner for each worker with worker_id and store.
+
+        This method is called once per worker in a distributed setup to provide
+        the worker with its ID and store connection.
+
+        Args:
+            worker_id: Unique identifier for this worker process.
+            store: The LightningStore instance for task coordination and data persistence.
+            **kwargs: Additional worker-specific initialization arguments (currently unused).
+        """
+        self._store = store
+        self.worker_id = worker_id
+
+        self._tracer.init_worker(worker_id)
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner and clean up all resources.
+
+        This method resets all internal state including the agent, store,
+        hooks, and worker ID, and calls the tracer's teardown method.
+
+        Args:
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self._agent = None
+        self._store = None
+        self.worker_id = None
+        self._hooks = []
+
+        self._tracer.teardown()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the runner for a specific worker.
+
+        This method cleans up worker-specific resources and resets the worker ID.
+
+        Args:
+            worker_id: The unique identifier of the worker being torn down.
+            *args: Additional teardown arguments (currently unused).
+            **kwargs: Additional teardown keyword arguments (currently unused).
+        """
+        self.worker_id = None
+
+        self._tracer.teardown_worker(worker_id)
+
+    @property
+    def tracer(self) -> BaseTracer:
+        """Get the tracer instance.
+
+        Returns:
+            The BaseTracer instance used by this runner.
+        """
+        return self._tracer
+
+    def get_agent(self) -> LitAgent[T_task]:
+        """Get the agent instance.
+
+        Returns:
+            The LitAgent instance managed by this runner.
+
+        Raises:
+            ValueError: If the agent has not been initialized via init().
+        """
+        if self._agent is None:
+            raise ValueError("Agent not initialized. Call init() first.")
+        return self._agent
+
+    def get_store(self) -> LightningStore:
+        """Get the store instance.
+
+        Returns:
+            The LightningStore instance for this worker.
+
+        Raises:
+            ValueError: If the store has not been initialized via init_worker().
+        """
+        if self._store is None:
+            raise ValueError("Store not initialized. Call init_worker() first.")
+        return self._store
+
+    def get_worker_id(self) -> str:
+        """Get the formatted worker ID string.
+
+        Returns:
+            A formatted string like "Worker-0" if initialized, or "Worker-Unknown"
+            if the worker ID has not been set.
+        """
+        return f"Worker-{self.worker_id}" if self.worker_id is not None else "Worker-Unknown"
+
+    def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
+        """Generate a standardized log prefix for the current worker.
+
+        This creates a consistent prefix format for log messages to identify
+        which worker and rollout the message is associated with.
+
+        Args:
+            rollout_id: Optional rollout ID to include in the prefix.
+
+        Returns:
+            A formatted log prefix string like "[Worker 0 | Rollout xyz]",
+            "[Worker 0]", "[Rollout xyz]", or "[Default Worker]".
+        """
+        if self.worker_id is not None:
+            if rollout_id:
+                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+            else:
+                return f"[Worker {self.worker_id}]"
+        if rollout_id:
+            return f"[Rollout {rollout_id}]"
+        return "[Default Worker]"
+
+    async def _trigger_hooks(
+        self,
+        hook_type: Literal["on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end"],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """Trigger all registered hooks of a specific type.
+
+        This method calls the specified hook method on all registered hooks,
+        catching and logging any exceptions that occur during hook execution
+        to prevent them from disrupting the main execution flow.
+
+        Args:
+            hook_type: The type of hook to trigger. Valid values are:
+                "on_trace_start", "on_trace_end", "on_rollout_start", "on_rollout_end".
+            *args: Positional arguments to pass to the hook methods.
+            **kwargs: Keyword arguments to pass to the hook methods.
+        """
+        for hook in self._hooks:
+            try:
+                await getattr(hook, hook_type)(*args, **kwargs)
+            except Exception:
+                logger.exception(f"{self._log_prefix()} Exception during {hook_type} hook {hook}.")
+
+    async def _post_process_rollout_result(
+        self, rollout: AttemptedRollout, raw_result: RolloutRawResult
+    ) -> List[ReadableSpan] | List[Span]:
+        """Standardizes the agent's return value and report what's needed to report to the store.
+
+        Args:
+            rollout: The rollout object for the current task.
+            raw_result: The output from the agent's rollout method.
+
+        Returns:
+            The spans that are assumed to be added to the store.
+            This only serves as an estimation for logging purposes. For precise tracking, use the store directly.
+        """
+        store = self.get_store()
+
+        trace_spans: list[ReadableSpan] | list[Span] = []
+
+        # Case 0: result is None
+        if raw_result is None:
+            trace_spans = self._tracer.get_last_trace()
+
+        # Case 1: result is a float (final reward)
+        if isinstance(raw_result, float):
+            # Preserve the existing spans before another span is emitted
+            trace_spans = list(self._tracer.get_last_trace())
+            # This will emit another span to the tracer
+            reward_span = emit_reward(raw_result)
+            await store.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, reward_span)
+            trace_spans.append(reward_span)
+
+        if isinstance(raw_result, list):
+            # For rollout methods that return a list, we assume that the returned spans
+            # are the complete span set from the whole rollout
+            trace_spans = raw_result
+
+            # Case 2: result is a list of ReadableSpan (OpenTelemetry spans)
+            if len(raw_result) > 0 and all(isinstance(t, ReadableSpan) for t in raw_result):
+
+                if not isinstance(
+                    self._tracer, AgentOpsTracer
+                ):  # TODO: this should be replaced with general OpenTelemetry tracer in next version
+                    for span in raw_result:
+                        await store.add_otel_span(
+                            rollout.rollout_id, rollout.attempt.attempt_id, cast(ReadableSpan, span)
+                        )
+                else:
+                    logger.warning(
+                        f"{self._log_prefix(rollout.rollout_id)} Tracer is already an OpenTelemetry tracer. "
+                        "The traces should have already been added to the store. "
+                        "No need to return anything from rollout."
+                    )
+
+            # Case 3: result is a list of Span (agentlightning spans)
+            elif len(raw_result) > 0 and all(isinstance(t, Span) for t in raw_result):
+                # Add the spans directly to the store
+                for span in raw_result:
+                    await store.add_span(cast(Span, span))
+                trace_spans = raw_result
+
+            # Left over cases for list
+            elif len(raw_result) == 0:
+                logger.warning(
+                    f"{self._log_prefix(rollout.rollout_id)} The rollout returns an empty list. "
+                    "Please check your rollout implementation."
+                )
+                trace_spans = raw_result
+
+            else:
+                types = [type(t).__name__ for t in raw_result][:10]
+                raise ValueError(
+                    f"Invalid raw result type. It's expected to be a list of ReadableSpan or Span, "
+                    f"but got: {', '.join(types)}..."
+                )
+
+        return trace_spans
+
+    async def _sleep_until_next_poll(self, event: Optional[ExecutionEvent] = None) -> None:
+        """Sleep until the next poll interval, with optional event-based interruption.
+
+        If an event is provided, the method will check it periodically (every 0.1s)
+        and return early if the event is set.
+
+        Args:
+            event: Optional ExecutionEvent object that can be used to interrupt the sleep.
+                If set during the sleep period, the method returns immediately.
+        """
+        if event is None:
+            await asyncio.sleep(self._poll_interval)
+            return
+        current_time = time.time()
+        next_time = current_time + self._poll_interval
+        while time.time() < next_time:
+            await asyncio.sleep(0.1)
+            if event.is_set():
+                return
+
+    async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: bool = False) -> str:
+        """Execute a single rollout implementation.
+
+        This is the core method that handles the execution of a single rollout,
+        including resource fetching, hook triggering, agent invocation, tracing,
+        and result processing.
+
+        Args:
+            next_rollout: The rollout to execute, containing input data, mode,
+                and resources information.
+            raise_on_exception: If True, exceptions during rollout execution will
+                be re-raised. If False, exceptions are logged but not propagated.
+        """
+        store = self.get_store()
+        agent = self.get_agent()
+
+        rollout_id = next_rollout.rollout_id
+
+        resources_id = next_rollout.resources_id
+        resources_update = None
+        if resources_id:
+            resources_update = await store.get_resources_by_id(resources_id)
+        else:
+            logger.debug(f"{self._log_prefix(rollout_id)} No 'resources_id'. Fetching latest resources.")
+            resources_update = await store.get_latest_resources()
+        if not resources_update:
+            if raise_on_exception:
+                raise RuntimeError(f"{self._log_prefix(rollout_id)} Failed to fetch resources")
+            else:
+                logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
+                return rollout_id
+
+        trace_spans: List[ReadableSpan] | List[Span] = []
+        has_exception: bool = False
+
+        try:
+            await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout)
+
+            start_time = time.time()
+            with self._tracer.trace_context(
+                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+            ):
+                await self._trigger_hooks(
+                    hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+                # NOTE: This is the most costly step in the whole function
+                # If the rollout method becomes unresponsive or timeouts, there is nothing we can do within the runner.
+                # We might need some mechanisms in execution strategy to restart the runner. But that's a future work.
+                if agent.is_async():
+                    rollout_method = (
+                        agent.training_rollout_async if next_rollout.mode == "train" else agent.validation_rollout_async
+                    )
+                    result = await rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+                else:
+                    rollout_method = (
+                        agent.training_rollout if next_rollout.mode == "train" else agent.validation_rollout
+                    )
+                    result = rollout_method(
+                        next_rollout.input, resources=resources_update.resources, rollout=next_rollout
+                    )
+
+                await self._trigger_hooks(
+                    hook_type="on_trace_end", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout
+                )
+
+            # Possible exceptions in post_process will be caught in the overall exception handler
+            trace_spans = await self._post_process_rollout_result(next_rollout, result)
+            last_reward = find_final_reward(trace_spans)
+
+            end_time = time.time()
+            logger.info(
+                f"{self._log_prefix(rollout_id)} Completed in "
+                f"{end_time - start_time:.2f}s. Collected {len(trace_spans)} span(s). "
+                f"Final reward: {last_reward}"
+            )
+
+        except Exception:
+            logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
+            has_exception = True
+
+            if raise_on_exception:
+                raise
+        finally:
+            try:
+                await self._trigger_hooks(
+                    hook_type="on_rollout_end", agent=agent, runner=self, rollout=next_rollout, spans=trace_spans
+                )
+            except Exception:
+                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+
+            try:
+                if has_exception:
+                    # possibly timed out and cancelled?
+                    await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="failed")
+                else:
+                    await store.update_attempt(rollout_id, next_rollout.attempt.attempt_id, status="succeeded")
+            except Exception:
+                logger.exception(
+                    f"{self._log_prefix(rollout_id)} Exception during update_attempt. Giving up the update."
+                )
+
+        return rollout_id
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method polls the store for new rollouts and executes them until:
+        - The event is set (if provided)
+        - The max_rollouts limit is reached (if configured)
+        - No more tasks are available
+
+        All exceptions during rollout execution are caught and logged but not
+        propagated, allowing the runner to continue processing subsequent tasks.
+
+        Args:
+            event: Optional ExecutionEvent object to signal the runner to stop. The runner
+                will check this event periodically and stop gracefully when set.
+        """
+        num_tasks_processed = 0
+        logger.info(f"{self._log_prefix()} Started async rollouts (max: {self._max_rollouts or 'unlimited'}).")
+        store = self.get_store()
+
+        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()
+                if next_rollout is None:
+                    logger.debug(f"{self._log_prefix()} No rollout to poll. Waiting for {self._poll_interval} seconds.")
+                    await self._sleep_until_next_poll(event)
+                else:
+                    break
+
+            if next_rollout is None:
+                return
+
+            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)
+
+            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'}")
+
+        logger.info(f"{self._log_prefix()} Finished async rollouts. Processed {num_tasks_processed} tasks.")
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
+        """Execute a single task directly, bypassing the task queue.
+
+        This method creates a new rollout for the given input and executes it
+        immediately. Unlike iter(), exceptions are propagated to the caller.
+
+        Args:
+            input: The task input to be processed by the agent.
+            resources: Optional named resources to be used for this specific task.
+                If provided, a new resources entry will be created in the store.
+                If not provided, the latest resources from the store will be used.
+            mode: Optional rollout mode ("train" or "validation"). If not provided,
+                the agent's default mode will be used.
+            event: Optional ExecutionEvent object to signal interruption (currently unused
+                but included for interface consistency).
+
+        Returns:
+            The completed rollout.
+
+        Raises:
+            Exception: Any exception that occurs during rollout execution will be
+                re-raised to the caller.
+        """
+        store = self.get_store()
+
+        if resources is not None:
+            resources_update = await store.add_resources(resources)
+            resources_id = resources_update.resources_id
+        else:
+            resources_id = None
+
+        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
+        rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
+
+        completed_rollout = await store.get_rollout_by_id(rollout_id)
+        if completed_rollout is None:
+            raise RuntimeError(f"{self._log_prefix()} Failed to fetch completed rollout by id after step: {rollout_id}")
+        return completed_rollout

agentlightning/runner/base.py

@@ -0,0 +1,203 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Base runner interface for executing agent tasks.
+
+This module defines the abstract base class for all runner implementations
+in the agent-lightning framework. Runners are responsible for managing the
+execution lifecycle of agents and coordinating with the store.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generic, Iterator, Optional, Sequence, TypeVar
+
+from agentlightning.execution.events import ExecutionEvent
+from agentlightning.litagent import LitAgent
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Hook, NamedResources, ParallelWorkerBase, Rollout, RolloutMode
+
+if TYPE_CHECKING:
+    from agentlightning.execution.events import ExecutionEvent
+
+
+T_task = TypeVar("T_task")
+
+logger = logging.getLogger(__name__)
+
+
+class BaseRunner(ParallelWorkerBase, Generic[T_task]):
+    """Base class for all runners.
+
+    This abstract base class defines the interface that all runner implementations
+    must follow. Runners are responsible for executing agent tasks, managing the
+    execution lifecycle, and coordinating with the store.
+    """
+
+    def init(self, agent: LitAgent[T_task], **kwargs: Any) -> None:
+        """Initialize the runner with the agent.
+
+        This method is called once during setup to configure the runner with
+        the agent it will execute.
+
+        Args:
+            agent: The LitAgent instance to be managed by this runner.
+            **kwargs: Additional initialization arguments specific to the runner implementation.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> None:
+        """Initialize the runner for each worker with worker_id and store.
+
+        This method is called once per worker process in a distributed setup.
+        It provides the worker with its unique ID and the store instance for
+        task coordination.
+
+        Args:
+            worker_id: Unique identifier for this worker process.
+            store: The LightningStore instance for task coordination and data persistence.
+            **kwargs: Additional worker-specific initialization arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Undefined method - use iter() or step() instead.
+
+        This method is intentionally not implemented as the execution behavior
+        should be defined through iter() for continuous execution or step()
+        for single-task execution.
+
+        Args:
+            *args: Unused positional arguments.
+            **kwargs: Unused keyword arguments.
+
+        Raises:
+            RuntimeError: Always raised to indicate this method should not be used.
+        """
+        raise RuntimeError("The behavior of run() of Runner is undefined. Use iter() or step() instead.")
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Clean up runner resources and reset state.
+
+        This method is called once during shutdown to clean up any resources
+        allocated during initialization and reset the runner state.
+
+        Args:
+            *args: Additional teardown arguments.
+            **kwargs: Additional teardown keyword arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Clean up worker-specific resources.
+
+        This method is called once per worker during shutdown to clean up
+        any resources specific to that worker.
+
+        Args:
+            worker_id: The unique identifier of the worker being torn down.
+            *args: Additional teardown arguments.
+            **kwargs: Additional teardown keyword arguments.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    @contextmanager
+    def run_context(
+        self,
+        *,
+        agent: LitAgent[T_task],
+        store: LightningStore,
+        hooks: Optional[Sequence[Hook]] = None,
+        worker_id: Optional[int] = None,
+    ) -> Iterator[BaseRunner[T_task]]:
+        """Context manager for quickly init and teardown the runner,
+        so that you can debug the runner without a trainer environment.
+
+        Args:
+            agent: The LitAgent instance to be managed by this runner.
+                   It should be the same agent that is to be run within the context.
+            store: The LightningStore instance for task coordination and data persistence.
+                   If you don't have one, you can easily create one with `InMemoryLightningStore()`.
+            hooks: Optional sequence of Hook instances to be used by the runner.
+                   Only some runners support hooks.
+            worker_id: Optional worker ID to be used by the runner.
+        """
+        _initialized: bool = False
+        _worker_initialized: bool = False
+        try:
+            self.init(agent=agent, hooks=hooks)
+            _initialized = True
+            self.init_worker(worker_id=0, store=store)
+            _worker_initialized = True
+            yield self
+        finally:
+            try:
+                if _worker_initialized:
+                    self.teardown_worker(worker_id=worker_id if worker_id is not None else 0)
+            except Exception:
+                logger.error("Error during runner worker teardown", exc_info=True)
+
+            try:
+                if _initialized:
+                    self.teardown()
+            except Exception:
+                logger.error("Error during runner teardown", exc_info=True)
+
+    async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
+        """Run the runner, continuously iterating over tasks in the store.
+
+        This method runs in a loop, polling the store for new tasks and executing
+        them until interrupted by the event or when no more tasks are available.
+
+        Args:
+            event: Optional ExecutionEvent object that can be used to signal the runner
+                to stop gracefully. When set, the runner should finish its current
+                task and exit the iteration loop.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()
+
+    async def step(
+        self,
+        input: T_task,
+        *,
+        resources: Optional[NamedResources] = None,
+        mode: Optional[RolloutMode] = None,
+        event: Optional[ExecutionEvent] = None,
+    ) -> Rollout:
+        """Execute a single task with the given input.
+
+        This method provides fine-grained control for executing individual tasks
+        directly, bypassing the store's task queue.
+
+        Args:
+            input: The task input to be processed by the agent.
+            resources: Optional named resources to be used for this specific task.
+                If not provided, the latest resources from the store will be used.
+            mode: Optional rollout mode (e.g., "train", "test"). If not provided,
+                the default mode will be used.
+            event: Optional ExecutionEvent object to signal interruption. When set, the
+                runner may abort the current execution.
+
+        Returns:
+            The completed rollout.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError()

agentlightning/runner/legacy.py

@@ -1,25 +1,29 @@
-import asyncio
+# Copyright (c) Microsoft. All rights reserved.
+
 import json
 import logging
-import os
 import time
-from contextlib import nullcontext
-from typing import List, Optional, Union, Dict, Any
-
-import agentops
+from typing import Any, Dict, List, Optional, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
-from .client import AgentLightningClient
-from .litagent import LitAgent
-from .types import Rollout, Task, Triplet, RolloutRawResult
-from .types import ParallelWorkerBase
-from .tracer.base import BaseTracer
-from .tracer import TripletExporter
+
+from agentlightning.adapter import TracerTraceToTriplet
+from agentlightning.client import AgentLightningClient
+from agentlightning.litagent import LitAgent
+from agentlightning.litagent.litagent import is_v0_1_rollout_api
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import RolloutLegacy, RolloutRawResultLegacy, Triplet
+
+from .base import BaseRunner
 
 logger = logging.getLogger(__name__)
 
+__all__ = [
+    "LegacyAgentRunner",
+]
 
-class AgentRunner(ParallelWorkerBase):
+
+class LegacyAgentRunner(BaseRunner[Any]):
     """Manages the agent's execution loop and integrates with AgentOps.
 
     This class orchestrates the interaction between the agent (`LitAgent`) and
@@ -37,10 +41,10 @@ class AgentRunner(ParallelWorkerBase):
 
     def __init__(
         self,
-        agent: LitAgent,
+        agent: LitAgent[Any],
         client: AgentLightningClient,
         tracer: BaseTracer,
-        triplet_exporter: TripletExporter,
+        triplet_exporter: TracerTraceToTriplet,
         worker_id: Optional[int] = None,
         max_tasks: Optional[int] = None,
     ):
@@ -54,30 +58,43 @@ class AgentRunner(ParallelWorkerBase):
         self.worker_id = worker_id
         self.max_tasks = max_tasks
 
+    # These methods are overridden by BaseRunner, getting them back to old behavior.
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
     def _log_prefix(self, rollout_id: Optional[str] = None) -> str:
         """Generates a standardized log prefix for the current worker."""
         if self.worker_id is not None:
             if rollout_id:
-                return f"[Worker {self.worker_id} | Rollout {rollout_id}]"
+                return f"[Worker {self.worker_id} | RolloutLegacy {rollout_id}]"
             else:
                 return f"[Worker {self.worker_id}]"
         if rollout_id:
-            return f"[Rollout {rollout_id}]"
+            return f"[RolloutLegacy {rollout_id}]"
         return "[Default Worker]"
 
     def _to_rollout_object(
         self,
-        result: RolloutRawResult,
+        result: RolloutRawResultLegacy,
         rollout_id: str,
-    ) -> Rollout:
-        """Standardizes the agent's return value into a Rollout object.
+    ) -> RolloutLegacy:
+        """Standardizes the agent's return value into a RolloutLegacy object.
 
         Args:
             result: The output from the agent's rollout method.
             rollout_id: The unique identifier for the current task.
 
         Returns:
-            A standardized `Rollout` object for reporting to the server.
+            A standardized `RolloutLegacy` object for reporting to the server.
         """
         trace: Any = None
         final_reward: Optional[float] = None
@@ -98,8 +115,8 @@ class AgentRunner(ParallelWorkerBase):
         # Case 4: result is a list of dict (trace JSON)
         if isinstance(result, list) and all(isinstance(t, dict) for t in result):
             trace = result
-        # Case 5: result is a Rollout object
-        if isinstance(result, Rollout):
+        # Case 5: result is a RolloutLegacy object
+        if isinstance(result, RolloutLegacy):
             final_reward = result.final_reward
             triplets = result.triplets
             trace = result.trace
@@ -111,15 +128,15 @@ class AgentRunner(ParallelWorkerBase):
                 trace = [json.loads(readable_span.to_json()) for readable_span in spans]
                 trace_spans = spans
 
-        # Always extract triplets from the trace using TripletExporter
+        # Always extract triplets from the trace using TracerTraceToTriplet
         if trace_spans:
-            triplets = self.triplet_exporter.export(trace_spans)
+            triplets = self.triplet_exporter(trace_spans)  # type: ignore
 
         # If the agent has triplets, use the last one for final reward if not set
         if triplets and triplets[-1].reward is not None and final_reward is None:
             final_reward = triplets[-1].reward
 
-        # Create the Rollout object with standardized fields
+        # Create the RolloutLegacy object with standardized fields
         result_dict: Dict[str, Any] = {
             "rollout_id": rollout_id,
         }
@@ -130,11 +147,11 @@ class AgentRunner(ParallelWorkerBase):
         if trace is not None:
             result_dict["trace"] = trace
 
-        if isinstance(result, Rollout):
+        if isinstance(result, RolloutLegacy):
             return result.model_copy(update=result_dict)
-        return Rollout(**result_dict)
+        return RolloutLegacy(**result_dict)
 
-    def run(self) -> bool:
+    def run(self) -> bool:  # type: ignore
         """Poll the task and rollout once synchronously."""
         self.agent.set_runner(self)  # Ensure the agent has a reference to this runner
 
@@ -155,7 +172,7 @@ class AgentRunner(ParallelWorkerBase):
             logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
             return False
 
-        rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
 
         try:
             try:
@@ -167,8 +184,16 @@ class AgentRunner(ParallelWorkerBase):
                 start_time = time.time()
                 rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
                 # Pass the task input, not the whole task object
-                result = rollout_method(task.input, task.rollout_id, resources_update.resources)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                 end_time = time.time()
                 logger.info(
                     f"{self._log_prefix(rollout_id)} Completed in "
@@ -181,14 +206,14 @@ class AgentRunner(ParallelWorkerBase):
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
             try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
             self.client.post_rollout(rollout_obj)
 
         return True
 
-    def iter(self) -> int:
+    def iter(self) -> int:  # type: ignore
         """Executes the synchronous polling and rollout loop."""
         num_tasks_processed = 0
         logger.info(f"{self._log_prefix()} Started sync rollouts (max: {self.max_tasks or 'unlimited'}).")
@@ -224,7 +249,7 @@ class AgentRunner(ParallelWorkerBase):
             logger.error(f"{self._log_prefix(rollout_id)} Failed to fetch resources. Skipping.")
             return False
 
-        rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
+        rollout_obj = RolloutLegacy(rollout_id=task.rollout_id, task=task)  # Default empty rollout
 
         try:
             try:
@@ -238,18 +263,28 @@ class AgentRunner(ParallelWorkerBase):
                     self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async
                 )
                 # Pass the task input, not the whole task object
-                result = await rollout_method(task.input, task.rollout_id, resources_update.resources)
-                rollout_obj = self._to_rollout_object(result, task.rollout_id)
+                if is_v0_1_rollout_api(rollout_method):
+                    result = cast(
+                        RolloutRawResultLegacy,
+                        await rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
+                    )  # type: ignore
+                else:
+                    result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)  # type: ignore
+                rollout_obj = self._to_rollout_object(result, task.rollout_id)  # type: ignore
                 end_time = time.time()
                 logger.info(
                     f"{self._log_prefix(rollout_id)} Completed in "
-                    f"{end_time - start_time:.2f}s. Reward: {rollout_obj.final_reward}"
+                    f"{end_time - start_time:.2f}s. Triplet length: "
+                    f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
+                    f"Reward: {rollout_obj.final_reward}"
                 )
         except Exception:
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
             try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
+                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)  # type: ignore
             except Exception:
                 logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
             await self.client.post_rollout_async(rollout_obj)

agentlightning/server.py

@@ -1,22 +1,28 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Legacy server for the Agent Lightning framework. Deprecated in favor of agentlightning.store."""
+
+from __future__ import annotations
+
 import asyncio
 import logging
+import threading
 import time
 import uuid
-import threading
+import warnings
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional, Literal
+from typing import Any, Dict, List, Literal, Optional
 
 import uvicorn
 from fastapi import FastAPI, HTTPException, Path
-from pydantic import Field
 
 from .types import (
-    Rollout,
+    GenericResponse,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutLegacy,
     Task,
     TaskIfAny,
-    NamedResources,
-    GenericResponse,
-    ResourcesUpdate,
 )
 
 logger = logging.getLogger(__name__)
@@ -31,7 +37,7 @@ class ServerDataStore:
     def __init__(self):
         self._task_queue: asyncio.Queue[Task] = asyncio.Queue()
         self._processing_tasks: Dict[str, Task] = {}  # Currently processing tasks
-        self._completed_rollouts: Dict[str, Rollout] = {}
+        self._completed_rollouts: Dict[str, RolloutLegacy] = {}
 
         # Store for versioned resources
         self._resource_versions: Dict[str, NamedResources] = {}
@@ -116,7 +122,7 @@ class ServerDataStore:
             return await self.get_resources_by_id(self._latest_resources_id)
         return None
 
-    async def store_rollout(self, rollout: Rollout):
+    async def store_rollout(self, rollout: RolloutLegacy):
         """
         Safely stores a completed rollout from a client.
         """
@@ -125,14 +131,14 @@ class ServerDataStore:
             self._completed_rollouts[rollout.rollout_id] = rollout
             logger.info(f"Rollout received and stored: {rollout.rollout_id}")
 
-    async def retrieve_rollout(self, rollout_id: str) -> Optional[Rollout]:
+    async def retrieve_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
         """
         Safely retrieves a single rollout by its ID, removing it from the store.
         """
         async with self._results_lock:
             return self._completed_rollouts.pop(rollout_id, None)
 
-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
         """
         Retrieves all completed rollouts and clears the store.
         """
@@ -171,6 +177,9 @@ class AgentLightningServer:
             port: The port to bind the server to.
             task_timeout_seconds: Time in seconds after which a claimed task is considered stale and requeued.
         """
+        warnings.warn(
+            "AgentLightningServer is deprecated. Please use LightningStoreServer instead.", DeprecationWarning
+        )
         self.host = host
         self.port = port
         self.endpoint = f"http://{host}:{port}"
@@ -216,7 +225,7 @@ class AgentLightningServer:
             return
         processing_tasks = self._store.get_processing_tasks()
 
-        for rollout_id, task in processing_tasks.items():
+        for _, task in processing_tasks.items():
             if task.last_claim_time and current_time - task.last_claim_time > self._task_timeout_seconds:
                 await self._store.requeue_task(task)
                 logger.warning(
@@ -227,7 +236,7 @@ class AgentLightningServer:
         """Setup FastAPI routes."""
 
         @self._app.get("/task", response_model=TaskIfAny)
-        async def next_task() -> TaskIfAny:
+        async def next_task() -> TaskIfAny:  # type: ignore
             """Endpoint for clients to poll for the next available task."""
             await self._check_and_requeue_stale_tasks()
 
@@ -243,7 +252,7 @@ class AgentLightningServer:
                 return TaskIfAny(is_available=False)
 
         @self._app.get("/resources/latest", response_model=ResourcesUpdate)
-        async def fetch_latest_resources() -> ResourcesUpdate:
+        async def fetch_latest_resources() -> ResourcesUpdate:  # type: ignore
             """Endpoint for clients to poll for the latest available resources."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
@@ -254,7 +263,7 @@ class AgentLightningServer:
             return resources_update
 
         @self._app.get("/resources/{resource_id}", response_model=ResourcesUpdate)
-        async def fetch_resources_by_id(
+        async def fetch_resources_by_id(  # type: ignore
             resource_id: str = Path(..., description="The unique identifier for the resource version.")
         ) -> ResourcesUpdate:
             """Endpoint for clients to fetch a specific version of resources."""
@@ -267,7 +276,7 @@ class AgentLightningServer:
             return resources_update
 
         @self._app.post("/rollout", response_model=GenericResponse)
-        async def post_rollout(payload: Rollout) -> GenericResponse:
+        async def post_rollout(payload: RolloutLegacy) -> GenericResponse:  # type: ignore
             """Endpoint for clients to report a completed rollout."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
@@ -323,7 +332,7 @@ class AgentLightningServer:
         await self._store.update_resources(update)
         return resources_id
 
-    async def get_completed_rollout(self, rollout_id: str) -> Optional[Rollout]:
+    async def get_completed_rollout(self, rollout_id: str) -> Optional[RolloutLegacy]:
         """
         Retrieves a specific completed rollout by its ID.
         """
@@ -331,7 +340,7 @@ class AgentLightningServer:
             raise RuntimeError("Store not initialized. The server may not be running.")
         return await self._store.retrieve_rollout(rollout_id)
 
-    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[Rollout]:
+    async def poll_completed_rollout(self, rollout_id: str, timeout: Optional[float] = None) -> Optional[RolloutLegacy]:
         """
         Polls for a completed rollout by its ID, waiting up to `timeout` seconds.
         """
@@ -344,7 +353,7 @@ class AgentLightningServer:
                 return None
             await asyncio.sleep(1)
 
-    async def retrieve_completed_rollouts(self) -> List[Rollout]:
+    async def retrieve_completed_rollouts(self) -> List[RolloutLegacy]:
         """
         Retrieves all available completed trajectories and clears the internal store.
         """

agentlightning/store/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .base import LightningStore
+from .client_server import LightningStoreClient, LightningStoreServer
+from .memory import InMemoryLightningStore
+from .threading import LightningStoreThreaded
+
+__all__ = [
+    "LightningStore",
+    "LightningStoreClient",
+    "LightningStoreServer",
+    "InMemoryLightningStore",
+    "LightningStoreThreaded",
+]

agentlightning/store/base.py

@@ -0,0 +1,265 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+
+def is_queuing(rollout: Rollout) -> bool:
+    return rollout.status == "queuing" or rollout.status == "requeuing"
+
+
+def is_running(rollout: Rollout) -> bool:
+    return rollout.status == "preparing" or rollout.status == "running"
+
+
+def is_finished(rollout: Rollout) -> bool:
+    return rollout.status == "failed" or rollout.status == "succeeded" or rollout.status == "cancelled"
+
+
+class _UnsetType:
+    """A sentinel type to indicate an unset value."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "UNSET"
+
+    def __reduce__(self):
+        return (_get_unset, ())
+
+
+def _get_unset() -> _UnsetType:
+    return UNSET
+
+
+UNSET = _UnsetType()
+Unset = _UnsetType  # Alias for convenience
+
+
+class LightningStore:
+    """
+    A centralized, thread-safe, async, data store for the lightning's state.
+    This holds the task queue, versioned resources, and completed rollouts.
+
+    The store has a built-in clock and it should be responsible for tracking the times.
+    All the time-based operations like retry, timeout, etc. should be handled by the store.
+    """
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """
+        Add one incomplete rollout to the store, and get an attempt created for it.
+        This will immediately sets the rollout to a preparing state, and should be
+        used by whoever is going to execute the rollout.
+
+        Return a special rollout with attempt object. Do not update it directly.
+
+        But if the rollout fails or timeouts, it's still possible that the watchdog
+        sends it back to the queue for retry.
+
+        To enqueue a rollout to the task queue, use `enqueue_rollout` instead.
+        """
+        raise NotImplementedError()
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        """
+        Adds a new task to the queue with specific metadata and
+        returns the rollout object with its unique ID.
+        """
+        raise NotImplementedError()
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """
+        Retrieves the next task from the queue without blocking.
+        Returns None if the queue is empty.
+
+        Will set the rollout status to preparing.
+        """
+        raise NotImplementedError()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """
+        Create a new attempt for a given rollout ID and return the attempt details.
+        """
+        raise NotImplementedError()
+
+    async def add_span(self, span: Span) -> Span:
+        """
+        Add a span to the store.
+
+        This method is responsible for updating the rollout/attempt status to "running" if needed.
+        """
+        raise NotImplementedError()
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        """
+        Add an opentelemetry span to the store.
+
+        If sequence_id is not provided, it will be fetched from `get_next_span_sequence_id` and assigned automatically.
+        """
+        raise NotImplementedError()
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        """
+        Query and retrieve rollouts filtered by their status.
+        If no status is provided, returns all rollouts.
+        """
+        raise NotImplementedError()
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """
+        Query and retrieve all attempts associated with a specific rollout ID.
+        Returns an empty list if no attempts are found.
+        """
+        raise NotImplementedError()
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        """
+        Safely retrieves a specific rollout by its ID.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """
+        Safely retrieves the latest attempt for a given rollout ID.
+        """
+        raise NotImplementedError()
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves a specific version of named resources by its ID.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves the latest version of named resources.
+        """
+        raise NotImplementedError()
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """
+        Get the next span sequence ID for a given rollout and attempt.
+        This should be used to assign a unique sequence ID to each span within an attempt.
+
+        Recommend getting the ID before the operation even begins to avoid racing conditions.
+        """
+        raise NotImplementedError()
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """
+        Wait for specified rollouts to complete with a timeout.
+        Returns the completed rollouts, potentially incomplete if timeout is reached.
+
+        TODO: Add support for waiting for 20 new rollouts, or wait until 80% of the pending ids are completed.
+        """
+        raise NotImplementedError()
+
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """
+        Query and retrieve all spans associated with a specific rollout ID.
+        Returns an empty list if no spans are found.
+        """
+        raise NotImplementedError()
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+        Not implemented by many stores yet.
+        """
+        raise NotImplementedError()
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version or updates an existing version of named resources and sets it as the latest.
+        """
+        raise NotImplementedError()
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        """
+        Update the rollout status and related metadata.
+
+        Not-listed fields here either cannot be updated, or should be auto-updated (e.g., end_time).
+
+        When status is updated to a finished / problematic state, other states like task
+        queues will be updated accordingly.
+
+        Args:
+            rollout_id: Unique identifier for the rollout to update
+            input: New input data for the rollout. If set, will be updated. Can be updated to None
+            mode: New mode for the rollout. If set, will be updated. Can be updated to None
+            resources_id: New resources ID for the rollout. If set, will be updated. Can be updated to None
+            status: New status for the rollout. If set, will be updated
+            config: New config for the rollout. If set, will be updated
+            metadata: Dictionary of additional metadata to update. If set, will replace the existing metadata
+        """
+        raise NotImplementedError()
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        """
+        Update a specific or latest attempt for a given rollout.
+
+        Update the latest attempt will NOT affect the corresponding rollout status.
+
+
+        Args:
+            rollout_id: Unique identifier for the rollout
+            attempt_id: Unique identifier for the attempt
+            status: Status to set for the attempt, update if provided
+            worker_id: Worker identifier, update if provided
+            last_heartbeat_time: Timestamp of the last heartbeat from the worker
+            metadata: Dictionary of additional metadata to update, will replace the existing metadata
+        """
+        raise NotImplementedError()

agentlightning/store/client_server.py

@@ -0,0 +1,963 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import threading
+import time
+import traceback
+from contextlib import suppress
+from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Sequence, Union
+
+import aiohttp
+import uvicorn
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import JSONResponse
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, Unset
+
+logger = logging.getLogger(__name__)
+
+
+class PydanticUnset(BaseModel):
+    _type: Literal["UNSET"] = "UNSET"
+
+
+class RolloutRequest(BaseModel):
+    input: TaskInput
+    mode: Optional[Literal["train", "val", "test"]] = None
+    resources_id: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class QueryRolloutsRequest(BaseModel):
+    status: Optional[List[RolloutStatus]] = None
+    rollout_ids: Optional[List[str]] = None
+
+
+class WaitForRolloutsRequest(BaseModel):
+    rollout_ids: List[str]
+    timeout: Optional[float] = None
+
+
+class RolloutId(BaseModel):
+    rollout_id: str
+
+
+class AddResourcesRequest(BaseModel):
+    resources: NamedResources
+
+
+class UpdateRolloutRequest(BaseModel):
+    rollout_id: str
+    input: Union[TaskInput, PydanticUnset] = Field(default_factory=PydanticUnset)
+    mode: Union[Optional[Literal["train", "val", "test"]], PydanticUnset] = Field(default_factory=PydanticUnset)
+    resources_id: Union[Optional[str], PydanticUnset] = Field(default_factory=PydanticUnset)
+    status: Union[RolloutStatus, PydanticUnset] = Field(default_factory=PydanticUnset)
+    config: Union[RolloutConfig, PydanticUnset] = Field(default_factory=PydanticUnset)
+    metadata: Union[Dict[str, Any], PydanticUnset] = Field(default_factory=PydanticUnset)
+
+
+class UpdateAttemptRequest(BaseModel):
+    rollout_id: str
+    attempt_id: Union[str, Literal["latest"]]
+    status: Union[AttemptStatus, PydanticUnset] = Field(default_factory=PydanticUnset)
+    worker_id: Union[str, PydanticUnset] = Field(default_factory=PydanticUnset)
+    last_heartbeat_time: Union[float, PydanticUnset] = Field(default_factory=PydanticUnset)
+    metadata: Union[Dict[str, Any], PydanticUnset] = Field(default_factory=PydanticUnset)
+
+
+class LightningStoreServer(LightningStore):
+    """
+    Server wrapper that exposes a LightningStore via HTTP API.
+    Delegates all operations to an underlying store implementation.
+
+    Healthcheck and watchdog relies on the underlying store.
+    """
+
+    def __init__(self, store: LightningStore, host: str, port: int):
+        super().__init__()
+        self.store = store
+        self.host = host
+        self.port = port
+        self.app: FastAPI | None = FastAPI(title="LightningStore Server")
+        self._setup_routes()
+        self._uvicorn_config: uvicorn.Config | None = uvicorn.Config(
+            self.app, host="0.0.0.0", port=self.port, log_level="error"
+        )
+        self._uvicorn_server: uvicorn.Server | None = uvicorn.Server(self._uvicorn_config)
+
+        self._serving_thread: Optional[threading.Thread] = None
+
+        # Process-awareness:
+        # LightningStoreServer holds a plain Python object (self.store) in one process
+        # (the process that runs uvicorn/FastAPI).
+        # When you multiprocessing.Process(...) and call methods on a different LightningStore instance
+        # (or on a copy inherited via fork), you’re mutating another process’s memory, not the server’s memory.
+        # So we need to track the owner process (whoever creates the server),
+        # and only mutate the store in that process.
+        self._owner_pid = os.getpid()
+        self._client: Optional[LightningStoreClient] = None
+
+    def __getstate__(self):
+        """
+        Control pickling to prevent server state from being sent to subprocesses.
+
+        When LightningStoreServer is pickled (e.g., passed to a subprocess), we only
+        serialize the underlying store and connection details. The FastAPI app and
+        uvicorn server are excluded as they should not be transferred between processes.
+
+        The subprocess should create its own server instance if needed.
+        """
+        return {
+            "store": self.store,
+            "host": self.host,
+            "port": self.port,
+            "_owner_pid": self._owner_pid,
+        }
+
+    def __setstate__(self, state: Dict[str, Any]):
+        """
+        Restore from pickle by reconstructing only the essential attributes.
+
+        Note: This creates a new server instance without FastAPI/uvicorn initialized.
+        Call __init__() pattern or create a new LightningStoreServer if you need
+        a fully functional server in the subprocess.
+        """
+        self.store = state["store"]
+        self.host = state["host"]
+        self.port = state["port"]
+        self._owner_pid = state["_owner_pid"]
+        self._client = None
+        # Do NOT reconstruct app, _uvicorn_config, _uvicorn_server
+        # to avoid transferring server state to subprocess
+
+    @property
+    def endpoint(self) -> str:
+        return f"http://{self.host}:{self.port}"
+
+    async def start(self):
+        """Starts the FastAPI server in the background.
+
+        You need to call this method in the same process as the server was created in.
+        """
+        assert self._uvicorn_server is not None
+        logger.info(f"Starting server at {self.endpoint}")
+
+        uvicorn_server = self._uvicorn_server
+
+        def run_server_forever():
+            asyncio.run(uvicorn_server.serve())
+
+        self._serving_thread = threading.Thread(target=run_server_forever, daemon=True)
+        self._serving_thread.start()
+
+        # Wait for /health to be available
+        if not await self._server_health_check():
+            raise RuntimeError("Server failed to start within the 10 seconds.")
+
+    async def _server_health_check(self) -> bool:
+        """Checks if the server is healthy."""
+        current_time = time.time()
+        while time.time() - current_time < 10:
+            async with aiohttp.ClientSession() as session:
+                with suppress(Exception):
+                    async with session.get(f"{self.endpoint}/health") as response:
+                        if response.status == 200:
+                            return True
+            await asyncio.sleep(0.1)
+        return False
+
+    async def run_forever(self):
+        """Runs the FastAPI server indefinitely.
+
+        You need to call this method in the same process as the server was created in.
+        """
+        assert self._uvicorn_server is not None
+
+        async def _wait_till_healthy():
+            health = await self._server_health_check()
+            if not health:
+                raise RuntimeError("Server did not become healthy within the 10 seconds.")
+            logger.info("Store server is online at %s", self.endpoint)
+
+        # We run _wait_till_healthy and self._uvicorn_server.serve in parallel
+        # until one of them raises an exception.
+        await asyncio.gather(_wait_till_healthy(), self._uvicorn_server.serve())
+
+    async def stop(self):
+        """Gracefully stops the running FastAPI server.
+
+        You need to call this method in the same process as the server was created in.
+        """
+        assert self._uvicorn_server is not None
+        if self._uvicorn_server.started:
+            logger.info("Stopping server...")
+            self._uvicorn_server.should_exit = True
+            if self._serving_thread is not None:
+                self._serving_thread.join(timeout=10)
+            self._serving_thread = None
+            logger.info("Server stopped.")
+
+    def _backend(self) -> LightningStore:
+        """Returns the object to delegate to in *this* process.
+
+        - In the owner process: delegate to the in-process store.
+        - In a different process: delegate to a HTTP client talking to the server.
+        """
+        if os.getpid() == self._owner_pid:
+            return self.store
+        if self._client is None:
+            self._client = LightningStoreClient(self.endpoint)
+        return self._client
+
+    def _setup_routes(self):
+        """Set up FastAPI routes for all store operations."""
+        assert self.app is not None
+
+        @self.app.exception_handler(Exception)
+        async def _app_exception_handler(request: Request, exc: Exception):  # pyright: ignore[reportUnusedFunction]
+            """
+            Convert unhandled application exceptions into 400 responses.
+
+            - Client needs a reliable signal to distinguish "app bug / bad request"
+              from transport/session failures.
+            - 400 here means "do not retry"; network issues will surface as aiohttp
+              exceptions or 5xx and will be retried by the client shield.
+            """
+            logger.exception("Unhandled application error", exc_info=exc)
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "detail": str(exc),
+                    "error_type": type(exc).__name__,
+                    "traceback": traceback.format_exc(),
+                },
+            )
+
+        @self.app.middleware("http")
+        async def _log_time(  # pyright: ignore[reportUnusedFunction]
+            request: Request, call_next: Callable[[Request], Awaitable[Response]]
+        ):
+            start = time.perf_counter()
+            response = await call_next(request)
+            duration = (time.perf_counter() - start) * 1000
+            client = request.client
+            if client is None:
+                client_address = "unknown"
+            else:
+                client_address = f"{client.host}:{client.port}"
+            logger.info(
+                f"{client_address} - "
+                f'"{request.method} {request.url.path} HTTP/{request.scope["http_version"]}" '
+                f"{response.status_code} in {duration:.2f} ms"
+            )
+            return response
+
+        @self.app.get("/health")
+        async def health():  # pyright: ignore[reportUnusedFunction]
+            return {"status": "ok"}
+
+        @self.app.post("/start_rollout", response_model=AttemptedRollout)
+        async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.start_rollout(
+                input=request.input,
+                mode=request.mode,
+                resources_id=request.resources_id,
+                metadata=request.metadata,
+            )
+
+        @self.app.post("/enqueue_rollout", response_model=Rollout)
+        async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.enqueue_rollout(
+                input=request.input,
+                mode=request.mode,
+                resources_id=request.resources_id,
+                metadata=request.metadata,
+            )
+
+        @self.app.get("/dequeue_rollout", response_model=Optional[AttemptedRollout])
+        async def dequeue_rollout():  # pyright: ignore[reportUnusedFunction]
+            return await self.store.dequeue_rollout()
+
+        @self.app.post("/start_attempt", response_model=AttemptedRollout)
+        async def start_attempt(request: RolloutId):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.start_attempt(request.rollout_id)
+
+        @self.app.post("/query_rollouts", response_model=List[Rollout])
+        async def query_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.query_rollouts(status=request.status)
+
+        @self.app.get("/query_attempts/{rollout_id}", response_model=List[Attempt])
+        async def query_attempts(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.query_attempts(rollout_id)
+
+        @self.app.get("/get_latest_attempt/{rollout_id}", response_model=Optional[Attempt])
+        async def get_latest_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_latest_attempt(rollout_id)
+
+        @self.app.get("/get_rollout_by_id/{rollout_id}", response_model=Optional[Rollout])
+        async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_rollout_by_id(rollout_id)
+
+        @self.app.post("/add_resources", response_model=ResourcesUpdate)
+        async def add_resources(resources: AddResourcesRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.add_resources(resources.resources)
+
+        @self.app.post("/update_resources", response_model=ResourcesUpdate)
+        async def update_resources(update: ResourcesUpdate):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_resources(update.resources_id, update.resources)
+
+        @self.app.get("/get_resources_by_id/{resources_id}", response_model=Optional[ResourcesUpdate])
+        async def get_resources_by_id(resources_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_resources_by_id(resources_id)
+
+        @self.app.get("/get_latest_resources", response_model=Optional[ResourcesUpdate])
+        async def get_latest_resources():  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_latest_resources()
+
+        @self.app.post("/add_span", response_model=Span)
+        async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! add_span received")
+            return await self.store.add_span(span)
+
+        @self.app.get("/get_next_span_sequence_id/{rollout_id}/{attempt_id}", response_model=int)
+        async def get_next_span_sequence_id(rollout_id: str, attempt_id: str):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! get_next_span_sequence_id received")
+            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+        @self.app.post("/wait_for_rollouts", response_model=List[Rollout])
+        async def wait_for_rollouts(request: WaitForRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.wait_for_rollouts(rollout_ids=request.rollout_ids, timeout=request.timeout)
+
+        @self.app.get("/query_spans/{rollout_id}", response_model=List[Span])
+        async def query_spans(  # pyright: ignore[reportUnusedFunction]
+            rollout_id: str, attempt_id: Optional[str] = None
+        ):
+            return await self.store.query_spans(rollout_id, attempt_id)
+
+        @self.app.post("/update_rollout", response_model=Rollout)
+        async def update_rollout(request: UpdateRolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_rollout(
+                rollout_id=request.rollout_id,
+                input=request.input if not isinstance(request.input, PydanticUnset) else UNSET,
+                mode=request.mode if not isinstance(request.mode, PydanticUnset) else UNSET,
+                resources_id=request.resources_id if not isinstance(request.resources_id, PydanticUnset) else UNSET,
+                status=request.status if not isinstance(request.status, PydanticUnset) else UNSET,
+                config=request.config if not isinstance(request.config, PydanticUnset) else UNSET,
+                metadata=request.metadata if not isinstance(request.metadata, PydanticUnset) else UNSET,
+            )
+
+        @self.app.post("/update_attempt", response_model=Attempt)
+        async def update_attempt(request: UpdateAttemptRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_attempt(
+                rollout_id=request.rollout_id,
+                attempt_id=request.attempt_id,
+                status=request.status if not isinstance(request.status, PydanticUnset) else UNSET,
+                worker_id=request.worker_id if not isinstance(request.worker_id, PydanticUnset) else UNSET,
+                last_heartbeat_time=(
+                    request.last_heartbeat_time if not isinstance(request.last_heartbeat_time, PydanticUnset) else UNSET
+                ),
+                metadata=request.metadata if not isinstance(request.metadata, PydanticUnset) else UNSET,
+            )
+
+    # Delegate methods
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        return await self._backend().start_rollout(input, mode, resources_id, metadata)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        return await self._backend().enqueue_rollout(input, mode, resources_id, metadata)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        return await self._backend().dequeue_rollout()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        return await self._backend().start_attempt(rollout_id)
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        return await self._backend().query_rollouts(status=status, rollout_ids=rollout_ids)
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        return await self._backend().query_attempts(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        return await self._backend().get_latest_attempt(rollout_id)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        return await self._backend().get_rollout_by_id(rollout_id)
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        return await self._backend().add_resources(resources)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        return await self._backend().update_resources(resources_id, resources)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        return await self._backend().get_resources_by_id(resources_id)
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        return await self._backend().get_latest_resources()
+
+    async def add_span(self, span: Span) -> Span:
+        return await self._backend().add_span(span)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        return await self._backend().get_next_span_sequence_id(rollout_id, attempt_id)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        return await self._backend().add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        return await self._backend().wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        return await self._backend().query_spans(rollout_id, attempt_id)
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        return await self._backend().update_rollout(
+            rollout_id=rollout_id,
+            input=input,
+            mode=mode,
+            resources_id=resources_id,
+            status=status,
+            config=config,
+            metadata=metadata,
+        )
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        return await self._backend().update_attempt(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            status=status,
+            worker_id=worker_id,
+            last_heartbeat_time=last_heartbeat_time,
+            metadata=metadata,
+        )
+
+
+# def _make_trace():
+#     tc = aiohttp.TraceConfig()
+#     # async def log_evt(session, context, params):
+#     #     print(f"[TRACE] {context}: {params}")
+#     tc.on_dns_resolvehost_start.append(lambda *a, **k: print("[TRACE] dns_start", a[-1].host))
+#     tc.on_connection_create_start.append(lambda *a, **k: print("[TRACE] conn_start"))
+#     tc.on_request_start.append(lambda *a, **k: print("[TRACE] req_start", a[-1].method, a[-1].url))
+#     tc.on_request_end.append(lambda *a, **k: print("[TRACE] req_end", a[-1].method, a[-1].url))
+#     tc.on_request_exception.append(lambda *a, **k: print("[TRACE] req_exc", a[-1].method, a[-1].url))
+#     return tc
+
+
+class LightningStoreClient(LightningStore):
+    """HTTP client that talks to a remote LightningStoreServer.
+
+    Args:
+        server_address: The address of the LightningStoreServer to connect to.
+        retry_delays:
+            Backoff schedule (seconds) used when the initial request fails for a
+            non-application reason. Each entry is a retry attempt.
+        health_retry_delays:
+            Delays between /health probes while waiting for the server to come back.
+    """
+
+    def __init__(
+        self,
+        server_address: str,
+        *,
+        retry_delays: Sequence[float] = (1.0, 2.0, 5.0),
+        health_retry_delays: Sequence[float] = (0.1, 0.2, 0.5),
+    ):
+        self.server_address = server_address.rstrip("/")
+        self._sessions: Dict[int, aiohttp.ClientSession] = {}  # id(loop) -> ClientSession
+        self._lock = threading.RLock()
+
+        # retry config
+        self._retry_delays = tuple(float(d) for d in retry_delays)
+        self._health_retry_delays = tuple(float(d) for d in health_retry_delays)
+
+        # Store whether the dequeue was successful in history
+        self._dequeue_was_successful: bool = False
+        self._dequeue_first_unsuccessful: bool = True
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        # In the proxy process, FastAPI middleware calls
+        # client_store.get_next_span_sequence_id(...). With
+        # reuse_session=True, _get_session() creates and caches a
+        # single ClientSession bound to the uvicorn event loop.
+        #
+        # Later, the OpenTelemetry exporter (LightningSpanExporter)
+        # runs its flush on its own private event loop (in a different
+        # thread) and calls client_store.add_otel_span(...) ->
+        # client_store.add_span(...).
+        #
+        # If we reuse one session across all, the exporter tries to reuse the
+        # same cached ClientSession that was created on the uvicorn
+        # loop. aiohttp.ClientSession is not loop-agnostic or
+        # thread-safe. Using it from another loop can hang on the
+        # first request. That's why we need a map from loop to session.
+
+        loop = asyncio.get_running_loop()
+        key = id(loop)
+        print("!!!!! _get_session received %s", key)
+        with self._lock:
+            print("!!!!! _get_session with lock")
+            sess = self._sessions.get(key)
+            if sess is None or sess.closed:
+                # connector = aiohttp.TCPConnector(
+                #     limit=64, limit_per_host=16, ttl_dns_cache=300, enable_cleanup_closed=True
+                # )
+                # sess = aiohttp.ClientSession(trace_configs=[_make_trace()])
+                sess = aiohttp.ClientSession()
+                self._sessions[key] = sess
+        print(self._sessions)
+        return sess
+
+    async def _wait_until_healthy(self, session: aiohttp.ClientSession) -> bool:
+        """
+        Probe the server's /health until it responds 200 or retries are exhausted.
+        Returns True if healthy, False otherwise.
+        """
+        logger.info(f"Waiting for server to be healthy at {self.server_address}/health")
+        for delay in [*self._health_retry_delays, 0.0]:
+            try:
+                async with session.get(f"{self.server_address}/health") as r:
+                    if r.status == 200:
+                        logger.info(f"Server is healthy at {self.server_address}/health")
+                        return True
+            except Exception:
+                # swallow and retry
+                if delay > 0.0:
+                    logger.warning(f"Server is not healthy yet. Retrying in {delay} seconds.")
+            if delay > 0.0:
+                await asyncio.sleep(delay)
+        logger.error(
+            f"Server is not healthy at {self.server_address}/health after {len(self._health_retry_delays)} retry attempts"
+        )
+        return False
+
+    async def _request_json(
+        self,
+        method: Literal["get", "post"],
+        path: str,
+        *,
+        json: Any | None = None,
+    ) -> Any:
+        """
+        Make an HTTP request with:
+
+        1) First attempt.
+        2) On network/session failures: probe /health until back, then retry
+           according to self._retry_delays.
+        3) On 4xx (e.g., 400 set by server exception handler): do not retry.
+
+        Returns parsed JSON (or raw JSON scalar like int).
+        Raises the last exception if all retries fail.
+        """
+        session = await self._get_session()
+        url = f"{self.server_address}{path if path.startswith('/') else '/'+path}"
+        print("$$$$$$ session acquired", url)
+
+        # attempt 0 is immediate, then follow retry schedule
+        attempts = (0.0,) + self._retry_delays
+        last_exc: Exception | None = None
+
+        for delay in attempts:
+            if delay:
+                logger.info(f"Waiting {delay} seconds before retrying {method}: {path}")
+                await asyncio.sleep(delay)
+            try:
+                http_call = getattr(session, method)
+                print("$$$$$$ http_call", http_call)
+                timeout = aiohttp.ClientTimeout(total=3.5, connect=1.0, sock_connect=1.0, sock_read=2.5)
+                async with http_call(url, json=json, timeout=timeout) as resp:
+                    print("$$$$$ resp", resp)
+                    resp.raise_for_status()
+                    return await resp.json()
+            except aiohttp.ClientResponseError as cre:
+                # Respect app-level 4xx as final (server marks app faults as 400)
+                # 4xx => application issue; do not retry (except 408 which is transient)
+                logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
+                if 400 <= cre.status < 500 and cre.status != 408:
+                    raise
+                # 5xx and others will be retried below if they raise
+                last_exc = cre
+                logger.info(f"5xx and other status codes will be retried. Retrying the request {method}: {path}")
+                # before next retry, ensure server is healthy
+                if not await self._wait_until_healthy(session):
+                    break  # server is not healthy, do not retry
+            except (
+                aiohttp.ServerDisconnectedError,
+                aiohttp.ClientConnectorError,
+                aiohttp.ClientOSError,
+                asyncio.TimeoutError,
+            ) as net_exc:
+                # Network/session issue: probe health before retrying
+                logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
+                last_exc = net_exc
+                logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
+                if not await self._wait_until_healthy(session):
+                    break  # server is not healthy, do not retry
+
+        # exhausted retries
+        assert last_exc is not None
+        raise last_exc
+
+    async def close(self):
+        """Close the HTTP session."""
+        with self._lock:
+            sessions = list(self._sessions.values())
+            self._sessions.clear()
+
+        # close them on their own loops to avoid warnings
+        async def _close(sess: aiohttp.ClientSession):
+            if not sess.closed:
+                await sess.close()
+
+        # If called from one loop, best-effort close here.
+        for s in sessions:
+            try:
+                await _close(s)
+            except RuntimeError:
+                # If created on a different loop/thread, schedule a thread-safe close
+                # Fallback: close without awaiting (library tolerates it in practice),
+                # or keep a per-loop shutdown hook where they were created.
+                pass
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        data = await self._request_json(
+            "post",
+            "/start_rollout",
+            json=RolloutRequest(input=input, mode=mode, resources_id=resources_id, metadata=metadata).model_dump(),
+        )
+        return AttemptedRollout.model_validate(data)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        data = await self._request_json(
+            "post",
+            "/enqueue_rollout",
+            json=RolloutRequest(input=input, mode=mode, resources_id=resources_id, metadata=metadata).model_dump(),
+        )
+        return Rollout.model_validate(data)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """
+        Dequeue a rollout from the server queue.
+
+        Returns:
+            AttemptedRollout if a rollout is available, None if queue is empty.
+
+        Note:
+            This method does NOT retry on failures. If any exception occurs (network error,
+            server error, etc.), it logs the error and returns None immediately.
+        """
+        session = await self._get_session()
+        url = f"{self.server_address}/dequeue_rollout"
+        try:
+            async with session.get(url) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+                self._dequeue_was_successful = True
+                return AttemptedRollout.model_validate(data) if data else None
+        except Exception as e:
+            if self._dequeue_was_successful:
+                if self._dequeue_first_unsuccessful:
+                    logger.error(f"dequeue_rollout failed with exception: {e}", exc_info=True)
+                    self._dequeue_first_unsuccessful = False
+            # Else ignore the exception because the server is not ready yet
+            return None
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        data = await self._request_json(
+            "post",
+            "/start_attempt",
+            json=RolloutId(rollout_id=rollout_id).model_dump(),
+        )
+        return AttemptedRollout.model_validate(data)
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        data = await self._request_json(
+            "post",
+            "/query_rollouts",
+            json=QueryRolloutsRequest(
+                status=list(status) if status else None,
+                rollout_ids=list(rollout_ids) if rollout_ids else None,
+            ).model_dump(),
+        )
+        return [Rollout.model_validate(item) for item in data]
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        data = await self._request_json("get", f"/query_attempts/{rollout_id}")
+        return [Attempt.model_validate(item) for item in data]
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """
+        Get the latest attempt for a rollout.
+
+        Args:
+            rollout_id: ID of the rollout to query.
+
+        Returns:
+            Attempt if found, None if not found or if all retries are exhausted.
+
+        Note:
+            This method retries on transient failures (network errors, 5xx status codes).
+            If all retries fail, it logs the error and returns None instead of raising an exception.
+        """
+        try:
+            data = await self._request_json("get", f"/get_latest_attempt/{rollout_id}")
+            return Attempt.model_validate(data) if data else None
+        except Exception as e:
+            logger.error(f"get_latest_attempt failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
+            return None
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        """
+        Get a rollout by its ID.
+
+        Args:
+            rollout_id: ID of the rollout to retrieve.
+
+        Returns:
+            Rollout if found, None if not found or if all retries are exhausted.
+
+        Note:
+            This method retries on transient failures (network errors, 5xx status codes).
+            If all retries fail, it logs the error and returns None instead of raising an exception.
+        """
+        try:
+            data = await self._request_json("get", f"/get_rollout_by_id/{rollout_id}")
+            return Rollout.model_validate(data) if data else None
+        except Exception as e:
+            logger.error(f"get_rollout_by_id failed after all retries for rollout_id={rollout_id}: {e}", exc_info=True)
+            return None
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        request = AddResourcesRequest(resources=resources)
+        data = await self._request_json("post", "/add_resources", json=request.model_dump())
+        return ResourcesUpdate.model_validate(data)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        data = await self._request_json(
+            "post",
+            "/update_resources",
+            json=ResourcesUpdate(resources_id=resources_id, resources=resources).model_dump(),
+        )
+        return ResourcesUpdate.model_validate(data)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """
+        Get resources by their ID.
+
+        Args:
+            resources_id: ID of the resources to retrieve.
+
+        Returns:
+            ResourcesUpdate if found, None if not found or if all retries are exhausted.
+
+        Note:
+            This method retries on transient failures (network errors, 5xx status codes).
+            If all retries fail, it logs the error and returns None instead of raising an exception.
+        """
+        try:
+            data = await self._request_json("get", f"/get_resources_by_id/{resources_id}")
+            return ResourcesUpdate.model_validate(data) if data else None
+        except Exception as e:
+            logger.error(
+                f"get_resources_by_id failed after all retries for resources_id={resources_id}: {e}", exc_info=True
+            )
+            return None
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """
+        Get the latest resources.
+
+        Returns:
+            ResourcesUpdate if found, None if not found or if all retries are exhausted.
+
+        Note:
+            This method retries on transient failures (network errors, 5xx status codes).
+            If all retries fail, it logs the error and returns None instead of raising an exception.
+        """
+        try:
+            data = await self._request_json("get", "/get_latest_resources")
+            return ResourcesUpdate.model_validate(data) if data else None
+        except Exception as e:
+            logger.error(f"get_latest_resources failed after all retries: {e}", exc_info=True)
+            return None
+
+    async def add_span(self, span: Span) -> Span:
+        print("$$$$$$ add_span received")
+        data = await self._request_json("post", "/add_span", json=span.model_dump(mode="json"))
+        return Span.model_validate(data)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        data = await self._request_json("get", f"/get_next_span_sequence_id/{rollout_id}/{attempt_id}")
+        # endpoint returns a plain JSON number
+        return int(data)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        # unchanged logic, now benefits from retries inside add_span/get_next_span_sequence_id
+        print("$$$$$$ add_otel_span received")
+        if sequence_id is None:
+            sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
+        span = Span.from_opentelemetry(
+            readable_span,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+        )
+        print("$$$$$$ span created")
+        await self.add_span(span)
+        return span
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        if timeout is not None and timeout > 0.1:
+            raise ValueError(
+                "Timeout must be less than 0.1 seconds in LightningStoreClient to avoid blocking the event loop"
+            )
+        data = await self._request_json(
+            "post",
+            "/wait_for_rollouts",
+            json=WaitForRolloutsRequest(rollout_ids=rollout_ids, timeout=timeout).model_dump(),
+        )
+        return [Rollout.model_validate(item) for item in data]
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        path = f"/query_spans/{rollout_id}"
+        if attempt_id is not None:
+            path += f"?attempt_id={attempt_id}"
+        data = await self._request_json("get", path)
+        return [Span.model_validate(item) for item in data]
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        payload: Dict[str, Any] = {"rollout_id": rollout_id}
+        if not isinstance(input, Unset):
+            payload["input"] = input
+        if not isinstance(mode, Unset):
+            payload["mode"] = mode
+        if not isinstance(resources_id, Unset):
+            payload["resources_id"] = resources_id
+        if not isinstance(status, Unset):
+            payload["status"] = status
+        if not isinstance(config, Unset):
+            payload["config"] = config.model_dump()
+        if not isinstance(metadata, Unset):
+            payload["metadata"] = metadata
+
+        data = await self._request_json("post", "/update_rollout", json=payload)
+        return Rollout.model_validate(data)
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        payload: Dict[str, Any] = {
+            "rollout_id": rollout_id,
+            "attempt_id": attempt_id,
+        }
+        if not isinstance(status, Unset):
+            payload["status"] = status
+        if not isinstance(worker_id, Unset):
+            payload["worker_id"] = worker_id
+        if not isinstance(last_heartbeat_time, Unset):
+            payload["last_heartbeat_time"] = last_heartbeat_time
+        if not isinstance(metadata, Unset):
+            payload["metadata"] = metadata
+
+        data = await self._request_json("post", "/update_attempt", json=payload)
+        return Attempt.model_validate(data)

agentlightning/store/memory.py

@@ -0,0 +1,703 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import hashlib
+import logging
+import threading
+import time
+import uuid
+from collections import deque
+from typing import Any, Callable, Counter, Dict, List, Literal, Optional, Sequence, TypeVar, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, Unset, is_finished, is_queuing
+from .utils import healthcheck, propagate_status
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def _healthcheck_wrapper(func: T_callable) -> T_callable:
+    """
+    Decorator to run the watchdog healthcheck **before** executing the decorated method.
+    Only runs if the store has a watchdog configured.
+    Prevents recursive healthcheck execution using a flag on the store instance.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: InMemoryLightningStore, *args: Any, **kwargs: Any) -> Any:
+        # Check if healthcheck is already running to prevent recursion
+        if getattr(self, "_healthcheck_running", False):
+            # Skip healthcheck if already running
+            return await func(self, *args, **kwargs)
+
+        # Set flag to prevent recursive healthcheck calls
+        # This flag is not asyncio/thread-safe, but it doesn't matter
+        self._healthcheck_running = True  # type: ignore
+        try:
+            # The following methods should live inside one lock.
+            await self._healthcheck()  # pyright: ignore[reportPrivateUsage]
+        finally:
+            # Always clear the flag, even if healthcheck fails
+            self._healthcheck_running = False  # type: ignore
+
+        # Execute the original method
+        # This should be outside the lock.
+        return await func(self, *args, **kwargs)
+
+    return cast(T_callable, wrapper)
+
+
+def _generate_resources_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "rs-" + short_id
+
+
+def _generate_rollout_id() -> str:
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:12]
+    return "ro-" + short_id
+
+
+def _generate_attempt_id() -> str:
+    """We don't need that long because attempts are limited to rollouts."""
+    short_id = hashlib.sha1(uuid.uuid4().bytes).hexdigest()[:8]
+    return "at-" + short_id
+
+
+class InMemoryLightningStore(LightningStore):
+    """
+    In-memory implementation of LightningStore using Python data structures.
+    Thread-safe and async-compatible but data is not persistent.
+
+    The methods in this class should generally not call each other,
+    especially those that are locked.
+    """
+
+    def __init__(self):
+        self._lock = asyncio.Lock()
+
+        # Task queue and rollouts storage
+        self._task_queue: deque[Rollout] = deque()
+        self._rollouts: Dict[str, Rollout] = {}
+
+        # Resources storage (similar to legacy server.py)
+        self._resources: Dict[str, ResourcesUpdate] = {}
+        self._latest_resources_id: Optional[str] = None
+
+        # Spans storage
+        self._spans: Dict[str, List[Span]] = {}  # rollout_id -> list of spans
+        self._span_sequence_ids: Dict[str, int] = Counter()  # rollout_id -> sequence_id
+
+        # Attempt tracking
+        self._attempts: Dict[str, List[Attempt]] = {}  # rollout_id -> list of attempts
+
+        # Completion tracking for wait_for_rollouts (cross-loop safe)
+        self._completion_events: Dict[str, threading.Event] = {}
+
+    @_healthcheck_wrapper
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """
+        Notify the store that I'm about to run a rollout.
+        """
+        async with self._lock:
+            rollout_id = _generate_rollout_id()
+            current_time = time.time()
+
+            rollout = Rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id or self._latest_resources_id,
+                start_time=current_time,
+                status="preparing",
+                metadata=metadata or {},
+            )
+
+            # Create the initial attempt
+            attempt_id = _generate_attempt_id()
+            attempt = Attempt(
+                rollout_id=rollout.rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=1,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            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
+            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        """
+        Adds a new task to the queue with specific metadata and returns its unique ID.
+        """
+        async with self._lock:
+            rollout_id = _generate_rollout_id()
+            current_time = time.time()
+
+            rollout = Rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id or self._latest_resources_id,
+                start_time=current_time,
+                status="queuing",  # should be queuing
+                metadata=metadata or {},
+            )
+
+            self._rollouts[rollout.rollout_id] = rollout
+            self._task_queue.append(rollout)  # add it to the end of the queue
+            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+
+            return rollout
+
+    @_healthcheck_wrapper
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """
+        Retrieves the next task from the queue without blocking.
+        Returns None if the queue is empty.
+
+        Will set the rollout status to preparing and create a new attempt.
+        """
+        async with self._lock:
+            # Keep looking until we find a rollout that's still in queuing status
+            # or the queue is empty
+            while self._task_queue:
+                rollout = self._task_queue.popleft()
+
+                # Check if rollout is still in a queuing state
+                # (it might have been updated to a different status while in queue)
+                if is_queuing(rollout):
+                    # Update status to preparing
+                    rollout.status = "preparing"
+
+                    # Create a new attempt (could be first attempt or retry)
+                    attempt_id = _generate_attempt_id()
+                    current_time = time.time()
+
+                    # Get existing attempts to determine sequence number
+                    existing_attempts = self._attempts.get(rollout.rollout_id, [])
+                    sequence_id = len(existing_attempts) + 1
+
+                    attempt = Attempt(
+                        rollout_id=rollout.rollout_id,
+                        attempt_id=attempt_id,
+                        sequence_id=sequence_id,
+                        start_time=current_time,
+                        status="preparing",
+                    )
+
+                    if rollout.rollout_id not in self._attempts:
+                        self._attempts[rollout.rollout_id] = []
+                    self._attempts[rollout.rollout_id].append(attempt)
+
+                    return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+                # If not in queuing state, skip this rollout and continue
+                # (it was updated externally and should not be processed)
+
+            # No valid rollouts found
+            return None
+
+    @_healthcheck_wrapper
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """
+        Create a new attempt for a given rollout ID and return the attempt details.
+        """
+        async with self._lock:
+            # Get the rollout
+            rollout = self._rollouts.get(rollout_id)
+            if not rollout:
+                raise ValueError(f"Rollout {rollout_id} not found")
+
+            # Get existing attempts to determine sequence number
+            existing_attempts = self._attempts.get(rollout_id, [])
+            sequence_id = len(existing_attempts) + 1
+
+            # We don't care whether the max attempts have reached or not
+            # This attempt is from user trigger
+
+            # Create new attempt
+            attempt_id = _generate_attempt_id()
+            current_time = time.time()
+
+            attempt = Attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=sequence_id,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            # Add attempt to storage
+            if rollout_id not in self._attempts:
+                self._attempts[rollout_id] = []
+            self._attempts[rollout_id].append(attempt)
+
+            self._completion_events.setdefault(rollout.rollout_id, threading.Event())
+
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[Rollout]:
+        """
+        Query and retrieve rollouts filtered by their status and rollout ids.
+        If no status is provided, returns all rollouts.
+        """
+        async with self._lock:
+            rollouts = list(self._rollouts.values())
+
+            # Filter by rollout_ids if provided
+            if rollout_ids is not None:
+                rollout_ids_set = set(rollout_ids)
+                rollouts = [rollout for rollout in rollouts if rollout.rollout_id in rollout_ids_set]
+
+            # Filter by status if provided
+            if status is not None:
+                status_set = set(status)
+                rollouts = [rollout for rollout in rollouts if rollout.status in status_set]
+
+            return rollouts
+
+    @_healthcheck_wrapper
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        """
+        Safely retrieves a specific rollout by its ID.
+        """
+        async with self._lock:
+            return self._rollouts.get(rollout_id)
+
+    @_healthcheck_wrapper
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """
+        Query and retrieve all attempts associated with a specific rollout ID.
+        Returns an empty list if no attempts are found.
+        """
+        async with self._lock:
+            return self._attempts.get(rollout_id, [])
+
+    @_healthcheck_wrapper
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """
+        Safely retrieves the latest attempt for a given rollout ID.
+        """
+        async with self._lock:
+            attempts = self._attempts.get(rollout_id, [])
+            if not attempts:
+                return None
+            return max(attempts, key=lambda a: a.sequence_id)
+
+    @_healthcheck_wrapper
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+        """
+        resources_id = _generate_resources_id()
+        async with self._lock:
+            update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+            self._resources[resources_id] = update
+            self._latest_resources_id = resources_id
+            return update
+
+    @_healthcheck_wrapper
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+        """
+        async with self._lock:
+            update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+            self._resources[resources_id] = update
+            self._latest_resources_id = resources_id
+            return update
+
+    @_healthcheck_wrapper
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves a specific version of named resources by its ID.
+        """
+        async with self._lock:
+            return self._resources.get(resources_id)
+
+    @_healthcheck_wrapper
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves the latest version of named resources.
+        """
+        async with self._lock:
+            if self._latest_resources_id:
+                return self._resources.get(self._latest_resources_id)
+            return None
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """
+        Get the next span sequence ID for a given rollout and attempt.
+        The number is strictly increasing for each rollout.
+        The store will not issue the same sequence ID twice.
+        """
+        async with self._lock:
+            self._span_sequence_ids[rollout_id] += 1
+            return self._span_sequence_ids[rollout_id]
+
+    async def add_span(self, span: Span) -> Span:
+        """Persist a pre-converted span."""
+        async with self._lock:
+            self._span_sequence_ids[span.rollout_id] = max(self._span_sequence_ids[span.rollout_id], span.sequence_id)
+            return await self._add_span_unlocked(span)
+
+    async def add_otel_span(
+        self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
+    ) -> Span:
+        """Add an opentelemetry span to the store."""
+        async with self._lock:
+            if sequence_id is None:
+                # Issue a new sequence ID for the rollout
+                self._span_sequence_ids[rollout_id] += 1
+                sequence_id = self._span_sequence_ids[rollout_id]
+            else:
+                # Comes from a provided sequence ID
+                # Make sure our counter is strictly increasing
+                self._span_sequence_ids[rollout_id] = max(self._span_sequence_ids[rollout_id], sequence_id)
+
+            span = Span.from_opentelemetry(
+                readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id
+            )
+            await self._add_span_unlocked(span)
+            return span
+
+    async def _add_span_unlocked(self, span: Span) -> Span:
+        rollout = self._rollouts.get(span.rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {span.rollout_id} not found")
+        attempts = self._attempts.get(span.rollout_id, [])
+        current_attempt = next((a for a in attempts if a.attempt_id == span.attempt_id), None)
+        latest_attempt = max(attempts, key=lambda a: a.sequence_id) if attempts else None
+        if not current_attempt:
+            raise ValueError(f"Attempt {span.attempt_id} not found for rollout {span.rollout_id}")
+        if not latest_attempt:
+            raise ValueError(f"No attempts found for rollout {span.rollout_id}")
+
+        if span.rollout_id not in self._spans:
+            self._spans[span.rollout_id] = []
+        self._spans[span.rollout_id].append(span)
+
+        # Update attempt heartbeat
+        current_attempt.last_heartbeat_time = time.time()
+        if current_attempt.status in ["preparing", "unresponsive", "timeout"]:
+            current_attempt.status = "running"
+
+        # If the status has already timed out or failed, do not change it
+
+        # Update rollout status if it's the latest attempt
+        if current_attempt == latest_attempt:
+            if rollout.status == "preparing":
+                rollout.status = "running"
+            elif rollout.status in ["queuing", "requeuing"]:
+                try:
+                    self._task_queue.remove(rollout)
+                except ValueError:
+                    logger.warning(
+                        f"Trying to remove rollout {rollout.rollout_id} from the queue but it's not in the queue."
+                    )
+                rollout.status = "running"
+
+        return span
+
+    @_healthcheck_wrapper
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        """
+        Wait for specified rollouts to complete with a timeout.
+        Returns the completed rollouts, potentially incomplete if timeout is reached.
+
+        This method does not change the state of the store.
+        """
+        completed_rollouts: List[Rollout] = []
+
+        async def wait_for_rollout(rollout_id: str):
+            # First check if already completed
+            async with self._lock:
+                rollout = self._rollouts.get(rollout_id)
+                if rollout and is_finished(rollout):
+                    completed_rollouts.append(rollout)
+                    return
+
+            # No timeout, return immediately
+            if timeout is not None and timeout <= 0:
+                return
+
+            # If not completed and we have an event, wait for completion
+            if rollout_id in self._completion_events:
+                evt = self._completion_events[rollout_id]
+
+                # Wait for the event with proper timeout handling
+                # evt.wait() returns True if event was set, False if timeout occurred
+                if timeout is None:
+                    # Wait indefinitely by polling with finite timeouts
+                    # This allows threads to exit cleanly on shutdown
+                    while True:
+                        result = await asyncio.to_thread(evt.wait, 10.0)  # Poll every 10 seconds
+                        if result:  # Event was set
+                            break
+                        # Loop and check again (continues indefinitely since timeout=None)
+                else:
+                    # Wait with the specified timeout
+                    result = await asyncio.to_thread(evt.wait, timeout)
+
+                # If event was set (not timeout), check if rollout is finished
+                if result:
+                    async with self._lock:
+                        rollout = self._rollouts.get(rollout_id)
+                        if rollout and is_finished(rollout):
+                            completed_rollouts.append(rollout)
+
+            # Rollout not found, return
+
+        # Wait for all rollouts concurrently
+        await asyncio.gather(*[wait_for_rollout(rid) for rid in rollout_ids], return_exceptions=True)
+
+        return completed_rollouts
+
+    @_healthcheck_wrapper
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """
+        Query and retrieve all spans associated with a specific rollout ID.
+        Returns an empty list if no spans are found.
+        """
+        async with self._lock:
+            spans = self._spans.get(rollout_id, [])
+            if attempt_id is None:
+                return spans
+            elif attempt_id == "latest":
+                # Find the latest attempt_id
+                if not spans:
+                    return []
+                latest_attempt = max(spans, key=lambda s: s.sequence_id if s.attempt_id else "").attempt_id
+                return [s for s in spans if s.attempt_id == latest_attempt]
+            else:
+                return [s for s in spans if s.attempt_id == attempt_id]
+
+    @_healthcheck_wrapper
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        """
+        Update the rollout status and related metadata.
+        """
+        async with self._lock:
+            return await self._update_rollout_unlocked(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    @_healthcheck_wrapper
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        """
+        Update a specific or latest attempt for a given rollout.
+        """
+        async with self._lock:
+            attempt = await self._update_attempt_unlocked(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )
+
+        return attempt
+
+    async def _update_rollout_unlocked(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        # No lock inside this one.
+        rollout = self._rollouts.get(rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        # Update fields if they are not UNSET
+        if not isinstance(input, Unset):
+            rollout.input = input
+        if not isinstance(mode, Unset):
+            rollout.mode = mode
+        if not isinstance(resources_id, Unset):
+            rollout.resources_id = resources_id
+        if not isinstance(status, Unset):
+            rollout.status = status
+        if not isinstance(config, Unset):
+            rollout.config = config
+        if not isinstance(metadata, Unset):
+            rollout.metadata = metadata
+
+        # Set end time for finished rollouts
+        # Rollout is only finished when it succeeded or fail with no more retries.
+        if not isinstance(status, Unset) and is_finished(rollout):
+            rollout.end_time = time.time()
+            # Signal completion
+            if rollout_id in self._completion_events:
+                self._completion_events[rollout_id].set()
+
+        # If requeuing, add back to queue
+        elif is_queuing(rollout) and rollout not in self._task_queue:
+            self._task_queue.append(rollout)
+
+        # 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:
+                self._task_queue.remove(rollout)
+            except ValueError:
+                # Another coroutine may have already removed the rollout from the queue.
+                logger.warning(
+                    f"Trying to remove rollout {rollout.rollout_id} from the queue but it's not in the queue."
+                )
+
+        # Re-validate the rollout to ensure legality
+        Rollout.model_validate(rollout.model_dump())
+
+        return rollout
+
+    async def _update_attempt_unlocked(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        # No lock, but with status propagation.
+        rollout = self._rollouts.get(rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        attempts = self._attempts.get(rollout_id, [])
+        if not attempts:
+            raise ValueError(f"No attempts found for rollout {rollout_id}")
+
+        latest_attempt = max(attempts, key=lambda a: a.sequence_id)
+
+        # Find the attempt to update
+        if attempt_id == "latest":
+            attempt = latest_attempt
+        else:
+            attempt = next((a for a in attempts if a.attempt_id == attempt_id), None)
+            if not attempt:
+                raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
+
+        # Update fields if they are not UNSET
+        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
+        if not isinstance(last_heartbeat_time, Unset):
+            attempt.last_heartbeat_time = last_heartbeat_time
+        if not isinstance(metadata, Unset):
+            attempt.metadata = metadata
+
+        # Re-validate the attempt to ensure legality
+        Attempt.model_validate(attempt.model_dump())
+
+        if attempt == latest_attempt:
+
+            async def _update_status(rollout_id: str, status: RolloutStatus) -> Rollout:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            # Propagate the status to the rollout
+            await propagate_status(
+                _update_status,
+                attempt,
+                rollout.config,
+            )
+
+        return attempt
+
+    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"]:
+                    all_attempts = self._attempts.get(rollout.rollout_id, [])
+                    if not all_attempts:
+                        # The rollout is running but has no attempts, this should not happen
+                        logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
+                        continue
+                    latest_attempt = max(all_attempts, key=lambda a: a.sequence_id)
+                    running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
+
+            async def _update_attempt_status(rollout_id: str, attempt_id: str, status: AttemptStatus) -> Attempt:
+                return await self._update_attempt_unlocked(rollout_id, attempt_id, status=status)
+
+            async def _update_rollout_status(rollout_id: str, status: RolloutStatus) -> Rollout:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            await healthcheck(
+                running_rollouts,
+                _update_rollout_status,
+                _update_attempt_status,
+            )

agentlightning/store/sqlite.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# TODO: Implement this

agentlightning/store/threading.py

@@ -0,0 +1,171 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, Unset
+
+
+class LightningStoreThreaded(LightningStore):
+    """Facade that delegates all store operations to a underlying store instance.
+
+    The operations are guaranteed to be thread-safe.
+    Make sure the threaded stores are instantiated before initializing the threads.
+    """
+
+    def __init__(self, store: LightningStore) -> None:
+        super().__init__()  # watchdog relies on the underlying store
+        self.store = store
+        self._lock = threading.Lock()
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_rollout(input, mode, resources_id, metadata)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_rollout()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id)
+
+    async def query_rollouts(
+        self,
+        *,
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> List[Rollout]:
+        with self._lock:
+            return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        with self._lock:
+            return await self.store.query_attempts(rollout_id)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        with self._lock:
+            return await self.store.get_rollout_by_id(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        with self._lock:
+            return await self.store.get_latest_attempt(rollout_id)
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.add_resources(resources)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.update_resources(resources_id, resources)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_resources_by_id(resources_id)
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_latest_resources()
+
+    async def add_span(self, span: Span) -> Span:
+        with self._lock:
+            return await self.store.add_span(span)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        with self._lock:
+            return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        # This method does not change the state of the store, and it's not thread-safe.
+        return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        with self._lock:
+            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        with self._lock:
+            return await self.store.query_spans(rollout_id, attempt_id)
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.update_rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        with self._lock:
+            return await self.store.update_attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )

agentlightning/store/utils.py

@@ -0,0 +1,126 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from typing import Awaitable, Callable, List, cast
+
+from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
+
+UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
+UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
+
+
+async def propagate_status(
+    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
+    attempt: Attempt,
+    config: RolloutConfig,
+) -> Rollout:
+    """
+    Propagate the status of an attempt to the rollout.
+
+    The rollout should be made sure in a state to be outdated.
+    Requeue the rollout if it should be retried.
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+    """
+    # Propagate the status directly to the rollout
+    if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
+        return await update_rollout_status(
+            attempt.rollout_id,
+            attempt.status,
+        )
+
+    if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
+        # Check if this status should trigger a retry
+        if attempt.status in config.retry_condition:
+            # If we haven't exceeded max attempts, retry
+            if attempt.sequence_id < config.max_attempts:
+                return await update_rollout_status(
+                    attempt.rollout_id,
+                    "requeuing",
+                )
+
+        # If we can't retry or shouldn't retry, mark as failed
+        return await update_rollout_status(
+            attempt.rollout_id,
+            "failed",
+        )
+
+    raise ValueError(f"Invalid attempt status: {attempt.status}")
+
+
+async def healthcheck(
+    rollouts: List[AttemptedRollout],
+    update_rollout_status: UpdateRolloutStatus,
+    update_attempt_status: UpdateAttemptStatus,
+) -> None:
+    """
+    Perform health check on all running rollouts in the store.
+
+    This method should be called periodically to:
+    1. Update rollout status to failed to succeeded when the attempt is done
+    2. Check for unresponsive attempts (no heartbeat or spans for a while)
+    3. Check for timed-out rollouts (running too long since start_time)
+    4. Update attempt/rollout status accordingly
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+
+    Args:
+        store: The LightningStore instance to check rollouts from
+    """
+    current_time = time.time()
+
+    for rollout in rollouts:
+        config = rollout.config  # policy for retry and timeout
+
+        # Get the latest attempt for this rollout
+        latest_attempt = rollout.attempt
+        if not latest_attempt:
+            continue
+
+        # Check if the attempt has already failed or succeeded
+        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
+            await propagate_status(update_rollout_status, latest_attempt, config)
+            continue
+
+        # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
+        if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "timeout",
+            )
+            continue
+
+        # Check for unresponsive condition (based on last heartbeat)
+        if latest_attempt.last_heartbeat_time:
+            if latest_attempt.status == "preparing":
+                # If still preparing, mark it as running
+                latest_attempt = await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "running",
+                )
+
+            # Haven't received heartbeat for a while
+            if (
+                config.unresponsive_seconds is not None
+                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
+            ):
+                await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "unresponsive",
+                )
+                continue
+
+        # Check if there's no last heartbeat (no spans) at all
+        if (
+            latest_attempt.last_heartbeat_time is None
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.start_time > config.unresponsive_seconds
+        ):
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "unresponsive",
+            )

agentlightning/tracer/__init__.py

@@ -1,3 +1,7 @@
-from .base import BaseTracer
+# Copyright (c) Microsoft. All rights reserved.
+
 from .agentops import AgentOpsTracer
-from .triplet import TripletExporter
+from .base import BaseTracer
+from .otel import OtelTracer
+
+__all__ = ["AgentOpsTracer", "BaseTracer", "OtelTracer"]

agentlightning/tracer/agentops.py

@@ -1,20 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
+import threading
 from contextlib import contextmanager
-from typing import List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Awaitable, Iterator, List, Optional
 
-import agentops.sdk.core
 import agentops
+import agentops.sdk.core
 from agentops.sdk.core import TracingCore
 from agentops.sdk.processors import SpanProcessor
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.instrumentation.agentops import AgentOpsServerManager
 from agentlightning.instrumentation import instrument_all, uninstrument_all
-from .base import BaseTracer
+from agentlightning.instrumentation.agentops import AgentOpsServerManager
+from agentlightning.store.base import LightningStore
 
+from .base import BaseTracer
 
 if TYPE_CHECKING:
     from agentops.integration.callbacks.langchain import LangchainCallbackHandler
@@ -65,12 +70,12 @@ class AgentOpsTracer(BaseTracer):
         logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
         return state
 
-    def __setstate__(self, 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, **kwargs):
+    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()
@@ -122,7 +127,7 @@ class AgentOpsTracer(BaseTracer):
                 )
 
             if not agentops.get_client().initialized:
-                agentops.init()
+                agentops.init()  # type: ignore
                 logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
             else:
                 logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
@@ -132,11 +137,13 @@ class AgentOpsTracer(BaseTracer):
         try:
             # new versions
             instance = agentops.sdk.core.tracer
-            instance.provider.add_span_processor(self._lightning_span_processor)
+            # 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()
-            instance._provider.add_span_processor(self._lightning_span_processor)
+            instance = TracingCore.get_instance()  # type: ignore
+            instance._provider.add_span_processor(self._lightning_span_processor)  # type: ignore
 
     def teardown_worker(self, worker_id: int) -> None:
         super().teardown_worker(worker_id)
@@ -146,12 +153,22 @@ class AgentOpsTracer(BaseTracer):
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None):
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[LightningSpanProcessor]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
         Args:
             name: Optional name for the tracing context.
+            store: Optional store to add the spans to.
+            rollout_id: Optional rollout ID to add the spans to.
+            attempt_id: Optional attempt ID to add the spans to.
 
         Yields:
             The LightningSpanProcessor instance to collect spans.
@@ -159,8 +176,15 @@ class AgentOpsTracer(BaseTracer):
         if not self._lightning_span_processor:
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
 
-        with self._lightning_span_processor:
-            yield self._lightning_span_processor
+        if store is not None and rollout_id is not None and attempt_id is not None:
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx as processor:
+                yield processor
+        elif store is None and rollout_id is None and attempt_id is None:
+            with self._lightning_span_processor:
+                yield self._lightning_span_processor
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
 
     def get_last_trace(self) -> List[ReadableSpan]:
         """
@@ -173,7 +197,7 @@ class AgentOpsTracer(BaseTracer):
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
         return self._lightning_span_processor.spans()
 
-    def get_langchain_callback_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
         """
         Get the Langchain callback handler for integrating with Langchain.
 
@@ -197,18 +221,125 @@ class AgentOpsTracer(BaseTracer):
             )
         return LangchainCallbackHandler(api_key=api_key, tags=tags)
 
+    get_langchain_callback_handler = get_langchain_handler  # alias
+
+
+async def heartbeat(name="exporter-loop", period=0.5):
+    import asyncio
+    import time
+
+    last = time.perf_counter()
+    while True:
+        await asyncio.sleep(period)
+        now = time.perf_counter()
+        dt = now - last
+        last = now
+        if dt > period * 4:  # e.g., >2s if period=0.5s
+            print("!!!!!!! [%s] loop stall detected: slept %.3fs (expected %.3fs)" % (name, dt, period))
+
+
+import asyncio
+import logging
+
+# logging.basicConfig(level=logging.DEBUG)
+# asyncio.get_event_loop().set_debug(True)
+import time
+
+
+def debug_dump(loop):
+    while True:
+        try:
+            print("=== Pending tasks ===")
+            for t in asyncio.all_tasks(loop):
+                if not t.done():
+                    print(t, "awaiting", t.get_coro())
+                    t.print_stack()
+        except Exception:
+            pass
+        time.sleep(5)
+
 
 class LightningSpanProcessor(SpanProcessor):
+    def __init__(self):
+        self._spans: List[ReadableSpan] = []
 
-    _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 = 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
+        self._loop.set_debug(True)
+        asyncio.set_event_loop(loop)
+        self._loop_ready.set()
+
+        thread = threading.Thread(target=debug_dump, args=(loop,), daemon=True)
+        thread.start()
+        # asyncio.create_task(heartbeat())
+        loop.run_forever()
+        loop.close()
 
     def __enter__(self):
         self._last_trace = None
         self._spans = []
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        pass
+    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]:
         """
@@ -220,6 +351,22 @@ class LightningSpanProcessor(SpanProcessor):
         """
         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.
@@ -227,14 +374,40 @@ class LightningSpanProcessor(SpanProcessor):
         Args:
             span: The span that has ended.
         """
+        import traceback
+
+        # print("ON_END")
+        # print(traceback.format_stack())
         # 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
+                print("!!! before,")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", len(self._loop._scheduled))
+                if self._loop._scheduled:
+                    print("First in the queue:", self._loop._scheduled[0])
+                print("..... Current thread: ", threading.current_thread())
+                print("..... Loop thread: ", self._loop_thread)
+                if self._loop_thread.ident == threading.current_thread().ident:
+                    traceback.print_stack()
+                    print("Span content: ", span.attributes)
+                from opentelemetry.instrumentation.utils import suppress_instrumentation
+
+                with suppress_instrumentation():
+                    self._await_in_loop(
+                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                        timeout=30.0,
+                    )
+                print("!!! after,")
+                print("All tasks")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", self._loop._scheduled)
+            except Exception:
+                # log; on_end MUST NOT raise
+                logger.exception(f"Error adding span to store: {span.name}")
+
         self._spans.append(span)
-
-    def shutdown(self) -> None:
-        pass
-
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
-        return True

agentlightning/tracer/base.py

@@ -1,9 +1,21 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
 from contextlib import contextmanager
-from typing import Iterator, List, Optional, Callable, Any, Awaitable
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, List, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.base import LightningStore
 from agentlightning.types import ParallelWorkerBase
 
+if TYPE_CHECKING:
+    from langchain.callbacks.base import BaseCallbackHandler
+
+logger = logging.getLogger(__name__)
+
 
 class BaseTracer(ParallelWorkerBase):
     """
@@ -35,13 +47,20 @@ class BaseTracer(ParallelWorkerBase):
 
     # Process the trace data
     if trace_tree:
-        rl_triplets = TripletExporter().export(spans)
+        rl_triplets = TracerTraceToTriplet().adapt(spans)
         # ... do something with the triplets
     ```
     """
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[Any]:
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[Any]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
@@ -50,8 +69,13 @@ class BaseTracer(ParallelWorkerBase):
         within the `with` block are collected and made available via
         `get_last_trace`.
 
+        If a store is provided, the spans will be added to the store when tracing.
+
         Args:
             name: The name for the root span of this trace context.
+            store: The store to add the spans to.
+            rollout_id: The rollout ID to add the spans to.
+            attempt_id: The attempt ID to add the spans to.
         """
         raise NotImplementedError()
 
@@ -64,7 +88,7 @@ class BaseTracer(ParallelWorkerBase):
         """
         raise NotImplementedError()
 
-    def trace_run(self, func: Callable, *args, **kwargs) -> Any:
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single synchronous function.
 
@@ -79,7 +103,7 @@ class BaseTracer(ParallelWorkerBase):
         with self.trace_context(name=func.__name__):
             return func(*args, **kwargs)
 
-    async def trace_run_async(self, func: Callable[..., Awaitable], *args, **kwargs) -> Any:
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single asynchronous function.
 
@@ -93,3 +117,11 @@ class BaseTracer(ParallelWorkerBase):
         """
         with self.trace_context(name=func.__name__):
             return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:
+        """Get a handler to install in langchain agent callback.
+
+        Agents are expected to use this handler in their agents to enable tracing.
+        """
+        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
+        return None

agentlightning/tracer/http.py

@@ -1,25 +1,25 @@
-from contextlib import contextmanager
-from typing import Iterator, List, Optional, Any, Dict, Callable, Awaitable
-import logging
-import uuid
-import pickle
-import multiprocessing
-import asyncio
-import queue
-from urllib.parse import urlparse
+# Copyright (c) Microsoft. All rights reserved.
 
-from .base import BaseTracer
+import asyncio
+import logging
+import multiprocessing
+import queue
+import uuid
+from contextlib import contextmanager
+from typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
+from urllib.parse import urlparse
 
 from httpdbg.hooks.all import httprecord
 from httpdbg.records import HTTPRecords
 from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import StatusCode, SpanKind, Status
+from opentelemetry.trace import SpanKind, Status, StatusCode
 from opentelemetry.trace.span import (
     SpanContext,
     TraceFlags,
     TraceState,
 )
 
+from .base import BaseTracer
 
 logger = logging.getLogger(__name__)
 
@@ -37,6 +37,9 @@ class HttpTracer(BaseTracer):
     and we do not recommend using it in production.
     It is primarily for demonstration and testing purposes.
 
+    Deprecated: This tracer is deprecated and will be removed in a future version.
+    Please use LLMProxy as an alternative.
+
     Attributes:
         include_headers: Whether to include HTTP headers in the spans.
             Headers may contain sensitive information. Use with caution.
@@ -58,14 +61,14 @@ class HttpTracer(BaseTracer):
         subprocess_timeout: float = 3600.0,
     ):
         super().__init__()
-        self._last_records = None
+        self._last_records: Optional[HTTPRecords] = None
         self.include_headers = include_headers
         self.include_body = include_body
         self.include_agentlightning_requests = include_agentlightning_requests
         self.subprocess_mode = subprocess_mode
         self.subprocess_timeout = subprocess_timeout
 
-    def init_worker(self, worker_id: int):
+    def init_worker(self, worker_id: int) -> None:
         """
         Initialize the tracer in a worker process.
 
@@ -76,7 +79,7 @@ class HttpTracer(BaseTracer):
         logger.info(f"[Worker {worker_id}] HttpTracer initialized.")
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None) -> Iterator[HTTPRecords]:
+    def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]:
         """
         Starts a new HTTP tracing context. This should be used as a context manager.
 
@@ -113,7 +116,7 @@ class HttpTracer(BaseTracer):
         Returns:
             A list of ReadableSpan objects representing the HTTP activities.
         """
-        spans = []
+        spans: List[ReadableSpan] = []
 
         # Create a trace ID that will be shared by all spans in this trace
         trace_id = int(uuid.uuid4().hex[:16], 16)
@@ -156,7 +159,7 @@ class HttpTracer(BaseTracer):
                 "http.host": parsed_url.netloc,
             }
 
-            if status_code is not None and status_code > 0:
+            if status_code is not None and status_code > 0:  # type: ignore
                 attributes["http.status_code"] = status_code
 
             # Calculate duration - from begin time to last update
@@ -220,7 +223,7 @@ class HttpTracer(BaseTracer):
 
         return spans
 
-    def trace_run(self, func: Callable, *args, **kwargs) -> Any:
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single synchronous function.
 
@@ -240,7 +243,7 @@ class HttpTracer(BaseTracer):
         else:
             return super().trace_run(func, *args, **kwargs)
 
-    async def trace_run_async(self, func: Callable[..., Awaitable], *args, **kwargs) -> Any:
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single asynchronous function.
 
@@ -263,7 +266,13 @@ class HttpTracer(BaseTracer):
         else:
             return await super().trace_run_async(func, *args, **kwargs)
 
-    def _trace_run_subprocess(self, func: Callable, args=None, kwargs=None, is_async: bool = False) -> Any:
+    def _trace_run_subprocess(
+        self,
+        func: Callable[..., Any],
+        args: Optional[Tuple[Any, ...]] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+        is_async: bool = False,
+    ) -> Any:
         """
         Execute a function in a subprocess with HTTP tracing.
 
@@ -282,23 +291,23 @@ class HttpTracer(BaseTracer):
             kwargs = {}
 
         # Create a queue to receive results from the subprocess
-        result_queue = multiprocessing.Queue()
+        result_queue = multiprocessing.Queue()  # type: ignore
 
         # Create and start the subprocess
         process = multiprocessing.Process(
-            target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async)
+            target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async)  # type: ignore
         )
         process.start()
 
         try:
             # Wait for the process to complete and get the result
             process.join(timeout=self.subprocess_timeout)
-            result = result_queue.get_nowait()
+            result = result_queue.get_nowait()  # type: ignore
 
             if result["success"]:
                 # Store the captured records for get_last_trace()
                 self._last_records = result["records"]
-                return result["return_value"]
+                return result["return_value"]  # type: ignore
             else:
                 if "records" in result:
                     self._last_records = result["records"]
@@ -316,7 +325,14 @@ class HttpTracer(BaseTracer):
                 process.terminate()
                 process.join()
 
-    def _subprocess_worker(self, func: Callable, args, kwargs, result_queue: multiprocessing.Queue, is_async: bool):
+    def _subprocess_worker(
+        self,
+        func: Callable[..., Any],
+        args: Tuple[Any, ...],
+        kwargs: Dict[str, Any],
+        result_queue: multiprocessing.Queue,  # type: ignore
+        is_async: bool,
+    ) -> None:
         """
         Worker function that runs in the subprocess to execute the traced function.
 
@@ -354,7 +370,7 @@ class HttpTracer(BaseTracer):
             records = subprocess_tracer._last_records
 
             # Send success result back to parent
-            result_queue.put({"success": True, "return_value": return_value, "records": records})
+            result_queue.put({"success": True, "return_value": return_value, "records": records})  # type: ignore
 
         except Exception as e:
             # Log the exception
@@ -363,4 +379,4 @@ class HttpTracer(BaseTracer):
             # Get the captured records even when there's an exception
             records = subprocess_tracer._last_records
             # Send error result back to parent
-            result_queue.put({"success": False, "exception": e, "records": records})
+            result_queue.put({"success": False, "exception": e, "records": records})  # type: ignore

agentlightning/tracer/otel.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import Iterator, List, Optional
+
+import opentelemetry.trace as trace_api
+from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+
+from agentlightning.store.base import LightningStore
+
+from .agentops import LightningSpanProcessor  # FIXME: This import should be from otel to agentops
+from .base import BaseTracer
+
+logger = logging.getLogger(__name__)
+
+
+class OtelTracer(BaseTracer):
+    """Tracer that provides a basic OpenTelemetry tracer provider.
+
+    You should be able to collect agent-lightning signals like rewards with this tracer,
+    but no other function instrumentations like `openai.chat.completion`.
+    """
+
+    def __init__(self):
+        super().__init__()
+        # This provider is only initialized when the worker is initialized.
+        self._tracer_provider: Optional[TracerProvider] = None
+        self._lightning_span_processor: Optional[LightningSpanProcessor] = None
+        self._initialized: bool = False
+
+    def init_worker(self, worker_id: int):
+        super().init_worker(worker_id)
+        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._lightning_span_processor = LightningSpanProcessor()
+        tracer_provider.add_span_processor(self._lightning_span_processor)
+        self._initialized = True
+
+    def teardown_worker(self, worker_id: int):
+        super().teardown_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
+        self._tracer_provider = None
+
+    @contextmanager
+    def trace_context(
+        self,
+        name: Optional[str] = None,
+        *,
+        store: Optional[LightningStore] = None,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+    ) -> Iterator[LightningSpanProcessor]:
+        """
+        Starts a new tracing context. This should be used as a context manager.
+
+        Args:
+            name: Optional name for the tracing context.
+            store: Optional store to add the spans to.
+            rollout_id: Optional rollout ID to add the spans to.
+            attempt_id: Optional attempt ID to add the spans to.
+
+        Yields:
+            The LightningSpanProcessor 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:
+            ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
+            with ctx as processor:
+                yield processor
+        elif store is None and rollout_id is None and attempt_id is None:
+            with self._lightning_span_processor:
+                yield self._lightning_span_processor
+        else:
+            raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
+
+    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()

agentlightning/trainer.py

@@ -1,311 +0,0 @@
-import asyncio
-import logging
-import multiprocessing
-import os
-import signal
-import time
-from typing import List, Optional, Union
-import importlib
-
-import agentops
-
-from .client import AgentLightningClient
-from .litagent import LitAgent
-from .runner import AgentRunner
-from .types import ParallelWorkerBase
-from .tracer.base import BaseTracer
-from .tracer.agentops import AgentOpsTracer
-from .tracer.triplet import TripletExporter
-
-
-logger = logging.getLogger(__name__)
-
-
-class Trainer(ParallelWorkerBase):
-    """Orchestrates the distributed execution of agent rollouts.
-
-    The Trainer is responsible for launching one or more worker processes
-    that run the agent's execution loop. It manages multiprocessing,
-    handles graceful shutdown, and serves as the main entry point for
-    running a client-side agent fleet.
-
-    Attributes:
-        dev: If True, rollouts are run against the dev endpoint provided in `fit`.
-        n_workers: Number of agent workers (processes) to run in parallel.
-        max_tasks: Maximum number of tasks to process per worker. If None,
-                   workers run until no more tasks are available.
-        daemon: Whether worker processes should be daemons. Daemon processes
-                are terminated automatically when the main process exits.
-        tracer: A tracer instance, or a string pointing to the class full name or a dictionary with a 'type' key
-                that specifies the class full name and other initialization parameters.
-                If None, a default `AgentOpsTracer` will be created with the current settings.
-        triplet_exporter: An instance of `TripletExporter` to export triplets from traces,
-                          or a dictionary with the initialization parameters for the exporter.
-    """
-
-    def __init__(
-        self,
-        *,
-        dev: bool = False,
-        n_workers: int = 1,
-        max_tasks: Optional[int] = None,
-        daemon: bool = True,
-        tracer: Union[BaseTracer, str, dict, None] = None,
-        triplet_exporter: Union[TripletExporter, dict, None] = None,
-    ):
-        super().__init__()
-        self.n_workers = n_workers
-        self.max_tasks = max_tasks
-        self.daemon = daemon
-        self.dev = dev
-        self._client: AgentLightningClient | None = None  # Will be initialized in fit method
-
-        self.tracer = self._make_tracer(tracer)
-        if isinstance(triplet_exporter, TripletExporter):
-            self.triplet_exporter = triplet_exporter
-        elif isinstance(triplet_exporter, dict):
-            self.triplet_exporter = TripletExporter(**triplet_exporter)
-        elif triplet_exporter is None:
-            self.triplet_exporter = TripletExporter()
-        else:
-            raise ValueError(
-                f"Invalid triplet_exporter type: {type(triplet_exporter)}. Expected TripletExporter, dict, or None."
-            )
-
-        if not self.daemon:
-            logger.warning(
-                "daemon=False. Worker processes are non-daemonic. "
-                "The worker processes will NOT be terminated when the main process exits. "
-                "The cleanup must be handled manually."
-            )
-
-    def _make_tracer(self, tracer: Union[BaseTracer, str, dict, None]) -> BaseTracer:
-        """Creates a tracer instance based on the provided configuration."""
-        if isinstance(tracer, BaseTracer):
-            return tracer
-        if isinstance(tracer, str):
-            module_name, class_name = tracer.rsplit(".", 1)
-            module = importlib.import_module(module_name)
-            tracer_cls = getattr(module, class_name)
-            return tracer_cls()
-        if isinstance(tracer, dict):
-            tracer_type = tracer.get("type")
-            if tracer_type is None:
-                raise ValueError("tracer dict must have a 'type' key with the class full name")
-            module_name, class_name = tracer_type.rsplit(".", 1)
-            module = importlib.import_module(module_name)
-            tracer_cls = getattr(module, class_name)
-            # Remove 'type' key and pass remaining keys as kwargs
-            tracer_kwargs = {k: v for k, v in tracer.items() if k != "type"}
-            return tracer_cls(**tracer_kwargs)
-        if tracer is None:
-            return AgentOpsTracer(agentops_managed=True, instrument_managed=True, daemon=self.daemon)
-        raise ValueError(f"Invalid tracer type: {type(tracer)}. Expected BaseTracer, str, dict, or None.")
-
-    def init(self, backend: Union[str, AgentLightningClient]) -> None:
-        logger.info(f"Initializing Trainer...")
-
-        self._init_client(backend)
-
-        self.tracer.init()
-
-        logger.info(f"Trainer main initialization complete.")
-
-    def teardown(self) -> None:
-        logger.info(f"Cleaning up Trainer...")
-        self.tracer.teardown()
-
-        self._client = None
-        logger.info(f"Trainer main cleanup complete.")
-
-    def client(self) -> AgentLightningClient:
-        """Returns the AgentLightningClient instance."""
-        if self._client is None:
-            raise RuntimeError("AgentLightningClient has not been initialized. Call `init` first.")
-        return self._client
-
-    def _init_client(self, backend: Union[str, AgentLightningClient]) -> AgentLightningClient:
-        if self._client is None:
-            if isinstance(backend, AgentLightningClient):
-                logger.info("Using provided AgentLightningClient instance.")
-                self._client = backend
-            else:
-                logger.info(f"Initializing AgentLightningClient with endpoint: {backend}")
-                if not isinstance(backend, str):
-                    raise ValueError("backend must be a string URL or an AgentLightningClient instance.")
-                if not backend.startswith("http://") and not backend.startswith("https://"):
-                    raise ValueError("backend must be a valid URL starting with http:// or https://")
-                # Initialize the client with the provided backend URL
-                self._client = AgentLightningClient(endpoint=backend)
-        else:
-            logger.warning("AgentLightningClient already initialized. Returning existing instance.")
-        return self._client
-
-    def _worker_main_loop(self, agent: LitAgent, worker_id: int, is_async: bool):
-        """The main function for each worker process.
-
-        This function initializes the client and the loop, then starts the
-        execution. It also configures process-specific settings like the
-        process title and signal handling.
-
-        Args:
-            agent: The `LitAgent` instance to run.
-            worker_id: The unique ID for this worker.
-            is_async: A boolean indicating if the async loop should be run.
-        """
-        if self.n_workers > 1:
-            import setproctitle
-
-            # Ignore Ctrl+C in worker processes; the main process handles it
-            signal.signal(signal.SIGINT, signal.SIG_IGN)
-            setproctitle.setproctitle(multiprocessing.current_process().name)
-
-        # Now we are in child processes, so we can safely set up the environment.
-        agent.set_trainer(self)
-        # TODO: this should be set elsewhere
-        if agent.trained_agents:
-            self.triplet_exporter.agent_match = agent.trained_agents
-        self._initialize_worker_env(worker_id)
-
-        mode = "Async" if is_async else "Sync"
-        logger.info(f"[Worker {worker_id}] {mode} worker process started.")
-
-        num_processed = 0
-
-        try:
-            client = self.client()
-            loop = AgentRunner(
-                agent=agent,
-                client=client,
-                tracer=self.tracer,
-                triplet_exporter=self.triplet_exporter,
-                max_tasks=self.max_tasks,
-                worker_id=worker_id,
-            )
-            loop.init_worker(worker_id)
-            if is_async:
-                num_processed = asyncio.run(loop.iter_async())
-            else:
-                num_processed = loop.iter()
-        except Exception:
-            logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
-        finally:
-            self._teardown_worker_env(worker_id)
-
-        return num_processed
-
-    def _initialize_worker_env(self, worker_id: int):
-        logger.info(f"[Worker {worker_id}] Setting up trainer environment...")  # worker_id included in process name
-        self.tracer.init_worker(worker_id)
-
-    def _teardown_worker_env(self, worker_id: int):
-        logger.info(f"[Worker {worker_id}] Cleaning up trainer environment...")
-        self.tracer.teardown_worker(worker_id)
-        logger.info(f"[Worker {worker_id}] Environment cleanup complete.")
-
-    @staticmethod
-    def kill_orphaned_processes() -> None:
-        """
-        Kill any orphaned processes that may have been left behind by previous runs.
-        This is useful for cleaning up after crashes or unexpected exits.
-        """
-        import psutil
-
-        for proc in psutil.process_iter():
-            # check whether the process name matches
-            if proc.name().startswith("AgentLightning-"):
-                proc.kill()
-
-    def fit(
-        self,
-        agent: LitAgent,
-        backend: Union[str, AgentLightningClient],
-        dev_backend: Union[str, AgentLightningClient, None] = None,
-    ):
-        if self.dev:
-            if dev_backend is None:
-                raise ValueError("dev_backend must be provided when dev=True.")
-            logger.warning(f"Running in dev mode. Using dev backend: {dev_backend}")
-            self.init(dev_backend)
-        else:
-            logger.debug(f"Running in non-dev mode. Using backend: {backend}")
-            self.init(backend)
-
-        processes: List[multiprocessing.Process] = []
-
-        # Determine if the agent is asynchronous.
-        is_async = (
-            hasattr(agent, "training_rollout_async")
-            and agent.__class__.training_rollout_async is not LitAgent.training_rollout_async
-        )
-
-        mode = "asynchronous" if is_async else "synchronous"
-
-        try:
-            if self.n_workers == 1:
-                logger.info(f"Running with n_workers=1 ({mode} in main process).")
-                num_tasks = self._worker_main_loop(agent, 0, is_async)
-                logger.info(f"Single worker mode finished. Tasks processed: {num_tasks}")
-            else:
-                logger.info(f"Running with n_workers={self.n_workers} ({mode} multiprocessing).")
-                for i in range(self.n_workers):
-                    process_name = f"AgentLightning-Worker-{i}"
-                    p = multiprocessing.Process(
-                        target=self._worker_main_loop,
-                        args=(agent, i, is_async),
-                        daemon=self.daemon,
-                        name=process_name,
-                    )
-                    processes.append(p)
-                    logger.info(f"Starting worker process {i} (name: {process_name})...")
-                    p.start()
-
-                if self.daemon:
-                    for i, p in enumerate(processes):
-                        p.join()  # Wait for the process to complete
-                        logger.info(
-                            f"Worker process {i} (name: {p.name}, PID: {p.pid}) joined with exit code {p.exitcode}."
-                        )
-                        if p.exitcode != 0:
-                            logger.warning(
-                                f"Worker process {i} (name: {p.name}, PID: {p.pid}) exited with non-zero code: {p.exitcode}."
-                            )
-
-                    logger.info(f"All {self.n_workers} worker processes have completed.")
-                else:
-                    logger.info("All worker processes started. Main process will not wait.")
-
-                    # A hack to stop the main process from waiting for child processes to finish.
-                    time.sleep(1)  # Give workers time to start
-                    import multiprocessing.process as multiprocessing_process
-
-                    multiprocessing_process._children.clear()  # type: ignore
-
-        except KeyboardInterrupt:
-            if self.n_workers > 1 and len(processes) > 0:
-                logger.info(f"KeyboardInterrupt received. Terminating workers...")
-                for i, p in enumerate(processes):
-                    if p.is_alive():
-                        logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
-                        p.terminate()
-                    else:
-                        logger.info(
-                            f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated."
-                        )
-                for i, p in enumerate(processes):
-                    if p.is_alive():
-                        p.join(timeout=10)  # Give some time to terminate
-                    if p.is_alive():  # If still alive, kill
-                        logger.warning(
-                            f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
-                        )
-                        p.kill()
-                        p.join(timeout=10)  # Ensure it's reaped
-            logger.info(f"Workers terminated or single worker interrupted.")
-        except Exception as e:
-            logger.exception(f"Unhandled exception in fit method.")
-        finally:
-            if self.daemon:
-                self.teardown()
-            else:
-                logger.info("Main process exiting. Please use Trainer.kill_orphaned_processes() for cleanup.")

agentlightning/trainer/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .init_utils import build_component
+from .trainer import Trainer
+
+__all__ = ["Trainer", "build_component"]

agentlightning/trainer/init_utils.py

@@ -0,0 +1,263 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utility helpers for dynamic component initialization within the trainer."""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+from typing import Any, Callable, Dict, Optional, TypeVar, Union, cast, overload
+
+OptionalDefaults = Dict[str, Callable[[], Any] | Any]
+T = TypeVar("T")
+
+
+def load_class(path: str) -> type[Any]:
+    """Load a class from its fully qualified import path."""
+    module_name, class_name = path.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)
+
+
+def instantiate_component(
+    cls: type[Any],
+    provided_kwargs: Optional[Dict[str, Any]] = None,
+    optional_defaults: Optional[OptionalDefaults] = None,
+) -> Any:
+    """Instantiate `cls`, filling optional kwargs when the constructor accepts them."""
+    kwargs = dict(provided_kwargs or {})
+    if optional_defaults:
+        signature = inspect.signature(cls.__init__)
+        for name, value in optional_defaults.items():
+            if name in kwargs or name not in signature.parameters:
+                continue
+            kwargs[name] = value() if callable(value) else value
+    return cls(**kwargs)
+
+
+def instantiate_from_spec(
+    spec: Union[str, Dict[str, Any]],
+    *,
+    spec_name: str,
+    optional_defaults: Optional[OptionalDefaults] = None,
+    dict_requires_type: bool = True,
+    dict_default_cls: type[Any] | None = None,
+    registry: Optional[Dict[str, str]] = None,
+) -> Any:
+    """Instantiate a component from a string or dict spec."""
+    if isinstance(spec, str):
+        type_path = registry.get(spec, spec) if registry else spec
+        cls = load_class(type_path)
+        return instantiate_component(cls, optional_defaults=optional_defaults)
+
+    if isinstance(spec, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+        spec_conf = dict(spec)
+        type_path = spec_conf.pop("type", None)
+        if type_path is None and registry and "name" in spec_conf:
+            type_path = registry.get(spec_conf.pop("name"))
+        elif registry and type_path is not None:
+            type_path = registry.get(type_path, type_path)
+        if type_path is None:
+            if dict_requires_type:
+                raise ValueError(f"{spec_name} dict must have a 'type' key with the class full name")
+            if dict_default_cls is None:
+                raise ValueError(f"{spec_name} dict missing 'type' and no default class provided")
+            cls = dict_default_cls
+        else:
+            cls = load_class(type_path)
+        return instantiate_component(cls, spec_conf, optional_defaults)
+
+    raise TypeError(f"{spec_name} spec must be a string or dict (got {type(spec)}).")
+
+
+def _ensure_expected_type(
+    instance: Any,
+    expected_type: type[T],
+    spec_name: str,
+    type_error_fmt: str | None,
+) -> T:
+    if not isinstance(instance, expected_type):
+        type_name = str(type(instance))  # type: ignore
+        if type_error_fmt:
+            raise TypeError(type_error_fmt.format(type_name=type_name, expected_type=expected_type.__name__))
+        raise TypeError(f"{spec_name} factory returned {type_name}, which is not a {expected_type.__name__} subclass.")
+    return instance
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: Callable[[], T],
+    allow_none: bool = ...,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T: ...
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: None = ...,
+    allow_none: bool,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T | None: ...
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: None = ...,
+    allow_none: bool = ...,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T | None: ...
+
+
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: Callable[[], T] | None = None,
+    allow_none: bool = False,
+    optional_defaults: Optional[OptionalDefaults] = None,
+    dict_requires_type: bool = True,
+    dict_default_cls: type[T] | None = None,
+    type_error_fmt: str | None = None,
+    invalid_spec_error_fmt: str | None = None,
+    registry: Optional[Dict[str, str]] = None,
+) -> T | None:
+    """Build and return a component instance from a flexible specification.
+
+    This function provides a flexible way to create component instances from various
+    input formats including direct instances, class types, factory functions, import
+    paths, or configuration dictionaries.
+
+    Args:
+        spec: The component specification. Can be:
+            - An instance of expected_type (returned as-is)
+            - A string import path (e.g., 'module.Class') or registry key
+            - A dict with 'type' key (import path or registry key) and constructor kwargs
+            - A class type (will be instantiated)
+            - A factory function (will be called)
+            - None (uses default_factory or returns None if allow_none=True)
+        expected_type: The type that the resulting instance must be or inherit from.
+        spec_name: Descriptive name for the spec, used in error messages.
+        default_factory: Optional factory function called when spec is None.
+        allow_none: If True, allows None to be returned when spec is None and
+            no default_factory is provided.
+        optional_defaults: Dict mapping parameter names to default values or factory
+            functions that will be injected if the constructor accepts them.
+        dict_requires_type: If True, dict specs must include a 'type' key.
+        dict_default_cls: Default class to use for dict specs without a 'type' key
+            (only used when dict_requires_type=False).
+        type_error_fmt: Custom format string for type validation errors. Should include
+            {type_name} and {expected_type} placeholders.
+        invalid_spec_error_fmt: Custom format string for invalid spec type errors.
+            Should include {actual_type} and {expected_type} placeholders.
+        registry: Optional mapping of short names to fully qualified import paths.
+            When provided, string specs or dict 'type'/'name' entries are first
+            resolved through this registry before attempting to import.
+
+    Returns:
+        An instance of expected_type, or None if allow_none=True and spec is None
+        without a default_factory.
+
+    Raises:
+        TypeError: If the instantiated object is not an instance of expected_type.
+        ValueError: If spec is None and neither default_factory nor allow_none is set,
+            or if spec type is invalid, or if dict spec is invalid.
+
+    Examples:
+        >>> # Direct instance
+        >>> optimizer = build_component(AdamW(), expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # String import path
+        >>> optimizer = build_component('torch.optim.AdamW', expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Dict with type and kwargs
+        >>> spec = {'type': 'torch.optim.AdamW', 'lr': 0.001}
+        >>> optimizer = build_component(spec, expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Class type
+        >>> optimizer = build_component(AdamW, expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Factory function
+        >>> optimizer = build_component(lambda: AdamW(lr=0.001), expected_type=Optimizer,
+        ...                            spec_name='optimizer')
+    """
+    if isinstance(spec, expected_type):
+        return cast(T, spec)
+
+    if spec is None:
+        if default_factory is not None:
+            instance = default_factory()
+            return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+        if allow_none:
+            return None
+        raise ValueError(
+            invalid_spec_error_fmt.format(actual_type=type(spec), expected_type=expected_type.__name__)
+            if invalid_spec_error_fmt
+            else f"{spec_name} cannot be None."
+        )
+
+    if isinstance(spec, type) and issubclass(spec, expected_type):
+        instance = instantiate_component(spec, optional_defaults=optional_defaults)
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if callable(spec) and not isinstance(spec, type):  # type: ignore
+        instance = spec()
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if isinstance(spec, str):
+        instance = instantiate_from_spec(
+            spec,
+            spec_name=spec_name,
+            optional_defaults=optional_defaults,
+            dict_requires_type=dict_requires_type,
+            dict_default_cls=dict_default_cls,
+            registry=registry,
+        )
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if isinstance(spec, dict):
+        instance = instantiate_from_spec(
+            spec,  # type: ignore
+            spec_name=spec_name,
+            optional_defaults=optional_defaults,
+            dict_requires_type=dict_requires_type,
+            dict_default_cls=dict_default_cls,
+            registry=registry,
+        )
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if invalid_spec_error_fmt:
+        raise ValueError(invalid_spec_error_fmt.format(actual_type=type(spec), expected_type=expected_type.__name__))  # type: ignore
+
+    type_name = str(type(spec))  # type: ignore
+    raise ValueError(f"Invalid {spec_name} type: {type_name}. Expected {expected_type.__name__}, str, dict, or None.")
+
+
+__all__ = ["OptionalDefaults", "build_component", "instantiate_component", "instantiate_from_spec", "load_class"]

agentlightning/trainer/legacy.py

@@ -0,0 +1,367 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import multiprocessing
+import signal
+import time
+import warnings
+from typing import Any, List, Optional, TypeVar, Union
+
+from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
+from agentlightning.algorithm import BaseAlgorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.litagent import LitAgent
+from agentlightning.runner import LegacyAgentRunner
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import Dataset, ParallelWorkerBase
+
+logger = logging.getLogger(__name__)
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class TrainerLegacy(ParallelWorkerBase):
+    """Trainer for legacy mode for v0.1 compatibility."""
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        """Initialize the TrainerLegacy.
+
+        This method is mainly to make type checker happy.
+        It won't be used in practice.
+        """
+        self._dev = kwargs.pop("dev", False)
+        self.algorithm: Optional[BaseAlgorithm] = kwargs.pop("algorithm", None)
+        self.tracer: BaseTracer = kwargs.pop("tracer", None)
+        self.n_workers: int = kwargs.pop("n_workers", None)
+        self.max_tasks: Optional[int] = kwargs.pop("max_tasks", None)
+        self.daemon: bool = kwargs.pop("daemon", True)
+        self.triplet_exporter: TraceAdapter[Any] = kwargs.pop("triplet_exporter", None)
+
+    def _extract_client_from_data(
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
+    ) -> Optional[AgentLightningClient]:
+        """Extract client from data if it's a string URL or AgentLightningClient."""
+        if isinstance(data, str):
+            if not data.startswith("http://") and not data.startswith("https://"):
+                raise ValueError("String data must be a valid URL starting with http:// or https://")
+            return AgentLightningClient(endpoint=data)
+        elif isinstance(data, AgentLightningClient):
+            return data
+        return None
+
+    def _extract_dataset_from_data(
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
+    ) -> Optional[Dataset[Any]]:
+        """Extract dataset from data if it's a Dataset."""
+        if isinstance(data, str) or isinstance(data, AgentLightningClient):
+            return None
+        return data
+
+    def _determine_backend(
+        self,
+        train_data: Union[str, AgentLightningClient, Dataset[Any]],
+        dev_data: Union[str, AgentLightningClient, Dataset[Any], None] = None,
+    ) -> Union[str, AgentLightningClient]:
+        """Determine which backend to use for initialization."""
+        if self._dev:
+            if dev_data is None:
+                raise ValueError("dev_data must be provided when dev=True.")
+            client = self._extract_client_from_data(dev_data)
+            if client is None:
+                raise ValueError("dev_data must be a string URL or AgentLightningClient when dev=True.")
+            return client
+        else:
+            client = self._extract_client_from_data(train_data)
+            if client is None and self.algorithm is None:
+                raise ValueError(
+                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
+                )
+            elif client is None and self.algorithm is not None:
+                # Algorithm will be responsible for creating the client
+                client = self.algorithm.get_client()
+                logger.info(f"Algorithm created client: {client}")
+                return client
+            if client is None:
+                raise ValueError(
+                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
+                )
+            return client
+
+    def init(self, backend: Union[str, AgentLightningClient]) -> None:
+        logger.info(f"Initializing Trainer...")
+
+        self._init_client(backend)
+
+        self.tracer.init()
+
+        logger.info(f"Trainer main initialization complete.")
+
+    def teardown(self) -> None:
+        logger.info(f"Cleaning up Trainer...")
+        self.tracer.teardown()
+
+        self._client = None
+        logger.info(f"Trainer main cleanup complete.")
+
+    def client(self) -> AgentLightningClient:
+        """Returns the AgentLightningClient instance."""
+        if self._client is None:
+            raise RuntimeError("AgentLightningClient has not been initialized. Call `init` first.")
+        return self._client
+
+    def _init_client(self, backend: Union[str, AgentLightningClient]) -> AgentLightningClient:
+        if self._client is None:
+            if isinstance(backend, AgentLightningClient):
+                logger.info("Using provided AgentLightningClient instance.")
+                self._client = backend
+            else:
+                logger.info(f"Initializing AgentLightningClient with endpoint: {backend}")
+                if not isinstance(backend, str):  # type: ignore
+                    raise ValueError("backend must be a string URL or an AgentLightningClient instance.")
+                if not backend.startswith("http://") and not backend.startswith("https://"):
+                    raise ValueError("backend must be a valid URL starting with http:// or https://")
+                # Initialize the client with the provided backend URL
+                self._client = AgentLightningClient(endpoint=backend)
+        else:
+            logger.warning("AgentLightningClient already initialized. Returning existing instance.")
+        return self._client
+
+    def _worker_main_loop(self, agent: LitAgent[Any], worker_id: int, is_async: bool):
+        """The main function for each worker process.
+
+        This function initializes the client and the loop, then starts the
+        execution. It also configures process-specific settings like the
+        process title and signal handling.
+
+        Args:
+            agent: The `LitAgent` instance to run.
+            worker_id: The unique ID for this worker.
+            is_async: A boolean indicating if the async loop should be run.
+        """
+        if self.n_workers > 1:
+            import setproctitle
+
+            # Ignore Ctrl+C in worker processes; the main process handles it
+            signal.signal(signal.SIGINT, signal.SIG_IGN)
+            setproctitle.setproctitle(multiprocessing.current_process().name)
+
+        # Now we are in child processes, so we can safely set up the environment.
+        agent.set_trainer(self)  # type: ignore
+        if not isinstance(self.triplet_exporter, TracerTraceToTriplet):  # type: ignore
+            raise ValueError("triplet_exporter must be a TracerTraceToTriplet for the legacy trainer.")
+        # TODO: this should be set elsewhere
+        if agent.trained_agents:
+            self.triplet_exporter.agent_match = agent.trained_agents
+        self._initialize_worker_env(worker_id)
+
+        mode = "Async" if is_async else "Sync"
+        logger.info(f"[Worker {worker_id}] {mode} worker process started.")
+
+        num_processed = 0
+
+        try:
+            client = self.client()
+            loop = LegacyAgentRunner(
+                agent=agent,
+                client=client,
+                tracer=self.tracer,
+                triplet_exporter=self.triplet_exporter,
+                max_tasks=self.max_tasks,
+                worker_id=worker_id,
+            )
+            loop.init_worker(worker_id)  # type: ignore
+            if is_async:
+                num_processed = asyncio.run(loop.iter_async())
+            else:
+                num_processed = loop.iter()
+        except Exception:
+            logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
+        finally:
+            self._teardown_worker_env(worker_id)
+
+        return num_processed
+
+    def _initialize_worker_env(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Setting up trainer environment...")  # worker_id included in process name
+        self.tracer.init_worker(worker_id)
+
+    def _teardown_worker_env(self, worker_id: int):
+        logger.info(f"[Worker {worker_id}] Cleaning up trainer environment...")
+        self.tracer.teardown_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Environment cleanup complete.")
+
+    @staticmethod
+    def kill_orphaned_processes() -> None:
+        """
+        Kill any orphaned processes that may have been left behind by previous runs.
+        This is useful for cleaning up after crashes or unexpected exits.
+        """
+        import psutil
+
+        for proc in psutil.process_iter():  # type: ignore
+            # check whether the process name matches
+            if proc.name().startswith("AgentLightning-"):
+                proc.kill()
+
+    def _terminate_processes(self, processes: List[multiprocessing.Process]) -> None:
+        if self.n_workers > 1 and len(processes) > 0:
+            for i, p in enumerate(processes):
+                if p.is_alive():
+                    logger.info(f"Terminating worker {i} (name: {p.name}, PID: {p.pid})...")
+                    p.terminate()
+                else:
+                    logger.info(f"Worker {i} (name: {p.name}, PID: {p.pid}) is not alive or has already terminated.")
+            for i, p in enumerate(processes):
+                if p.is_alive():
+                    p.join(timeout=10)  # Give some time to terminate
+                if p.is_alive():  # If still alive, kill
+                    logger.warning(
+                        f"Worker {i} (name: {p.name}, PID: {p.pid}) did not terminate gracefully, killing..."
+                    )
+                    p.kill()
+                    p.join(timeout=10)  # Ensure it's reaped
+
+    def fit_v0(
+        self,
+        agent: LitAgent[T_co],
+        train_data: Union[str, AgentLightningClient, Dataset[T_co]],
+        *,
+        val_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
+        dev_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
+        dev_backend: Union[str, AgentLightningClient, None] = None,
+    ):
+        """Train the agent using the provided data.
+
+        Each data argument can be a string URL connecting to a agent-lightning server,
+        or an AgentLightningClient instance connecting to a server (or mock server), or a dataset.
+        If no algorithm is provided when instantiating the trainer, the data must be
+        provided to connecting a server. Otherwise, dataset is also allowed and will be
+        passed to the algorithm.
+
+        If the algorithm is instantiated and there is no URL/client provided,
+        the algorithm will be responsible for creating a client that will connect to itself.
+        It can also create a mock client if the algorithm does not require a server.
+        """
+
+        if dev_backend is not None:
+            warnings.warn("dev_backend is deprecated. Use dev_data instead.")
+            if dev_data is not None:
+                raise ValueError("dev_data and dev_backend cannot be provided at the same time.")
+            dev_data = dev_backend
+
+        # Extract datasets for algorithm if available
+        train_dataset = self._extract_dataset_from_data(train_data)
+        val_dataset = self._extract_dataset_from_data(val_data) if val_data else None
+
+        # Initialize the algorithm with trainer if provided
+        if self.algorithm is not None:
+            self.algorithm.set_trainer(self)  # type: ignore
+            # DO NOT RUN TRAINING HERE. Need to spawn the worker first.
+
+        # Determine the backend to use for client-server mode
+        backend = self._determine_backend(train_data, dev_data)
+
+        if self._dev:
+            logger.warning(f"Running in dev mode. Using dev backend: {backend}")
+        else:
+            logger.debug(f"Running in non-dev mode. Using backend: {backend}")
+
+        self.init(backend)
+
+        processes: List[multiprocessing.Process] = []
+
+        # Determine if the agent is asynchronous
+
+        mode = "asynchronous" if agent.is_async() else "synchronous"
+
+        try:
+            if self.n_workers == 1:
+                logger.info(f"Running with n_workers=1 ({mode} in main process).")
+
+                # Warn if algorithm is set with single worker mode
+                if self.algorithm is not None:
+                    logger.warning(
+                        "Algorithm is set but using single worker mode. Algorithm will never get the chance to run."
+                    )
+                    # Ideally the single worker should be run in a separate thread or process.
+
+                num_tasks = self._worker_main_loop(agent, 0, agent.is_async())
+                logger.info(f"Single worker mode finished. Tasks processed: {num_tasks}")
+
+                # If algorithm is provided and we have datasets, run algorithm after worker completes
+                if self.algorithm is not None and train_dataset is not None:
+                    logger.info("Running algorithm training after worker completion.")
+                    self.algorithm.run(
+                        train_dataset=train_dataset,
+                        val_dataset=val_dataset,
+                    )
+            else:
+                logger.info(f"Running with n_workers={self.n_workers} ({mode} multiprocessing).")
+                for i in range(self.n_workers):
+                    process_name = f"AgentLightning-Worker-{i}"
+                    p = multiprocessing.Process(
+                        target=self._worker_main_loop,
+                        args=(agent, i, agent.is_async()),
+                        daemon=self.daemon,
+                        name=process_name,
+                    )
+                    processes.append(p)
+                    logger.info(f"Starting worker process {i} (name: {process_name})...")
+                    p.start()
+
+                if self.daemon:
+                    # If algorithm is provided and we have datasets, pass them to the algorithm
+                    if self.algorithm is not None:
+                        logger.info("All workers have been spawned. Running algorithm training with provided datasets.")
+                        self.algorithm.run(
+                            train_dataset=train_dataset,
+                            val_dataset=val_dataset,
+                        )
+                        logger.info("Algorithm exits. Killing the workers.")
+                        self._terminate_processes(processes)
+
+                    for i, p in enumerate(processes):
+                        p.join()  # Wait for the process to complete
+                        logger.info(
+                            f"Worker process {i} (name: {p.name}, PID: {p.pid}) joined with exit code {p.exitcode}."
+                        )
+                        if p.exitcode != 0:
+                            logger.warning(
+                                f"Worker process {i} (name: {p.name}, PID: {p.pid}) exited with non-zero code: {p.exitcode}."
+                            )
+
+                    logger.info(f"All {self.n_workers} worker processes have completed.")
+                else:
+                    logger.info("All worker processes started. Main process will not wait.")
+
+                    # A hack to stop the main process from waiting for child processes to finish.
+                    time.sleep(1)  # Give workers time to start
+                    import multiprocessing.process as multiprocessing_process
+
+                    multiprocessing_process._children.clear()  # type: ignore
+
+                    if self.algorithm is not None:
+                        logger.info("Main process continues to run algorithm.")
+                        self.algorithm.run(
+                            train_dataset=train_dataset,
+                            val_dataset=val_dataset,
+                        )
+                        logger.info("Algorithm exits. Killing the workers.")
+                        self._terminate_processes(processes)
+
+        except KeyboardInterrupt:
+            logger.info("KeyboardInterrupt received. Killing the workers.")
+            self._terminate_processes(processes)
+            logger.info(f"Workers terminated or single worker interrupted.")
+            raise
+        except Exception:
+            logger.exception(f"Unhandled exception in fit method.")
+            self._terminate_processes(processes)
+            logger.info(f"Workers terminated or single worker interrupted.")
+            raise
+        finally:
+            if self.daemon:
+                self.teardown()
+            else:
+                logger.info("Main process exiting. Please use Trainer.kill_orphaned_processes() for cleanup.")

agentlightning/trainer/registry.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Put components in this file to make them available to the Trainer.
+
+Currently only used for ExecutionStrategy.
+"""
+
+ExecutionStrategyRegistry = {
+    "shm": "agentlightning.execution.shared_memory.SharedMemoryExecutionStrategy",
+    # "ipc": "agentlightning.execution.inter_process.InterProcessExecutionStrategy",
+    "cs": "agentlightning.execution.client_server.ClientServerExecutionStrategy",
+}

agentlightning/trainer/trainer.py

@@ -0,0 +1,438 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import functools
+import logging
+import warnings
+from typing import Any, Callable, Dict, Optional, Sequence, TypeVar, Union
+
+from agentlightning.adapter import TraceAdapter, TracerTraceToTriplet
+from agentlightning.algorithm import BaseAlgorithm, Baseline, FastAlgorithm
+from agentlightning.client import AgentLightningClient
+from agentlightning.execution.base import ExecutionStrategy
+from agentlightning.execution.client_server import ClientServerExecutionStrategy
+from agentlightning.execution.events import ExecutionEvent
+from agentlightning.litagent import LitAgent
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.runner import BaseRunner, LitAgentRunner
+from agentlightning.store.base import LightningStore
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer.agentops import AgentOpsTracer
+from agentlightning.tracer.base import BaseTracer
+from agentlightning.types import Dataset, Hook, NamedResources
+
+from .init_utils import build_component, instantiate_component
+from .legacy import TrainerLegacy
+from .registry import ExecutionStrategyRegistry
+
+logger = logging.getLogger(__name__)
+
+T_co = TypeVar("T_co", covariant=True)
+T = TypeVar("T")
+
+ComponentSpec = Union[T, type[T], Callable[[], T], str, Dict[str, Any], None]
+
+
+class Trainer(TrainerLegacy):
+    """Orchestrates the distributed execution of agent rollouts.
+
+    The Trainer is responsible for launching one or more worker processes
+    that run the agent's execution loop. It manages multiprocessing,
+    handles graceful shutdown, and serves as the main entry point for
+    running a client-side agent fleet.
+
+    Attributes:
+        algorithm: An instance of `BaseAlgorithm` to use for training.
+        store: An instance of `LightningStore` to use for storing tasks and traces.
+        runner: An instance of `BaseRunner` to use for running the agent.
+        initial_resources: An instance of `Resources` to use for bootstrapping the fit/dev process.
+            The resources will be handed over to the algorithm.
+            Note that not all algorithms support seeding resources.
+        n_runners: Number of agent runners to run in parallel.
+        max_rollouts: Maximum number of rollouts to process per runner. If None,
+                      workers run until no more rollouts are available.
+        strategy: An instance of `ExecutionStrategy` to use for spawning the algorithm and runners.
+        tracer: A tracer instance, or a string pointing to the class full name or a dictionary with a 'type' key
+                that specifies the class full name and other initialization parameters.
+                If None, a default `AgentOpsTracer` will be created with the current settings.
+        hooks: A sequence of `Hook` instances to be called at various lifecycle stages (e.g., on_trace_start,
+               on_trace_end, on_rollout_start, on_rollout_end).
+        adapter: An instance of `TracerTraceToTriplet` to export data consumble by algorithms from traces.
+        llm_proxy: An instance of `LLMProxy` to use for intercepting the LLM calls.
+                   If not provided, algorithm will create one on its own.
+        n_workers: Number of agent workers to run in parallel. Deprecated in favor of `n_runners`.
+        max_tasks: Maximum number of tasks to process per runner. Deprecated in favor of `max_rollouts`.
+        daemon: Whether worker processes should be daemons. Daemon processes
+                are terminated automatically when the main process exits. Deprecated.
+                Only have effect with `fit_v0`.
+        triplet_exporter: An instance of `TracerTraceToTriplet` to export triplets from traces,
+                          or a dictionary with the initialization parameters for the exporter.
+                          Deprecated. Use `adapter` instead.
+        dev: If True, rollouts are run against the dev endpoint provided in `fit`.
+             Deprecated in favor of `dev()` method.
+    """
+
+    def __init__(
+        self,
+        *,
+        dev: bool = False,
+        n_runners: Optional[int] = None,
+        max_rollouts: Optional[int] = None,
+        initial_resources: Optional[NamedResources] = None,
+        tracer: ComponentSpec[BaseTracer] = None,
+        adapter: ComponentSpec[TraceAdapter[Any]] = None,
+        store: ComponentSpec[LightningStore] = None,
+        runner: ComponentSpec[BaseRunner[Any]] = None,
+        strategy: ComponentSpec[ExecutionStrategy] = None,
+        algorithm: ComponentSpec[BaseAlgorithm] = None,
+        llm_proxy: ComponentSpec[LLMProxy] = None,
+        n_workers: Optional[int] = None,
+        max_tasks: Optional[int] = None,
+        daemon: bool = True,
+        triplet_exporter: ComponentSpec[TracerTraceToTriplet] = None,
+        hooks: Optional[Union[Hook, Sequence[Hook]]] = None,
+    ):
+        # Do not call super().__init__() here.
+        # super().__init__() will call TrainerLegacy's initialization, which is not intended.
+        self.worker_id: Optional[int] = None
+
+        self._dev = dev
+        self.daemon = daemon
+        self._client: AgentLightningClient | None = None  # Will be initialized in fit or fit_v0
+
+        if n_workers is not None:
+            warnings.warn(
+                "`n_workers` is deprecated. Please use `n_runners`.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if n_runners is None:
+            n_runners = n_workers if n_workers is not None else 1
+        else:
+            if n_workers is not None and n_workers != n_runners:
+                warnings.warn(
+                    "`n_workers` is ignored when `n_runners` is provided.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+
+        self.n_runners = n_runners
+        self.n_workers = n_runners  # Backwards compatibility for fit_v0
+
+        if max_tasks is not None:
+            warnings.warn(
+                "`max_tasks` is deprecated. Please use `max_rollouts`.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        if max_rollouts is None:
+            max_rollouts = max_tasks
+        elif max_tasks is not None and max_tasks != max_rollouts:
+            warnings.warn(
+                "`max_tasks` is ignored when `max_rollouts` is provided.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        self.max_rollouts = max_rollouts
+        self.max_tasks = max_tasks if max_tasks is not None else max_rollouts
+
+        self.tracer = self._make_tracer(tracer)
+
+        if adapter is not None and triplet_exporter is not None:
+            warnings.warn(
+                "`triplet_exporter` is deprecated and ignored because `adapter` is provided.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        adapter_spec = adapter if adapter is not None else triplet_exporter
+        self.adapter = self._make_adapter(adapter_spec)
+        self.triplet_exporter = self.adapter  # Backwards compatibility
+
+        self.algorithm = self._make_algorithm(algorithm)
+
+        # We might be able to support a list of resources in future.
+        self.initial_resources = initial_resources
+
+        # The active store for the current execution context
+        self.store = self._make_store(store)
+        self.runner = self._make_runner(runner)
+
+        self.strategy = self._make_strategy(strategy, n_runners=self.n_runners)
+        if hasattr(self.strategy, "n_runners"):
+            strategy_runners = getattr(self.strategy, "n_runners")
+            if isinstance(strategy_runners, int) and strategy_runners > 0:
+                self.n_runners = strategy_runners
+                self.n_workers = strategy_runners
+
+        self.llm_proxy = self._make_llm_proxy(llm_proxy, store=self.store)
+
+        self.hooks = self._normalize_hooks(hooks)
+
+        if not self.daemon:
+            logger.warning(
+                "daemon=False. Worker processes are non-daemonic. "
+                "The worker processes will NOT be terminated when the main process exits. "
+                "The cleanup must be handled manually."
+            )
+
+    def _make_tracer(self, tracer: ComponentSpec[BaseTracer]) -> BaseTracer:
+        """Creates a tracer instance based on the provided configuration."""
+        default_factory = lambda: AgentOpsTracer(
+            agentops_managed=True,
+            instrument_managed=True,
+            daemon=self.daemon,
+        )
+        return build_component(
+            tracer,
+            expected_type=BaseTracer,
+            spec_name="tracer",
+            default_factory=default_factory,
+            dict_requires_type=True,
+            invalid_spec_error_fmt="Invalid tracer type: {actual_type}. Expected BaseTracer, str, dict, or None.",
+            type_error_fmt="Tracer factory returned {type_name}, which is not a BaseTracer subclass.",
+        )
+
+    def _make_algorithm(self, algorithm: ComponentSpec[BaseAlgorithm]) -> Optional[BaseAlgorithm]:
+        """Creates an algorithm instance based on the provided configuration."""
+        return build_component(
+            algorithm,
+            expected_type=BaseAlgorithm,
+            spec_name="algorithm",
+            allow_none=True,
+            invalid_spec_error_fmt="Invalid algorithm type: {actual_type}. Expected BaseAlgorithm, str, dict, or None.",
+            type_error_fmt="Algorithm factory returned {type_name}, which is not a BaseAlgorithm subclass.",
+        )
+
+    def _make_adapter(self, adapter: ComponentSpec[TraceAdapter[Any]]) -> TraceAdapter[Any]:
+        return build_component(
+            adapter,
+            expected_type=TraceAdapter,
+            spec_name="adapter",
+            default_factory=TracerTraceToTriplet,
+            dict_requires_type=False,
+            dict_default_cls=TracerTraceToTriplet,
+            invalid_spec_error_fmt="Invalid adapter type: {actual_type}. Expected TraceAdapter, dict, or None.",
+            type_error_fmt="Adapter factory returned {type_name}, which is not a TraceAdapter subclass.",
+        )
+
+    def _make_store(self, store: ComponentSpec[LightningStore]) -> LightningStore:
+        return build_component(
+            store,
+            expected_type=LightningStore,
+            spec_name="store",
+            default_factory=InMemoryLightningStore,
+            invalid_spec_error_fmt="Invalid store type: {actual_type}. Expected LightningStore, str, dict, or None.",
+            type_error_fmt="Store factory returned {type_name}, which is not a LightningStore subclass.",
+        )
+
+    def _make_strategy(
+        self,
+        strategy: ComponentSpec[ExecutionStrategy],
+        *,
+        n_runners: int,
+    ) -> ExecutionStrategy:
+        if isinstance(strategy, ExecutionStrategy):
+            return strategy
+        optional_defaults: Dict[str, Callable[[], Any]] = {"n_runners": lambda: n_runners}
+
+        def default_factory() -> ExecutionStrategy:
+            return ClientServerExecutionStrategy(n_runners=n_runners, role="both")
+
+        return build_component(
+            strategy,
+            expected_type=ExecutionStrategy,
+            spec_name="strategy",
+            default_factory=default_factory,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid strategy type: {actual_type}. Expected ExecutionStrategy, str, dict, or None.",
+            type_error_fmt="Strategy factory returned {type_name}, which is not an ExecutionStrategy subclass.",
+            registry=ExecutionStrategyRegistry,
+        )
+
+    def _make_llm_proxy(
+        self,
+        llm_proxy: ComponentSpec[LLMProxy],
+        *,
+        store: LightningStore,
+    ) -> Optional[LLMProxy]:
+        if isinstance(llm_proxy, LLMProxy):
+            return llm_proxy
+
+        optional_defaults: Dict[str, Callable[[], Any]] = {"store": lambda: store}
+        if isinstance(llm_proxy, dict):
+            llm_proxy = {**llm_proxy}
+            llm_proxy.setdefault("store", store)
+
+        return build_component(
+            llm_proxy,
+            expected_type=LLMProxy,
+            spec_name="llm_proxy",
+            allow_none=True,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid llm_proxy type: {actual_type}. Expected LLMProxy, dict, str, or None.",
+            type_error_fmt="llm_proxy factory returned {type_name}, which is not an LLMProxy subclass.",
+        )
+
+    def _make_runner(self, runner: ComponentSpec[BaseRunner[Any]]) -> BaseRunner[Any]:
+        optional_defaults: Dict[str, Callable[[], Any]] = {"tracer": lambda: self.tracer}
+        if self.max_rollouts is not None:
+            optional_defaults["max_rollouts"] = lambda: self.max_rollouts
+
+        def default_runner_factory() -> BaseRunner[Any]:
+            return instantiate_component(LitAgentRunner, optional_defaults=optional_defaults)
+
+        return build_component(
+            runner,
+            expected_type=BaseRunner,
+            spec_name="runner",
+            default_factory=default_runner_factory,
+            optional_defaults=optional_defaults,
+            invalid_spec_error_fmt="Invalid runner type: {actual_type}. Expected BaseRunner, callable, str, dict, or None.",
+            type_error_fmt="Runner factory returned {type_name}, which is not a BaseRunner subclass.",
+        )
+
+    def _normalize_hooks(self, hooks: Optional[Union[Hook, Sequence[Hook]]]) -> Sequence[Hook]:
+        if hooks is None:
+            return ()
+        if isinstance(hooks, Hook):
+            return (hooks,)
+        return tuple(hooks)
+
+    def fit(
+        self,
+        agent: LitAgent[T_co],
+        train_dataset: Optional[Dataset[T_co]] = None,
+        *,
+        val_dataset: Optional[Dataset[T_co]] = None,
+    ) -> None:
+        """Run the training loop using the configured strategy, store, and runner.
+
+        Args:
+            agent: The LitAgent instance to be trained on.
+            train_dataset: The dataset to train on.
+            val_dataset: The dataset to validate on.
+        """
+        agent.set_trainer(self)
+
+        algorithm_bundle = functools.partial(
+            self._algorithm_bundle,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            algorithm=self.algorithm,
+        )
+        runner_bundle = functools.partial(self._runner_bundle, agent=agent)
+
+        self.strategy.execute(algorithm_bundle, runner_bundle, self.store)
+
+    def dev(
+        self,
+        agent: LitAgent[T_co],
+        train_dataset: Optional[Dataset[T_co]] = None,
+        *,
+        val_dataset: Optional[Dataset[T_co]] = None,
+    ) -> None:
+        """Dry run the training loop with a FastAlgorithm and the real runner.
+
+        Args:
+            agent: The LitAgent instance to be trained on.
+            train_dataset: The dataset to train on.
+            val_dataset: The dataset to validate on.
+
+        Raises:
+            TypeError: If the configured algorithm is not a :class:`FastAlgorithm`.
+        """
+        agent.set_trainer(self)
+
+        # Sanity check
+        if self.algorithm is None:
+            algorithm = Baseline()
+        else:
+            algorithm = self.algorithm
+
+        if not isinstance(algorithm, FastAlgorithm):
+            raise TypeError(
+                "Trainer.dev() requires an algorithm that inherits from FastAlgorithm. "
+                f"Received {type(algorithm).__name__}."
+            )
+
+        algorithm_bundle = functools.partial(
+            self._algorithm_bundle,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            algorithm=algorithm,
+        )
+        runner_bundle = functools.partial(self._runner_bundle, agent=agent)
+        self.strategy.execute(algorithm_bundle, runner_bundle, self.store)
+
+    async def _algorithm_bundle(
+        self,
+        store: LightningStore,
+        event: ExecutionEvent,
+        train_dataset: Optional[Dataset[T_co]],
+        val_dataset: Optional[Dataset[T_co]],
+        algorithm: Optional[BaseAlgorithm],
+    ) -> None:
+        if algorithm is not None:
+            algorithm.set_trainer(self)
+            algorithm.set_store(store)
+            algorithm.set_adapter(self.adapter)
+            if self.initial_resources is not None:
+                algorithm.set_initial_resources(self.initial_resources)
+            if self.llm_proxy is not None:
+                self.llm_proxy.set_store(store)
+                algorithm.set_llm_proxy(self.llm_proxy)
+
+        if algorithm is None:
+            while not event.is_set():
+                await asyncio.sleep(0.1)
+            return
+        try:
+            if algorithm.is_async():
+                await algorithm.run(  # type: ignore
+                    train_dataset=train_dataset,
+                    val_dataset=val_dataset,
+                )
+            else:
+                # This will block the event loop to maximize the debugging experience
+                # It's the responsibility of the execution strategy to enable async execution
+                algorithm.run(
+                    train_dataset=train_dataset,
+                    val_dataset=val_dataset,
+                )
+        except Exception:
+            logger.exception("Algorithm bundle encountered an error.")
+            raise
+
+    async def _runner_bundle(
+        self, store: LightningStore, worker_id: int, event: ExecutionEvent, agent: LitAgent[T_co]
+    ) -> None:
+        runner_instance: BaseRunner[Any] | None = None
+        runner_initialized = False
+        worker_initialized = False
+        try:
+            # If not using shm execution strategy, we are already in the forked process
+            runner_instance = self.runner
+            runner_instance.init(agent=agent, hooks=self.hooks)
+            runner_initialized = True
+            runner_instance.init_worker(worker_id, store)
+            worker_initialized = True
+            await runner_instance.iter(event=event)
+        except Exception:
+            logger.exception("Runner bundle encountered an error (worker_id=%s).", worker_id)
+            raise
+        finally:
+            if runner_instance is not None:
+                if worker_initialized:
+                    try:
+                        runner_instance.teardown_worker(worker_id)
+                    except Exception:
+                        logger.exception("Error during runner worker teardown (worker_id=%s).", worker_id)
+                if runner_initialized:
+                    try:
+                        runner_instance.teardown()
+                    except Exception:
+                        logger.exception("Error during runner teardown (worker_id=%s).", worker_id)

agentlightning/types.py

@@ -1,204 +0,0 @@
-from typing import Any, Dict, List, Optional, Union, Literal, Annotated
-
-from pydantic import BaseModel, Field, Discriminator
-from opentelemetry.sdk.trace import ReadableSpan
-
-__all__ = [
-    "Triplet",
-    "Rollout",
-    "Task",
-    "TaskInput",
-    "TaskIfAny",
-    "RolloutRawResult",
-    "Resource",
-    "LLM",
-    "PromptTemplate",
-    "ResourceUnion",
-    "NamedResources",
-    "ResourcesUpdate",
-    "GenericResponse",
-    "ParallelWorkerBase",
-]
-
-
-class Triplet(BaseModel):
-    """A standard structure for a single turn in a trajectory."""
-
-    prompt: Any
-    response: Any
-    reward: Optional[float] = None
-    metadata: Dict[str, Any] = Field(default_factory=dict)
-
-
-class Rollout(BaseModel):
-    """The standard reporting object from client to server."""
-
-    rollout_id: str
-
-    # Primary, high-level feedback
-    final_reward: Optional[float] = None
-
-    # Structured, sequential feedback for RL-style optimization
-    triplets: Optional[List[Triplet]] = None
-
-    # Optional, rich-context data for deep analysis
-    trace: Optional[List[Dict[str, Any]]] = Field(
-        default=None,
-        description="A list of spans that conform to the OpenTelemetry JSON format. "
-        "Users of the opentelemetry-sdk can generate this by calling "
-        "json.loads(readable_span.to_json()).",
-    )
-    logs: Optional[List[str]] = None
-
-    # A bucket for any other relevant information
-    metadata: Dict[str, Any] = Field(default_factory=dict)
-
-
-TaskInput = Any
-
-
-class Task(BaseModel):
-    """A task (rollout request) to be processed by the client agent."""
-
-    rollout_id: str
-    input: TaskInput
-
-    mode: Optional[Literal["train", "val", "test"]] = None
-    resources_id: Optional[str] = None
-
-    # Optional fields for tracking task lifecycle
-    create_time: Optional[float] = None
-    last_claim_time: Optional[float] = None
-    num_claims: Optional[int] = None
-
-    # Allow additional metadata fields
-    metadata: Dict[str, Any] = Field(default_factory=dict)
-
-
-class TaskIfAny(BaseModel):
-    is_available: bool
-    task: Optional[Task] = None
-
-
-RolloutRawResult = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], Rollout]
-
-
-class Resource(BaseModel):
-    """
-    Base class for all tunable resources.
-    """
-
-    resource_type: Any
-
-
-class LLM(Resource):
-    """
-    Provide an LLM endpoint and model name as a resource.
-
-    Attributes:
-        endpoint (str): The URL of the LLM API endpoint.
-        model (str): The identifier for the model to be used (e.g., 'gpt-4o').
-        sampling_parameters (SamplingParameters): A dictionary of hyperparameters
-            for model inference, such as temperature, top_p, etc.
-    """
-
-    resource_type: Literal["llm"] = "llm"
-    endpoint: str
-    model: str
-    sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
-
-
-class PromptTemplate(Resource):
-    """
-    A prompt template as a resource.
-
-    Attributes:
-        template (str): The template string. The format depends on the engine.
-        engine (Literal['jinja', 'f-string', 'poml']): The templating engine
-            to use for rendering the prompt. I imagine users can use their own
-            customized engines, but algos can only well operate on a subset of them.
-    """
-
-    resource_type: Literal["prompt_template"] = "prompt_template"
-    template: str
-    engine: Literal["jinja", "f-string", "poml"]
-
-
-# Use discriminated union for proper deserialization
-ResourceUnion = Annotated[Union[LLM, PromptTemplate], Field(discriminator="resource_type")]
-NamedResources = Dict[str, ResourceUnion]
-"""
-A dictionary-like class to hold named resources.
-
-Example:
-    resources: NamedResources = {
-        'main_llm': LLM(
-            endpoint="http://localhost:8080",
-            model="llama3",
-            sampling_parameters={'temperature': 0.7, 'max_tokens': 100}
-        ),
-        'system_prompt': PromptTemplate(
-            template="You are a helpful assistant.",
-            engine='f-string'
-        )
-    }
-"""
-
-
-class ResourcesUpdate(BaseModel):
-    """
-    A resource update message to be sent from the server to clients.
-
-    This message contains a dictionary of resources that clients should use
-    for subsequent tasks. It is used to update the resources available to
-    clients dynamically.
-    """
-
-    resources_id: str
-    resources: NamedResources
-
-
-class GenericResponse(BaseModel):
-    """
-    A generic response message that can be used for various purposes.
-    """
-
-    status: str = "success"
-    message: Optional[str] = None
-    data: Optional[Dict[str, Any]] = None
-
-
-class ParallelWorkerBase:
-    """Base class for objects that can be parallelized across multiple worker processes.
-
-    This class defines the standard lifecycle for parallel processing:
-
-    Main Process:
-        1. init() - Initialize the object in the main process
-        2. spawn workers and call init_worker() in each worker
-        3. run() - Execute the main workload in parallel across workers
-        4. teardown_worker() - Clean up resources in each worker
-        5. teardown() - Final cleanup in the main process
-
-    Subclasses should implement the run() method and optionally override
-    the lifecycle methods for custom initialization and cleanup behavior.
-    """
-
-    def __init__(self) -> None:
-        """Initialize the base class. This method can be overridden by subclasses."""
-        self.worker_id: Optional[int] = None
-
-    def init(self, *args: Any, **kwargs: Any) -> None:
-        pass
-
-    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
-        self.worker_id = worker_id
-
-    def run(self, *args: Any, **kwargs: Any) -> Any:
-        pass
-
-    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
-        pass
-
-    def teardown(self, *args: Any, **kwargs: Any) -> None:
-        pass

agentlightning/types/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .core import *
+from .resources import *
+from .tracer import *

agentlightning/types/core.py

@@ -0,0 +1,338 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Protocol,
+    SupportsIndex,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field, model_validator
+
+from .tracer import Span
+
+if TYPE_CHECKING:
+    from agentlightning.litagent import LitAgent
+    from agentlightning.runner.base import BaseRunner
+    from agentlightning.tracer.base import BaseTracer
+
+__all__ = [
+    "Triplet",
+    "RolloutLegacy",
+    "Task",
+    "TaskInput",
+    "TaskIfAny",
+    "RolloutRawResultLegacy",
+    "RolloutRawResult",
+    "RolloutMode",
+    "GenericResponse",
+    "ParallelWorkerBase",
+    "Dataset",
+    "AttemptStatus",
+    "RolloutStatus",
+    "RolloutConfig",
+    "Rollout",
+    "Attempt",
+    "AttemptedRollout",
+    "Hook",
+]
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class Triplet(BaseModel):
+    """A standard structure for a single turn in a trajectory."""
+
+    prompt: Any
+    response: Any
+    reward: Optional[float] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class RolloutLegacy(BaseModel):
+    """The standard reporting object from client to server."""
+
+    rollout_id: str
+
+    # Echoing the input task
+    task: Optional[Task] = None
+
+    # Primary, high-level feedback
+    final_reward: Optional[float] = None
+
+    # Structured, sequential feedback for RL-style optimization
+    triplets: Optional[List[Triplet]] = None
+
+    # Optional, rich-context data for deep analysis
+    trace: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="A list of spans that conform to the OpenTelemetry JSON format. "
+        "Users of the opentelemetry-sdk can generate this by calling "
+        "json.loads(readable_span.to_json()).",
+    )
+    logs: Optional[List[str]] = None
+
+    # A bucket for any other relevant information
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+RolloutStatus = Literal[
+    "queuing",  # initial status
+    "preparing",  # after the trace is claimed
+    "running",  # after receiving the first trace
+    "failed",  # crashed
+    "succeeded",  # status OK
+    "cancelled",  # cancelled by user (or watchdog)
+    "requeuing",  # retrying
+]
+
+AttemptStatus = Literal[
+    # A status is essentially a process.
+    # It should not have scheduling/management statuses like "queuing" or "cancelled".
+    "preparing",
+    "running",
+    "failed",
+    "succeeded",
+    "unresponsive",  # the worker has not reported results for a while
+    "timeout",  # the worker has been emitting new logs, but have been working on the task for too long
+]
+
+RolloutMode = Literal["train", "val", "test"]
+
+
+class Attempt(BaseModel):
+    """An attempt to execute a rollout. A rollout can have multiple attempts if retries are needed."""
+
+    rollout_id: str  # the rollout this attempt belongs to
+    attempt_id: str  # the universal id for current attempt
+    sequence_id: int  # the sequence number of the attempt, starting from 1
+    start_time: float  # time when the attempt has started
+    end_time: Optional[float] = None  # time when the attempt has ended
+
+    status: AttemptStatus = "preparing"
+    # The rollout worker which is executing this attempt
+    worker_id: Optional[str] = None
+
+    last_heartbeat_time: Optional[float] = None  # last time when the worker has reported progress
+
+    # A bucket for any other relevant information
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class RolloutConfig(BaseModel):
+    """Configurations for rollout execution."""
+
+    timeout_seconds: Optional[float] = None  # none indicates no timeout
+    unresponsive_seconds: Optional[float] = None  # none indicates no unresponsive timeout
+    max_attempts: int = Field(default=1, ge=1)  # including the first attempt
+    retry_condition: List[AttemptStatus] = Field(
+        default_factory=cast(Callable[[], List[AttemptStatus]], list)
+    )  # list of statuses that should trigger a retry
+
+
+class Rollout(BaseModel):
+    rollout_id: str
+
+    # Inputs
+    input: TaskInput
+
+    # Time to track the lifecycle of the rollout
+    start_time: float
+    end_time: Optional[float] = None
+
+    mode: Optional[RolloutMode] = None
+    resources_id: Optional[str] = None
+
+    # Overall scheduling/running information
+    status: RolloutStatus = "queuing"
+
+    config: RolloutConfig = Field(default_factory=RolloutConfig)
+
+    # A bucket for any other relevant information
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class AttemptedRollout(Rollout):
+    """A rollout along with its active attempt."""
+
+    attempt: Attempt
+
+    @model_validator(mode="after")
+    def check_consistency(self) -> AttemptedRollout:
+        if self.attempt.rollout_id != self.rollout_id:
+            raise ValueError("Inconsistent rollout_id between Rollout and Attempt")
+        return self
+
+
+TaskInput = Any
+"""Task input type. Can be any type."""
+
+
+class Task(BaseModel):
+    """A task (rollout request) to be processed by the client agent. Deprecated."""
+
+    rollout_id: str
+    input: TaskInput
+
+    mode: Optional[RolloutMode] = None
+    resources_id: Optional[str] = None
+
+    # Optional fields for tracking task lifecycle
+    create_time: Optional[float] = None
+    last_claim_time: Optional[float] = None
+    num_claims: Optional[int] = None
+
+    # Allow additional metadata fields
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class TaskIfAny(BaseModel):
+    is_available: bool
+    task: Optional[Task] = None
+
+
+RolloutRawResultLegacy = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], RolloutLegacy]
+
+RolloutRawResult = Union[
+    None,  # nothing (relies on tracer)
+    float,  # only final reward
+    List[ReadableSpan],  # constructed OTEL spans by user
+    List[Span],  # constructed Span objects by user
+]
+
+
+class GenericResponse(BaseModel):
+    """
+    A generic response message that can be used for various purposes.
+    """
+
+    status: str = "success"
+    message: Optional[str] = None
+    data: Optional[Dict[str, Any]] = None
+
+
+class ParallelWorkerBase:
+    """Base class for objects that can be parallelized across multiple worker processes.
+
+    This class defines the standard lifecycle for parallel processing:
+
+    Main Process:
+        1. init() - Initialize the object in the main process
+        2. spawn workers and call init_worker() in each worker
+        3. run() - Execute the main workload in parallel across workers
+        4. teardown_worker() - Clean up resources in each worker
+        5. teardown() - Final cleanup in the main process
+
+    Subclasses should implement the run() method and optionally override
+    the lifecycle methods for custom initialization and cleanup behavior.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the base class. This method can be overridden by subclasses."""
+        self.worker_id: Optional[int] = None
+
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        self.worker_id = worker_id
+
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        pass
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+
+class Dataset(Protocol, Generic[T_co]):
+    """The general interface for a dataset.
+
+    It's currently implemented as a protocol, having a similar interface to torch.utils.data.Dataset.
+    You don't have to inherit from this class; you can use a simple list if you want to.
+    """
+
+    def __getitem__(self, index: SupportsIndex, /) -> T_co: ...
+
+    def __len__(self) -> int: ...
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    async def on_trace_start(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the tracer enters the trace context but before the rollout begins.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The :class:`BaseTracer` instance associated with the runner.
+            rollout: The :class:`Rollout` object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. By default, this is a no-op.
+        """
+
+    async def on_trace_end(
+        self, *, agent: LitAgent[Any], runner: BaseRunner[Any], tracer: BaseTracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the rollout completes but before the tracer exits the trace context.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            tracer: The :class:`BaseTracer` instance associated with the runner.
+            rollout: The :class:`Rollout` object that has been processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource cleanup. By default, this is a no-op.
+        """
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: BaseRunner[Any], rollout: Rollout) -> None:
+        """Hook called immediately before a rollout *attempt* begins.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            rollout: The :class:`Rollout` object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    async def on_rollout_end(
+        self,
+        *,
+        agent: LitAgent[Any],
+        runner: BaseRunner[Any],
+        rollout: Rollout,
+        spans: Union[List[ReadableSpan], List[Span]],
+    ) -> None:
+        """Hook called after a rollout *attempt* completes.
+
+        Args:
+            agent: The :class:`LitAgent` instance associated with the runner.
+            runner: The :class:`BaseRunner` managing the rollout.
+            rollout: The :class:`Rollout` object that has been processed.
+            spans: The spans that have been added to the store.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """

agentlightning/types/resources.py

@@ -0,0 +1,189 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import logging
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    Literal,
+    Optional,
+    Union,
+)
+
+from pydantic import BaseModel, Field
+
+from .core import AttemptedRollout
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "Resource",
+    "LLM",
+    "ProxyLLM",
+    "PromptTemplate",
+    "ResourceUnion",
+    "NamedResources",
+    "ResourcesUpdate",
+]
+
+
+class Resource(BaseModel):
+    """
+    Base class for all tunable resources.
+    """
+
+    resource_type: Any
+
+
+class LLM(Resource):
+    """
+    Provide an LLM endpoint and model name as a resource.
+
+    Attributes:
+        endpoint (str): The URL of the LLM API endpoint.
+        model (str): The identifier for the model to be used (e.g., 'gpt-4o').
+        sampling_parameters (SamplingParameters): A dictionary of hyperparameters
+            for model inference, such as temperature, top_p, etc.
+    """
+
+    resource_type: Literal["llm"] = "llm"
+    endpoint: str
+    model: str
+    api_key: Optional[str] = None
+    sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
+
+    def get_base_url(self, *args: Any, **kwargs: Any) -> str:
+        """The base_url to put into openai.OpenAI.
+
+        Users are encouraged to use `base_url` to get the LLM endpoint instead of accessing `endpoint` directly.
+        """
+        return self.endpoint
+
+
+class ProxyLLM(LLM):
+    """Proxy LLM resource that is tailored by `llm_proxy.LLMProxy`."""
+
+    resource_type: Literal["proxy_llm"] = "proxy_llm"  # type: ignore
+    _initialized: bool = False
+
+    def model_post_init(self, __context: Any) -> None:
+        """Mark initialization as complete after Pydantic finishes setup."""
+        super().model_post_init(__context)
+        object.__setattr__(self, "_initialized", True)
+
+    def __getattribute__(self, name: str) -> Any:
+        """Override to emit a warning when endpoint is accessed directly."""
+        # Check if we're accessing endpoint after initialization and not from base_url
+        if name == "endpoint":
+            try:
+                initialized = object.__getattribute__(self, "_initialized")
+            except AttributeError:
+                initialized = False
+
+            if initialized:
+                # Check the call stack to see if we're being called from base_url
+                frame = inspect.currentframe()
+                if frame and frame.f_back:
+                    caller_name = frame.f_back.f_code.co_name
+                    if caller_name != "get_base_url":
+                        logger.warning(
+                            "Accessing 'endpoint' directly on ProxyLLM is discouraged. "
+                            "Use 'get_base_url(rollout_id, attempt_id)' instead to get the properly formatted endpoint."
+                        )
+        return super().__getattribute__(name)
+
+    def with_attempted_rollout(self, rollout: AttemptedRollout) -> LLM:
+        """Bake the rollout and attempt id into the endpoint."""
+        return LLM(
+            endpoint=self.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+            model=self.model,
+            sampling_parameters=self.sampling_parameters,
+            api_key=self.api_key,
+        )
+
+    def get_base_url(self, rollout_id: Optional[str], attempt_id: Optional[str]) -> str:
+        if rollout_id is None and attempt_id is None:
+            return self.endpoint
+
+        if not (isinstance(rollout_id, str) and isinstance(attempt_id, str)):
+            raise ValueError("rollout_id and attempt_id must be strings or all be empty")
+
+        prefix = self.endpoint
+        if prefix.endswith("/"):
+            prefix = prefix[:-1]
+        if prefix.endswith("/v1"):
+            prefix = prefix[:-3]
+            has_v1 = True
+        else:
+            has_v1 = False
+        # Now the prefix should look like "http://localhost:11434"
+
+        # Append the rollout and attempt id to the prefix
+        prefix = prefix + f"/rollout/{rollout_id}/attempt/{attempt_id}"
+        if has_v1:
+            prefix = prefix + "/v1"
+        return prefix
+
+
+class PromptTemplate(Resource):
+    """
+    A prompt template as a resource.
+
+    Attributes:
+        template (str): The template string. The format depends on the engine.
+        engine (Literal['jinja', 'f-string', 'poml']): The templating engine
+            to use for rendering the prompt. I imagine users can use their own
+            customized engines, but algos can only well operate on a subset of them.
+    """
+
+    resource_type: Literal["prompt_template"] = "prompt_template"
+    template: str
+    engine: Literal["jinja", "f-string", "poml"]
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt template with the given kwargs."""
+        if self.engine == "f-string":
+            return self.template.format(**kwargs)
+        else:
+            raise NotImplementedError(
+                "Formatting prompt templates for non-f-string engines with format() helper is not supported yet."
+            )
+
+
+# Use discriminated union for proper deserialization
+# TODO: migrate to use a registry
+ResourceUnion = Annotated[Union[LLM, ProxyLLM, PromptTemplate], Field(discriminator="resource_type")]
+NamedResources = Dict[str, ResourceUnion]
+"""
+A dictionary-like class to hold named resources.
+
+Example:
+    resources: NamedResources = {
+        'main_llm': LLM(
+            endpoint="http://localhost:8080",
+            model="llama3",
+            sampling_parameters={'temperature': 0.7, 'max_tokens': 100}
+        ),
+        'system_prompt': PromptTemplate(
+            template="You are a helpful assistant.",
+            engine='f-string'
+        )
+    }
+"""
+
+
+class ResourcesUpdate(BaseModel):
+    """
+    A resource update message to be sent from the server to clients.
+
+    This message contains a dictionary of resources that clients should use
+    for subsequent tasks. It is used to update the resources available to
+    clients dynamically.
+    """
+
+    resources_id: str
+    resources: NamedResources

agentlightning/types/tracer.py

@@ -0,0 +1,326 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from enum import Enum
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.resources import Resource as OtelResource
+from opentelemetry.sdk.trace import Event as OtelEvent
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
+from opentelemetry.trace.status import Status as OtelStatus
+from pydantic import BaseModel
+
+__all__ = [
+    "AttributeValue",
+    "Attributes",
+    "TraceState",
+    "SpanContext",
+    "TraceStatus",
+    "Event",
+    "Link",
+    "Resource",
+    "Span",
+    "SpanNames",
+    "SpanAttributeNames",
+    "SpanLike",
+]
+
+
+def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:
+    """Convert timestamp from nanoseconds to seconds if needed.
+
+    Auto-detects format: if > 1e12, assumes nanoseconds; otherwise seconds.
+    """
+    if not timestamp:
+        return None
+    return timestamp / 1_000_000_000 if timestamp > 1e12 else timestamp
+
+
+def extract_extra_fields(src: Any, excluded_fields: List[str]) -> Dict[str, Any]:
+    """Extract extra fields from source object, excluding specified fields and private fields."""
+    excluded_fields_set = set(excluded_fields) | set(["_" + k for k in excluded_fields])
+    # Exclude the function fields
+    excluded_fields_set |= set(src.__class__.__dict__.keys())
+    stripped_dict = {k.lstrip("_"): v for k, v in src.__dict__.items()}
+    candidates = {k: v for k, v in stripped_dict.items() if k not in excluded_fields_set and not k.startswith("_")}
+    # This should strip or flatten the unserializable fields
+    candidates_serialized = json.dumps(candidates, default=str)
+    return json.loads(candidates_serialized)
+
+
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+Attributes = Dict[str, AttributeValue]
+TraceState = Dict[str, str]
+
+
+class SpanContext(BaseModel):
+    """Corresponding to opentelemetry.trace.SpanContext"""
+
+    trace_id: str
+    span_id: str
+    is_remote: bool
+    trace_state: TraceState
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.SpanContext) -> "SpanContext":
+        return cls(
+            trace_id=trace_api.format_trace_id(src.trace_id),
+            span_id=trace_api.format_span_id(src.span_id),
+            is_remote=src.is_remote,
+            trace_state={k: v for k, v in src.trace_state.items()} if src.trace_state else {},
+            **extract_extra_fields(src, ["trace_id", "span_id", "is_remote", "trace_state"]),
+        )
+
+
+class TraceStatus(BaseModel):
+    """Corresponding to opentelemetry.trace.Status"""
+
+    status_code: str
+    description: Optional[str] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelStatus) -> "TraceStatus":
+        return cls(
+            status_code=src.status_code.name,
+            description=src.description,
+            **extract_extra_fields(src, ["status_code", "description"]),
+        )
+
+
+class Event(BaseModel):
+    """Corresponding to opentelemetry.trace.Event"""
+
+    name: str
+    attributes: Attributes
+    timestamp: Optional[float] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelEvent) -> "Event":
+        return cls(
+            name=src.name,
+            attributes=dict(src.attributes) if src.attributes else {},
+            timestamp=convert_timestamp(src.timestamp),
+            **extract_extra_fields(src, ["name", "attributes", "timestamp"]),
+        )
+
+
+class Link(BaseModel):
+    """Corresponding to opentelemetry.trace.Link"""
+
+    context: SpanContext
+    attributes: Optional[Attributes] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.Link) -> "Link":
+        return cls(
+            context=SpanContext.from_opentelemetry(src.context),
+            attributes=dict(src.attributes) if src.attributes else None,
+            **extract_extra_fields(src, ["context", "attributes"]),
+        )
+
+
+class Resource(BaseModel):
+    """Corresponding to opentelemetry.sdk.resources.Resource"""
+
+    attributes: Attributes
+    schema_url: str
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelResource) -> "Resource":
+        return cls(
+            attributes=dict(src.attributes) if src.attributes else {},
+            schema_url=src.schema_url if src.schema_url else "",
+            **extract_extra_fields(src, ["attributes", "schema_url"]),
+        )
+
+
+class Span(BaseModel):
+
+    class Config:
+        allow_extra = True  # allow extra fields if needed
+
+    rollout_id: str
+    attempt_id: str
+    # The ID to make spans ordered within a single attempt
+    sequence_id: int
+
+    # Current ID (in hex, formatted via trace_api.format_*)
+    trace_id: str  # one rollout can have traces coming from multiple places
+    span_id: str
+    parent_id: Optional[str]
+
+    # Core ReadableSpan fields
+    name: str
+    status: TraceStatus
+    attributes: Attributes
+    events: List[Event]
+    links: List[Link]
+
+    # Timestamps
+    start_time: Optional[float]
+    end_time: Optional[float]
+
+    # Other parsable fields
+    context: Optional[SpanContext]
+    parent: Optional[SpanContext]
+    resource: Resource
+
+    # Preserve other fields in the readable span as extra fields
+    # Make sure that are json serializable (so no bytes, complex objects, ...)
+
+    @classmethod
+    def from_opentelemetry(
+        cls,
+        src: ReadableSpan,
+        rollout_id: str,
+        attempt_id: str,
+        sequence_id: int,
+    ) -> "Span":
+        context = src.get_span_context()
+        if context is None:
+            trace_id = span_id = 0
+        else:
+            trace_id = context.trace_id
+            span_id = context.span_id
+        return cls(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            trace_id=trace_api.format_trace_id(trace_id),
+            span_id=trace_api.format_span_id(span_id),
+            parent_id=(trace_api.format_span_id(src.parent.span_id) if src.parent else None),
+            name=src.name,
+            status=TraceStatus.from_opentelemetry(src.status),
+            attributes=dict(src.attributes) if src.attributes else {},
+            events=[Event.from_opentelemetry(event) for event in src.events] if src.events else [],
+            links=[Link.from_opentelemetry(link) for link in src.links] if src.links else [],
+            start_time=convert_timestamp(src.start_time),
+            end_time=convert_timestamp(src.end_time),
+            context=SpanContext.from_opentelemetry(context) if context else None,
+            parent=(SpanContext.from_opentelemetry(src.parent) if src.parent else None),
+            resource=Resource.from_opentelemetry(src.resource),
+            **extract_extra_fields(
+                src,
+                [
+                    "name",
+                    "context",
+                    "parent",
+                    "resource",
+                    "attributes",
+                    "events",
+                    "links",
+                    "start_time",
+                    "end_time",
+                    "status",
+                    "span_processor",
+                    "rollout_id",
+                    "attempt_id",
+                    "trace_id",
+                    "span_id",
+                    "parent_id",
+                ],
+            ),
+        )
+
+    @classmethod
+    def from_attributes(
+        cls,
+        *,
+        attributes: Attributes,
+        rollout_id: Optional[str] = None,
+        attempt_id: Optional[str] = None,
+        sequence_id: Optional[int] = None,
+        name: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        start_time: Optional[float] = None,
+        end_time: Optional[float] = None,
+        resource: Optional[Resource] = None,
+    ) -> "Span":
+
+        id_generator = RandomIdGenerator()
+        trace_id = trace_id or trace_api.format_trace_id(id_generator.generate_trace_id())
+        span_id = span_id or trace_api.format_span_id(id_generator.generate_span_id())
+
+        return cls(
+            rollout_id=rollout_id or "",
+            attempt_id=attempt_id or "",
+            sequence_id=sequence_id or 0,
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_id=parent_id,
+            start_time=start_time,
+            end_time=end_time,
+            context=SpanContext(
+                trace_id=trace_id,
+                span_id=span_id,
+                is_remote=False,
+                trace_state={},
+            ),
+            name=name or SpanNames.VIRTUAL.value,
+            resource=resource or Resource(attributes={}, schema_url=""),
+            attributes=attributes,
+            status=TraceStatus(status_code="OK"),
+            events=[],
+            links=[],
+            parent=(
+                SpanContext(
+                    trace_id=trace_id,
+                    span_id=parent_id,
+                    is_remote=False,
+                    trace_state={},
+                )
+                if parent_id
+                else None
+            ),
+        )
+
+
+class SpanNames(str, Enum):
+    """Standard span name values for AgentLightning.
+
+    Currently reward, message, object and exception spans are supported.
+    We will add more spans related to error handling in the future.
+    """
+
+    REWARD = "agentlightning.reward"
+    MESSAGE = "agentlightning.message"
+    OBJECT = "agentlightning.object"
+    EXCEPTION = "agentlightning.exception"
+    VIRTUAL = "agentlightning.virtual"
+
+
+class SpanAttributeNames(str, Enum):
+    """Standard attribute names for AgentLightning spans."""
+
+    MESSAGE = "message"
+    OBJECT = "object"
+
+
+SpanLike = Union[ReadableSpan, Span]

agentlightning/verl/__init__.py

@@ -1,3 +1,8 @@
-from .trainer import *
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This package contains a *hacky* integration of VERL with Agent Lightning."""
+
 from .daemon import *
 from .dataset import *
+from .entrypoint import *
+from .trainer import *

agentlightning/verl/__main__.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 from .entrypoint import main
 
 if __name__ == "__main__":

agentlightning/verl/async_server.py

@@ -1,11 +1,16 @@
-import ray
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
 from copy import deepcopy
 
-from agentlightning.instrumentation.vllm import instrument_vllm, ChatCompletionResponsePatched
+import ray
 from starlette.requests import Request
 from starlette.responses import JSONResponse, StreamingResponse
-from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
 from verl.workers.rollout.vllm_rollout.vllm_async_server import AsyncvLLMServer
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
+
+from agentlightning.instrumentation.vllm import ChatCompletionResponsePatched, instrument_vllm
 
 
 def _unwrap_ray_remote(cls):

agentlightning/verl/daemon.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 import asyncio
 import json
 import random
@@ -5,22 +7,34 @@ import socket
 import threading
 import time
 import uuid
-from typing import Dict, List, Optional
+from collections.abc import Mapping
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
 import requests
 import torch
-from agentlightning import LLM, AgentLightningServer, NamedResources, Rollout, configure_logger
 from flask import Flask, Response, abort, request
-from openai.types.chat.chat_completion import ChatCompletion
 from tensordict import TensorDict
-
 from verl import DataProto
 
+from agentlightning import LLM, AgentLightningServer, NamedResources, RolloutLegacy, configure_logger
+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()
 
+__all__ = [
+    "AgentModeDaemon",
+    "get_left_padded_ids_and_attention_mask",
+    "get_right_padded_ids_and_attention_mask",
+]
 
-def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
+
+def get_left_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
     """
     Left-pad (or truncate) a sequence of token IDs to a fixed length,
     and build the corresponding attention mask.
@@ -49,7 +63,9 @@ def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_
     return padded_ids, attention_mask
 
 
-def get_right_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
+def get_right_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
     """
     Right-pad (or truncate) a sequence of token IDs to a fixed length,
     and build the corresponding attention mask.
@@ -84,6 +100,28 @@ def _find_available_port() -> int:
         return s.getsockname()[1]
 
 
+def _to_native(obj: Any) -> Any:
+    """Convert data retrieved from Parquet to data usable in AGL server."""
+    # 1) Arrays -> list (then recurse)
+    if isinstance(obj, np.ndarray):
+        return _to_native(obj.tolist())
+
+    # 2) NumPy scalar types -> Python scalars
+    if isinstance(obj, np.generic):
+        return _to_native(obj.item())
+
+    # 3) Dict-like -> dict
+    if isinstance(obj, Mapping):
+        return {_to_native(k): _to_native(v) for k, v in obj.items()}  # type: ignore
+
+    # 4) Lists/Tuples/Sets -> list
+    if isinstance(obj, (list, tuple, set)):
+        return [_to_native(x) for x in obj]  # type: ignore
+
+    # 5) Anything else: leave as-is
+    return obj
+
+
 class AgentModeDaemon:
     """
     AgentModeDaemon using the AgentLightningServer SDK.
@@ -95,22 +133,50 @@ class AgentModeDaemon:
 
     def __init__(
         self,
-        port,
-        train_rollout_n,
-        train_information,
-        tokenizer,
-        mini_batch_size,
-        pad_token_id,
-        reward_fillna_value=0.0,
-        llm_timeout_seconds=600.0,
+        port: Optional[int],
+        train_rollout_n: int,
+        train_information: Dict[str, Any],
+        tokenizer: Any,
+        mini_batch_size: int,
+        pad_token_id: int,
+        reward_fillna_value: float = 0.0,
+        llm_timeout_seconds: float = 1200.0,
+        mode: Literal["v0", "v1"] = "v1",
+        llm_proxy: LLMProxy | None = None,
+        store: LightningStore | None = None,
+        adapter: TraceToTripletBase | None = None,
     ):
-        # Server and Task Configuration
-        self.server_port = port
+        self.mode = mode
         self.llm_timeout_seconds = llm_timeout_seconds
-        self.server = AgentLightningServer(
-            host="0.0.0.0", port=self.server_port, task_timeout_seconds=self.llm_timeout_seconds
-        )
-        self.proxy_port = _find_available_port()  # Run proxy on a different port
+
+        # Server and Task Configuration
+        if mode == "v0":
+            assert port is not None
+            self.server_port = port
+            self.server = AgentLightningServer(
+                host="0.0.0.0", port=self.server_port, task_timeout_seconds=self.llm_timeout_seconds
+            )
+            self.proxy_port = _find_available_port()  # Run proxy on a different port
+        else:
+            assert store is not None
+            self.store = store
+            if llm_proxy is None:
+                self.llm_proxy = LLMProxy(
+                    port=_find_available_port(),
+                    model_list=[],
+                    store=store,
+                )
+            else:
+                # Reuse the existing LLM proxy (probably configured by user)
+                self.llm_proxy = llm_proxy
+            if adapter is None:
+                self.adapter = TracerTraceToTriplet()
+            else:
+                # Reuse the one from trainer
+                self.adapter = adapter
+            self._internal_loop: Optional[asyncio.AbstractEventLoop] = None
+            self._internal_loop_thread = threading.Thread(target=self._internal_loop_runner, daemon=True)
+            self._internal_loop_thread.start()
 
         # Training and Data Configuration
         self.train_rollout_n = train_rollout_n
@@ -123,13 +189,21 @@ class AgentModeDaemon:
         # Internal State
         self.backend_llm_server_addresses: List[str] = []
         self._total_tasks_queued = 0
-        self._completed_rollouts: Dict[str, Rollout] = {}
-        self._task_id_to_original_sample: Dict[str, Dict] = {}
+        self._completed_rollouts_v0: Dict[str, RolloutLegacy] = {}
+        self._task_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
         self._server_thread: Optional[threading.Thread] = None
         self._proxy_thread: Optional[threading.Thread] = None
         self.is_train = True
 
-    def _start_proxy_server(self):
+    def _internal_loop_runner(self):
+        """Run the internal loop."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        self._internal_loop = loop
+        loop.run_forever()
+        loop.close()
+
+    def _start_proxy_server_v0(self):
         """
         Initializes and runs a Flask-based proxy server in a separate thread.
         This proxy load-balances requests to the actual backend LLM servers.
@@ -140,7 +214,7 @@ class AgentModeDaemon:
         last_request_time = 0
 
         @app.route("/v1/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
-        def proxy(path):
+        def proxy(path: str):  # type: ignore
             if not self.backend_llm_server_addresses:
                 abort(503, description="No backend LLM servers available.")
 
@@ -165,7 +239,7 @@ class AgentModeDaemon:
                     method=request.method,
                     url=target_url,
                     headers=headers,
-                    params=request.args,
+                    params=request.args,  # type: ignore
                     data=request.get_data(),
                     cookies=request.cookies,
                     allow_redirects=False,
@@ -219,40 +293,90 @@ class AgentModeDaemon:
         self._proxy_thread.start()
         print(f"Proxy server running on port {self.proxy_port}")
 
+    def _update_proxy_server_v1(self):
+        model_name = self.train_information.get("model")
+        if not model_name:
+            raise ValueError("Model name is not set.")
+        self.llm_proxy.update_model_list(
+            [
+                ModelConfig(
+                    {
+                        "model_name": model_name,
+                        "litellm_params": {
+                            "model": "hosted_vllm/" + model_name,
+                            "api_base": f"http://{address}/v1/",
+                        },
+                    }
+                )
+                for address in self.backend_llm_server_addresses
+            ],
+        )
+
+        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()
+
     def start(self):
         """Starts the main AgentLightningServer and the proxy server."""
 
-        def run_server():
-            """Run the AgentLightningServer in a separate thread."""
-            asyncio.run(self.server.run_forever())
+        if self.mode == "v0":
 
-        self._server_thread = threading.Thread(target=run_server, daemon=True)
-        self._server_thread.start()
+            def run_server():
+                """Run the AgentLightningServer in a separate thread."""
+                asyncio.run(self.server.run_forever())
 
-        # Wait for the server's internal startup event to be set.
-        print("Waiting for AgentLightningServer to start...")
-        is_ready = self.server.startup_event.wait(timeout=20.0)  # Wait up to 20s
-        if not is_ready:
-            raise RuntimeError("AgentLightningServer failed to start within the timeout period.")
+            self._server_thread = threading.Thread(target=run_server, daemon=True)
+            self._server_thread.start()
 
-        print(f"AgentLightningServer control plane running on port {self.server_port}")
+            # Wait for the server's internal startup event to be set.
+            print("Waiting for AgentLightningServer to start...")
+            is_ready = self.server.startup_event.wait(timeout=20.0)  # Wait up to 20s
+            if not is_ready:
+                raise RuntimeError("AgentLightningServer failed to start within the timeout period.")
 
-        self._start_proxy_server()
+            print(f"AgentLightningServer control plane running on port {self.server_port}")
 
-    async def _async_set_up(self, data, server_addresses, is_train=True):
+            self._start_proxy_server_v0()
+        else:
+            # Agent lightning server is no longer needed;
+            # Start proxy server in _async_set_up
+            pass
+
+    async def _async_set_up(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Async helper to set up data and resources on the server."""
         self.clear_data_and_server()
-        self.backend_llm_server_addresses = server_addresses
+        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()
         self.is_train = is_train
 
         # 1. Update resources on the server for clients to use
-        llm_resource = LLM(
-            endpoint=f"http://127.0.0.1:{self.proxy_port}/v1",
-            model=self.train_information.get("model", "default-model"),
-            sampling_parameters={"temperature": self.train_information.get("temperature", 0.7)},
-        )
+        if self.mode == "v0":
+            llm_resource = LLM(
+                endpoint=f"http://127.0.0.1:{self.proxy_port}/v1",
+                model=self.train_information.get("model", "default-model"),
+                sampling_parameters={
+                    "temperature": self.train_information.get("temperature", 0.7 if is_train else 0.0)
+                },
+            )
+        else:
+            llm_resource = self.llm_proxy.as_resource(
+                sampling_parameters={
+                    "temperature": self.train_information.get("temperature", 0.7 if is_train else 0.0)
+                },
+            )
+
         resources: NamedResources = {"main_llm": llm_resource}
-        resources_id = await self.server.update_resources(resources)
+
+        if self.mode == "v0":
+            resources_id = await self.server.update_resources(resources)
+        else:
+            resources_update = await self.store.add_resources(resources)
+            resources_id = resources_update.resources_id
 
         # 2. Queue tasks for agents to process
         keys = list(data.keys())
@@ -265,34 +389,58 @@ class AgentModeDaemon:
             original_sample["data_id"] = data_id
 
             # For training, each sample is rolled out multiple times
-            for j in range(rollouts_per_sample):
+            for _ in range(rollouts_per_sample):
                 task_metadata = {"data_id": data_id, "is_train": is_train}
 
                 # Data ID is different from Rollout ID, as one data can have multiple rollouts.
-                rollout_id = await self.server.queue_task(
-                    sample=original_sample,
-                    mode="train" if is_train else "val",
-                    resources_id=resources_id,
-                    metadata=task_metadata,
-                )
+                if self.mode == "v0":
+                    rollout_id = await self.server.queue_task(
+                        sample=_to_native(original_sample),
+                        mode="train" if is_train else "val",
+                        resources_id=resources_id,
+                        metadata=task_metadata,
+                    )
+                else:
+                    rollout = await self.store.enqueue_rollout(
+                        input=_to_native(original_sample),
+                        mode="train" if is_train else "val",
+                        resources_id=resources_id,
+                        metadata=task_metadata,
+                    )
+                    await self.store.update_rollout(
+                        rollout_id=rollout.rollout_id,
+                        config=RolloutConfig(
+                            unresponsive_seconds=self.llm_timeout_seconds,
+                            timeout_seconds=self.llm_timeout_seconds,
+                        ),
+                    )
+                    rollout_id = rollout.rollout_id
+
                 # Store original sample data to reconstruct batch information later
                 self._task_id_to_original_sample[rollout_id] = original_sample
                 self._total_tasks_queued += 1
 
-    def set_up_data_and_server(self, data, server_addresses, is_train=True):
+    def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Synchronous wrapper for setting up data and server resources."""
-        if not self.server.loop or not self.server.startup_event.is_set():
-            raise RuntimeError("Server is not running or ready.")
-
         coro = self._async_set_up(data, server_addresses, is_train)
-        future = asyncio.run_coroutine_threadsafe(coro, self.server.loop)
+
+        if self.mode == "v0":
+            if not self.server.loop or not self.server.startup_event.is_set():
+                raise RuntimeError("Server is not running or ready.")
+
+            future = asyncio.run_coroutine_threadsafe(coro, self.server.loop)
+
+        else:
+            if self._internal_loop is None:
+                raise RuntimeError("Internal loop is not running.")
+            future = asyncio.run_coroutine_threadsafe(coro, self._internal_loop)
         try:
             future.result(timeout=60)  # Wait for completion with a timeout
         except Exception as e:
             print(f"Failed to set up data on server: {e}")
             raise
 
-    def _validate_data(self, rollout: Rollout):
+    def _validate_data(self, rollout: RolloutLegacy):
         if rollout.final_reward is None:
             print(
                 f"Warning: Reward is None for rollout {rollout.rollout_id}, will be auto-set to {self.reward_fillna_value}."
@@ -306,29 +454,98 @@ class AgentModeDaemon:
         elif any(not r.prompt.get("token_ids", []) for r in rollout.triplets):
             print(f"Warning: Rollout {rollout.rollout_id} contains empty prompt: {rollout.triplets}")
 
-    async def _async_run_until_finished(self, verbose=True):
+    async def _validate_data_v1(self, rollout: Rollout) -> RolloutLegacy:
+        """Convert Rollout to RolloutLegacy and validate.
+
+        1. Task: construct from Rollout
+        2. Triplets: obtained by querying spans and feeding into the adapter
+        3. Final reward: extracted from last triplet's reward, searching backwards if not found
+        """
+        # Query spans for this rollout (latest attempt)
+        spans = await self.store.query_spans(rollout.rollout_id, attempt_id="latest")
+
+        # Convert spans to triplets using the adapter
+        if not spans:
+            # No triplets found, will emit a warning later.
+            triplets = []
+        else:
+            triplets = self.adapter.adapt(spans)
+
+        # Extract final reward from triplets
+        final_reward: Optional[float] = None
+        if triplets:
+            # Search backwards through triplets for the first non-None reward
+            for triplet in reversed(triplets):
+                if triplet.reward is not None:
+                    final_reward = triplet.reward
+                    break
+
+        # Construct the Task object from Rollout
+        task = Task(
+            rollout_id=rollout.rollout_id,
+            input=rollout.input,
+            mode=rollout.mode,
+            resources_id=rollout.resources_id,
+            metadata=rollout.metadata or {},
+        )
+
+        # Create the Rollout object (without trace and logs as per user's note)
+        result_rollout = RolloutLegacy(
+            rollout_id=rollout.rollout_id,
+            task=task,
+            final_reward=final_reward,
+            triplets=triplets,
+            metadata=rollout.metadata or {},
+        )
+
+        # Run the same validation as v0
+        self._validate_data(result_rollout)
+
+        return result_rollout
+
+    async def _async_run_until_finished(self, verbose: bool = True):
         """Async helper to wait for all tasks to complete."""
-        while len(self._completed_rollouts) < self._total_tasks_queued:
-            completed_batch = await self.server.retrieve_completed_rollouts()
+        while len(self._completed_rollouts_v0) < self._total_tasks_queued:
+            if self.mode == "v0":
+                completed_batch = await self.server.retrieve_completed_rollouts()
+            else:
+                completed_batch = await self.store.wait_for_rollouts(
+                    rollout_ids=list(self._task_id_to_original_sample.keys()), timeout=0
+                )
             for rollout in completed_batch:
-                self._validate_data(rollout)
-                self._completed_rollouts[rollout.rollout_id] = rollout
+                if rollout.rollout_id in self._completed_rollouts_v0:
+                    # Already processed, skip
+                    continue
+                if isinstance(rollout, Rollout):
+                    rollout = await self._validate_data_v1(rollout)
+                else:
+                    self._validate_data(rollout)
+                if rollout.rollout_id not in self._task_id_to_original_sample:
+                    print(f"Warning: Received unknown rollout ID {rollout.rollout_id}, skipping.")
+                else:
+                    self._completed_rollouts_v0[rollout.rollout_id] = rollout
             if verbose:
-                print(f"Completed {len(self._completed_rollouts)}/{self._total_tasks_queued} tasks...")
+                print(f"Completed {len(self._completed_rollouts_v0)}/{self._total_tasks_queued} tasks...")
             await asyncio.sleep(5)
+
         print("All tasks finished.")
 
-    def run_until_all_finished(self, verbose=True):
+    def run_until_all_finished(self, verbose: bool = True):
         """Synchronously waits for all queued tasks to be completed and reported."""
         if self._total_tasks_queued == 0:
             print("Warning: No tasks were queued.")
             return
 
-        if not self.server.loop or not self.server.startup_event.is_set():
-            raise RuntimeError("Server is not running or ready.")
+        if self.mode == "v0":
+            if not self.server.loop or not self.server.startup_event.is_set():
+                raise RuntimeError("Server is not running or ready.")
+            loop = self.server.loop
+        else:
+            loop = self._internal_loop
+            assert loop is not None
 
         coro = self._async_run_until_finished(verbose)
-        future = asyncio.run_coroutine_threadsafe(coro, self.server.loop)
+        future = asyncio.run_coroutine_threadsafe(coro, loop)
         try:
             future.result()  # Wait indefinitely for all tasks to complete
         except Exception as e:
@@ -338,14 +555,16 @@ class AgentModeDaemon:
     def get_test_metrics(self):
         """Calculates and returns metrics for a validation run."""
         assert not self.is_train, "This method should only be called during validation."
-        assert len(self._completed_rollouts) == self._total_tasks_queued
+        assert len(self._completed_rollouts_v0) == self._total_tasks_queued
 
-        sample_stat_list = []
-        for rollout_id, rollout in self._completed_rollouts.items():
+        sample_stat_list: List[Dict[str, Any]] = []
+        for _, rollout in self._completed_rollouts_v0.items():
+            final_reward = self._fillna_reward(rollout)
             if not rollout.triplets:
+                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
+                sample_stat_list.append({"reward": final_reward})
                 continue
             response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
-            final_reward = self._fillna_reward(rollout)
             sample_stat_list.append(
                 {
                     "sum_response_length": np.sum(response_length_list),
@@ -355,14 +574,19 @@ class AgentModeDaemon:
                 }
             )
 
+        stats_w_trace = [stat for stat in sample_stat_list if "sum_response_length" in stat]
         return {
-            "val/reward": np.mean([stat["reward"] for stat in sample_stat_list]),
-            "val/mean_response_length": np.mean([stat["mean_response_length"] for stat in sample_stat_list]),
-            "val/sum_response_length": np.mean([stat["sum_response_length"] for stat in sample_stat_list]),
-            "val/turn_count": np.mean([stat["turn_count"] for stat in sample_stat_list]),
+            "val/n_rollouts": len(sample_stat_list),
+            "val/n_rollouts_w_trace": len(stats_w_trace),
+            "val/reward": np.mean(
+                [stat["reward"] for stat in sample_stat_list]
+            ),  # each rollout must have a reward (fillna if missing)
+            "val/mean_response_length": np.mean([stat["mean_response_length"] for stat in stats_w_trace]),
+            "val/sum_response_length": np.mean([stat["sum_response_length"] for stat in stats_w_trace]),
+            "val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
         }
 
-    def get_train_data_batch(self, max_prompt_length, max_response_length, device):
+    def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
         """
         Processes completed rollouts to generate a training data batch.
 
@@ -371,14 +595,19 @@ class AgentModeDaemon:
         truncation, and tensor creation for the PPO training loop.
         """
         assert self.is_train, "This method should only be called during training."
-        assert len(self._completed_rollouts) == self._total_tasks_queued
+        assert len(self._completed_rollouts_v0) == self._total_tasks_queued
 
         # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
-        finished_id_to_sample_info = {}
-        for rollout_id, rollout in self._completed_rollouts.items():
+        finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {}
+        finished_id_to_final_reward: Dict[str, float] = {}
+        for rollout_id, rollout in self._completed_rollouts_v0.items():
             original_sample = self._task_id_to_original_sample[rollout_id]
 
+            final_reward = self._fillna_reward(rollout)
+
             if not rollout.triplets:
+                finished_id_to_final_reward[rollout_id] = final_reward
+                print(f"Warning: No triplets found for training rollout {rollout.rollout_id}, skipping.")
                 continue
 
             # The client should report triplets that contain prompt_ids and response_ids.
@@ -388,14 +617,13 @@ class AgentModeDaemon:
                 {"prompt_ids": t.prompt.get("token_ids", []), "response_ids": t.response.get("token_ids", [])}
                 for t in rollout.triplets
             ]
-
-            final_reward = self._fillna_reward(rollout)
             info = {
                 "reward": final_reward,
                 "trace_list": trace_list,
                 "data_id": original_sample["data_id"],
             }
             finished_id_to_sample_info[rollout_id] = info
+            finished_id_to_final_reward[rollout_id] = final_reward
         #
         # --- Data processing and tensor creation logic ---
         # Get all the reported data.
@@ -407,9 +635,15 @@ class AgentModeDaemon:
         #     discarded here. They are only truncated and marked, to be discarded later.
         #     This is for the correctness of the advantage calculation.
         #   - The discard for the PPO mini-batch should also be handled this way.
-        input_ids_list, input_attention_mask_list = [], []
-        response_ids_list, response_attention_mask_list = [], []
-        reward_list, data_id_list, rollout_id_list, turn_index_list, is_drop_list = [], [], [], [], []
+        input_ids_list: List[List[int]] = []
+        input_attention_mask_list: List[List[int]] = []
+        response_ids_list: List[List[int]] = []
+        response_attention_mask_list: List[List[int]] = []
+        reward_list: List[float] = []
+        data_id_list: List[str] = []
+        rollout_id_list: List[str] = []
+        turn_index_list: List[int] = []
+        is_drop_list: List[bool] = []
         n_trunc_sample_because_of_response = 0
 
         for rollout_id, sample_info in finished_id_to_sample_info.items():
@@ -484,30 +718,33 @@ class AgentModeDaemon:
         data_proto = DataProto(batch=batch)
 
         data_metrics = {
-            "agent_mode/n_trunc_sample_because_of_response": n_trunc_sample_because_of_response,
-            "agent_mode/n_sample_to_train": n_transition,
+            "training/reward": np.mean(list(finished_id_to_final_reward.values())),
+            "training/n_rollouts": len(finished_id_to_final_reward),
+            "training/n_rollouts_w_trace": len(finished_id_to_sample_info),
+            "training/n_truncated_triplets": n_trunc_sample_because_of_response,
+            "training/n_triplets": n_transition,
         }
 
         # Add non-tensor data for advantage calculation and logging
-        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)
-        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)
-        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)
+        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)  # type: ignore
+        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)  # type: ignore
+        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
 
         return data_proto, data_metrics
 
     def clear_data_and_server(self):
         """Resets the internal state of the daemon for the next run."""
         self.backend_llm_server_addresses = []
-        self._completed_rollouts.clear()
+        self._completed_rollouts_v0.clear()
         self._task_id_to_original_sample.clear()
         self._total_tasks_queued = 0
         # For a true reset, the server's internal queues would also need clearing.
         # This implementation assumes that `set_up_data_and_server` is called
         # for each new run, effectively starting a fresh batch.
 
-    def _fillna_reward(self, rollout):
+    def _fillna_reward(self, rollout: RolloutLegacy):
         if rollout.final_reward is None:
-            if self.reward_fillna_value is not None:
+            if self.reward_fillna_value is not None:  # type: ignore
                 final_reward = self.reward_fillna_value
             else:
                 raise ValueError(f"Reward is None for rollout {rollout.rollout_id}, please check the reward function.")

agentlightning/verl/dataset.py

@@ -1,6 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
 import torch
+from datasets import Dataset as HuggingFaceDataset
+from omegaconf import DictConfig
 from verl.utils.dataset.rl_dataset import RLHFDataset
 
+from agentlightning.types import Dataset
+
+__all__ = [
+    "AgentDataset",
+    "LoadedDataset",
+]
+
 
 class AgentDataset(RLHFDataset):
 
@@ -18,3 +31,14 @@ class AgentDataset(RLHFDataset):
         # Workaround for data proto. At least one tensor is needed.
         row_dict["fake_ids"] = torch.ones(1, dtype=torch.int)
         return row_dict
+
+
+class LoadedDataset(AgentDataset):
+
+    def __init__(self, dataset: Dataset):
+        super().__init__([], None, DictConfig({}))  # type: ignore
+        dataset_copy = [dataset[i] for i in range(len(dataset))]
+        self.dataframe = HuggingFaceDataset.from_list(dataset_copy)
+
+    def _read_files_and_tokenize(self):
+        pass

agentlightning/verl/entrypoint.py

@@ -1,18 +1,42 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+from typing import Any
+
 import hydra
 import ray
-
-from .dataset import AgentDataset
-from .trainer import AgentLightningTrainer
-from verl.trainer.ppo.reward import load_reward_manager
 from verl.trainer.main_ppo import create_rl_sampler
+from verl.trainer.ppo.reward import load_reward_manager
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset
+
+from .dataset import AgentDataset, LoadedDataset
+from .trainer import AgentLightningTrainer
+
+__all__ = [
+    "main",
+    "run_ppo",
+    "TaskRunner",
+]
 
 
 @hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
 def main(config):
-    run_ppo(config)
+    run_ppo(config, train_dataset=None, val_dataset=None, store=None, llm_proxy=None, adapter=None)
 
 
-def run_ppo(config) -> None:
+def run_ppo(
+    config: Any,
+    train_dataset: Dataset[Any] | None,
+    val_dataset: Dataset[Any] | None,
+    store: LightningStore | None,
+    llm_proxy: LLMProxy | None,
+    adapter: TraceAdapter[Any] | None,
+) -> None:
     if not ray.is_initialized():
         # this is for local ray cluster
         ray.init(
@@ -23,17 +47,33 @@ def run_ppo(config) -> None:
         )
 
     runner = TaskRunner.remote()
-    ray.get(runner.run.remote(config))
+    ray.get(
+        runner.run.remote(
+            config=config,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            store=store,
+            llm_proxy=llm_proxy,
+            adapter=adapter,
+        )
+    )
 
 
 @ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
 class TaskRunner:
-    def run(self, config):
+    def run(
+        self,
+        config: Any,
+        train_dataset: Dataset | None,
+        val_dataset: Dataset | None,
+        store: LightningStore | None,
+        llm_proxy: LLMProxy | None,
+        adapter: TraceAdapter | None,
+    ):
         # print initial config
         from pprint import pprint
 
         from omegaconf import OmegaConf
-
         from verl.utils.fs import copy_to_local
 
         pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values
@@ -121,18 +161,26 @@ class TaskRunner:
         from verl.utils.dataset.rl_dataset import collate_fn
 
         # Use our special dataset
-        train_dataset = AgentDataset(
-            data_files=config.data.train_files,
-            tokenizer=tokenizer,
-            processor=processor,
-            config=config.data,
-        )
-        val_dataset = AgentDataset(
-            data_files=config.data.val_files,
-            tokenizer=tokenizer,
-            processor=processor,
-            config=config.data,
-        )
+        if train_dataset is None:
+            train_dataset = AgentDataset(
+                data_files=config.data.train_files,
+                tokenizer=tokenizer,
+                processor=processor,
+                config=config.data,
+            )
+        else:
+            train_dataset = LoadedDataset(train_dataset)
+
+        if val_dataset is None:
+            val_dataset = AgentDataset(
+                data_files=config.data.val_files,
+                tokenizer=tokenizer,
+                processor=processor,
+                config=config.data,
+            )
+        else:
+            val_dataset = LoadedDataset(val_dataset)
+
         train_sampler = create_rl_sampler(config.data, train_dataset)
         trainer = AgentLightningTrainer(
             config=config,
@@ -147,6 +195,9 @@ class TaskRunner:
             val_dataset=val_dataset,
             collate_fn=collate_fn,
             train_sampler=train_sampler,
+            store=store,
+            llm_proxy=llm_proxy,
+            adapter=adapter,
         )
         trainer.init_workers()
         trainer.fit()

agentlightning/verl/trainer.py

@@ -1,35 +1,48 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+from __future__ import annotations
+
 import random
 from contextlib import contextmanager
 from copy import deepcopy
+from pprint import pprint
 from typing import Dict, Tuple
 
 import numpy as np
 import torch
-from omegaconf import OmegaConf
-from pprint import pprint
-from tqdm import tqdm
-
 from codetiming import Timer
+from omegaconf import OmegaConf
+from tqdm import tqdm
 from verl import DataProto
 from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
-from verl.trainer.ppo.ray_trainer import (
-    RayPPOTrainer,
-    AdvantageEstimator,
-    apply_kl_penalty,
-    compute_advantage,
-    compute_response_mask,
-)
 from verl.trainer.ppo.core_algos import agg_loss
 from verl.trainer.ppo.metric_utils import (
     compute_data_metrics,
     compute_throughout_metrics,
     compute_timing_metrics,
 )
+from verl.trainer.ppo.ray_trainer import (
+    AdvantageEstimator,
+    RayPPOTrainer,
+    apply_kl_penalty,
+    compute_advantage,
+    compute_response_mask,
+)
 from verl.utils.metric import reduce_metrics
 from verl.utils.tracking import Tracking
 
+from agentlightning.adapter import TraceAdapter, TraceToTripletBase
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+
 from .daemon import AgentModeDaemon
 
+__all__ = [
+    "AgentLightningTrainer",
+]
+
 
 @contextmanager
 def _timer(name: str, timing_raw: Dict[str, float]):
@@ -56,6 +69,14 @@ class AgentLightningTrainer(RayPPOTrainer):
     4. Streamlined validation using agent_mode validation
     """
 
+    def __init__(
+        self, store: LightningStore | None, llm_proxy: LLMProxy | None, adapter: TraceAdapter | None, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.store = store
+        self.llm_proxy = llm_proxy
+        self.adapter = adapter
+
     def _validate(self):
         assert len(self.val_dataloader) == 1, "Please set val_batch_size to None for better throughput."
 
@@ -105,7 +126,7 @@ class AgentLightningTrainer(RayPPOTrainer):
                 with _timer("gen_max", timing_raw):
                     gen_baseline_batch = deepcopy(gen_batch)
                     gen_baseline_batch.meta_info["do_sample"] = False
-                    gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
+                    gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)
 
                     batch = batch.union(gen_baseline_output)
                     reward_baseline_tensor = self.reward_fn(batch)
@@ -195,7 +216,7 @@ class AgentLightningTrainer(RayPPOTrainer):
 
             # 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["agent_mode/n_dropped_sample_because_of_prompt"] = (
+            metrics["training/n_triplets_prompt_too_long"] = (
                 batch.batch["is_drop_mask"].shape[0] - keep_indices.shape[0]
             )
             batch = batch[keep_indices]
@@ -207,7 +228,7 @@ class AgentLightningTrainer(RayPPOTrainer):
             batch.reorder(torch.tensor(random_indices).type(torch.int32))
             n_remained_transition = n_transition // mini_batch_size * mini_batch_size
             batch = batch[list(range(n_remained_transition))]
-            metrics["agent_mode/n_dropped_sample_because_of_mini_batch"] = n_transition - n_remained_transition
+            metrics["training/n_triplets_dropped_remainder"] = n_transition - n_remained_transition
 
             # Agent mode note: Change the order of balance batch;
             #     1. first calculate advantage
@@ -274,6 +295,8 @@ class AgentLightningTrainer(RayPPOTrainer):
         self._load_checkpoint()
 
         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.")
         self.agent_mode_daemon = AgentModeDaemon(
             self.config.agentlightning.port,
             self.config.actor_rollout_ref.rollout.n,
@@ -287,6 +310,10 @@ class AgentLightningTrainer(RayPPOTrainer):
             tokenizer=self.tokenizer,
             mini_batch_size=self.config.actor_rollout_ref.actor.ppo_mini_batch_size,
             pad_token_id=self.tokenizer.pad_token_id,
+            mode="v1" if self.store is not None else "v0",
+            store=self.store,
+            llm_proxy=self.llm_proxy,
+            adapter=self.adapter,
         )
         self.agent_mode_daemon.start()
 

docs/algorithm-zoo/apo.md

@@ -0,0 +1,25 @@
+# APO
+
+!!! tip "Shortcut"
+
+    You can use the shortcut `agl.APO(...)` to create an APO instance.
+
+    ```python
+    import agentlightning as agl
+
+    agl.APO(...)
+    ```
+
+## Installation
+
+```bash
+pip install agentlightning[apo]
+```
+
+## Tutorials Using APO
+
+TBD
+
+## References
+
+::: agentlightning.algorithm.apo

docs/algorithm-zoo/index.md

@@ -0,0 +1,10 @@
+# Algorithm Zoo
+
+AgentLightning includes several popular and frequently requested algorithms in its built-in library, allowing agent developers to use them directly. These algorithms are designed to be compatible with most agent scenarios.
+
+For customizing algorithms, see [Algorithm-side References](../reference/algorithm.md).
+
+| Algorithm | Optimizing Resources | Description |
+| --------- | ------------------- | ----------- |
+| [APO](./apo.md) | [PromptTemplate][agentlightning.PromptTemplate] | Automatic Prompt Optimization (APO) algorithm using textual gradients and beam search. |
+| [VERL](./verl.md) | [LLM][agentlightning.LLM] | Reinforcement Learning with [VERL framework](https://github.com/volcengine/verl). |

docs/algorithm-zoo/verl.md

@@ -0,0 +1,37 @@
+# VERL
+
+!!! tip "Shortcut"
+
+    You can use the shortcut `agl.VERL(...)` to create a VERL instance.
+
+    ```python
+    import agentlightning as agl
+
+    agl.VERL(...)
+    ```
+
+!!! warning "Customization note"
+
+    Customization of VERL is not supported as of current version. We recommend copying the source code from VERL and modifying it as needed to suit your requirements.
+
+## Installation
+
+```bash
+pip install agentlightning[verl]
+```
+
+!!! warning
+
+    For best results, follow the steps in the [installation guide](../quickstart/installation.md) to set up VERL and its dependencies. Installing VERL directly with `pip install agentlightning[verl]` can cause issues unless you already have a compatible version of PyTorch installed.
+
+## Tutorials Using VERL
+
+TBD
+
+## References - Entrypoint
+
+::: agentlightning.algorithm.verl
+
+## References - Implementation
+
+::: agentlightning.verl