ProtoRAIL: A Risk-cognizant Imitation Agent for Adaptive vCPU Oversubscription in the Cloud

ProtoRAIL: A Risk-cognizant Imitation Agent for Adaptive vCPU Oversubscription in the Cloud

Speakers: Lu Wang, Mayukh Das, Fangkai Yang, Íñigo Goiri, Saravan Rajmohan, Dongmei Zhang

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the highly competitive and resource-intensive landscape of cloud computing, optimizing resource utilization is paramount for both operational efficiency and profitability. The talk, "ProtoRAIL: A Risk-cognizant Imitation Agent for Adaptive vCPU Oversubscription in the Cloud," presented by Mayukh Das from Microsoft's MC65 Research team, delves into a critical but often overlooked aspect of cloud infrastructure management: vCPU oversubscription. This work introduces ProtoRAIL, a novel AI agent designed to dynamically adjust the allocation of virtual CPUs (vCPUs) to virtual machines (VMs), aiming to maximize resource utilization while meticulously managing the risk of performance degradation due to resource contention.

The core motivation behind ProtoRAIL is to address the "cloud efficiency" problem, which the speaker likens to the essential, albeit less glamorous, "plumbing" of the digital world. By ensuring that cloud resources are used as effectively as possible, ProtoRAIL directly impacts a cloud provider's Cost of Goods Sold (COGS), driving down operational expenditure and boosting revenue margins. The talk highlights that traditional, static approaches to oversubscription often lead to either wasted resources or unacceptable levels of risk. ProtoRAIL distinguishes itself by employing a prototypical imitation learning approach that learns from real-world telemetry, incorporates domain knowledge for safety, and uses interpretable prototypes to handle the inherent granularity challenges of cloud resource management.

Background

▶ Watch: Introduction to ProtoRAIL and talk outline (0:00)

The fundamental challenge in cloud resource management is the optimal mapping of a finite pool of physical resources to a diverse and dynamic pool of customer demands. This encompasses a range of complex problems, including forecasting, packing optimization, and balancing competing objectives. At its heart lies capacity or resource efficiency, where oversubscription is a key technique. Oversubscription entails offering more virtual resources than the physical capacity available, based on the statistical assumption that not all users will simultaneously utilize their full allocated capacity. This practice is crucial for diminishing unutilized resources, increasing service gains, and reducing COGS for cloud providers – a concept analogous to airlines overbooking flights to ensure full capacity despite no-shows.

The focus of this work is vCPU oversubscription, as CPU is a primary billable unit in cloud services. Consider a scenario where users request 16-core VMs, but a physical server only has 60 physical CPUs, and each vCPU translates to 1.2 physical CPUs. Without oversubscription, only three such VMs can be placed, leading to 28 standard cores being unutilized and "blocked" by allocation. The solution proposed is to opaquely allocate less than what is requested, a percentage known as the oversubscription ratio (zeta). This allows packing more VMs onto a physical machine, assuming temporary bursts can be accommodated. The critical challenge then becomes determining this ideal zeta: too tight, and the physical machine overloads; too loose, and resources are wasted.

Practically, directly optimizing VM placement and the oversubscription ratio simultaneously is intractable. Therefore, the problem is decoupled: the VM scheduler is assumed to handle placement efficiently, while ProtoRAIL focuses on optimizing the zeta ratio. This creates a competing objective problem: minimizing the risk of overloaded physical machines (leading to contention, jitter, or performance degradation) versus maximizing the benefits from harvested cores (i.e., the resources freed up by oversubscription). Several challenges complicate this: aggressive policies increase contention, conservative policies waste resources, demands vary significantly over time, and decisions need to be made at appropriate granularities (per VM, per client, per service, or per region). Safety is paramount, as overloading can directly impact customer experience and service level agreements (SLAs).

Prior attempts to solve this problem include:

• Naive Statistical Models: Platforms like Azure Resource Central use heuristics to classify VMs as high or low usage, applying a constant oversubscription ratio. This approach yields trivial savings and lacks explicit risk control.
• Reinforcement Learning (RL) Methods: While RL can adaptively choose ratios, it faces significant practical hurdles. Training an RL agent in a live production environment is infeasible and unsafe. Training in test environments lacks real traffic diversity. Designing an effective reward function that accurately reflects oversubscription success is difficult. Furthermore, RL methods struggle with decision granularity and safety, as constrained RL, often used for safety, is known to have convergence issues.
• Offline Methods (e.g., LSTM-based, Coin): These methods analyze past telemetry to derive policies. While beneficial, they contend with noisy data, which introduces safety risks. A major limitation is the absence of ground truth for an ideal oversubscription ratio; researchers often resort to pseudo ground truth by observing actual usage patterns with some buffer. Granularity remains a persistent problem.

Key Findings

▶ Watch: VCPU oversubscription problem with an example (4:00)

ProtoRAIL emerges as a significant advancement in adaptive cloud resource management, primarily through its novel application of prototypical imitation learning. The core findings and contributions are:

• Novel Prototypical Imitation Learning: ProtoRAIL introduces a new paradigm for learning dynamic oversubscription policies. Instead of direct RL, it learns from historical usage telemetry, mimicking observed patterns to derive optimal zeta ratios.
• Risk-Cognizant and Domain Knowledge Integration: A critical feature of ProtoRAIL is its explicit awareness and management of risk. It integrates Knowledge-in-the-Loop (KITL), allowing the model to actively query and incorporate domain expertise to refine its policies, thereby minimizing the likelihood of physical machine overloading.
• Granularity Solved by Prototypes: The agent addresses the long-standing challenge of decision granularity by discovering and leveraging prototypes. These prototypes represent equivalence classes of approximately symmetric usage patterns, enabling the system to apply tailored policies based on the observed behavior type.
• Superior Performance in Both Risk and Benefit: Experimental evaluations on real-world Microsoft internal workloads demonstrate that ProtoRAIL is a "clear winner" compared to existing methods. It simultaneously achieves substantially higher sellable core hours (a measure of benefit) while minimizing overloading risk.
• Interpretable Policies: The use of prototypes not only solves granularity but also enhances the interpretability of the learned policies. Cloud operators can inspect the discovered prototypes to understand the underlying usage patterns driving specific oversubscription decisions.
• Efficient and Practical Deployment: The KITL module operates with minimal queries (as few as 6-7 per training run), making the learning process efficient. ProtoRAIL has been tested on real problems within Microsoft and is actively "on the path to production," underscoring its practical viability.
• Generalizability: The underlying prototypical imitation learning method is not restricted to vCPU oversubscription but is designed to generalize to "any resource optimization problem under the sun," suggesting broad applicability across different cloud resource types and optimization challenges.

Technical Deep Dive

▶ Watch: Concept of allocating less to pack more VMs (5:00)

ProtoRAIL is architected as an intelligent agent comprising three interconnected modules, designed to collaboratively learn and apply adaptive vCPU oversubscription policies.

1. Module 1: Prototype Discovery:

This module is foundational to ProtoRAIL's ability to handle the diverse and dynamic nature of cloud workloads. It is responsible for identifying prototypes, which are essentially representative samples that encapsulate distinct, recurring patterns of vCPU usage. These prototypes act as centroids for "equivalence classes of approximately symmetric patterns" in workload trajectories. The process involves:

• Trajectory Embedding: Historical vCPU usage data (telemetry) for various VMs or services is transformed into fixed-size numerical embeddings. This step often utilizes techniques like autoencoders or recurrent neural networks to capture temporal dependencies and salient features of the usage patterns.
• Clustering/Discovery: These embeddings are then clustered into groups, with each cluster's centroid or a representative sample forming a prototype. The goal is to identify a compact set of prototypes that can effectively characterize the vast spectrum of observed workload behaviors.
• Similarity Function: A crucial component is the similarity function, which can quantify how closely a new, unseen workload's trajectory embedding matches any of the discovered prototypes. This function allows ProtoRAIL to classify new workloads into existing prototype categories or identify novel patterns.

1. Module 2: Policy Learning:

Once prototypes are established, this module learns the actual oversubscription policy. Unlike traditional RL that explores actions, ProtoRAIL uses imitation learning, specifically behavior cloning, to learn from past "expert" behavior derived from telemetry. The "expert" here is not a human, but rather the observed, potentially sub-optimal, usage patterns and a derived pseudo ground truth for safe oversubscription.

• Individual Prototype Policies: For each discovered prototype, an individual policy is learned. This policy maps the characteristics of a workload belonging to that prototype to an optimal oversubscription ratio (zeta).
• Combined Policy: When a new workload arrives, its embedding is compared against all known prototypes using the similarity function. The final oversubscription ratio for this workload is then determined by combining the individual policies of the prototypes, weighted by their similarity to the incoming workload. This allows for fine-grained, adaptive decisions.
• Behavior Cloning Loss: The learning process is optimized using a behavior cloning loss. This loss function minimizes the difference between the zeta ratio predicted by ProtoRAIL and the zeta ratio derived from the pseudo ground truth (e.g., actual usage plus a safety buffer) for historical data. The goal is to mimic the "safe" and "efficient" oversubscription decisions that would ideally have been made in the past.

1. Module 3: Knowledge-in-the-Loop (KITL):

This module is ProtoRAIL's safeguard and refinement mechanism, directly addressing the safety concerns inherent in oversubscription and the "no ground truth" problem. KITL incorporates domain knowledge to guide and correct the learning process, particularly in situations where the model's predictions might lead to high risk.

• Active Querying: The model is designed to actively identify scenarios where its current policy might be problematic or uncertain (e.g., predicting a zeta that could lead to overloading based on historical patterns). In these cases, it "queries" for feedback.
• Feedback Integration: This feedback, representing domain expertise or predefined safety constraints, is then integrated into the learning loop. This feedback could be a hard constraint (e.g., "never oversubscribe this type of workload beyond X%") or a soft penalty.
• Loss Scaling and Refinement: The feedback is "exponentiated" and used to scale the behavior cloning loss. This means that errors or risky predictions in problematic scenarios incur a much higher penalty, forcing the model to prioritize safety and refine its policy in those areas. This iterative refinement helps the model converge towards safer and more robust policies, even with limited external queries. The speaker noted its efficiency, requiring "only like six or seven queries at most in a training run." KITL also helps refine the prototypes themselves, potentially removing or merging less useful ones, as demonstrated by one prototype being "gone" because it wasn't contributing much.

The entire ProtoRAIL agent operates by continuously observing vCPU usage trajectories, identifying their closest prototype, applying a dynamically weighted policy derived from these prototypes, and refining this policy through imitation learning and domain knowledge feedback. The oversubscription ratio (zeta), defined as the allocated-to-requested resource ratio, is the primary decision variable optimized by the agent. This adaptive, risk-aware approach allows cloud providers to maximize resource utilization while mitigating the critical risk of performance degradation for their customers.

Experimental Setup & Results

▶ Watch: Challenges: competing objectives and varying demands (7:00)

The evaluation of ProtoRAIL was conducted using rigorous real-world data, underscoring its practical applicability and effectiveness.

Datasets:

• Microsoft Internal Workloads: The primary evaluation leveraged proprietary, real-world data from Microsoft's internal cloud infrastructure. This dataset encompassed usage telemetry from two distinct geographical regions, spanning over 300 clusters, and collected over a period of two weeks. This extensive and realistic dataset provides a robust testbed for ProtoRAIL, reflecting the actual complexities and variability of cloud operations.
• Alternate Dataset (Airline Data): To further validate the generalizability and conceptual soundness of ProtoRAIL, an alternate dataset was utilized. This data was sourced from the US Department of Transportation website, focusing on airline overbooking statistics. This dataset served as an external validation, drawing a parallel to the classic overbooking problem that inspired the foundational concept of cloud oversubscription.

Baselines:

While the talk implies comparison against existing industry practices and research (e.g., naive statistical models, general RL methods, and other offline approaches like LSTM-based methods or Coin), specific named baselines for direct quantitative comparison in the results section were not explicitly detailed. However, the consistent assertion that "ProtoRAIL is a clear winner" suggests a strong comparative performance against contemporary or baseline approaches to vCPU oversubscription.

