Move circuit weights into metadata (#52)

Author: millioner

docs/custom_circuit_integrations.md

@@ -73,17 +73,18 @@ There are three primary functions which must be implemented.
 
 This file is responsible for describing the circuit's properties, and contains both optional and required fields.
 
-| Field             | Type            | Description                                                                          | Required |
-| ----------------- | --------------- | ------------------------------------------------------------------------------------ | -------- |
-| `name`            | `string`        | The name of the circuit.                                                             | ✅       |
-| `description`     | `string`        | A short description of the circuit.                                                  | ✅       |
-| `author`          | `string`        | The author of the circuit.                                                           | ✅       |
-| `version`         | `string`        | The version of the circuit.                                                          | ✅       |
-| `proof_system`    | [`ProofSystem`] | The proof system used by the circuit.                                                | ✅       |
-| `type`            | [`CircuitType`] | The type of circuit. New circuits should use `PROOF_OF_COMPUTATION` by default       | ✅       |
-| `external_files`  | `dict`          | A dictionary of external files required by the circuit.                              | ✅       |
-| `netuid`          | `int`           | For Proof of Weights, the netuid of the target subnet.                               | ❌       |
-| `weights_version` | `int`           | For Proof of Weights, the version of subnet weights that the circuit corresponds to. | ❌       |
+| Field                     | Type            | Description                                                                          | Required |
+| ------------------------- | --------------- | ------------------------------------------------------------------------------------ | -------- |
+| `name`                    | `string`        | The name of the circuit.                                                             | ✅       |
+| `description`             | `string`        | A short description of the circuit.                                                  | ✅       |
+| `author`                  | `string`        | The author of the circuit.                                                           | ✅       |
+| `version`                 | `string`        | The version of the circuit.                                                          | ✅       |
+| `proof_system`            | [`ProofSystem`] | The proof system used by the circuit.                                                | ✅       |
+| `type`                    | [`CircuitType`] | The type of circuit. New circuits should use `PROOF_OF_COMPUTATION` by default       | ✅       |
+| `external_files`           | `dict`          | A dictionary of external files required by the circuit.                               | ✅       |
+| `netuid`                  | `int`           | For Proof of Weights, the netuid of the target subnet.                               | ❌       |
+| `weights_version`         | `int`           | For Proof of Weights, the version of subnet weights that the circuit corresponds to. | ❌       |
+| `benchmark_choice_weight` | `float`          | The probability of the circuit being selected for benchmark request.                 | ❌       |
 
 [See it's class definition here for more information](https://github.com/inference-labs-inc/omron-subnet/blob/main/neurons/execution_layer/circuit.py#L100)
 
@@ -120,7 +121,7 @@ Ensure all the following files are copied into this directory or are listed in t
 
 ### Adding benchmark requests
 
-The last step after adding your circuit is to add an allocation of benchmark requests to the circuit. This is done through a modification to the `neurons/constants.py` file. In the `CIRCUIT_WEIGHTS` dictionary, add a new key-value pair where the key is the `SHA256` hash of the circuit's verification key, and the value is the percentage of requests that should be allocated to the circuit. Percentages should be between 0 and 1 and do not need to sum to 1. A reasonable starting point is 0.20 for each circuit. The omron team will adjust these weights as necessary to ensure a healthy distribution of requests across all circuits.
+The last step after adding your circuit is to add an allocation of benchmark requests to the circuit. This is done through a modification to the circuit `metadata.json` file. The `benchmark_choice_weight` value is the percentage of requests that should be allocated to the circuit. Percentages should be between 0 and 1 and sum of them for all models do not need to be equal to 1. A reasonable starting point is 0.20 for each circuit. The omron team will adjust these weights as necessary to ensure a healthy distribution of requests across all circuits.
 
 A higher percentage will result in more requests being allocated to the circuit, increasing the circuit's optimization pressure.
 
@@ -129,9 +130,8 @@ A higher percentage will result in more requests being allocated to the circuit,
 To update an existing circuit:
 
 1. Follow the same steps as adding a new circuit, which will result in a new verification key hash
-2. In the `CIRCUIT_WEIGHTS` dictionary, set the old circuit's allocation to 0.0
-3. Add the new circuit with the desired benchmark allocation
-4. Open a PR with these changes
+2. Add the new circuit with the desired benchmark allocation
+3. Open a PR with these changes
 
 ## Querying a circuit
 

neurons/_validator/core/request_pipeline.py

@@ -1,25 +1,26 @@
 from __future__ import annotations
-from _validator.pow.proof_of_weights_handler import ProofOfWeightsHandler
-import bittensor as bt
+
+import copy
 import random
-from protocol import ProofOfWeightsSynapse, QueryZkProof
 
-from _validator.scoring.score_manager import ScoreManager
+import bittensor as bt
+
 from _validator.api import ValidatorAPI
 from _validator.config import ValidatorConfig
+from _validator.core.request import Request
+from _validator.models.request_type import RequestType
+from _validator.pow.proof_of_weights_handler import ProofOfWeightsHandler
+from _validator.scoring.score_manager import ScoreManager
+from _validator.utils.hash_guard import HashGuard
 from constants import (
     BATCHED_PROOF_OF_WEIGHTS_MODEL_ID,
-    CIRCUIT_WEIGHTS,
     SINGLE_PROOF_OF_WEIGHTS_MODEL_ID,
 )
 from deployment_layer.circuit_store import circuit_store
 from execution_layer.circuit import Circuit, CircuitType
 from execution_layer.generic_input import GenericInput
-from _validator.utils.hash_guard import HashGuard
-from _validator.core.request import Request
+from protocol import ProofOfWeightsSynapse, QueryZkProof
 from utils.wandb_logger import safe_log
-from _validator.models.request_type import RequestType
-import copy
 
 
 class RequestPipeline:
@@ -124,13 +125,16 @@ def select_circuit_for_benchmark(self) -> Circuit:
         """
         Select a circuit for benchmarking using weighted random selection.
         """
-
-        circuit_id = random.choices(
-            list(CIRCUIT_WEIGHTS.keys()), weights=list(CIRCUIT_WEIGHTS.values()), k=1
+        circuits = list(circuit_store.circuits.values())
+
+        return random.choices(
+            circuits,
+            weights=[
+                (circuit.metadata.benchmark_choice_weight or 0) for circuit in circuits
+            ],
+            k=1,
         )[0]
 
-        return circuit_store.get_circuit(circuit_id)
-
     def format_for_query(
         self, inputs: dict[str, object], circuit: Circuit
     ) -> dict[str, object]:

neurons/cli_parser.py

@@ -87,6 +87,12 @@ def init_config(role: Optional[str] = None):
         default=False,
         help="Whether to run the miner in localnet mode.",
     )
+    parser.add_argument(
+        "--timeout",
+        default=120,
+        type=int,
+        help="Timeout for requests in seconds (default: 120)",
+    )
     parser.add_argument(
         "--external-model-dir",
         default=None,
@@ -123,7 +129,6 @@ def init_config(role: Optional[str] = None):
         config.eth_wallet = (
             config.eth_wallet if config.eth_wallet is not None else "0x002"
         )
-        config.timeout = config.timeout if config.timeout is None else 120
         config.disable_wandb = True
         config.verbose = config.verbose if config.verbose is None else True
         config.max_workers = config.max_workers or 1
@@ -140,6 +145,9 @@ def init_config(role: Optional[str] = None):
     else:
         config.full_path_models = os.path.join(config.full_path, "models")
 
+    if config.whitelisted_public_keys:
+        config.whitelisted_public_keys = config.whitelisted_public_keys.split(",")
+
     os.makedirs(config.full_path, exist_ok=True)
     os.makedirs(config.full_path_score, exist_ok=True)
     os.makedirs(config.full_path_models, exist_ok=True)
@@ -267,6 +275,13 @@ def _validator_config():
         ),
     )
 
+    parser.add_argument(
+        "--whitelisted-public-keys",
+        type=str,
+        default=None,
+        help="A comma-separated list of public keys to whitelist for external requests.",
+    )
+
     parser.add_argument(
         "--certificate-path",
         type=str,

neurons/constants.py

@@ -68,16 +68,6 @@ class Roles:
 MAX_PROOFS_TO_LOG = 0
 # Era period for proof of weights (mortality of the pow log)
 PROOF_OF_WEIGHTS_LIFESPAN = 2
-# Weights that determine the probability of selecting a circuit for benchmarking
-CIRCUIT_WEIGHTS = {
-    "1d60d545b7c5123fd60524dcbaf57081ca7dc4a9ec36c892927a3153328d17c0": 0,
-    "33b92394b18412622adad75733a6fc659b4e202b01ee8a5465958a6bad8ded62": 0.20,
-    "37320fc74fec80805eedc8e92baf3c58842a2cb2a4ae127ad6e930f0c8441c7a": 0,
-    "a4bcecaf699fd9212600a1f2fcaa40c444e1aeaab409ea240a38c33ed356f4e2": 0.20,
-    "a849500803abdbb86a9460e18684a6411dc7ae0b75f1f6330e3028081a497dea": 0.20,
-    "e84b2e5f223621fa20078eb9f920d8d4d3a4ff95fa6e2357646fdbb43a2557c9": 0.20,
-    "8dcff627a782525ea86196941a694ffbead179905f0cd4550ddc3df9e2b90924": 0.20,
-}
 # Maximum signature lifespan for WebSocket requests
 MAX_SIGNATURE_LIFESPAN = 300
 # Whitelisted public keys (ss58 addresses) we accept external requests from by default

neurons/deployment_layer/model_1d60d545b7c5123fd60524dcbaf57081ca7dc4a9ec36c892927a3153328d17c0/metadata.json

@@ -6,5 +6,6 @@
     "proof_system": "JOLT",
     "netuid": 27,
     "type": "proof_of_weights",
-    "external_files": {}
+    "external_files": {},
+    "benchmark_choice_weight": 0
 }

neurons/deployment_layer/model_33b92394b18412622adad75733a6fc659b4e202b01ee8a5465958a6bad8ded62/metadata.json

@@ -9,5 +9,6 @@
     "type": "proof_of_weights",
     "external_files": {
         "pk.key": "https://storage.omron.ai/33b92394b18412622adad75733a6fc659b4e202b01ee8a5465958a6bad8ded62/pk.key"
-    }
+    },
+    "benchmark_choice_weight": 0.2
 }

