run_streaming.py

Rama main · réplica de lectura del repositorio original.

import asyncio
import os
import subprocess

import structlog
import typer

from skyvern.forge import app
from skyvern.forge.forge_app_initializer import start_streaming_worker_app
from skyvern.forge.sdk.api.files import get_skyvern_temp_dir
from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus
from skyvern.utils.files import get_json_from_file, get_skyvern_state_file_path, initialize_skyvern_state_file

INTERVAL = 1
LOG = structlog.get_logger()


async def run() -> None:
    start_streaming_worker_app()
    await initialize_skyvern_state_file(task_id=None, workflow_run_id=None, organization_id=None)

    try:
        while True:
            await asyncio.sleep(INTERVAL)
            try:
                current_json = get_json_from_file(get_skyvern_state_file_path())
            except Exception:
                continue

            task_id = current_json.get("task_id")
            workflow_run_id = current_json.get("workflow_run_id")
            organization_id = current_json.get("organization_id")
            if not organization_id or (not task_id and not workflow_run_id):
                continue

            try:
                if workflow_run_id:
                    workflow_run = await app.DATABASE.workflow_runs.get_workflow_run(workflow_run_id=workflow_run_id)
                    if not workflow_run or workflow_run.status in [
                        WorkflowRunStatus.completed,
                        WorkflowRunStatus.failed,
                        WorkflowRunStatus.terminated,
                    ]:
                        continue
                    file_name = f"{workflow_run_id}.png"

                elif task_id:
                    task = await app.DATABASE.tasks.get_task(task_id=task_id, organization_id=organization_id)
                    if not task or task.status.is_final():
                        continue
                    file_name = f"{task_id}.png"
                else:
                    continue
            except Exception:
                LOG.exception(
                    "Failed to get task or workflow run while taking streaming screenshot in worker",
                    task_id=task_id,
                    workflow_run_id=workflow_run_id,
                    organization_id=organization_id,
                )
                continue

            # create f"{get_skyvern_temp_dir()}/{organization_id}" directory if it does not exists
            os.makedirs(f"{get_skyvern_temp_dir()}/{organization_id}", exist_ok=True)
            png_file_path = f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}"

            # run subprocess to take screenshot
            subprocess.run(
                f"xwd -root | xwdtopnm 2>/dev/null | pnmtopng > {png_file_path}", shell=True, env={"DISPLAY": ":99"}
            )

            try:
                await app.STORAGE.save_streaming_file(organization_id, file_name)
            except Exception:
                LOG.debug("Failed to upload screenshot", organization_id=organization_id, file_name=file_name)
    finally:
        # Dispose the DB engine/pool this worker owns. The minimal streaming-worker
        # runtime makes this worker the sole owner of app.DATABASE, so releasing it
        # when the loop is cancelled or exits (task cancellation / SIGINT / an
        # unhandled error) frees the connection pool deterministically. A SIGKILL
        # container stop bypasses this, but the OS reclaims the pool with the process.
        try:
            await app.DATABASE.engine.dispose()
        except Exception:
            LOG.debug("Failed to dispose streaming worker database engine")


def main() -> None:
    asyncio.run(run())


if __name__ == "__main__":
    typer.run(main)