JAX-based high-performance discrete event simulation method and system

Abstract

The invention provides a JAX-based high-performance discrete event simulation method and system, and the method comprises the steps: firstly completing environment configuration and initialization through an immutable data structure, and generating an EnvParas object with fixed static parameters; secondly, processing event logic by adopting a pure function steepv, and realizing side-effect-free functional event processing and state iteration by combining JAX condition control primitive and lax.scan iterative loop; then, a batch processing task is generated through job packaging, and job-level parallel acceleration is achieved on multiple GPU / TPU devices by means of jax. Jit compiling and jax. Vmap vectorization; and finally, packaging a standardized AI interactive interface, supporting external algorithm calling and finishing result post-processing. Functional programming and JAX hardware acceleration are combined, the simulation speed is increased, meanwhile, the event dynamic accuracy is guaranteed, reinforcement learning large-scale sampling can be seamlessly supported, and the method is suitable for AI application scenes in the field of operational research.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of discrete event simulation, in particular to a high-performance discrete event simulation method and system based on JAX. BACKGROUND

[0002] Discrete event simulation is a core tool for studying complex dynamic systems such as queuing systems, supply chain networks, and traffic flow, which simulates system behavior by processing discrete events in chronological order, providing data support for system optimization and decision-making. Currently, traditional discrete event simulators implemented in Python (such as Salabim) have significant performance bottlenecks and compatibility issues in parallel execution, GPU compatibility, and AI integration.

[0003] To solve the parallelization problem, existing parallel discrete event simulation (PDES) methods, such as PDES based on multi-core CPUs or distributed systems. However, this method has poor compatibility in the Python environment and is difficult to seamlessly integrate with AI workflows such as reinforcement learning.

[0004] Therefore, there is an urgent need in the market for a high-performance discrete event simulation method and system based on JAX that can balance Python ecosystem compatibility, high-performance parallel computing capabilities, and seamless integration with AI workflows. SUMMARY

[0005] In view of the defects in the prior art, the purpose of the present application is to provide a high-performance discrete event simulation method and system based on JAX.

[0006] According to the high-performance discrete event simulation system based on JAX provided by the present application, it comprises a discrete event simulation engine module, a JAX acceleration module and a standardized external interface module; The discrete event simulation engine module obtains the simulation parameters provided by the user or the upper application and the system state data generated dynamically in the simulation process, and obtains the updated state data and event trajectory data after scheduling and processing the events; The JAX acceleration module performs hardware-accelerated parallel computing according to the updated state data and event trajectory data, and outputs the optimized and parallelized execution results to the discrete event simulation engine module to update the simulation state in synchronization, while also being called by the standardized external interface module to complete efficient environment stepping operations; The standardized external interface module obtains the bottom layer state data provided by the discrete event simulation engine module and the accelerated execution results provided by the JAX acceleration module, generates observation, reward and termination signals conforming to the standard of machine learning methods, and is called by the external AI algorithm framework to perform policy training or reasoning.

[0007] Preferably, the discrete event simulation engine module comprises an environment configuration submodule, an event scheduling submodule, an event resolution submodule and a job-level tracking submodule; The environment configuration submodule is configured to define and solidify the static parameters of the simulation environment; The event scheduling submodule is responsible for driving the simulation forward according to the current state; The event resolution submodule is configured to implement side-effect-free event logic; The job-level tracking submodule is configured to encapsulate and manage multiple parallel independent simulation instances.

[0008] Preferably, the environment configuration submodule uses immutable data classes to obtain user input configuration parameters, generates an EnvParams parameter structure, and provides the event scheduling submodule and the event resolution submodule with the initial simulation environment; The event scheduling submodule generates a new immutable EnvState object by calling the core step_env function according to the EnvParams parameter structure and the current EnvState, outputs new state data, and provides the event resolution submodule with the next event type to be processed and the job-level tracking submodule with the state of each independent simulation in batch processing mode; The event resolution submodule compares the occurrence times of customer arrival events and service completion events using JAX conditional control primitives based on the new state data, resolves and returns the updated state, and provides the event scheduling submodule with the single-step simulation loop, which defines the function logic for the JAX acceleration module to apply the lax.scan primitive for loop iteration; The job-level tracking submodule obtains a batch of random seeds and a single EnvParams specified by the user, generates job batch data containing multiple independent simulation initial states, and provides the JAX acceleration module with vectorized parallel execution and the standardized external interface module with batch observation data for multiple-instance environments.

[0009] Preferably, the configuration parameters include the total simulation duration, server processing time distribution, and customer arrival time interval; The EnvState includes the number of customers in the queue, the busy state of each server, and the current simulation clock.

[0010] Preferably, the JAX acceleration module comprises an instant compilation submodule, an automatic vectorization submodule and a multi-device support submodule; The instant compilation submodule is configured to convert Python functions into highly optimized machine code; The automatic vectorization submodule is configured to apply operations for a single instance to multiple instances in batches; The multi-device support submodule is configured to manage and schedule the computing hardware resources.

[0011] Preferably, the just-in-time compilation submodule obtains the XLA (Accelerated Linear Algebra) optimized code generated by the vectorized simulation function compilation processed by the automatic vectorization submodule, and sends the XLA (Accelerated Linear Algebra) optimized code to the multi-device support submodule for execution on the specified hardware. The automatic vectorization submodule obtains the single-instance simulation function defined by the discrete event simulation engine module and the batch data generated by the job-level tracking submodule, and generates a new vectorized simulation function by applying a jax.vmap transformation. The multi-device support submodule obtains the optimized code generated by the just-in-time compilation submodule and the vectorized data array prepared by the automatic vectorization submodule, executes the computation on the GPU or TPU, and returns the final result, which is sent to the discrete event simulation engine module to update the state of all simulations in the batch.

[0012] Preferably, the standardized external interface module includes a state observation conversion submodule, a reward calculation submodule, and a batch adaptation and demultiplexing submodule. The state observation conversion submodule is configured to convert the internal complex simulation state into a numerical observation space understandable by the AI algorithm. The reward calculation submodule is configured to evaluate the pros and cons of the state change quantitative simulation result. The batch adaptation and demultiplexing submodule is configured to bridge the external single-step operation and the internal parallel computation.

[0013] Preferably, the state observation conversion submodule obtains the EnvState object output by the discrete event simulation engine module, extracts and flattens the key information according to the predefined rules to generate an observation vector, and the observation vector is called by the batch adaptation and demultiplexing submodule to be returned as part of the standard step function. The reward calculation submodule obtains the EnvState object of the previous step and the EnvState of the current step, calculates the scalar reward value by comparing the state changes, and the scalar reward value is called by the batch adaptation and demultiplexing submodule to evaluate the pros and cons of the previous action. The batch adaptation and demultiplexing submodule obtains a single action from the external AI algorithm and obtains the execution result of the entire batch from the JAX acceleration module, wraps or injects the single action into the batch process, and after receiving the batch computation result, extracts the state, observation, and reward of the corresponding single instance, and the single-instance tuple generated after demultiplexing is directly returned to the caller of the external AI algorithm, thereby completing a standard environment interaction.

[0014] The application provides a high-performance discrete event simulation method based on JAX, which is realized through cooperation of a discrete-time simulation engine module, a JAX acceleration module and a standardized external interface module, and comprises the following steps: Step S1: all static parameters of a simulation environment are defined by using an immutable data structure through an environment configuration submodule in the discrete-time simulation engine module, and are solidified into an EnvParams object; the EnvParams object contains constant parameters of a simulation scene, a parameter value input by a user is received to instantiate the EnvParams object, and the EnvParams object is transmitted to subsequent modules as a static context, used for initializing a simulation starting state EnvState and participating in calculation as an invariable basis in the whole simulation process; Step S2: a pure function step_env is used to process the logic of a single event through an event scheduling submodule and an event analysis submodule in the discrete-time simulation engine module, the step_env function receives a current EnvState and an EnvParams as input and returns a brand-new immutable EnvState, event logic judgment is realized by using a conditional control primitive of JAX, a system dynamic state is modeled as a JAX array and is updated by creating a new array, and an iteration loop is constructed based on lax.scan, the step_env function is repeatedly called, the EnvState output in the last step is used as input in the next step to promote the simulation, and a final EnvState is obtained; Step S3: single and serial simulation logic defined in step S2 is promoted to large-scale parallel execution through cooperation of a job-level tracking submodule of the discrete-time simulation engine module; Step S4: parallel simulation capability of step S3 is encapsulated into an environment interaction interface conforming to an AI field standard through a standardized interface module, an external AI algorithm calling interface triggers the parallel execution process of step S3, batch results are received and analyzed into a data structure required by an AI framework, meanwhile, data post-processing and visualization functions are integrated, a key performance indicator is calculated, and a statistical chart is generated.

[0015] Preferably, step S3 comprises the following steps: Step S3.1: a pseudo-random number generator of JAX is used to assign a unique random seed to each simulation instance through the job-level tracking submodule of the discrete-time simulation engine module, all simulation instances share the same EnvParams object, and independent simulation tasks are organized into a batch data structure;

[0016] Step S3.2: compile the simulation iteration function constructed based on lax.scan in step S2 into XLA optimized code by the just-in-time compilation submodule of the JAX acceleration module using the jax.jit decorator; and transform the compiled function to process the batch data generated in step S3.1 by the automatic vectorization submodule using jax.vmap;

[0017] Step S3.3: automatically distribute the batch data processed by the parallelized function transformed by vmap to multiple computing hardware for parallel execution by the multi-device support submodule of the JAX acceleration module, and return the batch data containing the results of all simulation instances.

[0018] Compared with the prior art, the present application has the following beneficial effects: 1、The present application overcomes the parallelization bottleneck and side effect problems caused by variable state management and sequential execution logic in traditional discrete event simulation by using immutable data structures, thereby significantly improving the compatibility of the system with the JAX framework and the parallel computing efficiency.

[0019] 2、The present application solves the performance deficiency problem of traditional simulators when processing large-scale, high-concurrency simulation tasks by integrating just-in-time compilation (JIT) and automatic vectorization technology in the JAX framework, and realizes seamless support and acceleration of heterogeneous hardware such as GPU / TPU.

[0020] 3、The present application solves the interface mismatch problem between the pure functional, batch processing simulation core and the standard reinforcement learning framework (usually stateful, single-step interaction mode) by designing an adaptive layer that efficiently cooperates with the JAX backend, and realizes high-performance, scalable environment interaction. BRIEF DESCRIPTION OF DRAWINGS

[0021] Other features, objects and advantages of the present application will become more apparent after reading the following detailed description of non-limiting embodiments with reference to the accompanying drawings: Figure 1 The working method flowchart of the present application is shown in the figure; Figure 2 The queuing system model based on discrete event simulation technology (DES) of the M / M / 1 single-server model in the present application; Figure 3 The queuing system model based on discrete event simulation technology (DES) of the M / M / C multi-server model in the present application; Figure 4 The structural diagram of the above two typical simulation scenarios, (a) represents the M / M / 1 single-server model, and (b) represents the M / M / C multi-server model. DETAILED DESCRIPTION

[0022] The application will be described in detail below in conjunction with specific embodiments. The following examples will help those skilled in the art to further understand the application, but do not limit the application in any form. It should be noted that those skilled in the art can make several changes and improvements without departing from the concept of the application. These are within the scope of the present application.

[0023] The application uses JAX parallel framework to build a high-performance discrete event simulation engine DES-Gymnax of native Python, which realizes parallelization and efficient execution of the simulation process by integrating JAX's just-in-time compilation (JIT), automatic vectorization and GPU acceleration functions. The core idea is to convert the sequential execution model of traditional discrete event simulation into a functional programming paradigm, using immutable data structures and conditional branches to handle event scheduling and parsing, while achieving large-scale parallel simulation through job-level tracking. The invention also provides a Gym-like API interface for seamless integration with reinforcement learning algorithms, supporting AI applications in the field of operations research.

[0024] The application first completes environment configuration and initialization through an immutable data structure to generate an EnvParams object with solidified static parameters; secondly, it uses a pure function step_env to handle event logic, combining JAX conditional control primitives and lax.scan iteration loops to achieve side-effect-free functional event processing and state iteration; then it generates batch processing tasks through job encapsulation, and uses jax.jit compilation and jax.vmap vectorization to achieve job-level parallel acceleration on multiple GPU / TPU devices; finally, it encapsulates a standardized AI interaction interface to support external algorithm calls and complete result post-processing. Its functional process scheduling based on the JAX parallel framework, job-level tracking parallelization, and Gym-like API structure achieve a simulation speed improvement of 10x to 100x compared to traditional Python-based discrete event simulators, accurate maintenance of event-driven dynamics, and seamless support for large-scale sampling of reinforcement learning.

[0025] Figure 2 and Figure 3 respectively show two typical queuing system models based on discrete event simulation technology (DES), where Figure 2 is an M / M / 1 single-server model, Figure 3 is an M / M / C multi-server model. In the latter, the parameter C is set to 4 to represent the existence of 4 parallel service windows in the system. Figure 4 Further structural diagrams of the above two typical simulation scenarios are given.

[0026] In terms of performance testing, experiments were conducted for small-scale and large-scale scenarios respectively: (1) Small-scale tests (less than 10 Workers) are performed in a computing environment configured with 16 GB of RAM and an 8-core CPU; (2) Large-scale tests (20 or more Workers) are completed on a server equipped with a 24-core CPU, 128 GB of memory, and an Nvidia A100 graphics card with 80 GB of video memory, as they involve GPU acceleration.

