add actions to workflow run profiling (#4610)

This commit is contained in:
Shuchang Zheng
2026-02-02 22:28:50 -08:00
committed by GitHub
parent fcbe7fe84f
commit 5fc16b6a5f

View File

@@ -3,6 +3,7 @@ Script to profile a workflow run by collecting and displaying all key timestamps
Usage: Usage:
python scripts/profile_workflow_run.py <workflow_run_id> python scripts/profile_workflow_run.py <workflow_run_id>
python scripts/profile_workflow_run.py <workflow_run_id> --include-actions
""" """
import asyncio import asyncio
@@ -16,6 +17,7 @@ from sqlalchemy import select
from skyvern.forge import app from skyvern.forge import app
from skyvern.forge.forge_app_initializer import start_forge_app from skyvern.forge.forge_app_initializer import start_forge_app
from skyvern.forge.sdk.db.models import ( from skyvern.forge.sdk.db.models import (
ActionModel,
StepModel, StepModel,
TaskModel, TaskModel,
WorkflowRunBlockModel, WorkflowRunBlockModel,
@@ -28,7 +30,7 @@ class TimestampEntry:
"""Represents a single timestamp entry for profiling.""" """Represents a single timestamp entry for profiling."""
timestamp: datetime timestamp: datetime
entity_type: str # "workflow_run", "workflow_run_block", "task", "step" entity_type: str # "workflow_run", "workflow_run_block", "task", "step", "action"
entity_id: str entity_id: str
field_name: str # "created_at", "started_at", "finished_at", etc. field_name: str # "created_at", "started_at", "finished_at", etc.
label: str | None = None # For blocks with labels label: str | None = None # For blocks with labels
@@ -40,7 +42,7 @@ class TimestampEntry:
return f"{self.timestamp.isoformat()} | {self.entity_type:20}{label_str} | {self.field_name:12} | {self.entity_id}{status_str}" return f"{self.timestamp.isoformat()} | {self.entity_type:20}{label_str} | {self.field_name:12} | {self.entity_id}{status_str}"
async def collect_timestamps(workflow_run_id: str) -> list[TimestampEntry]: async def collect_timestamps(workflow_run_id: str, include_actions: bool = False) -> list[TimestampEntry]:
"""Collect all timestamps from the workflow run and its children.""" """Collect all timestamps from the workflow run and its children."""
entries: list[TimestampEntry] = [] entries: list[TimestampEntry] = []
@@ -136,6 +138,29 @@ async def collect_timestamps(workflow_run_id: str) -> list[TimestampEntry]:
) )
) )
# 5. Fetch all actions for all tasks (optional)
if include_actions and task_ids:
actions = (
await session.scalars(
select(ActionModel).filter(ActionModel.task_id.in_(task_ids)).order_by(ActionModel.created_at)
)
).all()
for action in actions:
for field in ["modified_at"]:
ts = getattr(action, field, None)
if ts:
entries.append(
TimestampEntry(
timestamp=ts,
entity_type="action",
entity_id=action.action_id,
field_name=field,
label=action.action_type,
status=action.status,
)
)
return entries return entries
@@ -186,20 +211,25 @@ def print_profile(entries: list[TimestampEntry]) -> None:
print("=" * 120 + "\n") print("=" * 120 + "\n")
async def profile_workflow_run(workflow_run_id: str) -> None: async def profile_workflow_run(workflow_run_id: str, include_actions: bool = False) -> None:
"""Main function to profile a workflow run.""" """Main function to profile a workflow run."""
print(f"Profiling workflow run: {workflow_run_id}") print(f"Profiling workflow run: {workflow_run_id}")
if include_actions:
print("(including actions)")
entries = await collect_timestamps(workflow_run_id) entries = await collect_timestamps(workflow_run_id, include_actions=include_actions)
print_profile(entries) print_profile(entries)
def main( def main(
workflow_run_id: Annotated[str, typer.Argument(help="The workflow run ID to profile")], workflow_run_id: Annotated[str, typer.Argument(help="The workflow run ID to profile")],
include_actions: Annotated[
bool, typer.Option("--include-actions", "-a", help="Include action timestamps (can be noisy)")
] = False,
) -> None: ) -> None:
"""Profile a workflow run by collecting and displaying all key timestamps.""" """Profile a workflow run by collecting and displaying all key timestamps."""
start_forge_app() start_forge_app()
asyncio.run(profile_workflow_run(workflow_run_id)) asyncio.run(profile_workflow_run(workflow_run_id, include_actions=include_actions))
if __name__ == "__main__": if __name__ == "__main__":