Skip to content

updates for stability and speed #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
teknium1 merged 2 commits into from
Jan 8, 2026
+448 −131
Merged

updates for stability and speed #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6af6ff2
updates for stability and speed
Jan 8, 2026
80d3263
Merge branch 'main' into speed-upgrades
teknium1 Jan 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 48 additions & 4 deletions batch_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,8 @@ def _process_single_prompt(
if config.get("verbose"):
print(f" Prompt {prompt_index}: Using toolsets {selected_toolsets}")

# Initialize agent with sampled toolsets
# Initialize agent with sampled toolsets and log prefix for identification
log_prefix = f"[B{batch_num}:P{prompt_index}]"
agent = AIAgent(
base_url=config.get("base_url"),
api_key=config.get("api_key"),
Expand All @@ -164,7 +165,12 @@ def _process_single_prompt(
save_trajectories=False, # We handle saving ourselves
verbose_logging=config.get("verbose", False),
ephemeral_system_prompt=config.get("ephemeral_system_prompt"),
log_prefix_chars=config.get("log_prefix_chars", 100)
log_prefix_chars=config.get("log_prefix_chars", 100),
log_prefix=log_prefix,
providers_allowed=config.get("providers_allowed"),
providers_ignored=config.get("providers_ignored"),
providers_order=config.get("providers_order"),
provider_sort=config.get("provider_sort"),
)

# Run the agent with task_id to ensure each task gets its own isolated VM
Expand Down Expand Up @@ -325,6 +331,10 @@ def __init__(
verbose: bool = False,
ephemeral_system_prompt: str = None,
log_prefix_chars: int = 100,
providers_allowed: List[str] = None,
providers_ignored: List[str] = None,
providers_order: List[str] = None,
provider_sort: str = None,
):
"""
Initialize the batch runner.
Expand All @@ -342,6 +352,10 @@ def __init__(
verbose (bool): Enable verbose logging
ephemeral_system_prompt (str): System prompt used during agent execution but NOT saved to trajectories (optional)
log_prefix_chars (int): Number of characters to show in log previews for tool calls/responses (default: 20)
providers_allowed (List[str]): OpenRouter providers to allow (optional)
providers_ignored (List[str]): OpenRouter providers to ignore (optional)
providers_order (List[str]): OpenRouter providers to try in order (optional)
provider_sort (str): Sort providers by price/throughput/latency (optional)
"""
self.dataset_file = Path(dataset_file)
self.batch_size = batch_size
Expand All @@ -355,6 +369,10 @@ def __init__(
self.verbose = verbose
self.ephemeral_system_prompt = ephemeral_system_prompt
self.log_prefix_chars = log_prefix_chars
self.providers_allowed = providers_allowed
self.providers_ignored = providers_ignored
self.providers_order = providers_order
self.provider_sort = provider_sort

# Validate distribution
if not validate_distribution(distribution):
Expand Down Expand Up @@ -511,7 +529,11 @@ def run(self, resume: bool = False):
"api_key": self.api_key,
"verbose": self.verbose,
"ephemeral_system_prompt": self.ephemeral_system_prompt,
"log_prefix_chars": self.log_prefix_chars
"log_prefix_chars": self.log_prefix_chars,
"providers_allowed": self.providers_allowed,
"providers_ignored": self.providers_ignored,
"providers_order": self.providers_order,
"provider_sort": self.provider_sort,
}

# Get completed prompts set
Expand All @@ -522,6 +544,8 @@ def run(self, resume: bool = False):

start_time = time.time()

print(f"\n🔧 Initializing {self.num_workers} worker processes...")

# Process batches in parallel
with Pool(processes=self.num_workers) as pool:
# Create tasks for each batch
Expand All @@ -536,6 +560,9 @@ def run(self, resume: bool = False):
for batch_num, batch_data in enumerate(self.batches)
]

print(f"✅ Created {len(tasks)} batch tasks")
print(f"🚀 Starting parallel batch processing...\n")

# Use map to process batches in parallel
results = pool.map(_process_batch_worker, tasks)

Expand Down Expand Up @@ -656,6 +683,10 @@ def main(
list_distributions: bool = False,
ephemeral_system_prompt: str = None,
log_prefix_chars: int = 100,
providers_allowed: str = None,
providers_ignored: str = None,
providers_order: str = None,
provider_sort: str = None,
):
"""
Run batch processing of agent prompts from a dataset.
Expand All @@ -675,6 +706,10 @@ def main(
list_distributions (bool): List available toolset distributions and exit
ephemeral_system_prompt (str): System prompt used during agent execution but NOT saved to trajectories (optional)
log_prefix_chars (int): Number of characters to show in log previews for tool calls/responses (default: 20)
providers_allowed (str): Comma-separated list of OpenRouter providers to allow (e.g. "anthropic,openai")
providers_ignored (str): Comma-separated list of OpenRouter providers to ignore (e.g. "together,deepinfra")
providers_order (str): Comma-separated list of OpenRouter providers to try in order (e.g. "anthropic,openai,google")
provider_sort (str): Sort providers by "price", "throughput", or "latency" (OpenRouter only)

Examples:
# Basic usage
Expand Down Expand Up @@ -722,6 +757,11 @@ def main(
print("❌ Error: --run_name is required")
return

# Parse provider preferences (comma-separated strings to lists)
providers_allowed_list = [p.strip() for p in providers_allowed.split(",")] if providers_allowed else None
providers_ignored_list = [p.strip() for p in providers_ignored.split(",")] if providers_ignored else None
providers_order_list = [p.strip() for p in providers_order.split(",")] if providers_order else None

# Initialize and run batch runner
try:
runner = BatchRunner(
Expand All @@ -736,7 +776,11 @@ def main(
num_workers=num_workers,
verbose=verbose,
ephemeral_system_prompt=ephemeral_system_prompt,
log_prefix_chars=log_prefix_chars
log_prefix_chars=log_prefix_chars,
providers_allowed=providers_allowed_list,
providers_ignored=providers_ignored_list,
providers_order=providers_order_list,
provider_sort=provider_sort,
)

runner.run(resume=resume)
Expand Down
Loading