diff --git a/alembic/versions/2025_12_10_1915-174dcd456325_add_memory_mb_to_task_runs.py b/alembic/versions/2025_12_10_1915-174dcd456325_add_memory_mb_to_task_runs.py new file mode 100644 index 00000000..010ae32f --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index b3c59a7e..03944b46 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/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: diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 83c29106..da0c50da 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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)