<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
import os
import ray

logger = logging.getLogger(__name__)


class TaskPool(object):
    """Helper class for tracking the status of many in-flight actor tasks."""

    def __init__(self):
        self._tasks = {}
        self._objects = {}
        self._fetching = []

    def add(self, worker, all_obj_ids):
        if isinstance(all_obj_ids, list):
            obj_id = all_obj_ids[0]
        else:
            obj_id = all_obj_ids
        self._tasks[obj_id] = worker
        self._objects[obj_id] = all_obj_ids

    def completed(self, blocking_wait=False):
        pending = list(self._tasks)
        if pending:
            ready, _ = ray.wait(pending, num_returns=len(pending), timeout=0)
            if not ready and blocking_wait:
                ready, _ = ray.wait(pending, num_returns=1, timeout=10.0)
            for obj_id in ready:
                yield (self._tasks.pop(obj_id), self._objects.pop(obj_id))

    def completed_prefetch(self, blocking_wait=False, max_yield=999):
        """Similar to completed but only returns once the object is local.

        Assumes obj_id only is one id."""

        for worker, obj_id in self.completed(blocking_wait=blocking_wait):
            self._fetching.append((worker, obj_id))

        remaining = []
        num_yielded = 0
        for worker, obj_id in self._fetching:
            if num_yielded < max_yield:<commit_msg>Revert "[rllib] remove exists call (#6168)" (#6616)<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging
import os
import ray

logger = logging.getLogger(__name__)


class TaskPool(object):
    """Helper class for tracking the status of many in-flight actor tasks."""

    def __init__(self):
        self._tasks = {}
        self._objects = {}
        self._fetching = []

    def add(self, worker, all_obj_ids):
        if isinstance(all_obj_ids, list):
            obj_id = all_obj_ids[0]
        else:
            obj_id = all_obj_ids
        self._tasks[obj_id] = worker
        self._objects[obj_id] = all_obj_ids

    def completed(self, blocking_wait=False):
        pending = list(self._tasks)
        if pending:
            ready, _ = ray.wait(pending, num_returns=len(pending), timeout=0)
            if not ready and blocking_wait:
                ready, _ = ray.wait(pending, num_returns=1, timeout=10.0)
            for obj_id in ready:
                yield (self._tasks.pop(obj_id), self._objects.pop(obj_id))

    def completed_prefetch(self, blocking_wait=False, max_yield=999):
        """Similar to completed but only returns once the object is local.

        Assumes obj_id only is one id."""

        for worker, obj_id in self.completed(blocking_wait=blocking_wait):
            (ray.worker.global_worker.raylet_client.fetch_or_reconstruct(
                [obj_id], True))
            self._fetching.append((worker, obj_id))

        remaining = []
        num_yielded = 0
        for worker, obj_id in self._fetching:
            if (num_yielded < max_yield
                    and ray.worker.global_worker.core_worker.object_exists(
                        obj_id)):
{'fraction_non_alphanumeric': 0.09112927641537996, 'fraction_numerical': 0.008477141992128368, 'mean_word_length': 2.481559536354057, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3475445', 'n_tokens_mistral': 1003, 'n_tokens_neox': 948, 'n_words': 251}