Python Memory Management: Doing More with Less

•⏱️17 days ago
DataPythonProgrammingETLBig-Data

Ever felt like your Python code is hogging all your RAM like a hungry tax office? You're not alone! In this guide, we'll dive deep into the wild world of Python memory management - where every byte counts and efficiency is king.

Think of it as teaching your code to go on a diet: we'll show you how to trim the fat from your applications and make them run smoother than a freshly waxed penguin (pun intended, we're talking Python after all! 🐍). From taming the infamous garbage collector to avoiding those pesky memory leaks that make your code slower than a turtle in molasses - we've got you covered.

Whether you're running code on a tiny IoT device or wrestling with massive datasets in the cloud, these tips and tricks will help you stop throwing money at cloud providers and start writing lean, mean Python machines. Because let's face it - nobody likes seeing that dreaded "Out of Memory" error, especially when your boss is watching!

Ready to become a memory management ninja? Let's dive in! 🚀

Understanding Memory Constraints

Memory constraints affect everything from IoT devices to large-scale applications. Python's memory management becomes crucial when working with limited resources. The problem with having memory at the bottom of the list of things to worry about is that it will sooner or later lead to out-of-memory errors, degraded performance and increased operational costs in cloud environments where memory usage directly affects billing. Remember: ignorantia juris non excusat

Key Optimization Strategies

I divide the strategies depending of the cases. Let's check it one by one.

Generator Expressions

Instead of loading entire lists into memory:

# Bad - loads full list
numbers = [x * x for x in range(1000000)]

# Good - generates values on demand
numbers = (x * x for x in range(1000000))

Chunked Processing

When handling large files you will find the issue of having some struggle loading all to memory. If isn't worth it the RAM increase of your instance or worker or pod you can try chunking the load.

In the following method you will see the route and how many likes will be loaded. Try to test first how many lines are worth to use

def process_large_binary(filename, chunk_size=1024):
    with open(filename, 'r') as f:
        while chunk := f.read(chunk_size):
            yield chunk

## Usage
import cv2
import numpy as np
temp_file = 'temp_video.mp4'
with open(temp_file, 'wb') as f:
    for chunk in process_large_binary(video_path):
        f.write(chunk)

cap = cv2.VideoCapture(temp_file)
frame_count = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    cv2.imwrite(f'{output_dir}/frame_{frame_count}.jpg', frame)
    frame_count += 1

cap.release()

There is a variation of this were you can use it to read the files of a file. That will be very helpful for CSV, TSV and TXT files.

def process_large_file(filename, chunk_size=1024):
    with open(filename, mode='r', newline='', encoding='utf-8') as file:
        reader = csv.reader(file)
        header = next(reader)  # Read the header row
        chunk = []
        for row in reader:
            chunk.append(row)
            if len(chunk) >= chunk_size:
                yield header, chunk
                chunk = []
        # Yield any remaining rows after the last full chunk
        if chunk:
            yield header, chunk

## Usage
for header, chunk in read_csv_in_chunks(file_path, chunk_size):
    # Load it into chunks in a DB
    load_to_postgres(values=chunk,columns=header)

Using slots

For classes with fixed attributes use __slots__ attribute. This allows you to explicitly state which instance attributes you expect your object instances to have, with the expected results

class Point:
    __slots__ = ['x', 'y']
    
    def __init__(self, x, y):
        self.x = x
        self.y = y

Memory-mapped files

For really large files operations with mmap allows you to use the operating system virtual memory system to access the data on the filesystem directly, instead of using normal I/O functions.

import mmap

with open('large_file.dat', 'rb') as f:
    mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    # Process file without loading into memory, it uses references

Monitoring Memory Usage

Another way is knowing how much ram are you using will help to identify any memory issue or identify how much RAM your instance will require.

import psutil
import os

def get_memory_usage():
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024  # MB

Final Tips

Use itertools for memory-efficient iterations

Verify if using itertools is a best for your case. Be careful because sometimes itertools can make you use more memory if you are using wrong. Remember that this thing is a generator.

from itertools import islice

list_ = [i for i in range(10, 100)]

def chunker(it, size):
    iterator = iter(it)
    while chunk := list(islice(iterator, size)):
        print(chunk)

Implement cleanup hooks using context managers

