diff --git a/src/aleph/sdk/client/vm_client.py b/src/aleph/sdk/client/vm_client.py index 4a738cba..a692d0b0 100644 --- a/src/aleph/sdk/client/vm_client.py +++ b/src/aleph/sdk/client/vm_client.py @@ -16,6 +16,7 @@ from aleph.sdk.chains.solana import SOLAccount from aleph.sdk.types import Account from aleph.sdk.utils import ( + create_control_payload, create_vm_control_payload, sign_vm_control_payload, to_0x_hex, @@ -333,6 +334,40 @@ async def notify_allocation(self, vm_id: ItemHash) -> Tuple[int, str]: return session.status, form_response_text + async def reserve_resources( + self, instance_content: Dict[str, Any] + ) -> Tuple[Optional[int], str]: + """Pre-check CRN capacity for an instance before creating it. + + Sends the instance content to the CRN for admission control. + Returns 200 with {"status": "reserved", "expires": ...} if resources + are available, 503 if the CRN cannot fit the request. + """ + path = "/control/reserve_resources" + payload = create_control_payload( + path=path, domain=self.node_domain, method="POST" + ) + signed_operation = sign_vm_control_payload(payload, self.ephemeral_key) + + if not self.pubkey_signature_header: + self.pubkey_signature_header = ( + await self._generate_pubkey_signature_header() + ) + + headers = { + "X-SignedPubKey": self.pubkey_signature_header, + "X-SignedOperation": signed_operation, + } + + try: + async with self.session.post( + f"{self.node_url}{path}", headers=headers, json=instance_content + ) as resp: + return resp.status, await resp.text() + except aiohttp.ClientError as e: + logger.error("HTTP error during reserve_resources: %s", str(e)) + return None, str(e) + async def manage_instance( self, vm_id: ItemHash, operations: List[Union[VmOperation, str]] ) -> Tuple[int, str]: diff --git a/src/aleph/sdk/utils.py b/src/aleph/sdk/utils.py index 94bc3bb9..0e23ced2 100644 --- a/src/aleph/sdk/utils.py +++ b/src/aleph/sdk/utils.py @@ -241,13 +241,16 @@ def create_vm_control_payload( vm_id: ItemHash, operation: str, domain: str, method: str ) -> Dict[str, str]: path = f"/control/machine/{vm_id}/{operation}" - payload = { + return create_control_payload(path=path, domain=domain, method=method) + + +def create_control_payload(path: str, domain: str, method: str) -> Dict[str, str]: + return { "time": datetime.utcnow().isoformat() + "Z", "method": method.upper(), "path": path, "domain": domain, } - return payload def sign_vm_control_payload(payload: Dict[str, str], ephemeral_key) -> str: diff --git a/tests/unit/test_vm_client.py b/tests/unit/test_vm_client.py index 42d9e0fd..b9c5f826 100644 --- a/tests/unit/test_vm_client.py +++ b/tests/unit/test_vm_client.py @@ -222,6 +222,53 @@ async def test_exit_rescue(): await vm_client.session.close() +@pytest.mark.asyncio +async def test_reserve_resources_success(): + account = ETHAccount(private_key=b"0x" + b"1" * 30) + instance_content = {"address": "0xabc", "time": 1700000000.0} + + with aioresponses() as m: + vm_client = VmClient( + account=account, + node_url="http://localhost", + session=aiohttp.ClientSession(), + ) + m.post( + "http://localhost/control/reserve_resources", + status=200, + payload={"status": "reserved", "expires": "2026-01-01T00:00:00Z"}, + ) + + status, response_text = await vm_client.reserve_resources(instance_content) + assert status == 200 + assert "reserved" in response_text + assert ("POST", URL("http://localhost/control/reserve_resources")) in m.requests + await vm_client.session.close() + + +@pytest.mark.asyncio +async def test_reserve_resources_insufficient_capacity(): + account = ETHAccount(private_key=b"0x" + b"1" * 30) + instance_content = {"address": "0xabc", "time": 1700000000.0} + + with aioresponses() as m: + vm_client = VmClient( + account=account, + node_url="http://localhost", + session=aiohttp.ClientSession(), + ) + m.post( + "http://localhost/control/reserve_resources", + status=503, + body="This CRN cannot reserve the requested resources at this time.", + ) + + status, response_text = await vm_client.reserve_resources(instance_content) + assert status == 503 + assert "cannot reserve" in response_text + await vm_client.session.close() + + @pytest.mark.asyncio async def test_get_logs(aiohttp_client): account = ETHAccount(private_key=b"0x" + b"1" * 30)