DEV Community: Suhas Mallesh

Azure ML Feature Store with Terraform: Managed Feature Materialization for Training and Inference 🗃️

Suhas Mallesh — Fri, 17 Apr 2026 07:00:00 +0000

Azure ML Feature Store is a specialized workspace that manages feature engineering, offline materialization to storage, and online serving with Redis. Terraform provisions the infrastructure, SDK defines feature sets. Here's how to build it.

In the previous posts, we set up the ML workspace and deployed endpoints. Now we need consistent features feeding those endpoints. Training uses historical features from batch sources. Inference needs the latest values in real time. When these diverge, your model's accuracy degrades silently.

Azure ML Feature Store is implemented as a special type of Azure ML workspace (kind = "FeatureStore"). It manages feature transformation pipelines, materializes features to offline storage (ADLS/Blob) and an online store (Redis), and provides point-in-time feature retrieval for training. Terraform provisions the infrastructure; the SDK defines entities, feature sets, and materialization schedules. 🎯

🏗️ Feature Store Architecture

Component	What It Does
Feature Store	Specialized ML workspace with `kind = "FeatureStore"`
Entity	Logical key (e.g., customer_id, account_id) shared across feature sets
Feature Set	Collection of features with transformation code and source definition
Offline Store	ADLS/Blob storage for materialized historical features
Online Store	Redis cache for low-latency inference lookups
Materialization	Spark jobs that compute and sync features on a schedule

The key concept: feature sets include transformation code. Raw data goes in, computed features come out. The same transformation runs for both offline materialization (training) and online materialization (inference), eliminating training-serving skew.

🔧 Terraform: Provision Feature Store Infrastructure

Feature Store Workspace

# feature_store/workspace.tf

resource "azurerm_machine_learning_workspace" "feature_store" {
  name                = "${var.environment}-feature-store"
  location            = azurerm_resource_group.ml.location
  resource_group_name = azurerm_resource_group.ml.name
  application_insights_id = azurerm_application_insights.ml.id
  key_vault_id            = azurerm_key_vault.ml.id
  storage_account_id      = azurerm_storage_account.ml.id

  kind = "FeatureStore"

  identity {
    type = "SystemAssigned"
  }

  tags = var.tags
}

kind = "FeatureStore" is the critical setting. This creates a workspace optimized for feature management rather than general ML development.

Offline Materialization Store

# feature_store/offline_store.tf

resource "azurerm_storage_account" "offline_store" {
  name                     = "${var.environment}fsoffline${random_string.suffix.result}"
  location                 = azurerm_resource_group.ml.location
  resource_group_name      = azurerm_resource_group.ml.name
  account_tier             = "Standard"
  account_replication_type = var.storage_replication
  is_hns_enabled           = true   # ADLS Gen2

  tags = var.tags
}

resource "azurerm_storage_container" "features" {
  name                  = "features"
  storage_account_id    = azurerm_storage_account.offline_store.id
  container_access_type = "private"
}

is_hns_enabled = true enables ADLS Gen2 hierarchical namespace, which is required for efficient feature materialization with Parquet files.

Online Store (Redis Cache)

# feature_store/online_store.tf

resource "azurerm_redis_cache" "online_store" {
  count               = var.enable_online_store ? 1 : 0
  name                = "${var.environment}-fs-redis"
  location            = azurerm_resource_group.ml.location
  resource_group_name = azurerm_resource_group.ml.name
  capacity            = var.redis_capacity
  family              = var.redis_family
  sku_name            = var.redis_sku
  minimum_tls_version = "1.2"

  redis_configuration {
    maxmemory_policy = "allkeys-lru"
  }

  tags = var.tags
}

The online store is optional. Enable it when you need low-latency feature lookups during inference. Skip it in dev if you only need offline features for training.

Compute for Materialization

# feature_store/compute.tf

resource "azurerm_machine_learning_compute_cluster" "materialization" {
  name                          = "${var.environment}-materialization"
  machine_learning_workspace_id = azurerm_machine_learning_workspace.feature_store.id
  location                      = azurerm_resource_group.ml.location
  vm_size                       = var.materialization_vm_size
  vm_priority                   = "LowPriority"

  identity {
    type = "SystemAssigned"
  }

  scale_settings {
    min_node_count                       = 0
    max_node_count                       = var.materialization_max_nodes
    scale_down_nodes_after_idle_duration  = "PT5M"
  }

  tags = var.tags
}

Materialization jobs run as Spark pipelines on this compute cluster. min_node_count = 0 means you pay nothing when no materialization is running.

🐍 Define Entities and Feature Sets (SDK)

Terraform provisions infrastructure. The SDK defines the feature engineering logic:

Create an Entity

from azure.ai.ml import MLClient
from azure.ai.ml.entities import FeatureStoreEntity, DataColumn
from azure.identity import DefaultAzureCredential

fs_client = MLClient(
    DefaultAzureCredential(),
    subscription_id="...",
    resource_group_name="...",
    workspace_name="prod-feature-store",
)

account_entity = FeatureStoreEntity(
    name="account",
    version="1",
    index_columns=[DataColumn(name="accountID", type="string")],
    description="Account entity for transaction features",
)

fs_client.feature_store_entities.begin_create_or_update(account_entity).result()

Entities define shared join keys. Multiple feature sets can reference the same entity, ensuring consistent joins.

Define Feature Set with Transformation Code

Feature set specification (YAML):

# featuresets/transactions/spec/FeaturesetSpec.yaml
$schema: https://azuremlschemas.azureedge.net/latest/featureSetSpec.schema.json

source:
  type: parquet
  path: abfss://data@storage.dfs.core.windows.net/transactions/
  timestamp_column:
    name: timestamp

feature_transformation_code:
  path: ./transformation_code
  transformer_class: transaction_transform.TransactionFeatureTransformer

features:
  - name: transaction_count_7d
    type: integer
  - name: avg_transaction_amount_7d
    type: float
  - name: total_spend_3d
    type: float
  - name: max_transaction_amount
    type: float

index_columns:
  - name: accountID
    type: string

Transformation code (Spark):

# transformation_code/transaction_transform.py
from pyspark.sql import DataFrame
from pyspark.sql import functions as F
from pyspark.sql.window import Window

class TransactionFeatureTransformer:
    def transform(self, raw_data: DataFrame) -> DataFrame:
        window_7d = Window.partitionBy("accountID").orderBy("timestamp").rangeBetween(-7*86400, 0)
        window_3d = Window.partitionBy("accountID").orderBy("timestamp").rangeBetween(-3*86400, 0)

        return raw_data.select(
            "accountID",
            "timestamp",
            F.count("*").over(window_7d).alias("transaction_count_7d"),
            F.avg("amount").over(window_7d).alias("avg_transaction_amount_7d"),
            F.sum("amount").over(window_3d).alias("total_spend_3d"),
            F.max("amount").over(window_7d).alias("max_transaction_amount"),
        )

Register and Materialize

from azure.ai.ml.entities import FeatureSet, FeatureSetSpecification

transaction_fset = FeatureSet(
    name="transactions",
    version="1",
    description="7-day and 3-day rolling transaction aggregations",
    entities=["azureml:account:1"],
    specification=FeatureSetSpecification(
        path="./featuresets/transactions/spec"
    ),
    tags={"data_type": "nonPII"},
)

fs_client.feature_sets.begin_create_or_update(transaction_fset).result()

Configure Materialization Schedule

from azure.ai.ml.entities import (
    MaterializationSettings,
    MaterializationComputeResource,
    RecurrenceTrigger,
)

materialization = MaterializationSettings(
    resource=MaterializationComputeResource(instance_type="Standard_E8s_v3"),
    schedule=RecurrenceTrigger(frequency="Hour", interval=6),
    offline_enabled=True,
    online_enabled=True,
)

fset = fs_client.feature_sets.get(name="transactions", version="1")
fset.materialization_settings = materialization
fs_client.feature_sets.begin_create_or_update(fset).result()

📐 Environment Configuration

# environments/dev.tfvars
environment              = "dev"
enable_online_store      = false        # No Redis in dev
storage_replication      = "LRS"
materialization_vm_size  = "Standard_E4s_v3"
materialization_max_nodes = 2

# environments/prod.tfvars
environment              = "prod"
enable_online_store      = true
redis_sku                = "Standard"
redis_capacity           = 1
redis_family             = "C"
storage_replication      = "GRS"
materialization_vm_size  = "Standard_E8s_v3"
materialization_max_nodes = 8

⚠️ Gotchas and Tips

Feature store is a workspace. It's implemented as kind = "FeatureStore" on azurerm_machine_learning_workspace. It needs the same dependencies (storage, KV, App Insights) as a regular workspace.

Transformation code runs as Spark. Feature transformations execute on the materialization compute cluster using PySpark. Test your transformations locally with a Spark session before registering.

Entities enforce consistent joins. Define entities once (e.g., "account" with key "accountID") and reuse across feature sets. This prevents mismatched join keys between teams.

Materialization costs. Each scheduled run spins up the compute cluster, runs the Spark job, and writes to storage. LowPriority VMs reduce cost. min_node_count = 0 ensures you pay nothing between runs.

Redis cost for online store. Standard Redis starts at ~$40/month. Premium with replication is ~$200/month. Skip online store in dev unless you're testing real-time inference.

Feature set versioning. Feature sets are versioned. Changing the transformation logic? Create version "2". This maintains backward compatibility for models still using version "1".

⏭️ What's Next

This is Post 3 of the Azure ML Pipelines & MLOps with Terraform series.

Post 1: Azure ML Workspace 🔬
Post 2: Azure ML Online Endpoints 🚀
Post 3: Azure ML Feature Store (you are here) 🗃️
Post 4: Azure ML Pipelines + Azure DevOps

Your features have a home. ADLS for offline training, Redis for online inference, Spark transformations that run the same code for both. No training-serving skew. Versioned feature sets with scheduled materialization, all provisioned with Terraform. 🗃️

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Vertex AI Feature Store with Terraform: BigQuery Offline + Bigtable Online Serving 🗃️

Suhas Mallesh — Thu, 16 Apr 2026 07:00:00 +0000

Feature Store on GCP uses BigQuery as the offline store and Bigtable for low-latency online serving. Feature groups register your data, feature views sync it to the online store. Here's how to provision the full stack with Terraform.

In the previous posts, we set up Workbench for development and deployed endpoints for inference. But the features feeding those models need a home. Training uses historical features from BigQuery. Inference needs the latest values with sub-millisecond latency. When these two sources diverge, you get training-serving skew.

Vertex AI Feature Store bridges this gap. BigQuery is the offline store - your features live in tables you already manage. Bigtable is the online store - an auto-scaling, low-latency serving layer that syncs from BigQuery on a schedule. You don't copy data to a separate system. Feature Store reads directly from BigQuery and syncs to Bigtable for serving. 🎯

🏗️ Feature Store Architecture

Component	What It Does
Feature Group	Registers a BigQuery table as a feature source
Feature	Individual column within a feature group
Feature Online Store	Bigtable instance for real-time serving
Feature View	Defines which features sync to the online store
Data Sync	Scheduled or continuous sync from BigQuery to Bigtable

The key insight: BigQuery is already your offline store. You don't move data. Feature Store registers your existing BigQuery tables, then syncs selected features to Bigtable for online serving.

🔧 Terraform: Create the Feature Online Store

APIs

# feature_store/apis.tf

resource "google_project_service" "required" {
  for_each = toset([
    "aiplatform.googleapis.com",
    "bigtable.googleapis.com",
    "bigtableadmin.googleapis.com",
    "bigquery.googleapis.com",
  ])
  project = var.project_id
  service = each.value
}

Feature Online Store (Bigtable-backed)

# feature_store/online_store.tf

resource "google_vertex_ai_feature_online_store" "this" {
  name     = "${var.environment}-feature-store"
  region   = var.region
  project  = var.project_id

  bigtable {
    auto_scaling {
      min_node_count         = var.bigtable_min_nodes
      max_node_count         = var.bigtable_max_nodes
      cpu_utilization_target = var.bigtable_cpu_target
    }
  }

  labels = {
    environment = var.environment
    managed_by  = "terraform"
  }
}

Bigtable autoscaling adjusts nodes based on CPU utilization. Set cpu_utilization_target to 50-60% for production workloads. The store scales up automatically during traffic spikes and scales down during quiet periods.

Feature Group (Register BigQuery Source)

# feature_store/feature_group.tf

resource "google_vertex_ai_feature_group" "customer_features" {
  name     = "${var.environment}-customer-features"
  region   = var.region
  project  = var.project_id

  big_query {
    big_query_source {
      input_uri = "bq://${var.project_id}.${var.dataset_id}.${var.customer_features_table}"
    }
    entity_id_columns = ["customer_id"]
  }

  labels = {
    domain = "customer"
  }
}

entity_id_columns defines the primary key for feature lookups. This is what you use to retrieve features for a specific customer during inference.

Register Individual Features

# feature_store/features.tf

resource "google_vertex_ai_feature_group_feature" "total_purchases" {
  name           = "total_purchases"
  region         = var.region
  feature_group  = google_vertex_ai_feature_group.customer_features.name
  project        = var.project_id
}

resource "google_vertex_ai_feature_group_feature" "avg_order_value" {
  name           = "avg_order_value"
  region         = var.region
  feature_group  = google_vertex_ai_feature_group.customer_features.name
  project        = var.project_id
}

resource "google_vertex_ai_feature_group_feature" "days_since_last_purchase" {
  name           = "days_since_last_purchase"
  region         = var.region
  feature_group  = google_vertex_ai_feature_group.customer_features.name
  project        = var.project_id
}

Each feature maps to a column in your BigQuery table. Registering features enables metadata tracking, drift monitoring, and controlled syncing to the online store.

Feature View (Sync to Online Store)

# feature_store/feature_view.tf

