Step-by-Step: Deploy a Llama 3.1 8B Model to GCP Vertex AI 2.0 for Real-Time Code Suggestions

# terraform/main.tf
# Provider configuration for GCP Vertex AI
terraform {
  required_providers {
    google = {
      source  = "hashicorp/google"
      version = "~> 5.0"
    }
  }
  required_version = "~> 1.7.0"
}

# Configure GCP provider with explicit project and region
provider "google" {
  project = var.gcp_project_id
  region  = var.gcp_region
}

# Variable definitions for environment-specific config
variable "gcp_project_id" {
  type        = string
  description = "GCP project ID for Vertex AI deployment"
  validation {
    condition     = can(regex("^[a-z0-9-]{6,30}$", var.gcp_project_id))
    error_message = "Project ID must be 6-30 characters, lowercase letters, numbers, hyphens."
  }
}

variable "gcp_region" {
  type        = string
  description = "GCP region for Vertex AI resources"
  default     = "us-central1"
}

variable "model_bucket_name" {
  type        = string
  description = "GCS bucket name for model artifacts"
  default     = "llama-3.1-8b-code-artifacts"
}

variable "model_display_name" {
  type        = string
  description = "Display name for the Llama 3.1 8B model in Vertex AI"
  default     = "llama-3.1-8b-code-suggest"
}

# Upload Llama 3.1 8B model to Vertex AI Model Registry
resource "google_vertex_ai_model" "llama_3_1_8b" {
  name                 = "llama-3-1-8b-code-suggest-${formatdate("YYYYMMDD-HHmmss", timestamp())}"
  display_name         = var.model_display_name
  description          = "Fine-tuned Llama 3.1 8B for real-time Python code suggestions"
  region               = var.gcp_region
  project              = var.gcp_project_id
  artifact_uri         = "gs://${var.model_bucket_name}/llama-3.1-8b-code/artifacts"
  # Use Hugging Face Transformers runtime for compatibility
  runtime_environment {
    python_version = "3.10"
    # Pinned Transformers version to avoid breaking changes
    dependencies = ["huggingface-hub==0.23.4", "transformers==4.41.0", "vllm==0.4.2"]
  }
  # Model metadata for discoverability
  labels = {
    task          = "code-suggestion"
    model-family  = "llama-3"
    deployment-env = "production"
  }
}

# Create Vertex AI endpoint for real-time inference
resource "google_vertex_ai_endpoint" "code_suggest_endpoint" {
  name         = "llama-3-1-8b-code-endpoint"
  display_name = "Llama 3.1 8B Code Suggestion Endpoint"
  description  = "Endpoint for low-latency real-time code suggestion requests"
  region       = var.gcp_region
  project      = var.gcp_project_id
  # Dedicated endpoint for predictable performance
  dedicated_endpoint_enabled = true
  labels = {
    task = "code-suggestion"
  }
}

# Deploy model to endpoint with autoscaling
resource "google_vertex_ai_model_deployment" "llama_deployment" {
  endpoint_id = google_vertex_ai_endpoint.code_suggest_endpoint.id
  model_id    = google_vertex_ai_model.llama_3_1_8b.id
  region      = var.gcp_region
  project     = var.gcp_project_id
  display_name = "llama-3.1-8b-code-deployment"
  # Autoscaling configuration to handle traffic spikes
  autoscaling {
    min_replica_count = 1
    max_replica_count = 5
    # Scale up when CPU utilization exceeds 70%
    target_cpu_utilization_percentage = 70
  }
  # Machine spec for Llama 3.1 8B: A100 40GB for optimal throughput
  machine_spec {
    machine_type = "a2-highgpu-1g"
    accelerator_type = "NVIDIA_TESLA_A100"
    accelerator_count = 1
  }
  # Traffic split: 100% to new deployment
  traffic_split = {
    "0" = 100
  }
}

