fix frozen for python backend

Author: reiase

Cargo.toml

@@ -17,9 +17,10 @@ pyo3 = { version = "0.18.1", features = [
     "abi3",
     "abi3-py37",
 ] }
+lazy_static = "1.4.0"
 
-[dev_dependencies]
-rspec = "1.0"
+# [dev_dependencies]
+# rspec = "1.0"
 
 [profile.dev]
-overflow-checks=false
+overflow-checks=false

tests/cxx_test.cc renamed to examples/cpp/cxx_test.cc

@@ -461,6 +461,11 @@ def current():
     def init(params=None):
         """init param_scope for a new thread."""
         param_scope(**params).__enter__()
+        
+    @staticmethod
+    def frozen():
+        with param_scope():
+            TLSKVStorage.frozen()
 
 
 _param_scope = param_scope._func

tests/cxx_test.py renamed to examples/cpp/cxx_test.py

@@ -1,7 +1,11 @@
+import os
 import threading
 from typing import Any, Callable, Dict, Iterable
 
 
+GLOBAL_STORAGE = {}
+
+
 class Storage:
     """Base class for all storage implementations"""
 
@@ -55,7 +59,13 @@ def __init__(self, parent=None) -> None:
         self._parent = parent
         super().__init__()
 
-        if hasattr(TLSKVStorage.tls, "his") and len(TLSKVStorage.tls.his) > 0:
+        if not hasattr(TLSKVStorage.tls, "his"):
+            TLSKVStorage.tls.his = [TLSKVStorage.__new__(TLSKVStorage)]
+            TLSKVStorage.tls.his[-1]._storage = GLOBAL_STORAGE
+            TLSKVStorage.tls.his[-1]._parent = None
+            self.update(GLOBAL_STORAGE)
+
+        elif hasattr(TLSKVStorage.tls, "his") and len(TLSKVStorage.tls.his) > 0:
             parent = TLSKVStorage.tls.his[-1]
             self.update(parent._storage)
 
@@ -86,6 +96,7 @@ def _update(values={}, prefix=None):
                     _update(v, prefix=key)
                 else:
                     storage[key] = v
+
         if kws is not None:
             return _update(kws, prefix=None)
 
@@ -123,10 +134,19 @@ def current():
             TLSKVStorage.tls.his = [TLSKVStorage()]
         return TLSKVStorage.tls.his[-1]
 
+    @staticmethod
+    def frozen():
+        GLOBAL_STORAGE.update(TLSKVStorage.tls.his[-1].storage())
+
+
 try:
-    from hyperparameter.rbackend import KVStorage 
-    TLSKVStorage = KVStorage
+    if os.environ.get("HYPERPARAMETER_BACKEND", "RUST") == "RUST":
+        from hyperparameter.rbackend import KVStorage
+        TLSKVStorage = KVStorage
+        print("using native hyperparameter backend")
+    else:
+        print("using python hyperparameter backend")
 except:
     import traceback
     traceback.print_exc()
-    pass
+    print("using python hyperparameter backend")

tests/cxx_test.sh renamed to examples/cpp/cxx_test.sh

@@ -36,4 +36,4 @@ source = ["hyperparameter"]
 [tool.pytest.ini_options]
 minversion = "6.0"
 addopts = "-ra -q --durations=5 --doctest-modules --doctest-glob=*.md"
-testpaths = ["hyperparameter/", "docs"]
+testpaths = ["hyperparameter/", "docs", "tests/"]

hyperparameter/api.py

@@ -14,9 +14,10 @@ use pyo3::types::PyString;
 use pyo3::FromPyPointer;
 
 use crate::entry::Value;
+use crate::storage::frozen_as_global_storage;
 use crate::storage::Storage;
-use crate::storage::MGR;
 use crate::storage::StorageManager;
+use crate::storage::MGR;
 
 #[pyclass]
 pub struct KVStorage {
@@ -148,6 +149,11 @@ impl KVStorage {
         kv.storage.isview = 1;
         kv
     }
+
+    #[staticmethod]
+    pub fn frozen() {
+        frozen_as_global_storage();
+    }
 }
 
 #[pymodule]

hyperparameter/storage.py

@@ -1,6 +1,3 @@
-// #![feature(local_key_cell_methods)]
-// #![feature(let_chains)]
-
 pub mod entry;
 pub mod storage;
 

pyproject.toml

@@ -2,6 +2,9 @@ use std::cell::{Ref, RefCell, RefMut};
 use std::collections::BTreeMap;
 use std::collections::HashSet;
 use std::rc::Rc;
+use std::sync::Mutex;
+
+use lazy_static::lazy_static;
 
 use crate::entry::{Entry, EntryValue, Value};
 use crate::xxh::xxhstr;
@@ -48,15 +51,65 @@ thread_local! {
 }
 
 pub fn init_storage_manager() -> RefCell<StorageManager> {
+    let mut tree = Tree::new();
+    global_storage_get(&mut tree);
     let sm = RefCell::new(StorageManager {
-        tls: Rc::new(RefCell::new(Tree::new())),
+        tls: Rc::new(RefCell::new(tree)),
         stack: Vec::new(),
     });
     sm.borrow_mut().stack.push(RefCell::new(HashSet::new()));
 
     return sm;
 }
 
+lazy_static! {
+    static ref GLOBAL_STORAGE: Mutex<u64> = {
+        let tree = Box::new(Tree::new());
+        Mutex::new(Box::into_raw(tree) as u64)
+    };
+}
+
+fn global_storage_set(t: &Tree) {
+    GLOBAL_STORAGE
+        .lock()
+        .and_then(|v| unsafe {
+            let ptr = v.clone() as *mut Tree;
+            match ptr.as_mut() {
+                Some(tree) => {
+                    tree.clear();
+                    tree.clone_from(t);
+                }
+                None => todo!(),
+            };
+            Ok(())
+        })
+        .unwrap();
+}
+
+pub fn frozen_as_global_storage() {
+    MGR.with(|mgr| {
+        let t = mgr.borrow().tls.borrow().clone();
+        global_storage_set(&t);
+    });
+}
+
+fn global_storage_get(t: &mut Tree) {
+    GLOBAL_STORAGE
+        .lock()
+        .and_then(|v| unsafe {
+            let ptr = v.clone() as *mut Tree;
+            match ptr.as_mut() {
+                Some(tree) => {
+                    t.clear();
+                    t.clone_from(tree);
+                }
+                None => todo!(),
+            };
+            Ok(())
+        })
+        .unwrap();
+}
+
 #[derive(Debug)]
 pub struct Storage {
     pub parent: Rc<RefCell<Tree>>,
@@ -114,10 +167,11 @@ impl Storage {
 
     pub fn exit(&mut self) {
         MGR.with(|m| {
-            if let Some(keys) = m.borrow_mut().stack.pop() {
+            let mut m = m.borrow_mut();
+            if let Some(keys) = m.stack.pop() {
                 keys.borrow()
                     .iter()
-                    .for_each(|k| tree_rollback(m.borrow_mut().tls.borrow_mut(), *k));
+                    .for_each(|k| tree_rollback(m.tls.borrow_mut(), *k));
             }
         });
         // MGR.with_borrow_mut(|m| {