resource "google_vertex_ai_feature_online_store_featureview" "customer_view" {
  name                 = "${var.environment}-customer-view"
  region               = var.region
  feature_online_store = google_vertex_ai_feature_online_store.this.name
  project              = var.project_id

  sync_config {
    cron = var.sync_schedule
  }

  feature_registry_source {
    feature_groups {
      feature_group_id = google_vertex_ai_feature_group.customer_features.name
      feature_ids      = [
        google_vertex_ai_feature_group_feature.total_purchases.name,
        google_vertex_ai_feature_group_feature.avg_order_value.name,
        google_vertex_ai_feature_group_feature.days_since_last_purchase.name,
      ]
    }
  }
}

The feature view selects which features from which groups sync to the online store. The cron schedule controls how frequently BigQuery data is synced to Bigtable.

📐 BigQuery Source Table Structure

Your BigQuery table needs an entity ID column and a feature timestamp:

CREATE TABLE `project.ml_features.customer_features` (
  customer_id STRING NOT NULL,
  feature_timestamp TIMESTAMP NOT NULL,
  total_purchases INT64,
  avg_order_value FLOAT64,
  days_since_last_purchase INT64,
  account_age_days INT64,
  is_premium BOOL
);

Feature Store reads this table directly. The feature_timestamp column enables point-in-time queries for training. The online store always serves the latest snapshot.

🐍 Read Features (SDK)

Online Store (Real-Time Inference)

from google.cloud import aiplatform

aiplatform.init(project="my-project", location="us-central1")

feature_online_store = aiplatform.FeatureOnlineStore("prod-feature-store")
feature_view = feature_online_store.get_feature_view("prod-customer-view")

# Fetch features for a specific customer
response = feature_view.fetch_feature_values(
    entity_ids=["cust-12345"],
)

for entity in response:
    print(entity.to_dict())
# {'customer_id': 'cust-12345', 'total_purchases': 47, 'avg_order_value': 89.5, ...}

Offline Store (Training via BigQuery)

from google.cloud import bigquery

client = bigquery.Client()

query = """
SELECT customer_id, total_purchases, avg_order_value, is_premium
FROM `project.ml_features.customer_features`
WHERE feature_timestamp BETWEEN '2025-01-01' AND '2025-12-31'
"""

training_df = client.query(query).to_dataframe()
print(f"Training data: {len(training_df)} rows")

No separate offline store to manage. Query BigQuery directly with standard SQL.

📐 Environment Configuration

# environments/dev.tfvars
environment          = "dev"
bigtable_min_nodes   = 1
bigtable_max_nodes   = 1
bigtable_cpu_target  = 80
sync_schedule        = "0 */6 * * *"    # Every 6 hours

# environments/prod.tfvars
environment          = "prod"
bigtable_min_nodes   = 1
bigtable_max_nodes   = 5
bigtable_cpu_target  = 50
sync_schedule        = "0 * * * *"      # Every hour

Sync frequency vs freshness: Hourly sync means online features can be up to 1 hour stale. For near-real-time features, use continuous data sync (requires Bigtable online serving and BigQuery source in specific regions).

⚠️ Gotchas and Tips

BigQuery is the source of truth. Unlike other feature stores where you ingest data into a proprietary system, Vertex AI Feature Store reads from BigQuery. Your existing ETL pipelines that write to BigQuery already feed the feature store.

Bigtable minimum cost. Even at 1 node, Bigtable costs roughly $0.65/hour (~$470/month). For dev environments, consider whether you need online serving at all, or if BigQuery direct queries suffice.

Optimized online serving is deprecated. As of May 2026, only Bigtable online serving is supported. Don't use optimized {} in new deployments. Migrate existing optimized stores to Bigtable.

Sync latency. Scheduled sync has an inherent delay based on your cron schedule. Continuous sync is near-real-time but only available in specific regions (us, eu, us-central1).

Feature monitoring. Register features through feature groups to enable drift detection and anomaly monitoring. Without registration, you lose this observability.

Bigtable serving latency. Expect ~30ms server-side latency at moderate load (~100 QPS). Client-side latency adds 5ms+. This is fast enough for most inference use cases but not sub-millisecond.

⏭️ What's Next

This is Post 3 of the GCP ML Pipelines & MLOps with Terraform series.

Post 1: Vertex AI Workbench 🔬
Post 2: Vertex AI Endpoints - Deploy to Prod 🚀
Post 3: Vertex AI Feature Store (you are here) 🗃️
Post 4: Vertex AI Pipelines + Cloud Build

Your features have a home. BigQuery for offline training, Bigtable for online serving, automatic sync between them. No data duplication. No training-serving skew. Your existing BigQuery tables are the source of truth, all provisioned with Terraform. 🗃️

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

SageMaker Feature Store with Terraform: Centralized ML Features for Training and Inference 🗃️

Suhas Mallesh — Wed, 15 Apr 2026 07:00:00 +0000

Features used for training must match features used for inference, or your model breaks silently. SageMaker Feature Store keeps them in sync with online (real-time) and offline (historical) stores. Here's how to provision it with Terraform.

In the previous posts, we set up the workspace and deployed endpoints. But there's a critical gap: features. Every ML model needs consistent, reliable feature data for both training (batch, historical) and inference (real-time, latest values). When training features and serving features diverge, you get training-serving skew, and your model's accuracy degrades silently.

SageMaker Feature Store solves this with a dual-store architecture. The online store provides low-latency access to the latest feature values for real-time inference. The offline store keeps the full history in S3 (Parquet format) for training and batch inference. When you write a feature, both stores sync automatically. One source of truth. 🎯

🏗️ Feature Store Architecture

Component	What It Does
Feature Group	A collection of related features (like a table)
Online Store	Low-latency key-value store for real-time lookups
Offline Store	Historical data in S3 (Parquet) for training
Record Identifier	Primary key for feature lookups
Event Time	Timestamp for point-in-time correctness
Glue Data Catalog	Auto-created metadata catalog for Athena queries

The online store always holds the latest snapshot. The offline store is append-only, keeping every version of every record. This enables point-in-time queries for training: "What did this customer's features look like 30 days ago?"

🔧 Terraform: Create Feature Groups

IAM Role

# feature_store/iam.tf

resource "aws_iam_role" "feature_store" {
  name = "${var.environment}-feature-store"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "sagemaker.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "feature_store_access" {
  name = "feature-store-s3-glue"
  role = aws_iam_role.feature_store.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect = "Allow"
        Action = ["s3:GetObject", "s3:PutObject", "s3:DeleteObject", "s3:GetBucketLocation"]
        Resource = [
          "${var.offline_store_bucket_arn}",
          "${var.offline_store_bucket_arn}/*"
        ]
      },
      {
        Effect = "Allow"
        Action = [
          "glue:CreateTable", "glue:UpdateTable", "glue:GetTable",
          "glue:GetDatabase", "glue:CreateDatabase"
        ]
        Resource = "*"
      }
    ]
  })
}

Feature Group Definition

# feature_store/feature_groups.tf

resource "aws_sagemaker_feature_group" "customer_features" {
  feature_group_name             = "${var.environment}-customer-features"
  record_identifier_feature_name = "customer_id"
  event_time_feature_name        = "event_time"
  role_arn                       = aws_iam_role.feature_store.arn

  # Feature schema
  feature_definition {
    feature_name = "customer_id"
    feature_type = "String"
  }

  feature_definition {
    feature_name = "event_time"
    feature_type = "Fractional"
  }

  feature_definition {
    feature_name = "total_purchases"
    feature_type = "Integral"
  }

  feature_definition {
    feature_name = "avg_order_value"
    feature_type = "Fractional"
  }

  feature_definition {
    feature_name = "days_since_last_purchase"
    feature_type = "Integral"
  }

  feature_definition {
    feature_name = "account_age_days"
    feature_type = "Integral"
  }

  feature_definition {
    feature_name = "is_premium"
    feature_type = "Integral"
  }

  # Enable both online and offline stores
  online_store_config {
    enable_online_store = true

    security_config {
      kms_key_id = var.kms_key_arn
    }
  }

  offline_store_config {
    s3_storage_config {
      s3_uri   = "s3://${var.offline_store_bucket}/${var.environment}/feature-store"
      kms_key_id = var.kms_key_arn
    }

    table_format = var.offline_table_format  # "Glue" or "Iceberg"
  }

  tags = {
    Environment = var.environment
    Domain      = "customer"
  }
}

Three feature types: String, Fractional (float), and Integral (integer). Everything else maps to String.

table_format: Choose Glue (default, Hive-compatible) or Iceberg (better for upserts and time travel). Iceberg is recommended for production workloads that need ACID transactions.

Multiple Feature Groups

resource "aws_sagemaker_feature_group" "transaction_features" {
  feature_group_name             = "${var.environment}-transaction-features"
  record_identifier_feature_name = "transaction_id"
  event_time_feature_name        = "event_time"
  role_arn                       = aws_iam_role.feature_store.arn

  feature_definition {
    feature_name = "transaction_id"
    feature_type = "String"
  }

  feature_definition {
    feature_name = "event_time"
    feature_type = "Fractional"
  }

  feature_definition {
    feature_name = "amount"
    feature_type = "Fractional"
  }

  feature_definition {
    feature_name = "merchant_category"
    feature_type = "String"
  }

  feature_definition {
    feature_name = "is_international"
    feature_type = "Integral"
  }

  online_store_config {
    enable_online_store = true
  }

  offline_store_config {
    s3_storage_config {
      s3_uri = "s3://${var.offline_store_bucket}/${var.environment}/feature-store"
    }
    table_format = var.offline_table_format
  }
}

🐍 Ingest Features (SDK)

Terraform defines the schema. The SDK ingests the data:

import boto3
import time

featurestore_runtime = boto3.client("sagemaker-featurestore-runtime")

# Write a single record (real-time)
featurestore_runtime.put_record(
    FeatureGroupName="prod-customer-features",
    Record=[
        {"FeatureName": "customer_id", "ValueAsString": "cust-12345"},
        {"FeatureName": "event_time", "ValueAsString": str(time.time())},
        {"FeatureName": "total_purchases", "ValueAsString": "47"},
        {"FeatureName": "avg_order_value", "ValueAsString": "89.50"},
        {"FeatureName": "days_since_last_purchase", "ValueAsString": "3"},
        {"FeatureName": "account_age_days", "ValueAsString": "730"},
        {"FeatureName": "is_premium", "ValueAsString": "1"},
    ],
)

Read Features for Inference (Online Store)

# Real-time feature lookup (single-digit ms latency)
response = featurestore_runtime.get_record(
    FeatureGroupName="prod-customer-features",
    RecordIdentifierValueAsString="cust-12345",
)

features = {r["FeatureName"]: r["ValueAsString"] for r in response["Record"]}
print(features)
# {'customer_id': 'cust-12345', 'total_purchases': '47', ...}

Query Features for Training (Offline Store via Athena)

import boto3

athena = boto3.client("athena")

query = """
SELECT customer_id, total_purchases, avg_order_value, is_premium
FROM "sagemaker_featurestore"."prod-customer-features"
WHERE event_time <= 1700000000
"""

response = athena.start_query_execution(
    QueryString=query,
    QueryExecutionContext={"Database": "sagemaker_featurestore"},
    ResultConfiguration={"OutputLocation": "s3://my-bucket/athena-results/"},
)

The offline store is automatically cataloged in Glue. Query with Athena for point-in-time training datasets.

📐 Environment Configuration

# environments/dev.tfvars
environment           = "dev"
offline_table_format  = "Glue"     # Simpler for dev
kms_key_arn           = null        # No encryption in dev

# environments/prod.tfvars
environment           = "prod"
offline_table_format  = "Iceberg"  # ACID transactions, time travel
kms_key_arn           = "arn:aws:kms:us-east-1:123456789012:key/abc-123"

⚠️ Gotchas and Tips

Offline store has a ~15 minute delay. Data written via PutRecord appears in the online store immediately but takes up to 15 minutes to land in the offline store (S3). Don't rely on the offline store for near-real-time analytics.

Feature groups are mutable. You can add new features to an existing feature group using the UpdateFeatureGroup API. You cannot remove or rename existing features.

Online store costs. The online store charges per read and write unit. High-throughput inference with thousands of feature lookups per second adds up. Monitor costs and batch lookups where possible using BatchGetRecord.

Point-in-time correctness. Always use event_time for training queries. Querying without time filtering risks data leakage, where future data appears in your training set.

KMS encryption for both stores. The online and offline stores support separate KMS keys. In production, encrypt both. The Glue Data Catalog metadata is not encrypted by Feature Store, manage it separately.

Schema planning matters. Feature types (String, Fractional, Integral) cannot be changed after creation. Plan your schema carefully. Use String for anything you're unsure about, since it's the most flexible.

⏭️ What's Next

This is Post 3 of the ML Pipelines & MLOps with Terraform series.

Post 1: SageMaker Studio Domain 🔬
Post 2: SageMaker Endpoints - Deploy to Prod 🚀
Post 3: SageMaker Feature Store (you are here) 🗃️
Post 4: SageMaker Pipelines - CI/CD for ML

Your features have a home. Online store for real-time inference, offline store for training, automatic sync between them. No more training-serving skew. No more duplicated feature pipelines. One source of truth, all in Terraform. 🗃️

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Azure ML Online Endpoints: Deploy Your Model to Production with Terraform 🚀

Suhas Mallesh — Mon, 13 Apr 2026 07:00:00 +0000

Your model is trained. Now deploy it to a managed online endpoint with traffic splitting for canary rollouts, autoscaling, health probes, and data collection. Here's how to deploy Azure ML endpoints with Terraform using azapi.

In the previous post, we set up the Azure ML workspace - the hub for experiments, models, and compute. Now comes the production side: taking a trained model and deploying it to a managed online endpoint that your applications can call for real-time predictions.

Azure ML online endpoints involve two resources: an Endpoint (the stable HTTPS URL with auth and traffic routing) and one or more Deployments (the model + compute behind it). Since the azurerm provider doesn't have native online endpoint resources yet, we use azapi to provision them directly via the Azure API. 🎯

🏗️ The Two-Layer Architecture

Endpoint (HTTPS URL, auth mode, traffic split)
    ↓
Deployment(s) (model, instance type, scaling, probes)

