src/service.py

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)