ML Workload Engine

An autonomous ADK agent delivering deterministic control, priority-based preemption, and real-time state tracking for ML compute workloads.

Comment

Inspiration

Modern AI teams face an intense infrastructure bottleneck: GPU contention, SLA deadline failures, and inefficient cluster utilization.

  • High-priority large language model (LLM) fine-tuning runs frequently arrive with strict SLA deadlines while GPU nodes are fully occupied by long-running batch embedding or data processing jobs.
  • Traditional schedulers rely on rigid, static rules or require manual human intervention (e.g. SREs waking up at 3 AM) to triage failed jobs or resolve hardware contention.
  • Existing tools lack multimodal situational awareness: telemetry metrics, workload state machines, cognitive agent reasoning traces, and audit logs are isolated across fragmented tools.

We built the Autonomous ML Workload Scheduler to bridge this gap: an agentic infrastructure management system that leverages Google's Agent Development Kit (ADK) and Gemini 3.5 to continuously observe cluster telemetry, evaluate deadlines, and execute safe, verified multi-step workload rescheduling on Google Cloud Platform.

What it does

The system provides autonomous, end-to-end management of heterogeneous compute clusters (NVIDIA A100/T4 GPUs and high-memory CPU nodes):

  1. Continuous Telemetry & Invariant Perception: Ingests live hardware utilization, GPU VRAM, power consumption, loss curve trajectories, and approaching SLA deadlines.
  2. Cognitive Policy Reasoning (Google ADK + Gemini 3.5): Evaluates multi-factor scheduling policies (workload protection, minimum progress thresholds, deadline urgency ranking) to formulate intelligent preemption and placement strategies.
  3. Strict Safety Boundary & Deterministic Control Plane: Enforces an un-bypassable mathematical verification layer. The agent is strictly bounded: all mutations must pass invariant checks ($\sum \text{allocated} \le \text{capacity}$) and valid state machine transitions before execution.
  4. Target-Agnostic Backend Interface: Seamlessly switches between an ultra-fast in-memory simulator and real Google Cloud Compute Engine VMs running a dedicated Python workload daemon (/opt/ml-scheduler/daemon.py) with zero changes to the control plane.
  5. Multimodal 4-Stream Observability Dashboard: A real-time cyberpunk-styled single-page UI built for high-stakes operational visibility, delivering four synchronized visual data streams and interactive scenario triggers.

How we built it

We built the system from the ground up using a modern Google Cloud Platform and Python architecture:

                                  GCP Cloud Run Service
                             (FastAPI Dashboard)
                                            │
        ┌───────────────────┬───────────────┴───────────────┬───────────────────┐
        │                   │                               │                   │
        ▼                   ▼                               ▼                   ▼
┌───────────────┐   ┌───────────────┐               ┌───────────────┐   ┌───────────────┐
│ Cloud Pub/Sub │   │  Google ADK   │               │Cloud Workflows│   │Cloud Firestore│
│ Ingress Queue │   │ Agent+Gemini  │               │ Orchestration │   │ Audit Memory  │
└───────┬───────┘   └───────┬───────┘               └───────┬───────┘   └───────┬───────┘
        │                   │                               │                   │
        └───────────────────┴───────────────┬───────────────┴───────────────────┘
                                            │
                                            ▼
                        ┌───────────────────────────────────────┐
                        │ Deterministic Scheduler Control Plane │
                        │      (Single Source of Truth)         │
                        └───────────────────┬───────────────────┘
                                            │
                                ┌───────────┴───────────┐
                                ▼                       ▼
                      ┌───────────────────┐   ┌───────────────────┐
                      │ Cluster Simulator │   │ Compute Engine VM │
                      │  (Local Engine)   │   │  Workload Daemon  │
                      └───────────────────┘   └───────────────────┘
  • Google Cloud Run: Hosts the containerized FastAPI web application and Server-Sent Events (SSE) streaming engine (<250MB image, non-root user execution).
  • Google Agent Development Kit (ADK) & Gemini 3.5: Provides structured reasoning, tool orchestration, and decision formulation.
  • Google Cloud Pub/Sub: Serves as the asynchronous event bus with idempotency deduplication and exponential retry backoff.
  • Google Cloud Workflows: Executes declarative multi-step rescheduling and failure-recovery pipelines.
  • Google Cloud Firestore: Stores persistent cluster state snapshots, workload history, and immutable decision audit logs with automatic credential sanitization.
  • Google Cloud Compute Engine: Powers real workload execution through a lightweight daemon managing subprocesses and training telemetry.
  • Terraform & Shell IaC: Fully reproducible infrastructure provisioning scripts (terraform/ and scripts/deploy.sh).

