[SYSTEMDS-3940] Improve Scuro Node Executor

src/main/python/systemds/scuro/drsearch/hyperparameter_tuner.py

@@ -297,6 +297,7 @@ def __init__(
         wandb_entity: Optional[str] = None,
         wandb_group: Optional[str] = None,
         wandb_tags: Optional[List[str]] = None,
+        enable_checkpointing: bool = False,
     ):
         self.tasks = tasks
         self.unimodal_optimization_results = optimization_results
@@ -333,6 +334,7 @@ def __init__(
         self.wandb_group = wandb_group
         self.wandb_tags = wandb_tags or []
         self._wandb_run = None
+        self.enable_checkpointing = enable_checkpointing
 
     def get_modalities_by_id(self, modality_ids: List[int]) -> Modality:
         modalities = []
@@ -415,14 +417,18 @@ def tune_unimodal_representations(self, max_eval_per_rep: Optional[int] = None):
                             )
                         )
                     self.optimization_results.add_result(results)
-                    self._checkpoint_manager.increment(task.model.name, len(results))
-                    self._checkpoint_manager.checkpoint_if_due(
-                        self.optimization_results.results,
-                    )
+                    if self.enable_checkpointing:
+                        self._checkpoint_manager.increment(
+                            task.model.name, len(results)
+                        )
+                        self._checkpoint_manager.checkpoint_if_due(
+                            self.optimization_results.results,
+                        )
                 except Exception:
-                    self._checkpoint_manager.save_checkpoint(
-                        self.optimization_results.results, {}
-                    )
+                    if self.enable_checkpointing:
+                        self._checkpoint_manager.save_checkpoint(
+                            self.optimization_results.results, {}
+                        )
                     raise
 
         if self.save_results:

src/main/python/systemds/scuro/drsearch/modality_result_cache.py

@@ -0,0 +1,136 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+from typing import Any, Dict, List, Optional
+
+from systemds.scuro.drsearch.modality_shared_memory import unlink_shm
+from systemds.scuro.utils.static_variables import DEBUG
+
+
+class RefCountResultCache:
+    def __init__(self):
+        self.cache: Dict[str, Any] = {}
+        self.ref_count: Dict[str, int] = {}
+        self.memory_usage_per_node: Dict[str, int] = {}
+        self.shared_memory_names: Dict[str, List[str]] = {}
+        self._shm_retain_count: Dict[str, int] = {}
+
+    def get(self, node_id: str) -> Any:
+        return self.cache[node_id]
+
+    def add_result(
+        self,
+        node_id: str,
+        result: Any,
+        shm_name: Optional[str] = None,
+        resident_bytes: Optional[int] = None,
+        shm_bytes: int = 0,
+    ):
+        if shm_name is not None:
+            self.shared_memory_names[node_id] = [shm_name]
+        self.cache[node_id] = result
+        self.memory_usage_per_node[node_id] = int(resident_bytes or 0) + int(
+            shm_bytes or 0
+        )
+        if DEBUG:
+            print(
+                f"Node {node_id} has a CPU memory usage of "
+                f"{self.memory_usage_per_node[node_id]/1024**3:.5f} GB"
+                + (
+                    f" ({int(shm_bytes or 0)/1024**3:.5f} GB of it shared memory)"
+                    if shm_name is not None
+                    else ""
+                )
+            )
+
+    def inc_ref(self, node_id: str):
+        self.ref_count[node_id] = self.ref_count.get(node_id, 0) + 1
+
+    def dec_ref(self, node_id: str):
+        if node_id not in self.ref_count:
+            return
+        self.ref_count[node_id] -= 1
+        if self.ref_count[node_id] <= 0:
+            self.ref_count[node_id] = 0
+            self._try_cleanup_node(node_id)
+
+    def clear(self, node_id: str):
+        self.ref_count[node_id] = 0
+        self._try_cleanup_node(node_id)
+
+    def retain_shm_names(self, shm_names: List[str]) -> List[str]:
+        retained: List[str] = []
+        for shm_name in shm_names:
+            if not shm_name:
+                continue
+            self._shm_retain_count[shm_name] = (
+                self._shm_retain_count.get(shm_name, 0) + 1
+            )
+            retained.append(shm_name)
+        return retained
+
+    def release_shm_names(self, shm_names: List[str]) -> None:
+        nodes_to_recheck: List[str] = []
+        for shm_name in shm_names:
+            if not shm_name:
+                continue
+            count = self._shm_retain_count.get(shm_name, 0) - 1
+            if count <= 0:
+                self._shm_retain_count.pop(shm_name, None)
+            else:
+                self._shm_retain_count[shm_name] = count
+            for node_id, node_names in self.shared_memory_names.items():
+                if shm_name in node_names and node_id not in nodes_to_recheck:
+                    nodes_to_recheck.append(node_id)
+        for node_id in nodes_to_recheck:
+            self._try_cleanup_node(node_id)
+
+    def __len__(self):
+        return len(self.cache)
+
+    def get_memory_total_memory_usage(self):
+        return sum(self.memory_usage_per_node.values())
+
+    def _shm_names_in_use(self, shm_names: List[str]) -> bool:
+        return any(self._shm_retain_count.get(name, 0) > 0 for name in shm_names)
+
+    def _try_cleanup_node(self, node_id: str) -> None:
+        if self.ref_count.get(node_id, 0) > 0:
+            return
+        shm_names = self.shared_memory_names.get(node_id, [])
+        if shm_names and self._shm_names_in_use(shm_names):
+            return
+        self.cache.pop(node_id, None)
+        self.ref_count.pop(node_id, None)
+        self.memory_usage_per_node.pop(node_id, None)
+        self._cleanup_shared_memory(node_id)
+
+    def _cleanup_shared_memory(self, node_id: str):
+        names = self.shared_memory_names.pop(node_id, [])
+        for shm_name in names:
+            unlink_shm(shm_name)
+
+    def cleanup_all(self):
+        self._shm_retain_count.clear()
+        for node_id in list(self.shared_memory_names.keys()):
+            self.ref_count.pop(node_id, None)
+            self.cache.pop(node_id, None)
+            self.memory_usage_per_node.pop(node_id, None)
+            self._cleanup_shared_memory(node_id)

