Lightning-AI/pytorch-lightning · error · NotImplementedError

The `serve_step` method needs to be overridden.

Error message

The `serve_step` method needs to be overridden.

What it means

The serve_step method is the core inference hook of ServableModule; the validator refuses to launch the test server if it is not overridden. is_overridden('serve_step', ...) fails in on_train_start and raises NotImplementedError.

Source

Thrown at src/lightning/pytorch/serve/servable_module_validator.py:92

    @override
    @rank_zero_only
    def on_train_start(self, trainer: "pl.Trainer", servable_module: "pl.LightningModule") -> None:
        if isinstance(trainer.strategy, _NOT_SUPPORTED_STRATEGIES):
            raise Exception(
                f"The current strategy {trainer.strategy.__class__.__qualname__} used "
                "by the trainer isn't supported for sanity serving yet."
            )

        if not isinstance(servable_module, ServableModule):
            raise TypeError(f"The provided model should be subclass of {ServableModule.__qualname__}.")

        if not is_overridden("configure_payload", servable_module, ServableModule):
            raise NotImplementedError("The `configure_payload` method needs to be overridden.")
        if not is_overridden("configure_serialization", servable_module, ServableModule):
            raise NotImplementedError("The `configure_serialization` method needs to be overridden.")
        if not is_overridden("serve_step", servable_module, ServableModule):
            raise NotImplementedError("The `serve_step` method needs to be overridden.")

        # Note: The Trainer needs to be detached from the pl_module before starting the process.
        # This would fail during the deepcopy with DDP.
        servable_module.trainer = None

        process = Process(target=self._start_server, args=(servable_module, self.host, self.port, self.optimization))
        process.start()

        servable_module.trainer = trainer

        ready = False
        t0 = time.time()
        while not ready:
            with contextlib.suppress(requests.exceptions.ConnectionError):
                resp = requests.get(f"http://{self.host}:{self.port}/ping")
                ready = resp.status_code == 200
            if time.time() - t0 > self.timeout:
                process.kill()

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Add serve_step(self, **kwargs) to your LightningModule that performs inference and returns a dict
2. If you already have predict_step logic, delegate: def serve_step(self, **kwargs): return self.predict_step(**kwargs)
3. Ensure the return value is a dict (keys must match output serializer names)

Example fix

// before
class MyModel(ServableModule):
    def predict_step(self, batch): ...
// after
class MyModel(ServableModule):
    def serve_step(self, **kwargs):
        x = kwargs["x"]
        return {"output": self(x)}

from lightning.pytorch.utilities import is_overridden
from lightning.pytorch.serve import ServableModule

assert is_overridden("serve_step", model, ServableModule), "serve_step must be overridden"

def has_serve_step(m) -> bool:
    return is_overridden("serve_step", m, ServableModule)

When it happens

Trigger: Model inherits ServableModule without defining serve_step, or defines a differently named method (e.g. predict_step) expecting it to be picked up.

Common situations: Assuming predict_step or validation_step doubles as the serving entry point; renaming methods during refactor so the override check no longer matches.

Related errors

• The `configure_payload` method needs to be overridden.
• The `configure_serialization` method needs to be overridden.
• Unsupported {cls}
• The server didn't start within {self.timeout} seconds.
• Your provided payload {payload} should have a field named "b

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