Merge pull request #288 from backend-developers-ltd/fix-types

Fix compute_horde type declarations
backend-developers-ltd · Oct 24, 2024 · 67a0980 · 67a0980
2 parents 117fc92 + 533c95f
commit 67a0980
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 10 deletions.
diff --git a/compute_horde/compute_horde/fv_protocol/facilitator_requests.py b/compute_horde/compute_horde/fv_protocol/facilitator_requests.py
@@ -125,6 +125,6 @@ def validate_at_least_docker_image_or_raw_script(self) -> Self:
 
 
 JobRequest = Annotated[
-    V0JobRequest | V1JobRequest,
+    V0JobRequest | V1JobRequest | V2JobRequest,
     pydantic.Field(discriminator="message_type"),
 ]
diff --git a/compute_horde/compute_horde/fv_protocol/validator_requests.py b/compute_horde/compute_horde/fv_protocol/validator_requests.py
@@ -1,4 +1,4 @@
-from typing import Any, Self
+from typing import Any, Literal, Self
 
 import bittensor
 from pydantic import BaseModel
@@ -7,13 +7,13 @@
 class V0Heartbeat(BaseModel, extra="forbid"):
     """Message sent from validator to facilitator to keep connection alive"""
 
-    message_type: str = "V0Heartbeat"
+    message_type: Literal["V0Heartbeat"] = "V0Heartbeat"
 
 
 class V0AuthenticationRequest(BaseModel, extra="forbid"):
     """Message sent from validator to facilitator to authenticate itself"""
 
-    message_type: str = "V0AuthenticationRequest"
+    message_type: Literal["V0AuthenticationRequest"] = "V0AuthenticationRequest"
     public_key: str
     signature: str
 
@@ -43,7 +43,7 @@ def ss58_address(self) -> str:
 class V0MachineSpecsUpdate(BaseModel, extra="forbid"):
     """Message sent from validator to facilitator to update miner specs"""
 
-    message_type: str = "V0MachineSpecsUpdate"
+    message_type: Literal["V0MachineSpecsUpdate"] = "V0MachineSpecsUpdate"
     miner_hotkey: str
     validator_hotkey: str
     specs: dict[str, Any]

diff --git a/validator/app/src/compute_horde_validator/validator/organic_jobs/facilitator_client.py b/validator/app/src/compute_horde_validator/validator/organic_jobs/facilitator_client.py
@@ -243,6 +243,7 @@ async def get_miner_axon_info(self, hotkey: str) -> bittensor.AxonInfo:
 
     async def miner_driver(self, job_request: JobRequest):
         """drive a miner client from job start to completion, then close miner connection"""
+        assert job_request.miner_hotkey is not None
         miner, _ = await Miner.objects.aget_or_create(hotkey=job_request.miner_hotkey)
         miner_axon_info = await self.get_miner_axon_info(job_request.miner_hotkey)
         job = await OrganicJob.objects.acreate(

diff --git a/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py b/validator/app/src/compute_horde_validator/validator/organic_jobs/miner_driver.py
@@ -4,10 +4,7 @@
 from typing import Literal
 
 from compute_horde.executor_class import ExecutorClass
-from compute_horde.fv_protocol.facilitator_requests import (
-    V0JobRequest,
-    V1JobRequest,
-)
+from compute_horde.fv_protocol.facilitator_requests import JobRequest
 from compute_horde.miner_client.organic import (
     FailureReason,
     OrganicJobDetails,
@@ -93,7 +90,7 @@ async def _dummy_notify_callback(_: JobStatusUpdate) -> None:
 async def execute_organic_job(
     miner_client: MinerClient,
     job: OrganicJob,
-    job_request: V0JobRequest | V1JobRequest | AdminJobRequest,
+    job_request: JobRequest | AdminJobRequest,
     total_job_timeout: int = 300,
     wait_timeout: int = 300,
     notify_callback: Callable[[JobStatusUpdate], Awaitable[None]] = _dummy_notify_callback,