diff --git a/tests/test_ali_trace_pipeline.py b/tests/test_ali_trace_pipeline.py index a484af4..588d8de 100644 --- a/tests/test_ali_trace_pipeline.py +++ b/tests/test_ali_trace_pipeline.py @@ -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) diff --git a/trace_formatter/formatting.py b/trace_formatter/formatting.py index ca9e4dd..0f87993 100644 --- a/trace_formatter/formatting.py +++ b/trace_formatter/formatting.py @@ -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))