Develop and deploy a machine learning model with Ray

kubectl -n <your-namespace> port-forward svc/<your-model-deployment>-head-svc 8000:8265

from ray import serve
from starlette.requests import Request

# These defaults are for the ray serve deployment
# when running simply from docker. The 'Create Ray Model Deployment'
# job can override these replicas and resources if needed.
@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 1})
class Deployment(object):
    def __init__(self):
        """
        Add any initialization parameters. Generally this is where you would load
        your model. This method will be called once when the deployment is created.
        """
        print("Initializing")
        self.model = load_model() #faux code

    # This can be named as any method which takes a dictionary as input and returns a dictionary
    # as output. In this example, we are using the encode method to encode the
    # input text into a vector.
    def encode(self, input_dict: Dict[str, Any]) -> Dict[str, Any]:
        """
        This method will be called when the deployment is queried. It will receive
        the input data and should return the output data.
        """
        text = input_dict["text"]
        embeddings = self.model.encode #faux code
        return { "vector": embeddings } # To use the 'Ray / Seldon Vectorize Field' stage, the output key should be `vector`, if using the 'Machine Learning' stage you must ensure the output key matches the output key in the 'Machine Learning' stage

    async def __call__(self, http_request: Request) -> Dict[str, Any]:
        input_dict: Dict[str, Any] = await http_request.json()
        return self.encode(input_dict=input_dict) # This will be the function you defined above, in this case encode


app = Deployment.bind()

import sys
from time import time
from typing import Any, Dict

import torch
import torch.nn.functional as F
from loguru import logger
from ray import serve
from starlette.requests import Request
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

HUB_MODEL_NAME = "intfloat/e5-small-v2"


@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 1})
class Deployment(object):
    def __init__(self):
        # Initializing logger
        logger.remove()
        logger.add(sys.stdout, level="INFO", serialize=False, colorize=True)

        # Initializing model
        logger.info("Loading model...")
        self.tokenizer = AutoTokenizer.from_pretrained(HUB_MODEL_NAME)
        self.model = AutoModel.from_pretrained(HUB_MODEL_NAME)
        self.model.eval()
        logger.info("Model initialization finished!")

    def encode(self, input_dict: Dict[str, Any]) -> Dict[str, Any]:
        _start_time = time()

        # Extracting text from input
        text = input_dict["text"]

        # Tokenization
        tokenized_texts = self.tokenizer(
            text,
            max_length=512,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )

        # Encoding
        with torch.inference_mode():
            # Forward pass of the model
            outputs = self.model(**tokenized_texts)

            # Average pooling the last hidden states
            embeddings = self.average_pool(
                outputs.last_hidden_state, tokenized_texts["attention_mask"]
            )

            # Normalizing embeddings
            embeddings = F.normalize(embeddings, p=2, dim=1)

            # Converting into output format
            output_dict = {"vector": embeddings.squeeze().tolist()}

        prediction_time = (time() - _start_time) * 1000
        logger.info(f"Time taken to make a prediction: {prediction_time:.0f}ms")
        return output_dict

    async def __call__(self, http_request: Request) -> Dict[str, Any]:
        input_dict: Dict[str, Any] = await http_request.json()
        return self.encode(input_dict=input_dict)

    @staticmethod
    def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
        last_hidden = last_hidden_states.masked_fill(
            ~attention_mask[..., None].bool(), 0.0
        )
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]


app = Deployment.bind()

#It is important that python version is 3.x-slim
FROM python:3.10-slim

# Install dependencies
RUN apt-get update && apt-get install -y wget

# Create working app directory
RUN mkdir -p /app
WORKDIR /app

# Copy the requirements file and install the dependencies
COPY requirements.txt /app
RUN pip install -r requirements.txt --no-cache-dir

# Copy source code
COPY deployment.py /app

# Expose serving port for HTTP communication with Fusion
EXPOSE 8000

# The end of the command follows module:application and the below value should be set in the RAY DEPLOYMENT IMPORT PATH field in 'Create Ray Model Deployment' job
CMD exec serve run deployment:app

torch -f https://download.pytorch.org/whl/torch_stable.html # Make sure that we download CPU version of PyTorch
transformers
loguru
ray[serve]

pip freeze > requirements.txt

DOCKER_BUILDKIT=1 docker build . -t [DOCKERHUB-USERNAME]/[REPOSITORY]:[VERSION-TAG]

docker push [DOCKERHUB USERNAME]/[REPOSITORY]:[VERSION-TAG]

DOCKER_BUILDKIT=1 docker build . -t jstrmec/e5-small-v2-ray:0.1

docker push jstrmec/e5-small-v2-ray:0.1