util.py

# Copyright 2026 NNAISENSE SA
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import random
from collections.abc import Callable, Iterator, Sequence
from typing import Any

from .core import _RemoteInstance, get


class ActorPool:
    """
    A mock implementation of `ray.util.ActorPool`.
    """

    SEED = 1

    def __init__(self, actors: Sequence[_RemoteInstance]):
        """
        Initialize the `ActorPool`.

        Args:
            actors: List of so-called remote actors.
        """
        self.__actors = actors
        self.__generator = random.Random(self.SEED)

    def map(self, fn: Callable[[_RemoteInstance, Any], Any], objects: Sequence[Any]) -> Iterator[Any]:
        """
        Pair actors with objects, and for each pair, call `fn(actor, object)`.

        Args:
            fn: The function which specifies how an object is to be used
                with its associated actor.
            objects: A sequence (e.g. list or tuple) of object to be processed
                by the pool of actors.

        Returns:
            Processed counterparts of the objects, as an iterator.
        """
        num_actors = len(self.__actors)
        result = []
        for i_obj, obj in enumerate(objects):
            i_actor = i_obj % num_actors
            actor = self.__actors[i_actor]
            result.append(fn(actor, obj))
        return iter(get(result))

    def map_unordered(self, fn: Callable[[_RemoteInstance, Any], Any], objects: Sequence[Any]) -> Iterator[Any]:
        """
        Like the `map` method, but the results are shuffled.

        The reasoning behind the shuffling is to mimic the fact that
        `map_unordered` of the actual 'ray' library returns the processed
        results in an undeterministic order.

        *Note:* for reproducibility during tests, the shuffling depends
        on a constant random seed, which is stored by class-level constant
        `SEED`.
        """
        result = list(self.map(fn, objects))
        self.__generator.shuffle(result)
        return iter(result)