Hooks in the engine core and the workers

    def get_scheduler_cls(self) -> type["SchedulerInterface"]:
        if self.scheduler_cls is None:
            if self.async_scheduling:
                from vllm.v1.core.sched.async_scheduler import AsyncScheduler

                return AsyncScheduler
            from vllm.v1.core.sched.scheduler import Scheduler

            return Scheduler

        # The first half of this warning can be removed once the Scheduler interface is
        # finalized and we can maintain support for scheduler classes that implement it
        if not (
            isinstance(self.scheduler_cls, str)
            and self.scheduler_cls.startswith("vllm.")
        ):
            logger.warning_once(
                "Using custom scheduler class %s. This scheduler interface is not "
                "public and compatibility may not be maintained. If you have "
                "subclassed Scheduler instead of AsyncScheduler, you will see "
                "degraded performance due to async scheduling being disabled.",
                self.scheduler_cls,  # type: ignore[arg-type]
            )
        if not isinstance(self.scheduler_cls, str):
            return cast(type["SchedulerInterface"], self.scheduler_cls)
        return resolve_obj_by_qualname(self.scheduler_cls)

class DummyV1Scheduler(Scheduler):
    def schedule(self, throttle_prefills: bool = False):
        raise Exception("Exception raised by DummyV1Scheduler")


def test_scheduler_plugins_v1(monkeypatch: pytest.MonkeyPatch):
    with monkeypatch.context() as m:
        # Explicitly turn off engine multiprocessing so
        # that the scheduler runs in this process
        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

        with pytest.raises(Exception) as exception_info:
            engine_args = EngineArgs(
                model="facebook/opt-125m",
                enforce_eager=True,  # reduce test time
                scheduler_cls=DummyV1Scheduler,
            )

            engine = LLMEngine.from_engine_args(engine_args=engine_args)

            sampling_params = SamplingParams(max_tokens=1)
            engine.add_request("0", "foo", sampling_params)
            engine.step()

        assert str(exception_info.value) == "Exception raised by DummyV1Scheduler"

        parallel_config = vllm_config.parallel_config
        if isinstance(parallel_config.worker_cls, str):
            worker_class: type[WorkerBase] = resolve_obj_by_qualname(
                parallel_config.worker_cls
            )
        else:
            raise ValueError(
                "passing worker_cls is no longer supported. "
                "Please pass keep the class in a separate module "
                "and pass the qualified name of the class as a string."
            )

        worker_extension_cls: type | None = None
        if parallel_config.worker_extension_cls:
            worker_extension_cls = resolve_obj_by_qualname(
                parallel_config.worker_extension_cls
            )
            extended_calls = []
            if worker_extension_cls not in worker_class.__bases__:
                # check any conflicts between worker and worker_extension_cls
                for attr in dir(worker_extension_cls):
                    if attr.startswith("__"):
                        continue
                    assert not hasattr(worker_class, attr), (
                        f"Worker class {worker_class} already has an attribute"
                        f" {attr}, which conflicts with the worker"
                        f" extension class {worker_extension_cls}."
                    )
                    if callable(getattr(worker_extension_cls, attr)):
                        extended_calls.append(attr)
                # dynamically inherit the worker extension class
                worker_class.__bases__ = worker_class.__bases__ + (
                    worker_extension_cls,
                )
                logger.info(
                    "Injected %s into %s for extended collective_rpc calls %s",
                    worker_extension_cls,
                    worker_class,
                    extended_calls,
                )