[0027] Each test is repeated 5 times, measuring the actual computation time required to simulate a system running for 100,000 seconds, and taking the average of the 5 results as the final evaluation index.

[0028] The test results show that the high-performance discrete event simulation system based on JAX proposed by the present application has a significant advantage in execution efficiency compared to the existing Python-based discrete event simulator Salabim. For example, in the test for the M / M / 1 model, Salabim single-threaded running simulation takes 3.0455 seconds, while the system of the present application only needs 0.0346 seconds to complete the same event simulation calculation under the condition of enabling 100 parallel Workers and using GPU acceleration. It is worth noting that the 0.0346 seconds is the average single simulation average time after taking the average of the total running time of 100 parallel simulation tasks (i.e. total time divided by 100), further demonstrating the superiority of the system in parallelization and acceleration performance.

[0029] According to the high-performance discrete event simulation system based on JAX provided by the present application, as shown in Figure 1 The discrete event simulation engine module is responsible for the execution of the core simulation logic, including the scheduling and processing of events. This module obtains simulation parameters provided by the user or the upper application (e.g. customer arrival rate, device service rate) and system state data generated dynamically during simulation (e.g. current queue length, simulation clock). The updated state data and event trajectory data output by this module are called by the JAX acceleration module to perform hardware-accelerated parallel computation, and by the standardized external interface module to generate observations and rewards that conform to the standards of machine learning methods such as reinforcement learning.

[0030] The discrete event simulation engine module includes an environment configuration submodule, an event scheduling submodule, an event resolution submodule, and a job-level tracking submodule. The environment configuration submodule is used to define and solidify the static parameters of the simulation environment. The submodule uses an immutable data class to obtain user input configuration parameters (e.g., simulation total duration, server processing time distribution, customer arrival time interval), and generates an EnvParams parameter structure that is called by the event scheduling submodule and the event resolution submodule to complete the initialization of the simulation environment.

[0031] The event scheduling submodule is responsible for driving the simulation forward according to the current state. The submodule obtains the EnvParams generated by the environment configuration submodule and the current EnvState (e.g., the number of customers in the queue, the busy state of each server, the current simulation clock), generates a completely new, immutable EnvState object by calling the core step_env function, and outputs the new state data to the event resolution submodule to determine the type of the next event to be processed, and to the job-level tracking submodule to manage the state of each independent simulation in batch processing mode.

[0032] The event resolution submodule is used to implement side-effect-free event logic judgment. The submodule obtains the state data and timestamp information output by the event scheduling submodule, compares the occurrence times of customer arrival events and service completion events using JAX conditional control primitives (e.g., lax.cond), and resolves and returns the updated state, which is called by the event scheduling submodule to complete a single-step simulation loop. At the same time, the function logic defined by the submodule is called by the JAX acceleration module to apply the lax.scan primitive for efficient loop iteration.

[0033] The job-level tracking submodule is used to encapsulate and manage multiple parallel independent simulation instances. The submodule obtains a batch of random seeds and a single EnvParams specified by the user, generates a job batch data containing multiple independent simulation initial states, which is called by the JAX acceleration module to achieve efficient vectorized parallel execution, and is called by the Gym-like API module to generate batch observation data of the multi-instance environment.

[0034] The JAX acceleration module is used to provide underlying computing acceleration for simulation operations. The module obtains the pure functional simulation logic and state data array provided by the discrete event simulation engine module, and outputs the optimized and parallelized execution results (e.g., a batch of simulation trajectories), which are returned to the discrete event simulation engine module to update the simulation state in synchronization, and are also called by the standardized external interface module to complete efficient environment stepping operations.

[0035] The JAX acceleration module includes a just-in-time compilation submodule, an automatic vectorization submodule, and a multi-device support submodule. The just-in-time compilation submodule is used to convert Python functions into highly optimized machine code, which takes the vectorized simulation function (e.g., vmapped_step_env) processed by the automatic vectorization submodule, compiles the generated XLA (Accelerated Linear Algebra) optimization code, and calls the multi-device support submodule to execute on the specified hardware, thereby greatly reducing the overhead of the Python interpreter.

[0036] The automatic vectorization submodule is used to apply operations for a single instance to multiple instances in batches, which takes the single-instance simulation function (e.g., step_env) defined by the discrete event simulation engine module and the batch data generated by the job-level tracking submodule, generates a vectorized new function by applying the jax.vmap transformation, and calls the just-in-time compilation submodule to compile the output of the vectorized function, and calls the standardized external interface module to implement parallel environment stepping with the output of the vectorized execution result.

