Kubernetes Scheduling

Kubernetes schedules pods independently by default (see Kubernetes Scheduler). When an experiment is submitted, each pod is placed on a node as soon as the resources it requests are available. There is no default mechanism to hold pods until the full set can start together (although this can be enabled through Gang scheduling). Whether this matters depends on the workload.

Asynchronous scheduling

For embarrassingly parallel and sequential workloads, pods starting at different times causes no problems.

Embarrassingly parallel: each job is independent. If only 4 of 16 GPUs are free, 4 trial pods start immediately and the rest queue. Each pod does useful work from the moment it starts.

Sequential batch processing: a single job with completions > 1 and parallelisms < completions. Pods run one after another anyway. Staggered scheduling simply means the queue drains as resources free up rather than all at once.

Concretely, this covers:

PyTorch: several independent training runs submitted as replicas of the same job (e.g. a sweep of model variants), each running its own single-process training loop with no dependency on the others.
JAX: several independent single-host JAX programs run as replicas (e.g. a parameter sweep), since there's no cross-replica rendezvous.
Ray (kuberay): additional workers an autoscaling cluster requests as load ramps up. Ray puts each new worker to use as soon as it lands.
XGBoost: several independent single-node training runs submitted as replicas (e.g. a hyperparameter sweep), each training its own model with no rendezvous between replicas.

See the JobSet examples for concrete manifests of both patterns.

Synchronous workloads need all pods ready

Some workloads require the full set of pods to be present before any useful work can begin:

Custom distributed workloads that coordinate via barriers or rendezvous points
Multi-node jobs where all workers must connect to each other or to a coordinator before starting

Concretely, this covers:

PyTorch: a DDP/FSDP job where every worker must rendezvous before training can start stepping.
JAX: a multi-host SPMD job where every host must be present before collective ops can execute.
Ray (kuberay): the head node and the minimum worker count a job needs to bootstrap.
XGBoost: a distributed job where every rank must connect to rank 0's tracker before boosting rounds run.

When pods start at different times in these cases, early pods block waiting for their peers and hold their resource allocation (including GPUs) while idle.

Gang scheduling

Gang scheduling solves the synchronous scheduling problem by treating a group of pods as an atomic unit: either all pods are scheduled together, or none are. Pods wait in a queue until all required resources are simultaneously available, then start together. Without it, pods that do get scheduled sit idle holding their resource allocation (including GPUs) while they wait for the rest of the group.

Gang scheduling requires a cluster-level scheduler plugin. On AIchor-managed engines this is backed by the KAI scheduler. On imported engines, availability depends on whether the underlying cluster has a compatible scheduler plugin installed, so contact the cluster administrator to confirm.

Having multiple replicas in a job does not automatically mean it needs gang scheduling. What matters is whether the replicas are tightly coupled (they must all be present at once to make progress) or independent (each replica can do useful work on its own, regardless of whether its siblings are scheduled yet).

When to use it

You're running tightly-coupled multi-pod jobs, meaning replicas that must rendezvous to make progress, such as PyTorch DDP/FSDP, JAX multi-host SPMD, a Ray cluster's initial head + minimum-worker bootstrap, or a distributed XGBoost job whose workers must all connect to the tracker.

When it's not relevant

Independent-replica workloads, such as a sweep of standalone PyTorch/JAX/XGBoost runs, or an autoscaling Ray cluster's incremental worker additions. Each replica is useful on its own, so all-or-nothing scheduling only adds latency for no benefit.
Single-pod jobs have nothing to "gang", since there's only one pod to schedule.

Enabling it

Set gangScheduling.enabled: true in the manifest. It's a spec-level setting that applies the same way across every operator: pytorch, jax, xgboost, jobset, and kuberay.

spec:
  gangScheduling:
    enabled: true

What changes when it's enabled

pytorch, jax, xgboost, jobset: the number of pods requested is the same either way. What changes is whether the experiment waits for all of them to be placeable together before any of them start. Disabled (the default), pods can start as soon as they're individually schedulable, so early pods may need to wait for late ones. See the fallback below. Enabled, either every pod starts together or none do.

kuberay: without gang scheduling, a worker pool's count is not a guaranteed starting size. It's a ceiling instead. The pool starts with a single worker, and Ray's own autoscaler grows it toward count as the cluster actually needs the extra capacity. With gang scheduling enabled, the full count is requested upfront and scheduled as one atomic unit alongside the head node.

Scaling `kuberay` workers back down

gangScheduling.autoScalingScaleDown sets how long, in seconds, an idle kuberay worker waits before Ray's autoscaler scales it back down. It only has an effect while gang scheduling is disabled, since that's the only case where workers scale up on demand in the first place.

spec:
  gangScheduling:
    enabled: false
    autoScalingScaleDown: 300 # scale an idle worker back down after 5 minutes

Leave it unset to keep idle workers around indefinitely, matching the behavior from before this autoscaling path existed.

Demos

The AIchor team maintains gang scheduling demo projects, one per operator, each showing a tightly-coupled workload that needs the feature and the rendezvous code that would otherwise be required to work around its absence:

Fallback: synchronizing without gang scheduling

Where gang scheduling isn't available, for example an imported engine without a compatible scheduler plugin, pods can probe peer hostnames via DNS and wait until all are reachable before starting distributed work themselves. Early pods spin idle (holding their resource allocation) for the duration of the wait. Where gang scheduling is available, enabling it (above) removes the need for this: the platform holds every pod until the whole group can start together, so there's nothing left for the pods themselves to wait for.

JobSet creates one headless Service per JobSet and sets each pod's hostname to {jobset-name}-{group}-{job-index}-{completion-index} with the subdomain defaulting to the JobSet name. This makes every peer's DNS address predictable at runtime from the injected environment variables.

Two coordination patterns are available:

Decentralized (every pod waits for all others)

Suitable for homogeneous groups with no designated coordinator. Each pod independently probes all its peers.

import os
import socket
import time


def wait_for_peers(timeout: int = 300) -> None:
    hostname = socket.gethostname()
    job_index = int(os.environ["JOB_INDEX"])
    completion_index = int(os.environ["JOB_COMPLETION_INDEX"])
    group_name = os.environ["REPLICATED_JOB_NAME"]
    replicas = int(os.environ["REPLICATED_JOB_REPLICAS"])

    # hostname = "{jobset-name}-{group}-{job-index}-{completion-index}"
    base = hostname.rsplit("-", 2)[0]             # "{jobset-name}-{group}"
    jobset_name = base[: -(len(group_name) + 1)]  # "{jobset-name}"

    peers = [
        f"{base}-{i}-{completion_index}.{jobset_name}"
        for i in range(replicas)
        if i != job_index
    ]

    print(f"[peers] I am {hostname} (job_index={job_index}); waiting for {len(peers)} peers", flush=True)

    deadline = time.time() + timeout
    remaining = set(peers)
    while remaining and time.time() < deadline:
        for host in list(remaining):
            try:
                socket.getaddrinfo(host, None)
                remaining.discard(host)
                print(f"[peers] {host} ready", flush=True)
            except socket.gaierror:
                pass
        if remaining:
            time.sleep(5)

    if remaining:
        raise RuntimeError(f"Timed out waiting for peers: {remaining}")


wait_for_peers()
# all peers are up, proceed with distributed work
print("[peers] all peers up, starting work", flush=True)
time.sleep(5)  # stand-in for the real distributed workload
print("[peers] done", flush=True)

Coordinator (workers wait for coordinator, coordinator waits for all workers)

Suitable for heterogeneous groups with a dedicated coordinator or master. Workers check a single hostname; the coordinator checks all workers.

Both roles can ship in a single script and one manifest command: the pod reads REPLICATED_JOB_NAME and dispatches to the worker or coordinator path.

import os
import socket
import time

COORDINATOR_GROUP = "coordinator"
WORKER_GROUP = "worker"


