Source code for schrodinger.application.desmond.memory

"""
Utilities for memory estimation and benchmarking

Copyright Schrodinger, LLC. All rights reserved.
"""

import os

import psutil


def estimate_memory(natoms: int) -> int:
    """
    :param: Number of atoms in the system
    :return: GPU memory usage in bytes
    """
    # 750 MB + 6 MB per thousand particles
    # Not sure how good this is now. Probably inaccurate for GCMC. -YW 11/1/2019
    return 786432000 + natoms * 6292


class MemoryGauger:
    """
    Gauge the memory usage of the process. Use example::

      mem_gauger = MemoryGauger()    # Point 1

      # code that uses a lot memory.
      ...

      mem_usage = mem_gauger.mem_usage()  # Memory usage change since Point 1.
    """

    def __init__(self):
        self._process = psutil.Process(os.getpid())
        self._vm_used_base = self._process.memory_info()[0]

    def mem_usage(self) -> float:
        """
        Returns the change of the memory usage (in MB) of this process since the
        creation of this `MemoryGauger` instance.
        """
        return (self._process.memory_info()[0] - self._vm_used_base) / 1048576