Implemented single threaded worker on writing the logs to the json file for controlled access to the resource on the file system now we are multi-threading the tests.

src/architectures.py

@@ -14,6 +14,8 @@ import traceback
 from abc import ABC, abstractmethod
 from enum import Enum
 from huggingface_hub import Repository
+from queue import Queue
+from threading import Thread, Timer
 from time import time
 from typing import List, Optional, Dict
 from better_profanity import profanity
@@ -185,6 +187,91 @@ class ArchitectureComponent(ABC):
         return ""
 
 
+class LogWorker(Thread):
+    instance = None
+    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)
+    queue = Queue()
+    commit_time = 10  # Number of seconds after which to commit with no activity
+    commit_after = 10  # Number of records after which to commit irrespective of time
+    commit_count = 0  # Current uncommitted records
+    commit_timer = None  # The actual commit timer - we will schedule the commit on this
+
+    def run(self):
+        while True:
+            request, trace, trace_tags, trace_comment = LogWorker.queue.get()
+            if request is None:
+                LogWorker.commit_repo()
+            else:
+                if LogWorker.commit_timer is not None and LogWorker.commit_timer.is_alive():
+                    LogWorker.commit_timer.cancel()
+                    LogWorker.commit_timer = None
+                try:
+                    save_dict = {
+                        'architecture': self.name,
+                        'request': request.as_dict(),
+                        'trace': trace.as_dict(),
+                        'test_tags': trace_tags,
+                        'test_comment': trace_comment
+                    }
+                    LogWorker.append_and_save_data_as_json(save_dict)
+                    LogWorker.commit_count += 1
+                    if LogWorker.commit_count >= LogWorker.commit_after:
+                        LogWorker.commit_repo()
+                except Exception as err:
+                    print(f"Request / trace save failed {err}")
+
+                LogWorker.commit_timer = Timer(LogWorker.commit_time, LogWorker.signal_commit)
+                LogWorker.commit_timer.start()
+
+    @classmethod
+    def append_and_save_data_as_json(cls, data: Dict):
+        print(f"LogWorker logging open record {LogWorker.commit_count + 1}")
+        if cls.save_repo is None and not cls.save_repo_load_error:
+            try:
+                hf_write_token = hf_api_token(write=True)
+                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")
+
+        if cls.save_repo is not None:
+            with open(cls.trace_file, 'r') as f:
+                test_json = json.load(f)
+            test_json['tests'].append(data)
+            with open(cls.trace_file, 'w') as f:
+                json.dump(test_json, f, indent=2)
+
+    @classmethod
+    def commit_repo(cls):
+        print(f"LogWorker committing {LogWorker.commit_count} open records")
+        cls.save_repo.push_to_hub()
+        LogWorker.commit_count = 0
+
+    @classmethod
+    def signal_commit(cls):
+        print("LogWorker signalling commit based on time elapsed")
+        cls.queue.put((None, None, None, None))
+
+    @classmethod
+    def write(cls, request: ArchitectureRequest, trace: ArchitectureTrace,
+              trace_tags: List[str] = None, trace_comment: str = None):
+        trace_tags = [] if trace_tags is None else trace_tags
+        trace_comment = "" if trace_comment is None else trace_comment
+        cls.queue.put((request, trace, trace_tags, trace_comment))
+
+
+# Instantiate and run worker on import
+if LogWorker.instance is None:
+    LogWorker.instance = LogWorker()
+    LogWorker.instance.start()
+
+
 class Architecture:
     """
     An architecture is built as a callable pipeline of steps. An
@@ -275,45 +362,6 @@ class Architecture:
                 return a
         raise ValueError(f"Could not find an architecture named {name}")
 
-    @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)
-                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")
-
-        if cls.save_repo is not None:
-            with open(cls.trace_file, 'r') as f:
-                test_json = json.load(f)
-            test_json['tests'].append(data)
-            with open(cls.trace_file, 'w') as f:
-                json.dump(test_json, f, indent=2)
-            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:
-        """
-        Attempt to save a request and trace pair to a json store on huggingface datasets
-        Catch any errors and simply print as non-fatal to functional flow
-        """
-        try:
-            if trace_tags is None:
-                trace_tags = []
-            if trace_comment is None:
-                trace_comment = ""
-            save_dict = {
-                'architecture': self.name,
-                'request': request.as_dict(),
-                'trace': trace.as_dict(),
-                'test_tags': trace_tags,
-                'test_comment': trace_comment
-            }
-            self.append_and_save_data_as_json(save_dict)
-        except Exception as err:
-            print(f"Request / trace save failed {err}")
 
     def __init__(self,
                  name: str,
@@ -355,7 +403,7 @@ class Architecture:
                 trace.steps[-1].exception = err
                 traceback.print_exc()
                 break
-        self.attempt_request_and_trace_save(request, trace, trace_tags, trace_comment)
+        LogWorker.write(request, trace, trace_tags, trace_comment)
         return trace
 
 

src/testing.py

@@ -9,14 +9,19 @@ import sys
 from huggingface_hub import Repository
 from queue import Queue
 from random import choices
-from threading import Thread
+from threading import Thread, Timer
 from typing import Dict, List, Optional, Tuple
 
-from src.architectures import Architecture, ArchitectureRequest
-from src.common import data_dir
+from src.architectures import Architecture, ArchitectureRequest, ArchitectureTrace
+from src.common import data_dir, hf_api_token
 
 
 class ArchitectureTestWorker(Thread):
+    """
+    This class is worker which takes a test request off the queue and passes
+    it to an architecture for execution.  Used to multi-thread the testing process
+    for speed as there is a tonne of i/o blocking waiting for the LLM
+    """
     def __init__(self, work_queue: Queue, worker_name: str, trace_tags: List[str], trace_comment: str):
         Thread.__init__(self)
         self.work_queue = work_queue
@@ -29,7 +34,7 @@ class ArchitectureTestWorker(Thread):
         while running:
             arch, request = self.work_queue.get()
             try:
-                if arch is None:
+                if arch is None:  # None passed to signal end of test requests
                     running = False
                 else:
                     print(f'{self.worker_name} running "{request.request}" through {arch}')