DSperce integration

.devcontainer/devcontainer.json

@@ -33,7 +33,7 @@
     }
   },
   // Install ezkl cli, done here instead of the Dockerfile to test new versions without rebuilding the image.
-  "onCreateCommand": "curl https://raw.githubusercontent.com/zkonduit/ezkl/main/install_ezkl_cli.sh | bash -s v19.0.7",
+  "onCreateCommand": "curl https://raw.githubusercontent.com/zkonduit/ezkl/main/install_ezkl_cli.sh | bash -s v22.2.1",
   "postCreateCommand": "uv tool install bittensor-cli",
   "remoteEnv": {
     "PATH": "${containerEnv:PATH}:/home/vscode/.ezkl"

neurons/_miner/miner_session.py

@@ -23,10 +23,12 @@
     SINGLE_PROOF_OF_WEIGHTS_MODEL_ID,
 )
 from deployment_layer.circuit_store import circuit_store
+from execution_layer.dsperse_manager import DSperseManager
 from execution_layer.generic_input import GenericInput
 from execution_layer.verified_model_session import VerifiedModelSession
 from protocol import (
     Competition,
+    DSliceProofGenerationDataModel,
     ProofOfWeightsDataModel,
     QueryForCapacities,
     QueryZkProof,
@@ -48,6 +50,7 @@ def __init__(self):
         self.configure()
         self.check_register(should_exit=True)
         self.auto_update = AutoUpdate()
+        self.dsperse_manager = DSperseManager()
         self.log_batch = []
         self.shuffled_uids = None
         self.last_shuffle_epoch = -1
@@ -80,6 +83,10 @@ def start_server(self) -> bool:
         self.server.register_route(
             path=f"/{QueryForCapacities.name}", endpoint=self.handleCapacityRequest
         )
+        self.server.register_route(
+            path=f"/{DSliceProofGenerationDataModel.name}",
+            endpoint=self.handleDSliceRequest,
+        )
         self.server.start()
 
         existing_miner = self.metagraph.axons[self.subnet_uid]
@@ -466,6 +473,24 @@ def handleCompetitionRequest(self, data: Competition) -> JSONResponse:
                 status_code=500,
             )
 