Resource	What It Defines
Endpoint	Stable URL, auth mode (key/AAD), traffic routing
Deployment	Model reference, instance type, scale settings, health probes

The endpoint URL stays fixed. You deploy new model versions as new deployments, shift traffic gradually, and delete old deployments after validation.

🔧 Terraform: Create the Online Endpoint

The Endpoint

# endpoint/main.tf

resource "azapi_resource" "online_endpoint" {
  type      = "Microsoft.MachineLearningServices/workspaces/onlineEndpoints@2025-06-01"
  name      = "${var.environment}-${var.model_name}"
  parent_id = azurerm_machine_learning_workspace.this.id
  location  = var.location

  identity {
    type = "SystemAssigned"
  }

  body = {
    properties = {
      authMode            = var.auth_mode
      publicNetworkAccess = var.public_network_access
      description         = "Production endpoint for ${var.model_name}"
      traffic = {
        (var.deployment_name) = 100
      }
    }
  }

  tags = {
    Environment = var.environment
    Model       = var.model_name
  }
}

authMode controls how clients authenticate: Key for API key auth (simpler), AADToken for Azure AD auth (more secure, no key rotation needed). Use AADToken in production.

The Deployment

# endpoint/deployment.tf

resource "azapi_resource" "deployment" {
  type      = "Microsoft.MachineLearningServices/workspaces/onlineEndpoints/deployments@2025-06-01"
  name      = var.deployment_name
  parent_id = azapi_resource.online_endpoint.id
  location  = var.location

  identity {
    type = "SystemAssigned"
  }

  body = {
    kind = "Managed"
    properties = {
      model             = var.model_uri
      environmentId     = var.environment_id
      instanceType      = var.instance_type
      appInsightsEnabled = true

      scaleSettings = {
        scaleType    = "Default"
        minInstances = var.min_instances
        maxInstances = var.max_instances
      }

      livenessProbe = {
        initialDelay   = "PT10S"
        period         = "PT10S"
        timeout        = "PT2S"
        failureThreshold = 30
        successThreshold = 1
      }

      readinessProbe = {
        initialDelay   = "PT10S"
        period         = "PT10S"
        timeout        = "PT2S"
        failureThreshold = 30
        successThreshold = 1
      }

      requestSettings = {
        requestTimeout       = "PT5S"
        maxConcurrentRequestsPerInstance = var.max_concurrent_requests
      }
    }
  }

  tags = {
    Environment = var.environment
    Model       = var.model_name
    Version     = var.model_version
  }
}

Key deployment properties:

model references the registered model in the workspace (e.g., azureml:fraud-detector:2)
environmentId references a curated or custom environment with your dependencies
scaleSettings controls autoscaling from minInstances to maxInstances
livenessProbe and readinessProbe configure health checks
requestSettings controls timeout and concurrency per instance

📐 Traffic Splitting for Canary Deployments

Deploy a new model version alongside the existing one, then gradually shift traffic:

# Endpoint with two deployments and traffic split
resource "azapi_resource" "online_endpoint" {
  type      = "Microsoft.MachineLearningServices/workspaces/onlineEndpoints@2025-06-01"
  name      = "${var.environment}-${var.model_name}"
  parent_id = azurerm_machine_learning_workspace.this.id
  location  = var.location

  identity {
    type = "SystemAssigned"
  }

  body = {
    properties = {
      authMode = var.auth_mode
      traffic = {
        "blue"  = 90   # Current stable version
        "green" = 10   # New canary version
      }
    }
  }
}

# Blue deployment (current version)
resource "azapi_resource" "blue" {
  type      = "Microsoft.MachineLearningServices/workspaces/onlineEndpoints/deployments@2025-06-01"
  name      = "blue"
  parent_id = azapi_resource.online_endpoint.id
  location  = var.location
  # ... deployment config for v1
}

# Green deployment (new version)
resource "azapi_resource" "green" {
  type      = "Microsoft.MachineLearningServices/workspaces/onlineEndpoints/deployments@2025-06-01"
  name      = "green"
  parent_id = azapi_resource.online_endpoint.id
  location  = var.location
  # ... deployment config for v2
}

Canary rollout workflow: Deploy green with 10% traffic. Monitor error rates and latency. If healthy, update traffic to "green" = 100, then delete blue. If unhealthy, set "blue" = 100 to roll back instantly.

Mirror Traffic for Shadow Testing

Test a new deployment without affecting users:

body = {
  properties = {
    authMode = var.auth_mode
    traffic = {
      "blue" = 100      # All live traffic to blue
    }
    mirrorTraffic = {
      "green" = 10       # Copy 10% of requests to green (responses discarded)
    }
  }
}

Mirror traffic sends a copy of live requests to the green deployment, but responses are discarded. This lets you validate the new model's behavior under real traffic patterns without any user impact.

📐 Data Collection for Model Monitoring

Enable request/response logging to catch model drift:

body = {
  kind = "Managed"
  properties = {
    model         = var.model_uri
    instanceType  = var.instance_type
    appInsightsEnabled = true

    dataCollector = {
      collections = {
        model_inputs = {
          dataCollectionMode = "Enabled"
          dataId             = var.data_asset_id
          samplingRate       = var.sampling_rate
        }
        model_outputs = {
          dataCollectionMode = "Enabled"
          dataId             = var.data_asset_id
          samplingRate       = var.sampling_rate
        }
      }
      requestLogging = {
        captureHeaders = ["Content-Type", "x-request-id"]
      }
      rollingRate = "Hour"
    }
  }
}

Data collection captures model inputs and outputs to a registered data asset. Use it for drift detection, fairness monitoring, and retraining triggers.

📐 Environment Configuration

# environments/dev.tfvars
model_name              = "fraud-detector"
deployment_name         = "blue"
model_uri               = "azureml:fraud-detector:2"
instance_type           = "Standard_DS2_v2"
min_instances           = 1
max_instances           = 2
max_concurrent_requests = 5
auth_mode               = "Key"
public_network_access   = "Enabled"

# environments/prod.tfvars
model_name              = "fraud-detector"
deployment_name         = "blue"
model_uri               = "azureml:fraud-detector:2"
instance_type           = "Standard_DS3_v2"
min_instances           = 2
max_instances           = 8
max_concurrent_requests = 10
auth_mode               = "AADToken"
public_network_access   = "Disabled"

🧪 Invoke the Endpoint

from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

ml_client = MLClient(
    DefaultAzureCredential(),
    subscription_id="...",
    resource_group_name="...",
    workspace_name="...",
)

result = ml_client.online_endpoints.invoke(
    endpoint_name="prod-fraud-detector",
    deployment_name="blue",
    request_file="sample_request.json",
)

print(result)

Or via REST:

curl -X POST \
  "https://prod-fraud-detector.eastus.inference.ml.azure.com/score" \
  -H "Authorization: Bearer $TOKEN" \
  -H "Content-Type: application/json" \
  -d '{"data": [[0.5, 1.2, 3.4, 0.8]]}'

⚠️ Gotchas and Tips

azapi is required. The azurerm provider doesn't have native azurerm_machine_learning_online_endpoint or azurerm_machine_learning_online_deployment resources. Use azapi_resource with the Microsoft.MachineLearningServices API.

Deployment creation takes 8-15 minutes. Provisioning VMs, pulling containers, loading models, and running health probes takes time. Plan for this in CI/CD pipelines.

Register models before deploying. The model field references a registered model in the workspace Model Registry (format: azureml:model-name:version). Register models via the SDK or CLI before running terraform apply.

Use AADToken in production. API keys work but require rotation. AAD token auth integrates with managed identity, eliminating key management entirely.

Scale to zero is not supported for managed endpoints. Unlike serverless compute, managed online endpoints require at least one instance running. If you need scale-to-zero, consider serverless endpoints (currently in preview).

Mirror traffic before canary. Mirror first (responses discarded, no user impact) to validate the model handles real request shapes correctly. Then switch to traffic splitting for a live canary test.

⏭️ What's Next

This is Post 2 of the Azure ML Pipelines & MLOps with Terraform series.

Post 1: Azure ML Workspace 🔬
Post 2: Azure ML Online Endpoints (you are here) 🚀
Post 3: Azure ML Feature Store
Post 4: Azure ML Pipelines + Azure DevOps

Your model is in production. Managed online endpoint with autoscaling, blue/green traffic splitting, mirror traffic for shadow testing, and data collection for drift monitoring. From workspace to production, all in Terraform. 🚀

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Vertex AI Endpoints: Deploy Your Model to Production with Terraform 🚀

Suhas Mallesh — Sun, 12 Apr 2026 07:00:00 +0000

Your model is trained. Now deploy it to a scalable endpoint with autoscaling, traffic splitting for canary rollouts, and request-response logging. Here's how to deploy Vertex AI endpoints with Terraform.

In the previous post, we set up the Vertex AI Workbench - the workspace where your team trains models. Now comes the production side: taking a trained model and deploying it to an endpoint that your applications can call for real-time predictions.

Vertex AI uses a three-step deployment: upload a Model to the Model Registry, create an Endpoint (the HTTPS URL), then deploy the model to the endpoint with compute resources. Terraform provisions the endpoint and configures autoscaling, traffic splitting, and logging. The SDK handles model upload and deployment. 🎯

🏗️ The Deployment Architecture

Model Registry (uploaded model + container)
    ↓
Endpoint (HTTPS URL, traffic routing)
    ↓
Deployed Model (machine type, replicas, autoscaling)

Component	What It Defines
Model	Container image + model artifacts in GCS
Endpoint	Stable HTTPS prediction URL
Deployed Model	Machine type, replica count, autoscaling
Traffic Split	Percentage routing between model versions

The endpoint URL stays stable across model versions. You update the deployed model behind it without changing your application code.

🔧 Terraform: Create the Endpoint

Endpoint Resource

# endpoint/main.tf

resource "google_vertex_ai_endpoint" "this" {
  name         = "${var.environment}-${var.model_name}-endpoint"
  display_name = "${var.environment}-${var.model_name}"
  description  = "Production endpoint for ${var.model_name}"
  location     = var.region
  project      = var.project_id

  labels = {
    environment = var.environment
    model       = var.model_name
    managed_by  = "terraform"
  }
}

Endpoint with Private Network (Production)

For production, deploy with a private endpoint inside your VPC:

resource "google_vertex_ai_endpoint" "private" {
  name         = "${var.environment}-${var.model_name}-endpoint"
  display_name = "${var.environment}-${var.model_name}"
  location     = var.region
  project      = var.project_id
  network      = "projects/${data.google_project.this.number}/global/networks/${var.vpc_network_name}"

  depends_on = [google_service_networking_connection.vertex_vpc]
}

resource "google_service_networking_connection" "vertex_vpc" {
  network                 = var.vpc_network_id
  service                 = "servicenetworking.googleapis.com"
  reserved_peering_ranges = [google_compute_global_address.vertex_range.name]
}

resource "google_compute_global_address" "vertex_range" {
  name          = "${var.environment}-vertex-range"
  purpose       = "VPC_PEERING"
  address_type  = "INTERNAL"
  prefix_length = 16
  network       = var.vpc_network_id
  project       = var.project_id
}

Prediction Logging to BigQuery

resource "google_bigquery_dataset" "predictions" {
  dataset_id = "${var.environment}_prediction_logs"
  location   = var.region
  project    = var.project_id
}

resource "google_vertex_ai_endpoint" "logged" {
  name         = "${var.environment}-${var.model_name}-endpoint"
  display_name = "${var.environment}-${var.model_name}"
  location     = var.region
  project      = var.project_id

  predict_request_response_logging_config {
    enabled       = var.enable_prediction_logging
    sampling_rate = var.logging_sample_rate

    bigquery_destination {
      output_uri = "bq://${var.project_id}.${google_bigquery_dataset.predictions.dataset_id}.request_response"
    }
  }
}

Log a sample of prediction requests and responses to BigQuery for model monitoring, drift detection, and debugging.

🐍 Model Upload and Deployment (SDK)

Terraform creates the endpoint. The Vertex AI SDK uploads the model and deploys it with compute resources:

# deploy.py

from google.cloud import aiplatform
import json

with open("deploy_config.json") as f:
    config = json.load(f)

aiplatform.init(
    project=config["project_id"],
    location=config["region"],
)

# Upload model to Model Registry
model = aiplatform.Model.upload(
    display_name=config["model_name"],
    artifact_uri=config["model_artifact_uri"],  # gs://bucket/model/
    serving_container_image_uri=config["serving_image"],
    serving_container_predict_route="/predict",
    serving_container_health_route="/health",
)

# Get the Terraform-created endpoint
endpoint = aiplatform.Endpoint(config["endpoint_resource_name"])

# Deploy model to endpoint
model.deploy(
    endpoint=endpoint,
    machine_type=config["machine_type"],
    min_replica_count=config["min_replicas"],
    max_replica_count=config["max_replicas"],
    traffic_percentage=100,
    deploy_request_timeout=1800,
)

print(f"Model deployed to: {endpoint.resource_name}")

Terraform Config Output for SDK

# endpoint/config.tf

resource "local_file" "deploy_config" {
  filename = "${path.module}/deploy_config.json"
  content = jsonencode({
    project_id             = var.project_id
    region                 = var.region
    model_name             = "${var.environment}-${var.model_name}"
    model_artifact_uri     = var.model_artifact_uri
    serving_image          = var.serving_container_image
    endpoint_resource_name = google_vertex_ai_endpoint.this.name
    machine_type           = var.machine_type
    min_replicas           = var.min_replicas
    max_replicas           = var.max_replicas
  })
}

📐 Traffic Splitting for Canary Deployments

Deploy a new model version alongside the existing one and gradually shift traffic:

# Canary: deploy new version with 10% traffic
new_model = aiplatform.Model.upload(
    display_name="fraud-detector-v3",
    artifact_uri="gs://ml-models/fraud-detector/v3/",
    serving_container_image_uri=config["serving_image"],
)

new_model.deploy(
    endpoint=endpoint,
    machine_type=config["machine_type"],
    min_replica_count=1,
    max_replica_count=4,
    traffic_percentage=10,  # 10% to new model
)

