Spaces:
Sleeping
Sleeping
File size: 5,751 Bytes
116524e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """ACERunner — shared runner infrastructure for all ACE runners."""
from __future__ import annotations
import logging
from collections.abc import Iterable, Sequence
from typing import Any
from pipeline import Pipeline
from pipeline.errors import CancellationToken
from pipeline.protocol import SampleResult
from ..core.context import ACEStepContext, SkillbookView
from ..core.skillbook import Skillbook
logger = logging.getLogger(__name__)
class ACERunner:
"""Shared runner infrastructure for all ACE runners.
Composes a ``Pipeline`` (does not extend it). Manages the epoch loop
and delegates per-sample iteration, error isolation, foreground/background
split, and concurrent workers to ``Pipeline.run()``.
Subclasses override two methods:
- ``run()`` — public API with a subclass-specific signature.
- ``_build_context()`` — maps a single input item to ``ACEStepContext``.
You can also construct an ``ACERunner`` directly with a hand-composed
pipeline::
from ace import Pipeline, ACERunner, AgentStep, learning_tail
pipe = Pipeline([AgentStep(agent, sb), *learning_tail(reflector, sm, sb)])
runner = ACERunner(pipeline=pipe, skillbook=sb)
Attributes:
pipeline: The composed ``Pipeline`` instance. Accessible for
inspection after construction.
skillbook: The ``Skillbook`` this runner operates on.
"""
def __init__(
self,
pipeline: Pipeline,
skillbook: Skillbook,
) -> None:
self.pipeline = pipeline
self.skillbook = skillbook
# ------------------------------------------------------------------
# Lifecycle helpers
# ------------------------------------------------------------------
def save(self, path: str) -> None:
"""Save the current skillbook to disk."""
self.skillbook.save_to_file(path)
def load(self, path: str) -> None:
"""Load a skillbook from disk, replacing the current one."""
self.skillbook = Skillbook.load_from_file(path)
def wait_for_background(self, timeout: float | None = None) -> None:
"""Block until all background learning tasks complete.
Delegates to ``Pipeline.wait_for_background()``. Call after
``run(wait=False)`` before saving the skillbook or reading final
results.
"""
self.pipeline.wait_for_background(timeout)
@property
def learning_stats(self) -> dict[str, int]:
"""Return background learning progress.
Delegates to ``Pipeline.background_stats()``.
"""
return self.pipeline.background_stats()
# ------------------------------------------------------------------
# Generic epoch loop (called by subclasses)
# ------------------------------------------------------------------
def _run(
self,
items: Sequence[Any] | Iterable[Any],
*,
epochs: int,
wait: bool = True,
cancel_token: CancellationToken | None = None,
**kwargs: Any,
) -> list[SampleResult]:
"""Generic run loop handling epochs and Iterable validation.
Returns when ``wait=True`` (default). Returns after foreground
steps when ``wait=False`` — background learning continues.
Args:
cancel_token: Optional cancellation signal. Forwarded to
``Pipeline.run()`` — checked between steps and inside
LLM calls (via contextvar).
Raises ``ValueError`` if ``epochs > 1`` and *items* is not a
``Sequence``.
"""
if epochs > 1 and not isinstance(items, Sequence):
raise ValueError(
"Multi-epoch requires a Sequence, not a consumed Iterable."
)
results: list[SampleResult] = []
n: int | None = len(items) if isinstance(items, Sequence) else None
for epoch in range(1, epochs + 1):
if cancel_token is not None and cancel_token.is_cancelled:
break
logger.info(
"Epoch %d/%d: processing %s samples",
epoch,
epochs,
n if n is not None else "unknown",
)
contexts: list[ACEStepContext] = [
self._build_context(
item,
epoch=epoch,
total_epochs=epochs,
index=idx,
total=n,
global_sample_index=(
(epoch - 1) * n + idx if n is not None else idx
),
**kwargs,
)
for idx, item in enumerate(items, start=1)
]
epoch_results = self.pipeline.run(contexts, cancel_token=cancel_token)
results.extend(epoch_results)
if wait:
self.pipeline.wait_for_background()
return results
# ------------------------------------------------------------------
# Subclass interface
# ------------------------------------------------------------------
def _build_context(
self,
item: Any,
*,
epoch: int,
total_epochs: int,
index: int,
total: int | None,
global_sample_index: int,
**kwargs: Any,
) -> ACEStepContext:
"""Map a single input item to an ``ACEStepContext``.
Must be overridden by subclasses. Stateless — depends only on
the item and the provided counters.
"""
raise NotImplementedError
|