Skip to content

Commit 1cf627d

Browse files
rchardxBruce-rl-hw
authored andcommitted
feat: introduces session-centric tracing APIs (inclusionAI#539)
* Introduces session-centric tracing APIs Replaces request tracing with a session-scoped tracer so rollout executors, engines, and workflows can log lifecycle and phase spans in a single record. Adds context-aware helpers and decorators to propagate session IDs through async code, enabling phase timing, counter aggregation, and structured flush thresholds. Updates configs, tests, docs, and example presets to align with the session tracer and tightens validation around workflow submission and tracing hooks. Adds request phase tracing helpers (inclusionAI#543) * Resolve review comments
1 parent f828bc9 commit 1cf627d

File tree

22 files changed

+1442
-468
lines changed

22 files changed

+1442
-468
lines changed

areal/api/cli_args.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,23 +1001,23 @@ class StatsLoggerConfig:
10011001

10021002

10031003
@dataclass
1004-
class RequestTracerConfig:
1005-
"""Configuration for per-request lifecycle tracing."""
1004+
class SessionTracerConfig:
1005+
"""Configuration for per-session lifecycle tracing."""
10061006

10071007
enabled: bool = field(
10081008
default=False,
10091009
metadata={
10101010
"help": (
1011-
"Enable per-request lifecycle tracing alongside perf events. "
1012-
"When true, request metadata is captured to requests.jsonl."
1011+
"Enable per-session lifecycle tracing alongside perf events. "
1012+
"When true, session metadata is captured to sessions.jsonl."
10131013
)
10141014
},
10151015
)
10161016
flush_threshold: int = field(
10171017
default=256,
10181018
metadata={
10191019
"help": (
1020-
"Flush request trace records once this many entries are ready. "
1020+
"Flush session trace records once this many entries are ready. "
10211021
"Values <= 0 fall back to 1."
10221022
)
10231023
},
@@ -1048,9 +1048,9 @@ class PerfTracerConfig:
10481048
)
10491049
},
10501050
)
1051-
request_tracer: RequestTracerConfig | None = field(
1051+
session_tracer: SessionTracerConfig | None = field(
10521052
default=None,
1053-
metadata={"help": "Request tracing configuration."},
1053+
metadata={"help": "Session tracing configuration."},
10541054
)
10551055

10561056

areal/api/io_struct.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class ModelRequest:
3030
tokenizer: PreTrainedTokenizerFast | None = None
3131

3232
# vlm
33-
image_data: list[ImageObject | str] | None = field(default_factory=list)
33+
image_data: list[str] | None = field(default_factory=list)
3434
processor: Optional["AutoProcessor"] = None
3535

3636
def copy(self):

areal/api/workflow_api.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations # noqa
22

3+
from abc import ABC, abstractmethod
34
from typing import TYPE_CHECKING, Any
45

56
from areal.experimental.openai.types import InteractionWithTokenLogpReward
@@ -8,7 +9,8 @@
89
from areal.api.engine_api import InferenceEngine
910

1011

11-
class RolloutWorkflow:
12+
class RolloutWorkflow(ABC):
13+
@abstractmethod
1214
async def arun_episode(
1315
self, engine: InferenceEngine, data: dict[str, Any]
1416
) -> dict[str, Any] | None | dict[str, InteractionWithTokenLogpReward]:

0 commit comments

Comments
 (0)