feat: add agentic pd hybrid benchmark prototype
This commit is contained in:
295
src/agentic_pd_hybrid/sampling.py
Normal file
295
src/agentic_pd_hybrid/sampling.py
Normal file
@@ -0,0 +1,295 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from agentic_pd_hybrid.trace import TraceRequest, load_trace
|
||||
|
||||
|
||||
SampleProfile = Literal["default", "small-append"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionSampleConfig:
|
||||
trace_path: Path
|
||||
output_path: Path
|
||||
target_duration_s: float = 600.0
|
||||
start_time_s: float = 0.0
|
||||
session_sample_rate: float = 1.0
|
||||
min_turns: int = 1
|
||||
max_requests: int | None = None
|
||||
profile: SampleProfile = "default"
|
||||
min_initial_input_tokens: int | None = None
|
||||
max_initial_input_tokens: int | None = None
|
||||
max_append_input_tokens: int | None = None
|
||||
max_output_tokens: int | None = None
|
||||
min_overlap_ratio: float | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SessionSampleSummary:
|
||||
input_trace_path: str
|
||||
output_trace_path: str
|
||||
request_count: int
|
||||
session_count: int
|
||||
multi_turn_session_count: int
|
||||
start_time_s: float
|
||||
end_time_s: float
|
||||
sampled_duration_s: float
|
||||
session_sample_rate: float
|
||||
min_turns: int
|
||||
profile: str
|
||||
min_initial_input_tokens: int | None
|
||||
max_initial_input_tokens: int | None
|
||||
max_append_input_tokens: int | None
|
||||
max_output_tokens: int | None
|
||||
min_overlap_ratio: float | None
|
||||
mean_append_input_tokens: float | None
|
||||
mean_turn_overlap_ratio: float | None
|
||||
|
||||
|
||||
def sample_trace_sessions(config: SessionSampleConfig) -> SessionSampleSummary:
|
||||
requests = load_trace(config.trace_path)
|
||||
sessions: dict[str, list[TraceRequest]] = defaultdict(list)
|
||||
for request in requests:
|
||||
sessions[request.session_id].append(request)
|
||||
|
||||
filters = _resolve_filters(config)
|
||||
eligible_sessions = {
|
||||
session_id: session_requests
|
||||
for session_id, session_requests in sessions.items()
|
||||
if len(session_requests) >= filters.min_turns
|
||||
and _session_matches_filters(session_requests, filters)
|
||||
and _keep_session(session_id, config.session_sample_rate)
|
||||
}
|
||||
ordered_sessions = sorted(
|
||||
eligible_sessions.values(),
|
||||
key=lambda session_requests: session_requests[0].timestamp_s,
|
||||
)
|
||||
|
||||
selected_requests: list[TraceRequest] = []
|
||||
sampled_start: float | None = None
|
||||
sampled_end: float | None = None
|
||||
for session_requests in ordered_sessions:
|
||||
session_first = session_requests[0].timestamp_s
|
||||
if session_first < config.start_time_s:
|
||||
continue
|
||||
|
||||
if sampled_start is None:
|
||||
sampled_start = session_first
|
||||
|
||||
selected_requests.extend(session_requests)
|
||||
sampled_end = max(request.timestamp_s for request in session_requests)
|
||||
|
||||
if config.max_requests is not None and len(selected_requests) >= config.max_requests:
|
||||
break
|
||||
if sampled_end - sampled_start >= config.target_duration_s:
|
||||
break
|
||||
|
||||
selected_requests.sort(key=lambda request: request.timestamp_s)
|
||||
if config.max_requests is not None:
|
||||
selected_requests = selected_requests[: config.max_requests]
|
||||
|
||||
if not selected_requests:
|
||||
raise ValueError("Sampling produced no requests; adjust the sampling arguments")
|
||||
|
||||
config.output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with config.output_path.open("w", encoding="utf-8") as handle:
|
||||
for request in selected_requests:
|
||||
payload = {
|
||||
"request_id": request.request_id,
|
||||
"session_id": request.session_id,
|
||||
"chat_id": request.chat_id,
|
||||
"parent_chat_id": request.parent_chat_id,
|
||||
"timestamp": request.timestamp_s,
|
||||
"input_length": request.input_length,
|
||||
"output_length": request.output_length,
|
||||
"type": request.request_type,
|
||||
"turn": request.turn_id,
|
||||
"hash_ids": list(request.hash_ids),
|
||||
}
|
||||
handle.write(json.dumps(payload, sort_keys=True) + "\n")
|
||||
|
||||
selected_session_ids = {request.session_id for request in selected_requests}
|
||||
selected_session_requests = [
|
||||
eligible_sessions[session_id] for session_id in selected_session_ids
|
||||
]
|
||||
append_lengths = [
|
||||
length
|
||||
for session_requests in selected_session_requests
|
||||
for length in _turn_append_lengths(session_requests)
|
||||
]
|
||||
overlap_ratios = [
|
||||
ratio
|
||||
for session_requests in selected_session_requests
|
||||
for ratio in _turn_overlap_ratios(session_requests)
|
||||
]
|
||||
summary = SessionSampleSummary(
|
||||
input_trace_path=str(config.trace_path),
|
||||
output_trace_path=str(config.output_path),
|
||||
request_count=len(selected_requests),
|
||||
session_count=len(selected_session_ids),
|
||||
multi_turn_session_count=sum(
|
||||
1
|
||||
for session_id in selected_session_ids
|
||||
if len(eligible_sessions[session_id]) > 1
|
||||
),
|
||||
start_time_s=selected_requests[0].timestamp_s,
|
||||
end_time_s=selected_requests[-1].timestamp_s,
|
||||
sampled_duration_s=selected_requests[-1].timestamp_s
|
||||
- selected_requests[0].timestamp_s,
|
||||
session_sample_rate=config.session_sample_rate,
|
||||
min_turns=filters.min_turns,
|
||||
profile=config.profile,
|
||||
min_initial_input_tokens=filters.min_initial_input_tokens,
|
||||
max_initial_input_tokens=filters.max_initial_input_tokens,
|
||||
max_append_input_tokens=filters.max_append_input_tokens,
|
||||
max_output_tokens=filters.max_output_tokens,
|
||||
min_overlap_ratio=filters.min_overlap_ratio,
|
||||
mean_append_input_tokens=_mean(append_lengths),
|
||||
mean_turn_overlap_ratio=_mean(overlap_ratios),
|
||||
)
|
||||
summary_path = config.output_path.with_suffix(config.output_path.suffix + ".summary.json")
|
||||
with summary_path.open("w", encoding="utf-8") as handle:
|
||||
json.dump(asdict(summary), handle, indent=2, sort_keys=True)
|
||||
return summary
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ResolvedFilters:
|
||||
min_turns: int
|
||||
min_initial_input_tokens: int | None
|
||||
max_initial_input_tokens: int | None
|
||||
max_append_input_tokens: int | None
|
||||
max_output_tokens: int | None
|
||||
min_overlap_ratio: float | None
|
||||
|
||||
|
||||
def _resolve_filters(config: SessionSampleConfig) -> _ResolvedFilters:
|
||||
if config.profile == "default":
|
||||
return _ResolvedFilters(
|
||||
min_turns=config.min_turns,
|
||||
min_initial_input_tokens=config.min_initial_input_tokens,
|
||||
max_initial_input_tokens=config.max_initial_input_tokens,
|
||||
max_append_input_tokens=config.max_append_input_tokens,
|
||||
max_output_tokens=config.max_output_tokens,
|
||||
min_overlap_ratio=config.min_overlap_ratio,
|
||||
)
|
||||
|
||||
if config.profile != "small-append":
|
||||
raise ValueError(f"Unsupported sample profile: {config.profile}")
|
||||
|
||||
return _ResolvedFilters(
|
||||
min_turns=max(config.min_turns, 2),
|
||||
min_initial_input_tokens=(
|
||||
2048
|
||||
if config.min_initial_input_tokens is None
|
||||
else config.min_initial_input_tokens
|
||||
),
|
||||
max_initial_input_tokens=(
|
||||
16000
|
||||
if config.max_initial_input_tokens is None
|
||||
else config.max_initial_input_tokens
|
||||
),
|
||||
max_append_input_tokens=(
|
||||
2048
|
||||
if config.max_append_input_tokens is None
|
||||
else config.max_append_input_tokens
|
||||
),
|
||||
max_output_tokens=(
|
||||
2048 if config.max_output_tokens is None else config.max_output_tokens
|
||||
),
|
||||
min_overlap_ratio=(
|
||||
0.75 if config.min_overlap_ratio is None else config.min_overlap_ratio
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _session_matches_filters(
|
||||
session_requests: list[TraceRequest],
|
||||
filters: _ResolvedFilters,
|
||||
) -> bool:
|
||||
ordered = sorted(
|
||||
session_requests,
|
||||
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
||||
)
|
||||
if not ordered:
|
||||
return False
|
||||
|
||||
initial = ordered[0]
|
||||
if (
|
||||
filters.min_initial_input_tokens is not None
|
||||
and initial.input_length < filters.min_initial_input_tokens
|
||||
):
|
||||
return False
|
||||
if (
|
||||
filters.max_initial_input_tokens is not None
|
||||
and initial.input_length > filters.max_initial_input_tokens
|
||||
):
|
||||
return False
|
||||
if filters.max_output_tokens is not None and any(
|
||||
request.output_length > filters.max_output_tokens for request in ordered
|
||||
):
|
||||
return False
|
||||
|
||||
append_lengths = _turn_append_lengths(ordered)
|
||||
if filters.max_append_input_tokens is not None and any(
|
||||
append_length <= 0 or append_length > filters.max_append_input_tokens
|
||||
for append_length in append_lengths
|
||||
):
|
||||
return False
|
||||
|
||||
overlap_ratios = _turn_overlap_ratios(ordered)
|
||||
if filters.min_overlap_ratio is not None and any(
|
||||
overlap_ratio < filters.min_overlap_ratio for overlap_ratio in overlap_ratios
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _turn_append_lengths(session_requests: list[TraceRequest]) -> list[int]:
|
||||
ordered = sorted(
|
||||
session_requests,
|
||||
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
||||
)
|
||||
return [
|
||||
current.input_length - (previous.input_length + previous.output_length)
|
||||
for previous, current in zip(ordered, ordered[1:], strict=False)
|
||||
]
|
||||
|
||||
|
||||
def _turn_overlap_ratios(session_requests: list[TraceRequest]) -> list[float]:
|
||||
ordered = sorted(
|
||||
session_requests,
|
||||
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
||||
)
|
||||
ratios: list[float] = []
|
||||
for previous, current in zip(ordered, ordered[1:], strict=False):
|
||||
if not current.hash_ids:
|
||||
ratios.append(0.0)
|
||||
continue
|
||||
previous_blocks = set(previous.hash_ids)
|
||||
overlap = sum(1 for block in current.hash_ids if block in previous_blocks)
|
||||
ratios.append(overlap / len(current.hash_ids))
|
||||
return ratios
|
||||
|
||||
|
||||
def _mean(values: list[int] | list[float]) -> float | None:
|
||||
if not values:
|
||||
return None
|
||||
return sum(values) / len(values)
|
||||
|
||||
|
||||
def _keep_session(session_id: str, sample_rate: float) -> bool:
|
||||
if sample_rate >= 1.0:
|
||||
return True
|
||||
if sample_rate <= 0.0:
|
||||
return False
|
||||
digest = hashlib.blake2b(session_id.encode("utf-8"), digest_size=8).digest()
|
||||
bucket = int.from_bytes(digest, byteorder="big", signed=False) / 2**64
|
||||
return bucket < sample_rate
|
||||
Reference in New Issue
Block a user