IntEnumValue — A multiprocessing safe enum value for Python

Announcing the initial release of the python-multiprocessing-intenum module, providing a multiprocessing safe enum value for the Python programming language.

Background

Within a large customer project involving multiple components running in a Kubernetes environment, we faced a specific problem with no generic solution available when implementing a daemon worker component.
As only one calculation can be running at a time, the component has an internal state to track whether it is idle and ready to accept new calculations, busy calculating or pending to shut itself down, e.g. during a rolling upgrade.
Besides the actual calculation process that is being spawned, the component provides a Django Rest Framework based API from multiple web server processes or threads that allows another scheduler component to query the current state of the worker and submit new calculation jobs. Thus the internal state needs to be safely accessed and modified from multiple processes and/or threads.

The natural data type to track a state in most high level programming languages including Python is a enum type that one can assign symbolic names corresponding to the states. When it comes to sharing data across multiple processes or threads in Python, the multiprocessing module comes into play, which provides shared ctypes objects, ie. basic data types such as character and integer, as well as locking primitives to handle sharing larger complex data structures.
It does not provide a ready to use way to share an enum value though.

The problem being pretty common, it called for a generic and reusable solution and so the idea for IntEnumValue was born, an implementation of a multiprocessing safe shared object for IntEnum enum values.

Features

To the user it appears and can be used much like a Python IntEnum, which happens to be an enum, whose named values each correspond to an unique integer.
Internally, it then uses a multiprocessing.Value shared ctypes integer object to store that integer value of the enum in a multiprocessing safe way.

To support a coding style that proactively guards against programming errors, the module is fully typed allowing static checking with mypy and implemented as a generic class taking the type of the designated underlying specific enum as a parameter.

The python-multiprocessing-intenum module comes with unit tests providing full coverage of its code.

Usage Examples

To illustrate the usage and better demonstrate the features, an example of a multithreaded worker is used, which has an internal state represented as an enum that should be tracked across its threads.

To use a multiprocessing safe IntEnum, first define the actual underlying IntEnum type:

class WorkerStatusEnum(IntEnum):
    UNAVAILABLE = enum.auto()
    IDLE = enum.auto()
    CALCULATING = enum.auto()

IntEnumValue is a generic class, accepting a type parameter for the type of the underlying IntEnum.
Thus define the corresponding IntEnumValue type for that IntEnum like this:

class WorkerStatus(IntEnumValue[WorkerStatusEnum]):
    pass

What this does, is define a new WorkerStatus type that is an IntEnumValue parametrized to handle the specific IntEnum type WorkerStatusEnum only.
It can be used like an enum to track the state:

>>> status = WorkerStatus(WorkerStatusEnum.IDLE)
>>> status.name
'IDLE'
>>> status.value
2
>>> with status.get_lock():
...     status.set(WorkerStatusEnum.CALCULATING)
>>> status.name
'CALCULATING'
>>> status.value
3

Trying to set it to a different type of IntEnum however is caught as a TypeError:

>>> class FrutEnum(IntEnum):
...     APPLE = enum.auto()
...     ORANGE = enum.auto()
>>> status.set(FrutEnum.APPLE)
Traceback (most recent call last):
  File "", line 1, in 
  File "/home/elho/python-multiprocessing-intenum/src/multiprocessing_intenum/__init__.py", line 60, in set
    raise TypeError(message)
TypeError: Can not set '<enum 'WorkerStatusEnum'>' to value of type '<enum 'FrutEnum'>'

This helps to guard against programming errors, where a different enum is erroneously being used.

Besides being used directly, the created WorkerStatus type can, of course, also be wrapped in a dedicated class, which is what the remaining examples will further expand on:

class WorkerState:
    def __init__(self) -> None:
        self.status = WorkerStatus(WorkerStatusEnum.IDLE)

When using multiple multiprocessing.Value instances (including IntEnumValue ones) that should share a lock to allow ensuring that they can only be changed in a consistent state, pass that shared lock as a keyword argument on instantiation:

class WorkerState:
    def __init__(self) -> None:
        self.lock = multiprocessing.RLock()
        self.status = WorkerStatus(WorkerStatusEnum.IDLE, lock=self.lock)
        self.job_id = multiprocessing.Value("i", -1, lock=self.lock)

    def get_lock(self) -> Lock | RLock:
        return self.lock

To avoid having to call the set() method to assign a value to the IntEnumValue attribute, it is suggested to keep the actual attribute private to the class and implement getter and setter methods for a public property that hides this implementation detail, e.g. as follows:

class WorkerState:
    def __init__(self) -> None:
        self._status = WorkerStatus(WorkerStatusEnum.IDLE)

    @property
    def status(self) -> WorkerStatusEnum:
        return self._status  # type: ignore[return-value]

    @status.setter
    def status(self, status: WorkerStatusEnum | str) -> None:
        self._status.set(status)

The result can be used in a more elegant manner by simply assigning to the status attribute:

>>> state = WorkerState()
>>> state.status.name
'IDLE'
>>> with state.get_lock():
...     state.status = WorkerStatusEnum.CALCULATING
>>> state.status.name
'CALCULATING'

The specific IntEnumValue type can override methods to add further functionality.

A common example is overriding the set() method to add logging:

class WorkerStatus(IntEnumValue[WorkerStatusEnum]):
    def set(self, value: WorkerStatusEnum | str) -> None:
        super().set(value)
        logger.info(f"WorkerStatus set to '{self.name}'")

Putting all these features and use cases together allows handling internal state of multiprocessing worker in an elegant, cohesive and robust way.