Autoscaling

The kubemq-rayserve adapter exposes queue depth as the scaling signal for ML inference workers. When tasks pile up faster than they are processed, more replicas come online; when the queue drains, replicas are released. There are two independent paths to wire this up.

Prerequisites

• kubemq-rayserve installed and a KubeMQTaskProcessorAdapter deployment already running (see Getting Started with Ray Serve)
• A running KubeMQ broker reachable from the policy or scaler
• For the KEDA path: KEDA v2.x and the kubemq-keda-scaler deployed in the cluster (see Two Autoscaling Paths below)

Two Autoscaling Paths

The in-process policy is the simplest to adopt — it is a single function you hand to Ray Serve. KEDA is the production-grade option when you want Kubernetes' Horizontal Pod Autoscaler (HPA) to own the scaling decision and you need scale-to-zero. Both read the same metric: the number of waiting messages in the task queue.

In-Process Policy

kubemq_queue_depth_policy is a Ray Serve custom autoscaling policy. Ray Serve invokes it on its autoscaling interval (roughly every 10 seconds), passing the per-deployment autoscaling contexts. The policy queries KubeMQ for the current queue depth and returns the desired replica count for each deployment.

kubemq_queue_depth_policy(
    ctxs: dict[str, Any],
    kubemq_address: str = "localhost:50000",
    queue_name: str = "",
    tasks_per_replica: int = 5,
    auth_token: str = "",
) -> tuple[dict[str, int], dict]

The scaling formula is:

desired_replicas = max(1, math.ceil(queue_depth / tasks_per_replica))

The same desired_replicas value is assigned to every deployment_id present in ctxs. The function returns a (decisions, state) tuple: decisions maps each deployment_id to its desired replica count, and state carries the observed queue_depth and computed desired_replicas for logging and inspection.

Basic Usage

The policy reads queue depth from a live broker, so a KubeMQ broker must be reachable. Start one locally if needed:

docker run -d \  --name kubemq \  -p 50000:50000 \  -p 9090:9090 \  -e KUBEMQ_TOKEN=YOUR_LICENSE_KEY \  europe-docker.pkg.dev/kubemq/images/kubemq:next

To see the policy compute a decision, enqueue some tasks, build a Ray Serve autoscaling context, and call the function directly:

import os
import uuid

from kubemq_rayserve import (
    KubeMQAdapterConfig,
    KubeMQTaskProcessorAdapter,
    kubemq_queue_depth_policy,
)

BROKER = os.environ.get("KUBEMQ_ADDRESS", "localhost:50000")


def dummy_handler(x: int) -> int:
    return x * 2


channel = f"example-autoscale-{uuid.uuid4().hex[:8]}"

config = KubeMQAdapterConfig(address=BROKER)
adapter = KubeMQTaskProcessorAdapter(config)


class _Cfg:
    queue_name = channel
    max_retries = 0
    failed_task_queue_name = ""
    unprocessable_task_queue_name = ""


adapter.initialize(consumer_concurrency=1, task_processor_config=_Cfg())

# Enqueue tasks WITHOUT starting the consumer -> builds queue depth
for i in range(12):
    adapter.enqueue_task_sync("dummy_handler", args=[i])


# Mock Ray Serve autoscaling context
class _MockContext:
    current_num_replicas = 1


ctxs = {"deployment-1": _MockContext()}

decisions, state = kubemq_queue_depth_policy(
    ctxs=ctxs,
    kubemq_address=BROKER,
    queue_name=channel,
    tasks_per_replica=5,
)

print(f"Queue depth: {state['queue_depth']}")          # 12
print(f"Desired replicas: {state['desired_replicas']}")  # max(1, ceil(12 / 5)) = 3
print(f"Decisions: {decisions}")                         # {"deployment-1": 3}

In a real Ray Serve deployment you register the function as the autoscaling policy rather than calling it yourself, binding the broker address and queue name with functools.partial so Ray Serve can invoke it with just ctxs.

Tuning tasks_per_replica

