Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SAIR: Cost-Efficient Multi-Stage ML Pipeline Autoscaling via In-Context Reinforcement Learning

Jianchang Su, Yifan Zhang, Shengkai Lin, Shizhen Zhao, Yusheng Zheng, Yiwei Yang, Wei Zhang

TL;DR

SAIR tackles autoscaling of multi-stage ML pipelines by deploying an in-context reinforcement learning controller that leverages an LLM to learn scaling policies online without gradient updates. The approach combines Pareto-frontier reward shaping, surprisal-based experience retrieval, and lightweight CUDA-based GPU rate control to coordinate across preprocessing, inference, and postprocessing stages. Key theoretical contributions include regret bounds that decompose errors into retrieval coverage and LLM selection terms, plus frontier-separation guarantees, while practical results show SAIR achieving the best or tied-best P99 latency and substantial cost reductions across four pipelines and three workload patterns, with 86% bottleneck-detection accuracy and zero offline training. This combination enables zero-shot deployment, robust generalization to distribution shifts, and fine-grained resource control suitable for shared GPU environments.

Abstract

Multi-stage ML inference pipelines are difficult to autoscale due to heterogeneous resources, cross-stage coupling, and dynamic bottleneck migration. We present SAIR, an autoscaling framework that uses an LLM as an in-context reinforcement learning controller, improving its policy online from reward-labeled interaction histories without gradient updates. SAIR combines Pareto-dominance reward shaping with a provable separation margin, surprisal-guided experience retrieval for context efficiency, and fine-grained GPU rate control via user-space CUDA interception. We provide regret analysis decomposing error into retrieval coverage and LLM selection components. On four ML serving pipelines under three workload patterns, SAIR achieves the best or tied-best P99 latency and effective resource cost among deployed baselines, improving P99 by up to 50% and reducing effective cost by up to 97% (under GPU rate-control assumptions), with 86% bottleneck detection accuracy and no offline training.

SAIR: Cost-Efficient Multi-Stage ML Pipeline Autoscaling via In-Context Reinforcement Learning

TL;DR

SAIR tackles autoscaling of multi-stage ML pipelines by deploying an in-context reinforcement learning controller that leverages an LLM to learn scaling policies online without gradient updates. The approach combines Pareto-frontier reward shaping, surprisal-based experience retrieval, and lightweight CUDA-based GPU rate control to coordinate across preprocessing, inference, and postprocessing stages. Key theoretical contributions include regret bounds that decompose errors into retrieval coverage and LLM selection terms, plus frontier-separation guarantees, while practical results show SAIR achieving the best or tied-best P99 latency and substantial cost reductions across four pipelines and three workload patterns, with 86% bottleneck-detection accuracy and zero offline training. This combination enables zero-shot deployment, robust generalization to distribution shifts, and fine-grained resource control suitable for shared GPU environments.

Abstract

Multi-stage ML inference pipelines are difficult to autoscale due to heterogeneous resources, cross-stage coupling, and dynamic bottleneck migration. We present SAIR, an autoscaling framework that uses an LLM as an in-context reinforcement learning controller, improving its policy online from reward-labeled interaction histories without gradient updates. SAIR combines Pareto-dominance reward shaping with a provable separation margin, surprisal-guided experience retrieval for context efficiency, and fine-grained GPU rate control via user-space CUDA interception. We provide regret analysis decomposing error into retrieval coverage and LLM selection components. On four ML serving pipelines under three workload patterns, SAIR achieves the best or tied-best P99 latency and effective resource cost among deployed baselines, improving P99 by up to 50% and reducing effective cost by up to 97% (under GPU rate-control assumptions), with 86% bottleneck detection accuracy and no offline training.
Paper Structure (24 sections, 7 theorems, 20 equations, 6 figures, 5 tables, 2 algorithms)

This paper contains 24 sections, 7 theorems, 20 equations, 6 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Problem Formulation
  4. SAIR: In-Context RL for Autoscaling
  5. Theoretical Analysis
  6. Experiments
  7. Experimental Setup
  8. Main Results
  9. Bottleneck Detection
  10. Ablation Study
  11. Discussion
  12. Conclusion
  13. Proof Details
  14. Proof of Theorem \ref{['thm:regret']} (SAIR Regret Bound)
  15. Proof of Proposition \ref{['prop:pareto']} (Frontier Separation)
  16. ...and 9 more sections

Key Result

Theorem 5.1

Under Assumptions ass:bounded--ass:llm, let $\epsilon_t$ be the exploration probability at round $t$. SAIR achieves: where $\xi_t$ is retrieval coverage gap, $\eta_t$ is LLM selection error, and $\sum_t \epsilon_t$ accounts for exploration rounds.

Figures (6)

  • Figure 1: P99 latency comparison across applications and patterns.
  • Figure 2: Normalized cost comparison (SAIR=1). Values $>$1 indicate higher cost.
  • Figure 3: Bottleneck detection confusion matrix (200 test cases). Overall accuracy: 86%. Per-class metrics in Appendix \ref{['app:confusion_metrics']}.
  • Figure 4: Ablation study: all ablations degrade from Full SAIR baseline (dashed line).
  • Figure 5: Learning dynamics. Full SAIR (solid) accumulates higher reward and achieves more stable throughput than ablated variants (dashed).
  • ...and 1 more figures

Theorems & Definitions (11)

  • Theorem 5.1: SAIR Regret Bound
  • Proposition 5.2: Pareto Reward: Frontier Separation
  • Theorem 5.3: Bottleneck Detection Sample Complexity
  • Lemma 5.4: Bottleneck Marginal Gain
  • Lemma 5.5: Retrieval Coverage Under Smoothness
  • Corollary 5.6: Sublinear Cumulative Coverage Regret
  • proof
  • proof
  • proof
  • Lemma B.1: Surprisal-Information Gain Connection
  • ...and 1 more