Skip to content

gpu #

GPU memory management utilities.

Functions:

cleanup_simulation #

cleanup_simulation(
    simulation: Simulation,
    integrator: Integrator | None = None,
) -> None

Clean up OpenMM simulation and free GPU memory.

This function properly releases GPU resources by deleting the simulation and integrator objects, synchronizing CUDA operations, and emptying the GPU cache. The synchronization step is critical - without it, the cache may be cleared while GPU operations are still in flight.

Parameters:

  • simulation (Simulation) –

    The OpenMM simulation object to clean up.

  • integrator (Integrator | None, default: None ) –

    The integrator object to clean up. If None, only the simulation is deleted.

Examples:

>>> integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 2*femtoseconds)
>>> simulation = Simulation(topology, system, integrator)
>>> # ... use simulation ...
>>> cleanup_simulation(simulation, integrator)
Source code in presto/utils/gpu.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def cleanup_simulation(
    simulation: Simulation, integrator: Integrator | None = None
) -> None:
    """Clean up OpenMM simulation and free GPU memory.

    This function properly releases GPU resources by deleting the simulation
    and integrator objects, synchronizing CUDA operations, and emptying the
    GPU cache. The synchronization step is critical - without it, the cache
    may be cleared while GPU operations are still in flight.

    Parameters
    ----------
    simulation : Simulation
        The OpenMM simulation object to clean up.
    integrator : Integrator | None, optional
        The integrator object to clean up. If None, only the simulation
        is deleted.

    Examples
    --------
    >>> integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 2*femtoseconds)
    >>> simulation = Simulation(topology, system, integrator)
    >>> # ... use simulation ...
    >>> cleanup_simulation(simulation, integrator)
    """
    del simulation
    if integrator is not None:
        del integrator
    if torch.cuda.is_available():
        torch.cuda.synchronize()  # Wait for all GPU operations to complete
        torch.cuda.empty_cache()  # Now GPU is idle, memory can actually be freed