# After validation, shift 100% to new model
endpoint.update(traffic_split={
    new_model.id: 100,
})

The endpoint URL doesn't change. Your application sends predictions to the same endpoint while you shift traffic from the old model to the new one.

📐 Model Garden Deployment (Terraform-Only)

For open models from Model Garden (Gemma, Llama, PaLiGemma), use the newer Terraform resource that handles everything in one step:

resource "google_vertex_ai_endpoint_with_model_garden_deployment" "gemma" {
  publisher_model_name = "publishers/google/models/gemma3@gemma-3-1b-it"
  location             = var.region

  model_config {
    accept_eula = true
  }

  deploy_config {
    dedicated_resources {
      machine_spec {
        machine_type      = "g2-standard-12"
        accelerator_type  = "NVIDIA_L4"
        accelerator_count = 1
      }
      min_replica_count = 1
    }
  }
}

One resource creates the endpoint, uploads the model, and deploys it. This works for Model Garden models only - custom models still need the endpoint + SDK pattern.

📐 Environment Configuration

# environments/dev.tfvars
model_name              = "fraud-detector"
machine_type            = "n1-standard-4"
min_replicas            = 1
max_replicas            = 2
enable_prediction_logging = false
logging_sample_rate     = 0.0

# environments/prod.tfvars
model_name              = "fraud-detector"
machine_type            = "n1-standard-8"
min_replicas            = 2
max_replicas            = 10
enable_prediction_logging = true
logging_sample_rate     = 0.1   # Log 10% of requests

🧪 Invoke the Endpoint

from google.cloud import aiplatform

aiplatform.init(project="my-project", location="us-central1")

endpoint = aiplatform.Endpoint("ENDPOINT_ID")

prediction = endpoint.predict(
    instances=[{"features": [0.5, 1.2, 3.4, 0.8]}]
)

print(prediction.predictions)

⚠️ Gotchas and Tips

Endpoint creation is fast, deployment is slow. Creating an endpoint takes seconds. Deploying a model (provisioning VMs, loading containers, health checks) takes 10-20 minutes. Plan for this in your CI/CD pipeline.

Autoscaling uses replica count. Set min_replica_count to your baseline and max_replica_count to your peak. Vertex AI scales based on CPU utilization and request queue depth automatically.

GPU quota must be requested. NVIDIA L4, T4, A100 accelerators require quota increases. Request early in your project setup.

Model Registry keeps versions. Every Model.upload creates a new version in the registry. Old versions remain available for rollback. Use the traffic_split to redirect traffic back to a previous version if needed.

Prediction logging costs. BigQuery logging at 10% sampling rate is manageable. At 100%, costs scale with request volume. Use sampling for high-traffic endpoints.

⏭️ What's Next

This is Post 2 of the GCP ML Pipelines & MLOps with Terraform series.

Post 1: Vertex AI Workbench 🔬
Post 2: Vertex AI Endpoints - Deploy to Prod (you are here) 🚀
Post 3: Vertex AI Feature Store
Post 4: Vertex AI Pipelines + Cloud Build

Your model is in production. Stable endpoint URL, autoscaling replicas, traffic splitting for canary rollouts, and prediction logging to BigQuery. From training to production, all in Terraform and Python. 🚀

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

SageMaker Endpoints: Deploy Your Model to Production with Terraform 🚀

Suhas Mallesh — Sat, 11 Apr 2026 00:48:06 +0000

Training a model is half the battle. Deploying it to a scalable, auto-scaling endpoint with blue/green deployment and rollback is the other half. Here's how to deploy SageMaker real-time endpoints with Terraform.

In the previous post, we set up the SageMaker Studio domain - the workspace where your team trains models. Now comes the production side: taking a trained model and deploying it to a scalable HTTPS endpoint that your applications can call for real-time predictions.

SageMaker endpoints involve three Terraform resources: a Model (what to serve), an Endpoint Configuration (how to serve it), and an Endpoint (the live HTTPS URL). Add autoscaling and deployment policies on top, and you have a production-grade inference system. 🎯

🏗️ The Three-Layer Architecture

Model (container image + model artifacts in S3)
    ↓
Endpoint Configuration (instance type, count, variants)
    ↓
Endpoint (HTTPS URL, auto-scaling, blue/green deployment)

Resource	What It Defines
Model	Container image + S3 model artifacts + IAM role
Endpoint Config	Instance type, initial count, production variants
Endpoint	Live endpoint with deployment policy and autoscaling

Separating these layers means you can update the model without touching the endpoint config, or change instance types without retraining.

🔧 Terraform: Deploy to Production

Model Configuration

# variables.tf

variable "model_config" {
  description = "Model deployment configuration. Change to deploy new models."
  type = object({
    name           = string
    image_uri      = string  # ECR container image
    model_data_url = string  # S3 path to model.tar.gz
    instance_type  = string
    instance_count = number
  })
}

IAM Role for Model Execution

# inference/iam.tf

resource "aws_iam_role" "sagemaker_execution" {
  name = "${var.environment}-sagemaker-inference"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "sagemaker.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "model_access" {
  name = "model-s3-ecr-access"
  role = aws_iam_role.sagemaker_execution.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect   = "Allow"
        Action   = ["s3:GetObject"]
        Resource = "${var.model_bucket_arn}/*"
      },
      {
        Effect = "Allow"
        Action = [
          "ecr:GetAuthorizationToken",
          "ecr:BatchGetImage",
          "ecr:GetDownloadUrlForLayer"
        ]
        Resource = "*"
      },
      {
        Effect   = "Allow"
        Action   = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"]
        Resource = "*"
      }
    ]
  })
}

The Model

# inference/model.tf

resource "aws_sagemaker_model" "this" {
  name               = "${var.environment}-${var.model_config.name}"
  execution_role_arn = aws_iam_role.sagemaker_execution.arn

  primary_container {
    image          = var.model_config.image_uri
    model_data_url = var.model_config.model_data_url
    environment = {
      SAGEMAKER_PROGRAM = "inference.py"
    }
  }

  tags = {
    Environment = var.environment
    Model       = var.model_config.name
  }
}

image is your serving container - either an AWS Deep Learning Container or your custom container from ECR. model_data_url points to model.tar.gz in S3 containing your trained weights and inference code.

Endpoint Configuration

# inference/endpoint_config.tf

resource "aws_sagemaker_endpoint_configuration" "this" {
  name = "${var.environment}-${var.model_config.name}-config"

  production_variants {
    variant_name           = "primary"
    model_name             = aws_sagemaker_model.this.name
    initial_instance_count = var.model_config.instance_count
    instance_type          = var.model_config.instance_type
    initial_variant_weight = 1.0
  }

  tags = {
    Environment = var.environment
    Model       = var.model_config.name
  }

  lifecycle {
    create_before_destroy = true
  }
}

create_before_destroy = true is important. When you update the endpoint config (new model version, different instance type), Terraform creates the new config before deleting the old one. This prevents downtime during updates.

The Endpoint

# inference/endpoint.tf

resource "aws_sagemaker_endpoint" "this" {
  name                 = "${var.environment}-${var.model_config.name}"
  endpoint_config_name = aws_sagemaker_endpoint_configuration.this.name

  deployment_config {
    blue_green_update_policy {
      traffic_routing_configuration {
        type                     = "CANARY"
        canary_size {
          type  = "INSTANCE_COUNT"
          value = 1
        }
        wait_interval_in_seconds = 300
      }

      maximum_execution_timeout_in_seconds = 1800
      termination_wait_in_seconds          = 120
    }

    auto_rollback_configuration {
      alarms {
        alarm_name = aws_cloudwatch_metric_alarm.endpoint_errors.alarm_name
      }
    }
  }

  tags = {
    Environment = var.environment
    Model       = var.model_config.name
  }
}

Canary deployment: Routes traffic to 1 instance on the new fleet first, waits 5 minutes, then shifts remaining traffic. If the CloudWatch alarm fires during the canary phase, SageMaker automatically rolls back.

Autoscaling

# inference/autoscaling.tf

resource "aws_appautoscaling_target" "endpoint" {
  max_capacity       = var.autoscaling_max
  min_capacity       = var.model_config.instance_count
  resource_id        = "endpoint/${aws_sagemaker_endpoint.this.name}/variant/primary"
  scalable_dimension = "sagemaker:variant:DesiredInstanceCount"
  service_namespace  = "sagemaker"
}

resource "aws_appautoscaling_policy" "endpoint" {
  name               = "${var.environment}-endpoint-scaling"
  policy_type        = "TargetTrackingScaling"
  resource_id        = aws_appautoscaling_target.endpoint.resource_id
  scalable_dimension = aws_appautoscaling_target.endpoint.scalable_dimension
  service_namespace  = aws_appautoscaling_target.endpoint.service_namespace

  target_tracking_scaling_policy_configuration {
    target_value = var.target_invocations_per_instance

    predefined_metric_specification {
      predefined_metric_type = "SageMakerVariantInvocationsPerInstance"
    }

    scale_in_cooldown  = 300
    scale_out_cooldown = 60
  }
}

SageMakerVariantInvocationsPerInstance scales based on requests per instance. Set the target to the maximum your instance can handle (determined by load testing). The policy adds instances when traffic exceeds the target and removes them when traffic drops.

Monitoring

# inference/monitoring.tf

resource "aws_cloudwatch_metric_alarm" "endpoint_errors" {
  alarm_name          = "${var.environment}-endpoint-5xx-errors"
  comparison_operator = "GreaterThanThreshold"
  evaluation_periods  = 2
  metric_name         = "Invocation5XXErrors"
  namespace           = "AWS/SageMaker"
  period              = 60
  statistic           = "Sum"
  threshold           = 5
  alarm_description   = "Endpoint returning 5xx errors"
  alarm_actions       = [var.sns_alert_topic_arn]

  dimensions = {
    EndpointName = aws_sagemaker_endpoint.this.name
    VariantName  = "primary"
  }
}

resource "aws_cloudwatch_metric_alarm" "endpoint_latency" {
  alarm_name          = "${var.environment}-endpoint-high-latency"
  comparison_operator = "GreaterThanThreshold"
  evaluation_periods  = 3
  metric_name         = "ModelLatency"
  namespace           = "AWS/SageMaker"
  period              = 60
  statistic           = "Average"
  threshold           = var.latency_threshold_ms * 1000  # microseconds
  alarm_description   = "Model inference latency too high"
  alarm_actions       = [var.sns_alert_topic_arn]

  dimensions = {
    EndpointName = aws_sagemaker_endpoint.this.name
    VariantName  = "primary"
  }
}

The error alarm also feeds into the blue/green deployment's auto_rollback_configuration. If 5xx errors spike during a deployment, SageMaker automatically rolls back to the previous fleet.

📐 Environment Configuration

# environments/dev.tfvars
model_config = {
  name           = "fraud-detector-v2"
  image_uri      = "123456789012.dkr.ecr.us-east-1.amazonaws.com/ml-serving:latest"
  model_data_url = "s3://ml-models-dev/fraud-detector/v2/model.tar.gz"
  instance_type  = "ml.t2.medium"
  instance_count = 1
}
autoscaling_max = 2
target_invocations_per_instance = 100

# environments/prod.tfvars
model_config = {
  name           = "fraud-detector-v2"
  image_uri      = "123456789012.dkr.ecr.us-east-1.amazonaws.com/ml-serving:v2.1.0"
  model_data_url = "s3://ml-models-prod/fraud-detector/v2/model.tar.gz"
  instance_type  = "ml.c5.xlarge"
  instance_count = 2
}
autoscaling_max = 10
target_invocations_per_instance = 500

Deploying a new model version: Update model_data_url to the new S3 path, run terraform apply. The canary deployment creates a new fleet, validates it, and shifts traffic automatically.

🧪 Invoke the Endpoint

import boto3
import json

client = boto3.client("sagemaker-runtime")

response = client.invoke_endpoint(
    EndpointName="prod-fraud-detector-v2",
    ContentType="application/json",
    Body=json.dumps({"features": [0.5, 1.2, 3.4, 0.8]})
)

prediction = json.loads(response["Body"].read())
print(prediction)

⚠️ Gotchas and Tips

Load test before setting autoscaling targets. The target_invocations_per_instance should be based on actual load testing, not guesswork. Use SageMaker Inference Recommender to find the optimal instance type and max throughput.

Endpoint creation takes 5-10 minutes. SageMaker provisions instances, downloads the container, loads the model, and runs health checks. Plan for this in your deployment pipeline.

Use create_before_destroy on endpoint configs. Without it, Terraform deletes the old config before creating the new one, causing downtime.

Pin container image tags. Use versioned tags (v2.1.0) instead of latest in production. This ensures reproducible deployments and meaningful rollbacks.

Serverless inference for low traffic. If your endpoint gets fewer than ~100 requests/hour, consider serverless_config on the endpoint configuration instead of instance-based. It scales to zero and charges per request.

⏭️ What's Next

This is Post 2 of the ML Pipelines & MLOps with Terraform series.

Post 1: SageMaker Studio Domain 🔬
Post 2: SageMaker Endpoints - Deploy to Prod (you are here) 🚀
Post 3: SageMaker Feature Store
Post 4: SageMaker Pipelines - CI/CD for ML

Your model is in production. Real-time HTTPS endpoint, autoscaling based on traffic, canary deployments with automatic rollback, and monitoring that pages you before customers notice. From notebook to production, all in Terraform. 🚀

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Azure ML Workspace with Terraform: Your ML Platform on Azure 🔬

Suhas Mallesh — Fri, 03 Apr 2026 07:00:00 +0000

Azure Machine Learning workspace is the hub for all ML activities - experiments, models, endpoints, pipelines. It requires four dependent services. Here's how to provision the entire platform with Terraform including compute instances and clusters.

In Series 1-3, we worked with managed AI services - AI Foundry for models, AI Search for RAG, Agent Service for orchestration. Series 5 shifts to custom ML - training your own models, deploying endpoints, managing features, and building CI/CD pipelines.

