Merge branch 'rewrite' into dynamic-scoring

womboai · Nov 29, 2024 · 90e9222 · 90e9222
2 parents 0073713 + 5428f08
commit 90e9222
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/validator/weight_setting/contest_state.py b/validator/weight_setting/contest_state.py
@@ -2,16 +2,17 @@
 from threading import Event
 
 from fiber.logging_utils import get_logger
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from base.checkpoint import Key, current_time, Submissions, Benchmarks
-from base.contest import Metrics
+from base.contest import Metrics, BenchmarkState
 from weight_setting.winner_selection import get_contestant_scores, get_contestant_ranks, calculate_rank_weights
 
 logger = get_logger(__name__)
 
 
 class ContestState(BaseModel):
+    model_config = ConfigDict(extra="ignore")
     step: int
     benchmarks_version: int
     submissions: Submissions
@@ -20,6 +21,7 @@ class ContestState(BaseModel):
     invalid_submissions: set[Key]
     last_benchmarks: Benchmarks
     average_benchmarking_time: float | None
+    benchmarking_state: BenchmarkState
     contest_end: datetime
 
     def start_new_contest(self, benchmarks_version: int, submissions: Submissions):
@@ -89,6 +91,7 @@ def create(cls, benchmarks_version: int):
             invalid_submissions=set(),
             last_benchmarks={},
             average_benchmarking_time=None,
+            benchmarking_state=BenchmarkState.NOT_STARTED,
             contest_end=current_time()
         )
         return state
diff --git a/validator/weight_setting/validator.py b/validator/weight_setting/validator.py
@@ -112,6 +112,7 @@ def __init__(self):
 
     @tracer.start_as_current_span("initialize_contest")
     def initialize_contest(self):
+        self.contest_state.benchmarking_state = BenchmarkState.NOT_STARTED
         for api in self.benchmarking_apis:
             api.initialize(
                 uid=self.uid,
@@ -150,6 +151,8 @@ def do_step(self):
         untested_submissions = self.contest_state.get_untested_submissions()
 
         if not untested_submissions:
+            self.contest_state.benchmarking_state = BenchmarkState.FINISHED
+            self.wandb_manager.send_metrics(self.contest_state)
             self.contest_state.sleep_to_next_contest(self._stop_flag)
             return
 
@@ -164,6 +167,7 @@ def do_step(self):
             return
 
         self.update_benchmarks(benchmarking_results)
+        self.contest_state.benchmarking_state = BenchmarkState.IN_PROGRESS
         self._stop_flag.wait(BENCHMARK_UPDATE_RATE_BLOCKS * 12)
 
     def update_benchmarks(self, benchmarking_results: list[BenchmarkingResults]):