From f0dcef727aba9f9adfa5fd15f5e20ecb1d8db26b Mon Sep 17 00:00:00 2001 From: Adal Chiriliuc Date: Fri, 25 Oct 2024 20:24:05 +0300 Subject: [PATCH 1/2] fix function name to signal it's public --- compute_horde/compute_horde/mv_protocol/validator_requests.py | 4 ++-- compute_horde/compute_horde/utils.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/compute_horde/compute_horde/mv_protocol/validator_requests.py b/compute_horde/compute_horde/mv_protocol/validator_requests.py index 129da6568..a3a76d0bb 100644 --- a/compute_horde/compute_horde/mv_protocol/validator_requests.py +++ b/compute_horde/compute_horde/mv_protocol/validator_requests.py @@ -12,7 +12,7 @@ from ..base.volume import Volume, VolumeType from ..base_requests import BaseRequest, JobMixin from ..executor_class import ExecutorClass -from ..utils import MachineSpecs, _json_dumps_default +from ..utils import MachineSpecs, json_dumps_default SAFE_DOMAIN_REGEX = re.compile(r".*") @@ -99,7 +99,7 @@ class ReceiptPayload(pydantic.BaseModel): def blob_for_signing(self): # pydantic v2 does not support sort_keys anymore. - return json.dumps(self.model_dump(), sort_keys=True, default=_json_dumps_default) + return json.dumps(self.model_dump(), sort_keys=True, default=json_dumps_default) class JobFinishedReceiptPayload(ReceiptPayload): diff --git a/compute_horde/compute_horde/utils.py b/compute_horde/compute_horde/utils.py index 1b95e343e..70dfc281b 100644 --- a/compute_horde/compute_horde/utils.py +++ b/compute_horde/compute_horde/utils.py @@ -52,7 +52,7 @@ def get_validators(netuid=12, network="finney", block: int | None = None) -> lis return neurons[:VALIDATORS_LIMIT] -def _json_dumps_default(obj): +def json_dumps_default(obj): if isinstance(obj, datetime.datetime): return obj.isoformat() From 49c6f18dcb61f566a477832ed415d8b25be3bc2b Mon Sep 17 00:00:00 2001 From: Adal Chiriliuc Date: Fri, 25 Oct 2024 20:33:46 +0300 Subject: [PATCH 2/2] removed separate payload from signed request --- .../fv_protocol/facilitator_requests.py | 25 ++++++--- compute_horde/tests/test_job_request.py | 53 +++++++++++++++++++ .../organic_jobs/facilitator_client.py | 4 +- 3 files changed, 72 insertions(+), 10 deletions(-) create mode 100644 compute_horde/tests/test_job_request.py diff --git a/compute_horde/compute_horde/fv_protocol/facilitator_requests.py b/compute_horde/compute_horde/fv_protocol/facilitator_requests.py index c5e2374c0..3c8a0a81b 100644 --- a/compute_horde/compute_horde/fv_protocol/facilitator_requests.py +++ b/compute_horde/compute_horde/fv_protocol/facilitator_requests.py @@ -21,12 +21,12 @@ class Response(BaseModel, extra="forbid"): errors: list[Error] = [] -class SignedRequest(BaseModel, extra="forbid"): - signature_type: str - signatory: str - timestamp_ns: int - signature: str - signed_payload: JsonValue +class Signature(BaseModel, extra="forbid"): + # has defaults to allow easy instantiation + signature_type: str = "" + signatory: str = "" + timestamp_ns: int = 0 + signature: str = "" class V0JobRequest(BaseModel, extra="forbid"): @@ -102,8 +102,10 @@ class V2JobRequest(BaseModel, extra="forbid"): # this points to a `ValidatorConsumer.job_new` handler (fuck you django-channels!) type: Literal["job.new"] = "job.new" message_type: Literal["V2JobRequest"] = "V2JobRequest" + signature: Signature | None = None + + # !!! all fields below are included in the signed json payload uuid: str - miner_hotkey: str | None executor_class: ExecutorClass docker_image: str raw_script: str @@ -112,11 +114,18 @@ class V2JobRequest(BaseModel, extra="forbid"): use_gpu: bool volume: Volume | None = None output_upload: OutputUpload | None = None - signed_request: SignedRequest + # !!! all fields above are included in the signed json payload def get_args(self): return self.args + def json_for_signing(self) -> JsonValue: + payload = self.model_dump(mode="json") + del payload["type"] + del payload["message_type"] + del payload["signature"] + return payload + @model_validator(mode="after") def validate_at_least_docker_image_or_raw_script(self) -> Self: if not (bool(self.docker_image) or bool(self.raw_script)): diff --git a/compute_horde/tests/test_job_request.py b/compute_horde/tests/test_job_request.py new file mode 100644 index 000000000..190084191 --- /dev/null +++ b/compute_horde/tests/test_job_request.py @@ -0,0 +1,53 @@ +import base64 +import uuid + +from compute_horde.base.volume import VolumeType, ZipUrlVolume +from compute_horde.executor_class import DEFAULT_EXECUTOR_CLASS +from compute_horde.fv_protocol.facilitator_requests import Signature, V2JobRequest +from compute_horde.signature import BittensorWalletSigner, BittensorWalletVerifier +from compute_horde.signature import Signature as RawSignature + + +def test_signed_job_roundtrip(signature_wallet): + volume = ZipUrlVolume( + volume_type=VolumeType.zip_url, + contents="https://example.com/input.zip", + relative_path="input", + ) + job = V2JobRequest( + uuid=str(uuid.uuid4()), + executor_class=DEFAULT_EXECUTOR_CLASS, + docker_image="hello-world", + raw_script="bash", + args=["--verbose", "--dry-run"], + env={"CUDA": "1"}, + use_gpu=False, + volume=volume, + output_upload=None, + ) + + signer = BittensorWalletSigner(signature_wallet) + payload = job.json_for_signing() + raw_signature = signer.sign(payload) + + job.signature = Signature( + signature_type=raw_signature.signature_type, + signatory=raw_signature.signatory, + timestamp_ns=raw_signature.timestamp_ns, + signature=base64.b64encode(raw_signature.signature).decode("utf8"), + ) + + job_json = job.model_dump_json() + deserialized_job = V2JobRequest.model_validate_json(job_json) + + assert deserialized_job.signature is not None + deserialized_raw_signature = RawSignature( + signature_type=deserialized_job.signature.signature_type, + signatory=deserialized_job.signature.signatory, + timestamp_ns=deserialized_job.signature.timestamp_ns, + signature=base64.b64decode(deserialized_job.signature.signature), + ) + + deserialized_payload = deserialized_job.json_for_signing() + verifier = BittensorWalletVerifier() + verifier.verify(deserialized_payload, deserialized_raw_signature) 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 index 605926942..5868beda4 100644 --- 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 @@ -8,7 +8,7 @@ import tenacity import websockets from channels.layers import get_channel_layer -from compute_horde.fv_protocol.facilitator_requests import Error, JobRequest, Response +from compute_horde.fv_protocol.facilitator_requests import Error, JobRequest, Response, V2JobRequest from compute_horde.fv_protocol.validator_requests import ( V0AuthenticationRequest, V0Heartbeat, @@ -243,7 +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 + assert not isinstance(job_request, V2JobRequest) 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(