def _load_v2_logitsprocs(
    logits_processors: Sequence[str | type] | None,
) -> list[type[LogitsProcessor]]:
    """Load all custom logits processors for the V2 model runner.

    Combines installed plugins (validated as V2 processors) with the
    user-specified list of processor types and FQCN strings.
    """
    return _load_v2_logitsprocs_plugins() + _load_v2_logitsprocs_by_fqcns(
        logits_processors or ()
    )

    def __init__(  # noqa: B027
        self, vllm_config: VllmConfig, req_states: LogitsProcRequestState
    ):
        """Capture what stays constant for the processor's lifetime.

        ``req_states`` exposes the on-device token history and batch constants a
        processor may read. Treat it as read-only.
        """

    @classmethod  # noqa: B027
    def validate_params(cls, sampling_params: SamplingParams) -> None:
        """Raise ``ValueError`` for invalid per-request arguments.

        Runs at request admission, so invalid arguments fail the request
        with an error instead of reaching the sampler.
        """

    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
        """Initialize per-slot state for a request entering the batch.

        The slot may hold a previous occupant's state; overwrite or
        neutralize all of it here.

        Returns whether this processor modifies logits for the request.
        """
        return True

    def apply_staged_writes(self) -> None:  # noqa: B027
        """Flush any host-side writes staged by ``add_request()`` to the device.

        Called once per step before the forward pass, after the model runner
        has flushed ``req_states``, so a processor that stages writes here can
        read the request's tokens back on device.
        """

    @abstractmethod
    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
        """Modify logits in place or return a new tensor. In-place modification
        is preferred for efficiency.

        ``apply()`` is called once for the whole batch, including rows of
        requests this processor declined in ``add_request()``, so filter rows
        via ``ctx.expanded_idx_mapping``.

        Args:
            logits: [num_logits_rows, vocab_size] float32 tensor.
            ctx: this step's batch layout.

        """
        raise NotImplementedError

    def _wrapper(model_loader_cls):
        if load_format in _LOAD_FORMAT_TO_MODEL_LOADER:
            logger.warning(
                "Load format `%s` is already registered, and will be "
                "overwritten by the new loader class `%s`.",
                load_format,
                model_loader_cls,
            )
        if not issubclass(model_loader_cls, BaseModelLoader):
            raise ValueError(
                "The model loader must be a subclass of `BaseModelLoader`."
            )
        _LOAD_FORMAT_TO_MODEL_LOADER[load_format] = model_loader_cls
        logger.info(
            "Registered model loader `%s` with load format `%s`",
            model_loader_cls,
            load_format,
        )
        return model_loader_cls

    return _wrapper


def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
    """Get a model loader based on the load format."""
    load_format = load_config.load_format
    if load_format not in _LOAD_FORMAT_TO_MODEL_LOADER:
        raise ValueError(f"Load format `{load_format}` is not supported")
    return _LOAD_FORMAT_TO_MODEL_LOADER[load_format](load_config)

class SleepModeBackendFactory:
    """Registry and resolver for sleep-mode backends.

    Mirrors ``KVConnectorFactory``: lazy module/class registration and a
    built-in registry populated at import time. Third-party backends register
    the same way from a ``vllm.general_plugins`` entry point.
    """

    _registry: dict[str, Callable[[], type[SleepModeBackend]]] = {}

    @classmethod
    def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
        """Register a backend with a lazy-loading module and class name."""
        if name in cls._registry:
            raise ValueError(f"Sleep-mode backend '{name}' is already registered.")

        def loader() -> type[SleepModeBackend]:
            module = importlib.import_module(module_path)
            return getattr(module, class_name)

        cls._registry[name] = loader

    @classmethod
    def get_backend_class(cls, name: str) -> type[SleepModeBackend]:
        """Resolve a registered backend class by name."""
        if name not in cls._registry:
            available = ", ".join(sorted(cls._registry)) or "<none>"
            raise ValueError(
                f"Unsupported sleep-mode backend '{name}'. "
                f"Registered backends: {available}."
            )
        return cls._registry[name]()

    @classmethod
    def create_backend(cls, model_config: ModelConfig) -> SleepModeBackend:
        """Instantiate the backend selected by ``model_config``."""
        name = model_config.sleep_mode_backend
        backend_cls = cls.get_backend_class(name)
        if not backend_cls.is_supported():
            raise ValueError(
                f"Sleep-mode backend '{name}' is not supported on this platform."
            )
        logger.info("Using sleep-mode backend: %s", name)
        return backend_cls()


# Register built-in backends here. Registration is lazy: only the module for the
# selected backend is imported. Third-party backends (CUDA checkpoint, CRIU,
# durable snapshot) register the same way through a vllm.general_plugins entry
# point, without changes to vLLM core.
SleepModeBackendFactory.register_backend(
    "cumem",
    "vllm.device_allocator.sleep_mode_backend",
    "CuMemBackend",
)

    def _decorate(cls):
        if name in KV_CACHE_DTYPES:
            logger.warning(
                "The kv-cache-dtype '%s' already exists and will be "
                "overwritten by handler %s.",
                name,
                cls,
            )
        else:
            KV_CACHE_DTYPES.append(name)
        inst = cls()
        # Inject into the shared dtype dict in-place so every existing
        # direct [] / `in` access site sees the new entry.
        STR_DTYPE_TO_TORCH_DTYPE[name] = inst.torch_dtype()
        _KV_CACHE_DTYPE_HANDLERS[name] = inst
        return cls

    return _decorate