GCP Vertex AI Endpoint Traffic Split Discrepancy Guard

id: gcp-vertex-ai-endpoint-traffic-split-discrepancy-guard
namespace: company.ai

description: |
  Audits Google Cloud Vertex AI model serving endpoints for skewed traffic splits, unrouted deployed models,
  and failed canary rollouts to safeguard production machine learning inference pipelines.

concurrency:
  limit: 1

inputs:
  - id: gcp_project_id
    type: STRING
    displayName: GCP Project ID
    defaults: "company-prod-mlops"
    description: Google Cloud project identifier containing the Vertex AI serving endpoints.

  - id: gcp_region
    type: STRING
    displayName: GCP Region
    defaults: "us-central1"
    description: Google Cloud location where Vertex AI endpoints are deployed.

  - id: endpoint_id
    type: STRING
    displayName: Vertex AI Endpoint ID
    defaults: "1234567890123456789"
    description: Numeric resource ID of the target Vertex AI model serving endpoint.

  - id: expected_primary_split_min_percent
    type: INT
    displayName: Min Primary Split (%)
    defaults: 80
    description: Minimum acceptable traffic split percentage allocated to the
      primary baseline model.

  - id: alert_on_discrepancy
    type: BOOL
    displayName: Alert on Discrepancy
    defaults: true
    description: Whether to dispatch automated Slack alerts when traffic split
      percentages deviate from safe deployment baselines.

tasks:
  - id: scan_vertex_traffic_split
    type: io.kestra.plugin.scripts.python.Script
    description: Queries GCP Vertex AI API to inspect endpoint model deployment
      configuration and traffic split ratios.
    containerImage: python:3.11-slim
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
    beforeCommands:
      - pip install google-cloud-aiplatform kestra
    env:
      GCP_SERVICE_ACCOUNT_KEY: "{{ secret('GCP_SERVICE_ACCOUNT_KEY') }}"
    script: |
      import os
      import json
      from google.cloud import aiplatform
      from google.oauth2 import service_account
      from kestra import Kestra

      project_id = "{{ inputs.gcp_project_id }}".strip()
      region = "{{ inputs.gcp_region }}".strip()
      endpoint_id = "{{ inputs.endpoint_id }}".strip()
      min_split = int("{{ inputs.expected_primary_split_min_percent }}")
      sa_json = os.environ.get("GCP_SERVICE_ACCOUNT_KEY", "")

      endpoint_name = endpoint_id
      traffic_split = {}
      deployed_models = []
      has_discrepancy = False

      try:
          if sa_json:
              creds = service_account.Credentials.from_service_account_info(json.loads(sa_json))
              aiplatform.init(project=project_id, location=region, credentials=creds)
          else:
              aiplatform.init(project=project_id, location=region)

          endpoint = aiplatform.Endpoint(endpoint_name=endpoint_id)
          endpoint_name = endpoint.display_name
          traffic_split = endpoint.traffic_split or {}

          for m in endpoint.deployed_models:
              deployed_models.append({
                  "id": m.id,
                  "model": m.model,
                  "display_name": m.display_name
              })

          # Evaluate traffic split consistency
          total_traffic = sum(traffic_split.values())
          if total_traffic != 100 and len(deployed_models) > 0:
              has_discrepancy = True

          # Check if any deployed model has 0% traffic (orphaned candidate model)
          for m in deployed_models:
              if m["id"] not in traffic_split or traffic_split[m["id"]] == 0:
                  has_discrepancy = True

          report = {
              "endpoint_id": endpoint_id,
              "endpoint_name": endpoint_name,
              "deployed_models_count": len(deployed_models),
              "traffic_split": traffic_split,
              "total_traffic": total_traffic,
              "has_discrepancy": has_discrepancy,
              "remediation_advice": f"Review deployed model splits on Vertex AI Endpoint {endpoint_id}. Verify canary model promotion."
          }

          with open("vertex_report.json", "w") as f:
              json.dump(report, f)

          Kestra.outputs({
              "endpoint_id": endpoint_id,
              "endpoint_name": endpoint_name,
              "has_discrepancy": has_discrepancy,
              "deployed_models_count": len(deployed_models)
          })

      except Exception as e:
          print(f"Error querying Vertex AI Endpoint: {e}")
          report = {
              "endpoint_id": endpoint_id,
              "error": str(e),
              "has_discrepancy": False
          }
          with open("vertex_report.json", "w") as f:
              json.dump(report, f)
          Kestra.outputs({
              "endpoint_id": endpoint_id,
              "endpoint_name": endpoint_name,
              "has_discrepancy": False,
              "deployed_models_count": 0
          })

  - id: check_traffic_split_anomaly
    type: io.kestra.plugin.core.flow.If
    description: Evaluates whether traffic split ratios or orphaned deployed models
      indicate a deployment discrepancy.
    condition: "{{ fromJson(read('vertex_report.json')).has_discrepancy and
      inputs.alert_on_discrepancy }}"
    then:
      - id: notify_slack_traffic_split_anomaly
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Dispatches high-priority Vertex AI model routing discrepancy alert
          to MLOps on Slack.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "channel": "#mlops-serving-alerts",
            "username": "VertexAI-Sentinel",
            "text": "WARNING: Google Cloud Vertex AI Endpoint `{{ fromJson(read('vertex_report.json')).endpoint_name }}` (`{{ fromJson(read('vertex_report.json')).endpoint_id }}`) has traffic split discrepancies! Deployed models: {{ fromJson(read('vertex_report.json')).deployed_models_count }}. Traffic split: {{ fromJson(read('vertex_report.json')).traffic_split }}. Runbook: {{ fromJson(read('vertex_report.json')).remediation_advice }}"
          }

  - id: log_serving_health_state
    type: io.kestra.plugin.core.log.Log
    description: Records an operational audit trail of Vertex AI endpoint traffic
      split evaluation.
    message: "Vertex AI audit completed for '{{
      fromJson(read('vertex_report.json')).endpoint_name }}'. Deployed models:
      {{ fromJson(read('vertex_report.json')).deployed_models_count }},
      Discrepancy: {{ fromJson(read('vertex_report.json')).has_discrepancy }}."

  - id: export_vertex_manifest
    type: io.kestra.plugin.core.debug.Return
    description: Emits structured manifest of Vertex AI model serving configurations
      for MLOps audit records.
    format: "{{ fromJson(read('vertex_report.json')) }}"

triggers:
  - id: hourly_vertex_traffic_audit
    type: io.kestra.plugin.core.trigger.Schedule
    description: Automated schedule trigger auditing Vertex AI model serving
      endpoints hourly.
    cron: "0 * * * *"