[0037] The multi-device support submodule is used to manage and schedule computing hardware resources, which takes the optimization code generated by the just-in-time compilation submodule and the vectorized data array prepared by the automatic vectorization submodule, executes the final result returned after computing on the GPU or TPU, and passes it to the discrete event simulation engine module to update the state of all simulations in the batch, and to the Gym-like API module for efficient generation of batch samples.

[0038] The standardized external interface module is used to provide a standardized interface for interacting with external artificial intelligence algorithms, which takes the underlying state data provided by the discrete event simulation engine module and the accelerated execution result provided by the JAX acceleration module, generates Observations, Rewards, and Done signals, which are called by external AI algorithm frameworks (e.g., reinforcement learning libraries) for policy training or inference. The standardized external interface module includes a state observation conversion submodule, a reward calculation submodule, and a batch adaptation and demultiplexing submodule. The state observation conversion submodule is used to convert internal complex simulation states into numerical observation spaces understandable by AI algorithms, which takes the structured EnvState object output by the discrete event simulation engine module, extracts and flattens key information (such as queue length, waiting time) according to predefined rules to generate an observation vector, and calls the batch adaptation and demultiplexing submodule to return as part of the standard step function.

[0039] a reward calculation submodule, which quantifies the pros and cons of the simulation results according to the state changes, obtains the EnvState of the previous step and the EnvState of the current step, and calculates a scalar reward value (Reward) by comparing the state changes (for example, the reduction of the queue length, the shortening of the average waiting time of customers), and the generated reward value is called by the batch processing adaptation and demultiplexing submodule to evaluate the pros and cons of the action of the previous step. In actual deployment, the design of the scalar reward value function can be freely designed by the user according to the task scene and the business needs, such as designing a service efficiency-oriented reward for the queuing scene:

[0040] wherein, represents the scalar reward received by the system at the moment, represents the change in the number of customers served in a unit simulation period, represents the change in the queue length; and is a weight coefficient, which can be set by the user according to the business scene.

[0041] a batch processing adaptation and demultiplexing submodule, which bridges external single-step operations and internal parallel computing, obtains a single action (Action) from an external AI algorithm, and obtains the execution results of the entire batch from the JAX acceleration module, wraps or injects the single action into the batch processing flow, and after receiving the batch calculation results, accurately extracts the state, observation and reward of the corresponding single instance from the batch calculation results, and the demultiplexed single instance (observation, reward, termination signal) tuple is directly returned to the caller of the external AI algorithm, thereby completing a standard environment interaction.

[0042] According to the high-performance discrete event simulation method based on JAX provided by the application, the cooperation of the discrete time simulation engine module, the JAX acceleration module and the standardized external interface module is realized, wherein the discrete time simulation engine module is responsible for core event management and state update, the JAX acceleration module provides calculation optimization and parallel support, and the standardized external interface module provides a standardized interface; the data flow starts from the environment configuration, is transmitted into the simulation engine for event processing, is optimized and executed in parallel through the JAX acceleration module, and finally the observation data and workflow results are output by the API module, supporting the calling of external AI algorithms, and the specific steps include: Step S1: All static parameters of the simulation environment are defined by the environment configuration submodule in the discrete-time simulation engine module using an immutable data structure (e.g., dataclass in Python) and are solidified into an EnvParams object. Take the M / M / 1 queuing model as an example, this object can contain constant parameters such as maximum simulation duration, average service rate, average arrival rate, etc. This design of immutability is a prerequisite for adapting the JAX functional programming paradigm, which can avoid the failure of just-in-time (JIT) compilation or unintended side effects caused by state variability from the root. The EnvParams object contains constant parameters of the simulation scenario, and the instantiation of the EnvParams object is received by the user input parameter value, which is passed as a static context to the subsequent module for initializing the simulation starting state EnvState and participating in the calculation as an invariant throughout the simulation process.

[0043] The environment configuration submodule is configured to receive parameter values from the user (e.g., through configuration files or command line transmission), and to instantiate the EnvParams object accordingly. This EnvParams object is then passed as a static context to the subsequent event processing and scheduling module, which is used to initialize the starting state of the simulation (EnvState) on the one hand, and to participate in the calculation as an invariant throughout the simulation process on the other hand. This step ensures that the initialization of the simulation environment is independent of subsequent execution, supports flexible parameter adjustment (such as specifying the maximum time step, processing time, and arrival time through command line parameters), and flows the configuration data to the event processing module for step S2.

