visible logs

agentlightning/litagent.py

@@ -94,32 +94,6 @@ class LitAgent:
             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.
 

agentlightning/logging.py

@@ -1,4 +1,16 @@
+from __future__ import annotations
+
+import json
 import logging
+from typing import Optional, TYPE_CHECKING
+import numpy as np
+
+from .types import Hook
+
+if TYPE_CHECKING:
+    from .types import Task, Rollout
+    from .tracer import BaseTracer
+    from .runner import AgentRunner
 
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
@@ -14,3 +26,192 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
     logger.setLevel(level)
     logger.propagate = False  # prevent double logging
     return logger
+
+
+logger = logging.getLogger(__name__)
+
+
+class LightningLogger(Hook):
+    """Agent-lightning logger that supports tracing events and metrics throughout the training."""
+
+    def log_event(self, event: str, data: dict):
+        """
+        Log an event with its associated data when something happens.
+        """
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        """
+        Log a metric with its value and an optional step.
+        """
+
+    def log_message(self, level: int, message: str):
+        """
+        Log a message at a specific logging level.
+        """
+
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer):
+        """
+        By default, each logger automatically logs the rollout event at the end of each rollout.
+        """
+        self.log_event("rollout", {"task": task.model_dump(), "rollout": rollout.model_dump()})
+
+
+class ConsoleLogger(LightningLogger):
+    """A simple logger that logs messages to the console using Python's logging module."""
+
+    def __init__(self, level: int = logging.INFO):
+        self.logger = configure_logger(level, name="agentlightning.ConsoleLogger")
+        self.default_level = level
+        self.worker_id: Optional[int] = None
+
+    def init_worker(self, worker_id: int):
+        super().init_worker(worker_id)
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int):
+        super().teardown_worker(worker_id)
+        self.worker_id = None
+
+    def log_event(self, event: str, data: dict):
+        data_str = str(data)
+        if len(data_str) > 512:
+            data_str = f"{data_str[:512]}... (truncated)"
+        message = f"Event: {event}, Data: {data_str}"
+        self.log_message(self.default_level, message)
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        step_str = f" at step {step}" if step is not None else ""
+        message = f"Metric: {metric} = {value}{step_str}"
+        self.log_message(self.default_level, message)
+
+    def log_message(self, level: int, message: str):
+        if level >= self.default_level:
+            if self.worker_id is not None:
+                message = f"(Worker-{self.worker_id}) {message}"
+            else:
+                message = f"(Main) {message}"
+            self.logger.log(level, message)
+        # else skip logging if below default level
+
+
+class WandbLogger(LightningLogger):
+
+    def __init__(
+        self,
+        project: str,
+        entity: Optional[str] = None,
+        name: Optional[str] = None,
+        config: Optional[dict] = None,
+        *,
+        flush_every_n_events: int = 128,
+        aggregate_every_n_metrics: int = 128,
+    ):
+        import wandb
+        from wandb.sdk.wandb_run import Run
+
+        self.wandb_run: Optional[Run] = None
+        self.wandb_run_id: Optional[str] = None
+
+        self.project = project
+        self.entity = entity
+        self.name = name or wandb.util.generate_id()
+        self.config = config or {}
+
+        self.event_table: Optional[wandb.Table] = None
+        self.flush_every_n_events = flush_every_n_events
+        self.aggregate_every_n_metrics = aggregate_every_n_metrics
+
+        self.metrics_buffer: dict[str, list[float]] = {}
+
+    def init_worker(self, worker_id: int):
+        import wandb
+
+        super().init_worker(worker_id)
+        self.wandb_run = wandb.init(
+            project=self.project,
+            entity=self.entity,
+            group=self.name,
+            job_type=f"worker_{worker_id}",
+            config=self.config,
+        )
+        logger.info(f"Wandb run initialized: {self.name} (Worker {worker_id})")
+        if self.wandb_run is None:
+            raise RuntimeError("Failed to initialize Wandb run.")
+        self.wandb_run_id = self.wandb_run.id
+
+    def teardown_worker(self, worker_id: int):
+        import wandb
+
+        super().teardown_worker(worker_id)
+
+        for metric in self.metrics_buffer:
+            if len(self.metrics_buffer[metric]) > 0:
+                self._log_aggregated_metrics(metric)
+
+        if len(self.event_table.data) > 0:
+            logger.info(f"Flushing {len(self.event_table.data)} events to Wandb before finishing...")
+            wandb.log({"client/events": self.event_table})
+            self.event_table = None
+
+        wandb.finish(exit_code=0)
+
+    def teardown(self):
+        import wandb
+
+        super().teardown()
+        if self.wandb_run is not None:
+            wandb.finish(exit_code=0)
+            self.wandb_run = None
+            self.wandb_run_id = None
+
+    def log_event(self, event: str, data: dict):
+        import wandb
+
+        if self.event_table is None:
+            self.event_table = wandb.Table(columns=["event", "data"])
+
+        try:
+            data_str = json.dumps(data)  # Ensure data is JSON serializable
+        except (TypeError, ValueError):
+            data_str = str(data)
+        self.event_table.add_data(event, data_str)
+
+        if len(self.event_table.data) % self.flush_every_n_events == 0:
+            logger.info(f"Flushing {len(self.event_table.data)} events to Wandb...")
+            wandb.log({"client/events": self.event_table})
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        import wandb
+
+        if step is not None:
+            wandb.log({"client_metric/" + metric: value}, step=step)
+        else:
+            wandb.log({"client_metric/" + metric: value})
+
+        if metric not in self.metrics_buffer:
+            self.metrics_buffer[metric] = []
+        self.metrics_buffer[metric].append(value)
+        if len(self.metrics_buffer[metric]) >= self.aggregate_every_n_metrics:
+            self._log_aggregated_metrics(metric, step)
+            self.metrics_buffer[metric] = []
+
+    def log_message(self, level: int, message: str):
+        pass  # Wandb handles logging internally, so we don't need to implement this
+
+    def _log_aggregated_metrics(self, metric, step: Optional[int] = None):
+        import wandb
+
+        arr = np.array(self.metrics_buffer[metric])
+        aggregated_value = {
+            "mean": float(np.mean(arr)),
+            "max": float(np.max(arr)),
+            "min": float(np.min(arr)),
+            "std": float(np.std(arr)),
+            "count": int((~np.isnan(arr)).sum()),
+        }
+        for key, value in aggregated_value.items():
+            if value is not None:
+                if step is not None:
+                    wandb.log({"client_agg/" + metric + "/" + key: value}, step=step)
+                else:
+                    wandb.log({"client_agg/" + metric + "/" + key: value})

