Merge branch 'main' into wandb-state-logging

.github/workflows/pytype-miner.yaml

@@ -1,9 +1,8 @@
 name: Test Miner Types
 
 on:
-  pull_request:
-    types:
-      - opened
+  - pull_request
+  - workflow_dispatch
 
 jobs:
   test-types:

.github/workflows/pytype-neuron.yaml

@@ -1,9 +1,8 @@
 name: Test Common Types
 
 on:
-  pull_request:
-    types:
-      - opened
+  - pull_request
+  - workflow_dispatch
 
 jobs:
   test-types:

.github/workflows/pytype-validator.yaml

@@ -1,9 +1,8 @@
 name: Test Validator Types
 
 on:
-  pull_request:
-    types:
-      - opened
+  - pull_request
+  - workflow_dispatch
 
 jobs:
   test-types:

neuron/neuron/contest.py

@@ -3,6 +3,7 @@
 from enum import Enum
 from functools import partial
 from io import BytesIO
+from typing import ContextManager
 
 from pydantic import BaseModel
 from transformers import CLIPProcessor, CLIPVisionModelWithProjection
@@ -19,7 +20,7 @@ class ContestId(Enum):
     FLUX_NVIDIA_4090 = 2
 
 
-class OutputComparator(ABC):
+class OutputComparator(ContextManager, ABC):
     @abstractmethod
     def compare(self, baseline: bytes, optimized: bytes) -> float:
         pass
@@ -87,7 +88,6 @@ def __exit__(self, exc_type, exc_value, traceback):
         getattr(torch, self.device).empty_cache()
 
 
-
 class Contest(ABC):
     id: ContestId
     output_comparator: Callable[[], OutputComparator]

validator/base_validator/metrics.py

@@ -1,7 +1,10 @@
 from enum import Enum
 
 from neuron import Key
-from neuron.submission_tester.metrics import *
+from neuron.submission_tester.metrics import *  # noqa
+from neuron.submission_tester.metrics import CheckpointBenchmark, MetricData
+
+from pydantic import BaseModel
 
 
 class BenchmarkState(Enum):