Merge pull request #523 from backend-developers-ltd/COM-485

Streaming support

Author: piotr-figwer-reef

.github/workflows/integration_ci.yml

@@ -31,6 +31,7 @@ jobs:
         run: |
           docker pull alpine
           docker pull backenddevelopersltd/compute-horde-executor:v1-latest
+          docker pull backenddevelopersltd/compute-horde-streaming-job-test:v0-latest
       - name: Start all services
         run: local_stack/run_and_await_readiness.sh /tmp/integration_test_logs/
 

compute_horde/compute_horde/fv_protocol/facilitator_requests.py

@@ -6,6 +6,7 @@
 from compute_horde_core.executor_class import ExecutorClass
 from compute_horde_core.output_upload import MultiUpload, OutputUpload
 from compute_horde_core.signature import SignedFields, SignedRequest
+from compute_horde_core.streaming import StreamingDetails
 from compute_horde_core.volume import MultiVolume, Volume
 from pydantic import BaseModel, JsonValue
 
@@ -62,6 +63,8 @@ class V2JobRequest(SignedRequest, BaseModel, extra="forbid"):
     download_time_limit: int
     execution_time_limit: int
     upload_time_limit: int
+    streaming_start_time_limit: int
+    streaming_details: StreamingDetails | None = None
     # !!! all fields above are included in the signed json payload
 
     def get_args(self):
@@ -100,6 +103,8 @@ def get_signed_fields(self) -> SignedFields:
             download_time_limit=self.download_time_limit,
             execution_time_limit=self.execution_time_limit,
             upload_time_limit=self.upload_time_limit,
+            streaming_start_time_limit=self.streaming_start_time_limit,
+            streaming_details=self.streaming_details,
         )
         return signed_fields
 

compute_horde/compute_horde/fv_protocol/validator_requests.py

@@ -48,9 +48,16 @@ class MinerResponse(BaseModel, extra="allow"):
     upload_results: dict[str, str] | None = None
 
 
+class StreamingServerDetails(BaseModel, extra="forbid"):
+    streaming_server_cert: str | None = None
+    streaming_server_address: str | None = None
+    streaming_server_port: int | None = None
+
+
 class JobStatusMetadata(BaseModel, extra="allow"):
     comment: str
     miner_response: MinerResponse | None = None
+    streaming_details: StreamingServerDetails | None = None
 
 
 class JobStatusUpdate(BaseModel, extra="forbid"):
@@ -67,6 +74,7 @@ class Status(StrEnum):
         COMPLETED = "completed"
         REJECTED = "rejected"
         FAILED = "failed"
+        STREAMING_READY = "streaming_ready"
 
     message_type: Literal["V0JobStatusUpdate"] = "V0JobStatusUpdate"
     uuid: str

compute_horde/compute_horde/miner_client/organic.py

@@ -10,6 +10,7 @@
 import bittensor
 from compute_horde_core.executor_class import ExecutorClass
 from compute_horde_core.output_upload import OutputUpload
+from compute_horde_core.streaming import StreamingDetails
 from compute_horde_core.volume import Volume
 from pydantic import TypeAdapter
 
@@ -176,7 +177,10 @@ async def notify_volumes_ready(self, msg: V0VolumesReadyRequest) -> None:
         """This method is called when miner sends executor ready message"""
 
     async def notify_execution_done(self, msg: V0ExecutionDoneRequest) -> None:
-        """This method is called when miner sends executor ready message"""
+        """This method is called when miner sends execution done message"""
+
+    async def notify_streaming_readiness(self, msg: V0StreamingJobReadyRequest) -> None:
+        """This method is called when miner sends streaming ready message"""
 
     async def handle_manifest_request(self, msg: V0ExecutorManifestRequest) -> None:
         try:
@@ -396,9 +400,11 @@ async def send_initial_job_request(
                     download_time_limit=job_details.job_timing.download_time_limit,
                     execution_time_limit=job_details.job_timing.execution_time_limit,
                     upload_time_limit=job_details.job_timing.upload_time_limit,
+                    streaming_start_time_limit=job_details.job_timing.streaming_start_time_limit,
                 )
                 if job_details.job_timing is not None
                 else None,
