Limitar la memoria de Python

Limitar la memoria de Python

Limitar la memoria de un programa de Python puede ser útil para depurar el código en situaciones diversas. A la postre, en la mayoría de las ocasiones no sabemos en qué hardware acabará corriendo nuestra aplicación, por lo cual es mejor estar preparado para los escenarios más variados. Por ejemplo, si nuestro código lee información a través de una red, o si carga los datos de un archivo del sistema, o si descarga un archivo vía HTTP, o cualquier otra situación similar donde se trabaje con grandes volúmenes de datos, ¿estamos preparados para manejar los errores que arrojará Python cuando no haya memoria RAM suficiente? Con algunas de las herramientas que mencionaremos a continuación, podremos limitar la memoria de un proceso de Python para ver cómo responde nuestro código en entornos de memoria limitada.

Para distribuciones de Linux y sistemas basados en Unix (como macOS) disponemos del módulo estándar resource, cuya función setrlimit() permite establecer un límite a la memoria RAM consumida por nuestro programa:

import resource

# Límite de memoria (expresado en bytes).
limit = 1024 * 1024 * 1024   # 1GB
# Aplicar la restricción de memoria.
resource.setrlimit(resource.RLIMIT_DATA, (limit, -1))
# Intentamos crear un objeto del tamaño de la memoria límite.
# Si la restricción se aplicó correctamente, debe lanzar
# la excepción MemoryError.
bytearray(limit)

Si el código se ejecuta correctamente, Python debe lanzar la excepción MemoryError al intentar crear un bytearray del tamaño del límite de memoria establecido.

En Windows podemos limitar la memoria del proceso de Python vía job objects. Una implementación posible usando pywin32 es la siguiente:

# Basado en https://stackoverflow.com/questions/54949110/limit-python-script-ram-usage-in-windows

import sys
import warnings

import winerror
import win32api
import win32job


g_hjob = None


def create_job(job_name='', breakaway='silent'):
    hjob = win32job.CreateJobObject(None, job_name)
    if breakaway:
        info = win32job.QueryInformationJobObject(hjob,
                    win32job.JobObjectExtendedLimitInformation)
        if breakaway == 'silent':
            info['BasicLimitInformation']['LimitFlags'] |= (
                win32job.JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK)
        else:
            info['BasicLimitInformation']['LimitFlags'] |= (
                win32job.JOB_OBJECT_LIMIT_BREAKAWAY_OK)
        win32job.SetInformationJobObject(hjob,
            win32job.JobObjectExtendedLimitInformation, info)
    return hjob


def assign_job(hjob):
    global g_hjob
    hprocess = win32api.GetCurrentProcess()
    try:
        win32job.AssignProcessToJobObject(hjob, hprocess)
        g_hjob = hjob
    except win32job.error as e:
        if (e.winerror != winerror.ERROR_ACCESS_DENIED or
            sys.getwindowsversion() >= (6, 2) or
            not win32job.IsProcessInJob(hprocess, None)):
            raise
        warnings.warn('The process is already in a job. Nested jobs are not '
            'supported prior to Windows 8.')


def limit_memory(memory_limit):
    if g_hjob is None:
        return
    info = win32job.QueryInformationJobObject(g_hjob,
                win32job.JobObjectExtendedLimitInformation)
    info['ProcessMemoryLimit'] = memory_limit
    info['BasicLimitInformation']['LimitFlags'] |= (
        win32job.JOB_OBJECT_LIMIT_PROCESS_MEMORY)
    win32job.SetInformationJobObject(g_hjob,
        win32job.JobObjectExtendedLimitInformation, info)


assign_job(create_job())
# Límite de memoria (expresado en bytes).
limit = 1024 * 1024 * 1024      # 1GB
# Aplicar la restricción de memoria.
limit_memory(limit)
# Intentamos crear un objeto del tamaño de la memoria límite.
# Si la restricción se aplicó correctamente, debe lanzar
# la excepción MemoryError.
bytearray(limit)

Otra posibilidad en Windows consiste en usar procgov, una herramienta de código abierto para limitar procesos de múltiples maneras y que internamente trabaja con la misma API de Windows que el código anterior. Desde Python podemos invocarlo vía subprocess para aplicar restricciones a la memoria RAM en nuestro programa:

import os
import subprocess


# Función que lee el texto impreso por el proceso (procgov64) hasta
# encontrar el texto indicado.
def read_until_find(process: subprocess.Popen, text: bytes):
    buffer: bytes = b""
    # Asegurarse de que el proceso tiene un `stdout` desde el
    # cual leer.
    if process.stdout is None:
        raise ValueError("no stdout to read from")
    while (byte := process.stdout.read(1)):
        buffer += byte
        # Luego de cada byte leído, chequear si se encuentra el texto.
        if buffer.endswith(text):
            return
    # Si el intérprete llegó a esta línea, el texto no se encuentra
    # en la salida del proceso.
    raise ValueError("text not found in stdout")


# Límite de memoria (expresado en megabytes).
limit = 1024
# Aplicar el límite de memoria al proceso actual, cuyo PID obtenemos
# vía os.getpid().
p = subprocess.Popen(
    ["procgov64", "-m", f"{limit}M", "-p", str(os.getpid())],
    stdout=subprocess.PIPE,
    stdin=subprocess.DEVNULL,
    stderr=subprocess.DEVNULL
)
# Esperar a que procgov64 aplique efectivamente el límite.
read_until_find(p, b"Press Ctrl-C to end execution without terminating the process.")

try:
    # Aquí va el código que queremos ejecutar bajo la restricción
    # de memoria.
    # Intentamos crear un objeto del tamaño de la memoria límite.
    # Si la restricción se aplicó correctamente, debe lanzar
    # la excepción MemoryError.
    bytearray(limit * 1024 * 1024)    # Lanza MemoryError
finally:
    # Finalizar la restricción: cerrar procgov64 y esperar
    # a que termine.
    p.kill()
    p.wait()

Este código supone que el archivo procgov64.exe (que puedes descargar desde este enlace) se encuentra en la misma carpeta (o, mejor dicho, en el mismo directorio actual de trabajo) que nuestro programa. Si estamos usando un intérprete de 32 bits, debemos reemplazar procgov64.exe por procgov32.exe en la llamada a subprocess.Popen().

Donar ❤️

¿Te gusta nuestro contenido? ¡Ayudanos a seguir creciendo con una donación!

Deja una respuesta