Enable Python API of memory-map IO for PECOS-HNSW (#212)

Co-authored-by: Wei-Cheng Chang <chanweic@amazon.com>

Author: OctoberChang

pecos/ann/hnsw/model.py

@@ -10,6 +10,7 @@
 #  and limitations under the License.
 from ctypes import (
     POINTER,
+    c_bool,
     c_float,
     c_uint32,
     c_char_p,
@@ -149,10 +150,11 @@ def train(cls, X, train_params=None, pred_params=None):
         return cls(model_ptr, pX.rows, pX.cols, fn_dict, pred_params)
 
     @classmethod
-    def load(cls, model_folder):
+    def load(cls, model_folder, lazy_load=False):
         """Load HNSW model from file
         Args:
             model_folder (str): model directory from which the model is loaded.
+            lazy_load (bool): whether to lazy_load memory-mapped files (default False).
         Returns:
             HNSWModel (pecos.ann.hnsw.HNSW): the loaded HNSW model
         """
@@ -168,7 +170,7 @@ def load(cls, model_folder):
         c_model_dir = f"{model_folder}/c_model"
         if not os.path.isdir(c_model_dir):
             raise ValueError(f"c_model_dir did not exist: {c_model_dir}")
-        model_ptr = fn_dict["load"](c_char_p(c_model_dir.encode("utf-8")))
+        model_ptr = fn_dict["load"](c_char_p(c_model_dir.encode("utf-8")), c_bool(lazy_load))
         pred_params = cls.PredParams.from_dict(param["pred_kwargs"])
         return cls(model_ptr, param["num_item"], param["feat_dim"], fn_dict, pred_params)
 

pecos/core/base.py

@@ -1608,7 +1608,10 @@ def link_ann_hnsw_methods(self):
                 c_fn_name = f"c_ann_hnsw_{fn_name}_{data_type}_{metric_type}_f32"
                 local_fn_dict[fn_name] = getattr(self.clib_float32, c_fn_name)
                 res_list = c_void_p  # pointer to C/C++ pecos::ann::HNSW
-                arg_list = [c_char_p]  # pointer to char* model_dir
+                arg_list = [
+                    c_char_p,  # pointer to C/C++ pecos:ann::hnsw
+                    c_bool,  # bool for lazy_load of mmap files
+                ]
                 corelib.fillprototype(local_fn_dict[fn_name], res_list, arg_list)
 
                 fn_name = "save"

pecos/core/libpecos.cpp

@@ -378,9 +378,9 @@ extern "C" {
     C_ANN_HNSW_TRAIN(_drm_l2_f32, ScipyDrmF32, pecos::drm_t, hnsw_drm_l2_t)
 
     #define C_ANN_HNSW_LOAD(SUFFIX, HNSW_T) \
-    void* c_ann_hnsw_load ## SUFFIX(const char* model_dir) { \
+    void* c_ann_hnsw_load ## SUFFIX(const char* model_dir, const bool lazy_load) { \
         HNSW_T *model_ptr = new HNSW_T(); \
-        model_ptr->load(model_dir); \
+        model_ptr->load(model_dir, lazy_load); \
         return static_cast<void*>(model_ptr); \
     }
     C_ANN_HNSW_LOAD(_drm_ip_f32, hnsw_drm_ip_t)

test/pecos/ann/test_hnsw.py

@@ -45,6 +45,14 @@ def test_save_and_load(tmpdir):
     assert Yp_from_mem == approx(
         Yp_from_file, abs=0.0
     ), f"save and load failed: Yp_from_mem != Yp_from_file"
+    del model
+
+    # test load memory-mapped files
+    model = HNSW.load(model_folder, lazy_load=True)
+    Yp_from_file, _ = model.predict(X_tst, pred_params=pred_params, ret_csr=False)
+    assert Yp_from_mem == approx(
+        Yp_from_file, abs=0.0
+    ), f"load mmap-file failed: Yp_from_mem != Yp_from_file"
 
 
 def test_predict_and_recall():