+                streaming_details=job_details.streaming_details,
             ),
         )
 
@@ -428,6 +434,7 @@ class FailureReason(enum.Enum):
     EXECUTOR_FAILED = enum.auto()
     STREAMING_JOB_READY_TIMED_OUT = enum.auto()
     JOB_FAILED = enum.auto()
+    STREAMING_FAILED = enum.auto()
 
 
 class OrganicJobError(Exception):
@@ -458,6 +465,7 @@ class TimingDetails:
         download_time_limit: int
         execution_time_limit: int
         upload_time_limit: int
+        streaming_start_time_limit: int
 
         @property
         def total(self):
@@ -473,6 +481,7 @@ def total(self):
     volume: Volume | None = None
     output: OutputUpload | None = None
     artifacts_dir: str | None = None
+    streaming_details: StreamingDetails | None = None
 
 
 async def execute_organic_job_on_miner(
@@ -483,7 +492,6 @@ async def execute_organic_job_on_miner(
 ) -> tuple[str, str, dict[str, str], dict[str, str]]:  # stdout, stderr, artifacts, upload_results
     """
     Run an organic job. This is a simpler way to use OrganicMinerClient.
-
     :param client: the organic miner client
     :param job_details: details specific to the job that needs to be run
     :param reservation_time_limit: time for the miner to report reservation success (or decline the job)
@@ -514,7 +522,6 @@ async def execute_organic_job_on_miner(
             executor_class=job_details.executor_class,
             ttl=reservation_time_limit,
         )
-
         await client.send_initial_job_request(job_details, receipt_payload, receipt_signature)
         logger.debug("Sent initial job request")
         await JobStartedReceipt.from_payload(
@@ -583,7 +590,7 @@ async def execute_organic_job_on_miner(
                 )
             )
 
-            ## STAGE: volume download
+            ## STAGE: Volume download
             try:
                 if executor_timing:
                     logger.debug(
@@ -604,6 +611,26 @@ async def execute_organic_job_on_miner(
                 raise OrganicJobError(FailureReason.VOLUMES_TIMED_OUT) from exc
             await client.notify_volumes_ready(volumes_ready_response)
 
+            ## STAGE: Start streaming
+            if job_details.streaming_details:
+                try:
+                    if executor_timing:
+                        logger.debug(
+                            f"Extending deadline by streaming_start_time_limit: +{executor_timing.streaming_start_time_limit}s"
+                        )
+                        deadline.extend_timeout(executor_timing.streaming_start_time_limit)
+                    logger.debug(f"Waiting for streaming (time left: {deadline.time_left():.2f}s)")
+                    streaming_response = await asyncio.wait_for(
+                        client.streaming_job_ready_or_not_future,
+                        timeout=deadline.time_left(),
+                    )
+                except TimeoutError as exc:
+                    raise OrganicJobError(FailureReason.STREAMING_JOB_READY_TIMED_OUT) from exc
+                if isinstance(streaming_response, V0StreamingJobNotReadyRequest):
+                    raise OrganicJobError(FailureReason.STREAMING_FAILED, streaming_response)
+
+                await client.notify_streaming_readiness(streaming_response)
+
             ## STAGE: execution
             try:
                 if executor_timing:

compute_horde/compute_horde/protocol_messages.py

@@ -3,6 +3,7 @@
 
 from compute_horde_core.executor_class import ExecutorClass
 from compute_horde_core.output_upload import OutputUpload
+from compute_horde_core.streaming import StreamingDetails
 from compute_horde_core.volume import Volume
 from pydantic import BaseModel, Field
 
@@ -62,15 +63,12 @@ def total_count(self) -> int:
 
 # validator -> miner.vc -> miner.ec -> executor
 class V0InitialJobRequest(BaseModel):
-    class StreamingDetails(BaseModel):
-        public_key: str
-        executor_ip: str | None = None  # set by miner before sending to executor
-
     class ExecutorTimingDetails(BaseModel):
         allowed_leeway: int
         download_time_limit: int
         execution_time_limit: int
         upload_time_limit: int
+        streaming_start_time_limit: int
 
     message_type: Literal["V0InitialJobRequest"] = "V0InitialJobRequest"
     job_uuid: str

compute_horde/tests/test_certificate.py

@@ -6,17 +6,16 @@
 
 import pytest
 import requests
-from cryptography import x509
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
-
-from compute_horde.certificate import (
+from compute_horde_core.certificate import (
     generate_certificate,
     generate_certificate_at,
     save_public_key,
     start_nginx,
     write_certificate,
 )
+from cryptography import x509
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
 NGINX_PORT = 8443
 NGINX_URI = f"https://localhost:{NGINX_PORT}"

compute_horde/tests/test_job_request.py

@@ -26,6 +26,7 @@ def test_signed_job_roundtrip(signature_wallet):
         download_time_limit=1,
         execution_time_limit=1,
         upload_time_limit=1,
+        streaming_start_time_limit=1,
     )
 
     signer = BittensorWalletSigner(signature_wallet)

compute_horde/tests/test_run_organic_job.py

@@ -1,5 +1,6 @@
 import pytest
 from compute_horde_core.executor_class import ExecutorClass
+from compute_horde_core.streaming import StreamingDetails
 from pydantic import TypeAdapter
 
 from compute_horde.miner_client.organic import (
@@ -16,6 +17,7 @@
     V0JobFinishedReceiptRequest,
     V0JobFinishedRequest,
     V0JobRequest,
+    V0StreamingJobReadyRequest,
     V0VolumesReadyRequest,
     ValidatorAuthForMiner,
     ValidatorToMinerMessage,
@@ -83,6 +85,58 @@ async def test_run_organic_job__success(keypair):
     ]
 
 
+@pytest.mark.asyncio
+@pytest.mark.django_db(transaction=True)
+async def test_run_organic_job_streaming__success(keypair):
+    mock_transport = MinerStubTransport(
+        "mock",
+        [
+            V0AcceptJobRequest(job_uuid=JOB_UUID).model_dump_json(),
+            V0ExecutorReadyRequest(job_uuid=JOB_UUID).model_dump_json(),
+            V0VolumesReadyRequest(job_uuid=JOB_UUID).model_dump_json(),
+            V0StreamingJobReadyRequest(
+                job_uuid=JOB_UUID,
+                public_key="dummy-cert",
+                port=12345,
+            ).model_dump_json(),
+            V0ExecutionDoneRequest(job_uuid=JOB_UUID).model_dump_json(),
+            V0JobFinishedRequest(
+                job_uuid=JOB_UUID,
+                docker_process_stdout="streaming-stdout",
+                docker_process_stderr="streaming-stderr",
+                artifacts={},
+            ).model_dump_json(),
+        ],
+    )
+    client = OrganicMinerClient(
+        miner_hotkey="mock",
+        miner_address="0.0.0.0",
+        miner_port=1234,
+        job_uuid=JOB_UUID,
+        my_keypair=keypair,
+        transport=mock_transport,
+    )
+    job_details = OrganicJobDetails(
+        job_uuid=JOB_UUID,
+        executor_class=ExecutorClass.always_on__llm__a6000,
+        docker_image="mock",
+        streaming_details=StreamingDetails(public_key="dummy-cert"),
+    )
+    await execute_organic_job_on_miner(
+        client, job_details, reservation_time_limit=2, executor_startup_time_limit=2
+    )
+
+    # Find the initial job request sent by the client
+    initial_job_request = next(
+        (m for m in mock_transport.sent_models if isinstance(m, V0InitialJobRequest)), None
+    )
+    assert initial_job_request is not None, "Initial job request was not sent"
+    assert initial_job_request.streaming_details is not None, (
+        "Streaming details not sent in initial job request"
+    )
+    assert initial_job_request.streaming_details.public_key == "dummy-cert"
+
+
 # TODO:
 #   - unhappy path
 #       - connection error

compute_horde/tests/test_signature.py

@@ -128,6 +128,7 @@ def test_signed_fields__missing_fields():
         "download_time_limit": 1,
         "execution_time_limit": 1,
         "upload_time_limit": 1,
+        "streaming_start_time_limit": 1,
     }
     facilitator_signed_fields = SignedFields.from_facilitator_sdk_json(facilitator_request_json)
 
@@ -143,6 +144,7 @@ def test_signed_fields__missing_fields():
         download_time_limit=1,
         execution_time_limit=1,
         upload_time_limit=1,
+        streaming_start_time_limit=1,
     )
     assert v2_job_request.get_signed_fields() == facilitator_signed_fields
 
@@ -181,6 +183,7 @@ def test_signed_fields__volumes_uploads():
         "download_time_limit": 1,
         "execution_time_limit": 1,
         "upload_time_limit": 1,
+        "streaming_start_time_limit": 1,
     }
     facilitator_signed_fields = SignedFields.from_facilitator_sdk_json(facilitator_request_json)
 
@@ -198,6 +201,7 @@ def test_signed_fields__volumes_uploads():
         download_time_limit=1,
         execution_time_limit=1,
         upload_time_limit=1,
+        streaming_start_time_limit=1,
     )
     assert v2_job_request.get_signed_fields() == facilitator_signed_fields
 

…mpute_horde/compute_horde/certificate.py …dk/src/compute_horde_core/certificate.pycompute_horde/compute_horde/certificate.py renamed to compute_horde_sdk/src/compute_horde_core/certificate.py compute_horde/compute_horde/certificate.py renamed to compute_horde_sdk/src/compute_horde_core/certificate.py

@@ -39,7 +39,7 @@ async def get_docker_container_ip(container_name: str, bridge_network: bool = Fa
         raise Exception(f"Failed to get IP of {container_name}")
 
 
-async def check_endpoint(url, timeout) -> bool:
+async def check_endpoint(url: str, timeout: int) -> bool:
     """
     Pings endpoint every second until it returns 200 or timeout is reached.
     """
@@ -62,7 +62,7 @@ async def start_nginx(
     job_network: str,
     container_name: str = "job-nginx",
     timeout: int = 10,
-):
+) -> None:
     nginx_conf_file = dir_path / "nginx.conf"
     nginx_conf_file.write_text(nginx_conf)
 
@@ -89,9 +89,7 @@ async def start_nginx(
     ip = await get_docker_container_ip(container_name)
 
     # connect to internal network for job communication
-    process = await asyncio.create_subprocess_exec(
-        "docker", "network", "connect", job_network, container_name
-    )
+    process = await asyncio.create_subprocess_exec("docker", "network", "connect", job_network, container_name)
     await process.wait()
 
     # wait for nginx to start
@@ -108,17 +106,13 @@ def generate_certificate(alternative_name: str) -> tuple[Certificate, RSAPrivate
     private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
 
     builder = x509.CertificateBuilder()
-    builder = builder.subject_name(
-        x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")])
-    )
+    builder = builder.subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]))
     builder = builder.issuer_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]))
     builder = builder.not_valid_before(datetime.now(UTC))
     builder = builder.not_valid_after(datetime.now(UTC) + timedelta(days=365))
     builder = builder.serial_number(x509.random_serial_number())
     builder = builder.public_key(private_key.public_key())
-    builder = builder.add_extension(
-        x509.BasicConstraints(ca=False, path_length=None), critical=True
-    )
+    builder = builder.add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
 
     alt_name: x509.GeneralName
     try:
@@ -128,9 +122,7 @@ def generate_certificate(alternative_name: str) -> tuple[Certificate, RSAPrivate
     else:
         alt_name = x509.IPAddress(_ip)
 
-    builder = builder.add_extension(
-        x509.SubjectAlternativeName([x509.DNSName("localhost"), alt_name]), critical=False
-    )
+    builder = builder.add_extension(x509.SubjectAlternativeName([x509.DNSName("localhost"), alt_name]), critical=False)
 
     certificate = builder.sign(private_key=private_key, algorithm=hashes.SHA256())