+    def handleDSliceRequest(self, data: DSliceProofGenerationDataModel) -> JSONResponse:
+        """
+        Handle DSlice proof generation requests from validators.
+        """
+        bt.logging.info(
+            f"Handling DSlice slice proof generation request for slice_num={data.slice_num} run_uid={data.run_uid}"
+        )
+
+        result = self.dsperse_manager.prove_slice(
+            circuit_id=data.circuit,
+            slice_num=data.slice_num,
+            inputs=data.inputs,
+            outputs=data.outputs,
+        )
+
+        # Implementation for handling DSlice slice requests goes here
+        return JSONResponse(content=result, status_code=200)
+
     def queryZkProof(self, data: QueryZkProof) -> JSONResponse:
         """
         This function run proof generation of the model (with its output as well)

neurons/_validator/api/__init__.py

@@ -23,6 +23,7 @@
 import bittensor as bt
 from _validator.models.poc_rpc_request import ProofOfComputationRPCRequest
 from _validator.models.pow_rpc_request import ProofOfWeightsRPCRequest
+from _validator.models.base_rpc_request import QueuedRequestDataModel
 import hashlib
 from constants import (
     MAX_SIGNATURE_LIFESPAN,
@@ -70,9 +71,10 @@ def _should_rate_limit(ip: str):
 class ValidatorAPI:
     def __init__(self, config: ValidatorConfig):
         self.config = config
-        self.external_requests_queue: list[
-            ProofOfWeightsRPCRequest | ProofOfComputationRPCRequest
-        ] = []
+        # a Queue of requests to be sent to miners
+        # consists of "real world requests" ProofOfWeightsRPCRequest and ProofOfComputationRPCRequest
+        # and a Request with one slice of a DSperse model (DSlice)
+        self.stacked_requests_queue: list[QueuedRequestDataModel] = []
         self.ws_manager = WebSocketManager()
         self.recent_requests: dict[str, int] = {}
         self.validator_keys_cache = ValidatorKeysCache(config)
@@ -282,7 +284,7 @@ async def handle_proof_of_weights(
                 return InvalidParams(str(e))
 
             self.pending_requests[external_request.hash] = asyncio.Event()
-            self.external_requests_queue.insert(0, external_request)
+            self.stacked_requests_queue.insert(0, external_request)
             bt.logging.success(
                 f"External request with hash {external_request.hash} added to queue"
             )
@@ -341,7 +343,7 @@ async def handle_proof_of_computation(
                 return InvalidParams(str(e))
 
             self.pending_requests[external_request.hash] = asyncio.Event()
-            self.external_requests_queue.insert(0, external_request)
+            self.stacked_requests_queue.insert(0, external_request)
             bt.logging.success(
                 f"External request with hash {external_request.hash} added to queue"
             )

neurons/_validator/core/request_pipeline.py

@@ -21,7 +21,11 @@
 from deployment_layer.circuit_store import circuit_store
 from execution_layer.circuit import Circuit, CircuitType
 from execution_layer.generic_input import GenericInput
-from protocol import ProofOfWeightsDataModel, QueryZkProof
+from protocol import (
+    ProofOfWeightsDataModel,
+    QueryZkProof,
+    DSliceProofGenerationDataModel,
+)
 from utils.wandb_logger import safe_log
 
 
@@ -34,25 +38,6 @@ def __init__(
         self.api = api
         self.hash_guard = HashGuard()
 
-    def prepare_requests(self, filtered_uids) -> list[Request]:
-        """
-        Prepare requests for the current validation step.
-        This includes both regular benchmark requests and any external requests.
-
-        Args:
-            filtered_uids (list): List of UIDs to send requests to.
-
-        Returns:
-            list[Request]: List of prepared requests.
-        """
-        if len(filtered_uids) == 0:
-            bt.logging.error("No UIDs to query")
-            return []
-
-        if self.api.external_requests_queue:
-            return self._prepare_real_world_requests(filtered_uids)
-        return self._prepare_benchmark_requests(filtered_uids)
-
     def _check_and_create_request(
         self,
         uid: int,
@@ -64,7 +49,9 @@ def _check_and_create_request(
     ) -> Request | None:
         """Check hash and create request if valid."""
         try:
-            if isinstance(request_data, ProofOfWeightsDataModel):
+            if isinstance(request_data, ProofOfWeightsDataModel) or isinstance(
+                request_data, DSliceProofGenerationDataModel
+            ):
                 input_data = request_data.inputs
             else:
                 input_data = request_data.query_input
@@ -91,25 +78,28 @@ def _check_and_create_request(
             circuit=circuit,
             request_type=request_type,
             # 'inputs' are used for verification later on validator side:
+            #   I suppose `RWR` passed here to prevent new data generation
             inputs=GenericInput(RequestType.RWR, input_data),
             request_hash=request_hash,
             save=save,
         )
 
-    def _prepare_real_world_requests(self, filtered_uids: list[int]) -> list[Request]:
-        external_request = self.api.external_requests_queue.pop()
+    def _prepare_queued_requests(self, filtered_uids: list[int]) -> list[Request]:
+        external_request = self.api.stacked_requests_queue.pop()
         requests = []
 
         for uid in filtered_uids:
             try:
                 request_data, save = self.get_request_data(
-                    RequestType.RWR, external_request.circuit, external_request
+                    external_request.request_type,
+                    external_request.circuit,
+                    external_request,
                 )
                 request = self._check_and_create_request(
                     uid=uid,
                     request_data=request_data,
                     circuit=external_request.circuit,
-                    request_type=RequestType.RWR,
+                    request_type=external_request.request_type,
                     request_hash=external_request.hash,
                     save=save,
                 )
@@ -169,7 +159,7 @@ def get_request_data(
             circuit.input_handler(request_type)
             if request_type == RequestType.BENCHMARK
             else circuit.input_handler(
-                RequestType.RWR,
+                request_type,
                 copy.deepcopy(request.inputs),
             )
         )
@@ -207,6 +197,17 @@ def get_request_data(
                 QueryZkProof(query_input=inputs, model_id=circuit.id, query_output=""),
                 False,
             )
+        elif circuit.metadata.type == CircuitType.DSPERSE_PROOF_GENERATION:
+            return (
+                DSliceProofGenerationDataModel(
+                    circuit=circuit.id,
+                    inputs=request.inputs,
+                    outputs=request.outputs,
+                    slice_num=request.slice_num,
+                    run_uid=request.run_uid,
+                ),
+                False,
+            )
 
         return (
             ProofOfWeightsDataModel(
@@ -230,8 +231,8 @@ def prepare_single_request(self, uid: int) -> Request | None:
         Returns:
             Request | None: The prepared request, or None if preparation failed.
         """
-        if self.api.external_requests_queue:
-            requests = self._prepare_real_world_requests([uid])
+        if self.api.stacked_requests_queue:
+            requests = self._prepare_queued_requests([uid])
         else:
             requests = self._prepare_benchmark_requests([uid])
 

neurons/_validator/core/validator_loop.py

@@ -47,6 +47,7 @@
     ONE_HOUR,
     ONE_MINUTE,
 )
+from execution_layer.dsperse_manager import DSperseManager
 from utils import AutoUpdate, clean_temp_files, with_rate_limit
 from utils.gc_logging import gc_log_competition_metrics
 from utils.gc_logging import log_responses as gc_log_responses
@@ -120,6 +121,7 @@ def __init__(self, config: ValidatorConfig):
         self.request_pipeline = RequestPipeline(
             self.config, self.score_manager, self.api
         )
+        self.dsperse_manager = DSperseManager()
 
         self.request_queue = asyncio.Queue()
         self.active_tasks: dict[int, asyncio.Task] = {}
@@ -342,6 +344,13 @@ async def maintain_request_pool(self):
                 slots_available = self.current_concurrency - len(self.active_tasks)
 
                 if slots_available > 0:
+                    if not self.api.stacked_requests_queue:
+                        # Refill the stacked requests queue from DSperse manager if needed
+                        for (
+                            dslice_request
+                        ) in self.dsperse_manager.generate_dslice_requests():
+                            self.api.stacked_requests_queue.insert(0, dslice_request)
+
                     available_uids = [
                         uid
                         for uid in self.queryable_uids

neurons/_validator/models/base_rpc_request.py

@@ -1,11 +1,19 @@
 from pydantic import BaseModel
-from execution_layer.circuit import Circuit
+
+from _validator.models.request_type import RequestType
 from _validator.utils.api import hash_inputs
+from execution_layer.circuit import Circuit
+
 
+class QueuedRequestDataModel(BaseModel):
+    """
+    Base model for requests that are stacked in the validator's queue and waiting to be sent to miners.
+    At the moment, that's a Real World Request (RWR) or a Request with one slice of a DSperse model (DSlice).
+    """
 
-class RealWorldRequest(BaseModel):
     circuit: Circuit
     inputs: dict
+    request_type: RequestType = RequestType.RWR
 
     model_config = {"arbitrary_types_allowed": True}
 

neurons/_validator/models/dslice_request.py

@@ -0,0 +1,15 @@
+from pydantic import Field
+
+from _validator.models.base_rpc_request import QueuedRequestDataModel
+from _validator.models.request_type import RequestType
+
+
+class DSliceQueuedProofRequest(QueuedRequestDataModel):
+    """
+    Request for a DSperse slice.
+    """
+
+    request_type: RequestType = RequestType.DSLICE
+    slice_num: str = Field(..., description="Num of the DSperse slice")
+    run_uid: str = Field(..., description="UID of the DSperse run")
+    outputs: dict = Field(..., description="Outputs of the DSperse slice")

neurons/_validator/models/poc_rpc_request.py

@@ -1,10 +1,10 @@
-from _validator.models.base_rpc_request import RealWorldRequest
+from _validator.models.base_rpc_request import QueuedRequestDataModel
 from pydantic import Field
 from deployment_layer.circuit_store import circuit_store
 from execution_layer.circuit import CircuitType
 
 
-class ProofOfComputationRPCRequest(RealWorldRequest):
+class ProofOfComputationRPCRequest(QueuedRequestDataModel):
     """
     Request for the Proof of Computation RPC method.
     """

neurons/_validator/models/pow_rpc_request.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
-from _validator.models.base_rpc_request import RealWorldRequest
+from _validator.models.base_rpc_request import QueuedRequestDataModel
 from pydantic import Field
 from deployment_layer.circuit_store import circuit_store
 
 
-class ProofOfWeightsRPCRequest(RealWorldRequest):
+class ProofOfWeightsRPCRequest(QueuedRequestDataModel):
     """
     Request for the Proof of Weights RPC method.
     """

neurons/_validator/models/request_type.py

@@ -2,14 +2,31 @@
 
 
 class RequestType(Enum):
+    """
+    Enumeration of different types of requests that the validator can send to miner.
+     - BENCHMARK: Requests with generated input data for benchmarking purposes.
+       In case of empty RWR queue the validator generates some input data to keep miners busy.
+     - RWR: Real World Requests with actual input data for real-world inference.
+       Validator collects such requests from external users and stacks them in a queue to be sent to miners.
+     - DSLICE: That's a tricky one. We use DSperse app for slicing large models into smaller parts.
+       And some requests involve sliced model and each slice is sent as a separate request to the miner.
+       That each slice request is of type DSLICE.
+       At the moment we just stack DSlices to the same RWR queue, and behave as normal RWR requests.
+       That means DSLICE request is a part of RWR or BENCHMARK request.
+       XXX: Yeah, not very elegant, but we'll improve it later, I promise.
+    """
+
     BENCHMARK = "benchmark_request"
     RWR = "real_world_request"
+    DSLICE = "dslice_request"
 
     def __str__(self) -> str:
         if self == RequestType.BENCHMARK:
             return "Benchmark"
         elif self == RequestType.RWR:
             return "Real World Request"
+        elif self == RequestType.DSPERSE:
+            return "DSperse Request (one slice)"
         else:
             raise ValueError(f"Unknown request type: {self}")
 

neurons/cli_parser.py

@@ -96,6 +96,11 @@ def init_config(role: Optional[str] = None):
         default=None,
         help="Custom location for storing models data (optional)",
     )
+    parser.add_argument(
+        "--dsperse-run-dir",
+        default=None,
+        help="Custom location for storing dsperse run data (optional)",
+    )
     if role == Roles.VALIDATOR:
         # CLI arguments specific to the validator
         _validator_config()
@@ -141,6 +146,10 @@ def init_config(role: Optional[str] = None):
     else:
         config.full_path_models = os.path.join(config.full_path, "models")
 
+    if not config.dsperse_run_dir:
+        config.dsperse_run_dir = os.path.join(config.full_path, "dsperse_runs")
+    os.makedirs(config.dsperse_run_dir, exist_ok=True)
+
     if config.whitelisted_public_keys:
         config.whitelisted_public_keys = config.whitelisted_public_keys.split(",")