83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
import asyncio
|
|
import os
|
|
import subprocess
|
|
|
|
import structlog
|
|
import typer
|
|
|
|
from skyvern.forge import app
|
|
from skyvern.forge.forge_app_initializer import start_forge_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_forge_app()
|
|
await initialize_skyvern_state_file(task_id=None, workflow_run_id=None, organization_id=None)
|
|
|
|
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.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.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)
|
|
|
|
|
|
def main() -> None:
|
|
asyncio.run(run())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|