Fix: Miss calculations / add Persistant vm

deployment/migrations/versions/0018_7bcb8e5fe186_fix_vm_cost_view.py

@@ -124,7 +124,7 @@ def upgrade() -> None:
                                 LEFT JOIN files ON volume.volume_to_use::text = files.hash::text
                        GROUP BY volume.vm_hash) file_volumes_size
                       ON vm_versions.current_version::text = file_volumes_size.vm_hash::text
-                 JOIN (SELECT instance_rootfs.instance_hash,
+                 LEFT JOIN (SELECT instance_rootfs.instance_hash,
                               instance_rootfs.size_mib::bigint * 1024 * 1024 AS rootfs_size
                        FROM instance_rootfs) rootfs_size ON vm_versions.vm_hash::text = rootfs_size.instance_hash::text
                  JOIN (SELECT vm_machine_volumes.vm_hash,
@@ -151,8 +151,7 @@ def upgrade() -> None:
                                   END AS base_compute_unit_price) bcp,
              LATERAL ( SELECT 1 + vms.environment_internet::integer AS compute_unit_price_multiplier) m,
              LATERAL ( SELECT cu.compute_units_required * m.compute_unit_price_multiplier::double precision *
-                          bcp.base_compute_unit_price::double precision *
-                          m.compute_unit_price_multiplier::double precision AS compute_unit_price) cpm,
+                          bcp.base_compute_unit_price::double precision AS compute_unit_price) cpm,
          LATERAL ( SELECT additional_disk.additional_disk_space * 20::double precision /
                           (1024 * 1024)::double precision AS disk_price) adp,
          LATERAL ( SELECT cpm.compute_unit_price + adp.disk_price AS total_price) tp

src/aleph/handlers/content/vm.py

