Scaling TraceWin beam dynamics simulations on Kubernetes for Reinforcement Learning training

Reinforcement learning is emerging as a viable technology to implement autonomous beam dynamics setup and optimization in particle accelerators. A Deep Learning agent can be trained to efficiently explore the parameter space of an accelerator control system and converge to the optimal beam setup much faster than traditional methods. Training these models requires programmatic execution of a high volume of simulations. This contribution introduces pytracewin, a Python wrapper of the TraceWin beam dynamics simulator, which exposes simple methods to run simulations and retrieve results. It can be easily combined with the large Python ecosystem of machine learning libraries to develop optimization models. Still, the training process is computationally constrained by the number of simulations that can be run in a reasonable time. It is thus crucial to scale such workload on a dedicated computing infrastructure while retaining a simple high-level user interface. We exploit Ray, an opensource library, to enable embarrassingly parallel execution of TraceWin simulations on Kubernetes (K8s), using a dynamically scalable number of workers and requiring minimal user code modifications. Workers are instantiated with a custom docker image combining Ray and pytracewin. The approach is validated using two K8s clusters on INFN Cloud to simulate the ADIGE (Acceleratore Di Ioni a Grande Carica Esotici) beam line at Legnaro National Laboratories.