# Output endpoint URL for client configuration
output "endpoint_url" {
  value       = google_vertex_ai_endpoint.code_suggest_endpoint.deployed_models[0].private_endpoints[0].predict_http_uri
  description = "HTTP endpoint URI for code suggestion inference"
}

# Output model ID for reference
output "model_id" {
  value       = google_vertex_ai_model.llama_3_1_8b.id
  description = "Vertex AI Model Registry ID for Llama 3.1 8B"
}

# finetune_llama.py
# Fine-tune Llama 3.1 8B on Python code suggestion pairs for Vertex AI deployment
import os
import torch
import logging
from datasets import load_dataset, Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
import pandas as pd

# Configure logging for debugging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Configuration constants
BASE_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
DATASET_PATH = "code_suggestions_dataset.jsonl"  # 12k Python (prompt, completion) pairs
OUTPUT_DIR = "./llama-3.1-8b-code-finetuned"
BATCH_SIZE = 2  # Adjust based on GPU memory (A100 40GB: batch size 2 with QLoRA)
EPOCHS = 3
LEARNING_RATE = 2e-4


def load_and_prepare_dataset(path: str) -> Dataset:
    """Load and tokenize code suggestion dataset for training."""
    try:
        logger.info(f"Loading dataset from {path}")
        df = pd.read_json(path, lines=True)
        if "prompt" not in df.columns or "completion" not in df.columns:
            raise ValueError("Dataset must contain 'prompt' and 'completion' columns")
        # Format as instruction-response pairs for Llama 3.1 Instruct format
        def format_example(row):
            return {
                "text": f"<|begin_of_text|><|user|>
{row['prompt']}
<|assistant|>
{row['completion']}<|end_of_text|>"
            }
        dataset = Dataset.from_pandas(df).map(format_example, remove_columns=df.columns)
        return dataset
    except FileNotFoundError:
        logger.error(f"Dataset file not found at {path}")
        raise
    except Exception as e:
        logger.error(f"Failed to load dataset: {str(e)}")
        raise


def configure_qlora_model(base_model_id: str):
    """Configure Llama 3.1 8B with QLoRA for memory-efficient fine-tuning."""
    try:
        logger.info(f"Loading base model {base_model_id}")
        # 4-bit quantization config for QLoRA
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        # Load model with quantization
        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True
        )
        # Prepare model for k-bit training
        model = prepare_model_for_kbit_training(model)
        # LoRA configuration for parameter-efficient fine-tuning
        lora_config = LoraConfig(
            r=64,  # Rank of LoRA update matrices
            l_alpha=128,  # Alpha parameter for LoRA scaling
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Attention layers
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM"
        )
        model = get_peft_model(model, lora_config)
        logger.info(f"Trainable parameters: {model.print_trainable_parameters()}")
        return model
    except Exception as e:
        logger.error(f"Failed to configure QLoRA model: {str(e)}")
        raise


