Lightning-AI/pytorch-lightning · error · ModuleNotFoundError

You are trying to use `ScheduleWrapper` which require kineto

Error message

You are trying to use `ScheduleWrapper` which require kineto install.

What it means

ScheduleWrapper wraps a torch.profiler schedule for step-level recording with the PyTorch (kineto/LibTorch) profiler. It requires torch.profiler.profiler.ProfilerAction, which only exists when kineto is available in the installed PyTorch build; if the availability check fails, __init__ raises ModuleNotFoundError.

Source

Thrown at src/lightning/pytorch/profilers/pytorch.py:109

                    partial(self._stop_recording_forward, record_name=record_name)
                )

                self._handles[module_name] = [pre_forward_handle, post_forward_handle]

    def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
        for handles in self._handles.values():
            for h in handles:
                h.remove()
        self._handles = {}


class ScheduleWrapper:
    """This class is used to override the schedule logic from the profiler and perform recording for both
    `training_step`, `validation_step`."""

    def __init__(self, schedule: Callable) -> None:
        if not _KINETO_AVAILABLE:
            raise ModuleNotFoundError("You are trying to use `ScheduleWrapper` which require kineto install.")
        self._schedule = schedule
        self.reset()

    def reset(self) -> None:
        # handle properly `fast_dev_run`. PyTorch Profiler will fail otherwise.
        self._num_training_step = 0
        self._num_validation_step = 0
        self._num_test_step = 0
        self._num_predict_step = 0
        self._training_step_reached_end = False
        self._validation_step_reached_end = False
        self._test_step_reached_end = False
        self._predict_step_reached_end = False
        # used to stop profiler when `ProfilerAction.RECORD_AND_SAVE` is reached.
        self._current_action: Optional[str] = None
        self._prev_schedule_action: Optional[ProfilerAction] = None
        self._start_action_name: Optional[str] = None

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Install a standard PyTorch build that includes kineto (standard pip/conda wheels do)
  2. Upgrade PyTorch to a recent version
  3. If kineto is unavailable in your environment, use AdvancedProfiler or SimpleProfiler instead of PyTorchProfiler

Example fix

# before
profiler = PyTorchProfiler(schedule=torch.profiler.schedule(wait=1, warmup=1, active=3))

# after (no kineto available)
from lightning.pytorch.profilers import AdvancedProfiler
profiler = AdvancedProfiler()
Defensive patterns

Strategy: type-guard

Validate before calling

from lightning.pytorch.profilers.pytorch import _KINETO_AVAILABLE
if not _KINETO_AVAILABLE:
    profiler = AdvancedProfiler()  # fallback
else:
    profiler = PyTorchProfiler(...)

Type guard

def has_kineto() -> bool:
    try:
        from torch.profiler.profiler import ProfilerAction  # noqa: F401
        return True
    except ImportError:
        return False

Try / catch

try:
    profiler = PyTorchProfiler(schedule=torch.profiler.schedule(wait=1, warmup=1, active=3))
except ModuleNotFoundError:
    profiler = AdvancedProfiler()

Prevention

When it happens

Trigger: Constructing ScheduleWrapper directly, or creating PyTorchProfiler with a schedule kwarg, on a PyTorch build without kineto support (e.g. some stripped/older builds or restricted environments).

Common situations: Running on unusual PyTorch installations (custom builds, very old versions, some mobile/embedded builds) where torch.profiler kineto components are absent; using lightning on such an environment with the default PyTorchProfiler setup.

Understand the failure class

Background: "X is not installed. Please install it with pip install Y": missing optional dependency errors — ImportError/ValueError raised when a library's optional extra was never installed — this error's family across 22 libraries.

Related errors


AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/b0b6f74867ed3d0a. Report an issue: GitHub.