agentlightning/runner.py

@@ -7,12 +7,11 @@ from contextlib import nullcontext
 from typing import List, Optional, Union, Dict, Any
 
 import agentops
-
 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 .types import Rollout, Task, Triplet, RolloutRawResult, ParallelWorkerBase, Hook
 from .tracer.base import BaseTracer
 from .tracer import TripletExporter
 
@@ -43,12 +42,14 @@ class AgentRunner(ParallelWorkerBase):
         triplet_exporter: TripletExporter,
         worker_id: Optional[int] = None,
         max_tasks: Optional[int] = None,
+        hooks: Optional[List[Hook]] = None,
     ):
         super().__init__()
         self.agent = agent
         self.client = client
         self.tracer = tracer
         self.triplet_exporter = triplet_exporter
+        self.hooks = hooks or []
 
         # Worker-specific attributes
         self.worker_id = worker_id
@@ -158,10 +159,11 @@ class AgentRunner(ParallelWorkerBase):
         rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
 
         try:
-            try:
-                self.agent.on_rollout_start(task, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_start(task, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook: {hook}.")
 
             with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                 start_time = time.time()
@@ -180,10 +182,11 @@ class AgentRunner(ParallelWorkerBase):
         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)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_end(task, rollout_obj, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook: {hook}.")
             self.client.post_rollout(rollout_obj)
 
         return True
@@ -227,10 +230,11 @@ class AgentRunner(ParallelWorkerBase):
         rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
 
         try:
-            try:
-                self.agent.on_rollout_start(task, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_start(task, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook: {hook}.")
 
             with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                 start_time = time.time()
@@ -248,10 +252,11 @@ class AgentRunner(ParallelWorkerBase):
         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)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_end(task, rollout_obj, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook: {hook}.")
             await self.client.post_rollout_async(rollout_obj)
 
         return True

agentlightning/trainer.py

@@ -1,16 +1,14 @@
 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 .logging import LightningLogger
 from .runner import AgentRunner
 from .types import ParallelWorkerBase
 from .tracer.base import BaseTracer
@@ -52,6 +50,7 @@ class Trainer(ParallelWorkerBase):
         daemon: bool = True,
         tracer: Union[BaseTracer, str, dict, None] = None,
         triplet_exporter: Union[TripletExporter, dict, None] = None,
+        loggers: Optional[List[LightningLogger]] = None,
     ):
         super().__init__()
         self.n_workers = n_workers
@@ -60,6 +59,8 @@ class Trainer(ParallelWorkerBase):
         self.dev = dev
         self._client: AgentLightningClient | None = None  # Will be initialized in fit method
 
+        self._loggers = loggers or []
+
         self.tracer = self._make_tracer(tracer)
         if isinstance(triplet_exporter, TripletExporter):
             self.triplet_exporter = triplet_exporter
@@ -109,12 +110,18 @@ class Trainer(ParallelWorkerBase):
 
         self.tracer.init()
 
+        for _logger in self.loggers:
+            _logger.init()
+
         logger.info(f"Trainer main initialization complete.")
 
     def teardown(self) -> None:
         logger.info(f"Cleaning up Trainer...")
         self.tracer.teardown()
 
+        for _logger in self.loggers:
+            _logger.teardown()
+
         self._client = None
         logger.info(f"Trainer main cleanup complete.")
 
@@ -187,6 +194,8 @@ class Trainer(ParallelWorkerBase):
                 num_processed = asyncio.run(loop.iter_async())
             else:
                 num_processed = loop.iter()
+        except KeyboardInterrupt:
+            logger.info(f"[Worker {worker_id}] KeyboardInterrupt received. Exiting worker loop...")
         except Exception:
             logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
         finally:
@@ -197,11 +206,21 @@ class Trainer(ParallelWorkerBase):
     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)
+        for _logger in self.loggers:
+            _logger.init_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Worker environment setup complete.")
 
     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.")
+        for _logger in self.loggers:
+            _logger.teardown_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Worker environments torn down.")
+
+    @property
+    def loggers(self) -> List[LightningLogger]:
+        return self._loggers
 
     @staticmethod
     def kill_orphaned_processes() -> None:

agentlightning/types.py

@@ -1,8 +1,15 @@
-from typing import Any, Dict, List, Optional, Union, Literal, Annotated
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union, Literal, Annotated, TYPE_CHECKING
 
 from pydantic import BaseModel, Field, Discriminator
 from opentelemetry.sdk.trace import ReadableSpan
 
+if TYPE_CHECKING:
+    from .runner import AgentRunner
+    from .tracer import BaseTracer
+
+
 __all__ = [
     "Triplet",
     "Rollout",
@@ -18,6 +25,7 @@ __all__ = [
     "ResourcesUpdate",
     "GenericResponse",
     "ParallelWorkerBase",
+    "Hook",
 ]
 
 
@@ -202,3 +210,33 @@ class ParallelWorkerBase:
 
     def teardown(self, *args: Any, **kwargs: Any) -> None:
         pass
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    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.
+        """

examples/apo/apo_client.py

@@ -35,5 +35,10 @@ if __name__ == "__main__":
     configure_logger()
     dotenv.load_dotenv()
     agent = SimpleAgent()
-    trainer = Trainer(n_workers=2)
+    from agentlightning.logging import ConsoleLogger, WandbLogger
+    loggers = [
+        ConsoleLogger(),
+        WandbLogger(project="agent-lightning-debug"),
+    ]
+    trainer = Trainer(n_workers=2, loggers=loggers)
     trainer.fit(agent, backend="http://127.0.0.1:9997")

pyproject.toml

@@ -14,6 +14,9 @@ dependencies = [
   "uvicorn",
   "fastapi",
   "aiohttp",
+  "opentelemetry-api",
+  "opentelemetry-sdk",
+  "numpy",
 ]
 
 [project.optional-dependencies]
@@ -49,6 +52,7 @@ agent = [
   "sqlparse",
   "nltk",
   "uv",
+  "wandb",
 ]
 
 [build-system]

scripts/validate_example_wandb.py

@@ -3,6 +3,7 @@ import sys
 
 if len(sys.argv) != 3:
     print("Usage: python validate_example_wandb.py <project> <run_name>")
+    sys.exit(1)
 
 project = sys.argv[1]
 run_name = sys.argv[2]

tests/test_litagent_hooks.py

@@ -1,8 +1,7 @@
 import pytest
 from contextlib import contextmanager
 
-
-from agentlightning import LitAgent, Task, ResourcesUpdate
+from agentlightning import Hook, LitAgent, Task, ResourcesUpdate
 from agentlightning.runner import AgentRunner
 from agentlightning.tracer import BaseTracer, TripletExporter
 
@@ -53,11 +52,6 @@ class DummyAsyncClient:
 
 
 class HookAgent(LitAgent):
-    def __init__(self):
-        super().__init__()
-        self.start_called = False
-        self.end_called = False
-        self.end_rollout = None
 
     def training_rollout(self, task, rollout_id, resources):
         return 0.5
@@ -65,6 +59,9 @@ class HookAgent(LitAgent):
     async def training_rollout_async(self, task, rollout_id, resources):
         return 0.5
 
+
+class TestHook(Hook):
+
     def on_rollout_start(self, task, runner, tracer):
         self.start_called = True
         self.start_task = task
@@ -78,12 +75,13 @@ def test_runner_calls_hooks():
     agent = HookAgent()
     client = DummyClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    hook = TestHook()
+    runner = AgentRunner(agent, client, tracer, TripletExporter(), hooks=[hook])
 
     assert runner.run() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert hook.start_called
+    assert hook.end_called
+    assert hook.end_rollout.final_reward == 0.5
 
 
 @pytest.mark.asyncio
@@ -91,9 +89,10 @@ async def test_runner_calls_hooks_async():
     agent = HookAgent()
     client = DummyAsyncClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    hook = TestHook()
+    runner = AgentRunner(agent, client, tracer, TripletExporter(), hooks=[hook])
 
     assert await runner.run_async() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert hook.start_called
+    assert hook.end_called
+    assert hook.end_rollout.final_reward == 0.5