It starts with an Azure Machine Learning workspace. The workspace is the top-level resource for all ML activities: experiments, datasets, models, compute targets, endpoints, and pipelines live here. Unlike a simple resource, the workspace requires four dependent services before it can be created: Storage Account, Key Vault, Application Insights, and Container Registry. Terraform provisions the entire stack. 🎯

🏗️ Workspace Architecture

Component	What It Does
Workspace	Central hub for ML experiments, models, and pipelines
Storage Account	Default datastore for datasets, model artifacts, logs
Key Vault	Secrets, connection strings, API keys
Application Insights	Experiment tracking, endpoint monitoring
Container Registry	Custom training images, model deployment containers
Compute Instance	Per-user notebook/IDE (JupyterLab, VS Code)
Compute Cluster	Auto-scaling training cluster (CPU or GPU)

All four dependent services must exist before the workspace. Terraform handles the dependency ordering automatically.

🔧 Terraform: The Full Workspace Setup

Dependent Services

# ml/dependencies.tf

resource "azurerm_storage_account" "ml" {
  name                     = "${var.environment}ml${random_string.suffix.result}"
  location                 = azurerm_resource_group.ml.location
  resource_group_name      = azurerm_resource_group.ml.name
  account_tier             = "Standard"
  account_replication_type = var.storage_replication
  min_tls_version          = "TLS1_2"
  allow_nested_items_to_be_public = false

  tags = var.tags
}

resource "azurerm_key_vault" "ml" {
  name                       = "${var.environment}ml${random_string.suffix.result}kv"
  location                   = azurerm_resource_group.ml.location
  resource_group_name        = azurerm_resource_group.ml.name
  tenant_id                  = data.azurerm_client_config.current.tenant_id
  sku_name                   = "standard"
  purge_protection_enabled   = true
  enable_rbac_authorization  = true

  tags = var.tags
}

resource "azurerm_application_insights" "ml" {
  name                = "${var.environment}-ml-insights"
  location            = azurerm_resource_group.ml.location
  resource_group_name = azurerm_resource_group.ml.name
  application_type    = "web"

  tags = var.tags
}

resource "azurerm_container_registry" "ml" {
  name                = "${var.environment}ml${random_string.suffix.result}acr"
  location            = azurerm_resource_group.ml.location
  resource_group_name = azurerm_resource_group.ml.name
  sku                 = var.acr_sku
  admin_enabled       = false

  tags = var.tags
}

Container Registry is optional but recommended. Without it, the workspace uses Azure-managed image building. With it, you control custom training images and model serving containers. Set admin_enabled = false and use managed identity instead.

The Workspace

# ml/workspace.tf

resource "azurerm_machine_learning_workspace" "this" {
  name                          = "${var.environment}-ml-workspace"
  location                      = azurerm_resource_group.ml.location
  resource_group_name           = azurerm_resource_group.ml.name
  application_insights_id       = azurerm_application_insights.ml.id
  key_vault_id                  = azurerm_key_vault.ml.id
  storage_account_id            = azurerm_storage_account.ml.id
  container_registry_id         = azurerm_container_registry.ml.id
  public_network_access_enabled = var.public_network_access

  identity {
    type = "SystemAssigned"
  }

  tags = var.tags
}

The workspace creates a system-assigned managed identity that accesses the dependent services. No keys or connection strings to manage.

Compute Instance (Per-User IDE)

# ml/compute_instance.tf

resource "azurerm_machine_learning_compute_instance" "this" {
  for_each = var.compute_instances

  name                          = each.key
  machine_learning_workspace_id = azurerm_machine_learning_workspace.this.id
  location                      = azurerm_resource_group.ml.location
  virtual_machine_size          = each.value.vm_size
  node_public_ip_enabled        = var.public_network_access

  identity {
    type = "SystemAssigned"
  }

  tags = merge(var.tags, {
    Team = each.value.team
    User = each.key
  })
}

Each data scientist gets their own compute instance with JupyterLab and VS Code access. Instances stop and start independently - you only pay when they're running.

Compute Cluster (Training)

# ml/compute_cluster.tf

resource "azurerm_machine_learning_compute_cluster" "training" {
  name                          = "${var.environment}-training"
  machine_learning_workspace_id = azurerm_machine_learning_workspace.this.id
  location                      = azurerm_resource_group.ml.location
  vm_size                       = var.training_vm_size
  vm_priority                   = var.training_vm_priority

  identity {
    type = "SystemAssigned"
  }

  scale_settings {
    min_node_count                       = 0
    max_node_count                       = var.training_max_nodes
    scale_down_nodes_after_idle_duration  = "PT${var.scale_down_minutes}M"
  }

  tags = var.tags
}

min_node_count = 0 is the cost control key. The cluster scales to zero when no training jobs are running. You pay nothing for idle compute. scale_down_nodes_after_idle_duration controls how quickly nodes are released after a job finishes.

vm_priority = "LowPriority" saves up to 80% on training costs. Low-priority VMs can be evicted, so use them for fault-tolerant training jobs with checkpointing.

📐 Environment Configuration

# environments/dev.tfvars
environment            = "dev"
public_network_access  = true
storage_replication    = "LRS"
acr_sku                = "Basic"
training_vm_size       = "Standard_DS3_v2"
training_vm_priority   = "Dedicated"
training_max_nodes     = 2
scale_down_minutes     = 15

compute_instances = {
  "ds-dev-1" = {
    vm_size = "Standard_DS3_v2"
    team    = "ml-team"
  }
}

# environments/prod.tfvars
environment            = "prod"
public_network_access  = false
storage_replication    = "GRS"
acr_sku                = "Standard"
training_vm_size       = "Standard_NC6s_v3"   # GPU
training_vm_priority   = "LowPriority"
training_max_nodes     = 8
scale_down_minutes     = 30

compute_instances = {
  "ds-lead" = {
    vm_size = "Standard_DS4_v2"
    team    = "ml-team"
  }
  "ds-engineer-1" = {
    vm_size = "Standard_DS3_v2"
    team    = "ml-team"
  }
  "ds-engineer-2" = {
    vm_size = "Standard_DS3_v2"
    team    = "ml-team"
  }
}

Dev: Public access, LRS storage, small dedicated cluster for quick iteration.
Prod: Private access, GRS storage, GPU cluster with low-priority VMs for cost-efficient training.

🔧 RBAC for Team Access

# ml/rbac.tf

# Data scientists get Contributor on the workspace
resource "azurerm_role_assignment" "ds_contributor" {
  for_each = var.data_scientist_principals

  scope                = azurerm_machine_learning_workspace.this.id
  role_definition_name = "AzureML Compute Operator"
  principal_id         = each.value
}

# ML engineers get full workspace access
resource "azurerm_role_assignment" "mle_contributor" {
  for_each = var.ml_engineer_principals

  scope                = azurerm_machine_learning_workspace.this.id
  role_definition_name = "Contributor"
  principal_id         = each.value
}

Use AzureML Compute Operator for data scientists who need to run experiments but shouldn't modify workspace settings. Use Contributor for ML engineers who manage the full lifecycle.

🔧 Security Hardening

Control	Dev	Prod
Network access	Public	Private (managed VNet)
Storage	LRS, TLS 1.2	GRS, TLS 1.2, no public blob
Key Vault	RBAC auth	RBAC auth, purge protection
Container Registry	Basic, no admin	Standard, managed identity
Compute	Public IP	No public IP, subnet attached
Identity	System-assigned	System-assigned + RBAC

For production, set public_network_access_enabled = false on the workspace and use managed virtual network isolation. All compute instances and clusters run inside a managed VNet with outbound rules controlled by the workspace.

⚠️ Gotchas and Tips

Globally unique names required. Storage account and ACR names must be globally unique. Use random_string suffixes to avoid conflicts across environments.

Container Registry costs. Basic SKU is $5/month. Standard is $20/month. If you're not building custom training images yet, you can skip the ACR initially and add it later.

Compute instance auto-stop. Unlike clusters, compute instances don't auto-stop by default. Set up an idle shutdown schedule through the workspace settings or Azure Policy to prevent overnight charges.

Workspace deletion is complex. Deleting a workspace doesn't automatically delete dependent resources (storage, KV, ACR). Terraform handles this correctly with depends_on, but manual deletion requires cleaning up each resource individually.

SDK v1 is deprecated. Azure ML SDK v1 reached end-of-support in March 2025. Use SDK v2 (azure-ai-ml) for all new development. Terraform provisions the infrastructure; SDK v2 handles experiments and models.

⏭️ What's Next

This is Post 1 of the Azure ML Pipelines & MLOps with Terraform series.

Post 1: Azure ML Workspace (you are here) 🔬
Post 2: Azure ML Online Endpoints - Deploy to Prod
Post 3: Azure ML Feature Store
Post 4: Azure ML Pipelines + Azure DevOps

Your ML platform is provisioned. Workspace, storage, key vault, ACR, compute instances for notebooks, auto-scaling clusters for training - all in Terraform. The foundation for model development, deployment, and production ML pipelines. 🔬

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Vertex AI Workbench with Terraform: Your ML Workspace on GCP 🔬

Suhas Mallesh — Thu, 02 Apr 2026 07:00:00 +0000

Vertex AI Workbench is the JupyterLab IDE for ML on GCP - pre-installed ML frameworks, Vertex AI integration, and GPU support out of the box. Here's how to provision instances with Terraform including networking, IAM, auto-shutdown, and custom containers.

In Series 1-3, we worked with managed AI services - Vertex AI for models, RAG Engine for retrieval, ADK for agents. Series 5 shifts to custom ML - training your own models, deploying endpoints, managing features, and building ML pipelines.

It starts with a development environment. Vertex AI Workbench provides JupyterLab instances backed by Compute Engine VMs, pre-loaded with ML frameworks (TensorFlow, PyTorch, JAX, scikit-learn), the Vertex AI SDK, and direct integration with GCS, BigQuery, and Vertex AI services. Each instance is a full ML workspace. Terraform provisions them consistently across your team. 🎯

🏗️ Workbench Architecture

Component	What It Does
Workbench Instance	JupyterLab on a Compute Engine VM
VM Image	Pre-built with ML frameworks and Vertex SDK
Custom Container	Your own Docker image with specific dependencies
Data Disk	Persistent storage for notebooks and datasets
Service Account	Controls access to GCS, BigQuery, Vertex AI
Network	VPC, subnet, firewall rules for isolation

Unlike SageMaker's domain-based approach, Workbench instances are standalone VMs. Each data scientist gets their own instance with full control over machine type, GPU, and installed packages. Terraform standardizes the configuration across the team.

🔧 Terraform: The Full Workbench Setup

APIs and Networking

# workbench/apis.tf

resource "google_project_service" "required" {
  for_each = toset([
    "notebooks.googleapis.com",
    "aiplatform.googleapis.com",
    "compute.googleapis.com",
    "storage.googleapis.com",
    "bigquery.googleapis.com",
  ])
  project = var.project_id
  service = each.value
}

# VPC for Workbench instances
resource "google_compute_network" "ml" {
  name                    = "${var.environment}-ml-network"
  auto_create_subnetworks = false
  project                 = var.project_id
}

resource "google_compute_subnetwork" "ml" {
  name                     = "${var.environment}-ml-subnet"
  ip_cidr_range            = var.subnet_cidr
  region                   = var.region
  network                  = google_compute_network.ml.id
  private_ip_google_access = true  # Access Google APIs without public IP
  project                  = var.project_id
}

# Allow internal traffic for JupyterLab proxy
resource "google_compute_firewall" "internal" {
  name    = "${var.environment}-ml-internal"
  network = google_compute_network.ml.id
  project = var.project_id

  allow {
    protocol = "tcp"
    ports    = ["8080", "443"]
  }

  source_ranges = [var.subnet_cidr]
}

private_ip_google_access = true is essential. It lets Workbench instances access Google APIs (GCS, BigQuery, Vertex AI) without a public IP address, keeping your ML environment network-isolated.

Service Account

# workbench/iam.tf

resource "google_service_account" "workbench" {
  account_id   = "${var.environment}-workbench-sa"
  display_name = "Workbench ML Service Account"
  project      = var.project_id
}

# Access to training data in GCS
resource "google_project_iam_member" "storage_user" {
  project = var.project_id
  role    = "roles/storage.objectUser"
  member  = "serviceAccount:${google_service_account.workbench.email}"
}

# Access to Vertex AI for training and prediction
resource "google_project_iam_member" "vertex_ai_user" {
  project = var.project_id
  role    = "roles/aiplatform.user"
  member  = "serviceAccount:${google_service_account.workbench.email}"
}

# Access to BigQuery for data exploration
resource "google_project_iam_member" "bigquery_user" {
  project = var.project_id
  role    = "roles/bigquery.user"
  member  = "serviceAccount:${google_service_account.workbench.email}"
}

# Read Artifact Registry for custom containers
resource "google_project_iam_member" "artifact_reader" {
  project = var.project_id
  role    = "roles/artifactregistry.reader"
  member  = "serviceAccount:${google_service_account.workbench.email}"
}

Workbench Instance

# workbench/instance.tf

resource "google_workbench_instance" "this" {
  for_each = var.workbench_instances

  name     = "${var.environment}-${each.key}"
  location = var.zone
  project  = var.project_id

  gce_setup {
    machine_type = each.value.machine_type

    vm_image {
      project = "cloud-notebooks-managed"
      family  = "workbench-instances"
    }

    # Optional GPU
    dynamic "accelerator_configs" {
      for_each = each.value.gpu_type != null ? [1] : []
      content {
        type       = each.value.gpu_type
        core_count = each.value.gpu_count
      }
    }

    boot_disk {
      disk_type    = "PD_SSD"
      disk_size_gb = each.value.boot_disk_gb
    }

    data_disks {
      disk_type    = "PD_SSD"
      disk_size_gb = each.value.data_disk_gb
    }

    network_interfaces {
      network  = google_compute_network.ml.id
      subnet   = google_compute_subnetwork.ml.id
      nic_type = "GVNIC"
    }

    service_accounts {
      email = google_service_account.workbench.email
    }

    metadata = {
      idle-timeout-seconds   = each.value.idle_timeout
      install-nvidia-driver  = each.value.gpu_type != null ? "true" : "false"
      post-startup-script    = each.value.post_startup_script
    }

    enable_ip_forwarding = false
  }

  disable_proxy_access = false  # Enable JupyterLab proxy access

  labels = {
    environment = var.environment
    team        = each.value.team
    managed_by  = "terraform"
  }
}

