Implementation:ArroyoSystems Arroyo Kubernetes Scheduler

Overview

KubernetesScheduler implements the Scheduler trait to manage Arroyo worker pods on a Kubernetes cluster. It creates pods with configurable resources, environment variables, volumes, and labels, and manages their lifecycle through the Kubernetes API.

Description

The KubernetesScheduler struct wraps a Kubernetes API client and a KubernetesSchedulerConfig. Its primary responsibilities are:

• Pod creation: The make_pod method constructs a Kubernetes Pod spec from the pipeline request, including resource calculations (per-slot or per-pod modes), environment variables, volume mounts, labels, annotations, and owner references.
• Resource scaling: In PerSlot mode, CPU and memory resource requests/limits are multiplied by the number of slots assigned to each pod. In PerPod mode, resources are used as-is.
• Worker start: Calculates the number of replicas needed based on total slots and max slots per pod, then creates pods via the Kubernetes API.
• Worker stop: Deletes pods by label selector (job ID and optional run ID), with optional force deletion (grace period 0), then polls until pods are terminated.
• Worker discovery: Lists pods by label selector and filters out those with a deletion timestamp.

The scheduler uses labels (cluster, job_id, run_id, job_name) for pod identification and management.

Usage

The Kubernetes scheduler is selected when Arroyo is deployed in a Kubernetes environment. It is instantiated with a default Kubernetes client and the cluster configuration.

Code Reference

Source Location

crates/arroyo-controller/src/schedulers/kubernetes/mod.rs

Signature

pub struct KubernetesScheduler {
    client: Option<Client>,
    config: KubernetesSchedulerConfig,
}

impl KubernetesScheduler {
    pub async fn new() -> Self
    pub fn with_config(client: Option<Client>, config: KubernetesSchedulerConfig) -> Self
    fn make_pod(&self, req: &StartPipelineReq, number: usize, slots: usize) -> Pod
}

#[async_trait]
impl Scheduler for KubernetesScheduler {
    async fn start_workers(&self, req: StartPipelineReq) -> Result<(), SchedulerError>
    async fn register_node(&self, _req: RegisterNodeReq)
    async fn heartbeat_node(&self, _req: HeartbeatNodeReq) -> Result<(), Status>
    async fn worker_finished(&self, _req: WorkerFinishedReq)
    async fn stop_workers(&self, job_id: &str, run_id: Option<u64>, force: bool) -> anyhow::Result<()>
    async fn workers_for_job(&self, job_id: &str, run_id: Option<u64>) -> anyhow::Result<Vec<WorkerId>>
}

Import

use crate::schedulers::kubernetes::KubernetesScheduler;

I/O Contract

Inputs

Name	Type	Description
StartPipelineReq	struct	Pipeline start request containing job ID, run ID, slots, env vars, and WASM path
job_id	&str	Job identifier for stopping or querying workers
run_id	Option<u64>	Optional run identifier for scoping operations to a specific run

Outputs

Name	Type	Description
Pod	Kubernetes Pod	Created Kubernetes pod spec with configured containers, resources, and labels
workers	Vec<WorkerId>	List of active worker IDs for a given job

Usage Examples

// Creating a Kubernetes scheduler with default config
let scheduler = KubernetesScheduler::new().await;

// Starting workers for a pipeline
let req = StartPipelineReq {
    name: "my_pipeline".to_string(),
    job_id: Arc::new("job_123".to_string()),
    run_id: 1,
    slots: 8,
    // ...
};
scheduler.start_workers(req).await?;

// Stopping workers for a job
scheduler.stop_workers("job_123", Some(1), false).await?;

Related Pages

• Implementation:ArroyoSystems_Arroyo_Scheduler_Trait - The Scheduler trait that KubernetesScheduler implements
• Implementation:ArroyoSystems_Arroyo_K8s_Quantities - Kubernetes quantity parsing used for resource calculations
• Implementation:ArroyoSystems_Arroyo_Controller_Server - The controller server that uses the scheduler