Chatting With Dataproc Clusters Using an AI Agent

Managing fleet-scale Apache Spark and Hadoop environments on Google Cloud Dataproc routinely requires juggling gcloud dataproc clusters CLI commands, Cloud Console tabs, and YARN/Spark configuration flags. What if platform engineers and data scientists could inspect, provision, resize, and tear down clusters simply by conversing with an AI infrastructure agent?

Using Google’s open-source Agent Development Kit (ADK) and the googleCloudPlatform/agent-starter-pack, I built a production-ready Dataproc AI Agent powered by Gemini on Vertex AI—and open-sourced the full implementation in github.com/datasherlock/ai-agents.

Architecture: Google ADK + Vertex AI + Async Dataproc SDK

Rather than brittle prompt-to-bash wrappers, the agent uses structured Python function tools backed by the official google-cloud-dataproc asynchronous gRPC client (ClusterControllerAsyncClient). Each tool exposes rich Python type hints and docstrings so Gemini accurately extracts project IDs, regions, machine types, and worker counts—and asks clarifying questions before executing mutating operations.

• Reasoning Core: Gemini on Vertex AI orchestrated via google.adk.agents.Agent.
• Tool Suite: Asynchronous cluster lifecycle tools—list_dataproc_clusters_async, get_dataproc_cluster_async, create_dataproc_cluster_async, update_dataproc_cluster_async, and delete_dataproc_cluster_async.
• Developer & Production Runtime: Local interactive testing via a Streamlit Playground (make playground) and one-command deployment to Vertex AI Agent Engine (make backend) with Terraform-managed IAM and Cloud Trace telemetry.

Defining the Root ADK Agent (agent.py)

With Google ADK, wiring the agent definition takes under 25 lines of clean Python:

import os
import google.auth
from google.adk.agents import Agent
from .utils.cluster_controller_tools import (
    list_dataproc_clusters,
    create_dataproc_cluster_async,
    get_dataproc_cluster_async,
    list_dataproc_clusters_async,
    update_dataproc_cluster_async,
    delete_dataproc_cluster_async,
)
from .utils.llm_config import root_model, root_prompt

_, project_id = google.auth.default()
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", project_id)
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "us-central1")
os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "True")

root_agent = Agent(
    name="root_agent",
    model=root_model,
    instruction=root_prompt,
    tools=[
        list_dataproc_clusters,
        create_dataproc_cluster_async,
        get_dataproc_cluster_async,
        list_dataproc_clusters_async,
        update_dataproc_cluster_async,
        delete_dataproc_cluster_async,
    ],
)

Writing Non-Blocking Long-Running Operation (LRO) Tools

Provisioning or resizing a Dataproc cluster takes 60–90 seconds. Blocking the agent’s event loop during cluster creation degrades UX, so our tool invokes ClusterControllerAsyncClient and returns the Long-Running Operation (LRO) handle immediately so the user can check progress conversationally:

from google.cloud import dataproc_v1 as dataproc
from google.api_core.exceptions import GoogleAPICallError

async def create_dataproc_cluster_async(
    project_id: str, region: str, cluster_name: str, cluster_config: dict
) -> dict:
    """Asynchronously creates a Dataproc cluster in the target GCP project and region."""
    try:
        cluster_client = dataproc.ClusterControllerAsyncClient(
            client_options={"api_endpoint": f"{region}-dataproc.googleapis.com:443"}
        )
        cluster = {
            "project_id": project_id,
            "cluster_name": cluster_name,
            "config": cluster_config,
        }
        operation = await cluster_client.create_cluster(
            request={"project_id": project_id, "region": region, "cluster": cluster}
        )
        return {
            "status": "creating",
            "operation_name": operation.operation.name,
            "metadata": str(operation.metadata),
        }
    except GoogleAPICallError as e:
        return {"error": f"API Error creating cluster async: {e.message}"}

Try It Yourself Locally or on Vertex AI Agent Engine

Clone the repository, authenticate with Application Default Credentials (gcloud auth application-default login), and spin up the local Streamlit playground using uv and make:

git clone https://github.com/datasherlock/ai-agents.git
cd ai-agents/dataproc-agent
make install && make playground

# Ready for production? Deploy straight to Vertex AI Agent Engine:
make backend

Once running, you can ask natural-language prompts like “List all Dataproc clusters in us-central1 and flag any sitting idle”, “Scale secondary workers on analytics-prod-cluster to 10”, or pair this agent with Dataproc tuning concepts from Understanding Driver Pools in Dataproc and Autoscaling in Dataproc.

Originally published on Medium / Dev Genius — Chatting With Dataproc Clusters Using an AI Agent. Source code: github.com/datasherlock/ai-agents.