Key configuration points:

vm_image.family = "workbench-instances" uses Google's pre-built image with ML frameworks
data_disks provides persistent storage separate from the boot disk
idle-timeout-seconds in metadata auto-stops the instance after inactivity
install-nvidia-driver auto-installs GPU drivers when GPU is attached
disable_proxy_access = false enables browser access to JupyterLab via IAP

Custom Container Alternative

For teams that need specific package versions or private packages, use a custom container:

resource "google_workbench_instance" "custom" {
  name     = "${var.environment}-custom-ml"
  location = var.zone
  project  = var.project_id

  gce_setup {
    machine_type = "n1-standard-8"

    container_image {
      repository = "${var.region}-docker.pkg.dev/${var.project_id}/ml-images/custom-workbench"
      tag        = "latest"
    }

    boot_disk {
      disk_type    = "PD_SSD"
      disk_size_gb = 150
    }

    network_interfaces {
      network  = google_compute_network.ml.id
      subnet   = google_compute_subnetwork.ml.id
    }

    service_accounts {
      email = google_service_account.workbench.email
    }
  }
}

Build your container from Google's base image, add your packages, push to Artifact Registry, and reference it in Terraform. Updates are a new tag and terraform apply.

📐 Environment Configuration

# environments/dev.tfvars
environment = "dev"
zone        = "us-central1-a"

workbench_instances = {
  "data-scientist-1" = {
    machine_type       = "n1-standard-4"
    gpu_type           = null          # No GPU for dev
    gpu_count          = 0
    boot_disk_gb       = 100
    data_disk_gb       = 200
    idle_timeout        = "3600"       # 1 hour auto-stop
    team               = "ml-team"
    post_startup_script = null
  }
}

# environments/prod.tfvars
environment = "prod"
zone        = "us-central1-a"

workbench_instances = {
  "ml-lead" = {
    machine_type       = "n1-standard-8"
    gpu_type           = "NVIDIA_TESLA_T4"
    gpu_count          = 1
    boot_disk_gb       = 150
    data_disk_gb       = 500
    idle_timeout        = "7200"       # 2 hour auto-stop
    team               = "ml-team"
    post_startup_script = "gs://my-bucket/scripts/setup.sh"
  }
  "ml-engineer-1" = {
    machine_type       = "n1-standard-8"
    gpu_type           = "NVIDIA_TESLA_T4"
    gpu_count          = 1
    boot_disk_gb       = 150
    data_disk_gb       = 500
    idle_timeout        = "7200"
    team               = "ml-team"
    post_startup_script = "gs://my-bucket/scripts/setup.sh"
  }
}

Cost control through idle_timeout. GPU instances cost $1-3/hour. A T4 instance left running overnight costs $24-72 in wasted spend. The idle-timeout-seconds metadata auto-stops the VM when JupyterLab kernels are idle.

🔧 Post-Startup Script for Team Standardization

Standardize packages and configs across all instances:

#!/bin/bash
# scripts/setup.sh - uploaded to GCS

# Install team-standard packages
pip install --upgrade \
  google-cloud-aiplatform \
  google-cloud-bigquery \
  pandas \
  scikit-learn \
  xgboost

# Clone team shared repos
git clone https://github.com/your-org/ml-utils.git /home/jupyter/ml-utils

# Set environment variables
echo 'export PROJECT_ID=your-project' >> /home/jupyter/.bashrc
echo 'export REGION=us-central1' >> /home/jupyter/.bashrc

Reference it in the instance metadata via post_startup_script. Runs once when the instance starts for the first time.

⚠️ Gotchas and Tips

User-managed notebooks are deprecated. As of January 2025, only Workbench Instances (google_workbench_instance) are supported. Don't use google_notebooks_instance - it's legacy.

GPU quota must be requested. T4, V100, A100 GPUs require quota increases in your GCP project. Request early - approval can take 1-2 business days.

Instance stops don't lose data. Stopping a Workbench instance preserves the boot and data disks. You only pay for disk storage when stopped. Starting it again restores your full environment.

Use data disks for large datasets. Boot disks are limited and get recreated on image updates. Store notebooks and datasets on the data disk (/home/jupyter) which persists independently.

Private Google Access vs public IP. With private_ip_google_access = true on the subnet, instances can reach GCS, BigQuery, and Vertex AI without a public IP. Combine with IAP for secure browser access to JupyterLab.

⏭️ What's Next

This is Post 1 of the GCP ML Pipelines & MLOps with Terraform series.

Post 1: Vertex AI Workbench (you are here) 🔬
Post 2: Vertex AI Endpoints - Deploy to Prod
Post 3: Vertex AI Feature Store
Post 4: Vertex AI Pipelines + Cloud Build

Your ML workspace is provisioned. Network-isolated, GPU-ready, auto-shutdown enabled, with standardized packages across your team. The foundation for model training, deployment, and production ML pipelines. 🔬

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

SageMaker Studio Domain with Terraform: Your ML Workspace on AWS 🔬

Suhas Mallesh — Wed, 01 Apr 2026 07:00:00 +0000

SageMaker Studio is the IDE for ML on AWS - notebooks, training, deployment, all in one place. Here's how to provision the entire domain with Terraform including VPC, IAM, user profiles, and security hardening.

In Series 1-3, we worked with managed AI services - Bedrock for models, Knowledge Bases for RAG, Agents for orchestration. Series 5 shifts to custom ML - training your own models, deploying them to endpoints, managing features, and building CI/CD pipelines.

It all starts with a SageMaker Studio Domain. The domain is the foundation for everything in SageMaker - it provides the IDE (JupyterLab, Code Editor), manages user profiles, attaches shared storage (EFS), and controls network access. Think of it as the workspace where your ML team lives. This post provisions the entire setup with Terraform. 🎯

🏗️ What a SageMaker Domain Contains

Component	What It Does
Domain	Top-level resource: VPC config, auth mode, encryption
User Profile	Per-user config: execution role, instance types, apps
Shared Space	Team collaboration spaces with shared notebooks
EFS Volume	Persistent storage for notebooks and data files
Apps	JupyterLab, Code Editor, Canvas, Studio Classic

One domain per team or environment. User profiles within the domain give each data scientist their own isolated workspace with shared access to the same data and models.

🔧 Terraform: The Full Domain Setup

VPC and Networking

SageMaker Studio requires a VPC. For production, use VPC-only mode with private subnets and VPC endpoints:

# network/main.tf

resource "aws_vpc" "ml" {
  cidr_block           = var.vpc_cidr
  enable_dns_support   = true
  enable_dns_hostnames = true

  tags = { Name = "${var.environment}-ml-vpc" }
}

resource "aws_subnet" "private" {
  count             = 2
  vpc_id            = aws_vpc.ml.id
  cidr_block        = cidrsubnet(var.vpc_cidr, 8, count.index)
  availability_zone = data.aws_availability_zones.available.names[count.index]

  tags = { Name = "${var.environment}-ml-private-${count.index}" }
}

# S3 Gateway endpoint (required for SageMaker in VPC-only mode)
resource "aws_vpc_endpoint" "s3" {
  vpc_id       = aws_vpc.ml.id
  service_name = "com.amazonaws.${var.region}.s3"

  route_table_ids = [aws_route_table.private.id]
}

# SageMaker API interface endpoint
resource "aws_vpc_endpoint" "sagemaker_api" {
  vpc_id              = aws_vpc.ml.id
  service_name        = "com.amazonaws.${var.region}.sagemaker.api"
  vpc_endpoint_type   = "Interface"
  subnet_ids          = aws_subnet.private[*].id
  security_group_ids  = [aws_security_group.vpc_endpoints.id]
  private_dns_enabled = true
}

# SageMaker Runtime interface endpoint
resource "aws_vpc_endpoint" "sagemaker_runtime" {
  vpc_id              = aws_vpc.ml.id
  service_name        = "com.amazonaws.${var.region}.sagemaker.runtime"
  vpc_endpoint_type   = "Interface"
  subnet_ids          = aws_subnet.private[*].id
  security_group_ids  = [aws_security_group.vpc_endpoints.id]
  private_dns_enabled = true
}

resource "aws_security_group" "vpc_endpoints" {
  name_prefix = "${var.environment}-vpce-"
  vpc_id      = aws_vpc.ml.id

  ingress {
    from_port   = 443
    to_port     = 443
    protocol    = "tcp"
    cidr_blocks = [var.vpc_cidr]
  }
}

VPC-only mode blocks all internet access. You need VPC endpoints for every AWS service SageMaker uses: S3 (gateway), SageMaker API, SageMaker Runtime, and optionally STS, CloudWatch Logs, and ECR for training jobs.

IAM Execution Role

# iam/main.tf

resource "aws_iam_role" "sagemaker_execution" {
  name = "${var.environment}-sagemaker-execution"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "sagemaker.amazonaws.com" }
    }]
  })
}

# Scoped permissions - not AmazonSageMakerFullAccess
resource "aws_iam_role_policy" "sagemaker_core" {
  name = "sagemaker-core"
  role = aws_iam_role.sagemaker_execution.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Effect = "Allow"
        Action = [
          "s3:GetObject",
          "s3:PutObject",
          "s3:ListBucket"
        ]
        Resource = [
          "${var.data_bucket_arn}",
          "${var.data_bucket_arn}/*",
          "${var.artifacts_bucket_arn}",
          "${var.artifacts_bucket_arn}/*"
        ]
      },
      {
        Effect = "Allow"
        Action = [
          "ecr:GetAuthorizationToken",
          "ecr:BatchGetImage",
          "ecr:GetDownloadUrlForLayer"
        ]
        Resource = "*"
      },
      {
        Effect = "Allow"
        Action = [
          "logs:CreateLogGroup",
          "logs:CreateLogStream",
          "logs:PutLogEvents"
        ]
        Resource = "arn:aws:logs:${var.region}:${data.aws_caller_identity.current.account_id}:*"
      }
    ]
  })
}

Avoid AmazonSageMakerFullAccess in production. It grants broad permissions across S3, ECR, Lambda, and more. Scope your policies to the specific buckets, registries, and services your team needs.

KMS Key for Encryption

# security/kms.tf

resource "aws_kms_key" "sagemaker" {
  description             = "SageMaker Studio EFS and EBS encryption"
  deletion_window_in_days = 7
  enable_key_rotation     = true
}

resource "aws_kms_alias" "sagemaker" {
  name          = "alias/${var.environment}-sagemaker"
  target_key_id = aws_kms_key.sagemaker.key_id
}

The Domain

# sagemaker/domain.tf

resource "aws_sagemaker_domain" "this" {
  domain_name = "${var.environment}-ml-studio"
  auth_mode   = "IAM"
  vpc_id      = aws_vpc.ml.id
  subnet_ids  = aws_subnet.private[*].id

  # VPC-only mode blocks internet access
  app_network_access_type = var.network_access_type

  # Encrypt EFS and EBS volumes
  kms_key_id = aws_kms_key.sagemaker.arn

  default_user_settings {
    execution_role = aws_iam_role.sagemaker_execution.arn

    security_groups = [aws_security_group.sagemaker_studio.id]

    jupyter_lab_app_settings {
      default_resource_spec {
        instance_type       = var.default_instance_type
        sagemaker_image_arn = var.jupyter_image_arn
      }
    }

    sharing_settings {
      notebook_output_option = "Allowed"
      s3_output_path         = "s3://${var.artifacts_bucket}/sharing"
    }
  }

  tags = {
    Environment = var.environment
    ManagedBy   = "terraform"
  }
}

User Profiles

# sagemaker/users.tf

resource "aws_sagemaker_user_profile" "this" {
  for_each = var.user_profiles

  domain_id         = aws_sagemaker_domain.this.id
  user_profile_name = each.key

  user_settings {
    execution_role = coalesce(
      each.value.execution_role_arn,
      aws_iam_role.sagemaker_execution.arn
    )
  }

  tags = {
    Team        = each.value.team
    Environment = var.environment
  }
}

📐 Environment Configuration

# environments/dev.tfvars
environment          = "dev"
network_access_type  = "PublicInternetOnly"  # Simpler for dev
default_instance_type = "ml.t3.medium"

user_profiles = {
  "data-scientist-1" = {
    team               = "ml-team"
    execution_role_arn = null  # Uses domain default
  }
}

# environments/prod.tfvars
environment          = "prod"
network_access_type  = "VpcOnly"  # Locked down
default_instance_type = "ml.t3.medium"

user_profiles = {
  "ds-lead" = {
    team               = "ml-team"
    execution_role_arn = null
  }
  "ds-engineer-1" = {
    team               = "ml-team"
    execution_role_arn = null
  }
  "ds-engineer-2" = {
    team               = "ml-team"
    execution_role_arn = null
  }
}

Dev uses PublicInternetOnly for easier setup - no VPC endpoints needed, pip install works out of the box. Prod uses VpcOnly for network isolation - all traffic stays within your VPC, requires VPC endpoints for AWS services.

🔧 Security Hardening Checklist

Control	Dev	Prod
Network access	PublicInternetOnly	VpcOnly
EFS encryption	KMS key	KMS key
IAM scope	Scoped to buckets	Least-privilege per team
Instance types	ml.t3.medium	Restricted via Service Quotas
Notebook sharing	Allowed (S3 output)	Allowed (encrypted S3)
Auto-shutdown	Optional	Lifecycle config enforced

🔧 Auto-Shutdown Lifecycle Config (Cost Control)

Idle notebooks burn money. Attach a lifecycle config to auto-shutdown:

resource "aws_sagemaker_studio_lifecycle_config" "auto_shutdown" {
  studio_lifecycle_config_name     = "${var.environment}-auto-shutdown"
  studio_lifecycle_config_app_type = "JupyterServer"
  studio_lifecycle_config_content  = base64encode(file("${path.module}/scripts/auto-shutdown.sh"))
}

