prediction_profiler

def stage(self, name):
    """Return a context manager that records elapsed time for *name*."""
    return _StageTimer(self, name)

def end_batch(self, batch_size):
    """Finalize the current batch record."""
    wall_clock = (
        time.perf_counter() - self._batch_wall_start
        if self._batch_wall_start is not None
        else 0.0
    )
    stage_sum = sum(self._current_stages.values())
    overhead = max(0.0, wall_clock - stage_sum)

    record = {
        "batch_idx": len(self._batch_records),
        "batch_size": batch_size,
        "wall_clock": wall_clock,
        "overhead": overhead,
        "stages": dict(self._current_stages),
    }

    # Memory snapshot every N batches
    if len(self._batch_records) % self.MEMORY_SNAPSHOT_EVERY_N_BATCHES == 0:
        mem = self._take_memory_snapshot()
        record["memory"] = mem
        for key in self._peak_memory:
            if mem.get(key) is not None:
                self._peak_memory[key] = max(self._peak_memory[key], mem[key])

    self._batch_records.append(record)
    self._total_images += batch_size
    self._current_stages = {}
    self._batch_wall_start = None

def record_finalization(self):
    """Record current stages as one-off finalization metrics.

    Call this instead of ``end_batch()`` for post-loop operations like
    ``result_save`` that should not inflate batch counts or throughput.
    """
    self._finalization_stages.update(self._current_stages)
    self._current_stages = {}
    self._batch_wall_start = None

