Feat: Support estimated_size_mib in program content schema for cost estimation endpoint (#715)

* Feat: Support estimated_size_mib in program content schema for cost estimation endpoint

* fix: lint fixes

* fix: lint errors

* pyproject version fix

* fix: lint fixes

* fix: cast estimated_size_mib to Decimal

Author: amalcaraz

pyproject.toml

@@ -45,7 +45,7 @@ dependencies = [
   "msgpack==1.0.8",                                                                                                        # required by aiocache
   "multiaddr==0.0.9",                                                                                                      # for libp2p-stubs
   "orjson>=3.7.7",                                                                                                         # Minimum version for Python 3.11
-  "psycopg2-binary==2.9.10",                                                                                                # Note: psycopg3 is not yet supported by SQLAlchemy
+  "psycopg2-binary==2.9.10",                                                                                               # Note: psycopg3 is not yet supported by SQLAlchemy
   "pycryptodome==3.22.0",                                                                                                  # for libp2p-stubs
   "pymultihash==0.8.2",                                                                                                    # for libp2p-stubs
   "pynacl==1.5",

src/aleph/schemas/cost_estimation_messages.py

@@ -7,6 +7,11 @@
     ProgramContent,
     StoreContent,
 )
+from aleph_message.models.execution.program import (
+    CodeContent,
+    DataContent,
+    FunctionRuntime,
+)
 from aleph_message.models.execution.volume import (
     EphemeralVolume,
     ImmutableVolume,
@@ -37,7 +42,26 @@ class CostEstimationInstanceContent(InstanceContent):
     )
 
 
+class CostEstimationCodeContent(CodeContent):
+    estimated_size_mib: Optional[int] = None
+
+
+class CostEstimationFunctionRuntime(FunctionRuntime):
+    estimated_size_mib: Optional[int] = None
+
+
+class CostEstimationDataContent(DataContent):
+    estimated_size_mib: Optional[int] = None
+
+
 class CostEstimationProgramContent(ProgramContent):
+    code: CostEstimationCodeContent = Field(description="Code to execute")
+    runtime: CostEstimationFunctionRuntime = Field(
+        description="Execution runtime (rootfs with Python interpreter)"
+    )
+    data: Optional[CostEstimationDataContent] = Field(
+        default=None, description="Data to use during computation"
+    )
     volumes: List[CostEstimationMachineVolume] = Field(
         default=[], description="Volumes to mount on the filesystem"
     )

src/aleph/services/cost.py

@@ -299,28 +299,67 @@ def _get_execution_volumes_costs(
         )
 
     elif isinstance(content, ProgramContent):
-        volumes += [
-            RefVolume(
-                CostType.EXECUTION_PROGRAM_VOLUME_CODE,
-                content.code.ref,
-                content.code.use_latest,
-            ),
-            RefVolume(
-                CostType.EXECUTION_PROGRAM_VOLUME_RUNTIME,
-                content.runtime.ref,
-                content.runtime.use_latest,
-            ),
-        ]
+        if (
+            isinstance(content, CostEstimationProgramContent)
+            and content.code.estimated_size_mib
+        ):
+            volumes.append(
+                SizedVolume(
+                    CostType.EXECUTION_PROGRAM_VOLUME_CODE,
+                    Decimal(content.code.estimated_size_mib),
+                    content.code.ref,
+                )
+            )
+        else:
+            volumes.append(
+                RefVolume(
+                    CostType.EXECUTION_PROGRAM_VOLUME_CODE,
+                    content.code.ref,
+                    content.code.use_latest,
+                )
+            )
 
-        if content.data:
+        if (
+            isinstance(content, CostEstimationProgramContent)
+            and content.runtime.estimated_size_mib
+        ):
+            volumes.append(
+                SizedVolume(
+                    CostType.EXECUTION_PROGRAM_VOLUME_RUNTIME,
+                    Decimal(content.runtime.estimated_size_mib),
+                    content.runtime.ref,
+                )
+            )
+        else:
             volumes.append(
                 RefVolume(
-                    CostType.EXECUTION_PROGRAM_VOLUME_DATA,
-                    content.data.ref,
-                    content.data.use_latest,
+                    CostType.EXECUTION_PROGRAM_VOLUME_RUNTIME,
+                    content.runtime.ref,
+                    content.runtime.use_latest,
                 ),
             )
 
+        if content.data:
+            if (
+                isinstance(content, CostEstimationProgramContent)
+                and content.data.estimated_size_mib
+            ):
+                volumes.append(
+                    SizedVolume(
+                        CostType.EXECUTION_PROGRAM_VOLUME_DATA,
+                        Decimal(content.data.estimated_size_mib),
+                        content.data.ref,
+                    )
+                )
+            else:
+                volumes.append(
+                    RefVolume(
+                        CostType.EXECUTION_PROGRAM_VOLUME_DATA,
+                        content.data.ref,
+                        content.data.use_latest,
+                    ),
+                )
+
     for i, volume in enumerate(content.volumes):
         # NOTE: There are legacy volumes with no "mount" property set
         # or with same values for different volumes causing unique key constraint errors

tests/services/test_cost_service.py

@@ -5,6 +5,7 @@
 from aleph_message.models import ExecutableContent, InstanceContent, PaymentType
 
 from aleph.db.models import AggregateDb
+from aleph.schemas.cost_estimation_messages import CostEstimationProgramContent
 from aleph.services.cost import (
     _get_additional_storage_price,
     _get_price_aggregate,
@@ -212,6 +213,63 @@ def fixture_flow_instance_message_complete() -> ExecutableContent:
     return InstanceContent.parse_obj(content)
 
 
+@pytest.fixture
+def fixture_hold_program_message_complete() -> ExecutableContent:
+    content = {
+        "on": {"http": True, "persistent": False},
+        "code": {
+            "ref": "cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe",
+            "encoding": "zip",
+            "entrypoint": "main:app",
+            "use_latest": True,
+            "estimated_size_mib": 2048,
+        },
+        "time": 1740986893.735,
+        "type": "vm-function",
+        "address": "0xAD8ac12Ae5bC9f6D902cBDd2f0Dd70F43e522BC2",
+        "payment": {"type": "hold", "chain": "ETH"},
+        "runtime": {
+            "ref": "cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe",
+            "comment": "Aleph Alpine Linux with Python 3.8",
+            "use_latest": True,
+            "estimated_size_mib": 1024,
+        },
+        "data": {
+            "ref": "cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe",
+            "mount": "/data",
+            "encoding": "zip",
+            "use_latest": True,
+            "estimated_size_mib": 2048,
+        },
+        "volumes": [
+            {
+                "mount": "/mount1",
+                "ref": "cafecafecafecafecafecafecafecafecafecafecafecafecafecafecafecafe",
+                "use_latest": True,
+                "estimated_size_mib": 512,
+            },
+            {
+                "mount": "/mount2",
+                "persistence": "host",
+                "name": "pers1",
+                "size_mib": 1024,
+            },
+        ],
+        "metadata": {"name": "My program", "description": "My program description"},
+        "resources": {"vcpus": 1, "memory": 128, "seconds": 30},
+        "variables": {},
+        "allow_amend": False,
+        "environment": {
+            "internet": True,
+            "aleph_api": True,
+            "reproducible": False,
+            "shared_cache": False,
+        },
+    }
+
+    return CostEstimationProgramContent.parse_obj(content)
+
+
 def test_compute_cost(
     session_factory: DbSessionFactory,
     fixture_product_prices_aggregate_in_db,
@@ -290,7 +348,7 @@ def test_get_additional_storage_price(
         assert additional_cost == 0
 
 
-def test_compute_cost_complete(
+def test_compute_cost_instance_complete(
     session_factory: DbSessionFactory,
     fixture_product_prices_aggregate_in_db,
     fixture_settings_aggregate_in_db,
@@ -309,6 +367,25 @@ def test_compute_cost_complete(
         assert cost == 1017.50
 
 
+def test_compute_cost_program_complete(
+    session_factory: DbSessionFactory,
+    fixture_product_prices_aggregate_in_db,
+    fixture_settings_aggregate_in_db,
+    fixture_hold_program_message_complete,
+):
+    file_db = StoredFileDb()
+    mock = Mock()
+    mock.patch("_get_file_from_ref", return_value=file_db)
+
+    with session_factory() as session:
+        cost, _ = get_total_and_detailed_costs(
+            session=session,
+            content=fixture_hold_program_message_complete,
+            item_hash="asdf",
+        )
+        assert cost == Decimal("630.400000000000000000")
+
+
 def test_compute_flow_cost(
     session_factory: DbSessionFactory,
     fixture_product_prices_aggregate_in_db,