@@ -320,15 +320,19 @@ class VmMessageHandler(ContentHandler):
     """
 
     async def check_balance(self, session: DbSession, message: MessageDb) -> None:
-        if message.type != MessageType.instance:
+        if not (message.type == MessageType.instance) and not (message.type == MessageType.program):
             return
+
         content = _get_vm_content(message)
-        if isinstance(content, ProgramContent):
-            return
+
+        if message.type == MessageType.program:
+            if not content.on.persistent:
+                return
 
         required_tokens = compute_cost(session=session, content=content)
+
         current_balance = (
-            get_total_balance(address=content.address, session=session) or 0
+                get_total_balance(address=content.address, session=session) or 0
         )
         current_instance_costs = get_total_cost_for_address(
             session=session, address=content.address

src/aleph/services/cost.py

@@ -1,3 +1,4 @@
+import math
 from decimal import Decimal
 from typing import Optional, Union
 
@@ -12,9 +13,8 @@
 
 
 def _get_file_from_ref(
-    session: DbSession, ref: str, use_latest: bool
+        session: DbSession, ref: str, use_latest: bool
 ) -> Optional[StoredFileDb]:
-
     tag_or_pin: Optional[Union[MessageFilePinDb, FileTagDb]]
 
     if use_latest:
@@ -46,12 +46,13 @@ def get_volume_size(session: DbSession, content: ExecutableContent) -> Decimal:
     total_volume_size: Decimal = Decimal(0)
 
     for volume in ref_volumes:
-        file = _get_file_from_ref(
-            session=session, ref=volume.ref, use_latest=volume.use_latest
-        )
-        if file is None:
-            raise RuntimeError(f"Could not find entry in file tags for {volume.ref}.")
-        total_volume_size += Decimal(file.size)
+        if hasattr(volume, "ref"):
+            file = _get_file_from_ref(
+                session=session, ref=volume.ref, use_latest=volume.use_latest
+            )
+            if file is None:
+                raise RuntimeError(f"Could not find entry in file tags for {volume.ref}.")
+            total_volume_size += Decimal(file.size)
 
     for volume in sized_volumes:
         total_volume_size += Decimal(volume.size_mib * MiB)
@@ -60,7 +61,7 @@ def get_volume_size(session: DbSession, content: ExecutableContent) -> Decimal:
 
 
 def get_additional_storage_price(
-    content: ExecutableContent, session: DbSession
+        content: ExecutableContent, session: DbSession
 ) -> Decimal:
     is_microvm = isinstance(content, ProgramContent) and not content.on.persistent
     nb_compute_units = content.resources.vcpus
@@ -77,7 +78,18 @@ def get_additional_storage_price(
 def compute_cost(session: DbSession, content: ExecutableContent) -> Decimal:
     is_microvm = isinstance(content, ProgramContent) and not content.on.persistent
     compute_unit_cost: Decimal = Decimal("200.0") if is_microvm else Decimal("2000.0")
+    cpu: Decimal = Decimal(content.resources.vcpus)
+    memory: Decimal = Decimal(math.ceil(content.resources.memory / 2000))
+    compute_unit_multiplier = cpu if cpu >= memory else memory
+    environment_internet: Decimal = Decimal(0)
+
+    if isinstance(content, ProgramContent) and content.environment.internet:
+        environment_internet = Decimal(2)
 
-    return (compute_unit_cost * content.resources.vcpus) + get_additional_storage_price(
+    compute_units_required = compute_unit_multiplier
+    compute_unit_multiplier = environment_internet or 1
+    compute_unit_price = compute_units_required * compute_unit_multiplier * compute_unit_cost
+
+    return compute_unit_price + get_additional_storage_price(
         content, session
-    )
+    )

tests/message_processing/test_process_programs.py

@@ -27,14 +27,14 @@
     ImmutableVolumeDb,
     EphemeralVolumeDb,
     PersistentVolumeDb,
-    StoredFileDb,
+    StoredFileDb, AlephBalanceDb,
 )
 from aleph.jobs.process_pending_messages import PendingMessageProcessor
 from aleph.toolkit.timestamp import timestamp_to_datetime
 from aleph.types.db_session import DbSessionFactory, DbSession
 from aleph.types.files import FileTag, FileType
 from aleph.types.message_status import MessageStatus, ErrorCode
-
+from decimal import Decimal
 
 @pytest.fixture
 def fixture_program_message(session_factory: DbSessionFactory) -> PendingMessageDb:
@@ -45,7 +45,7 @@ def fixture_program_message(session_factory: DbSessionFactory) -> PendingMessage
         sender="0x7083b90eBA420832A03C6ac7e6328d37c72e0260",
         signature="0x5c5c757f35403e9b6d6b1c5dc0be349284d76ded4cfa9edc6ad8522212f4235448bc80432ea7f8c5f4da315d6ce6902072d4eb7bbbebd969841f08a3df61c2971c",
         item_type=ItemType.inline,
-        item_content='{"address":"0x7083b90eBA420832A03C6ac7e6328d37c72e0260","time":1655123939.12433,"type":"vm-function","allow_amend":false,"code":{"encoding":"squashfs","entrypoint":"python run.py","ref":"53ee77caeb7d6e0e982abf010b3d6ea2dbc1225e157e09283e3a9d7da757e193","use_latest":true},"variables":{"LD_LIBRARY_PATH":"/opt/extra_lib","DB_FOLDER":"/data","RPC_ENDPOINT":"https://rpc.tzkt.io/ithacanet","TRUSTED_RPC_ENDPOINT":"https://rpc.tzkt.io/ithacanet","WELL_CONTRACT":"KT1ReVgfaUqHzWWiNRfPXQxf7TaBLVbxrztw","PORT":"8080","CONCURRENT_JOB":"5","BATCH_SIZE":"10","UNTIL_BLOCK":"201396","PUBSUB":"{\\"namespace\\": \\"tznms\\",\\"uuid\\": \\"tz_uid_1\\",\\"hook_url\\": \\"_domain_or_ip_addess\\",\\"pubsub_server\\": \\"domain_or_ip_address\\",\\"secret_shared_key\\": \\"112secret_key\\",\\"channel\\": \\"storage\\",\\"running_mode\\": \\"readonly\\"}"},"on":{"http":true},"environment":{"reproducible":false,"internet":true,"aleph_api":true,"shared_cache":false},"resources":{"vcpus":4,"memory":4000,"seconds":300},"runtime":{"ref":"bd79839bf96e595a06da5ac0b6ba51dea6f7e2591bb913deccded04d831d29f4","use_latest":true,"comment":"Aleph Alpine Linux with Python 3.8"},"volumes":[{"comment":"Extra lib","mount":"/opt/extra_lib","ref":"d7cecdeccc916280f8bcbf0c0e82c3638332da69ece2cbc806f9103a0f8befea","use_latest":true},{"comment":"Python Virtual Environment","mount":"/opt/packages","ref":"1000ebe0b61e41d5e23c10f6eb140e837188158598049829f2820f830139fc7d","use_latest":true},{"comment":"Data storage","mount":"/data","persistence":"host","name":"data","size_mib":128}]}',
+        item_content='{"address":"0x7083b90eBA420832A03C6ac7e6328d37c72e0260","time":1655123939.12433,"type":"vm-function","allow_amend":false,"code":{"encoding":"squashfs","entrypoint":"python run.py","ref":"53ee77caeb7d6e0e982abf010b3d6ea2dbc1225e157e09283e3a9d7da757e193","use_latest":true},"variables":{"LD_LIBRARY_PATH":"/opt/extra_lib","DB_FOLDER":"/data","RPC_ENDPOINT":"https://rpc.tzkt.io/ithacanet","TRUSTED_RPC_ENDPOINT":"https://rpc.tzkt.io/ithacanet","WELL_CONTRACT":"KT1ReVgfaUqHzWWiNRfPXQxf7TaBLVbxrztw","PORT":"8080","CONCURRENT_JOB":"5","BATCH_SIZE":"10","UNTIL_BLOCK":"201396","PUBSUB":"{\\"namespace\\": \\"tznms\\",\\"uuid\\": \\"tz_uid_1\\",\\"hook_url\\": \\"_domain_or_ip_addess\\",\\"pubsub_server\\": \\"domain_or_ip_address\\",\\"secret_shared_key\\": \\"112secret_key\\",\\"channel\\": \\"storage\\",\\"running_mode\\": \\"readonly\\"}"},"on":{"http":false, "persistent": true},"environment":{"reproducible":false,"internet":true,"aleph_api":true,"shared_cache":false},"resources":{"vcpus":4,"memory":4095,"seconds":300},"runtime":{"ref":"bd79839bf96e595a06da5ac0b6ba51dea6f7e2591bb913deccded04d831d29f4","use_latest":true,"comment":"Aleph Alpine Linux with Python 3.8"},"volumes":[{"comment":"Extra lib","mount":"/opt/extra_lib","ref":"5f31b0706f59404fad3d0bff97ef89ddf24da4761608ea0646329362c662ba51","use_latest":true},{"comment":"Python Virtual Environment","mount":"/opt/packages","ref":"1000ebe0b61e41d5e23c10f6eb140e837188158598049829f2820f830139fc7d","use_latest":true},{"comment":"Data storage","mount":"/data","persistence":"host","name":"data","size_mib":128}]}',
         time=timestamp_to_datetime(1671637391),
         channel=None,
         reception_time=timestamp_to_datetime(1671637391),
@@ -148,9 +148,23 @@ def insert_volume_refs(session: DbSession, message: PendingMessageDb):
                 last_updated=created,
             )
 
+@pytest.fixture
+def user_balance(session_factory: DbSessionFactory) -> AlephBalanceDb:
+    balance = AlephBalanceDb(
+        address="0x7083b90eBA420832A03C6ac7e6328d37c72e0260",
+        chain=Chain.ETH,
+        balance=Decimal(22_192),
+        eth_height=0,
+    )
+
+    with session_factory() as session:
+        session.add(balance)
+        session.commit()
+    return balance
 
 @pytest.mark.asyncio
 async def test_process_program(
+    user_balance: AlephBalanceDb,
     session_factory: DbSessionFactory,
     message_processor: PendingMessageProcessor,
     fixture_program_message: PendingMessageDb,
@@ -176,7 +190,7 @@ async def test_process_program(
         assert program.program_type == MachineType.vm_function
         assert not program.allow_amend
         assert program.replaces is None
-        assert program.http_trigger
+        assert program.persistent
 
         assert program.resources_vcpus == content_dict["resources"]["vcpus"]
         assert program.resources_memory == content_dict["resources"]["memory"]