temporal_v2: add memory cost reporting (#4257)

Signed-off-by: Benji Visser <benji@093b.org>
Co-authored-by: Benji Visser <benji@093b.org>

alembic/versions/2025_12_10_1915-174dcd456325_add_memory_mb_to_task_runs.py

@@ -0,0 +1,31 @@
+"""add memory_mb to task_runs
+
+Revision ID: 174dcd456325
+Revises: 67784da0203e
+Create Date: 2025-12-10 19:15:24.067819+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "174dcd456325"
+down_revision: Union[str, None] = "67784da0203e"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("task_runs", sa.Column("memory_mb", sa.Integer(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("task_runs", "memory_mb")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/client.py

@@ -4567,6 +4567,7 @@ class AgentDB:
         run_id: str,
         instance_type: str | None = None,
         vcpu_millicores: int | None = None,
+        memory_mb: int | None = None,
         duration_ms: int | None = None,
         compute_cost: float | None = None,
     ) -> None:
@@ -4589,6 +4590,8 @@ class AgentDB:
                 task_run.instance_type = instance_type
             if vcpu_millicores is not None:
                 task_run.vcpu_millicores = vcpu_millicores
+            if memory_mb is not None:
+                task_run.memory_mb = memory_mb
             if duration_ms is not None:
                 task_run.duration_ms = duration_ms
             if compute_cost is not None:

skyvern/forge/sdk/db/models.py

@@ -871,6 +871,7 @@ class TaskRunModel(Base):
     # Compute cost tracking fields
     instance_type = Column(String, nullable=True)
     vcpu_millicores = Column(Integer, nullable=True)
+    memory_mb = Column(Integer, nullable=True)
     duration_ms = Column(BigInteger, nullable=True)
     compute_cost = Column(Numeric, nullable=True)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)