[App] Introduce auto scaler

docs/source-app/api_references.rst

@@ -37,6 +37,7 @@ ___________________
     ~training.LightningTrainerScript
     ~serve.gradio.ServeGradio
     ~serve.serve.ModelInferenceAPI
+    ~auto_scaler.AutoScaler
 
 ----
 

examples/app_server_with_auto_scaler/app.py

@@ -0,0 +1,86 @@
+from typing import Any, List
+
+import torch
+import torchvision
+from pydantic import BaseModel
+
+import lightning as L
+
+
+class RequestModel(BaseModel):
+    image: str  # bytecode
+
+
+class BatchRequestModel(BaseModel):
+    inputs: List[RequestModel]
+
+
+class BatchResponse(BaseModel):
+    outputs: List[Any]
+
+
+class PyTorchServer(L.app.components.PythonServer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(
+            port=L.app.utilities.network.find_free_network_port(),
+            input_type=BatchRequestModel,
+            output_type=BatchResponse,
+            cloud_compute=L.CloudCompute("gpu"),
+        )
+
+    def setup(self):
+        self._device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self._model = torchvision.models.resnet18(pretrained=True).to(self._device)
+
+    def predict(self, requests: BatchRequestModel):
+        transforms = torchvision.transforms.Compose(
+            [
+                torchvision.transforms.Resize(224),
+                torchvision.transforms.ToTensor(),
+                torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+        images = []
+        for request in requests.inputs:
+            image = L.app.components.serve.types.image.Image.deserialize(request.image)
+            image = transforms(image).unsqueeze(0)
+            images.append(image)
+        images = torch.cat(images)
+        images = images.to(self._device)
+        predictions = self._model(images)
+        results = predictions.argmax(1).cpu().numpy().tolist()
+        return BatchResponse(outputs=[{"prediction": pred} for pred in results])
+
+
+class MyAutoScaler(L.app.components.AutoScaler):
+    def scale(self, replicas: int, metrics: dict) -> int:
+        """The default scaling logic that users can override."""
+        # scale out if the number of pending requests exceeds max batch size.
+        max_requests_per_work = self.max_batch_size
+        pending_requests_per_running_or_pending_work = metrics["pending_requests"] / (
+            replicas + metrics["pending_works"]
+        )
+        if pending_requests_per_running_or_pending_work >= max_requests_per_work:
+            return replicas + 1
+
+        # scale in if the number of pending requests is below 25% of max_requests_per_work
+        min_requests_per_work = max_requests_per_work * 0.25
+        pending_requests_per_running_work = metrics["pending_requests"] / replicas
+        if pending_requests_per_running_work < min_requests_per_work:
+            return replicas - 1
+
+        return replicas
+
+
+app = L.LightningApp(
+    MyAutoScaler(
+        PyTorchServer,
+        min_replicas=2,
+        max_replicas=4,
+        autoscale_interval=10,
+        endpoint="predict",
+        input_type=RequestModel,
+        output_type=Any,
+        timeout_batching=1,
+    )
+)

pyproject.toml

@@ -80,6 +80,7 @@ module = [
     "lightning_app.components.serve.types.type",
     "lightning_app.components.serve.python_server",
     "lightning_app.components.training",
+    "lightning_app.components.auto_scaler",
     "lightning_app.core.api",
     "lightning_app.core.app",
     "lightning_app.core.flow",

requirements/app/base.txt

@@ -12,3 +12,4 @@ beautifulsoup4>=4.8.0, <4.11.2
 inquirer>=2.10.0
 psutil<5.9.4
 click<=8.1.3
+aiohttp>=3.8.0, <=3.8.3

src/lightning_app/CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added the CLI command `lightning run model` to launch a `LightningLite` accelerated script ([#15506](https://github.com/Lightning-AI/lightning/pull/15506))
+
 - Added the CLI command `lightning delete app` to delete a lightning app on the cloud ([#15783](https://github.com/Lightning-AI/lightning/pull/15783))
 
 - Show a message when `BuildConfig(requirements=[...])` is passed but a `requirements.txt` file is already present in the Work ([#15799](https://github.com/Lightning-AI/lightning/pull/15799))
@@ -17,6 +18,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added a CloudMultiProcessBackend which enables running a child App from within the Flow in the cloud ([#15800](https://github.com/Lightning-AI/lightning/pull/15800))
 
+- Added `AutoScaler` component ([#15769](https://github.com/Lightning-AI/lightning/pull/15769))
+
 - Added the property `ready` of the LightningFlow to inform when the `Open App` should be visible ([#15921](https://github.com/Lightning-AI/lightning/pull/15921))
 
 - Added private work attributed `_start_method` to customize how to start the works ([#15923](https://github.com/Lightning-AI/lightning/pull/15923))

src/lightning_app/components/__init__.py

@@ -1,3 +1,4 @@
+from lightning_app.components.auto_scaler import AutoScaler
 from lightning_app.components.database.client import DatabaseClient
 from lightning_app.components.database.server import Database
 from lightning_app.components.multi_node import (
@@ -15,6 +16,7 @@
 from lightning_app.components.training import LightningTrainerScript, PyTorchLightningScriptRunner
 
 __all__ = [
+    "AutoScaler",
     "DatabaseClient",
     "Database",
     "PopenPythonScript",