src/main/python/systemds/scuro/drsearch/modality_shared_memory.py

@@ -20,11 +20,29 @@
 # -------------------------------------------------------------
 from typing import Any, List, Tuple
 import numpy as np
-from multiprocessing import shared_memory
+from multiprocessing import shared_memory, resource_tracker
 
 SHARED_MEMORY_MIN_BYTES = 1 * 1024 * 1024
 
 
+def _untrack(shm: shared_memory.SharedMemory) -> None:
+    try:
+        resource_tracker.unregister(shm._name, "shared_memory")
+    except Exception:
+        pass
+
+
+def unlink_shm(name: str) -> None:
+    try:
+        shm = shared_memory.SharedMemory(name=name)
+        shm.close()
+        shm.unlink()
+    except FileNotFoundError:
+        pass
+    except Exception:
+        pass
+
+
 class SharedStringList:
     def __init__(
         self, shm_name: str, offsets: List[Tuple[int, int]], payload_nbytes: int
@@ -37,6 +55,7 @@ def __init__(
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
 
     def __len__(self):
         return len(self.offsets)
@@ -91,6 +110,7 @@ def __init__(
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._buffer = np.ndarray(
                 (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
             )
@@ -146,6 +166,7 @@ def __init__(self, shm_name: str, dtype_str: str, shape: tuple):
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._arr = np.ndarray(self.shape, dtype=self._dtype, buffer=self._shm.buf)
             self._arr.setflags(write=False)
 
@@ -215,6 +236,7 @@ def __init__(
     def _ensure_attached(self):
         if self._shm is None:
             self._shm = shared_memory.SharedMemory(name=self.shm_name)
+            _untrack(self._shm)
             self._buffer = np.ndarray(
                 (self.total_elems,), dtype=self._dtype, buffer=self._shm.buf
             )
@@ -330,6 +352,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool:
                 resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_shared_ndarray_candidate(data):
         arr = data
@@ -344,6 +367,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool:
             resident_bytes = min(resident_bytes, 2 * 1024 * 1024)
 
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_nested_shared_memory_candidate(data):
         leaves: List[np.ndarray] = []
@@ -379,6 +403,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool:
                 resident_bytes, max(2 * 1024 * 1024, len(offsets) * 64)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
     elif _is_string_list_shared_memory_candidate(data):
         encoded = [s.encode("utf-8") for s in data]
@@ -398,6 +423,7 @@ def add_shared_memory_candidate(data: Any, resident_bytes: int = 0) -> bool:
                 resident_bytes, max(2 * 1024 * 1024, len(str_offsets) * 32)
             )
             shm.close()
+            _untrack(shm)
             return data, shm.name, data_nbytes, resident_bytes
 
     return None, None, 0, resident_bytes