The script checks for idle kernels and shuts down the notebook instance after a configurable timeout (typically 60-120 minutes of inactivity).

⚠️ Gotchas and Tips

One domain per region per account (soft limit). You can request an increase, but the default is one. Use user profiles to separate team members within a single domain.

EFS is created automatically. The domain creates an EFS volume and a home directory per user profile. You don't manage this directly, but you should encrypt it with KMS.

VPC endpoint costs add up. Each interface VPC endpoint costs roughly $7/month plus data processing charges. In VPC-only mode, you may need 5-10 endpoints. Budget $50-100/month for networking in prod.

Domain deletion is slow and complex. Deleting a domain requires deleting all apps, user profiles, and spaces first. Terraform handles this, but terraform destroy can take 15-20 minutes.

Instance type quotas. GPU instances (ml.g5, ml.p4d) require Service Quota increases. Request these early - approval can take days.

⏭️ What's Next

This is Post 1 of the ML Pipelines & MLOps with Terraform series.

Post 1: SageMaker Studio Domain (you are here) 🔬
Post 2: SageMaker Endpoints - Deploy to Prod
Post 3: SageMaker Feature Store
Post 4: SageMaker Pipelines - CI/CD for ML

Your ML workspace is provisioned. VPC-isolated, KMS-encrypted, IAM-scoped, with user profiles for every team member. The foundation for everything that follows: model training, deployment, feature engineering, and ML pipelines. 🔬

Found this helpful? Follow for the full ML Pipelines & MLOps with Terraform series! 💬

Azure AI Agent + Bing Grounding: Real-Time Web Search with Terraform 📚

Suhas Mallesh — Sat, 28 Mar 2026 07:00:00 +0000

An agent without grounding hallucinates about current events. Bing Grounding gives your Azure AI agent real-time web search with source citations. Here's how to provision the Bing resource, connect it to your Foundry project, and attach it to your agent with Terraform.

Through this series, we've built Azure AI agents that can reason, call functions, and coordinate multi-agent workflows. But ask one about today's news, current stock prices, or this week's weather, and it either hallucinates or admits its training data has a cutoff.

Bing Grounding connects your agent to real-time web search. When the agent determines it needs current information, it searches Bing, retrieves relevant results, and generates a response grounded in those results with URL citations. The agent decides when to search. Bing provides the facts. Your users get accurate, sourced answers. 🎯

🏗️ How Bing Grounding Works

User: "What's the weather in Seattle today?"
    ↓
Agent analyzes request → determines it needs current data
    ↓
Agent invokes Bing Grounding tool
    ↓
Bing searches the web, returns relevant results
    ↓
Agent generates response grounded in Bing results
    ↓
User: "It's 55°F and partly cloudy in Seattle today."
       [Citation: weather.gov/seattle]

The agent doesn't call the Bing API directly. The Foundry Agent Service handles the Bing integration through a connection resource. You provision the Bing resource, create the connection, and attach the tool to your agent.

🔧 Two Bing Grounding Options

Tool	Scope	Best For
Grounding with Bing Search	Full public web	General current events, weather, news, live data
Grounding with Bing Custom Search	Specific domains you define	Restricted searches (e.g., only your docs site, partner sites)

Bing Search searches the full web. Bing Custom Search lets you restrict results to specific domains, useful when you want grounding but only from trusted sources.

🔧 Terraform: Provision Bing Grounding

Step 1: Create the Bing Grounding Resource

# grounding/bing.tf

resource "azapi_resource" "bing_grounding" {
  type      = "Microsoft.Bing/accounts@2025-05-01-preview"
  name      = "${var.environment}-bing-grounding"
  location  = "global"
  parent_id = azurerm_resource_group.this.id

  body = {
    kind = "Grounding"
    sku = {
      name = "G1"
    }
  }
}

Why azapi? The Bing Grounding resource isn't available in the azurerm provider yet. The azapi provider gives access to the latest Azure resource types directly.

Step 2: Connect Bing to Your Foundry Project

# grounding/connection.tf

resource "azapi_resource" "bing_connection" {
  type      = "Microsoft.CognitiveServices/accounts/connections@2025-06-01"
  name      = "${var.environment}-bing-connection"
  parent_id = azurerm_cognitive_account.ai_foundry.id

  schema_validation_enabled = false

  body = {
    properties = {
      category  = "ApiKey"
      target    = azapi_resource.bing_grounding.output.properties.endpoint
      authType  = "ApiKey"
      isSharedToAll = true
      credentials = {
        key = azapi_resource.bing_grounding.output.properties.accessKeys.key1
      }
      metadata = {
        ApiType    = "Bing"
        ResourceId = azapi_resource.bing_grounding.id
      }
    }
  }
}

This creates a connection between your Foundry project and the Bing resource. The connection ID is what your agent code references when attaching the Bing tool.

Step 3: Output the Connection ID

# grounding/outputs.tf

resource "local_file" "agent_config" {
  filename = "${path.module}/agent_code/agent_config.json"
  content = jsonencode({
    foundry_endpoint  = azurerm_cognitive_account.ai_foundry.endpoint
    deployment_name   = azurerm_cognitive_deployment.agent_model.name
    agent_name        = "${var.environment}-grounded-agent"
    instruction       = var.agent_instruction
    bing_connection_id = azapi_resource.bing_connection.id
  })
}

🐍 Create Agent with Bing Grounding

# agent_code/create_agent.py

from azure.ai.agents import AgentsClient
from azure.ai.agents.models import BingGroundingTool
from azure.identity import DefaultAzureCredential
import json

with open("agent_config.json") as f:
    config = json.load(f)

client = AgentsClient(
    endpoint=config["foundry_endpoint"],
    credential=DefaultAzureCredential(),
)

# Create Bing grounding tool from connection
bing_tool = BingGroundingTool(
    connection_id=config["bing_connection_id"]
)

# Create agent with Bing grounding
agent = client.create_agent(
    model=config["deployment_name"],
    name=config["agent_name"],
    instructions="""You are a helpful assistant with access to real-time 
    web search. When answering questions about current events, weather, 
    news, stock prices, or any time-sensitive information, use Bing 
    Search to ground your response. Always include source citations.""",
    tools=bing_tool.definitions,
)

print(f"Agent created: {agent.id}")

🧪 Query the Grounded Agent

# Create thread and send message
thread = client.threads.create()
client.messages.create(
    thread_id=thread.id,
    role="user",
    content="What are the top tech headlines today?"
)

# Run the agent
run = client.runs.create_and_process(
    thread_id=thread.id,
    agent_id=agent.id,
)

# Get response with citations
if run.status == "completed":
    messages = client.messages.list(thread_id=thread.id)
    for msg in messages:
        if msg.role == "assistant":
            # Print the response text
            print(msg.content[0].text.value)

            # Print URL citations
            if hasattr(msg.content[0].text, "annotations"):
                for annotation in msg.content[0].text.annotations:
                    if hasattr(annotation, "url_citation"):
                        print(f"  Source: {annotation.url_citation.url}")

The response includes inline URL citations. When Bing contributes to the answer, annotations link back to the source websites. These citations must be displayed to end users per Microsoft's Use and Display Requirements.

📐 Combining Bing Grounding with Other Tools

An agent can have Bing Grounding alongside function calling tools:

from azure.ai.agents.models import BingGroundingTool, FunctionTool

bing_tool = BingGroundingTool(connection_id=config["bing_connection_id"])

exchange_rate_tool = FunctionTool(
    name="get_exchange_rate",
    description="Get live exchange rates between currencies.",
    parameters={
        "type": "object",
        "properties": {
            "currency_from": {"type": "string", "description": "Source currency code"},
            "currency_to": {"type": "string", "description": "Target currency code"},
        },
        "required": ["currency_from", "currency_to"]
    }
)

# Agent with both grounding and function tools
agent = client.create_agent(
    model=config["deployment_name"],
    name="full-agent",
    instructions="""You are a helpful assistant.
    Use Bing Search for current events, news, and general web queries.
    Use get_exchange_rate for live currency conversion rates.
    Choose the right tool based on the question.""",
    tools=bing_tool.definitions + [exchange_rate_tool],
)

The agent routes between Bing for general web queries and your custom functions for specific operations. No conflicts between tool types.

📐 Adding Azure AI Search (Private Docs)

For grounding against both public web and private documents, combine Bing with Azure AI Search:

from azure.ai.agents.models import (
    BingGroundingTool,
    AzureAISearchTool,
)

bing_tool = BingGroundingTool(connection_id=config["bing_connection_id"])

search_tool = AzureAISearchTool(
    index_connection_id=config["search_connection_id"],
    index_name=config["search_index_name"],
)

agent = client.create_agent(
    model=config["deployment_name"],
    name="fully-grounded-agent",
    instructions="""You are a company assistant.
    Use Bing for current events and public information.
    Use Azure AI Search for company policies and internal documents.
    Answer general knowledge questions directly.""",
    tools=bing_tool.definitions + search_tool.definitions,
)

The agent now has three paths: Bing for public web, AI Search for private docs, and direct answers for general knowledge.

⚠️ Gotchas and Tips

Citation display is required. When Bing Grounding contributes to a response, you must display the URL citations to end users. This is a requirement from Microsoft's Use and Display Requirements, not optional.

Data crosses compliance boundary. Bing queries are sent outside the Azure compliance boundary to the Bing service. Assess whether this meets your compliance requirements, especially in regulated industries.

Bing Grounding has per-transaction pricing. Each time the agent invokes the Bing tool, it counts as a grounding transaction. Check the Bing Grounding pricing page for current rates.

Agent decides when to search. The agent uses Bing only when it determines web search is needed. Clear instructions about when to use Bing help prevent unnecessary searches on questions the model can answer from general knowledge.

Connection-based architecture. The Bing resource connects to your Foundry project via a connection resource. If you move to a different project, recreate the connection. The Bing resource itself is reusable across projects.

⏭️ Series Complete!

This is Post 4 of the Azure AI Agents with Terraform series, and the final post of the entire multi-cloud blog series.

Post 1: Deploy First Azure AI Agent 🤖
Post 2: Function Calling - Connect to APIs 🔌
Post 3: Multi-Agent Orchestration 🧠
Post 4: Agent + Bing Grounding (you are here) 📚

Your agent is fully grounded. Bing for real-time web data. AI Search for private documents. Function calling for live APIs. From a simple model endpoint to a production-ready, multi-agent, grounded AI system - all defined in Terraform and Python. 📚

Vertex AI Agent + Google Search Grounding: Real-Time Facts with ADK and Terraform 📚

Suhas Mallesh — Fri, 27 Mar 2026 07:00:00 +0000

An agent without grounding hallucinates about current events. Google Search Grounding and Vertex AI Search give your ADK agent real-time web data and private document access with source citations. Here's how to wire both up with Terraform.

Through this series, we've built ADK agents that can reason, call tools, and coordinate with other agents. But ask one about today's stock prices or your company's internal policies, and it either hallucinates or admits it doesn't know. The model's training data has a cutoff, and it has no access to your documents.

Grounding fixes this with two built-in ADK tools: Google Search Grounding connects your agent to real-time web data - current events, live prices, recent news. Vertex AI Search Grounding connects your agent to your private documents - company policies, product specs, internal knowledge bases. Both return source-attributed responses with citations. 🎯

🏗️ Two Grounding Approaches

Tool	Data Source	Best For
Google Search	Public web (Google Search index)	Current events, live data, general facts
Vertex AI Search	Your private documents (indexed datastore)	Company policies, internal docs, product info

You can use one or both. Google Search handles what's public and current. Vertex AI Search handles what's private and specific to your organization.

🔧 Google Search Grounding

The simplest grounding option. Add the built-in google_search tool to your agent:

# agent_source/agent.py

from google.adk.agents import Agent
from google.adk.models.vertexai import VertexAi

model = VertexAi(model=config["model_id"])

search_agent = Agent(
    name="search-grounded-agent",
    model=model,
    instruction="""You are a helpful assistant that provides accurate,
    current information. When answering questions about recent events, 
    current data, or anything that requires up-to-date information,
    use Google Search to ground your response. Always cite your sources.""",
    tools=["google_search"],
)

That's it. One line: tools=["google_search"]. The agent automatically searches the web when it determines the question needs current information.

What the Response Looks Like

When the agent uses Google Search, the response includes groundingMetadata with source URLs, search queries executed, and confidence scores per segment:

async for event in runner.run_async(
    user_id="user-123",
    session_id=session.id,
    new_message=types.Content(
        role="user",
        parts=[types.Part(text="What were today's major tech news?")]
    )
):
    if event.content and event.content.parts:
        for part in event.content.parts:
            if hasattr(part, "text") and part.text:
                print(part.text)

    # Access grounding metadata for citations
    if hasattr(event, "grounding_metadata"):
        for chunk in event.grounding_metadata.grounding_chunks:
            print(f"Source: {chunk.web.uri} - {chunk.web.title}")

Important: When using Google Search Grounding in production, you must display the Search Suggestions provided in the response metadata. This is a requirement from Google's terms of service.

🔧 Vertex AI Search Grounding (Private Documents)

For grounding against your own documents, use the VertexAiSearchTool:

from google.adk.agents import Agent
from google.adk.tools import VertexAiSearchTool

# Your Vertex AI Search datastore path
DATASTORE_PATH = (
    f"projects/{config['project_id']}/locations/global"
    f"/collections/default_collection"
    f"/dataStores/{config['datastore_id']}"
)

doc_agent = Agent(
    name="doc-grounded-agent",
    model=model,
    instruction="""You are a company policy expert. Answer questions 
    about internal policies, procedures, and guidelines using the 
    provided document store. Always cite the source document.""",
    tools=[VertexAiSearchTool(data_store_id=DATASTORE_PATH)],
)

The agent queries your indexed documents, retrieves relevant chunks with semantic search, and generates a response grounded in your actual content. Source attributions include document titles and URIs.

Tool Limitation

Both google_search and VertexAiSearchTool are exclusive tools - they cannot be combined with other tools in the same agent instance. To use grounding alongside custom function tools, wrap the grounded agent as a sub-agent or use AgentTool:

