gahow/ali-trace-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(formatter): align release hash blocks to provider input length

Browse Source 
Local tokenization does not always agree with the input_length the provider
reports, which left the release row carrying a block count derived from our own
count. Truncate or pad digest_hexes to the block count implied by the provider
value. Padding digests are derived per request, so a padded block can never
collide with another request and be scored as shared prefix reuse.
This commit is contained in:
Gahow Wang
2026-09-11 06:30:57 +00:00
parent 5ee98f9e4d
commit 03d4828818
2 changed files with 117 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
tests/test_ali_trace_pipeline.py
View File

@@ -817,6 +817,74 @@ class AliTracePipelineTest(unittest.TestCase):
self.assertEqual(Path(captured["dir"]), release_output_dir)
def test_export_release_ready_trace_aligns_hash_blocks_to_provider_input_length(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
raw_input_path = root / "trace-raw.jsonl"
release_output_path = root / "trace.jsonl"
rows = [
{
"schema_version": "2026.04.21",
"sort_time_ms": utc_ms("2026-04-17 15:00:01.000"),
"meta": {
"request_id": "req-pad",
"model_family": "glm5",
"request_ready_time_ms": utc_ms("2026-04-17 15:00:01.000"),
"chat_id": 0,
"parent_chat_id": -1,
"turn": 1,
},
"canonical_prompt": "hello",
"usage": {"input_tokens": 17, "output_tokens": 1},
},
{
"schema_version": "2026.04.21",
"sort_time_ms": utc_ms("2026-04-17 15:00:02.000"),
"meta": {
"request_id": "req-trim",
"model_family": "glm5",
"request_ready_time_ms": utc_ms("2026-04-17 15:00:02.000"),
"chat_id": 1,
"parent_chat_id": -1,
"turn": 1,
},
"canonical_prompt": " ".join(str(index) for index in range(100)),
"usage": {"input_tokens": 1, "output_tokens": 1},
},
{
"schema_version": "2026.04.21",
"sort_time_ms": utc_ms("2026-04-17 15:00:03.000"),
"meta": {
"request_id": "req-zero",
"model_family": "glm5",
"request_ready_time_ms": utc_ms("2026-04-17 15:00:03.000"),
"chat_id": 2,
"parent_chat_id": -1,
"turn": 1,
},
"canonical_prompt": "hello",
"usage": {"input_tokens": 0, "output_tokens": 1},
},
]
raw_input_path.write_text(
"".join(json.dumps(row) + "\n" for row in rows),
encoding="utf-8",
)
export_release_ready_trace(
raw_input_path=raw_input_path,
release_output_path=release_output_path,
jobs=1,
block_size=8,
)
release_rows = [json.loads(line) for line in release_output_path.read_text(encoding="utf-8").splitlines()]
self.assertEqual([len(row["hash_ids"]) for row in release_rows], [3, 1, 0])
for row in release_rows:
input_length = row["input_length"]
expected_blocks = 0 if input_length <= 0 else (input_length + 7) // 8
self.assertEqual(len(row["hash_ids"]), expected_blocks)
def test_format_and_sort_trace_infers_window_in_ready_time_scale_when_wall_clock_has_timezone_offset(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)

49
trace_formatter/formatting.py
View File

@@ -539,6 +539,45 @@ def _block_digest(block: list[int]) -> bytes:
return digest.digest()
def _prompt_block_count_for_input_length(input_length: int, block_size: int) -> int:
if input_length <= 0:
return 0
return (input_length + block_size - 1) // block_size
def _padding_block_digest_hex(*, request_identity: str, input_length: int, block_index: int) -> str:
digest = hashlib.blake2b(digest_size=16)
digest.update(b"trace_formatter.provider_input_length_padding\x00")
digest.update(request_identity.encode("utf-8", "ignore"))
digest.update(b"\x00")
digest.update(str(input_length).encode("ascii"))
digest.update(b"\x00")
digest.update(str(block_index).encode("ascii"))
return digest.hexdigest()
def _align_digest_hexes_to_provider_input_length(
digest_hexes: list[str],
*,
input_length: int,
block_size: int,
request_identity: str,
) -> list[str]:
target_block_count = _prompt_block_count_for_input_length(input_length, block_size)
if len(digest_hexes) >= target_block_count:
return digest_hexes[:target_block_count]
aligned = list(digest_hexes)
for block_index in range(len(aligned), target_block_count):
aligned.append(
_padding_block_digest_hex(
request_identity=request_identity,
input_length=input_length,
block_index=block_index,
)
)
return aligned
def _load_release_tokenizer(model_family: str) -> Tokenizer:
resolved = Path(resolve_tokenizer_path(model_family=model_family))
tokenizer_file = resolved / "tokenizer.json" if resolved.is_dir() else resolved
@@ -637,6 +676,16 @@ def _build_release_shard(
digest_hexes.append(_block_digest(block).hex())
core = _compute_release_row_core(row, base_ms=base_ms)
request_identity = str(
meta.get("request_id")
or f"{core.get('chat_id')}:{core.get('timestamp')}:{core.get('turn')}:{shard_index}:{row_count}"
)
digest_hexes = _align_digest_hexes_to_provider_input_length(
digest_hexes,
input_length=safe_int(core.get("input_length")),
block_size=block_size,
request_identity=request_identity,
)
destination.write(json.dumps(core, ensure_ascii=False, separators=(",", ":")))
destination.write("\t")
destination.write(",".join(digest_hexes))