Refactored to bring common variables together. Also added a utility to get all the trace records as a list of records

src/architectures.py

@@ -197,27 +197,39 @@ class Architecture:
     architectures = None
     save_repo = None
     save_repo_load_error = False
+    save_repo_url = "https://huggingface.co/datasets/alfraser/llm-arch-trace"
+    trace_dir = "trace"
+    trace_file_name = "trace.json"
+    trace_file = os.path.join(trace_dir, trace_file_name)
 
     @classmethod
     def wipe_trace(cls):
-        repo_url = "https://huggingface.co/datasets/alfraser/llm-arch-trace"
-        trace_dir = "trace"
-        trace_file_name = "trace.json"
-        data_file = os.path.join(trace_dir, trace_file_name)
-        if os.path.exists(trace_dir):
-            shutil.rmtree(trace_dir)
-
+        if os.path.exists(cls.trace_dir):
+            shutil.rmtree(cls.trace_dir)
         try:
             hf_write_token = hf_api_token(write=True)
-            Architecture.save_repo = Repository(local_dir=trace_dir, clone_from=repo_url, token=hf_write_token)
+            cls.save_repo = Repository(local_dir=cls.trace_dir, clone_from=cls.save_repo_url, token=hf_write_token)
             test_json = {'tests': []}
-            with open(data_file, 'w') as f:
+            with open(cls.trace_file, 'w') as f:
                 json.dump(test_json, f, indent=2)
-            Architecture.save_repo.push_to_hub()
+            cls.save_repo.push_to_hub()
         except Exception as err:
-            Architecture.save_repo_load_error = True
+            cls.save_repo_load_error = True
             print(f"Error connecting to the save repo {err} - persistence now disabled")
 
+    @classmethod
+    def get_trace_records(cls) -> List[Dict]:
+        if not os.path.isfile(cls.trace_file):
+            hf_write_token = hf_api_token(write=True)
+            try:
+                cls.save_repo = Repository(local_dir=cls.trace_dir, clone_from=cls.save_repo_url, token=hf_write_token)
+            except Exception as err:
+                cls.save_repo_load_error = True
+                print(f"Error connecting to the save repo {err} - persistence now disabled")
+                return []
+        with open(cls.trace_file, 'r') as f:
+            test_json = json.load(f)
+            return test_json['tests']
 
     @classmethod
     def load_architectures(cls, force_reload: bool = False) -> None:
@@ -262,26 +274,23 @@ class Architecture:
                 return a
         raise ValueError(f"Could not find an architecture named {name}")
 
-    def append_and_save_data_as_json(self, data: Dict):
-        repo_url = "https://huggingface.co/datasets/alfraser/llm-arch-trace"
-        trace_dir = "trace"
-        trace_file_name = "trace.json"
-        data_file = os.path.join(trace_dir, trace_file_name)
-        if Architecture.save_repo is None and not Architecture.save_repo_load_error:
+    @classmethod
+    def append_and_save_data_as_json(cls, data: Dict):
+        if cls.save_repo is None and not cls.save_repo_load_error:
             try:
                 hf_write_token = hf_api_token(write=True)
-                Architecture.save_repo = Repository(local_dir=trace_dir, clone_from=repo_url, token=hf_write_token)
+                cls.save_repo = Repository(local_dir=cls.trace_dir, clone_from=cls.repo_url, token=hf_write_token)
             except Exception as err:
-                Architecture.save_repo_load_error = True
+                cls.save_repo_load_error = True
                 print(f"Error connecting to the save repo {err} - persistence now disabled")
 
-        if Architecture.save_repo is not None:
-            with open(data_file, 'r') as f:
+        if cls.save_repo is not None:
+            with open(cls.data_file, 'r') as f:
                 test_json = json.load(f)
             test_json['tests'].append(data)
-            with open(data_file, 'w') as f:
+            with open(cls.data_file, 'w') as f:
                 json.dump(test_json, f, indent=2)
-            Architecture.save_repo.push_to_hub()
+            cls.save_repo.push_to_hub()
 
     def attempt_request_and_trace_save(self, request: ArchitectureRequest, trace: ArchitectureTrace,
                                        trace_tags: List[str] = None, trace_comment: str = None) -> None: