Python curses: multiprocessing issue with Pool.map?

724 views Asked by At

I have an issue with Pool.map in combination with the curses module of Python. Whenever I compute bigger workloads with Pool.map my curses UI breaks: it does not react upon the default screen's getch anymore. Instead of reading in any pressed key instantly (and continuing with parsing it) I can press any amount of keys until I hit enter. Sometimes (in addition to that) even the UI breaks (like showing a fraction of my normal shell).


Curses UI wrapper

This is a wrapper class (Screen) that handles the curses UI stuff for me:

# -*- coding: utf-8 -*-
import curses



class Screen(object):

    def __init__(self):
        # create a default screen
        self.__mainscr = curses.initscr()
        self.__stdscr = curses.newwin(curses.LINES - 2, curses.COLS - 2, 1, 1)
        self.__max_height, self.__max_width = self.__stdscr.getmaxyx()

        # start coloring
        curses.start_color()
        curses.use_default_colors()

        # define colors
        curses.init_pair(1, 197, -1)  # red
        curses.init_pair(2, 227, -1)  # yellow
        curses.init_pair(3, curses.COLOR_MAGENTA, -1)
        curses.init_pair(4, curses.COLOR_GREEN, -1)  # darkgreen
        curses.init_pair(5, curses.COLOR_BLUE, -1)
        curses.init_pair(6, curses.COLOR_BLACK, curses.COLOR_WHITE)
        curses.init_pair(7, curses.COLOR_WHITE, -1)
        curses.init_pair(8, curses.COLOR_CYAN, -1)
        curses.init_pair(9, 209, -1)  # orange
        curses.init_pair(10, 47, -1)  # green


    def add_str_to_scr(self, add_str: str, colorpair: int = 7):
        self.__stdscr.addstr(str(add_str), curses.color_pair(colorpair))


    def linebreak(self):
        self.__stdscr.addstr("\n")


    def clear_screen(self):
        self.__stdscr.clear()


    def refresh_screen(self):
        self.__stdscr.refresh()


    def wait_for_enter_or_esc(self):
        curses.noecho()
        while True:
            c = self.__stdscr.getch()
            if c == 10 or c == 27:  # 10: Enter, 27: ESC
                break
        curses.echo()


    def get_user_input_chr(self) -> str:
        return chr(self.__stdscr.getch())


    def get_user_input_str(self) -> str:
        return self.__stdscr.getstr().decode(encoding="utf-8")

Actual program

I've written a small example, since the mentioned failure always happen when I combine Pool.map within a curses UI and have high workload. The code just computes some useless mult and add stuff on a numpy array.

import curses
from screen import Screen
from multiprocessing import Pool, cpu_count
import numpy as np

s = Screen()           # initializing my Screen wrapper
np.random.seed(1234)   # setting the rng fixed to make results comparable

# worker function to simulate workload
def worker(arr):
    return arr * 2 + 1

s.clear_screen()    # cleans the screen
s.refresh_screen()  # displays current buffer's content

s.add_str_to_scr("Start processing data...")
s.linebreak()
s.linebreak()
s.refresh_screen()

# data to feed worker function with (sliced by rows)
data_arr = np.random.rand(8, int(1e7))   # <-- big array for high workload

with Pool(cpu_count()) as p:
    buffer = p.map(worker, [data_arr[row] for row in np.ndindex(data_arr.shape[0])])

s.add_str_to_scr("...finished processing:")
s.linebreak()
s.linebreak()
s.refresh_screen()
for row in buffer:
    s.add_str_to_scr(row[0:3])
    s.linebreak()
    s.refresh_screen()

# *Here* the program should wait until the user presses *any* key
# and continue INSTANTLY when any key gets pressed.
# However, for big workloads, it does not react to single key presses,
# but wait for any amount of keys pressed until you hit 'Enter'
s.get_user_input_chr()
curses.endwin()

Now, when I execute the code with a high workload (i.e. crunching an array of shape (8, int(1e7) equals 8 rows with 10,000,000 columns) curse's getch breaks and I get this behavior:

enter image description here

As you can see, I can hit q (or any other key) as often as I want, but curse's getch does not react. I have to press the Enter key to make it recognize the input.

Moreover the first line gets overwritten with my original shell's output for some reason.

This behavior only happens, when the calculation of Pool.map roughly needs 1 second or longer.

When I set data_arr to a small array like np.random.rand(8, 100) everything works like a charm, but as soon as I feed big arrays where the computation takes like >= 1second, this weird bug appears and breaks my curses UI.

Any ideas?

Is Pool.map not joining the worker processes correctly somehow?

1

There are 1 answers

5
Thomas Dickey On

The program is doing what you told it to do:

  • you're calling initscr (and ignoring the fact that curses creates a top-level window),
  • then creating a subwindow which covers most of the screen,
  • printing several lines on the screen, refreshing the display after each line, and
  • finally at the end, waiting for input from the subwindow.

However, your program doesn't call cbreak, raw, etc., which would let you read an unbuffered (no "Enter" pressed) character. Also, the program doesn't turn off echo. If the load is light, you won't notice, since response is fast. But under heavy load, e.g., swapping or high memory/CPU usage, it'll still be recovering when it gets to the prompt. So you notice.

Regarding the screen size, perhaps you meant

self.__stdscr = curses.newwin(curses.LINES - 1, curses.COLS - 1, 0, 0)

but supposing that you intended to leave "empty" space around the window, you could improve things by doing

        self.__mainscr.refresh()

immediately after initscr (which would erase the screen).