Still debugging, but turned off background loading and continuing to load cache entries

Author: josephmaa

bridger/logging_utils/cache_entry_database.py

@@ -1,4 +1,5 @@
 from bridger.go_explore_phase_1 import CacheEntry
+import time
 
 
 class CacheEntryDatabase:
@@ -7,27 +8,68 @@ class CacheEntryDatabase:
     """
 
     def __init__(self, cache_entries: list[CacheEntry]):
+        print(f"Initializing CacheEntryDatabase with {len(cache_entries)} entries")
+        start_time = time.time()
+
+        # Debug info about the entries
+        if cache_entries:
+            print(f"First entry type: {type(cache_entries[0])}")
+            print(f"First entry attributes: {dir(cache_entries[0])}")
+            try:
+                print(
+                    f"First entry state shape: {cache_entries[0].state_representative.shape}"
+                )
+            except Exception as e:
+                print(f"Error getting state shape: {e}")
+
         self.cache_entries = cache_entries
+        end_time = time.time()
+        print(
+            f"CacheEntryDatabase initialization took {end_time - start_time:.2f} seconds"
+        )
 
     def sort_by_key(self, sort_key: "SortKey"):
         """
         Sorts the cache entries using the provided sort key.
         """
-        self.cache_entries = sorted(self.cache_entries, key=sort_key)
+        print(f"Sorting cache entries by {sort_key.key}")
+        start_time = time.time()
+        try:
+            self.cache_entries = sorted(self.cache_entries, key=sort_key)
+            end_time = time.time()
+            print(f"Sorting completed in {end_time - start_time:.2f} seconds")
+        except Exception as e:
+            print(f"Error during sorting: {e}")
+            raise
 
     def get_top_n_by_sort_key(self, sort_key: "SortKey", n: int):
         """
         Returns the top n cache entries.
         """
-        return self.sort_by_key(sort_key)[:n]
+        print(f"Getting top {n} entries by {sort_key.key}")
+        start_time = time.time()
+        try:
+            sorted_entries = self.sort_by_key(sort_key)
+            result = sorted_entries[:n]
+            end_time = time.time()
+            print(f"Retrieved top {n} entries in {end_time - start_time:.2f} seconds")
+            return result
+        except Exception as e:
+            print(f"Error getting top entries: {e}")
+            raise
 
 
 class SortKey:
     def __init__(self, key: str):
         self.key = key
 
     def __call__(self, cache_entry: CacheEntry):
-        return getattr(cache_entry, self.key)
+        try:
+            value = getattr(cache_entry, self.key)
+            return value
+        except Exception as e:
+            print(f"Error accessing {self.key} on cache entry: {e}")
+            raise
 
 
 class TrajectorySortKey(SortKey):

bridger/logging_utils/object_log_readers.py

@@ -28,23 +28,10 @@
 
 def read_object_log(log_filepath: str):
     with open(log_filepath, "rb") as f:
-        buffer = None
         while True:
             try:
-                # Try to load with torch.load first to handle CUDA tensors
-                try:
-                    buffer = torch.load(f, map_location=torch.device("cpu"))
-                except:
-                    # If that fails, try regular pickle.load
-                    f.seek(0)  # Reset file pointer
-                    buffer = pickle.load(f)
-                    # Move any tensors to CPU
-                    if isinstance(buffer, list):
-                        for i, item in enumerate(buffer):
-                            if isinstance(item, torch.Tensor):
-                                buffer[i] = item.cpu()
-                    elif isinstance(buffer, torch.Tensor):
-                        buffer = buffer.cpu()
+                print(f"This is the type: {type(f)}")
+                buffer = pickle.load(f)
 
                 for element in buffer:
                     yield element
@@ -311,6 +298,7 @@ def __init__(self, dirname: str):
         for entry in _read_object_log(
             os.path.dirname(dirname), log_entry.STATE_NORMALIZED_LOG_ENTRY
         ):
+            print(f"")
             self._states[entry.id] = entry.object
 
         # Store visited states sorted by visit count.

tools/web/sibyl.py

@@ -21,6 +21,7 @@
     StepsSinceLedToSomethingNewResetCountSortKey,
 )
 from bridger.logging_utils.object_log_readers import read_object_log
+from bridger.go_explore_phase_1 import StateSampler
 
 app = flask.Flask(__name__)
 
@@ -80,30 +81,32 @@ def n_fewest_steps_since_led_to_something_new_go_explore_plot_data():
     """
     Provides plot data for the n states that have most recently seen a new cell.
     """
-    experiment_name = _get_string_or_default(_EXPERIMENT_NAME)
     n = _get_int_or_default("n", 10)
-
-    cache_entry_database = CacheEntryDatabase(
-        list(read_object_log(os.path.join(_LOG_DIR, "go_explore", "start_entries.pkl")))
+    _CACHE_ENTRY_DATABASE = list(
+        read_object_log(os.path.join(_LOG_DIR, "go_explore", "state_cache-6.pkl"))
     )
+    print(_CACHE_ENTRY_DATABASE)
 
     return {
         "states": [
-            cache_entry.state for cache_entry in cache_entry_database.cache_entries
+            cache_entry.object.state
+            for cache_entry in _CACHE_ENTRY_DATABASE.cache_entries
         ],
-        "trajectory_length": cache_entry_database.get_top_n_by_sort_key(
-            TrajectorySortKey, n
-        ),
-        "steps_since_led_to_something_new": cache_entry_database.get_top_n_by_sort_key(
-            StepsSinceLedToSomethingNewSortKey, n
-        ),
-        "steps_since_led_to_something_new_reset_count": cache_entry_database.get_top_n_by_sort_key(
-            StepsSinceLedToSomethingNewResetCountSortKey, n
-        ),
-        "sample_count": cache_entry_database.get_top_n_by_sort_key(
-            SampleCountSortKey, n
-        ),
-        "visit_count": cache_entry_database.get_top_n_by_sort_key(VisitCountSortKey, n),
+        # "trajectory_length": _CACHE_ENTRY_DATABASE.get_top_n_by_sort_key(
+        #     TrajectorySortKey, n
+        # ),
+        # "steps_since_led_to_something_new": _CACHE_ENTRY_DATABASE.get_top_n_by_sort_key(
+        #     StepsSinceLedToSomethingNewSortKey, n
+        # ),
+        # "steps_since_led_to_something_new_reset_count": _CACHE_ENTRY_DATABASE.get_top_n_by_sort_key(
+        #     StepsSinceLedToSomethingNewResetCountSortKey, n
+        # ),
+        # "sample_count": _CACHE_ENTRY_DATABASE.get_top_n_by_sort_key(
+        #     SampleCountSortKey, n
+        # ),
+        # "visit_count": _CACHE_ENTRY_DATABASE.get_top_n_by_sort_key(
+        #     VisitCountSortKey, n
+        # ),
     }
 
 
@@ -447,6 +450,6 @@ def go_explore():
         target=_OBJECT_LOG_CACHE.convert_logs_to_saved_databases,
         args=(_get_experiment_names(),),
     )
-    convert_logs_background_thread.start()
+    # convert_logs_background_thread.start()
 
     app.run(host="0.0.0.0", port=6006)