Hardware:

The specific hardware (e.g., GPU types, server configurations) used for training and inference of the ProtoRAIL agent was not detailed in the transcript. This is a common omission in high-level conference talks that focus on the algorithmic and systemic innovation rather than the specific compute infrastructure.

Metrics:

The evaluation focused on two primary, often competing, objectives:

• Risk: Measured by the incidence of "overloaded physical machines," which translates to resource contention, performance jitter, and potential SLA violations. The goal is to minimize this.
• Benefit: Quantified by "harvested cores" or, more concretely, "sellable core hours." This metric directly reflects the increase in resource utilization and the subsequent reduction in COGS and increase in revenue margins for the cloud provider. The goal is to maximize this.
• Benefit-to-Risk Ratio: This composite metric provides a holistic view of the agent's effectiveness in balancing the two competing objectives.

Headline Numbers & Results:

The results presented unequivocally positioned ProtoRAIL as a superior solution:

• Clear Winner: ProtoRAIL demonstrated a "clear winner" status across both dimensions of risk and benefit. This indicates that it effectively minimizes overloading risks while simultaneously maximizing resource harvesting.
• Core Hour Savings: On actual partial deployments within Microsoft's limited number of clusters, ProtoRAIL yielded significant "savings we get in core hours." While specific percentage or absolute numbers were not disclosed in the talk, the qualitative description implies substantial gains.
• Improved Benefit-to-Risk Ratio: The system achieved a higher "benefit to risk ratio" on these deployments, confirming its ability to make intelligent trade-offs.
• Interpretable Prototypes: The prototypes learned by ProtoRAIL were highlighted as "very interpretable." This allows operators to "drill down on the prototypes and find out what these pattern of usages look like," enhancing trust and debuggability.
• KITL Efficiency: The Knowledge-in-the-Loop (KITL) module was shown to be highly efficient, utilizing "minimal queries, like six or seven queries are used at most in a training run." This low query count ensures that incorporating domain knowledge does not become a bottleneck for the training process. Furthermore, KITL actively refines the prototypes, as evidenced by one prototype being "gone because that wasn't contributing much," indicating an adaptive and optimized set of learned patterns.

In summary, the experimental results from real-world Microsoft cloud environments validate ProtoRAIL's capability to deliver substantial improvements in cloud resource efficiency and safety, demonstrating its readiness for broader deployment.

Practical Implications

▶ Watch: Existing solutions: Azure Resource Central platform (8:00)

ProtoRAIL's development and successful evaluation carry significant practical implications for various stakeholders within the cloud computing ecosystem, from infrastructure teams to model builders and deployers.

For cloud infrastructure teams and practitioners, ProtoRAIL offers an adaptive, risk-aware, intelligent, and multi-granular oversubscription framework. This translates directly into tangible operational and financial benefits. The ability to dynamically adjust the oversubscription ratio (zeta) based on real-time telemetry and learned usage patterns means that physical machines can be packed more densely without compromising performance. This leads to substantially higher sellable core hours, directly impacting the cloud provider's bottom line by reducing Cost of Goods Sold (COGS) and increasing revenue margins. The explicit focus on minimized overloading risk is crucial, as avoiding resource contention and performance jitter is paramount for customer satisfaction and adherence to Service Level Agreements (SLAs). The speaker's analogy to "plumbers of this world" highlights the fundamental role of such solutions in maintaining the smooth operation of complex cloud businesses.

For model builders and deployers, ProtoRAIL showcases a powerful application of prototypical imitation learning. The approach of learning from pseudo ground truth derived from telemetry, rather than relying on difficult-to-design reward functions (as in RL) or perfectly labeled data, provides a robust methodology for resource optimization problems. The Knowledge-in-the-Loop (KITL) module is a particularly valuable innovation, demonstrating how to effectively inject domain expertise and safety constraints into an ML model with minimal overhead ("six or seven queries at most"). This makes the deployment of such critical systems safer and more trustworthy, bridging the gap between theoretical ML models and production-grade reliability. The interpretability offered by the learned prototypes also aids in debugging, understanding model decisions, and building confidence in the automated system.

