Implementation:SqueezeAILab ETS Result Serialization

Knowledge Sources	ETS
Domains	Data_Management, Experiment_Logging
Last Updated	2026-02-14 02:00 GMT

Overview

Concrete tool for serializing tree search results and statistics to JSON and log files, provided by the main() function in rebase.py.

Description

After reward_guided_search.run_batch() completes, the main() function collects all return values, writes the combined results to answers.json via json.dump(), and records aggregate statistics (KV cache size, model calls, tokens generated, wall-clock time) to stats.log.

Usage

Automatically executed at the end of a tree search run. The output path is controlled by --output_path.

Code Reference

Source Location

• Repository: ETS
• File: rebase.py
• Lines: 748-783

Signature

# Result serialization (inside main())
json.dump(results, open(args.output_path, "w"), indent=4)

# Statistics logging
with open(log_file_path, 'w') as log_file:
    log_file.write(f'Total KV cache size: {global_kv_size}\n')
    log_file.write(f'Number of Model Calls: {global_num_model_calls}\n')
    log_file.write(f'Number of Tokens Generated: {global_num_tokens_generated}\n')
    log_file.write(f'Time: {t2-t1}\n')

Import

import json
import time

I/O Contract

Inputs

Name	Type	Required	Description
results	list[dict]	Yes	List of per-question result dicts from reward_guided_search return values
args.output_path	str	Yes	Path for combined JSON output file
global_kv_size	int	Yes	Total KV cache size across all questions (global counter)
global_num_model_calls	int	Yes	Total number of model generation calls (global counter)
global_num_tokens_generated	int	Yes	Total tokens generated (global counter)

Outputs

Name	Type	Description
answers.json	File	JSON array of per-question results with candidates and scores
stats.log	File	Text log with KV cache size, model calls, tokens, and wall-clock time

answers.json element structure:

{
    "id": 0,
    "question": "Find the value of x such that ...",
    "model_answer": [
        {
            "text": "Full trajectory text including all reasoning steps...",
            "step_scores": [0.95, 0.87, 0.92]
        },
        # ... more candidates
    ],
    "ground_truth_answer": "42",
    "total_tokens": 1523
}

Usage Examples

Full Result Collection Flow

import json
import time
import os

t1 = time.time()

# Run batch tree search
states = reward_guided_search.run_batch(input_list_dict, backend=policy_endpoint,
                                         num_threads=paras["num_threads"], progress_bar=True)

# Collect results
results = []
total_gen_tokens = 0
for s in states:
    answer = s.ret_value
    total_gen_tokens += answer["total_tokens"]
    results.append(answer)

# Write combined results
json.dump(results, open(args.output_path, "w"), indent=4)

# Write statistics
t2 = time.time()
output_dir = os.path.dirname(args.output_path)
with open(os.path.join(output_dir, 'stats.log'), 'w') as log_file:
    log_file.write(f'Total KV cache size: {global_kv_size}\n')
    log_file.write(f'Number of Model Calls: {global_num_model_calls}\n')
    log_file.write(f'Number of Tokens Generated: {global_num_tokens_generated}\n')
    log_file.write(f'Time: {t2-t1}\n')

Related Pages

Implements Principle

• Principle:SqueezeAILab_ETS_Result_Collection

Requires Environment

• Environment:SqueezeAILab_ETS_Multi_GPU_Sglang_Runtime