support get_by_hash in auto_param (#4)

Author: reiase

hyperparameter/api.py

@@ -2,7 +2,8 @@
 import inspect
 from typing import Any, Callable, Dict
 
-from hyperparameter.storage import TLSKVStorage
+from hyperparameter.storage import TLSKVStorage, has_rust_backend
+from hyperparameter.storage import xxh64
 from .tune import Suggester
 
 
@@ -461,7 +462,7 @@ def current():
     def init(params=None):
         """init param_scope for a new thread."""
         param_scope(**params).__enter__()
-        
+
     @staticmethod
     def frozen():
         with param_scope():
@@ -520,6 +521,38 @@ def auto_param(name_or_func):
     if callable(name_or_func):
         return auto_param(None)(name_or_func)
 
+    if has_rust_backend:
+
+        def hashed_wrapper(func):
+            predef_kws = {}
+
+            if name_or_func is None:
+                namespace = func.__name__
+            else:
+                namespace = name_or_func
+
+            signature = inspect.signature(func)
+            for k, v in signature.parameters.items():
+                if v.default != v.empty:
+                    name = "{}.{}".format(namespace, k)
+                    predef_kws[k] = xxh64(name)
+
+            @functools.wraps(func)
+            def inner(*arg, **kws):
+                with param_scope() as hp:
+                    for k, v in predef_kws.items():
+                        if k not in kws:
+                            try:
+                                val = hp._storage.get_by_hash(v)
+                                kws[k] = val
+                            except ValueError:
+                                pass
+                    return func(*arg, **kws)
+
+            return inner
+
+        return hashed_wrapper
+
     def wrapper(func):
         predef_kws = {}
         predef_val = {}

hyperparameter/storage.py

@@ -102,6 +102,9 @@ def _update(values={}, prefix=None):
 
     def clear(self):
         self._storage.clear()
+        
+    def get_by_hash(self, *args, **kwargs):
+        raise RuntimeError("hyperparameter is not build with rust backend")
 
     def get(self, name: str, accessor: Callable = None) -> Any:
         if name in self.__slots__:
@@ -139,10 +142,21 @@ def frozen():
         GLOBAL_STORAGE.update(TLSKVStorage.tls.his[-1].storage())
 
 
+has_rust_backend = False
+
+
+def xxh64(*args, **kwargs):
+    raise RuntimeError("hyperparameter is not build with rust backend")
+
+
 try:
     if os.environ.get("HYPERPARAMETER_BACKEND", "RUST") == "RUST":
         from hyperparameter.rbackend import KVStorage
+        from hyperparameter.rbackend import xxh64
+
         TLSKVStorage = KVStorage
+        has_rust_backend = True
 except:
     import traceback
+
     traceback.print_exc()

src/ext.rs

@@ -18,6 +18,7 @@ use crate::storage::frozen_as_global_storage;
 use crate::storage::Storage;
 use crate::storage::StorageManager;
 use crate::storage::MGR;
+use crate::xxh::xxhstr;
 
 #[repr(C)]
 enum UserDefinedType {
@@ -127,6 +128,28 @@ impl KVStorage {
         }
     }
 
+    pub unsafe fn get_by_hash(&mut self, py: Python<'_>, hkey: u64) -> PyResult<Option<PyObject>> {
+        match self.storage.get_by_hash(hkey) {
+            Some(val) => match val {
+                Value::Empty => Err(PyValueError::new_err("not found")),
+                Value::Int(v) => Ok(Some(v.into_py(py))),
+                Value::Float(v) => Ok(Some(v.into_py(py))),
+                Value::Text(v) => Ok(Some(v.into_py(py))),
+                Value::Boolen(v) => Ok(Some(v.into_py(py))),
+                Value::UserDefined(v, k, _) => {
+                    if k == UserDefinedType::PyObjectType as i32 {
+                        Ok(Some(
+                            PyAny::from_borrowed_ptr(py, v as *mut pyo3::ffi::PyObject).into(),
+                        ))
+                    } else {
+                        Ok(Some((v as u64).into_py(py)))
+                    }
+                }
+            },
+            None => Err(PyValueError::new_err("not found")),
+        }
+    }
+
     pub unsafe fn put(&mut self, key: String, val: &PyAny) -> PyResult<()> {
         if self.isview {
             MGR.with(|mgr: &RefCell<StorageManager>| mgr.borrow_mut().put_key(key.clone()));
@@ -172,8 +195,14 @@ impl KVStorage {
     }
 }
 
+#[pyfunction]
+pub fn xxh64(s: &str) -> u64 {
+    xxhstr(s)
+}
+
 #[pymodule]
 fn rbackend(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<KVStorage>()?;
+    m.add_function(wrap_pyfunction!(xxh64, m)?)?;
     Ok(())
 }

tests/test_auto_param.py

@@ -0,0 +1,28 @@
+from unittest import TestCase
+
+from hyperparameter import auto_param, param_scope
+
+
+class TestAutoParam(TestCase):
+    def test_auto_param_func(self):
+        @auto_param("foo")
+        def foo(a, b=1, c=2.0, d=False, e="str"):
+            return a, b, c, d, e
+
+        with param_scope(**{"foo.b": 2}):
+            self.assertEqual(foo(1), (1, 2, 2.0, False, "str"))
+
+        with param_scope(**{"foo.c": 3.0}):
+            self.assertEqual(foo(1), (1, 1, 3.0, False, "str"))
+
+    def test_auto_param_func2(self):
+        @auto_param("foo")
+        def foo(a, b=1, c=2.0, d=False, e="str"):
+            return a, b, c, d, e
+
+        with param_scope():
+            param_scope.foo.b = 2
+            self.assertEqual(foo(1), (1, 2, 2.0, False, "str"))
+            param_scope.foo.c = 3.0
+            self.assertEqual(foo(1), (1, 2, 3.0, False, "str"))
+        self.assertEqual(foo(1), (1, 1, 2.0, False, "str"))