def save_partial_report(self):
    """Write a per-process partial JSON report to the profiling subdirectory."""
    total_wall = time.perf_counter() - self._start_time
    throughput = self._total_images / total_wall if total_wall > 0 else 0.0

    # Aggregate per-stage totals
    stage_totals = {}
    stage_batch_times = {}
    for rec in self._batch_records:
        for stage_name, elapsed in rec["stages"].items():
            stage_totals[stage_name] = stage_totals.get(stage_name, 0.0) + elapsed
            stage_batch_times.setdefault(stage_name, []).append(elapsed)

    total_stage_time = sum(stage_totals.values())
    stages = {}
    for name, total_s in stage_totals.items():
        pct = (total_s / total_stage_time * 100) if total_stage_time > 0 else 0.0
        stages[name] = {
            "total_s": round(total_s, 6),
            "percentage": round(pct, 2),
            "batch_times": [round(t, 6) for t in stage_batch_times[name]],
        }

    report = {
        "process_idx": self._process_idx,
        "total_images": self._total_images,
        "total_wall_clock_s": round(total_wall, 6),
        "throughput_images_per_sec": round(throughput, 2),
        "stages": stages,
        "peak_memory": {k: round(v, 2) for k, v in self._peak_memory.items()},
        "num_batches": len(self._batch_records),
        "batch_sizes": [r["batch_size"] for r in self._batch_records],
        "script": self._script,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    if self._finalization_stages:
        report["finalization_stages"] = {
            k: round(v, 6) for k, v in self._finalization_stages.items()
        }

    path = os.path.join(
        self._profiling_dir,
        self.PARTIAL_FILENAME_TEMPLATE.format(self._process_idx),
    )
    # The profiling directory is created in __init__, but a concurrent
    # process may rmtree the session dir (label cache rebuilds, session
    # rotation, user clearing output_dir) between profiler construction
    # and report save at end-of-batch.  Re-create defensively so the
    # subprocess doesn't exit non-zero on an otherwise-successful run.
    os.makedirs(self._profiling_dir, exist_ok=True)
    with open(path, "w") as f:
        json.dump(report, f, indent=2)
    logger.info(f"Profiler partial report saved to {path}")
    return path

@staticmethod
def merge_reports(profiling_dir, spawn_times=None):
    """Merge per-process partial reports into a single aggregate report."""
    pattern = os.path.join(
        profiling_dir,
        PredictionProfiler.PARTIAL_FILENAME_TEMPLATE.format("*"),
    )
    partial_files = sorted(glob_module.glob(pattern))

    if not partial_files:
        report = {
            "total_images": 0,
            "total_wall_clock_s": 0.0,
            "throughput_images_per_sec": 0.0,
            "stages": {},
            "per_process": [],
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        os.makedirs(profiling_dir, exist_ok=True)
        out_path = os.path.join(profiling_dir, PredictionProfiler.MERGED_FILENAME)
        with open(out_path, "w") as f:
            json.dump(report, f, indent=2)
        return report

    partials = []
    for pf in partial_files:
        with open(pf) as f:
            partials.append(json.load(f))
    partials.sort(key=lambda p: p["process_idx"])

    # Aggregate totals — include spawn overhead for true end-to-end wall time
    total_images = sum(p["total_images"] for p in partials)
    total_process_wall = sum(p["total_wall_clock_s"] for p in partials)
    spawn_overhead = sum(spawn_times) if spawn_times else 0.0
    total_wall = total_process_wall + spawn_overhead
    throughput = total_images / total_wall if total_wall > 0 else 0.0

    # Aggregate stages
    all_stage_names = set()
    for p in partials:
        all_stage_names.update(p["stages"].keys())

    overall_total = sum(sum(s["total_s"] for s in p["stages"].values()) for p in partials)

    stages = {}
    for name in sorted(all_stage_names):
        combined_times = []
        total_s = 0.0
        for p in partials:
            if name in p["stages"]:
                total_s += p["stages"][name]["total_s"]
                combined_times.extend(p["stages"][name]["batch_times"])

        percentiles = {}
        if combined_times:
            arr = np.array(combined_times)
            for pct in (50, 95, 99):
                percentiles[f"p{pct}"] = round(float(np.percentile(arr, pct)), 6)

        pct = (total_s / overall_total * 100) if overall_total > 0 else 0.0

        stages[name] = {
            "total_s": round(total_s, 6),
            "percentage": round(pct, 2),
            "percentiles": percentiles,
            "num_batches": len(combined_times),
        }

    # Aggregate finalization stages
    finalization_stages = {}
    for p in partials:
        for name, elapsed in p.get("finalization_stages", {}).items():
            finalization_stages[name] = finalization_stages.get(name, 0.0) + elapsed

    # Aggregate peak memory
    peak_memory = {
        "resident_set_size_mb": 0.0,
        "gpu_allocated_mb": 0.0,
        "gpu_reserved_mb": 0.0,
    }
    for p in partials:
        for key in peak_memory:
            peak_memory[key] = max(peak_memory[key], p.get("peak_memory", {}).get(key, 0.0))

    # Spawn times
    process_spawn = None
    if spawn_times:
        process_spawn = {
            "spawn_times_s": [round(t, 6) for t in spawn_times],
            "total_spawn_overhead_s": round(spawn_overhead, 6),
        }

    # Per-process summaries
    per_process = []
    for p in partials:
        per_process.append(
            {
                "process_idx": p["process_idx"],
                "total_images": p["total_images"],
                "total_wall_clock_s": p["total_wall_clock_s"],
                "throughput_images_per_sec": p["throughput_images_per_sec"],
            }
        )

    # Batch timeseries (sampled uniformly per process x stage combination)
    all_batch_times = _build_batch_timeseries(partials, max_entries=200)

    report = {
        "total_images": total_images,
        "total_wall_clock_s": round(total_wall, 6),
        "throughput_images_per_sec": round(throughput, 2),
        "stages": stages,
        "peak_memory": {k: round(v, 2) for k, v in peak_memory.items()},
        "per_process": per_process,
        "batch_timeseries": all_batch_times,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    if finalization_stages:
        report["finalization_stages"] = {k: round(v, 6) for k, v in finalization_stages.items()}
    if process_spawn is not None:
        report["process_spawn"] = process_spawn

    out_path = os.path.join(profiling_dir, PredictionProfiler.MERGED_FILENAME)
    with open(out_path, "w") as f:
        json.dump(report, f, indent=2)
    logger.info(f"Profiler merged report saved to {out_path}")
    return report

@staticmethod
def cleanup_partial_reports(profiling_dir):
    """Remove partial report files, preserving the merged report."""
    pattern = os.path.join(
        profiling_dir,
        PredictionProfiler.PARTIAL_FILENAME_TEMPLATE.format("*"),
    )
    for path in glob_module.glob(pattern):
        os.remove(path)