Python: Allow configuration of storage thresholds.

pspoerri · pspoerri · commit faa14aa71fa7 · 2024-04-17T20:05:36.000+02:00
Signed-off-by: Pascal Spo&#776;rri &lt;psp@zurich.ibm.com&gt;
diff --git a/src/python/geds_smart_open/src/geds_smart_open/geds.py b/src/python/geds_smart_open/src/geds_smart_open/geds.py
@@ -225,8 +225,15 @@ def init_geds(cls):
             if BLOCK_SIZE is not None:
                 config.cache_block_size = int(BLOCK_SIZE)
 
-            # Init GEDS
+            GEDS_AVAILABLE_STORAGE = os.environ.get("GEDS_AVAILABLE_STORAGE")
+            if GEDS_AVAILABLE_STORAGE is not None:
+                config.available_local_storage = int(GEDS_AVAILABLE_STORAGE)
+
+            GEDS_AVAILABLE_MEMORY = os.environ.get("GEDS_AVAILABLE_MEMORY")
+            if GEDS_AVAILABLE_MEMORY is not None:
+                config.available_local_memory = int(GEDS_AVAILABLE_MEMORY)
 
+            # Init GEDS
             cls._geds = pygeds.GEDS(config)
             try:
                 cls._geds.start()
@@ -266,10 +273,12 @@ def register_object_store(
     def object_store_mapped(cls, bucket: str) -> bool:
         return bucket in cls._known_s3_buckets
 
+
 @atexit.register
 def handle_shutdown():
     GEDSInstance.handle_shutdown()
 
+
 def register_object_store(
     bucket: str, endpoint_url: str, access_key: str, secret_key: str
 ):
diff --git a/src/python/wrapper.cpp b/src/python/wrapper.cpp
@@ -45,7 +45,9 @@ PYBIND11_MODULE(pygeds, m) {
       .def_readwrite("port_http_server", &GEDSConfig::portHttpServer)
       .def_readwrite("local_storage_path", &GEDSConfig::localStoragePath)
       .def_readwrite("cache_block_size", &GEDSConfig::cacheBlockSize)
-      .def_readwrite("cache_objects_from_s3", &GEDSConfig::cache_objects_from_s3);
+      .def_readwrite("cache_objects_from_s3", &GEDSConfig::cache_objects_from_s3)
+      .def_readwrite("available_local_storage", &GEDSConfig::available_local_storage)
+      .def_readwrite("available_local_memory", &GEDSConfig::available_local_memory);
 
   py::class_<GEDS, std::shared_ptr<GEDS>>(m, "GEDS")
       .def_property_readonly_static(
