profiler ¶

def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
    """Create or reuse the process-wide torch profiler controller."""

    existing = get_global_controller()
    if existing is not None:
        if trace_dir:
            logger.info("Reusing existing global torch profiler controller")
        return existing

    if not trace_dir:
        logger.info("Torch profiler disabled; returning no-op controller")
        return TorchProfilerController(None, _DEFAULT_ACTIVITIES, disabled=True)

    logger.info("Profiling enabled. Traces will be saved to: %s", trace_dir)
    logger.info(
        "Profiler config: record_shapes=%s, profile_memory=%s, with_stack=%s, with_flops=%s",
        envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
        envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
        envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
        envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
    )
    logger.info("FASTVIDEO_TORCH_PROFILE_REGIONS=%s", envs.FASTVIDEO_TORCH_PROFILE_REGIONS)

    profiler = torch.profiler.profile(
        activities=_DEFAULT_ACTIVITIES,
        record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
        profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
        with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
        with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
        # No schedule: nothing in the codebase calls profiler.step(), so a
        # wait/warmup schedule never advances and the profiler records nothing.
        # Region toggling gates collection; the single trace exports at stop().
        on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True),
    )
    controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES)
    controller._trace_dir = trace_dir
    controller.start()
    # The trace only exports at stop(); inference paths have no shutdown hook
    # that calls it, so register one. stop() is idempotent.
    import atexit
    atexit.register(controller.stop)
    logger.info("Torch profiler started")
    return controller

def list_profiler_regions() -> list[ProfilerRegion]:
    """Return all registered profiler regions sorted by canonical name."""

    return [_REGISTERED_REGIONS[name] for name in sorted(_REGISTERED_REGIONS)]

@contextlib.contextmanager
def nvtx_range(name: str):
    """Emit one optional NVTX range for an external CUDA profiler.

    ``FASTVIDEO_NVTX_PROFILE=1`` enables the marker. The context manager stays
    a no-op without CUDA so call sites can remain shared with CPU tests.
    """
    enabled = envs.FASTVIDEO_NVTX_PROFILE and torch.cuda.is_available()
    if not enabled:
        yield
        return

    torch.cuda.nvtx.range_push(name)
    try:
        yield
    finally:
        torch.cuda.nvtx.range_pop()

def profile_region(region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Wrap a bound method so it runs inside a profiler region if available."""

    def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:

        @functools.wraps(fn)
        def wrapped(self, *args, **kwargs):
            controller = getattr(self, "profiler_controller", None)
            if controller is None or not controller.has_profiler:
                return fn(self, *args, **kwargs)
            with controller.region(region):
                return fn(self, *args, **kwargs)

        return wrapped

    return decorator

@contextlib.contextmanager
def profiler_region(region: str):
    """Module-level region context manager bound to the global controller.

    A no-op when no profiler is configured — stages can use this without the
    get_global_controller()/nullcontext dance.
    """
    controller = get_global_controller()
    if controller is None:
        yield
        return
    with controller.region(region):
        yield

def register_profiler_region(
    name: str,
    description: str,
) -> None:
    """Register a profiler region so configuration can validate inputs."""

    canonical = _normalize_token(name)
    if canonical in _REGISTERED_REGIONS:
        raise ValueError(f"Profiler region {name!r} is already registered")

    region = ProfilerRegion(
        name=canonical,
        description=description,
    )
    _REGISTERED_REGIONS[canonical] = region

def resolve_profiler_region(name: str) -> ProfilerRegion | None:
    """Return the registered region for ``name`` (long or short form).

    Accepts both the canonical name and the short form without the
    ``profiler_region_`` prefix, so ``REGIONS=inference_denoising`` works.
    """

    canonical = _normalize_token(name)
    region = _REGISTERED_REGIONS.get(canonical)
    if region is None and not canonical.startswith("profiler_region_"):
        region = _REGISTERED_REGIONS.get(f"profiler_region_{canonical}")
    return region