From 3693fd9cc699d7d03aa168881f4b555f8eeb5312 Mon Sep 17 00:00:00 2001 From: elijahgjacob <78339643+elijahgjacob@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:00:19 -0700 Subject: [PATCH 1/3] Refresh Large_Scale_Tasks_Recipe.py for the current Task Group API The old recipe was a product-matching script on the beta SDK surface with a hard-coded task spec. This replaces it with a generic, resumable batch runner on client.task_group: plan (no API calls), paced submit at 1,000 runs per request under the caller's quota with run ids checkpointed before the next request, status across groups, and JSONL export that validates every input row came back exactly once (exit 2 otherwise). Verified end to end against the live API on 3 lite runs: submit, resume no-op, status, export with basis, and a tampered run log failing validation. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BqEaXBEuFVsj4Cz5iHyy2q --- README.md | 2 +- python-recipes/Large_Scale_Tasks_Recipe.py | 875 +++++++++------------ python-recipes/README.md | 15 +- 3 files changed, 364 insertions(+), 528 deletions(-) diff --git a/README.md b/README.md index d6a56a6..097d143 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ Take a thin input (a name, a domain) and return structured, cited fields. | Recipe | Description | APIs | Stack | Demo | | --- | --- | --- | --- | --- | | [**Supabase Enrichment**](typescript-recipes/parallel-supabase-enrichment) | Real-time enrichment pipeline — INSERT a company, an Edge Function fires a Task, results stream back via Supabase Realtime. | `Task` `Webhooks` | Next.js · Supabase Edge Functions · Postgres | – | -| [**Large-Scale Tasks**](python-recipes/Large_Scale_Tasks_Recipe.py) | Production batch script for 1k+ row CSVs — three-stage enqueue → fetch → merge with retry, dry-run, and idempotent state. | `Task Group` | Python | – | +| [**Large-Scale Tasks**](python-recipes/Large_Scale_Tasks_Recipe.py) | Resumable batch runner for CSVs of any size — plan, paced submit at 1,000 runs/request, status, JSONL export with row-level validation. | `Task Group` | Python | – | | [**Task Group + Temporal**](python-recipes/Task_Group_Temporal_Recipe.py) | Combine Task Groups with Temporal workflow orchestration for enterprise-grade durability. | `Task Group` | Python · Temporal | – | ### Realtime Streaming (SSE) diff --git a/python-recipes/Large_Scale_Tasks_Recipe.py b/python-recipes/Large_Scale_Tasks_Recipe.py index 35c9a83..30c4130 100644 --- a/python-recipes/Large_Scale_Tasks_Recipe.py +++ b/python-recipes/Large_Scale_Tasks_Recipe.py @@ -1,519 +1,356 @@ - -"""Script to run large scale batches of tasks. - -## Overview - -An overview of the workflow (when running end-to-end): -1. Enqueue all files. - - Artifacts created: output_dir/{input_file_name}_runs.csv - This file is used to store the runs. Each file corresponds to one task group. - - Note: Non validated files are not enqueued -2. Fetch all results. - - Artifacts created: output_dir/results/{input_file_name}_results.csv and - output_dir/runs/{input_file_name}_runs.csv - The runs directory is used to store the results. Results for each input file is stored in - the results directory at output_dir/{input_file_name}_runs.csv. -3. Merge results. - - Artifacts created: output_dir/merged_results.csv - This merges results from the input and output directories and combines them into a single file. - The output columns will contain dictionary values that you will have to parse. - -## Setup: -### Inputs: -- make sure that the input folder is defined -- if input size is >1k rows, split it into multiple files so that no file has more than 1k rows -- Each input file should be a valid csv. -- all input files should be in the same input directory -- Input files shouldn't end with the _tgrp_runs.csv suffix. -### Outputs: -- Make sure that the output directory has been created -### Environment variables -- PARALLEL_API_KEY is set as an environment variable. - -## Dependencies: -- pandas -- parallel-web>=0.2.0 - -## Running the script - -It is strongly recommended to start a dry run first to verify all inputs are correct. - -To run, use: `python3 combined_script.py --input-dir --output-dir --processor --dry-run` -Once you don't receive any errors, remove the --dry-run flag and run again. -In case there are csv files that are not valid and are not fixable, you can skip them by adding the --skip-invalid flag. -""" - -import os -from typing import Literal -import pandas as pd -from collections.abc import Iterator - - -import time -import argparse -import logging - -import pandas as pd -from parallel import Parallel -from parallel.types import TaskRunJsonOutput -from typing import Any -import pandas as pd - - -from parallel.types.beta import BetaRunInputParam - -######################### UTILITY FUNCTIONS ######################### - -class ValidationError(Exception): - """Error raised when a file fails validation.""" - - pass - -class NonRetryableError(Exception): - """Error raised when a file fails validation.""" - - pass - -def load_csv(file: str) -> pd.DataFrame: - """Load a CSV file safely. - - Empty columns should not be treated as a float. - """ - return pd.read_csv(file, na_filter=False) - -def iter_files(input_dir: str) -> Iterator[str]: - """Iterate over all files in a directory.""" - if not os.path.exists(input_dir): - raise ValidationError(f"Directory {input_dir} does not exist.") - file_found = False - for file in os.listdir(input_dir): - if file.endswith(".csv"): - file_found = True - yield os.path.join(input_dir, file) - if not file_found: - raise ValidationError(f"No CSV files found in {input_dir}.") - -Stages = Literal["runs", "results"] - -class FileManager: - """Manages file names and paths for a given output directory.""" - - def __init__(self, output_dir: str): - self.output_dir = output_dir - - def output_file_path(self, input_filename: str, *, stage: Stages) -> str: - """Get the output file path for a given input file and stage.""" - RUN_STATE_FILE_SUFFIX = "_tgrp_runs.csv" - if input_filename.endswith(RUN_STATE_FILE_SUFFIX): - input_filename = os.path.basename(input_filename).rstrip(RUN_STATE_FILE_SUFFIX) - else: - input_filename = os.path.basename(input_filename).rstrip(".csv") - RESULT_STATE_FILE_SUFFIX = "_results.csv" - subdir = stage - if not os.path.exists(os.path.join(self.output_dir, subdir)): - os.makedirs(os.path.join(self.output_dir, subdir)) - match stage: - case "results": - filename = f"{input_filename}{RESULT_STATE_FILE_SUFFIX}" - case "runs": - filename = f"{input_filename}{RUN_STATE_FILE_SUFFIX}" - return os.path.join(self.output_dir, subdir, filename) - - - @staticmethod - def validate_file(file: str): - """Validate a file.""" - if not os.path.exists(file): - raise ValidationError(f"File {file} does not exist.") - FILE_LENGTH_LIMIT = 1000 - try: - input_df = pd.read_csv(file) - except Exception as e: - raise NonRetryableError(f"Incorrect csv file {file}: {e}") - if len(input_df) > FILE_LENGTH_LIMIT: - raise ValidationError( - f"File {file} has more than {FILE_LENGTH_LIMIT} rows (len={len(input_df)})." - ) - - - def already_enqueued(self, input_file_name: str) -> bool: - """Check if a file has already been submitted for a run. - - This is proxied by the presence of a file in the output directory. - """ - run_file_name = self.output_file_path(input_file_name, stage="runs") - return os.path.exists(run_file_name) - - - def already_fetched(self, run_file_name: str) -> bool: - """Check if a file has already been fetched. - - This is proxied by the presence of a file in the output directory. - """ - result_file_name = self.output_file_path(run_file_name, stage="results") - return os.path.exists(result_file_name) - - def get_output_dir(self, stage: Stages) -> str: - """Get the output directory.""" - match stage: - case "runs": - return os.path.join(self.output_dir, "runs") - case "results": - return os.path.join(self.output_dir, "results") - - -######################### Task configuration ######################### - - -OUTPUT_COLS = ["match_1", "match_2", "match_3", "match_4", "match_5"] - -def build_task_spec(domains: list[str], source_rootdomain_name: str) -> dict[str, Any]: - if len(domains) != len(OUTPUT_COLS): - raise ValidationError(f"Number of domains ({len(domains)}) does not match number of output columns ({len(OUTPUT_COLS)})") - - return { - "input_schema": { - "json_schema": { - "type": "object", - "required": ["ManufacturerPartID", "SKU", "ManufacturerPartNumber", "OptionName", "UPC", "AdditionalUPC", "PrName", "ProductDescription", "MarketingCategory", "Class", "Manufacturer", "URL"], - "properties": { - "ManufacturerPartID": { - "description": "Manufacturer part ID of the product to find matches for.", - "type": "string", - }, - "SKU": { - "description": "SKU identifier of the product to find match for.", - "type": "string", - }, - "ManufacturerPartNumber": { - "description": "Manufacturer part number (MPN) of the product to find matches for.", - "type": "string", - }, - "OptionName": { - "description": "", - "type": "string", - }, - "UPC": { - "description": "The UPC of the product to find matches for.", - "type": "string", - }, - "AdditionalUPC": { - "description": "", - "type": "string", - }, - "PrName": { - "description": "The name of the product to find matches for.", - "type": "string", - }, - "ProductDescription": { - "description": "The description of the product to find matches for.", - "type": "string", - }, - "MarketingCategory": { - "description": "The category of the product to find matches for.", - "type": "string", - }, - "Class": { - "description": "The class of the product to find matches for.", - "type": "string", - }, - "Manufacturer": { - "description": "Name of manufacturer of the product to find matches for.", - "type": "string", - }, - "URL": { - "description": f"The direct URL to the {source_rootdomain_name} product page. Use this URL to first extract all the product details including manufacturer, part number, product name, specifications, dimensions, weight, price, etc. before matching.", - "type": "string", - } - }, - } - }, - "output_schema": { - "json_schema": { - "type": "object", - "required": OUTPUT_COLS, - "description": ( - f"An exact match to the given product on target domains: {domains}. The match must have the same make and model -- i.e. the same manufacturer name and other details.\n" - "Matching Criteria (in order of priority):\n" - "1. **UPC (Universal Product Code) - Exact match** 2. **Manufacturer Part Number (MPN) - Exact match** 3. **Manufacturer Name** – Must match or be a known alias/brand variation 4. **Product Title/Option Name** – High similarity 5. **Product Class/Category** – Must be consistent 6. ** **Visual Match** (if available) – Product images should be visually identical " - ), - "properties": { - OUTPUT_COLS[i]: { - "description": f"The exact match to the original product on {domains[i]}.", - "properties": { - "product_url": { - "description": f"The direct URL of the matched {domains[i]} product page (must be from {domains[i]}). URL that is not from {domains[i]} is invalid and not considered a match. Must be a valid URL that opens up to the actual product page directly. If no match, return empty string.", - "type": "string", - }, - "product_description": { - "description": f"The description of the matched {domains[i]} product page. If a description is not available, return 'Description unavailable.'. If no match, return empty string.", - "type": "string", - }, - "product_price": { - "description": "The price of the matched product, including currency symbol (e.g., '$5.99'). If unavailable, return 'Price Not Available'. If no match, return empty string.", - "type": "string", - }, - "product_in_stock": { - "description": "An indication whether the matched product is in-stock or not. If no match, return empty string.", - "enum": ["yes", "no", ""], - "type": "string", - } - }, - "type": "object", - } for i in range(len(OUTPUT_COLS)) - }, - } - }, - } - -def create_run_payloads( - chunk_df: pd.DataFrame, source_rootdomain_name: str, processor: str -) -> dict[str, BetaRunInputParam]: - # Build inputs across domains - run_map : dict[str, BetaRunInputParam] = {} - for row in chunk_df.itertuples(): - row_domains = [row.competitor1, row.competitor2, row.competitor3, row.competitor4, row.competitor5] - product_data = { - "ManufacturerPartID": row.ManufacturerPartID, - "SKU": row.SKU, - "ManufacturerPartNumber": row.ManufacturerPartNumber, - "OptionName": row.OptionName, - "UPC": row.UPC, - "AdditionalUPC": row.AdditionalUPC, - "PrName": row.PrName, - "ProductDescription": row.ProductDescription, - "MarketingCategory": row.MarketingCategory, - "Class": row.Class, - "Manufacturer": row.Manufacturer, - "URL": row.URL, - "domains": row_domains - } - mpn_str = str(row.ManufacturerPartNumber) - - task_spec = build_task_spec(row_domains, source_rootdomain_name) - run_map[mpn_str] = BetaRunInputParam( - input=product_data, - processor=processor, - task_spec=task_spec, - metadata={"manufacturerPartId": product_data.get("ManufacturerPartID", ""), "taskType": "match_search"}, - source_policy={"include_domains": row_domains}, - ) - - return run_map - -######################### Execution script ######################### - -# Flags -DRY_RUN = False -SKIP_INVALID = False -ENQUEUE_SLEEP_TIME = 5 -FETCH_SLEEP_TIME = 60 - -SOURCE_ROOTDOMAIN_NAME = "Wayfair" - -OUTPUT_DIR: str | None = None - -# initialize logging -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') -logger = logging.getLogger(__name__) - -# set the env variable PARALLEL_API_KEY or specify the api key explicitly -# via client = Parallel(api_key="your_api_key") -client = Parallel() - -taskGroupIdCol = "TaskGroupID" -runIdCol = "RunId" -mergeIdCol = "ManufacturerPartNumber" - -def enqueue_one(filepath: str, processor: str, file_manager: FileManager) -> bool: - """Enqueue a file after validation. Wraps errors for retries.""" - try: - return _enqueue_one(filepath, processor, file_manager) - except ValidationError: - raise - except Exception as e: - logger.error(f"Error enqueuing file {filepath}: {e}. Will retry.") - return False - -def _enqueue_one(filepath: str, processor: str, file_manager: FileManager) -> bool: - """Enqueue a file after validation. - - Each file corresponds to one task group. - """ - try: - file_manager.validate_file(filepath) - except NonRetryableError as e: - if not SKIP_INVALID: - raise ValidationError(f"File {filepath} failed .") from e - logger.error(f"File {filepath} failed validation: {e}. It will be skipped.") - return True - except ValidationError as e: - logger.error(f"File {filepath} failed validation: {e}") - raise ValidationError("Files failed validation. Please check logs for more details.") from e - input_df = load_csv(filepath) - if DRY_RUN: - logger.info(f"Skipping enqueue for file {filepath} due to dry run.") - return True - - # enqueue - input_map = create_run_payloads(input_df, SOURCE_ROOTDOMAIN_NAME, processor) - tgroup = client.beta.task_group.create() - run_responses = client.beta.task_group.add_runs(tgroup.task_group_id, inputs=[v for _,v in input_map.items()]) - - # write to state file in output directory - state_map: list[dict[str, str]] = [] - for i, run_key in enumerate(input_map): - state_map.append({ - mergeIdCol: run_key, - runIdCol: run_responses.run_ids[i], - taskGroupIdCol: tgroup.task_group_id - }) - - logger.info(f"Processed file {filepath} with {len(state_map)} runs.") - pd.DataFrame(state_map).to_csv(file_manager.output_file_path(filepath, stage="runs"), index=False) - return True - - -def enqueue_all(input_dir: str, processor: str, file_manager: FileManager): - """Enqueue all files in the input directory. - - Depending on the mode, it might raise an error. - """ - logger.info("DRY RUN" if DRY_RUN else "Live Run") - while True: - all_completed = True - for file in iter_files(input_dir): - if file_manager.already_enqueued(file): - logger.info(f"File {file} is already enqueued, skipping.") - continue - completed = enqueue_one(file, processor, file_manager) - if not completed: - all_completed = False - - if all_completed: - logger.info("All files enqueued successfully.") - break - logger.info(f"Some files failed to enqueue. Waiting {ENQUEUE_SLEEP_TIME} seconds before retrying.") - if DRY_RUN: - logger.info("Breaking out of enqueue loop due to dry run.") - break - time.sleep(ENQUEUE_SLEEP_TIME) - - -def fetch_all(file_manager: FileManager): - """Fetch all results from the output directory.""" - while True: - all_completed = True - for file in iter_files(file_manager.get_output_dir("runs")): - if file_manager.already_fetched(file): - logger.info(f"Results for file {file} already fetched, skipping") - continue - completed = fetch_one(file, file_manager) - # heuristic to reduce poll count - # early exit and sleep - if not completed: # still active - all_completed = False - break - if all_completed: - break - time.sleep(FETCH_SLEEP_TIME) - - - -def fetch_one(run_file: str, file_manager: FileManager) -> bool: - """Fetch a single result from the output directory. Wraps errors for retries.""" - try: - logger.info(f"Fetching result from file {run_file}") - return _fetch_one(run_file, file_manager) - except Exception as e: - logger.error(f"Error fetching result from file {run_file}: {e}") - return False - -def _fetch_one(run_file: str, file_manager: FileManager) -> bool: - """Fetch a single result from the output directory. - - Poll until the task group is complete. Once it is finished, fetch the result. - """ - run_df = load_csv(run_file) - # each file has just one task group - tgroup_id = str(run_df[taskGroupIdCol][0]) - tgroup = client.beta.task_group.retrieve(tgroup_id) - if tgroup.status.is_active: - logger.info(f"File {run_file} (Task group {tgroup_id}) is still active, skipping") - return False - results: list[dict[str, str]] = [] - for _, row in run_df.iterrows(): - run_id = str(row[runIdCol]) - result = None - try: - result = client.task_run.result(run_id) - except Exception as e: - # taskgroup is done, which means the run failed. - logger.error(f"Run {run_id} in file {run_file} failed. Most likely failed. Error: {e}") - continue - if not isinstance(result.output, TaskRunJsonOutput): - logger.error(f"Result for run {run_id} in file {run_file} is not a JSON output, skipping") - continue - - results.append({ - mergeIdCol: str(row[mergeIdCol]), - runIdCol: run_id, - taskGroupIdCol: tgroup_id, - **{OUTPUT_COLS[i]: result.output.content.get(OUTPUT_COLS[i], None) for i in range(len(OUTPUT_COLS))} # pyright: ignore[reportArgumentType] - }) - pd.DataFrame(results).to_csv(file_manager.output_file_path(run_file, stage="results"), index=False) - return True - - -def merge_results(input_dir: str, file_manager: FileManager): - """Merge results from the input and output directory.""" - df_list: list[pd.DataFrame] = [] - for input_file in iter_files(input_dir): - if not file_manager.already_fetched(input_file): - logger.info(f"Results for file {input_file} not fetched, skipping for merge.") - continue - result_df = load_csv(file_manager.output_file_path(input_file, stage="results")) - input_df = load_csv(input_file) - print(file_manager.output_file_path(input_file, stage="results")) - merged_df = pd.merge(input_df, result_df, on=mergeIdCol, how="left") - df_list.append(merged_df) - merged_df = pd.concat(df_list) - merged_df.to_csv(file_manager.output_file_path("merged", stage="results"), index=False) - -def run_batch(input_dir: str, output_dir: str, processor: str): - """Run a batch of files in the input directory. - - This is a three step process: - 1. Enqueue all files in the input directory. - 2. Fetch results until complete. - 3. Merge results. - - This should be an idempotent operation, meaning that running it on the same input, - should not incur any additional cost/time. - """ - file_manager = FileManager(output_dir) - enqueue_all(input_dir, processor, file_manager) - if DRY_RUN: - logger.info("Dry run complete.") - return - logger.info("Fetching results.") - fetch_all(file_manager) - logger.info("Merging results.") - merge_results(input_dir, file_manager) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--input-dir", type=str, required=True) - parser.add_argument("--output-dir", type=str, required=True) - parser.add_argument("--processor", type=str, required=True) - parser.add_argument("--dry-run", action="store_true") - parser.add_argument("--skip-invalid", action="store_true") - args = parser.parse_args() - - if args.dry_run: - DRY_RUN = True - if args.skip_invalid: - SKIP_INVALID = True - run_batch(args.input_dir, args.output_dir, args.processor) - +"""Run a large CSV through the Parallel Task Group API. + +One file, four commands, resumable. Built from the pattern we use for +customer batches in the hundreds of thousands to millions of runs. + + python Large_Scale_Tasks_Recipe.py plan --input rows.csv --rate-limit 2000 + python Large_Scale_Tasks_Recipe.py submit --input rows.csv --task-spec spec.json \\ + --processor core --id-column row_id --work-dir jobs/run1 + python Large_Scale_Tasks_Recipe.py status --work-dir jobs/run1 --wait + python Large_Scale_Tasks_Recipe.py export --work-dir jobs/run1 --output results.jsonl + +How it works +- One CSV row becomes one run. Every column except --id-column is sent as a + string field of the input object, or pass --input-json-column to send one + column as a JSON payload. +- Runs are added 1,000 per request (the API maximum) with refresh_status=False, + sharded into Task Groups of --runs-per-group. +- Submission is paced at 90% of --rate-limit (runs per minute), because the + quota counts runs, not requests. A steady rate beats one burst: it enqueues + cleanly and surfaces a bad spec after thousands of rows, not millions. +- Every add_runs response is appended to work-dir/runs.jsonl before the next + request goes out, so a crash or re-run never resubmits a paid run. +- status polls each group summary (one cheap GET per group). export streams + each group's runs to disk with include_output and checks that every input + row came back exactly once. It exits 2 if anything is missing or duplicated. + +Things the API will not do for you +- Runs cannot be cancelled once created. Run `plan`, then a small pilot, then + the full job. Submit high-priority rows first. +- Your rate limit controls intake, not throughput. Measure a ~5k pilot to + learn runs/hour for your processor and spec, then extrapolate. + +Requires Python 3.11+ and `pip install parallel-web>=1.3`. Set PARALLEL_API_KEY. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import sys +import time +from pathlib import Path +from typing import Any + +MAX_RUNS_PER_REQUEST = 1_000 # API limit, do not raise +RATE_LIMIT_TARGET = 0.9 # fraction of the quota to use +POLL_INTERVAL_S = 60 + +csv.field_size_limit(10**9) + + +# ----------------------------------------------------------------------------- input + + +def read_rows(path: str, id_column: str) -> list[dict[str, str]]: + with open(path, newline="", encoding="utf-8") as f: + rows = list(csv.DictReader(f)) + if not rows: + sys.exit(f"{path}: no rows") + if id_column not in rows[0]: + sys.exit(f"{path}: missing id column {id_column!r}. Columns: {list(rows[0])}") + ids = [r[id_column] for r in rows] + if len(set(ids)) != len(ids): + sys.exit(f"{path}: {len(ids) - len(set(ids))} duplicate values in {id_column!r}") + return rows + + +def build_input(row: dict[str, str], id_column: str, json_column: str | None) -> Any: + if json_column: + try: + return json.loads(row[json_column]) + except (KeyError, json.JSONDecodeError) as e: + sys.exit(f"row {row.get(id_column)}: bad JSON in {json_column!r}: {e}") + return {k: v for k, v in row.items() if k != id_column} + + +# ----------------------------------------------------------------------------- state + + +class Job: + """Append-only state in a work directory. Safe to re-run any command.""" + + def __init__(self, work_dir: str): + self.dir = Path(work_dir) + self.dir.mkdir(parents=True, exist_ok=True) + self.config_path = self.dir / "config.json" + self.groups_path = self.dir / "groups.jsonl" + self.runs_path = self.dir / "runs.jsonl" + + def save_config(self, cfg: dict[str, Any]) -> None: + if self.config_path.exists(): + existing = json.loads(self.config_path.read_text()) + if existing != cfg: + sys.exit(f"{self.config_path} already exists with different settings. Use a new --work-dir.") + self.config_path.write_text(json.dumps(cfg, indent=2)) + + def config(self) -> dict[str, Any]: + if not self.config_path.exists(): + sys.exit(f"no job at {self.dir}; run submit first") + return json.loads(self.config_path.read_text()) + + def groups(self) -> list[str]: + return [json.loads(l)["task_group_id"] for l in _lines(self.groups_path)] + + def add_group(self, task_group_id: str) -> None: + _append(self.groups_path, {"task_group_id": task_group_id, "created_at": time.time()}) + + def submitted(self) -> dict[str, dict[str, str]]: + """row_id -> {run_id, task_group_id}""" + out: dict[str, dict[str, str]] = {} + for l in _lines(self.runs_path): + rec = json.loads(l) + out[rec["row_id"]] = rec + return out + + def add_runs(self, task_group_id: str, pairs: list[tuple[str, str]]) -> None: + with open(self.runs_path, "a", encoding="utf-8") as f: + for row_id, run_id in pairs: + f.write(json.dumps({"row_id": row_id, "run_id": run_id, "task_group_id": task_group_id}) + "\n") + f.flush() + + +def _lines(path: Path) -> list[str]: + return [l for l in path.read_text().splitlines() if l.strip()] if path.exists() else [] + + +def _append(path: Path, rec: dict[str, Any]) -> None: + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(rec) + "\n") + f.flush() + + +# ----------------------------------------------------------------------------- plan + + +def plan(n_runs: int, rate_limit: float, runs_per_group: int) -> dict[str, Any]: + per_min = rate_limit * RATE_LIMIT_TARGET + posts = math.ceil(n_runs / MAX_RUNS_PER_REQUEST) + return { + "runs": n_runs, + "posts": posts, + "task_groups": math.ceil(n_runs / runs_per_group), + "submit_rate_runs_per_min": per_min, + "enqueue_minutes": round(n_runs / per_min, 1), + "note": "Enqueue time only. Execution time depends on processor, spec, and platform load: run a ~5k pilot and extrapolate.", + } + + +# ----------------------------------------------------------------------------- submit + + +def cmd_submit(args: argparse.Namespace) -> None: + from parallel import Parallel + + rows = read_rows(args.input, args.id_column) + task_spec = json.loads(Path(args.task_spec).read_text()) + job = Job(args.work_dir) + job.save_config( + { + "input": str(Path(args.input).resolve()), + "id_column": args.id_column, + "input_json_column": args.input_json_column, + "processor": args.processor, + "runs_per_group": args.runs_per_group, + "label": args.label, + } + ) + done = job.submitted() + pending = [r for r in rows if r[args.id_column] not in done] + print(f"{len(rows)} rows, {len(done)} already submitted, {len(pending)} to go") + if args.dry_run or not pending: + print(json.dumps(plan(len(pending), args.rate_limit, args.runs_per_group), indent=2)) + return + + client = Parallel() + per_request_s = 60.0 * MAX_RUNS_PER_REQUEST / (args.rate_limit * RATE_LIMIT_TARGET) + groups = job.groups() + group_fill = _group_fill(job) + + i = 0 + while i < len(pending): + # pick a group with room, or create one + gid = next((g for g in groups if group_fill.get(g, 0) < args.runs_per_group), None) + if gid is None: + gid = client.task_group.create(metadata={"label": args.label} if args.label else None).task_group_id + job.add_group(gid) + groups.append(gid) + room = args.runs_per_group - group_fill.get(gid, 0) + batch = pending[i : i + min(MAX_RUNS_PER_REQUEST, room)] + inputs = [ + { + "input": build_input(r, args.id_column, args.input_json_column), + "processor": args.processor, + "metadata": {"row_id": r[args.id_column], **({"label": args.label} if args.label else {})}, + } + for r in batch + ] + started = time.monotonic() + resp = client.task_group.add_runs(gid, inputs=inputs, default_task_spec=task_spec, refresh_status=False) + if len(resp.run_ids) != len(batch): + sys.exit(f"server returned {len(resp.run_ids)} run ids for {len(batch)} inputs; reconcile {gid} before continuing") + job.add_runs(gid, list(zip((r[args.id_column] for r in batch), resp.run_ids))) + group_fill[gid] = group_fill.get(gid, 0) + len(batch) + i += len(batch) + print(f" {i}/{len(pending)} submitted ({gid})") + time.sleep(max(0.0, per_request_s - (time.monotonic() - started))) + print("done. next: status --wait, then export") + + +def _group_fill(job: Job) -> dict[str, int]: + fill: dict[str, int] = {} + for rec in job.submitted().values(): + fill[rec["task_group_id"]] = fill.get(rec["task_group_id"], 0) + 1 + return fill + + +# ----------------------------------------------------------------------------- status + + +def cmd_status(args: argparse.Namespace) -> None: + from parallel import Parallel + + job = Job(args.work_dir) + client = Parallel() + while True: + counts: dict[str, int] = {} + active = 0 + for gid in job.groups(): + st = client.task_group.retrieve(gid).status + active += st.is_active + for k, v in (st.task_run_status_counts or {}).items(): + counts[k] = counts.get(k, 0) + v + total = sum(counts.values()) + done = counts.get("completed", 0) + counts.get("failed", 0) + counts.get("cancelled", 0) + print(f"{time.strftime('%H:%M:%S')} {done}/{total} finished {counts} active_groups={active}") + if not active or not args.wait: + return + time.sleep(POLL_INTERVAL_S) + + +# ----------------------------------------------------------------------------- export + + +def cmd_export(args: argparse.Namespace) -> None: + from parallel import Parallel + + job = Job(args.work_dir) + client = Parallel() + expected = job.submitted() + run_to_row = {rec["run_id"]: row_id for row_id, rec in expected.items()} + seen: dict[str, int] = {} + n_written = n_failed = n_active = 0 + + with open(args.output, "w", encoding="utf-8") as out: + for gid in job.groups(): + for event in client.task_group.get_runs(gid, include_input=args.include_input, include_output=True): + if event.type != "task_run.state": + print(f" stream error in {gid}: {event}", file=sys.stderr) + continue + run = event.run + row_id = run_to_row.get(run.run_id) + if row_id is None: + print(f" unexpected run {run.run_id} in {gid}", file=sys.stderr) + seen["__unexpected__"] = seen.get("__unexpected__", 0) + 1 + continue + if run.is_active: + n_active += 1 + continue + seen[row_id] = seen.get(row_id, 0) + 1 + rec: dict[str, Any] = {"row_id": row_id, "run_id": run.run_id, "status": run.status} + if event.output is not None: + rec["output"] = event.output.content + rec["basis"] = [b.model_dump() for b in event.output.basis] + if run.status == "failed": + n_failed += 1 + rec["error"] = run.error.model_dump() if run.error else None + if args.include_input and event.input is not None: + rec["input"] = event.input.input + out.write(json.dumps(rec, default=str) + "\n") + n_written += 1 + + missing = [r for r in expected if r not in seen] + duplicated = {r: c for r, c in seen.items() if c > 1 and r != "__unexpected__"} + report = { + "expected_rows": len(expected), + "written": n_written, + "failed": n_failed, + "still_active": n_active, + "missing": len(missing), + "duplicated": len(duplicated), + "unexpected": seen.get("__unexpected__", 0), + "ok": not missing and not duplicated and not n_active and not seen.get("__unexpected__"), + } + Path(args.output + ".validation.json").write_text(json.dumps({**report, "missing_row_ids": missing[:1000]}, indent=2)) + print(json.dumps(report, indent=2)) + if not report["ok"]: + sys.exit(2) + + +# ----------------------------------------------------------------------------- cli + + +def cmd_plan(args: argparse.Namespace) -> None: + n = args.runs if args.runs else len(read_rows(args.input, args.id_column)) + print(json.dumps(plan(n, args.rate_limit, args.runs_per_group), indent=2)) + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + sub = p.add_subparsers(dest="command", required=True) + + def common(sp: argparse.ArgumentParser) -> None: + sp.add_argument("--rate-limit", type=float, default=2000, help="your Tasks quota in runs per minute (default 2000)") + sp.add_argument("--runs-per-group", type=int, default=10_000, help="runs per Task Group (default 10000)") + + a = sub.add_parser("plan", help="count runs, requests, groups and enqueue time. No API calls") + a.add_argument("--input") + a.add_argument("--runs", type=int, help="instead of --input, plan for this many runs") + a.add_argument("--id-column", default="row_id") + common(a) + a.set_defaults(func=cmd_plan) + + s = sub.add_parser("submit", help="create runs; re-run to resume") + s.add_argument("--input", required=True) + s.add_argument("--task-spec", required=True, help="JSON file with output_schema (and optional input_schema)") + s.add_argument("--processor", required=True) + s.add_argument("--work-dir", required=True) + s.add_argument("--id-column", default="row_id") + s.add_argument("--input-json-column", help="send this column parsed as JSON instead of all columns") + s.add_argument("--label", help="batch label stored in run and group metadata") + s.add_argument("--dry-run", action="store_true", help="validate the CSV and print the plan only") + common(s) + s.set_defaults(func=cmd_submit) + + t = sub.add_parser("status", help="progress across all groups") + t.add_argument("--work-dir", required=True) + t.add_argument("--wait", action="store_true", help=f"poll every {POLL_INTERVAL_S}s until no group is active") + t.set_defaults(func=cmd_status) + + e = sub.add_parser("export", help="write results.jsonl and validate; exit 2 if rows are missing") + e.add_argument("--work-dir", required=True) + e.add_argument("--output", required=True) + e.add_argument("--include-input", action="store_true") + e.set_defaults(func=cmd_export) + return p + + +def main(argv: list[str] | None = None) -> None: + args = build_parser().parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/python-recipes/README.md b/python-recipes/README.md index 200c03d..80c6a10 100644 --- a/python-recipes/README.md +++ b/python-recipes/README.md @@ -34,18 +34,17 @@ Shows how to combine Parallel's Task Group API with Temporal's workflow orchestr ### 📊 [Large Scale Tasks Recipe](./Large_Scale_Tasks_Recipe.py) -Production-ready batch processing for large datasets +Resumable batch processing for large CSVs with Task Groups -A robust script for processing large batches of CSV files using Parallel's Task Group API. Handles product matching across multiple e-commerce domains with comprehensive error handling, validation, and resumable operations. +One file, four commands: `plan` sizes the job with no API calls, `submit` adds runs 1,000 per request at a steady rate under your quota and checkpoints every run id before the next request, `status` polls group summaries, and `export` streams results to JSONL and checks that every input row came back exactly once. Re-running any command is safe. **Key Features:** -- Batch processing of 1000+ row CSV files -- Three-stage pipeline: enqueue → fetch → merge -- Comprehensive error handling and retry logic -- Dry-run mode for validation -- Idempotent operations for production reliability -- File validation and state management +- Paced submission against your Tasks rate limit (runs per minute) +- Crash-safe resume from an append-only run log +- Task Group sharding with `refresh_status=False` +- JSONL export with per-field basis and a validation report +- Notes on what the API will not do (no cancel, rate limit is intake not throughput) ## Getting Started From 8c8e8db26abfea993c7edd08a58497154f882cf9 Mon Sep 17 00:00:00 2001 From: elijahgjacob <78339643+elijahgjacob@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:58:45 -0700 Subject: [PATCH 2/3] Pilot guidance: at least 5k runs, measure the plateau not the ramp Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BqEaXBEuFVsj4Cz5iHyy2q --- python-recipes/Large_Scale_Tasks_Recipe.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python-recipes/Large_Scale_Tasks_Recipe.py b/python-recipes/Large_Scale_Tasks_Recipe.py index 30c4130..939bb21 100644 --- a/python-recipes/Large_Scale_Tasks_Recipe.py +++ b/python-recipes/Large_Scale_Tasks_Recipe.py @@ -27,8 +27,10 @@ Things the API will not do for you - Runs cannot be cancelled once created. Run `plan`, then a small pilot, then the full job. Submit high-priority rows first. -- Your rate limit controls intake, not throughput. Measure a ~5k pilot to - learn runs/hour for your processor and spec, then extrapolate. +- Your rate limit controls intake, not throughput. Pilot at least 5k runs on + your real processor and spec, watch `status` until the running count stops + climbing, and measure runs/hour from that plateau. Concurrency ramps up, so + a small pilot timed end to end understates steady-state throughput. Requires Python 3.11+ and `pip install parallel-web>=1.3`. Set PARALLEL_API_KEY. """ @@ -144,7 +146,7 @@ def plan(n_runs: int, rate_limit: float, runs_per_group: int) -> dict[str, Any]: "task_groups": math.ceil(n_runs / runs_per_group), "submit_rate_runs_per_min": per_min, "enqueue_minutes": round(n_runs / per_min, 1), - "note": "Enqueue time only. Execution time depends on processor, spec, and platform load: run a ~5k pilot and extrapolate.", + "note": "Enqueue time only. Execution time depends on processor, spec, and platform load: pilot at least 5k runs, measure runs/hour after the running count plateaus, then extrapolate.", } From 2bc3e378e726feeac6d513c20f03a41eba34e8c3 Mon Sep 17 00:00:00 2001 From: George Pickett Date: Thu, 1 Oct 2026 15:16:01 -0700 Subject: [PATCH 3/3] fix: prevent duplicate batch submissions and incomplete exports --- python-recipes/Large_Scale_Tasks_Recipe.py | 93 +++++- python-recipes/README.md | 6 +- .../test_large_scale_tasks_recipe.py | 287 ++++++++++++++++++ 3 files changed, 376 insertions(+), 10 deletions(-) create mode 100644 python-recipes/test_large_scale_tasks_recipe.py diff --git a/python-recipes/Large_Scale_Tasks_Recipe.py b/python-recipes/Large_Scale_Tasks_Recipe.py index 939bb21..d5e5a33 100644 --- a/python-recipes/Large_Scale_Tasks_Recipe.py +++ b/python-recipes/Large_Scale_Tasks_Recipe.py @@ -18,8 +18,10 @@ - Submission is paced at 90% of --rate-limit (runs per minute), because the quota counts runs, not requests. A steady rate beats one burst: it enqueues cleanly and surfaces a bad spec after thousands of rows, not millions. -- Every add_runs response is appended to work-dir/runs.jsonl before the next - request goes out, so a crash or re-run never resubmits a paid run. +- Each batch is recorded in work-dir/pending.json before submission, with SDK + retries disabled. On resume, its runs are recovered from server metadata; + an incomplete or ambiguous batch stops submission instead of risking duplicates. + Keep the work directory and use only one submit process per job. - status polls each group summary (one cheap GET per group). export streams each group's runs to disk with include_output and checks that every input row came back exactly once. It exits 2 if anything is missing or duplicated. @@ -41,6 +43,7 @@ import csv import json import math +import os import sys import time from pathlib import Path @@ -82,7 +85,7 @@ def build_input(row: dict[str, str], id_column: str, json_column: str | None) -> class Job: - """Append-only state in a work directory. Safe to re-run any command.""" + """Submission receipts and an in-flight batch in a work directory.""" def __init__(self, work_dir: str): self.dir = Path(work_dir) @@ -90,13 +93,15 @@ def __init__(self, work_dir: str): self.config_path = self.dir / "config.json" self.groups_path = self.dir / "groups.jsonl" self.runs_path = self.dir / "runs.jsonl" + self.pending_path = self.dir / "pending.json" def save_config(self, cfg: dict[str, Any]) -> None: if self.config_path.exists(): existing = json.loads(self.config_path.read_text()) if existing != cfg: sys.exit(f"{self.config_path} already exists with different settings. Use a new --work-dir.") - self.config_path.write_text(json.dumps(cfg, indent=2)) + return + _write_new(self.config_path, cfg) def config(self) -> dict[str, Any]: if not self.config_path.exists(): @@ -122,6 +127,67 @@ def add_runs(self, task_group_id: str, pairs: list[tuple[str, str]]) -> None: for row_id, run_id in pairs: f.write(json.dumps({"row_id": row_id, "run_id": run_id, "task_group_id": task_group_id}) + "\n") f.flush() + os.fsync(f.fileno()) + _sync_directory(self.dir) + + def begin_batch(self, task_group_id: str, row_ids: list[str]) -> None: + _write_new(self.pending_path, {"task_group_id": task_group_id, "row_ids": row_ids}) + + def finish_batch(self) -> None: + self.pending_path.unlink() + + def recover_batch(self, client: Any) -> None: + if not self.pending_path.exists(): + return + pending = json.loads(self.pending_path.read_text()) + gid = pending["task_group_id"] + row_ids = set(pending["row_ids"]) + recovered: dict[str, str] = {} + with client.task_group.get_runs(gid) as events: + for event in events: + if event.type != "task_run.state": + sys.exit(f"stream error reconciling {gid}; keep {self.pending_path} and retry submit later") + row_id = (event.run.metadata or {}).get("row_id") + if row_id not in row_ids: + continue + if row_id in recovered: + sys.exit(f"duplicate runs for row {row_id!r} in {gid}; manual reconciliation required") + recovered[row_id] = event.run.run_id + if recovered.keys() != row_ids: + sys.exit( + f"unresolved batch in {gid}: found {len(recovered)}/{len(row_ids)} rows. " + f"Keep {self.pending_path} and retry submit later. If it remains incomplete, " + "reconcile the group manually; missing runs may still be created. Do not resubmit the batch." + ) + submitted = self.submitted() + pairs = [] + for row_id, run_id in recovered.items(): + if row_id in submitted: + rec = submitted[row_id] + if rec["run_id"] != run_id or rec["task_group_id"] != gid: + sys.exit(f"conflicting receipt for row {row_id!r}; manual reconciliation required") + else: + pairs.append((row_id, run_id)) + self.add_runs(gid, pairs) + self.finish_batch() + + +def _write_new(path: Path, rec: dict[str, Any]) -> None: + # A partial file must stop resume, never authorize another paid POST. + with open(path, "x", encoding="utf-8") as f: + f.write(json.dumps(rec)) + f.flush() + os.fsync(f.fileno()) + _sync_directory(path.parent) + + +def _sync_directory(path: Path) -> None: + if os.name == "posix": + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) def _lines(path: Path) -> list[str]: @@ -132,6 +198,7 @@ def _append(path: Path, rec: dict[str, Any]) -> None: with open(path, "a", encoding="utf-8") as f: f.write(json.dumps(rec) + "\n") f.flush() + os.fsync(f.fileno()) # ----------------------------------------------------------------------------- plan @@ -167,8 +234,11 @@ def cmd_submit(args: argparse.Namespace) -> None: "processor": args.processor, "runs_per_group": args.runs_per_group, "label": args.label, + "expected_row_ids": [r[args.id_column] for r in rows], } ) + if not args.dry_run and job.pending_path.exists(): + job.recover_batch(Parallel(max_retries=0)) done = job.submitted() pending = [r for r in rows if r[args.id_column] not in done] print(f"{len(rows)} rows, {len(done)} already submitted, {len(pending)} to go") @@ -176,7 +246,7 @@ def cmd_submit(args: argparse.Namespace) -> None: print(json.dumps(plan(len(pending), args.rate_limit, args.runs_per_group), indent=2)) return - client = Parallel() + client = Parallel(max_retries=0) per_request_s = 60.0 * MAX_RUNS_PER_REQUEST / (args.rate_limit * RATE_LIMIT_TARGET) groups = job.groups() group_fill = _group_fill(job) @@ -199,11 +269,13 @@ def cmd_submit(args: argparse.Namespace) -> None: } for r in batch ] + job.begin_batch(gid, [r[args.id_column] for r in batch]) started = time.monotonic() resp = client.task_group.add_runs(gid, inputs=inputs, default_task_spec=task_spec, refresh_status=False) if len(resp.run_ids) != len(batch): sys.exit(f"server returned {len(resp.run_ids)} run ids for {len(batch)} inputs; reconcile {gid} before continuing") job.add_runs(gid, list(zip((r[args.id_column] for r in batch), resp.run_ids))) + job.finish_batch() group_fill[gid] = group_fill.get(gid, 0) + len(batch) i += len(batch) print(f" {i}/{len(pending)} submitted ({gid})") @@ -250,8 +322,12 @@ def cmd_export(args: argparse.Namespace) -> None: job = Job(args.work_dir) client = Parallel() - expected = job.submitted() - run_to_row = {rec["run_id"]: row_id for row_id, rec in expected.items()} + cfg = job.config() + if "expected_row_ids" not in cfg: + sys.exit("job has no original row manifest; reconcile it against the original CSV before exporting") + expected = set(cfg["expected_row_ids"]) + submitted = job.submitted() + run_to_row = {rec["run_id"]: row_id for row_id, rec in submitted.items() if row_id in expected} seen: dict[str, int] = {} n_written = n_failed = n_active = 0 @@ -293,7 +369,8 @@ def cmd_export(args: argparse.Namespace) -> None: "missing": len(missing), "duplicated": len(duplicated), "unexpected": seen.get("__unexpected__", 0), - "ok": not missing and not duplicated and not n_active and not seen.get("__unexpected__"), + "unresolved_batch": job.pending_path.exists(), + "ok": not missing and not duplicated and not n_active and not seen.get("__unexpected__") and not job.pending_path.exists(), } Path(args.output + ".validation.json").write_text(json.dumps({**report, "missing_row_ids": missing[:1000]}, indent=2)) print(json.dumps(report, indent=2)) diff --git a/python-recipes/README.md b/python-recipes/README.md index 80c6a10..06b9f78 100644 --- a/python-recipes/README.md +++ b/python-recipes/README.md @@ -36,12 +36,14 @@ Shows how to combine Parallel's Task Group API with Temporal's workflow orchestr Resumable batch processing for large CSVs with Task Groups -One file, four commands: `plan` sizes the job with no API calls, `submit` adds runs 1,000 per request at a steady rate under your quota and checkpoints every run id before the next request, `status` polls group summaries, and `export` streams results to JSONL and checks that every input row came back exactly once. Re-running any command is safe. +One file, four commands: `plan` sizes the job with no API calls, `submit` adds runs 1,000 per request at a steady rate under your quota and checkpoints every run id before the next request, `status` polls group summaries, and `export` streams results to JSONL and checks that every original input row came back exactly once, including rows never submitted. Use one submit process per work directory, and keep the original CSV and task spec unchanged when resuming. + +Submission records each batch before sending it and disables SDK retries for non-idempotent POSTs. After an interruption, re-run `submit` to recover the batch by its server-side row metadata. If any rows are missing or duplicated, submission stops: keep `pending.json`, retry later, and manually reconcile the group if it remains unresolved. Do not delete the checkpoint or start those rows in another work directory; the original request may still be creating paid runs. Jobs created before the original-row manifest was introduced require manual reconciliation before reuse. **Key Features:** - Paced submission against your Tasks rate limit (runs per minute) -- Crash-safe resume from an append-only run log +- Interrupted-batch recovery with an append-only run log - Task Group sharding with `refresh_status=False` - JSONL export with per-field basis and a validation report - Notes on what the API will not do (no cancel, rate limit is intake not throughput) diff --git a/python-recipes/test_large_scale_tasks_recipe.py b/python-recipes/test_large_scale_tasks_recipe.py new file mode 100644 index 0000000..e1c626c --- /dev/null +++ b/python-recipes/test_large_scale_tasks_recipe.py @@ -0,0 +1,287 @@ +"""Offline regressions: python -m unittest discover -s python-recipes -p 'test_large_scale_tasks_recipe.py'.""" + +import contextlib +import importlib.util +import io +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import httpx +from parallel import APITimeoutError, Parallel + + +spec = importlib.util.spec_from_file_location( + "recipe", os.environ.get("RECIPE_UNDER_TEST", str(Path(__file__).with_name("Large_Scale_Tasks_Recipe.py"))) +) +recipe = importlib.util.module_from_spec(spec) +spec.loader.exec_module(recipe) + + +class RecipeTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.root = Path(self.tmp.name) + self.csv = self.root / "rows.csv" + self.csv.write_text("row_id,name\na,Alpha\nb,Beta\nc,Gamma\n") + self.task_spec = self.root / "spec.json" + self.task_spec.write_text('{"output_schema":"Research this company"}') + self.work = self.root / "job" + self.output = self.root / "results.jsonl" + self.groups = {} + self.posts = [] + self.mode = "success" + self.stream_error = False + self.clients = [] + self.addCleanup(lambda: [c.close() for c in self.clients]) + self.stack = contextlib.ExitStack() + self.addCleanup(self.stack.close) + self.stack.enter_context(contextlib.redirect_stdout(io.StringIO())) + self.stack.enter_context(contextlib.redirect_stderr(io.StringIO())) + self.stack.enter_context(patch.object(recipe.time, "sleep")) + self.stack.enter_context(patch("parallel.Parallel", side_effect=self.client)) + + def client(self, **kwargs): + client = Parallel( + api_key="offline-test-key", + http_client=httpx.Client(transport=httpx.MockTransport(self.request)), + **kwargs, + ) + self.clients.append(client) + return client + + def request(self, request): + path = request.url.path + if request.method == "POST" and path.endswith("/groups"): + gid = f"group-{len(self.groups)}" + self.groups[gid] = [] + return httpx.Response(200, json={"taskgroup_id": gid, "status": self.status(gid)}) + gid = path.split("/")[-2] if path.endswith("/runs") else path.split("/")[-1] + if request.method == "POST": + self.assertEqual(request.url.params.get("refresh_status"), "false") + self.posts.append(json.loads(request.content)) + inputs = self.posts[-1]["inputs"] + if self.mode == "partial": + inputs = inputs[:1] + ids = [] + for value in inputs: + rid = f"{gid}-run-{len(self.groups[gid])}" + ids.append(rid) + self.groups[gid].append({ + "type": "task_run.state", "event_id": rid, + "run": { + "run_id": rid, "interaction_id": rid, "processor": "core", + "status": "completed", "is_active": False, + "metadata": value["metadata"], "taskgroup_id": gid, + }, + "input": value, + "output": {"type": "json", "content": {"name": value["input"]["name"]}, "basis": []}, + }) + if self.mode == "timeout": + raise httpx.ReadTimeout("response lost after accepting runs", request=request) + return httpx.Response(200, json={"run_ids": ids, "status": self.status(gid)}) + if path.endswith("/runs"): + events = list(self.groups[gid]) + if self.stream_error: + events.append({"type": "error", "error": {"message": "stream interrupted"}}) + data = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + return httpx.Response(200, text=data, headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json={"taskgroup_id": gid, "status": self.status(gid)}) + + def status(self, gid): + return {"is_active": False, "num_task_runs": len(self.groups[gid]), + "task_run_status_counts": {"completed": len(self.groups[gid])}} + + def submit(self, *extra): + recipe.main([ + "submit", "--input", str(self.csv), "--task-spec", str(self.task_spec), + "--processor", "core", "--work-dir", str(self.work), *extra, + ]) + + def export(self): + recipe.main(["export", "--work-dir", str(self.work), "--output", str(self.output), "--include-input"]) + + def report(self): + return json.loads(Path(str(self.output) + ".validation.json").read_text()) + + def test_lost_response_is_not_retried_and_resume_recovers_original_runs(self): + self.mode = "timeout" + with self.assertRaises(APITimeoutError): + self.submit() + self.assertEqual(len(self.posts), 1) + self.mode = "success" + self.submit() + self.assertEqual(len(self.posts), 1) + self.assertEqual(len(recipe.Job(str(self.work)).submitted()), 3) + self.export() + self.assertTrue(self.report()["ok"]) + + def test_crash_before_receipt_recovers_by_metadata_not_stream_order(self): + with patch.object(recipe.Job, "add_runs", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + self.groups["group-0"].reverse() + self.submit() + self.assertEqual(len(self.posts), 1) + self.export() + rows = [json.loads(line) for line in self.output.read_text().splitlines()] + self.assertEqual({r["row_id"]: r["output"]["name"] for r in rows}, {"a": "Alpha", "b": "Beta", "c": "Gamma"}) + + def test_partial_receipt_and_crash_after_receipt_do_not_duplicate(self): + original = recipe.Job.add_runs + def interrupted(job, gid, pairs): + original(job, gid, pairs[:1]) + raise KeyboardInterrupt + with patch.object(recipe.Job, "add_runs", interrupted): + with self.assertRaises(KeyboardInterrupt): + self.submit() + with patch.object(recipe.Job, "finish_batch", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + self.submit() + job = recipe.Job(str(self.work)) + self.assertEqual(len(self.posts), 1) + self.assertEqual(len(job.runs_path.read_text().splitlines()), 3) + self.assertFalse(job.pending_path.exists()) + + def test_partial_server_acceptance_stays_blocked_on_resume(self): + self.mode = "partial" + with self.assertRaises(SystemExit): + self.submit() + self.mode = "success" + with self.assertRaisesRegex(SystemExit, "unresolved batch"): + self.submit() + self.assertEqual(len(self.posts), 1) + self.assertTrue((self.work / "pending.json").exists()) + + def test_checkpoint_precedes_post_and_missing_remote_rows_are_not_retried(self): + original = recipe.Job.begin_batch + def interrupted(job, gid, rows): + original(job, gid, rows) + raise KeyboardInterrupt + with patch.object(recipe.Job, "begin_batch", interrupted): + with self.assertRaises(KeyboardInterrupt): + self.submit() + with self.assertRaisesRegex(SystemExit, "unresolved batch"): + self.submit() + self.assertEqual(self.posts, []) + + def test_duplicate_remote_rows_and_stream_errors_block_resume(self): + with patch.object(recipe.Job, "add_runs", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + self.stream_error = True + with self.assertRaisesRegex(SystemExit, "stream error"): + self.submit() + self.stream_error = False + self.groups["group-0"].append(self.groups["group-0"][0]) + with self.assertRaisesRegex(SystemExit, "duplicate runs"): + self.submit() + self.assertEqual(len(self.posts), 1) + + def test_incomplete_input_export_fails_even_when_all_submitted_runs_finished(self): + with patch.object(recipe.time, "sleep", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit("--runs-per-group", "2") + with self.assertRaises(SystemExit) as exc: + self.export() + self.assertEqual(exc.exception.code, 2) + self.assertEqual(self.report()["expected_rows"], 3) + self.assertEqual(self.report()["missing_row_ids"], ["c"]) + self.submit("--runs-per-group", "2") + self.export() + self.assertTrue(self.report()["ok"]) + self.assertEqual([len(p["inputs"]) for p in self.posts], [2, 1]) + + def test_dry_run_needs_no_client_and_zero_submission_export_fails(self): + with patch("parallel.Parallel", side_effect=AssertionError("dry-run must not need credentials")): + self.submit("--dry-run") + with self.assertRaises(SystemExit) as exc: + self.export() + self.assertEqual(exc.exception.code, 2) + self.assertEqual(self.report()["missing"], 3) + self.assertEqual(self.posts, []) + + def test_export_rejects_legacy_job_without_original_manifest(self): + self.work.mkdir() + (self.work / "config.json").write_text("{}") + with self.assertRaisesRegex(SystemExit, "original row manifest"): + self.export() + + def test_export_stays_invalid_until_pending_batch_is_reconciled(self): + with patch.object(recipe.Job, "finish_batch", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + with self.assertRaises(SystemExit): + self.export() + self.assertEqual(self.report()["written"], 3) + self.assertTrue(self.report()["unresolved_batch"]) + self.submit() + self.export() + self.assertTrue(self.report()["ok"]) + + def test_export_active_duplicate_unexpected_and_missing_runs(self): + self.submit() + originals = list(self.groups["group-0"]) + for case in ("active", "duplicate", "unexpected", "missing"): + with self.subTest(case=case): + self.groups["group-0"] = json.loads(json.dumps(originals)) + events = self.groups["group-0"] + if case == "active": + events[0]["run"].update(status="running", is_active=True) + elif case == "duplicate": + events.append(events[0]) + elif case == "unexpected": + events[0]["run"]["run_id"] = "unknown" + else: + events.pop() + with self.assertRaises(SystemExit) as exc: + self.export() + self.assertEqual(exc.exception.code, 2) + self.assertFalse(self.report()["ok"]) + + def test_submit_status_export_roundtrip_and_rerun(self): + self.submit() + self.submit() + recipe.main(["status", "--work-dir", str(self.work), "--wait"]) + self.export() + self.export() + self.assertEqual(len(self.posts), 1) + self.assertTrue(self.report()["ok"]) + self.assertEqual(self.report()["written"], 3) + + def test_request_cap_and_group_sharding_preserve_every_row(self): + self.csv.write_text("row_id,name\n" + "".join(f"{i},Company {i}\n" for i in range(1003))) + self.submit("--runs-per-group", "1001") + self.assertEqual([len(p["inputs"]) for p in self.posts], [1000, 1, 2]) + self.assertEqual([len(runs) for runs in self.groups.values()], [1001, 2]) + self.export() + self.assertTrue(self.report()["ok"]) + self.assertEqual(self.report()["written"], 1003) + + def test_torn_receipt_stops_resume_without_another_post(self): + with patch.object(recipe.Job, "add_runs", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + (self.work / "runs.jsonl").write_text('{"row_id":') + with self.assertRaises(json.JSONDecodeError): + self.submit() + self.assertEqual(len(self.posts), 1) + self.assertTrue((self.work / "pending.json").exists()) + + def test_conflicting_receipt_stops_recovery(self): + with patch.object(recipe.Job, "finish_batch", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + self.submit() + recipe.Job(str(self.work)).add_runs("group-0", [("a", "different-run")]) + with self.assertRaisesRegex(SystemExit, "conflicting receipt"): + self.submit() + self.assertEqual(len(self.posts), 1) + + +if __name__ == "__main__": + unittest.main()