Fix: Refactor to unify getting execution price method making it compatible with multiple payment methods.

Author: nesitor

src/aleph/vm/orchestrator/payment.py

@@ -2,6 +2,7 @@
 import logging
 from collections.abc import Iterable
 from decimal import Decimal
+from typing import List
 
 import aiohttp
 from aleph_message.models import ItemHash, PaymentType
@@ -70,48 +71,9 @@ async def fetch_credit_balance_of_address(address: str) -> Decimal:
         return resp_data["credits"]
 
 
-async def fetch_execution_flow_price(item_hash: ItemHash) -> Decimal:
-    """Fetch the flow price of an execution from the reference API server."""
-    async with aiohttp.ClientSession() as session:
-        url = f"{settings.API_SERVER}/api/v0/price/{item_hash}"
-        resp = await session.get(url)
-        # Raise an error if the request failed
-        resp.raise_for_status()
-
-        resp_data = await resp.json()
-        required_flow: float = resp_data["required_tokens"]
-        payment_type: str | None = resp_data["payment_type"]
-
-        if payment_type is None:
-            msg = "Payment type must be specified in the message"
-            raise ValueError(msg)
-        elif payment_type != PaymentType.superfluid:
-            msg = f"Payment type {payment_type} is not supported"
-            raise ValueError(msg)
-
-        return Decimal(required_flow)
-
-
-async def fetch_execution_hold_price(item_hash: ItemHash) -> Decimal:
-    """Fetch the hold price of an execution from the reference API server."""
-    async with aiohttp.ClientSession() as session:
-        url = f"{settings.API_SERVER}/api/v0/price/{item_hash}"
-        resp = await session.get(url)
-        # Raise an error if the request failed
-        resp.raise_for_status()
-
-        resp_data = await resp.json()
-        required_hold: float = resp_data["required_tokens"]
-        payment_type: str | None = resp_data["payment_type"]
-
-        if payment_type not in (None, PaymentType.hold):
-            msg = f"Payment type {payment_type} is not supported"
-            raise ValueError(msg)
-
-        return Decimal(required_hold)
-
-
-async def fetch_execution_credit_price(item_hash: ItemHash) -> Decimal:
+async def fetch_execution_price(
+    item_hash: ItemHash, allowed_payments: List[PaymentType], payment_type_required: bool = True
+) -> Decimal:
     """Fetch the credit price of an execution from the reference API server."""
     async with aiohttp.ClientSession() as session:
         url = f"{settings.API_SERVER}/api/v0/price/{item_hash}"
@@ -123,10 +85,15 @@ async def fetch_execution_credit_price(item_hash: ItemHash) -> Decimal:
         required_credits: float = resp_data["required_credits"]  # Field not defined yet on API side.
         payment_type: str | None = resp_data["payment_type"]
 
-        if payment_type not in (None, PaymentType.credit):
-            msg = f"Payment type {payment_type} is not supported"
+        if payment_type_required and payment_type is None:
+            msg = "Payment type must be specified in the message"
             raise ValueError(msg)
 
+        if payment_type:
+            if payment_type not in allowed_payments:
+                msg = f"Payment type {payment_type} is not supported"
+                raise ValueError(msg)
+
         return Decimal(required_credits)
 
 
@@ -178,17 +145,26 @@ async def get_stream(sender: str, receiver: str, chain: str) -> Decimal:
 
 async def compute_required_balance(executions: Iterable[VmExecution]) -> Decimal:
     """Get the balance required for the resources of the user from the messages and the pricing aggregate."""
-    costs = await asyncio.gather(*(fetch_execution_hold_price(execution.vm_hash) for execution in executions))
+    costs = await asyncio.gather(
+        *(
+            fetch_execution_price(execution.vm_hash, [PaymentType.hold], payment_type_required=False)
+            for execution in executions
+        )
+    )
     return sum(costs, Decimal(0))
 
 
 async def compute_required_credit_balance(executions: Iterable[VmExecution]) -> Decimal:
     """Get the balance required for the resources of the user from the messages and the pricing aggregate."""
-    costs = await asyncio.gather(*(fetch_execution_credit_price(execution.vm_hash) for execution in executions))
+    costs = await asyncio.gather(
+        *(fetch_execution_price(execution.vm_hash, [PaymentType.credit]) for execution in executions)
+    )
     return sum(costs, Decimal(0))
 
 
 async def compute_required_flow(executions: Iterable[VmExecution]) -> Decimal:
     """Compute the flow required for a collection of executions, typically all executions from a specific address"""
-    flows = await asyncio.gather(*(fetch_execution_flow_price(execution.vm_hash) for execution in executions))
+    flows = await asyncio.gather(
+        *(fetch_execution_price(execution.vm_hash, [PaymentType.superfluid]) for execution in executions)
+    )
     return sum(flows, Decimal(0))

src/aleph/vm/orchestrator/views/__init__.py

@@ -35,7 +35,7 @@
 from aleph.vm.orchestrator.payment import (
     InvalidAddressError,
     InvalidChainError,
-    fetch_execution_flow_price,
+    fetch_execution_price,
     get_stream,
 )
 from aleph.vm.orchestrator.pubsub import PubSub
@@ -577,7 +577,7 @@ async def notify_allocation(request: web.Request):
         if have_gpu:
             logger.debug(f"GPU Instance {item_hash} not using PAYG")
         user_balance = await payment.fetch_balance_of_address(message.sender)
-        hold_price = await payment.fetch_execution_hold_price(item_hash)
+        hold_price = await payment.fetch_execution_price(item_hash, [PaymentType.hold], False)
         logger.debug(f"Address {message.sender} Balance: {user_balance}, Price: {hold_price}")
         if hold_price > user_balance:
             return web.HTTPPaymentRequired(
@@ -606,7 +606,7 @@ async def notify_allocation(request: web.Request):
         if not active_flow:
             raise web.HTTPPaymentRequired(reason="Empty payment stream for this instance")
 
-        required_flow: Decimal = await fetch_execution_flow_price(item_hash)
+        required_flow: Decimal = await fetch_execution_price(item_hash, [PaymentType.superfluid])
         community_wallet = await get_community_wallet_address()
         required_crn_stream: Decimal
         required_community_stream: Decimal