Repository: ray-project/ray
File: python/ray/data/_internal/execution/operators/actor_pool_map_operator.py (Python)

Code before:


class _ActorPool(AutoscalingActorPool):
    """A pool of actors for map task execution.

    This class is in charge of tracking the number of in-flight tasks per actor,
    providing the least heavily loaded actor to the operator, and killing idle
    actors when the operator is done submitting work to the pool.
    """

    _ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S = 10
    _ACTOR_POOL_GRACEFUL_SHUTDOWN_TIMEOUT_S = 30
    _LOGICAL_ACTOR_ID_LABEL_KEY = "__ray_data_logical_actor_id"

    def __init__(
        self,
        create_actor_fn: "Callable[[Dict[str, str]], Tuple[ActorHandle, ObjectRef[Any]]]",
        per_actor_resource_usage: ExecutionResources,
        *,
        min_size: int,
        max_size: int,
        initial_size: int,
        max_actor_concurrency: int,
        max_tasks_in_flight_per_actor: int,
        _enable_actor_pool_on_exit_hook: bool = False,
        _map_worker_cls_name: Optional[str] = None,
    ):
        """Initialize the actor pool.

        Args:
            create_actor_fn: This function should take key-value labels as input, and
                create an actor with those labels. The function should return the actor
                handle and a reference to the actor's node ID.
            per_actor_resource_usage: The resource usage per actor.
            min_size: The minimum number of running actors to be maintained
                in the pool. Note, that this constraint could be violated when
                no new work is available for scheduling in the actor pool (ie
                when operator completes execution).
            max_size: The maximum number of running actors to be maintained
                in the pool.
            initial_size: The initial number of actors to start with.
            max_actor_concurrency: The maximum number of concurrent tasks a
                single actor can execute (derived from `ray_remote_args`
                passed to the operator).
            max_tasks_in_flight_per_actor: The maximum number of tasks that can
                be submitted to a single actor at any given time.
            _enable_actor_pool_on_exit_hook: Whether to enable the actor pool
                on exit hook.
            _map_worker_cls_name: Name of the map worker class for logging purposes.
        """


Review comment (suggestion):
Also, when will callers pass in `None`? If doesn't happen in practice (either in the source code or the tests), this might be premature generality, and it might be simpler to either require a `str` or do a `str` and default to something like `MapWorker`

Code after:


class _ActorPool(AutoscalingActorPool):
    """A pool of actors for map task execution.

    This class is in charge of tracking the number of in-flight tasks per actor,
    providing the least heavily loaded actor to the operator, and killing idle
    actors when the operator is done submitting work to the pool.
    """

    _ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S = 10
    _ACTOR_POOL_GRACEFUL_SHUTDOWN_TIMEOUT_S = 30
    _LOGICAL_ACTOR_ID_LABEL_KEY = "__ray_data_logical_actor_id"

    def __init__(
        self,
        create_actor_fn: "Callable[[Dict[str, str]], Tuple[ActorHandle, ObjectRef[Any]]]",
        per_actor_resource_usage: ExecutionResources,
        *,
        min_size: int,
        max_size: int,
        initial_size: int,
        max_actor_concurrency: int,
        max_tasks_in_flight_per_actor: int,
        map_worker_cls_name: str = "MapWorker",
        debounce_period_s: int = _ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S,
        _enable_actor_pool_on_exit_hook: bool = False,
    ):
        """Initialize the actor pool.

        Args:
            create_actor_fn: This function should take key-value labels as input, and
                create an actor with those labels. The function should return the actor
                handle and a reference to the actor's node ID.
            per_actor_resource_usage: The resource usage per actor.
            min_size: The minimum number of running actors to be maintained
                in the pool. Note, that this constraint could be violated when
                no new work is available for scheduling in the actor pool (ie
                when operator completes execution).
            max_size: The maximum number of running actors to be maintained
                in the pool.
            initial_size: The initial number of actors to start with.
            max_actor_concurrency: The maximum number of concurrent tasks a
                single actor can execute (derived from `ray_remote_args`
                passed to the operator).
            max_tasks_in_flight_per_actor: The maximum number of tasks that can
                be submitted to a single actor at any given time.
            map_worker_cls_name: Name of the map worker class for logging purposes.
            debounce_period_s: Debounce period for scaling down after scaling up
            _enable_actor_pool_on_exit_hook: Whether to enable the actor pool
                on exit hook.