tasks_per_replica is the throughput knob. It states how many queued tasks a single replica is expected to absorb before another replica is warranted. For a fixed queue depth, lowering the value scales out more aggressively; raising it scales more conservatively:

import math

# queue_depth = 20 waiting tasks
for tpr in [1, 3, 5, 10, 20, 50]:
    desired = max(1, math.ceil(20 / tpr))
    print(f"tasks_per_replica={tpr:>2} -> {desired} replicas")

# tasks_per_replica= 1 -> 20 replicas   (one replica per pending task)
# tasks_per_replica= 3 -> 7  replicas
# tasks_per_replica= 5 -> 4  replicas   (default)
# tasks_per_replica=10 -> 2  replicas
# tasks_per_replica=20 -> 1  replica
# tasks_per_replica=50 -> 1  replica

As a rule of thumb: lower tasks_per_replica for latency-sensitive, fast inference where you want maximum parallelism; raise it for heavyweight models where each replica is expensive (for example a GPU worker) and you would rather batch work onto fewer replicas.

Authenticated Brokers

If your broker enforces JWT authentication, pass the token through auth_token. The policy forwards it to the SDK client it uses to read queue depth:

import os

from kubemq_rayserve import kubemq_queue_depth_policy

BROKER = os.environ.get("KUBEMQ_ADDRESS", "localhost:50000")
AUTH_TOKEN = os.environ.get("KUBEMQ_AUTH_TOKEN", "")


class _MockContext:
    current_num_replicas = 2


ctxs = {"deployment-1": _MockContext()}

decisions, state = kubemq_queue_depth_policy(
    ctxs=ctxs,
    kubemq_address=BROKER,
    queue_name="inference-tasks",
    tasks_per_replica=5,
    auth_token=AUTH_TOKEN,
)

Observing Decisions Over Time

To watch scaling decisions evolve as a consumer drains the queue, poll the policy on an interval and compare each desired_replicas against the previous count:

import time

from kubemq_rayserve import kubemq_queue_depth_policy

ctxs = {"deployment-1": _MockContext()}

for tick in range(10):
    decisions, state = kubemq_queue_depth_policy(
        ctxs=ctxs,
        kubemq_address="localhost:50000",
        queue_name="inference-tasks",
        tasks_per_replica=5,
    )

    depth = state.get("queue_depth", 0)
    desired = state.get("desired_replicas", 1)
    current = _MockContext.current_num_replicas

    if desired > current:
        action = f"SCALE UP ({current} -> {desired})"
    elif desired < current:
        action = f"SCALE DOWN ({current} -> {desired})"
    else:
        action = "NO CHANGE"

    print(f"{tick:>2}s | depth={depth:>3} | desired={desired:>2} | {action}")

    # Simulate Ray Serve applying the decision
    _MockContext.current_num_replicas = desired
    time.sleep(1.0)

As the consumer processes tasks, queue_depth falls, desired_replicas follows the formula down toward 1, and the printed decisions transition from SCALE UP through NO CHANGE to SCALE DOWN.

How the Client Cache Works

Ray Serve calls the policy frequently, so the module avoids opening a fresh gRPC connection on every invocation. A module-level QueuesClient cache is keyed by address:sha256(auth_token)[:16], which lets multiple brokers (and multiple credentials) coexist without embedding raw tokens in dictionary keys. The cache is guarded by a threading.Lock so concurrent autoscaling ticks never create duplicate clients, and an atexit hook closes every cached client on process exit.

token_hash = hashlib.sha256(auth_token.encode()).hexdigest()[:16] if auth_token else ""
key = f"{address}:{token_hash}"
with _client_cache_lock:
    if key not in _client_cache:
        config = ClientConfig(
            address=address,
            client_id=f"rayserve-autoscaler-{token_hash[:8] or 'noauth'}",
            auth_token=auth_token or None,
        )
        _client_cache[key] = QueuesClient(config=config)
    return _client_cache[key]