The principial idea of this approach is to free memory when its not needed anymore for certain reasons: The task is finished, the task failed but some resources like connections or a file loaded can be cleanned easier.

This is an example of a resource manager that you can use it for loading big files or manage database connections

class ResourceManager:
    def __init__(self, resource_name):
        self.resource_name = resource_name
        self.resource = None

    def __enter__(self):
        # Allocate the resource
        print(f"Allocating resource: {self.resource_name}")
        self.resource = f"Resource({self.resource_name})"
        return self.resource

    def __exit__(self, exc_type, exc_value, traceback):
        # Cleanup the resource
        print(f"Cleaning up resource: {self.resource_name}")
        self.resource = None
        return False  # Propagate exceptions if any

# Usage
with ResourceManager("Database Connection") as resource:
    print(f"Using {resource}")
    # Simulate some operations
    # If an exception occurs, cleanup is still performed

print("Resource cleanup complete.")

There is another option with a decorator instead of using class and follow the pattern.

from contextlib import contextmanager

@contextmanager
def managed_resource(resource_name):
    print(f"Allocating resource: {resource_name}")
    resource = f"Resource({resource_name})"
    try:
        yield resource  # Provide the resource to the calling context
    finally:
        print(f"Cleaning up resource: {resource_name}")
        # Cleanup code goes here

# Usage
with managed_resource("File Handler") as resource:
    print(f"Using {resource}")
    # Simulate operations
    # If an exception occurs, cleanup is still performed

print("Resource cleanup complete.")

Real example

This context manager is designed to use temporary files

import tempfile

@contextmanager
def temporary_file():
    temp = tempfile.NamedTemporaryFile(delete=False)
    print(f"Created temporary file: {temp.name}")
    try:
        yield temp
    finally:
        temp.close()
        print(f"Deleting temporary file: {temp.name}")

# Usage
with temporary_file() as temp_file:
    temp_file.write(b"Temporary data")
    temp_file.flush()
    print(f"Temporary file in use: {temp_file.name}")

print("Temporary file cleanup complete.")

Consider using NumPy's memmap for large array operations

The advantages of using numpy memmap is different from other solutions that I showed you before. You can use another features from numpy lib that will you make accomplish to not load everything to RAM memory and also saving the information sequentially is solved or proccessed. This will let to work arrays that are larger than your current RAM.

Principal advantages:

  • Automatic memory management
  • Compatible with other numpy operations
  • Changes are written directly to disk
import numpy as np

# Create memory-mapped array
data = np.memmap('large_array.dat', 
                dtype='float64', 
                mode='w+', 
                shape=(1000000, 100))

# Process in chunks
chunk_size = 1000
for i in range(0, data.shape[0], chunk_size):
   chunk = data[i:i+chunk_size]
   # Process chunk
   chunk *= 2
   # Changes are written to disk automatically

# Using memmap for matrix operations
def process_large_matrix():
   matrix = np.memmap('matrix.dat',
                     dtype='float64',
                     mode='r+',
                     shape=(100000, 100000))
   
   # Process by rows
   for i in range(matrix.shape[0]):
       matrix[i] = matrix[i] * 2
       
   # Flush changes
   matrix.flush()

Profile memory usage with memory_profiler decorator

If you are identifyng memory leaks you can use the following decorator @profile from memory_profiler lib. Use it when you are trying to hunt something related with memory leaks or bad memory optimizations in a Django, FastAPI or Data Science project or anything built in python.

from memory_profiler import profile
import numpy as np

@profile
def analyze_matrix(size):
   # Create large matrix
   matrix = np.random.random((size, size))
   
   # Perform operations
   result = np.dot(matrix, matrix.T)
   eigenvalues = np.linalg.eigvals(result)
   
   return eigenvalues

@profile
def process_data_chunks(filename, chunk_size=1000):
   data = []
   with open(filename, 'r') as f:
       while chunk := f.readlines(chunk_size):
           processed = [line.strip().upper() for line in chunk]
           data.extend(processed)
   return data

# Run with:
# python -m memory_profiler script.py

if __name__ == '__main__':
   analyze_matrix(1000)
   process_data_chunks('large_file.txt')

Thanks for reading!

This post was inspired on work experiences that kindapped a lot of sleep hours thinking what it can be the problem.