[0044] Step S2: The logic of a single event is handled by the event scheduling and event resolution sub-modules in the discrete-time simulation engine module, using a pure function step_env. The pure function step_env receives the current environment state EnvState and environment parameters EnvParams as inputs, and returns a brand new, immutable EnvState object as its only output. The logic of the event (e.g., comparing the order of customer arrival and service completion) is implemented through JAX's conditional control primitive (e.g., lax.cond), ensuring that the entire process has no side effects. The dynamic state of the system (e.g., event queue, system clock) is itself modeled as a JAX array, whose updates are done by creating new arrays rather than modifying the original. The entire simulation process is constructed as an efficient iteration loop based on lax.scan. The lax.scan function is configured to repeatedly call the step_env function and use the new EnvState output from the previous step as input for the next step, efficiently advancing the simulation while reducing Python loop overhead. This functional iteration approach not only naturally fits with the JIT compiler but also provides the necessary computation graph structure for subsequent automatic vectorization. The final output of this step is the final EnvState after a series of event processing, which flows to the JAX-accelerated module in Step S3 to achieve job-level parallelization and hardware acceleration for discrete-event simulation.

[0045] Step S3: The single, serial simulation logic defined in Step S2 is lifted to large-scale parallel execution by working with the job-level tracking sub-module of the discrete-time simulation engine module. The step S3 includes: Step S3.1: Job packaging and batch generation. The job-level tracking sub-module is responsible for generating a batch of independent simulation tasks for parallel discrete-event simulation jobs. It assigns a unique random seed to each simulation instance (Worker) to be executed using JAX's pseudo-random number generator (e.g., jax.random.split). All tasks share the same EnvParams object but have independent random processes due to different random seeds, ensuring the diversity and statistical independence of simulation results. These independent tasks are organized into a batch (Batch) data structure, waiting for parallel processing. Data flows from the EnvState and EnvParams of Step S2 into the batch data structure, and then flows to Step S3.2 after the batch data is generated. The workflow emphasizes independence to support scaling in high-worker scenarios.

[0046] Step S3.2: Just-in-time compilation and automatic vectorization. The just-in-time (JIT) submodule of the JAX acceleration module first compiles the core simulation iteration function constructed in step S2 based on lax.scan with the jax.jit decorator. This operation converts the Python function into highly optimized XLA (Accelerated Linear Algebra) code, greatly reducing the interpreter overhead at runtime. Subsequently, the automatic vectorization submodule transforms the compiled function with jax.vmap to enable it to directly process the batch data generated in step S3.1. vmap automatically maps the single-instance operations of the original function to single-instruction, multiple-data (SIMD) operations that are efficiently executed on the data batch. The data flow and workflow direct the compiled results to step S3.3, supporting batch-level optimization. Technical notes include monitoring execution time (such as evaluating the average batch and individual rollout time through time measurement functions).

[0047] Step S3.3: Multi-device collaborative execution. The parallelized function after vmap transformation can automatically distribute the batch data it processes to one or more available computing hardware (such as multiple GPU cores or TPU cores) for parallel execution by the multi-device support submodule of the JAX acceleration module. The JAX runtime system handles data distribution and result retrieval across devices in the background, completely transparent to the developer, without the need for complex MPI or multi-threaded programming. After execution is complete, the system returns a batch data containing the results of all simulation instances, such as the batched final observation state, action sequence, and reward sequence. The data flow returns to the API module of step S4. The workflow ensures transparent acceleration, supporting smooth expansion from CPU to GPU.

[0048] Step S4: achieved through a standardized interface module, which serves as an abstraction layer and adapter for interaction between the present application and external AI applications. This module encapsulates the parallel simulation capabilities of step S3 into a set of environment interaction interfaces conforming to mainstream artificial intelligence standards, including but not limited to reset_env() for initializing or resetting a batch of parallel simulation environments, step_env() for executing a batch of actions in all parallel environments and returning the next batch of observations, rewards, and completion flags, and get_obs() for extracting observation data meaningful to AI algorithms from the current batch of environment states. External reinforcement learning algorithms interact with the simulation environment by calling these standard APIs. For example, when the RL algorithm calls step_env(), the module internally triggers the complete parallel execution process of step S3, receives the returned batch results, parses and formats them into the standard data structures required by the RL framework (obs, reward, done, etc.), and returns them to the RL algorithm. To optimize AI training efficiency, the lax.stop_gradient instruction can be used to prevent gradient information from flowing into the simulator, as the simulation process itself is usually not differentiable. Result output and application: In addition to serving AI algorithms, this module can also integrate data post-processing and visualization functions. For example, based on the batch results returned by parallel simulation, it automatically calculates key performance indicators (KPIs) such as average waiting time and system throughput, and calls visualization libraries (such as Matplotlib) to generate statistical charts (such as PDF or GIF) for in-depth analysis and intuitive display of simulation results.

