Generative AI Google Brain

open-models-serving-vertex-ai-pytorch-inference-pllum-with-custom-handler

3 min read

Hugging Face DLCs: Serving PLLuM using Pytorch Inference on Vertex AI with Custom Handler

Source notebook

Repo path: open-models/serving/vertex_ai_pytorch_inference_pllum_with_custom_handler.ipynb · Open on GitHub · advanced

Deploys CYFRAGOVPL/PLLuM-12B-chat to Vertex AI with a Hugging Face PyTorch custom handler.

Summary

This notebook teaches how to serve the Polish PLLuM language model from Hugging Face Hub on Vertex AI using a Hugging Face PyTorch inference DLC and a custom handler. It walks through local handler testing, copying model artifacts to Cloud Storage, uploading the model to Vertex AI Model Registry, deploying it to a GPU endpoint, and sending predictions through Python, gcloud, and cURL.

Key code patterns

Initialize Vertex AI

import vertexai
 
vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=BUCKET_URI,
)

Sets project, region, and staging bucket before model upload and deployment.

Load PLLuM locally

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.eval()

Tests the Hugging Face model with lower memory dtype and automatic device placement.

Custom handler

class EndpointHandler:
    def __init__(self, model_dir="/opt/huggingface/model", **kwargs):
        self.processor = AutoTokenizer.from_pretrained(model_dir)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_dir, torch_dtype=torch.bfloat16, device_map="auto"
        ).eval()

Defines how the PyTorch inference container loads tokenizer and model artifacts.

Handler prediction loop

for instance in data["instances"]:
    if "prompt" not in instance:
        raise ValueError("Missing prompt in request body")
    inputs = self.processor(instance["prompt"], return_tensors="pt", return_token_type_ids=False).to(self.model.device)
    generation = self.model.generate(**inputs, **generation_kwargs)
    predictions.append(self.processor.decode(generation[0][input_len:], skip_special_tokens=True))

Implements Vertex AI request parsing, generation, and response formatting.

Upload model

model = Model.upload(
    display_name="cyfragovpl--pllum-12b-it",
    artifact_uri=str(model_uri),
    serving_container_image_uri=IMAGE_URI,
    serving_container_ports=[8080],
)
model.wait()

Registers the Cloud Storage model directory and Hugging Face DLC in Vertex AI Model Registry.

Deploy endpoint

deployed_model = model.deploy(
    endpoint=Endpoint.create(display_name="cyfragovpl--pllum-12b-it-endpoint"),
    machine_type="g2-standard-8",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

Allocates GPU serving resources for online prediction.

Models & APIs used

  • Models: CYFRAGOVPL/PLLuM-12B-chat
  • APIs / services: Vertex AI, Cloud Storage, Artifact Registry
  • SDKs / libraries: google-cloud-aiplatform, vertexai, transformers, torch, etils

When to use this

Use this pattern when serving a Hugging Face causal language model on Vertex AI requires custom preprocessing or generation logic.

Gotchas & caveats

  • The notebook requires enabling aiplatform.googleapis.com and artifactregistry.googleapis.com.
  • The service account needs roles/aiplatform.user, roles/storage.objectAdmin, and roles/artifactregistry.reader.
  • The runtime must be restarted after installing torch, transformers, and google-cloud-aiplatform[prediction].
  • LocalModel testing requires a local Docker installation.
  • The local PLLuM test is described as resource exhausting and suggests capable hardware such as g2-standard-8.
  • The deployment can take around 15 to 25 minutes.
  • The selected machine type and accelerator type must be compatible.
  • The handler uses torch.bfloat16, which the notebook says should be enabled only when hardware supports bfloat16.

Best practices

  • Test the custom handler logic locally before packaging it as handler.py.
  • Use HF_HOME to keep Hugging Face cache artifacts in the tutorial directory.
  • Copy the model snapshot and handler.py to Cloud Storage before Vertex AI upload.
  • Use gsutil -m and parallel composite upload settings for large model directories.
  • Validate the serving container locally with LocalModel before deploying to Vertex AI.
  • Monitor local container deployment with docker container ls and docker logs.
  • Provide cleanup flags for endpoint, model, and bucket resources.

Graph View

Backlinks