import ipyparallel as ipp

rc = ipp.Cluster(n=4).start_and_connect_sync()
rc.wait_for_engines(4)
rc.activate()

%%px --local

import psutil
from threading import Thread


class MemInfo:
    """Simple memory info sampler
    
    Runs in a background thread,
    reports highest value
    
    - start() starts samplnig
    """
    def __init__(self, interval=0.1):

        self._done = False
        self._peak = -1
        self.interval = interval
        self.process = psutil.Process()
        self._thread = None

    def start(self):
        if self._thread:
            self.stop()
        self._peak = -1
        self._done = False
        self._thread = Thread(target=self.sample, daemon=True, name="memory sampler")
        self._thread.start()
    
    def sample(self):
        self._peak = max(self._peak, self.process.memory_info().rss / float(2 ** 20))

    def collect_samples(self):
        self.sample()
        while not self._done:
            time.sleep(self.interval)
            self.sample()


    def stop(self):
        self._done = True
        if self._thread:
            self._thread.join()
            self._thread = None
        self.sample()
        return self._peak


mem_info = MemInfo()

# not really necessary, just here for interactive development,
# unregister previous instance
try:
    ip.events.unregister("pre_run_cell", start_sampling)
    ip.events.unregister("post_run_cell", stop_sampling)
except NameError:
    pass


def start_sampling():
    mem_info.start()


def stop_sampling():
    rss = mem_info.stop()
    if rss > 0:
        print(f"Memory usage: {rss:.0f} MB")


ip.events.register("pre_run_cell", start_sampling)
ip.events.register("post_run_cell", stop_sampling)

%px pass

%%px
import numpy as np
a = np.random.random((10, 1024, 1024))

%px pass