[0049] Those skilled in the art know that in addition to implementing the system provided by the present application and each device, module, unit thereof in pure computer readable program code, the same functions can be achieved by logically programming the method steps in the form of logic gates, switches, application specific integrated circuits, programmable logic controllers, and embedded microcontrollers. Therefore, the system provided by the present application and each device, module, unit thereof can be considered as a hardware component, and the devices, modules, units included therein for implementing various functions can also be considered as structures within the hardware component; the devices, modules, units for implementing various functions can also be considered as both software modules implementing methods and structures within hardware components.

[0050] The specific embodiments of the present application are described above. It should be understood that the present application is not limited to the above specific embodiments, and those skilled in the art can make various changes or modifications within the scope of the claims, which does not affect the essential content of the present application. The embodiments of the present application and the features in the embodiments can be arbitrarily combined with each other without conflict.

Claims

1. A high performance discrete event simulation system based on JAX, characterized in that, The application relates to a discrete event simulation engine module, a JAX acceleration module and a standardized external interface module. The discrete event simulation engine module obtains simulation parameters provided by a user or an upper application and system state data dynamically generated in a simulation process, performs scheduling and processing of events, and obtains updated state data and event trajectory data. The JAX acceleration module performs hardware-accelerated parallel calculation according to the updated state data and event trajectory data, and outputs optimized and parallelized execution results to the discrete event simulation engine module to update the simulation state, and the execution results are also called by the standardized external interface module to complete efficient environment stepping operations. The standardized external interface module obtains bottom state data provided by the discrete event simulation engine module and accelerated execution results provided by the JAX acceleration module, generates observation, reward and termination signals conforming to a machine learning method standard, and is called by an external AI algorithm framework to perform policy training or reasoning. The discrete event simulation engine module comprises an environment configuration submodule, an event scheduling submodule, an event analysis submodule and a job-level tracking submodule.

2. The JAX-based high-performance discrete-event simulation system of claim 1, wherein, The environment configuration submodule is used for defining and solidifying static parameters of a simulation environment. The event scheduling submodule is responsible for driving simulation forward evolution according to a current state. The event analysis submodule is used for realizing event logic judgment without side effects. The job-level tracking submodule is used for encapsulating and managing multiple parallel independent simulation instances. The environment configuration submodule uses an immutable data class to obtain configuration parameters input by a user, generates an EnvParams parameter structure, and provides the parameter structure for the event scheduling submodule and the event analysis submodule to initialize the simulation environment.

3. The JAX-based high-performance discrete-event simulation system of claim 2, wherein, The event scheduling submodule generates a new immutable EnvState object by calling a core step_env function according to the EnvParams parameter structure and a current EnvState, outputs new state data, and provides the new state data for the event analysis submodule to judge a next event type to be processed and for the job-level tracking submodule to manage independent simulation states in a batch processing mode. The event analysis submodule compares the occurrence time of a customer arrival event and a service completion event by using a JAX conditional control primitive according to the new state data, analyzes and returns updated state, and provides the updated state for the event scheduling submodule to complete a single-step simulation cycle. The job-level tracking submodule obtains a batch of random seeds and a single EnvParams specified by a user, generates job batch data containing multiple independent simulation initial states, and provides the job batch data for the JAX acceleration module to realize vectorized parallel execution and for the standardized external interface module to generate batch observation data of a multi-instance environment. The configuration parameters comprise a total simulation time length, server processing time distribution and customer arrival time interval.

4. The JAX-based high-performance discrete-event simulation system of claim 3, wherein, The EnvState comprises the number of customers in a queue, the busy state of each server and a current simulation clock. ​ 5. The JAX-based high-performance discrete-event simulation system of claim 1, wherein, The JAX acceleration module includes an instant compilation submodule, an automatic vectorization submodule, and a multi-device support submodule; The instant compilation submodule is used to convert Python functions into highly optimized machine code; The automatic vectorization submodule is used to apply operations for a single instance to multiple instances in batches; The multi-device support submodule is used to manage and schedule computing hardware resources.