Challenges we ran into

  1. Safety Boundary Enforcement: Challenge: LLMs can hallucinate parameters or suggest invalid state transitions (e.g. attempting to resume a completed job, or allocating more VRAM than physically exists). Solution: We enforced a strict unidirectional safety architecture. Gemini can only invoke typed ADK tools that route to the Deterministic Scheduler. The scheduler mathematically verifies capacity, node compatibility, and state machine validity before applying any mutation.
  2. Asynchronous Race-Condition Prevention: Challenge: Multiple high-frequency Pub/Sub events could cause conflicting scheduling decisions. Solution: Implemented an idempotent event processing layer with event deduplication hashes, thread-safe memory locks, and correlation-tracked Cloud Workflows.
  3. Zero-Leakage Audit Memory: Challenge: Persisting raw runtime state in Firestore risks accidentally storing environment variables or Gemini API keys. Solution: Implemented a recursive credential safety sanitizer that strips sensitive tokens and secrets before writing to storage.

Accomplishments that we're proud of

  • 103 Automated Tests with 100% Pass Rate: Complete test coverage spanning unit tests, simulator invariants, ADK tools, Pub/Sub pipelines, Firestore adapters, Cloud Workflows runners, Compute Engine VM daemons, dashboard APIs, and E2E Cloud Run deployments.
  • True Multi-Stage Cold-Start Optimization: Multi-stage Docker image under 250MB with instant startup responsiveness.
  • Unified 4-Stream UX: Delivering a cohesive visual experience that merges hardware telemetry, state machines, LLM reasoning, and audit logging into a single synchronized view.

What we learned

1. Concrete Compute vs. Hallucinated Telemetry

  • The Reality Gap: Synthetic or hardcoded UI dashboards quickly drift into logical contradictions (e.g., CPU aggregated readouts not matching individual node sums, or preempting a CPU batch job to "free up" GPU VRAM).
  • The Fix: Hooking telemetry directly into live OS metrics (psutil running inside real Compute Engine VMs) guarantees zero metric hallucination. When a daemon allocates 2.5 GB of RAM, system metrics and Firestore snapshots reflect it accurately.

2. Dynamic Preemption Mechanics

  • Resource Isolation: Preemption is only valid if the evicted workload actually releases the target resource required by the incoming high-priority job.
  • State Machine Invariants: To maintain cluster safety:

$$\text{Priority}(\text{Incoming Job}) > \text{Priority}(\text{Running Job})$$

  • Graceful Eviction: Cloud Workflows must signal the worker daemon via persistent state (Firestore) to terminate its sub-processes (SIGTERM), ensuring system RAM/VRAM drops back to baseline before the critical job is scheduled.

3. Real-Time Hardware Integration

  • Lightweight Workload Simulation: You don't need heavy ML frameworks or expensive GPU drivers to test scheduling systems. A simple Python daemon utilizing bytearray memory allocation and CPU loop spikes provides real hardware load for testing system behavior under pressure.
  • Audit Trail Correlation: End-to-end reliability requires linking every Pub/Sub event, Gemini decision rationale, Cloud Workflow execution, and final Firestore state check under a unified correlation_id.

What's next for ML Workload Engine

  • Dynamic GKE & Cloud TPU Node Autoscaling: Extending the Compute Engine adapter to dynamically spin up preemptible TPU v5e/v6e slices on demand.
  • Reinforcement Learning from Scheduler Decisions: Training lightweight reward models on historical Firestore audit records to continually refine Gemini's preemption policies.
  • Multi-Region Bursting: Enabling cross-region workload migration when local GPU clusters experience regional capacity exhaustion.

Built With

  • cloud-firestore
  • cloud-pub/sub
  • cloud-workflows
  • compute-engine
  • fastapi
  • gemini-3.5-flash
  • google-agent-development-kit-(adk)
  • google-cloud-run
  • server-sent-events-(sse)
Share this project:

Updates