Store Benchmark - Part 4 (#356)

Author: ultmaster

.github/workflows/benchmark.yml

@@ -127,6 +127,7 @@ jobs:
             sleep 1
           done
           echo "Store did not become ready in time" >&2
+          docker compose -f "$COMPOSE_FILE" logs app
           exit 1
 
       - name: Prepare artifact directory
@@ -177,6 +178,9 @@ jobs:
           if [ -d docker/data/prometheus ]; then
             tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-${SCENARIO_ID}-${BACKEND_ID}.tar.gz" prometheus
           fi
+          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
+            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-${SCENARIO_ID}-${BACKEND_ID}.log" || true
+          fi
 
       - name: Upload benchmark artifacts
         if: ${{ always() }}
@@ -185,3 +189,146 @@ jobs:
           name: benchmark-${{ matrix.scenario.id }}-${{ matrix.backend.id }}
           path: ${{ env.ARTIFACT_DIR }}
           if-no-files-found: error
+
+  micro-benchmark:
+    name: Micro-benchmark (${{ matrix.backend.id }}, ${{ matrix.mode.display }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      matrix:
+        backend:
+          - id: memory
+            compose_file: compose.prometheus-memory-store.yml
+          - id: mongo
+            compose_file: compose.prometheus-mongo-store.yml
+        mode:
+          - id: worker
+            display: Update worker throughput
+            cli: worker
+          - id: dequeue-empty
+            display: Dequeue empty throughput
+            cli: dequeue-empty
+          - id: rollout
+            display: Rollout + span throughput
+            cli: rollout
+    env:
+      STORE_URL: http://localhost:4747
+      STORE_API_URL: http://localhost:4747/v1/agl
+      PROM_URL: http://localhost:9090
+      BACKEND_ID: ${{ matrix.backend.id }}
+      MODE_ID: ${{ matrix.mode.id }}
+      ARTIFACT_DIR: artifacts/micro-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+      COMPOSE_FILE: ${{ matrix.backend.compose_file }}
+      AGL_STORE_N_WORKERS: 8
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - name: Sync dependencies
+        run: uv sync --frozen --extra mongo --group core-stable --group dev
+
+      - name: Reset benchmark data directories
+        run: |
+          set -euo pipefail
+          cd docker
+          rm -rf data
+          bash setup.sh
+
+      - name: Launch ${{ matrix.backend.id }} Prometheus stack
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+          docker compose -f "$COMPOSE_FILE" up -d --quiet-pull
+
+      - name: Wait for store readiness
+        run: |
+          set -euo pipefail
+          for attempt in {1..60}; do
+            if curl -fsS "$STORE_API_URL/health" >/dev/null 2>&1; then
+              exit 0
+            fi
+            sleep 1
+          done
+          echo "Store did not become ready in time" >&2
+          cd docker && docker compose -f "$COMPOSE_FILE" logs app
+          exit 1
+
+      - name: Prepare artifact directory
+        run: mkdir -p "$ARTIFACT_DIR"
+
+      - name: Record micro benchmark start
+        run: echo "BENCHMARK_START=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run ${{ matrix.mode.display }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          uv run --locked --no-sync python -m tests.benchmark.micro_benchmark \
+            --store-url "$STORE_URL" \
+            --summary-file "$ARTIFACT_DIR/summary-${MODE_ID}.txt" \
+            "${{ matrix.mode.cli }}" | tee "$ARTIFACT_DIR/micro-${MODE_ID}.txt"
+
+      - name: Record micro benchmark end
+        if: ${{ always() }}
+        run: echo "BENCHMARK_END=$(date -u +%FT%TZ)" >> "$GITHUB_ENV"
+
+      - name: Run micro benchmark analysis
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          if [ -z "${BENCHMARK_START:-}" ] || [ -z "${BENCHMARK_END:-}" ]; then
+            echo "Analysis skipped: benchmark window not recorded." > "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+            exit 1
+          fi
+          uv run --locked --no-sync python -m tests.benchmark.analysis \
+            --prom-url "$PROM_URL" \
+            --store-url "$STORE_API_URL" \
+            --start "$BENCHMARK_START" \
+            --end "$BENCHMARK_END" \
+            | tee "$ARTIFACT_DIR/analysis-${MODE_ID}.txt"
+
+      - name: Show micro benchmark summary
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          summary_file="$ARTIFACT_DIR/summary-${MODE_ID}.txt"
+          if [ -f "$summary_file" ]; then
+            echo "Micro benchmark summary ($MODE_ID/$BACKEND_ID):"
+            cat "$summary_file"
+          else
+            echo "Summary file not found: $summary_file"
+          fi
+
+      - name: Stop ${{ matrix.backend.id }} Prometheus stack
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          cd docker
+          docker compose -f "$COMPOSE_FILE" down -v || true
+
+      - name: Archive Prometheus metrics
+        if: ${{ always() }}
+        run: |
+          set -euo pipefail
+          mkdir -p "$ARTIFACT_DIR"
+          if [ -d docker/data/prometheus ]; then
+            tar -C docker/data -czf "$ARTIFACT_DIR/prometheus-micro-${MODE_ID}-${BACKEND_ID}.tar.gz" prometheus
+          fi
+          if docker compose -f "$COMPOSE_FILE" ps --format '{{.Name}}' >/dev/null 2>&1; then
+            docker compose -f "$COMPOSE_FILE" logs app > "$ARTIFACT_DIR/docker-micro-${MODE_ID}-${BACKEND_ID}.log" || true
+          fi
+
+      - name: Upload micro benchmark artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: micro-benchmark-${{ matrix.mode.id }}-${{ matrix.backend.id }}
+          path: ${{ env.ARTIFACT_DIR }}
+          if-no-files-found: error

agentlightning/cli/store.py

@@ -64,7 +64,9 @@ def main(argv: Iterable[str] | None = None) -> int:
     setup_logging(args.log_level)
 
     if args.backend == "memory":
-        store = InMemoryLightningStore(prometheus=args.prometheus)
+        store = InMemoryLightningStore(
+            prometheus=args.prometheus, thread_safe=True
+        )  # Using thread_safe store for server
     elif args.backend == "mongo":
         from agentlightning.store.mongo import MongoLightningStore
 

agentlightning/runner/agent.py

@@ -73,7 +73,7 @@ def __init__(
         max_rollouts: Optional[int] = None,
         poll_interval: float = 5.0,
         heartbeat_interval: float = 10.0,
-        interval_jitter: float = 0.1,
+        interval_jitter: float = 0.5,
         heartbeat_launch_mode: Literal["asyncio", "thread"] = "asyncio",
     ) -> None:
         """Initialize the agent runner.
@@ -577,16 +577,6 @@ async def iter(self, *, event: Optional[ExecutionEvent] = None) -> None:
                 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)
 
@@ -640,12 +630,8 @@ async def step(
         else:
             resources_id = None
 
-        attempted_rollout = await self.get_store().start_rollout(input=input, mode=mode, resources_id=resources_id)
-        # Register the attempt as running by the current worker
-        await self.get_store().update_attempt(
-            attempted_rollout.rollout_id,
-            attempted_rollout.attempt.attempt_id,
-            worker_id=self.get_worker_id(),
+        attempted_rollout = await self.get_store().start_rollout(
+            input=input, mode=mode, resources_id=resources_id, worker_id=self.get_worker_id()
         )
         rollout_id = await self._step_impl(attempted_rollout, raise_on_exception=True)
 

agentlightning/store/base.py

@@ -10,6 +10,7 @@
     Attempt,
     AttemptedRollout,
     AttemptStatus,
+    EnqueueRolloutRequest,
     NamedResources,
     ResourcesUpdate,
     Rollout,
@@ -100,19 +101,6 @@ class LightningStoreStatistics(TypedDict, total=False):
     """Memory capacity of the store in bytes."""
 
 
-class _EnqueueRolloutRequestRequired(TypedDict):
-    input: TaskInput
-
-
-class EnqueueRolloutRequest(_EnqueueRolloutRequestRequired, total=False):
-    """Payload describing a rollout to be queued via `enqueue_rollout`."""
-
-    mode: Optional[RolloutMode]
-    resources_id: Optional[str]
-    config: Optional[RolloutConfig]
-    metadata: Optional[Dict[str, Any]]
-
-
 class LightningStore:
     """Contract for the persistent control-plane that coordinates training rollouts.
 
@@ -174,6 +162,7 @@ async def start_rollout(
         resources_id: str | None = None,
         config: RolloutConfig | None = None,
         metadata: Dict[str, Any] | None = None,
+        worker_id: str | None = None,
     ) -> AttemptedRollout:
         """Register a rollout and immediately create its first attempt.
 
@@ -196,6 +185,7 @@ async def start_rollout(
             resources_id: Concrete resource snapshot to execute against; defaults to the latest stored snapshot.
             config: Rollout retry/timeout policy. Should default to a fresh [`RolloutConfig`][agentlightning.RolloutConfig].
             metadata: Free-form metadata persisted verbatim with the rollout.
+            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The fully-populated [`AttemptedRollout`][agentlightning.AttemptedRollout] including
@@ -241,19 +231,19 @@ async def enqueue_rollout(
         """
         raise NotImplementedError()
 
-    async def enqueue_many_rollouts(self, inputs: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
         """Persist multiple rollouts in `queuing` state.
 
         The implementation can delegate to [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]
         per request and preserves the input ordering. Subclasses can override to provide
         more efficient bulk enqueue semantics.
 
         Args:
-            inputs: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
+            rollouts: Rollout submission payloads mirroring [`enqueue_rollout()`][agentlightning.LightningStore.enqueue_rollout]'s
                 parameters. Each entry requires `input` and can optionally include other fields.
 
         Returns:
-            Rollouts enqueued in the same order as `inputs`.
+            Rollouts enqueued in the same order as `rollouts`.
         """
         raise NotImplementedError()
 
@@ -273,6 +263,9 @@ async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[Att
         * Optionally refresh the caller's [`Worker`][agentlightning.Worker] telemetry
           (e.g., `last_dequeue_time`) when `worker_id` is provided.
 
+        Args:
+            worker_id: Optional worker identifier to associate the claimed attempt with.
+
         Returns:
             The next attempt to execute, or `None` when no eligible rollouts are queued.
 
@@ -304,7 +297,7 @@ async def dequeue_many_rollouts(
         """
         raise NotImplementedError()
 
-    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
         """Create a manual retry attempt for an existing rollout.
 
         This is typically invoked by runners that wish to retry outside of the
@@ -315,6 +308,7 @@ async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
 
         Args:
             rollout_id: Unique identifier of the rollout receiving a new attempt.
+            worker_id: Optional worker identifier to associate the new attempt with.
 
         Returns:
             The rollout paired with its newly-created attempt.