6. The JAX-based high-performance discrete-event simulation system of claim 5, wherein, The instant compilation submodule obtains XLA optimized code generated by the automatic vectorization submodule after processing the vectorized simulation function, and sends it to the multi-device support submodule for execution on specified hardware; The automatic vectorization submodule obtains single-instance simulation functions defined by the discrete event simulation engine module and batch data generated by the job-level tracking submodule, and generates new vectorized simulation functions by applying jax.vmap transformation; The multi-device support submodule obtains optimized code generated by the instant compilation submodule and vectorized data arrays prepared by the automatic vectorization submodule, and returns the final result after executing the calculation on the GPU or TPU, which is sent to the discrete event simulation engine module to update the state of all simulations in the batch.

7. The JAX-based high-performance discrete-event simulation system of claim 1, wherein, The standardized external interface module includes a state observation conversion submodule, a reward calculation submodule, and a batch adaptation and demultiplexing submodule; The state observation conversion submodule is used to convert internal complex simulation states into numerical observation spaces understandable by AI algorithms; The reward calculation submodule is used to evaluate the pros and cons of state change quantitative simulation results; The batch adaptation and demultiplexing submodule is used to bridge external single-step operations and internal parallel computing.

8. The JAX-based high-performance discrete event simulation system of claim 7, wherein The state observation conversion submodule obtains the EnvState object output by the discrete event simulation engine module, extracts and flattens the key information according to the predefined rules to generate an observation vector, which is called by the batch adaptation and demultiplexing submodule as part of the standard step function and returned; The reward calculation submodule obtains the EnvState object of the previous step and the EnvState of the current step, calculates the scalar reward value by comparing the state changes, which is called by the batch adaptation and demultiplexing submodule to evaluate the goodness of the previous action; The batch adaptation and demultiplexing submodule obtains a single action from the external AI algorithm and the entire batch execution result from the JAX acceleration module, wraps or injects the single action into the batch process, and extracts the state, observation, and reward corresponding to the single instance from the batch calculation result after receiving the batch calculation result. The single-instance tuple generated after demultiplexing is directly returned to the caller of the external AI algorithm, thereby completing a standard environment interaction.

9. A high performance discrete event simulation method based on JAX, characterized in that, Through the cooperation of the discrete-time simulation engine module, the JAX acceleration module, and the standardized external interface module, including: Step S1: All static parameters of the simulation environment are defined by the environment configuration submodule in the discrete-time simulation engine module using an immutable data structure, and are solidified into an EnvParams object; the EnvParams object contains constant parameters of the simulation scene, and the EnvParams object is instantiated by receiving parameter values input by the user, and is passed as a static context to subsequent modules for initializing the simulation starting state EnvState and participating in calculations as an invariant throughout the entire simulation process; Step S2: The event scheduling submodule and the event analysis submodule in the discrete-time simulation engine module use a pure function step_env to process the logic of a single event; the step_env function receives the current EnvState and EnvParams as input and returns a completely new immutable EnvState; the event logic judgment is implemented using the conditional control primitive of JAX, the system dynamic state is modeled as a JAX array, and the update is completed by creating a new array; an iterative loop is constructed based on lax.scan, the step_env function is repeatedly called, the EnvState output in the last step is used as the input in the next step to advance the simulation, and the final EnvState is obtained; Step S3: The single and serial simulation logic defined in step S2 is promoted to large-scale parallel execution by cooperating with the job-level tracking submodule of the discrete-time simulation engine module; Step S4: The parallel simulation capability of step S3 is packaged into an environment interaction interface conforming to the AI field standard through a standardized interface module; the interface triggered by the external AI algorithm calls the parallel execution process of step S3, receives the batch results and parses them into the data structure required by the AI framework, and integrates data post-processing and visualization functions, calculates key performance indicators, and generates statistical charts.

10. The JAX-based high-performance discrete-event simulation method of claim 9, wherein, The step S3 includes: Step S3.1: Assigning a unique random seed to each simulation instance using JAX's pseudo-random number generator via the job-level tracking submodule of the discrete-time simulation engine module, all simulation instances sharing the same EnvParams object, organizing independent simulation tasks into a batch data structure; Step S3.2: through the just-in-time compilation submodule of the JAX acceleration module, the simulation iteration function constructed based on lax.scan in step S2 is compiled into XLA optimized code using the jax.jit decorator; and through the automatic vectorization submodule, the compiled function is transformed using jax.vmap to process the batch data generated in step S3.1; Step S3.3: The batch data processed by the parallelization function after vmap transformation is automatically distributed to multiple computing hardware for parallel execution by the multi-device support submodule of the JAX acceleration module, and the batch data containing the results of all simulation instances is returned.

Images