Tradeoffs and Limitations:

While ProtoRAIL presents a compelling solution, understanding its tradeoffs and current limitations is important:

• Risk vs. Benefit: The system is explicitly designed to balance this tradeoff. While it minimizes risk, an inherent tension remains between maximizing utilization and guaranteeing absolute performance in all edge cases. The definition of "acceptable risk" is a business decision that ProtoRAIL helps manage quantitatively.
• Opaque Allocation: The oversubscription is "opaque" to the customer. While agreed upon in terms of maximum requested cores, the dynamic assignment of physical resources below that maximum carries an implicit risk. As highlighted in the Q&A, if a customer's workload temporarily exceeds the currently assigned (oversubscribed) resources and the physical machine cannot provide more, that constitutes a "risk" that ProtoRAIL aims to minimize.
• Reliance on VM Scheduler: ProtoRAIL operates under the assumption that the VM scheduler is performing its job fairly and efficiently in terms of initial placement. Any inefficiencies or biases in the underlying scheduler could potentially impact the effectiveness of ProtoRAIL's zeta optimization.
• Generalizability Beyond vCPU: While the authors claim generalizability to "any resource optimization problem," the current demonstration is focused on vCPU. Applying it to other resource types (e.g., memory, network I/O, storage) might require adapting the telemetry, prototype definitions, and domain knowledge.
• Real-time Adaptation: The talk mentions "adaptive choice" and "varying demands," implying responsiveness. However, the exact latency of adaptation and how frequently the policies are re-evaluated or models retrained are not explicitly detailed. This is crucial for highly dynamic cloud environments.
• No Ground Truth: The reliance on pseudo ground truth for behavior cloning, while practical, means the system learns from historical patterns that may not always represent the absolute ideal oversubscription. KITL helps mitigate this by incorporating safety, but the inherent lack of perfect ground truth is a persistent challenge in this domain.

Despite these considerations, ProtoRAIL's successful testing on real Microsoft internal workloads and its progression "on the path to production" signify its maturity and readiness to deliver substantial value in enhancing cloud efficiency and profitability.

Key Takeaways

• Critical Cloud Efficiency Problem: ProtoRAIL addresses the fundamental challenge of vCPU oversubscription in cloud computing, a vital mechanism for reducing Cost of Goods Sold (COGS) and increasing revenue margins for cloud providers.
• Novel Prototypical Imitation Learning: The system employs a unique prototypical imitation learning approach that learns dynamic oversubscription policies from usage telemetry, moving beyond static heuristics and the practical limitations of traditional Reinforcement Learning.
• Risk-Cognizant and Safe by Design: ProtoRAIL explicitly manages the competing objectives of maximizing resource utilization and minimizing the risk of physical machine overloading, integrating Knowledge-in-the-Loop (KITL) to incorporate domain knowledge and enhance safety with minimal queries.
• Multi-Granular and Interpretable: The use of prototypes effectively solves the problem of decision granularity, allowing the system to tailor policies to specific workload patterns, while also making the learned behaviors highly interpretable for operators.
• Demonstrated Real-World Impact: Evaluated on extensive Microsoft internal workloads, ProtoRAIL has shown to be a "clear winner," delivering substantially higher sellable core hours and an improved benefit-to-risk ratio compared to existing methods.
• Production-Ready and Generalizable: The solution is on the "path to production" within Microsoft and its core methodology is designed to be applicable to a broad range of resource optimization problems beyond just vCPU oversubscription.

About the Speaker(s)

The primary speaker for this talk was Mayukh Das, who holds the position of Senior Researcher at Microsoft, specifically within the MC65 Research team. The work presented, ProtoRAIL, is a collaborative effort, with co-authors including Lu Wang, Fangkai Yang, Íñigo Goiri, Saravan Rajmohan, and Dongmei Zhang, all implied to be part of Microsoft's research or engineering teams given the context of the talk and the "Microsoft internal workloads" used for evaluation. Their collective expertise focuses on AI operations (AI Ops) for cloud efficiency, emphasizing practical, impactful research that keeps cloud infrastructure running smoothly.