def wait_for_coordinator(coordinator_group: str = COORDINATOR_GROUP, timeout: int = 300) -> None:
    hostname = socket.gethostname()
    job_index = int(os.environ["JOB_INDEX"])
    completion_index = int(os.environ["JOB_COMPLETION_INDEX"])
    group_name = os.environ["REPLICATED_JOB_NAME"]

    base = hostname.rsplit("-", 2)[0]
    jobset_name = base[: -(len(group_name) + 1)]

    # coordinator is always job 0, pod 0 of its group
    coordinator_host = f"{jobset_name}-{coordinator_group}-0-0.{jobset_name}"

    print(f"[worker] I am {hostname}; waiting for coordinator {coordinator_host}", flush=True)

    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            socket.getaddrinfo(coordinator_host, None)
            print(f"[worker] coordinator ready", flush=True)
            return
        except socket.gaierror:
            time.sleep(5)

    raise RuntimeError(f"Timed out waiting for coordinator: {coordinator_host}")


def wait_for_workers(worker_group: str = WORKER_GROUP, timeout: int = 300) -> None:
    hostname = socket.gethostname()
    job_index = int(os.environ["JOB_INDEX"])
    completion_index = int(os.environ["JOB_COMPLETION_INDEX"])
    group_name = os.environ["REPLICATED_JOB_NAME"]
    global_replicas = int(os.environ["GLOBAL_REPLICAS"])

    base = hostname.rsplit("-", 2)[0]
    jobset_name = base[: -(len(group_name) + 1)]

    # GLOBAL_REPLICAS counts all jobs across all groups; subtract the coordinator
    worker_replicas = global_replicas - 1

    workers = [
        f"{jobset_name}-{worker_group}-{i}-0.{jobset_name}"
        for i in range(worker_replicas)
    ]

    print(f"[coordinator] I am {hostname}; waiting for {len(workers)} workers", flush=True)

    deadline = time.time() + timeout
    remaining = set(workers)
    while remaining and time.time() < deadline:
        for host in list(remaining):
            try:
                socket.getaddrinfo(host, None)
                remaining.discard(host)
                print(f"[coordinator] {host} ready", flush=True)
            except socket.gaierror:
                pass
        if remaining:
            time.sleep(5)

    if remaining:
        raise RuntimeError(f"Timed out waiting for workers: {remaining}")


# Dispatch on the JobSet group this pod belongs to. Both roles run from the
# same image and command; REPLICATED_JOB_NAME tells each pod which it is.
group = os.environ["REPLICATED_JOB_NAME"]
if group == COORDINATOR_GROUP:
    wait_for_workers()
    print("[coordinator] all workers up, begin coordination", flush=True)
else:
    wait_for_coordinator()
    print("[worker] coordinator is up, proceed", flush=True)

The worker_replicas = global_replicas - 1 calculation assumes one coordinator job. For setups with multiple groups, pass the exact worker count explicitly rather than deriving it from GLOBAL_REPLICAS.

Recovering from eviction

A pod is evicted when the platform terminates it before the workload finishes, most commonly when a spot instance is reclaimed, but also when a node runs out of memory or disk, becomes unhealthy, or is drained for maintenance. The jobset and kuberay operators support automatic recovery when this happens. The number of allowed restarts is set via spec.restartPolicy.backoffLimit:

spec:
  operator: "jobset"  # or kuberay
  restartPolicy:
    backoffLimit: 5

In the snippet above, spec.restartPolicy.backoffLimit is the number of allowed restarts: this experiment can handle 5 failures (including evictions) before being marked as failed.

Asynchronous scheduling​

Synchronous workloads need all pods ready​

Gang scheduling​

When to use it​

When it's not relevant​

Enabling it​

What changes when it's enabled​

Scaling kuberay workers back down​

Demos​

Fallback: synchronizing without gang scheduling​

Decentralized (every pod waits for all others)​

Coordinator (workers wait for coordinator, coordinator waits for all workers)​

Recovering from eviction​