from __future__ import annotations import json import re import time from contextlib import nullcontext from datetime import datetime, timezone import traceback import zipfile from pathlib import Path import torch from .backend import create_separator, release_accelerators from .config import CPU_MAX_SECONDS, MODEL_ROOT from .downloads import ( file_sha256, CustomModelRecord, ) from .gpu_timing import GpuCallbackTiming, synchronized_wall_time from .errors import BUNDLE_PACKAGING, GPU_INFERENCE, MODEL_LOAD, OUTPUT_ENCODING, PREPARED_JOB_VALIDATION, StageError, stage_error from .preparation import ( BUNDLE_LAYOUT_BY_INPUT, BITRATE_AUTO, parse_prepared_state, public_reproducibility_config, ) from .onnx_session_inspection import inspect_onnx_sessions from .observability import ( DEFAULT_LOG_LEVEL, JobLogCapture, ProgressBridge, ProgressReporter, StructuredTqdmCapture, normalize_log_level, read_log_tail, ) EXECUTION_MODE_INITIAL = "initial" EXECUTION_MODE_RETRY_INCOMPLETE = "retry-incomplete" _CANCEL_REQUEST_NAME = "cancel_request.json" def _json_read_dict(path: Path) -> dict: if not path.is_file(): return {} try: payload = json.loads(path.read_text(encoding="utf-8")) except Exception: return {} return payload if isinstance(payload, dict) else {} def _cancel_request_path(job) -> Path: return job.config_dir / _CANCEL_REQUEST_NAME def request_prepared_job_cancel(prepared_state_value) -> str: """Request cooperative cancellation for a prepared job. Gradio can cancel queued functions, but a regular function already running is allowed to finish. The GPU callback therefore checks this marker between Batch items and stops before starting the next item. """ try: job, _ = parse_prepared_state(prepared_state_value) except StageError as exc: return f"Cancellation was not requested: {exc.message}" payload = { "schema": "sesa-cancel-request-v1", "job_id": job.job_id, "requested_at_utc": datetime.now(timezone.utc).isoformat(), "scope": "stop-before-next-batch-item", } destination = _cancel_request_path(job) temporary = destination.with_suffix(".json.tmp") temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") temporary.replace(destination) return "Cancellation requested. A queued job will be cancelled; a running Batch stops before the next item." def _cancel_requested(job, *, run_started_epoch: float) -> bool: path = _cancel_request_path(job) if not path.is_file(): return False try: # Accept a marker created during the run and a small pre-start window to # cover the race between a queued event starting and the stop click. return path.stat().st_mtime >= run_started_epoch - 5.0 except OSError: return False def _clear_cancel_request(job) -> None: try: _cancel_request_path(job).unlink(missing_ok=True) except OSError: pass def _existing_outputs_from_manifest(job, manifest: dict) -> tuple[list[Path], list[dict]]: paths: list[Path] = [] records: list[dict] = [] root = job.root.resolve() output_root = job.output_dir.resolve() for record in list(manifest.get("outputs") or []): if not isinstance(record, dict): continue relative = str(record.get("relative_path") or "") candidate = (job.root / relative).resolve() if relative else (job.output_dir / str(record.get("name") or "")).resolve() if candidate.is_file() and root in candidate.parents and candidate.parent == output_root: paths.append(candidate) records.append(dict(record)) return paths, records def _retryable_indexes(manifest: dict, input_count: int) -> list[int]: completed = { int(item.get("input_index")) for item in list(manifest.get("items") or []) if isinstance(item, dict) and item.get("status") == "COMPLETED" and item.get("input_index") is not None } return [index for index in range(max(0, int(input_count))) if index not in completed] def _existing_result_response(job, manifest: dict, message: str): outputs, _ = _existing_outputs_from_manifest(job, manifest) first, second = _pick_previews(outputs) batch_manifest = job.config_dir / "batch_manifest.json" returned = [str(path) for path in outputs] if batch_manifest.is_file(): returned.append(str(batch_manifest)) archive = job.bundle_dir / f"sesa_{job.job_id}_results.zip" runtime_log = job.logs_dir / f"sesa_job_{job.job_id}.log" reproducibility = job.config_dir / "reproducibility.json" return ( message, first, second, returned, str(archive) if archive.is_file() else None, read_log_tail(runtime_log) if runtime_log.is_file() else "", str(runtime_log) if runtime_log.is_file() else None, str(reproducibility) if reproducibility.is_file() else None, ) def _pick_previews(outputs: list[Path]) -> tuple[str | None, str | None]: if not outputs: return None, None vocals = next((p for p in outputs if "vocal" in p.name.lower()), None) instrumental = next( (p for p in outputs if any(word in p.name.lower() for word in ("instrumental", "karaoke", "no_vocals"))), None, ) first = vocals or outputs[0] second = instrumental or next((p for p in outputs if p != first), None) return str(first) if first else None, str(second) if second else None _STEM_TOKENS = ( "instrumental", "vocals", "drums", "bass", "other", "guitar", "piano", "karaoke", "no_vocals", "reverb", "dry", "noise", "denoise", ) def _prepared_path(job, relative: str) -> Path: candidate = (job.root / str(relative)).resolve() if job.root.resolve() not in candidate.parents or not candidate.is_file(): raise StageError(PREPARED_JOB_VALIDATION, "A prepared input file is missing or escaped the job workspace.", hint="Prepare the job again.") return candidate def _stem_label(path: Path) -> str: text = path.stem.lower() for token in _STEM_TOKENS: if token in text: return token clean = re.sub(r"[^a-z0-9]+", "_", path.stem.lower()).strip("_") return clean[-40:] or "stem" def _uuidize_item_outputs( job, outputs: list[Path], *, input_index: int, input_record: dict, sequence_start: int, ) -> tuple[list[Path], list[dict], int]: """Rename outputs with an explicit Batch item mapping. Mapping is based on the individual ``separator.separate([input])`` call, not filename heuristics. This remains unambiguous when source names repeat. """ renamed: list[Path] = [] records: list[dict] = [] sequence = max(1, int(sequence_start)) for path in sorted(outputs, key=lambda item: item.name.lower()): label = _stem_label(path) target = job.output_dir / ( f"sesa_{job.job_id}_input_{int(input_index):02d}_{sequence:03d}_{label}" f"{path.suffix.lower()}" ) collision = 1 while target.exists() and target.resolve() != path.resolve(): target = job.output_dir / ( f"sesa_{job.job_id}_input_{int(input_index):02d}_{sequence:03d}_{label}_{collision}" f"{path.suffix.lower()}" ) collision += 1 if target.resolve() != path.resolve(): path.replace(target) renamed.append(target) records.append({ "name": target.name, "relative_path": target.relative_to(job.root).as_posix(), "input_index": int(input_index), "input_display_name": str(input_record.get("display_name") or f"input_{input_index:02d}"), "stem": label, "size_bytes": target.stat().st_size, "sha256": file_sha256(target), "format": target.suffix.lower().lstrip("."), }) sequence += 1 return renamed, records, sequence def _uuidize_outputs(job, outputs: list[Path], prepared_inputs: list[Path]) -> tuple[list[Path], list[dict]]: """Compatibility wrapper for older tests/callers using one inferred batch.""" renamed: list[Path] = [] records: list[dict] = [] sequence = 1 for input_index, prepared in enumerate(prepared_inputs): matching = [path for path in outputs if prepared.stem.lower() in path.stem.lower()] if not matching and input_index == 0: matching = [path for path in outputs if path not in renamed] item_outputs, item_records, sequence = _uuidize_item_outputs( job, matching, input_index=input_index, input_record={"display_name": prepared.name}, sequence_start=sequence, ) renamed.extend(item_outputs) records.extend(item_records) return renamed, records def _write_reproducibility(job, config: dict, *, execution: dict, outputs: list[dict], error: dict | None) -> Path: payload = public_reproducibility_config(config) batch_status = str((execution.get("batch") or {}).get("status") or "") if error is not None: status = "FAILED_EXECUTION" elif batch_status in {"COMPLETED_WITH_ITEM_FAILURES", "CANCELED_WITH_PARTIAL_RESULTS", "CANCELED"}: status = batch_status else: status = "COMPLETED" payload["status"] = status payload["execution"] = execution payload["outputs"] = outputs payload["error"] = error payload["completed_at_utc"] = datetime.now(timezone.utc).isoformat() destination = job.config_dir / "reproducibility.json" temporary = destination.with_suffix(".json.tmp") temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8") temporary.replace(destination) return destination def _write_batch_manifest(job, config: dict, execution: dict, outputs: list[dict]) -> Path: batch = dict(execution.get("batch") or {}) payload = { "schema": "sesa-batch-manifest-v2", "job_id": job.job_id, "prepared_batch_plan": config.get("batch"), "status": batch.get("status"), "input_count": batch.get("input_count"), "completed_count": batch.get("completed_count"), "failed_count": batch.get("failed_count"), "canceled_count": batch.get("canceled_count"), "pending_count": batch.get("pending_count"), "model_load_count": batch.get("model_load_count"), "total_model_load_count": batch.get("total_model_load_count"), "execution_order": batch.get("execution_order"), "execution_mode": batch.get("execution_mode"), "attempt_number": batch.get("attempt_number"), "attempts": batch.get("attempts") or [], "retryable_input_indexes": batch.get("retryable_input_indexes") or [], "cooperative_stop_boundary": batch.get("cooperative_stop_boundary"), "items": batch.get("items") or [], "outputs": outputs, "created_at_utc": datetime.now(timezone.utc).isoformat(), } destination = job.config_dir / "batch_manifest.json" temporary = destination.with_suffix(".json.tmp") temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8") temporary.replace(destination) return destination def _archive_prepared_result(job, outputs: list[Path], output_records: list[dict], config: dict, diagnostics: list[Path]) -> Path: archive = job.bundle_dir / f"sesa_{job.job_id}_results.zip" layout = str(config.get("output", {}).get("bundle_layout") or "Flat outputs") with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED, allowZip64=True) as bundle: for path, record in zip(outputs, output_records): if layout == BUNDLE_LAYOUT_BY_INPUT: arcname = f"outputs/input_{int(record.get('input_index', 0)):02d}/{path.name}" else: arcname = f"outputs/{path.name}" bundle.write(path, arcname=arcname) for path in diagnostics: if path.is_file(): bundle.write(path, arcname=f"diagnostics/{path.name}") return archive def _custom_records_from_config(config: dict) -> list[CustomModelRecord]: records: list[CustomModelRecord] = [] for item in config.get("models", {}).get("custom_records", []): model_path = Path(str(item.get("model_path") or "")) config_path = Path(str(item.get("config_path") or "")) if not model_path.is_file() or not config_path.is_file(): raise StageError(MODEL_LOAD, "A prefetched custom model file is unavailable.", hint="Prepare the job again.") records.append(CustomModelRecord( key=str(item.get("key") or model_path.stem), display_name=str(item.get("display_name") or model_path.name), filename=str(item.get("filename") or model_path.name), model_path=model_path, config_path=config_path, architecture=str(item.get("architecture") or "MDXC"), )) return records def run_prepared_job( prepared_state_value, *, progress=None, callback_timing: GpuCallbackTiming | None = None, execution_mode: str = EXECUTION_MODE_INITIAL, ): try: job, config = parse_prepared_state(prepared_state_value) except StageError as exc: status = ( "### Failed before GPU work\n" f"- Stage: `{exc.stage}`\n" f"- Error: `{exc.message}`" ) return status, None, None, [], None, "", None, None mode = str(execution_mode or EXECUTION_MODE_INITIAL) if mode not in {EXECUTION_MODE_INITIAL, EXECUTION_MODE_RETRY_INCOMPLETE}: mode = EXECUTION_MODE_INITIAL prior_manifest = _json_read_dict(job.config_dir / "batch_manifest.json") input_count = len(config.get("inputs") or []) if mode == EXECUTION_MODE_INITIAL and prior_manifest.get("items"): return _existing_result_response( job, prior_manifest, "### Already executed\nUse **Retry failed / unprocessed items** for incomplete Batch items, or run Preparation again for a fresh job.", ) if mode == EXECUTION_MODE_RETRY_INCOMPLETE and prior_manifest.get("items") and not _retryable_indexes(prior_manifest, input_count): return _existing_result_response( job, prior_manifest, "### Nothing to retry\nAll prepared Batch items are already complete.", ) output_cfg = config.get("output", {}) params = config.get("parameters", {}) models = config.get("models", {}) duration_estimate = config.get("duration", {}) level_name = normalize_log_level(params.get("log_level")) runtime_log = job.logs_dir / f"sesa_job_{job.job_id}.log" progress_bridge = ProgressBridge(progress) selected_models = list(models.get("selected_filenames") or []) selected_onnx = any(str(name).lower().endswith(".onnx") for name in selected_models) provider_recheck = { "applicability": "required" if selected_onnx else "not-applicable", "status": "pending" if selected_onnx else "not-applicable", "reason": ( "At least one prepared model is ONNX." if selected_onnx else "No prepared model filename ends with .onnx." ), "available_providers_at_callback_entry": None, "requested_providers": None, "session_registered_providers": None, "session_count": 0, "sessions": [], "cuda_available": None, "cuda_registered_in_session": None, "inspection_errors": [], } execution: dict = { "gpu_callback": callback_timing.to_dict() if callback_timing is not None else None, "preparation_wall_seconds": config.get("preparation_wall_seconds"), "zerogpu_declared_seconds": duration_estimate.get("seconds"), "onnx_provider_recheck": provider_recheck, "execution_mode": mode, } separator = None outputs, output_records = ( _existing_outputs_from_manifest(job, prior_manifest) if mode == EXECUTION_MODE_RETRY_INCOMPLETE else ([], []) ) archive: Path | None = None error_record: dict | None = None current_stage = PREPARED_JOB_VALIDATION using_cuda = False started = time.perf_counter() run_started_epoch = time.time() with JobLogCapture(runtime_log, job.job_id, level=level_name) as job_log: reporter = ProgressReporter(progress_bridge, job_log.event, job_log.debug_event, level=level_name) try: reporter.update(0.02, "Loading frozen Preparation plan", stage="prepared-job", force_info=True) prepared_inputs = [_prepared_path(job, item["prepared_path"]) for item in config.get("inputs", [])] if not prepared_inputs or not selected_models: raise StageError(PREPARED_JOB_VALIDATION, "Prepared job has no inputs or models.", hint="Prepare the job again.") custom_records = _custom_records_from_config(config) job_log.event( "prepared-job-loaded", input_count=len(prepared_inputs), model_count=len(selected_models), config_sha256=config.get("config_sha256"), declared_seconds=duration_estimate.get("seconds"), ) using_cuda = bool(torch.cuda.is_available()) if not using_cuda and not bool(params.get("allow_cpu_fallback")): raise StageError( GPU_INFERENCE, "CUDA was unavailable inside the execution callback and CPU fallback is disabled.", hint="Retry later or enable CPU fallback for a small job.", ) if not using_cuda and ( len(prepared_inputs) > 1 or len(selected_models) > 1 or float(config.get("input_total_seconds") or 0) > CPU_MAX_SECONDS ): raise StageError( GPU_INFERENCE, "CPU fallback is limited to one model, one file, and at most two minutes.", hint="Use ZeroGPU or reduce the job.", ) device_text = "CUDA" if using_cuda else "CPU fallback" device_name = torch.cuda.get_device_name(0) if using_cuda else "CPU" execution["device"] = {"kind": device_text, "name": device_name} job_log.event("device-selected", device=device_text, name=device_name) if selected_onnx: try: import onnxruntime as ort callback_providers = list(ort.get_available_providers()) except Exception as exc: callback_providers = [] provider_recheck["inspection_error"] = f"{type(exc).__name__}: {exc}" provider_recheck["available_providers_at_callback_entry"] = callback_providers provider_recheck["cuda_available"] = "CUDAExecutionProvider" in callback_providers execution["onnx_available_providers_inside_callback"] = callback_providers execution["onnx_cuda_provider_available_inside_callback"] = provider_recheck["cuda_available"] job_log.event("onnx-providers-inside-callback", providers=callback_providers) if not provider_recheck["cuda_available"] and not bool(params.get("allow_cpu_fallback")): provider_recheck["status"] = "failed" raise StageError( GPU_INFERENCE, "An ONNX model is selected but CUDAExecutionProvider is unavailable inside the execution callback.", hint="Retry later, repair ONNX Runtime GPU, or explicitly enable limited CPU fallback.", ) current_stage = MODEL_LOAD reporter.update(0.10, f"Initializing separator on {device_text}", stage="model-load", force_info=True) bitrate = output_cfg.get("bitrate") if bitrate == BITRATE_AUTO: bitrate = None separator = create_separator( output_dir=job.output_dir, model_dir=MODEL_ROOT, custom_models=custom_records, output_format=str(output_cfg.get("format") or "FLAC").upper(), ensemble_algorithm=str(models.get("ensemble_algorithm") or "avg_wave"), single_stem=None if output_cfg.get("single_stem") == "All stems" else output_cfg.get("single_stem"), pitch_shift=int(params.get("pitch_shift") or 0), chunk_duration=int(params.get("chunk_duration") or 0) or None, output_bitrate=bitrate, sample_rate=int(output_cfg.get("sample_rate") or 44100), normalization_threshold=float(output_cfg.get("normalization_threshold", 0.9)), amplification_threshold=float(output_cfg.get("amplification_threshold", 0.0)), log_level=level_name, job_id=job.job_id, ) _, model_load_seconds, model_load_sync_error = synchronized_wall_time( lambda: separator.load_model(selected_models if len(selected_models) > 1 else selected_models[0]), use_cuda=using_cuda, ) execution["model_load_synchronized_wall_seconds"] = round(model_load_seconds, 6) execution["model_load_cuda_sync_error"] = model_load_sync_error job_log.event("model-load-complete", elapsed_seconds=f"{model_load_seconds:.6f}", sync_error=model_load_sync_error) if selected_onnx: session_inspection = inspect_onnx_sessions(separator) provider_recheck.update(session_inspection) provider_recheck["cuda_registered_in_session"] = ( "CUDAExecutionProvider" in session_inspection["session_registered_providers"] ) if session_inspection["session_count"] == 0: provider_recheck["status"] = "unobserved" raise StageError( MODEL_LOAD, "The prepared ONNX model loaded, but its loaded InferenceSession registered providers could not be observed.", hint="Review audio-separator internals or the ONNX session inspector before accepting this job.", ) if provider_recheck["cuda_registered_in_session"]: provider_recheck["status"] = "passed" elif bool(params.get("allow_cpu_fallback")): provider_recheck["status"] = "cpu-fallback" else: provider_recheck["status"] = "failed" raise StageError( MODEL_LOAD, "The loaded ONNX InferenceSession is not using CUDAExecutionProvider.", hint="Repair ONNX Runtime GPU/provider ordering or explicitly enable limited CPU fallback.", ) job_log.event( "onnx-session-registered-providers", providers=session_inspection["session_registered_providers"], session_count=session_inspection["session_count"], status=provider_recheck["status"], ) current_stage = GPU_INFERENCE batch_plan = dict(config.get("batch") or {}) continue_on_item_error = bool(batch_plan.get("continue_on_item_error", True)) input_records = list(config.get("inputs") or []) prior_items_by_index = { int(item.get("input_index")): dict(item) for item in list(prior_manifest.get("items") or []) if isinstance(item, dict) and item.get("input_index") is not None } batch_items_by_index = dict(prior_items_by_index) if mode == EXECUTION_MODE_RETRY_INCOMPLETE else {} target_input_indexes = ( _retryable_indexes(prior_manifest, len(prepared_inputs)) if mode == EXECUTION_MODE_RETRY_INCOMPLETE else list(range(len(prepared_inputs))) ) total_inference_seconds = 0.0 inference_sync_errors: list[str] = [] output_sequence = len(output_records) + 1 tqdm_capture = StructuredTqdmCapture(job_log.debug_event) if level_name == "DEBUG" else nullcontext() previous_attempts = [dict(item) for item in list(prior_manifest.get("attempts") or []) if isinstance(item, dict)] attempt_number = len(previous_attempts) + 1 attempt_started_utc = datetime.now(timezone.utc).isoformat() attempted_indexes: list[int] = [] cancel_observed = False for target_position, input_index in enumerate(target_input_indexes): if _cancel_requested(job, run_started_epoch=run_started_epoch): cancel_observed = True job_log.event("batch-cancel-observed", before_input_index=input_index, attempt_number=attempt_number) break prepared_input = prepared_inputs[input_index] source_record = input_records[input_index] if input_index < len(input_records) else {} item_started = time.perf_counter() item_progress = 0.28 + (0.52 * target_position / max(1, len(target_input_indexes))) item_record = { "input_index": input_index, "display_name": str(source_record.get("display_name") or prepared_input.name), "prepared_path": str(source_record.get("prepared_path") or prepared_input.name), "prepared_sha256": source_record.get("prepared_sha256"), "prepared_duration_seconds": source_record.get("prepared_duration_seconds"), "status": "RUNNING", "attempt_number": attempt_number, "outputs": [], "error": None, } attempted_indexes.append(input_index) job_log.event( "batch-item-start", input_index=input_index, input_name=item_record["display_name"], total_inputs=len(prepared_inputs), attempt_number=attempt_number, execution_mode=mode, ) try: with reporter.long_operation( item_progress, f"Separating Batch item {input_index + 1}/{len(prepared_inputs)}: {item_record['display_name']}", stage="gpu-inference", ): with tqdm_capture: generated, item_inference_seconds, item_sync_error = synchronized_wall_time( lambda path=prepared_input: separator.separate([str(path)]), use_cuda=using_cuda, ) total_inference_seconds += float(item_inference_seconds) if item_sync_error: inference_sync_errors.append(str(item_sync_error)) current_stage = OUTPUT_ENCODING candidates: list[Path] = [] for generated_item in generated or []: candidate = Path(generated_item) if not candidate.is_absolute(): candidate = job.output_dir / candidate candidate = candidate.resolve() if candidate.is_file() and job.output_dir.resolve() in candidate.parents: candidates.append(candidate) candidates = sorted(set(candidates), key=lambda path: path.name.lower()) if not candidates: raise StageError( OUTPUT_ENCODING, f"Batch item {input_index + 1} returned no readable output files.", hint="Review the runtime log and model selection.", ) item_outputs, item_output_records, output_sequence = _uuidize_item_outputs( job, candidates, input_index=input_index, input_record=source_record, sequence_start=output_sequence, ) outputs.extend(item_outputs) output_records.extend(item_output_records) item_record.update({ "status": "COMPLETED", "inference_synchronized_wall_seconds": round(float(item_inference_seconds), 6), "inference_cuda_sync_error": item_sync_error, "output_count": len(item_output_records), "outputs": item_output_records, }) job_log.event( "batch-item-complete", input_index=input_index, output_count=len(item_output_records), elapsed_seconds=f"{float(item_inference_seconds):.6f}", attempt_number=attempt_number, ) except Exception as item_exc: staged_item = stage_error(current_stage, item_exc) item_error_name = ( f"sesa_{job.job_id}_batch_item_{input_index:02d}_error.txt" if attempt_number == 1 else f"sesa_{job.job_id}_batch_item_{input_index:02d}_attempt_{attempt_number:02d}_error.txt" ) item_error_path = job.logs_dir / item_error_name item_error_path.write_text(traceback.format_exc(limit=30), encoding="utf-8") item_record.update({ "status": "FAILED", "output_count": 0, "error": staged_item.to_dict(), "error_log": item_error_path.relative_to(job.root).as_posix(), }) job_log.event( "batch-item-failed", input_index=input_index, error_stage=staged_item.stage, error_type=type(item_exc).__name__, error=str(item_exc), attempt_number=attempt_number, ) if not continue_on_item_error: raise finally: item_record["item_wall_seconds"] = round(time.perf_counter() - item_started, 6) batch_items_by_index[input_index] = item_record current_stage = GPU_INFERENCE if cancel_observed: for input_index in target_input_indexes: if input_index in attempted_indexes: continue source_record = input_records[input_index] if input_index < len(input_records) else {} batch_items_by_index[input_index] = { "input_index": input_index, "display_name": str(source_record.get("display_name") or prepared_inputs[input_index].name), "prepared_path": str(source_record.get("prepared_path") or prepared_inputs[input_index].name), "prepared_sha256": source_record.get("prepared_sha256"), "prepared_duration_seconds": source_record.get("prepared_duration_seconds"), "status": "CANCELED_BEFORE_START", "attempt_number": attempt_number, "output_count": 0, "outputs": [], "error": None, } batch_items = [ batch_items_by_index.get(index, { "input_index": index, "display_name": str((input_records[index] if index < len(input_records) else {}).get("display_name") or prepared_inputs[index].name), "status": "NOT_RUN", "output_count": 0, "outputs": [], "error": None, }) for index in range(len(prepared_inputs)) ] completed_count = sum(1 for item in batch_items if item.get("status") == "COMPLETED") failed_count = sum(1 for item in batch_items if item.get("status") == "FAILED") canceled_count = sum(1 for item in batch_items if str(item.get("status") or "").startswith("CANCELED")) pending_count = len(prepared_inputs) - completed_count - failed_count - canceled_count if completed_count == len(prepared_inputs): batch_status = "COMPLETED" elif cancel_observed: batch_status = "CANCELED_WITH_PARTIAL_RESULTS" if completed_count else "CANCELED" elif completed_count == 0: batch_status = "FAILED" else: batch_status = "COMPLETED_WITH_ITEM_FAILURES" attempt_record = { "attempt_number": attempt_number, "mode": mode, "started_at_utc": attempt_started_utc, "completed_at_utc": datetime.now(timezone.utc).isoformat(), "target_input_indexes": target_input_indexes, "attempted_input_indexes": attempted_indexes, "status": batch_status, "cancel_observed": cancel_observed, "completed_input_indexes": [int(item["input_index"]) for item in batch_items if item.get("status") == "COMPLETED"], "failed_input_indexes": [int(item["input_index"]) for item in batch_items if item.get("status") == "FAILED"], "canceled_input_indexes": [int(item["input_index"]) for item in batch_items if str(item.get("status") or "").startswith("CANCELED")], "model_load_count": 1, } attempts = [*previous_attempts, attempt_record] execution["batch"] = { "schema": "sesa-batch-execution-v2", "status": batch_status, "enabled": len(prepared_inputs) > 1, "input_count": len(prepared_inputs), "completed_count": completed_count, "failed_count": failed_count, "canceled_count": canceled_count, "pending_count": pending_count, "model_load_count": 1, "total_model_load_count": sum(int(item.get("model_load_count") or 0) for item in attempts), "execution_order": "sequential-inputs-shared-loaded-model", "continue_on_item_error": continue_on_item_error, "output_mapping": "explicit-per-separate-call", "execution_mode": mode, "attempt_number": attempt_number, "attempts": attempts, "retryable_input_indexes": [int(item["input_index"]) for item in batch_items if item.get("status") != "COMPLETED"], "cooperative_stop_boundary": "between-batch-items", "items": batch_items, } if completed_count == 0 and not cancel_observed: raise StageError( GPU_INFERENCE, "All Batch items failed.", hint="Review the per-item error records in the runtime bundle.", ) execution["inference_synchronized_wall_seconds"] = round(total_inference_seconds, 6) execution["inference_cuda_sync_error"] = inference_sync_errors or None execution["gpu_synchronized_work_window_seconds"] = round(model_load_seconds + total_inference_seconds, 6) job_log.event( "batch-complete", input_count=len(prepared_inputs), completed_count=completed_count, failed_count=failed_count, output_count=len(output_records), elapsed_seconds=f"{total_inference_seconds:.6f}", ) reporter.update(0.86, "Batch outputs mapped and UUID-renamed", stage="output-encoding", force_info=True) job_log.event("outputs-ready", count=len(outputs), bytes=sum(path.stat().st_size for path in outputs)) reporter.finish("GPU work completed; packaging diagnostics", stage="gpu-execution") except Exception as exc: if provider_recheck["applicability"] == "required" and provider_recheck["status"] == "pending": provider_recheck["status"] = "not-completed" provider_recheck["reason"] = ( "ONNX provider recheck did not complete before the execution failure." ) staged = stage_error(current_stage, exc) error_record = staged.to_dict() detail = traceback.format_exc(limit=40) error_path = job.logs_dir / f"sesa_{job.job_id}_error.txt" error_path.write_text(detail, encoding="utf-8") job_log.event("job-failed", error_stage=staged.stage, error_type=type(exc).__name__, error=str(exc)) outputs = [error_path] finally: release_started = time.perf_counter() release_accelerators(separator) execution["accelerator_release_wall_seconds"] = round(time.perf_counter() - release_started, 6) execution["callback_body_wall_seconds"] = round(time.perf_counter() - started, 6) if callback_timing is not None: callback_timing.checkpoint("prepared_job_body_complete") execution["gpu_callback"] = callback_timing.to_dict( body_seconds=callback_timing.elapsed() ) job_log.event("accelerators-released", cuda=using_cuda, elapsed_seconds=execution["accelerator_release_wall_seconds"]) batch_manifest = _write_batch_manifest(job, config, execution, output_records) reproducibility = _write_reproducibility(job, config, execution=execution, outputs=output_records, error=error_record) diagnostics = [ runtime_log, reproducibility, batch_manifest, job.config_dir / "preflight.json", job.logs_dir / f"sesa_prepare_{job.job_id}.log", *sorted(job.logs_dir.glob(f"sesa_{job.job_id}_batch_item_*_error.txt")), *[path for path in outputs if path.parent != job.output_dir], ] try: archive = _archive_prepared_result(job, [path for path in outputs if path.parent == job.output_dir], output_records, config, diagnostics) except Exception as exc: packaging_error = stage_error(BUNDLE_PACKAGING, exc) error_record = packaging_error.to_dict() reproducibility = _write_reproducibility(job, config, execution=execution, outputs=output_records, error=error_record) archive = None _clear_cancel_request(job) log_tail = read_log_tail(runtime_log) if error_record is None: first, second = _pick_previews([path for path in outputs if path.parent == job.output_dir]) batch_execution = execution.get("batch") or {} batch_status = str(batch_execution.get("status") or "COMPLETED") heading = { "COMPLETED": "### Completed\n", "COMPLETED_WITH_ITEM_FAILURES": "### Completed with item failures\n", "CANCELED_WITH_PARTIAL_RESULTS": "### Stopped with partial results\n", "CANCELED": "### Stopped before any item completed\n", }.get(batch_status, f"### {batch_status}\n") status = ( heading + f"- Job: `{job.job_id}`\n" f"- Device: {execution.get('device', {}).get('kind')}\n" f"- Execution mode: {batch_execution.get('execution_mode')}\n" f"- Attempt: {batch_execution.get('attempt_number')}\n" f"- Inputs: {len(config.get('inputs', []))}\n" f"- Completed/failed/canceled: {batch_execution.get('completed_count', 0)}/" f"{batch_execution.get('failed_count', 0)}/{batch_execution.get('canceled_count', 0)}\n" f"- Retryable input indexes: {batch_execution.get('retryable_input_indexes', [])}\n" f"- Models loaded this attempt/total: {batch_execution.get('model_load_count', 0)}/" f"{batch_execution.get('total_model_load_count', 0)}\n" f"- Continue after item failure: {'ON' if batch_execution.get('continue_on_item_error') else 'OFF'}\n" f"- Output files: {len(output_records)}\n" f"- ZeroGPU declaration: {duration_estimate.get('seconds')} seconds" ) returned_files = [str(path) for path in outputs] if batch_manifest.is_file(): returned_files.append(str(batch_manifest)) return status, first, second, returned_files, str(archive) if archive else None, log_tail, str(runtime_log), str(reproducibility) status = ( "### Failed\n" f"- Stage: `{error_record.get('stage')}`\n" f"- Error: `{error_record.get('message')}`\n" f"- Job: `{job.job_id}`" ) returned_files = [str(path) for path in outputs] if batch_manifest.is_file(): returned_files.append(str(batch_manifest)) return status, None, None, returned_files, str(archive) if archive else None, log_tail, str(runtime_log), str(reproducibility)