[feat] Update hkv to support more evict strategy

tensorflow · Mar 19, 2024 · 7a2e790 · 7a2e790
1 parent 095db80
commit 7a2e790
Show file tree

Hide file tree

Showing 20 changed files with 1,016 additions and 259 deletions.
diff --git a/.github/workflows/make_wheel_Linux_x86.sh b/.github/workflows/make_wheel_Linux_x86.sh
@@ -13,7 +13,7 @@ fi
 
 # if tensorflow version >= 2.6.0 and <= 2.11.9
 if [[ "$TF_VERSION" =~ ^2\.([6-9]|10|11)\.[0-9]$ ]] ; then
-  export BUILD_IMAGE="tfra/nosla-cuda11.2.1-cudnn8-ubuntu20.04-manylinux2014-python$PY_VERSION"
+  export BUILD_IMAGE="tfra/nosla-cuda11.2.2-cudnn8-ubuntu20.04-manylinux2014-python$PY_VERSION"
   export TF_CUDA_VERSION="11.2"
   export TF_CUDNN_VERSION="8.1"
 elif [ $TF_VERSION == "2.4.1" ] ; then

diff --git a/WORKSPACE b/WORKSPACE
@@ -57,55 +57,9 @@ http_archive(
 http_archive(
     name = "hkv",
     build_file = "//build_deps/toolchains/hkv:hkv.BUILD",
-    # TODO(LinGeLin) remove this when update hkv
-    patch_cmds = [
-        """sed -i.bak '1772i\\'$'\\n    ThrustAllocator<uint8_t> thrust_allocator_;\\n' include/merlin_hashtable.cuh""",
-        """sed -i.bak '225i\\'$'\\n    thrust_allocator_.set_allocator(allocator_);\\n' include/merlin_hashtable.cuh""",
-        "sed -i.bak 's/thrust::sort_by_key(thrust_par.on(stream)/thrust::sort_by_key(thrust_par(thrust_allocator_).on(stream)/' include/merlin_hashtable.cuh",
-        "sed -i.bak 's/reduce(thrust_par.on(stream)/reduce(thrust_par(thrust_allocator_).on(stream)/' include/merlin_hashtable.cuh",
-        """sed -i.bak '125i\\'$'\\n    template <typename T>\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '126i\\'$'\\n    struct ThrustAllocator : thrust::device_malloc_allocator<T> {\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '127i\\'$'\\n     public:\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '128i\\'$'\\n      typedef thrust::device_malloc_allocator<T> super_t;\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '129i\\'$'\\n      typedef typename super_t::pointer pointer;\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '130i\\'$'\\n      typedef typename super_t::size_type size_type;\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '131i\\'$'\\n     public:\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '132i\\'$'\\n      pointer allocate(size_type n) {\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '133i\\'$'\\n        void* ptr = nullptr;\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '134i\\'$'\\n        MERLIN_CHECK(\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '135i\\'$'\\n            allocator_ != nullptr,\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '136i\\'$'\\n            "[ThrustAllocator] set_allocator should be called in advance!");\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '137i\\'$'\\n        allocator_->alloc(MemoryType::Device, &ptr, sizeof(T) * n);\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '138i\\'$'\\n        return pointer(reinterpret_cast<T*>(ptr));\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '139i\\'$'\\n      }\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '140i\\'$'\\n      void deallocate(pointer p, size_type n) {\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '141i\\'$'\\n        MERLIN_CHECK(\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '142i\\'$'\\n            allocator_ != nullptr,\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '143i\\'$'\\n            "[ThrustAllocator] set_allocator should be called in advance!");\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '144i\\'$'\\n        allocator_->free(MemoryType::Device, reinterpret_cast<void*>(p.get()));\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '145i\\'$'\\n      }\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '146i\\'$'\\n      void set_allocator(BaseAllocator* allocator) { allocator_ = allocator; }\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '147i\\'$'\\n     public:\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '148i\\'$'\\n      BaseAllocator* allocator_ = nullptr;\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '149i\\'$'\\n     };\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '20i\\'$'\\n     #include <thrust/device_malloc_allocator.h>\\n' include/merlin/allocator.cuh""",
-        """sed -i.bak '367i\\'$'\\n  for (auto addr : (*table)->buckets_address) {\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '368i\\'$'\\n    allocator->free(MemoryType::Device, addr);\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '369i\\'$'\\n  }\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '370i\\'$'\\n  /*\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '382i\\'$'\\n  */\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '224i\\'$'\\n  uint8_t* address = nullptr;\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '225i\\'$'\\n  allocator->alloc(MemoryType::Device, (void**)&(address), bucket_memory_size * (end - start));\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '226i\\'$'\\n  (*table)->buckets_address.push_back(address);\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '228i\\'$'\\n    allocate_bucket_others<K, V, S><<<1, 1>>>((*table)->buckets, i, address + (bucket_memory_size * (i-start)), reserve_size, bucket_max_size);\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '229i\\'$'\\n    /*\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '235i\\'$'\\n    */\\n' include/merlin/core_kernels.cuh""",
-        """sed -i.bak '22i\\'$'\\n#include <vector>\\n' include/merlin/types.cuh""",
-        """sed -i.bak '143i\\'$'\\n  std::vector<uint8_t*> buckets_address;\\n' include/merlin/types.cuh""",
-    ],
-    sha256 = "f8179c445a06a558262946cda4d8ae7252d313e73f792586be9b1bc0c993b1cf",
-    strip_prefix = "HierarchicalKV-0.1.0-beta.6",
-    url = "https://github.com/NVIDIA-Merlin/HierarchicalKV/archive/refs/tags/v0.1.0-beta.6.tar.gz",
+    sha256 = "841be4cfb4059e5745838a23a32c776cfff1d38306b95b1ac2659df0d4d9709b",
+    strip_prefix = "HierarchicalKV-master",
+    url = "https://github.com/NVIDIA-Merlin/HierarchicalKV/archive/refs/heads/master.zip",
 )
 
 tf_configure(

diff --git a/tensorflow_recommenders_addons/dynamic_embedding/__init__.py b/tensorflow_recommenders_addons/dynamic_embedding/__init__.py
@@ -18,6 +18,7 @@
     'CuckooHashTable',
     'CuckooHashTableConfig',
     'CuckooHashTableCreator',
+    'HkvEvictStrategy',
     'HkvHashTable',
     'HkvHashTableConfig',
     'HkvHashTableCreator',
@@ -55,7 +56,7 @@
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops import data_flow_ops as data_flow
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.dynamic_embedding_creator import (
     KVCreator, CuckooHashTableConfig, CuckooHashTableCreator,
-    HkvHashTableConfig, HkvHashTableCreator, RedisTableConfig,
+    HkvHashTableConfig, HkvHashTableCreator, HkvEvictStrategy, RedisTableConfig,
     RedisTableCreator, FileSystemSaver)
 from tensorflow_recommenders_addons.dynamic_embedding.python.ops.cuckoo_hashtable_ops import (
     CuckooHashTable,)