This is an internal detail — you do not manage the cache yourself — but it explains why the policy is cheap to call on a tight interval and why no explicit client cleanup is required in your code.

KEDA Path

When you want Kubernetes to own scaling — including true scale-to-zero — use KEDA. KEDA polls the KubeMQ external scaler over gRPC and drives the target's replica count through a standard HPA. You scale the underlying Deployment, StatefulSet, or Ray worker group rather than Ray Serve replicas directly.

ScaledObject

The ScaledObject declares the scaling bounds and an external trigger. The trigger points at two endpoints: scalerAddress is the standalone external scaler's Service (port 9090), and kubemqAddress is the KubeMQ broker (gRPC 50000) the scaler queries for queue depth. KEDA computes desiredReplicas = queueDepth / targetWaiting, mirroring the tasks_per_replica math of the in-process policy — targetWaiting: "5" means five waiting messages per replica.

apiVersion: keda.sh/v1alpha1
kind: ScaledObject
metadata:
  name: kubemq-rayserve-scaledobject
  namespace: default
spec:
  scaleTargetRef:
    name: ray-worker-deployment   # Replace with your Deployment name

  minReplicaCount: 1    # Set to 0 for scale-to-zero
  maxReplicaCount: 20
  pollingInterval: 15   # How often KEDA checks the scaler (seconds)
  cooldownPeriod: 60    # Seconds to wait before scaling down after the last trigger

  advanced:
    restoreToOriginalReplicaCount: false
    horizontalPodAutoscalerConfig:
      behavior:
        scaleUp:
          stabilizationWindowSeconds: 0    # Scale up immediately
          policies:
            - type: Pods
              value: 4                      # Add up to 4 pods per 15s window
              periodSeconds: 15
        scaleDown:
          stabilizationWindowSeconds: 120   # Wait 2 min before scaling down
          policies:
            - type: Percent
              value: 25                     # Remove at most 25% per period
              periodSeconds: 60

  triggers:
    - type: external
      metadata:
        scalerAddress: "kubemq-keda-scaler.default.svc.cluster.local:9090"  # External scaler Service
        kubemqAddress: "kubemq-cluster-service.default.svc.cluster.local:50000"  # Broker (gRPC)
        queueName: "rayserve-tasks"   # Must match the adapter's queue_name
        targetWaiting: "5"             # desiredReplicas = queueDepth / targetWaiting

The trigger has three required metadata fields — scalerAddress (the external scaler Service FQDN on port 9090), kubemqAddress (the broker's host:port on gRPC 50000, which the scaler dials to read the live Waiting count), and queueName — plus targetWaiting, the per-replica target. Unlike the in-process policy, KEDA's minReplicaCount may be 0, so the HPA can scale the workload all the way down when the queue is empty and back up the moment a task arrives. The behavior block tunes how quickly that happens — here, scale-up is immediate (add up to four pods every 15 seconds) while scale-down is deliberately damped (a two-minute stabilization window, removing at most 25% per minute) to avoid thrashing GPU workers.

Credentials via TriggerAuthentication

For production, do not inline an authToken in the trigger metadata. Reference a Kubernetes Secret through a TriggerAuthentication and secretTargetRef so credentials never appear in the ScaledObject:

apiVersion: keda.sh/v1alpha1
kind: TriggerAuthentication
metadata:
  name: kubemq-trigger-auth
  namespace: default
spec:
  secretTargetRef:
    - parameter: authToken
      name: kubemq-adapter-auth   # The Secret holding the token
      key: auth-token

Reference it from the trigger with authenticationRef.name: kubemq-trigger-auth, and KEDA injects the secret's auth-token value as the scaler's authToken parameter at poll time.

Choosing Between Them

Related

• Kubernetes Production Deployment — the full production scenario: broker, adapter config, Ray cluster, and KEDA wired together.
• KEDA integration — the cluster-level autoscaler that drives the KEDA path above.
• API Reference — the kubemq_queue_depth_policy signature, adapter methods, and the metrics dict.