def main():
    # Check for GPU availability
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU required for fine-tuning Llama 3.1 8B")
    logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
    tokenizer.pad_token = tokenizer.eos_token  # Set pad token to EOS for causal LM
    tokenizer.padding_side = "right"  # Required for batch processing

    # Load and tokenize dataset
    dataset = load_and_prepare_dataset(DATASET_PATH)
    def tokenize_function(examples):
        return tokenizer(examples["text"], truncation=True, max_length=512, padding="max_length")
    tokenized_dataset = dataset.map(tokenize_function, batched=True)

    # Configure QLoRA model
    model = configure_qlora_model(BASE_MODEL_ID)

    # Training arguments for Vertex AI compatibility
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=4,  # Effective batch size 8
        num_train_epochs=EPOCHS,
        learning_rate=LEARNING_RATE,
        bf16=True,  # Use bfloat16 for A100 GPUs
        logging_steps=10,
        save_steps=500,
        save_total_limit=2,
        # Optimizer settings for stable training
        optim="adamw_torch",
        weight_decay=0.01,
        warmup_ratio=0.03,
        # Disable Hugging Face Hub push (we upload to GCS for Vertex AI)
        push_to_hub=False
    )

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
    )

    # Start fine-tuning
    logger.info("Starting fine-tuning...")
    trainer.train()

    # Save fine-tuned model and tokenizer
    logger.info(f"Saving fine-tuned model to {OUTPUT_DIR}")
    trainer.save_model(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

    # Upload to GCS for Vertex AI Model Registry
    model_bucket = os.getenv("MODEL_BUCKET_NAME")
    if not model_bucket:
        raise ValueError("MODEL_BUCKET_NAME environment variable not set")
    os.system(f"gsutil cp -r {OUTPUT_DIR} gs://{model_bucket}/llama-3.1-8b-code/artifacts")
    logger.info(f"Model uploaded to gs://{model_bucket}/llama-3.1-8b-code/artifacts")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.error(f"Fine-tuning failed: {str(e)}")
        exit(1)

# inference_client.py
# Real-time code suggestion client for Vertex AI Llama 3.1 8B endpoint
import os
import time
import json
import logging
from typing import List, Dict, Optional
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import google.auth
from google.auth.transport.requests import Request
import requests

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Configuration
ENDPOINT_URL = os.getenv("VERTEX_ENDPOINT_URL")
AUTH_SCOPE = ["https://www.googleapis.com/auth/cloud-platform"]
MAX_RETRIES = 3
TIMEOUT = 10  # Seconds for request timeout
MAX_TOKENS = 256  # Max completion tokens for code suggestions
TEMPERATURE = 0.2  # Low temperature for deterministic code suggestions


class CodeSuggestionClient:
    def __init__(self, endpoint_url: str):
        if not endpoint_url:
            raise ValueError("VERTEX_ENDPOINT_URL environment variable not set")
        self.endpoint_url = endpoint_url
        self.session = requests.Session()
        self.creds = None
        self._refresh_credentials()

    def _refresh_credentials(self):
        """Refresh GCP credentials for Vertex AI API access."""
        try:
            self.creds, _ = google.auth.default(scopes=AUTH_SCOPE)
            if not self.creds.valid:
                self.creds.refresh(Request())
            logger.info("GCP credentials refreshed successfully")
        except Exception as e:
            logger.error(f"Failed to refresh credentials: {str(e)}")
            raise

    @retry(
        stop=stop_after_attempt(MAX_RETRIES),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((requests.exceptions.RequestException, TimeoutError))
    )
    def get_suggestion(self, prompt: str, context: Optional[List[str]] = None) -> str:
        """
        Get real-time code suggestion from Llama 3.1 8B endpoint.

        Args:
            prompt: Code context/prompt for suggestion (e.g., partial function definition)
            context: Optional list of previous code lines for additional context

        Returns:
            Generated code suggestion string
        """
        start_time = time.time()
        # Format prompt with Llama 3.1 Instruct template
        full_prompt = f"<|begin_of_text|><|user|>
Complete the following Python code:
{prompt}"
        if context:
            full_prompt += f"
Previous context:
{chr(10).join(context)}"
        full_prompt += "
<|assistant|>"

        # Prepare request payload for Vertex AI Prediction API
        payload = {
            "instances": [
                {
                    "inputs": full_prompt,
                    "parameters": {
                        "max_new_tokens": MAX_TOKENS,
                        "temperature": TEMPERATURE,
                        "top_p": 0.95,
                        "do_sample": True
                    }
                }
            ]
        }

        # Add auth header to request
        headers = {
            "Authorization": f"Bearer {self.creds.token}",
            "Content-Type": "application/json"
        }

        try:
            logger.debug(f"Sending request to {self.endpoint_url}")
            response = self.session.post(
                self.endpoint_url,
                headers=headers,
                json=payload,
                timeout=TIMEOUT
            )
            response.raise_for_status()  # Raise exception for 4xx/5xx responses

            # Parse response
            result = response.json()
            suggestion = result["predictions"][0]["outputs"]

            # Log latency metrics
            latency_ms = (time.time() - start_time) * 1000
            logger.info(f"Got suggestion in {latency_ms:.2f}ms")
            return suggestion
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 401:
                logger.warning("Auth token expired, refreshing credentials")
                self._refresh_credentials()
            logger.error(f"HTTP error: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"Failed to get suggestion: {str(e)}")
            raise


def main():
    if not ENDPOINT_URL:
        raise ValueError("Set VERTEX_ENDPOINT_URL environment variable")
    client = CodeSuggestionClient(ENDPOINT_URL)

    # Example prompt: partial FastAPI endpoint
    test_prompt = "def calculate_discount(price: float, discount_percent: float) -> float:"
    test_context = ["# Calculate discounted price for e-commerce orders"]

    try:
        suggestion = client.get_suggestion(test_prompt, test_context)
        print(f"Generated suggestion:
{suggestion}")
    except Exception as e:
        logger.error(f"Failed to get suggestion: {str(e)}")
        exit(1)


if __name__ == "__main__":
    main()

# Add to Terraform runtime_environment dependencies
dependencies = ["vllm==0.4.2", "huggingface-hub==0.23.4"]

# Set environment variable in deployment (via model container config)
# In Vertex AI, this is done by adding to the model's runtime_environment.env_variables
env_variables = {
  "VLLM_USE_MODELSCOPE" = "false"
}

def get_batch_suggestions(client: CodeSuggestionClient, prompts: List[str]) -> List[str]:
    """Get suggestions for multiple prompts in a single batch request."""
    payload = {
        "instances": [
            {"inputs": f"<|begin_of_text|><|user|>
Complete: {p}
<|assistant|>", "parameters": {...}}
            for p in prompts
        ]
    }
    response = client.session.post(client.endpoint_url, json=payload, headers=client.headers)
    return [pred["outputs"] for pred in response.json()["predictions"]]

from google.cloud import monitoring_v3

def emit_acceptance_metric(acceptance_rate: float):
    client = monitoring_v3.MetricServiceClient()
    project_name = f"projects/{os.getenv('GCP_PROJECT_ID')}"
    series = monitoring_v3.TimeSeries()
    series.metric.type = "custom.googleapis.com/code_suggestion/acceptance_rate"
    series.resource.type = "aiplatform.googleapis.com/Endpoint"
    series.resource.labels["endpoint_id"] = ENDPOINT_ID
    point = series.points.add()
    point.value.double_value = acceptance_rate
    point.interval.end_time.seconds = int(time.time())
    client.create_time_series(name=project_name, time_series=[series])

llama-3.1-8b-code-suggestions/
├── terraform/                # Infrastructure as Code for Vertex AI resources
│   ├── main.tf               # GCP provider, model, endpoint, deployment config
│   ├── variables.tf          # Environment variable definitions
│   ├── outputs.tf            # Endpoint URL and model ID outputs
│   └── terraform.tfvars.example  # Example config file
├── finetune/                 # Fine-tuning scripts and dataset
│   ├── finetune_llama.py     # QLoRA fine-tuning script (code example 2)
│   ├── requirements.txt      # Python dependencies for fine-tuning
│   └── sample_dataset.jsonl  # 100 sample code suggestion pairs
├── client/                   # Inference client and IDE integrations
│   ├── inference_client.py   # Production inference client (code example 3)
│   ├── vscode-extension/     # VS Code extension for real-time suggestions
│   └── requirements.txt      # Client dependencies
├── benchmarks/               # Latency and accuracy benchmark scripts
│   ├── run_benchmarks.py     # HumanEval and latency benchmark runner
│   └── results/              # Benchmark results (CSV and JSON)
└── README.md                 # Setup and deployment instructions