Mistill: Distilling Distributed Network Protocols from Examples

New applications and use-cases in data center networks require the design of Traffic Engineering ( TE ) algorithms that account for application-specific traffic patterns . TE makes forwarding decisions from the global state of the network . Thus , new TE algorithms require the design and implementation of effective information exchange and efficient algorithms to compute forwarding decisions . This is a challenging and labor and time-intensive process . To automate and simplify this process , we propose Mistill . Mistill distills the forwarding behavior of TE policies from exemplary forwarding decisions into a Neural Network . Mistill learns how to process local state to send it over the network , which network devices must exchange state with each other , and how to map the exchanged state into forwarding decisions . We show the abilities of Mistill by learning three exemplary policies and verify their performance in simulations on synthetic and real world traffic patterns . We show that the learned protocols closely implement the desired policies , and generalize to unseen traffic patterns . 1 INTRODUCTION . Data center workloads are diverse and range from client-server applications , machine learning , and web-traffic to high-performance computing applications ( Benson et al . ( 2010 ) ; Roy et al . ( 2015 ) ; Zhang et al . ( 2017 ) ) . Each workload has different requirements towards Traffic Engineering ( TE ) . For example , optimizing the completion time of individual flows is desirable for client-server applications ( Dukkipati & McKeown ( 2006 ) ) , but leads to decreased performance in machine learning workloads ( Chowdhury et al . ( 2014 ) ) . To achieve the best performance , each workload needs a TE scheme that is tailored towards its specific requirements . Indeed , new applications regularly trigger the development of new TE schemes in the networking community ( Katta et al . ( 2016 ) ; Chowdhury et al . ( 2014 ) ; Chen et al . ( 2018 ) ; Benson et al . ( 2011 ) ) . Obtaining a deployable distributed protocol from an initial TE design is challenging . A design includes the specification of update messages , the processing of update messages , and the algorithm to compute the forwarding decisions that result in the desired policy ( Katta et al . ( 2016 ) ) . The design of each component is challenging since the final protocol must meet line-rate , exchange information at ( sub ) millisecond granularity , and cope with limited computational requirements . All steps must be efficient , and the design holistic , requiring the exploitation of patterns in the network and the traffic . For instance , CONGA integrates a substantial part of the calculation for the forwarding decision into the exchange of update messages ( Hsu et al . ( 2020 ) ) . The complexity makes customizing TE schemes hardly feasible for smaller enterprises , potentially resulting in competitive disadvantages . To automate and simplify the translation of a TE policy to a distributed protocol , we propose Mistill . Mistill distills the forwarding behavior of a TE policy from exemplary forwarding decisions into a Neural Network ( NN ) . Mistill learns the encoding of the switch local state in update messages , the exchange of update messages , and the computation of forwarding decisions from update messages with Machine Learning ( ML ) . Mistill removes the need to manually design update messages , information exchange , and the calculation of forwarding decisions . ML allows Mistill to automatically detect helpful patterns in traffic and network and exploit the patterns to learn a protocol . Further , our NN design makes the learned patterns accessible to humans and can thus aid the manual design process . Our main contributions in this work are : 1 ) The design of an ML model that can learn forwarding behavior , makes learned patterns accessible to humans , and has the potential to be deployed on real hardware . 2 ) We show in simulations that Mistill can learn the forwarding behavior of three TE policies and generalizes to previously unseen traffic patterns . 3 ) We analyze the learned messages and their exchange . We find that Mistill can learn distributed protocols that closely implement the desired policy and are robust towards distributional shift in the input . We analyze the learned messages ’ content , and show that the NN learns an information exchange resembling an edge-cover , which can be used to optimize the information exchange of other TE schemes . The document is organized as follows . Sec . 2 describes the application scenarios and the requirements towards ML models . Sec . 3 describes the neural architecture and justifies design choices . Sec . 4 describes experiments and presents results and visualizations of the learned protocols . Sec . 5 discusses related work , and Sec . 6 concludes the paper . 2 APPLICATION SCENARIO . Mistill learns a distributed protocol from examples . Fig . 1 illustrates the process and expected outcome . The first component is a centralized implementation of a TE policy . For example , a packet or flow level simulation using an Integer Linear Program or heuristic algorithm to compute forwarding decisions based on the entire network state . The forwarding decisions specify how a switch forwards traffic to its neighbors , given the current state of the network . Usually , new TE schemes are evaluated in such a setting first to establish their performance and correctness . Here , Mistill hooks in and uses the forwarding decisions to train an ML model , i.e. , a NN . Mistill trains the model offline , i.e. , not on the network devices . The NN learns three aspects : 1 ) how to encode local state in update messages , 2 ) how to exchange the update messages , 3 ) how to map the exchanged state into forwarding decisions . After training , Mistill deploys the NN to switches , e.g. , with P4 ( Siracusano & Bifulco ( 2018 ) ) . All switches in the network have the same NN and use it in inference mode , i.e. , do forward passes only . The NN distinguishes switches through an identifier . During training , the NN learns how to condition the forwarding decisions on the identifiers , thus learning the correct behavior for each switch . To compute forwarding decisions , switches do not need to know the network topology . This work does not try to outperform existing TE strategies in their specific application scenarios . Instead , this work aims to facilitate the automatic generation of novel TE policies , taking direct advantage of switch architectures like Taurus ( Swamy et al . ( 2020 ) ) without asking network operators to take the time or acquire the expertise to design highly optimized protocols and implement them , e.g. , in P4 ( Bosshart et al . ( 2014 ) ) . 2.1 DATA CENTER NETWORK REQUIREMENTS . The goal of Mistill is to distill network protocols for data-center networks from exemplary data . We focus on Clos topologies due to their popularity and simple structure ( Zhang et al . ( 2019 ) ) . See Appendix B for details . TE protocols in data-centers exchange update messages at ( sub- ) millisecond scale ( Katta et al . ( 2016 ) ) . Thus , the updates must be small to keep the resulting overhead low . To effectively utilize the frequent updates , switches must make forwarding decisions frequently , e.g. , for each new flow . Thus , the decision logic should reside in the data-plane ( Alizadeh et al . ( 2014 ) ) . This limits the computational demand the decision logic can express towards the switch . The learned decision logic must thus be computationally efficient to meet the limited computational capacities . 2.2 WHY NEURAL NETWORKS ? . Mistill uses a NN to learn a distributed protocol , which is a good choice for three reasons . 1 ) NNs are general function approximators ( Glorot et al . ( 2011 ) ) . This is important because forwarding decisions can depend on the network state in a non-linear way . For example , minimizing the maximum link utilization results in forwarding decisions that do not depend linearly on the network state . A small change in the utilization can lead to a sudden change in the forwarding decision . 2 ) NNs can be executed with the limited computational resources in the data-plane ( Sanvito et al . ( 2018 ) ; Swamy et al . ( 2020 ) ) . Here , the choice of activation function is important . The ReLU activation requires only a check for negativity ( Glorot et al . ( 2011 ) ) . A NN with ReLU activations uses addition , multiplication and a check for negative values . Further , NNs can be quantized to integer or binary computations ( Gholami et al . ( 2021 ) ) . 3 ) The execution of NNs is easy to parallelize and accelerators for NN inference already exist for network equipment ( Luinaud et al . ( 2020 ) ; Swamy et al . ( 2020 ) ) . In addition , research efforts show that NNs can be executed in the data-plane at line rate already on today ’ s ASICs ( Siracusano & Bifulco ( 2018 ) ) , and on end-hosts with SmartNICs ( Sanvito et al . ( 2018 ) ) . We thus believe that switches will be increasingly capable of executing NNs at line rate . 3 NEURAL ARCHITECTURE . The NN architecture consists of four modules : the Localization Module ( LOCMOD ) , the HNSA Module ( HNSAMOD ) , the Network State Module ( NSMOD ) , and the Forwarding Module ( FWD MOD ) . Fig . 2 shows how the modules interact and represents inputs and outputs with hashes . See Appendix C for sketches of technical realizations of Mistill . Output . The NN has three outputs : the out-port , Hidden Node State Advertisements ( HNSAs ) 1 , i.e. , update messages that inform other network devices about the local state of the sending device , and State Interest ( SI ) , signaling which devices need to exchange information with each other . Input . The NN has four inputs : LinkState , AllNodes , Destination and Location . LinkState is the switch local state , e.g. , the availability and utilization of incident links . The NN transforms the LinkState into HNSAs . Location is the identifier of a switch , Destination the destination of a flow , and AllNodes the identifiers of all switches . The NN uses Location , Destination , and AllNodes to learn how to condition forwarding decisions on a specific switch and destination and which switches ’ HNSAs are needed to make a forwarding decision . Training . The NN is trained end-to-end . Training end-to-end tailors HNSAs and SI to a forwarding policy . The loss function depends on the forwarding objective . We interpret objectives as multi-label classification or as learning of a multinomial distribution . We use : Lmn = − ∑N i=1 p ( xi ) log p ( yi ) as loss function for the multinomial distribution , and Lml = − ∑N i=1 ( p ( xi ) log p ( yi ) + ( 1− p ( xi ) ) log ( 1− p ( yi ) ) ) as loss function for the multi-label classification . N corresponds to the number of outputs , p ( xi ) is the ground truth and p ( yi ) is the prediction of the NN . The loss Lmn can learn distributional targets to split traffic over adjacent nodes . Lml can learn policies that have one or multiple equally good paths . Changes in the network that lead to changes in the forwarding policy require a re-training of the model . For example , new networking hardware , or changes to the physical network topology . As we will show , re-training is short compared to the deployment process of new hardware , and can be integrated into the roll-out process . Inference . After training finishes , the NN is deployed to the switches . The switches use the NN for three tasks : 1 ) to compute HNSAs from their local state , 2 ) to select the HNSAs necessary to calculate the out-port , and 3 ) to calculate the out-port . The switches do not change the weights of the NN . The switch also never execute the full NN . To compute HNSAs , the switch executes the HNSAMOD . To compute the out-port , the switch executes the LOCMOD , NSMOD and FWD MOD , but not the HNSAMOD . The switch uses the HNSAs it received from other nodes instead .

Summary: The paper proposes an ML-based approach Mistill to help automatically deploy distributed protocol from a given TE policy. The approach learns forwarding decisions together with intermediate results such as exchange information and LinkState encoding from exemplary policies. The resulting network can later be deployed in switches to encode link-state and compute forwarding decisions. The technical contributions are the following. First, they leverage a reparameterization trick to handle categorical distributions when encoding link states into a binary vector. Second, they use the scaled dot-product attention method to help switches learn to exploit the encodings mentioned above and make the forwarding decision. In the evaluation part, the paper compares with three custom baselines, i.e., LCP, MinMax, WCMP, and exceeds all of them in terms of corresponding metrics.