from google.adk.tools import AgentTool

# Grounded agent as a tool for the main agent
search_assistant = Agent(
    name="searcher",
    model=model,
    instruction="Search the web for current information when asked.",
    tools=["google_search"],
)

doc_assistant = Agent(
    name="doc_lookup",
    model=model,
    instruction="Look up company policies and internal documents.",
    tools=[VertexAiSearchTool(data_store_id=DATASTORE_PATH)],
)

# Main agent uses grounded agents as tools alongside custom tools
main_agent = Agent(
    name="assistant",
    model=model,
    instruction="""You are a helpful assistant.
    Use the searcher for current events and live data.
    Use doc_lookup for company policies and internal info.
    Use get_exchange_rate for currency conversions.""",
    tools=[
        AgentTool(search_assistant),
        AgentTool(doc_assistant),
        get_exchange_rate,
    ],
)

This is the production pattern: dedicated grounding agents wrapped as tools for a main orchestrator agent.

🔧 Terraform: Infrastructure for Grounding

APIs and Vertex AI Search Datastore

# grounding/apis.tf

resource "google_project_service" "grounding" {
  for_each = toset([
    "aiplatform.googleapis.com",
    "discoveryengine.googleapis.com",  # For Vertex AI Search
    "storage.googleapis.com",
  ])
  project = var.project_id
  service = each.value
}

Vertex AI Search Datastore (for Private Docs)

# grounding/datastore.tf

resource "google_discovery_engine_data_store" "docs" {
  location                    = "global"
  data_store_id               = "${var.environment}-company-docs"
  display_name                = "Company Documents"
  industry_vertical           = "GENERIC"
  content_config              = "CONTENT_REQUIRED"
  solution_types              = ["SOLUTION_TYPE_SEARCH"]
  create_advanced_site_search = false
  project                     = var.project_id
}

IAM for Grounding

# grounding/iam.tf

# Agent SA needs Discovery Engine access for Vertex AI Search
resource "google_project_iam_member" "discovery_viewer" {
  project = var.project_id
  role    = "roles/discoveryengine.viewer"
  member  = "serviceAccount:${google_service_account.agent.email}"
}

Config with Datastore ID

# grounding/config.tf

resource "local_file" "agent_config" {
  filename = "${path.module}/agent_source/config.json"
  content = jsonencode({
    project_id   = var.project_id
    location     = var.region
    model_id     = var.agent_model.id
    agent_name   = "${var.environment}-grounded-agent"
    datastore_id = google_discovery_engine_data_store.docs.data_store_id
    instruction  = var.agent_instruction
  })
}

📐 Combining Google Search + Vertex AI Search

The production pattern uses both grounding sources through the sub-agent approach:

# Public data: current events, live information
web_agent = Agent(
    name="web_search",
    model=model,
    instruction="Search the public web for current, real-time information.",
    tools=["google_search"],
)

# Private data: company documents, internal policies
docs_agent = Agent(
    name="doc_search",
    model=model,
    instruction="Search company documents for internal policies and procedures.",
    tools=[VertexAiSearchTool(data_store_id=DATASTORE_PATH)],
)

# Orchestrator decides which source to use
grounded_assistant = Agent(
    name="grounded_assistant",
    model=model,
    instruction="""You are a knowledgeable assistant.
    For current events or public information, use web_search.
    For company policies or internal documents, use doc_search.
    For general knowledge questions, answer directly.""",
    tools=[
        AgentTool(web_agent),
        AgentTool(docs_agent),
    ],
)

The orchestrator routes to the right grounding source based on the question type. Public questions go to Google Search. Internal questions go to Vertex AI Search. General knowledge questions get answered directly without grounding.

⚠️ Gotchas and Tips

Exclusive tool limitation. google_search and VertexAiSearchTool can't be mixed with other tools in the same agent. Use the AgentTool wrapper pattern shown above.

Search Suggestions display requirement. When using Google Search Grounding in production, you must display the searchEntryPoint HTML/CSS from the response metadata. This is not optional.

Datastore indexing time. Vertex AI Search datastores need time to index your documents after upload. Don't test grounding immediately after creating the datastore - wait for indexing to complete.

Grounding adds latency. Each grounded response involves a search query plus retrieval. Expect 1-3 seconds of additional latency compared to non-grounded responses.

Confidence scores. The groundingMetadata includes confidence scores per response segment. Use these to flag low-confidence answers for human review in production.

⏭️ Series Complete!

This is Post 4 of the GCP AI Agents with Terraform series, and the final GCP post of the entire multi-cloud series.

Post 1: Deploy First Vertex AI Agent 🤖
Post 2: Agent Tools - Connect to APIs 🔌
Post 3: Multi-Agent Systems 🧠
Post 4: Agent + Grounding (you are here) 📚

Your agent is now grounded in reality. Google Search for real-time web data. Vertex AI Search for your private documents. Source citations on every response. No hallucinations, just facts. 📚

Thanks for following the full series! Follow for what comes next. 💬

Bedrock Agent + Knowledge Base: Ground Your Agent in Real Data with Terraform 📚

Suhas Mallesh — Thu, 26 Mar 2026 07:00:00 +0000

An agent without grounding makes things up. Associate a Knowledge Base with your Bedrock Agent and every response is backed by your actual documents - with citations. Here's how to wire it up with Terraform.

Through this series, we've built Bedrock Agents that can reason, call APIs via action groups, and coordinate with other agents. But ask one a question about your company's policies, and it either hallucinates or admits it doesn't know. The model has no access to your data.

Knowledge Base association changes that. When you attach a Knowledge Base to your agent, the agent automatically queries your vector store when it needs domain-specific information. It retrieves relevant chunks, uses them as context, and generates a grounded response with citations back to your source documents. No hallucination. No guessing. Real answers from real data. 🎯

🏗️ How Agent + Knowledge Base Works

User: "What is our return policy for electronics?"
    ↓
Agent analyzes request → determines it needs domain knowledge
    ↓
Agent queries Knowledge Base (vector search)
    ↓
Knowledge Base returns relevant chunks from policies.pdf
    ↓
Agent generates response grounded in retrieved context
    ↓
User: "Electronics can be returned within 30 days with original 
       receipt. Items must be in original packaging. [Source: 
       return-policy-2024.pdf, page 3]"

The agent decides when to query the Knowledge Base based on the KB instruction you provide. If the question can be answered without domain data (general knowledge, action group results), the agent skips the KB query.

🔧 What You Need

This post assumes you have:

A Bedrock Agent deployed (from Day 9)
A Knowledge Base with synced data (from Series 2, RAG Posts)

We're connecting the two. If you haven't built either, refer back to the earlier posts for the full setup.

🔧 Terraform: Associate Knowledge Base with Agent

The association is a single Terraform resource:

# agent/kb_association.tf

resource "aws_bedrockagent_agent_knowledge_base_association" "this" {
  agent_id             = aws_bedrockagent_agent.this.agent_id
  description          = var.kb_instruction
  knowledge_base_id    = var.knowledge_base_id
  knowledge_base_state = "ENABLED"
}

That's it. One resource. The description field is the most important part - it tells the agent when and how to use the Knowledge Base.

The KB Instruction

# variables.tf

variable "kb_instruction" {
  description = "Instruction telling the agent when to use the Knowledge Base"
  type        = string
  default     = <<-EOT
    Use this knowledge base to answer questions about company policies, 
    product information, and internal procedures. Query the knowledge 
    base when the user asks about policies, returns, warranties, 
    pricing, or any company-specific information. Do not answer these 
    questions from general knowledge.
  EOT
}

This instruction drives retrieval quality. A vague instruction like "use this for questions" means the agent queries the KB for everything, wasting tokens on general knowledge questions. A specific instruction like the one above targets retrieval to domain-specific queries only.

Re-Prepare After Association

The agent must be re-prepared after associating a Knowledge Base:

resource "null_resource" "prepare_after_kb" {
  triggers = {
    kb_association = aws_bedrockagent_agent_knowledge_base_association.this.id
    kb_instruction = var.kb_instruction
  }

  provisioner "local-exec" {
    command = "aws bedrock-agent prepare-agent --agent-id ${aws_bedrockagent_agent.this.agent_id} --region ${var.region}"
  }

  depends_on = [aws_bedrockagent_agent_knowledge_base_association.this]
}

IAM: Agent Needs KB Access

The agent's service role needs permission to retrieve from the Knowledge Base:

# agent/iam.tf

resource "aws_iam_role_policy" "agent_kb_retrieve" {
  name = "kb-retrieve"
  role = aws_iam_role.agent.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = [
        "bedrock:Retrieve",
        "bedrock:RetrieveAndGenerate"
      ]
      Resource = "arn:aws:bedrock:${var.region}:${data.aws_caller_identity.current.account_id}:knowledge-base/${var.knowledge_base_id}"
    }]
  })
}

📐 Full Agent with Action Groups + Knowledge Base

Here's the complete picture - an agent with both tools and grounding:

# agent/main.tf

resource "aws_bedrockagent_agent" "this" {
  agent_name              = "${var.environment}-${var.agent_name}"
  agent_resource_role_arn = aws_iam_role.agent.arn
  foundation_model        = var.agent_model.id
  instruction             = var.agent_instruction
  prepare_agent           = false  # Prepare manually after KB + action groups
  idle_session_ttl_in_seconds = var.idle_session_ttl

  tags = {
    Environment = var.environment
    Model       = var.agent_model.display
  }
}

# Action group for real-time API calls
resource "aws_bedrockagent_agent_action_group" "api_tools" {
  action_group_name = "APITools"
  agent_id          = aws_bedrockagent_agent.this.agent_id
  agent_version     = "DRAFT"
  description       = "Real-time API operations"

  action_group_executor {
    lambda = aws_lambda_function.api_tools.arn
  }

  api_schema {
    payload = file("${path.module}/schemas/api_tools.json")
  }
}

# Knowledge Base for document grounding
resource "aws_bedrockagent_agent_knowledge_base_association" "docs" {
  agent_id             = aws_bedrockagent_agent.this.agent_id
  description          = var.kb_instruction
  knowledge_base_id    = var.knowledge_base_id
  knowledge_base_state = "ENABLED"
}

# Prepare after everything is wired up
resource "null_resource" "prepare_agent" {
  triggers = {
    action_group = aws_bedrockagent_agent_action_group.api_tools.id
    kb           = aws_bedrockagent_agent_knowledge_base_association.docs.id
  }

  provisioner "local-exec" {
    command = "aws bedrock-agent prepare-agent --agent-id ${aws_bedrockagent_agent.this.agent_id} --region ${var.region}"
  }

  depends_on = [
    aws_bedrockagent_agent_action_group.api_tools,
    aws_bedrockagent_agent_knowledge_base_association.docs,
  ]
}

resource "time_sleep" "wait_for_preparation" {
  depends_on      = [null_resource.prepare_agent]
  create_duration = "15s"
}

resource "aws_bedrockagent_agent_alias" "live" {
  agent_alias_name = "${var.environment}-live"
  agent_id         = aws_bedrockagent_agent.this.agent_id
  depends_on       = [time_sleep.wait_for_preparation]
}

The agent now has two capabilities: action groups for real-time API calls (live exchange rates, order lookups) and Knowledge Base for document-grounded answers (policies, product info). The agent decides which to use based on the user's question.

🧪 How the Agent Decides

When a user asks a question, the agent follows this decision flow:

Can I answer from general knowledge? If yes, respond directly.
Does this need a tool call? If the question requires live data (exchange rates, weather), invoke the action group.
Does this need domain knowledge? If the question is about company-specific information, query the Knowledge Base.
Do I need both? For complex questions, the agent can combine action group results with KB context.

The KB instruction and action group descriptions guide these decisions. Clear, non-overlapping descriptions produce the best routing.

📐 Multiple Knowledge Bases

An agent can have multiple Knowledge Bases, each for a different domain:

resource "aws_bedrockagent_agent_knowledge_base_association" "policies" {
  agent_id          = aws_bedrockagent_agent.this.agent_id
  description       = "Use for questions about company policies, returns, and warranties."
  knowledge_base_id = var.policies_kb_id
  knowledge_base_state = "ENABLED"
}

resource "aws_bedrockagent_agent_knowledge_base_association" "products" {
  agent_id          = aws_bedrockagent_agent.this.agent_id
  description       = "Use for questions about product specifications, features, and compatibility."
  knowledge_base_id = var.products_kb_id
  knowledge_base_state = "ENABLED"
}

Each KB instruction should clearly define its domain so the agent routes to the right KB. Overlapping instructions lead to incorrect retrieval.

⚠️ Gotchas and Tips

KB instruction quality is retrieval quality. A vague instruction means the agent queries the KB for everything, adding latency and cost to questions it could answer from general knowledge. Be specific about what types of questions should trigger a KB query.

Citations come free. When the agent uses KB context, the response includes source attributions with document names and page numbers. You don't need to build citation logic.

KB + action groups don't conflict. The agent can use both in a single conversation turn. For example: "Check my order status and tell me the return policy" might trigger an action group for order status and a KB query for the return policy.

Sync your KB before testing. The association works immediately, but the agent retrieves from whatever is currently indexed. If you haven't synced your data source recently, you'll get stale or empty results.

KB instruction updates require re-preparation. Changing the description field on the association triggers a Terraform update, but the agent must be re-prepared for the change to take effect.

⏭️ Series Complete!

This is Post 4 of the AWS AI Agents with Terraform series, and the final post of the entire multi-cloud series.

Post 1: Deploy First Bedrock Agent 🤖
Post 2: Action Groups - Connect to APIs 🔌
Post 3: Multi-Agent Orchestration 🧠
Post 4: Agent + Knowledge Base Grounding (you are here) 📚

Your agent is now grounded. It reasons through complex requests, calls APIs for live data, coordinates with specialist agents, and answers domain questions from your actual documents with citations. From a simple endpoint to a fully autonomous, grounded AI agent - all in Terraform. 📚

Thanks for following the full series! Follow for what comes next. 💬