Add Actors for Ray Dedup. (modelscope#526)

Author: chenyushuo

configs/config_all.yaml

@@ -782,15 +782,15 @@ process:
   - video_deduplicator:                                     # deduplicator to deduplicate samples at document-level using exact matching of videos between documents.
       consider_text: false                                    # whether to consider text hash together with video hash when applying deduplication.
   - ray_video_deduplicator:                                 # the simple video deduplicator that can run on multi-nodes using md5 hashing exact matching method
-      redis_host: 'redis_host'                              # the host of the redis instance
-      redis_port: 6380                                      # the port of redis instance, please note that the default port of redis is 6379 which is the same as default port for ray, so we need to modify the default redis config to use it in other port
+      backend: 'ray_actor'                                    # the backend for dedup, either 'ray_actor' or 'redis'
+      redis_address: 'redis://localhost:6379'                 # the address of redis server
   - ray_image_deduplicator:                                 # the simple image deduplicator that can deduplicate samples at document-level using exact matching of images between documents.
-      redis_host: 'redis_host'                              # the host of the redis instance
-      redis_port: 6380                                      # the port of redis instance, please note that the default port of redis is 6379 which is the same as default port for ray, so we need to modify the default redis config to use it in other port
-      method: phash                                         # hash method for image. One of [phash, dhash, whash, ahash]
+      backend: 'ray_actor'                                    # the backend for dedup, either 'ray_actor' or 'redis'
+      redis_address: 'redis://localhost:6379'                 # the address of redis server
+      method: phash                                           # hash method for image. One of [phash, dhash, whash, ahash]
   - ray_document_deduplicator:                              # the simple document deduplicator that can run on multi-nodes using md5 hashing exact matching method
-      redis_host: 'redis_host'                              # the host of the redis instance
-      redis_port: 6380                                      # the port of redis instance, please note that the default port of redis is 6379 which is the same as default port for ray, so we need to modify the default redis config to use it in other port
+      backend: 'ray_actor'                                    # the backend for dedup, either 'ray_actor' or 'redis'
+      redis_address: 'redis://localhost:6379'                 # the address of redis server
       lowercase: false                                        # whether to convert text to lower case
       ignore_non_character: false                             # whether to ignore non-alphabet characters, including whitespaces, digits, and punctuations
   - ray_bts_minhash_deduplicator:                            # the document deduplicator that can run on multi-nodes using minhashLSH algorithm

data_juicer/ops/deduplicator/ray_basic_deduplicator.py

@@ -1,4 +1,6 @@
-from pydantic import PositiveInt
+from abc import ABC, abstractmethod
+
+import ray
 
 from data_juicer.utils.constant import HashKeys
 from data_juicer.utils.lazy_loader import LazyLoader
@@ -7,6 +9,69 @@
 
 redis = LazyLoader('redis', 'redis')
 
+MERSENNE_PRIME = (1 << 61) - 1
+
+
+@ray.remote(scheduling_strategy='SPREAD')
+class DedupSet:
+
+    def __init__(self):
+        self.hash_record = set()
+
+    def is_unique(self, key):
+        if key not in self.hash_record:
+            self.hash_record.add(key)
+            return True
+        else:
+            return False
+
+
+class Backend(ABC):
+    """
+    Backend for deduplicator.
+    """
+
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def is_unique(self, md5_value: str):
+        pass
+
+
+class ActorBackend(Backend):
+    """
+    Ray actor backend for deduplicator.
+    """
+
+    def __init__(self, dedup_set_num: int):
+        self.dedup_set_num = dedup_set_num
+        self.dedup_sets = [
+            DedupSet.remote() for _ in range(self.dedup_set_num)
+        ]
+
+    def is_unique(self, md5_value: str):
+        dedup_set_id = int.from_bytes(
+            md5_value.encode(),
+            byteorder='little') % MERSENNE_PRIME % self.dedup_set_num
+        return ray.get(
+            self.dedup_sets[dedup_set_id].is_unique.remote(md5_value))
+
+
+class RedisBackend(Backend):
+    """
+    Redis backend for deduplicator.
+    """
+
+    def __init__(self, redis_address: str):
+        self.redis_address = redis_address
+        self.redis_client = redis.from_url(url=self.redis_address)
+        self.redis_client.flushdb(0)
+
+    def is_unique(self, md5_value: str):
+        return self.redis_client.setnx(md5_value, 1)
+
 
 class RayBasicDeduplicator(Filter):
     """
@@ -19,37 +84,40 @@ class RayBasicDeduplicator(Filter):
     EMPTY_HASH_VALUE = 'EMPTY'
 
     def __init__(self,
-                 redis_host: str = 'localhost',
-                 redis_port: PositiveInt = 6380,
+                 backend: str = 'ray_actor',
+                 redis_address: str = 'redis://localhost:6379',
                  *args,
                  **kwargs):
         """
         Initialization.
-        :param redis_host: the hostname of redis server
-        :param redis_port: the port of redis server
+        :param backend: the backend for dedup, either 'ray_actor' or 'redis'
+        :param redis_address: the address of redis server
         :param args: extra args
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
-        self.redis_host = redis_host
-        self.redis_port = redis_port
-        # TODO: add a barrier to ensure that flushdb is performed before
-        # the operator is called
-        r = redis.StrictRedis(host=self.redis_host, port=self.redis_port, db=0)
-        r.flushdb(0)
+        self.redis_address = redis_address
+        self.backend = backend
+        if backend == 'ray_actor':
+            dedup_set_num = int(ray.cluster_resources().get('CPU') / 2)
+            self.backend = ActorBackend(dedup_set_num)
+        elif backend == 'redis':
+            # TODO: add a barrier to ensure that flushdb is performed before
+            # the operator is called
+            self.backend = RedisBackend(redis_address)
+        else:
+            raise ValueError(f'Unknown backend: {backend}')
 
     def calculate_hash(self, sample, context=False):
         """Calculate hash value for the sample."""
         raise NotImplementedError
 
     def compute_stats_single(self, sample, context=False):
-        # init redis client
-        r = redis.StrictRedis(host=self.redis_host, port=self.redis_port, db=0)
         # compute hash
         md5_value = self.calculate_hash(sample, context)
         # check existing
-        sample[HashKeys.is_duplicate] = r.setnx(md5_value, 1)
+        sample[HashKeys.is_unique] = self.backend.is_unique(md5_value)
         return sample
 
     def process_single(self, sample):
-        return sample[HashKeys.is_duplicate]
+        return sample[HashKeys.is_unique]

data_juicer/ops/deduplicator/ray_document_deduplicator.py

@@ -2,7 +2,6 @@
 import string
 
 import regex as re
-from pydantic import PositiveInt
 
 from ..base_op import OPERATORS
 from .ray_basic_deduplicator import RayBasicDeduplicator
@@ -17,24 +16,24 @@ class RayDocumentDeduplicator(RayBasicDeduplicator):
     """
 
     def __init__(self,
-                 redis_host: str = 'localhost',
-                 redis_port: PositiveInt = 6380,
+                 backend: str = 'ray_actor',
+                 redis_address: str = 'redis://localhost:6379',
                  lowercase: bool = False,
                  ignore_non_character: bool = False,
                  *args,
                  **kwargs):
         """
         Initialization method.
-        :param redis_host: the hostname of redis server
-        :param redis_port: the port of redis server
+        :param backend: the backend for dedup, either 'ray_actor' or 'redis'
+        :param redis_address: the address of redis server
         :param lowercase: Whether to convert sample text to lower case
         :param ignore_non_character: Whether to ignore non-alphabet
         characters, including whitespaces, digits, and punctuations
         :param args: extra args
         :param kwargs: extra args.
         """
-        super().__init__(redis_host=redis_host,
-                         redis_port=redis_port,
+        super().__init__(backend=backend,
+                         redis_address=redis_address,
                          *args,
                          **kwargs)
         self.lowercase = lowercase

data_juicer/ops/deduplicator/ray_image_deduplicator.py

@@ -1,5 +1,4 @@
 import numpy as np
-from pydantic import PositiveInt
 
 from data_juicer.utils.lazy_loader import LazyLoader
 from data_juicer.utils.mm_utils import load_data_with_context, load_image
@@ -36,20 +35,20 @@ class RayImageDeduplicator(RayBasicDeduplicator):
     """
 
     def __init__(self,
-                 redis_host: str = 'localhost',
-                 redis_port: PositiveInt = 6380,
+                 backend: str = 'ray_actor',
+                 redis_address: str = 'redis://localhost:6379',
                  method: str = 'phash',
                  *args,
                  **kwargs):
         """
         Initialization.
-        :param redis_host: the hostname of redis server
-        :param redis_port: the port of redis server
+        :param backend: the backend for dedup, either 'ray_actor' or 'redis'
+        :param redis_address: the address of redis server
         :param args: extra args
         :param kwargs: extra args
         """
-        super().__init__(redis_host=redis_host,
-                         redis_port=redis_port,
+        super().__init__(backend=backend,
+                         redis_address=redis_address,
                          *args,
                          **kwargs)
         if method not in HASH_METHOD:

data_juicer/ops/deduplicator/ray_video_deduplicator.py

@@ -1,7 +1,5 @@
 import hashlib
 
-from pydantic import PositiveInt
-
 from data_juicer.utils.mm_utils import (close_video, load_data_with_context,
                                         load_video)
 
@@ -21,19 +19,19 @@ class RayVideoDeduplicator(RayBasicDeduplicator):
     """
 
     def __init__(self,
-                 redis_host: str = 'localhost',
-                 redis_port: PositiveInt = 6380,
+                 backend: str = 'ray_actor',
+                 redis_address: str = 'redis://localhost:6379',
                  *args,
                  **kwargs):
         """
         Initialization.
-        :param redis_host: the hostname of redis server
-        :param redis_port: the port of redis server
+        :param backend: the backend for dedup, either 'ray_actor' or 'redis'
+        :param redis_address: the address of redis server
         :param args: extra args
         :param kwargs: extra args
         """
-        super().__init__(redis_host=redis_host,
-                         redis_port=redis_port,
+        super().__init__(backend=backend,
+                         redis_address=redis_address,
                          *args,
                          **kwargs)
 

data_juicer/utils/constant.py

@@ -268,7 +268,7 @@ class HashKeys(object):
     videohash = DEFAULT_PREFIX + 'videohash'
 
     # duplicate flag
-    is_duplicate = DEFAULT_PREFIX + 'is_duplicate'
+    is_unique = DEFAULT_PREFIX + 'is_unique'
 
 
 class InterVars(object):