neurons/deployment_layer/model_37320fc74fec80805eedc8e92baf3c58842a2cb2a4ae127ad6e930f0c8441c7a/metadata.json

@@ -5,5 +5,6 @@
     "version": "0.0.5",
     "proof_system": "JOLT",
     "type": "proof_of_computation",
-    "external_files": {}
+    "external_files": {},
+    "benchmark_choice_weight": 0
 }

neurons/deployment_layer/model_8dcff627a782525ea86196941a694ffbead179905f0cd4550ddc3df9e2b90924/metadata.json

@@ -7,5 +7,6 @@
     "type": "proof_of_computation",
     "external_files": {
         "pk.key": "https://storage.omron.ai/8dcff627a782525ea86196941a694ffbead179905f0cd4550ddc3df9e2b90924/pk.key"
-    }
+    },
+    "benchmark_choice_weight": 0.2
 }

neurons/deployment_layer/model_a4bcecaf699fd9212600a1f2fcaa40c444e1aeaab409ea240a38c33ed356f4e2/metadata.json

@@ -9,5 +9,6 @@
     "type": "proof_of_weights",
     "external_files": {
         "pk.key": "https://storage.omron.ai/a4bcecaf699fd9212600a1f2fcaa40c444e1aeaab409ea240a38c33ed356f4e2/pk.key"
-    }
+    },
+    "benchmark_choice_weight": 0.2
 }

neurons/deployment_layer/model_a849500803abdbb86a9460e18684a6411dc7ae0b75f1f6330e3028081a497dea/metadata.json

@@ -7,5 +7,6 @@
     "type": "proof_of_computation",
     "external_files": {
         "circuit.zkey": "https://storage.omron.ai/a849500803abdbb86a9460e18684a6411dc7ae0b75f1f6330e3028081a497dea/circuit.zkey"
-    }
+    },
+    "benchmark_choice_weight": 0.2
 }

neurons/deployment_layer/model_e84b2e5f223621fa20078eb9f920d8d4d3a4ff95fa6e2357646fdbb43a2557c9/metadata.json

@@ -9,5 +9,6 @@
     "type": "proof_of_weights",
     "external_files": {
         "circuit.zkey": "https://storage.omron.ai/e84b2e5f223621fa20078eb9f920d8d4d3a4ff95fa6e2357646fdbb43a2557c9/circuit.zkey"
-    }
+    },
+    "benchmark_choice_weight": 0.2
 }

neurons/execution_layer/circuit.py

@@ -122,6 +122,7 @@ class CircuitMetadata:
     external_files: dict[str, str]
     netuid: int | None = None
     weights_version: int | None = None
+    benchmark_choice_weight: float | None = None
 
     @classmethod
     def from_file(cls, metadata_path: str) -> CircuitMetadata: