stf/stf-api-alternative/pytriton/examples/online_learning_mnist/server.py

#!/usr/bin/env python3
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Server shows online learning model concept."""
import argparse
import logging
import threading
from queue import Queue
from threading import Lock

import numpy as np
import torch  # pytype: disable=import-error
import torch.nn.functional as functional  # pytype: disable=import-error
import torch.optim as optim  # pytype: disable=import-error
from torch.optim.lr_scheduler import StepLR  # pytype: disable=import-error

from pytriton.decorators import batch
from pytriton.model_config import ModelConfig, Tensor
from pytriton.triton import Triton, TritonConfig

from model import Net  # pytype: disable=import-error # isort:skip

LOGGER = logging.getLogger("examples.online_learning_mnist.server")
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s: %(message)s")


class Trainer:
    """Trainer class for MNIST model.
    It is used to train the model and to keep track of the training progress.
    It defines the learning rate scheduler and optimizer.
    It organizes the training process in epochs.
    """

    def __init__(self, model, lr, gamma, epoch_size):
        self.model = model
        self.optimizer = optim.Adadelta(model.parameters(), lr=lr)
        self.scheduler = StepLR(self.optimizer, step_size=1, gamma=gamma)
        self.iter = 0
        self.epoch = 0
        self.epoch_size = epoch_size

    def train_batch(self, data, target):
        self.optimizer.zero_grad()
        output = self.model(data)
        loss = functional.nll_loss(output, target)
        loss.backward()
        self.optimizer.step()
        self.iter += 1
        return loss.item()

    def ready(self):
        return self.iter >= self.epoch_size

    def next_epoch(self):
        self.iter = 0
        self.epoch += 1
        self.scheduler.step()


class OnlineLearning(threading.Thread):
    """Online learning class that implements two infer functions: train and infer.
    Infer function is used in inference endpoint and train function is used in training endpoint.
    Train function collects data and trains model in background thread.
    Infer function uses trained model to make inference.
    When trained model is ready, it is swapped with infer model.
    """

    def __init__(self, device, lr, gamma, epoch_size, max_queue_size):
        super().__init__()
        self.device = device

        self.trained_model = Net().to(self.device)
        self.trained_model.train()
        self.infer_model = Net().to(self.device)
        self.infer_model.eval()
        self.stopped = False
        self.train_data_queue = Queue(maxsize=max_queue_size)

        self.lock = Lock()
        self.trainer = Trainer(self.trained_model, lr, gamma, epoch_size)
        self.last_loss = 0.0

    def run(self) -> None:
        while not self.stopped:
            image, target = self.train_data_queue.get()
            if self.stopped:
                return

            data_tensor = torch.from_numpy(image).to(self.device)
            labels = target.reshape((target.shape[0],))
            labels_tensor = torch.from_numpy(labels).to(self.device)
            self.last_loss = self.trainer.train_batch(data_tensor, labels_tensor)

            if self.trainer.ready():
                self.replace_inference_model()
                self.trainer.next_epoch()

    def stop(self):
        self.stopped = True
        self.train_data_queue.put((None, None))
        self.join()

    def replace_inference_model(self):
        with self.lock:
            self.infer_model.load_state_dict(self.trained_model.state_dict())

    def train(self, requests):
        """Train function is used in training endpoint."""
        # concatenate all requests into one batch. No need for padding due to fixed image dimensions
        images = np.concatenate([request["image"] for request in requests], axis=0)
        targets = np.concatenate([request["target"] for request in requests], axis=0)
        self.train_data_queue.put((images, targets))
        return [{"last_loss": np.array([[self.last_loss]]).astype(np.float32)} for _ in requests]

    @batch
    def infer(self, image):
        """Infer function is used in inference endpoint."""
        data_tensor = torch.from_numpy(image).to(self.device)
        with self.lock:
            res = self.infer_model(data_tensor)
        res = res.numpy(force=True)
        return {"predictions": res}


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging in debug mode.",
    )
    return parser.parse_args()


def main():
    args = _parse_args()
    log_verbose = 1 if args.verbose else 0
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format="%(asctime)s - %(levelname)s - %(name)s: %(message)s")

    online_learning_model = OnlineLearning(
        device=torch.device("cuda"), lr=1.0, gamma=0.7, epoch_size=134, max_queue_size=1000
    )
    online_learning_model.start()
    try:
        with Triton(config=TritonConfig(log_verbose=log_verbose)) as triton:
            LOGGER.info("Loading OnlineLearning model")
            triton.bind(
                model_name="MnistTrain",
                infer_func=online_learning_model.train,
                inputs=[
                    # image for training
                    Tensor(name="image", dtype=np.float32, shape=(1, 28, 28)),
                    # target class corresponding to image (class index from 0 to 9)
                    Tensor(name="target", dtype=np.int64, shape=(1,)),
                ],
                outputs=[
                    # last loss value batch
                    Tensor(name="last_loss", dtype=np.float32, shape=(1,)),
                ],
                config=ModelConfig(max_batch_size=64),
                strict=True,
            )
            triton.bind(
                model_name="MnistInfer",
                infer_func=online_learning_model.infer,
                inputs=[
                    # image for classification
                    Tensor(name="image", dtype=np.float32, shape=(1, 28, 28)),
                ],
                outputs=[
                    # predictions taken from softmax layer
                    Tensor(name="predictions", dtype=np.float32, shape=(-1,)),
                ],
                config=ModelConfig(max_batch_size=64),
                strict=True,
            )

            LOGGER.info("Serving model")
            triton.serve()
    finally:
        LOGGER.info("Stopping online learning model")
        online_learning_model.stop()


if __name__ == "__main__":
    main()