ray: Memory leak when simple worker returns a numpy object

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04 LTS x86_64
• Ray installed from (source or binary): binary (pip install ray)
• Ray version: 0.5.0
• Python version: Python 3.6.0 :: Continuum Analytics, Inc.
• Exact command to reproduce: python main.py:

Describe the problem

The code below results in a memory leak. The ring buffer in the main process containing numpy images should take about 10GB. However after the buffer cycles the memory continues increasing. With the parameters in the script below the memory will increase to 16GB in about 500,000 steps.

I also reproduced this behavior on a server with 126GB RAM and 40 threads. The system will be out of memory after a few tens of millions of steps (common on RL tasks).

I first noticed this with workers returning atari screens and tried to provide a minimal working example that reproduces the issue.

edit: the output of the script will seem fine, however top or similar will show the increasing memory usage.

Source code / logs

import os
import psutil
import gc

import random
import numpy as np
import ray

import pyarrow
import lz4.frame
import base64


def pack(data):
    data = pyarrow.serialize(data).to_buffer().to_pybytes()
    data = lz4.frame.compress(data)
    return base64.b64encode(data)


def unpack(data):
    data = base64.b64decode(data)
    data = lz4.frame.decompress(data)
    return pyarrow.deserialize(data)


@ray.remote(num_cpus=1, num_gpus=0)
def sample(step_cnt):
    # some pseudo-random "images" that can still be compressed.
    state = np.full((128, 128, 3), step_cnt % 255, dtype=np.uint8)
    state[24:48, 64:96, :] = np.random.randint(0, 255, (24, 32, 3), np.uint8)

    # display process memory
    if step_cnt % 50000 == 0:
        process = psutil.Process(os.getpid())
        proc_mem = process.memory_info().rss / (1024**2)
        print(f'actor_pid={process.pid} \t mem={proc_mem:6.1f} MB.')

    return pack(state)


def main():
    ray.init()

    actor_no = 10
    step_no = 5000000
    buff_sz = 200000

    # init buffer
    buffer = [None for i in range(buff_sz)]

    step_cnt = 0
    while step_cnt < step_no:

        states = ray.get([sample.remote(step_cnt) for i in range(actor_no)])

        for state in states:
            # add to buffer
            buffer[step_cnt % buff_sz] = unpack(state)
            step_cnt += 1

        # display process memory
        if step_cnt % 5000 == 0:
            print('-' * 30)
            print(f'steps={step_cnt} / buffer={len(buffer)}.')

            process = psutil.Process(os.getpid())
            proc_mem = process.memory_info().rss / (1024**2)
            print(f'main_pid={process.pid} \t mem={proc_mem:6.1f} MB.')
            print('-' * 30)

if __name__ == '__main__':
    main()