fix task workflow cancel issue (#1111)

This commit is contained in:
Shuchang Zheng
2024-11-01 15:13:41 -07:00
committed by GitHub
parent d03ad38630
commit ad0bd0b4f5
6 changed files with 204 additions and 111 deletions

View File

@@ -10,7 +10,7 @@ from typing import Any, Awaitable, Callable, Protocol
import aiofiles
import structlog
from playwright.async_api import BrowserContext, ConsoleMessage, Error, Page, Playwright, async_playwright
from playwright.async_api import BrowserContext, ConsoleMessage, Error, Page, Playwright
from pydantic import BaseModel, PrivateAttr
from skyvern.config import settings
@@ -41,10 +41,21 @@ def get_download_dir(workflow_run_id: str | None, task_id: str | None) -> str:
return download_dir
def set_browser_console_log(browser_context: BrowserContext, browser_artifacts: BrowserArtifacts) -> str:
def set_browser_console_log(browser_context: BrowserContext, browser_artifacts: BrowserArtifacts) -> None:
if browser_artifacts.browser_console_log_path is None:
log_path = f"{settings.LOG_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}/{uuid.uuid4()}.log"
os.makedirs(os.path.dirname(log_path), exist_ok=True)
try:
log_path = f"{settings.LOG_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}/{uuid.uuid4()}.log"
os.makedirs(os.path.dirname(log_path), exist_ok=True)
# create the empty log file
with open(log_path, "w") as _:
pass
except Exception:
LOG.warning(
"Failed to create browser log file",
log_path=log_path,
exc_info=True,
)
return
browser_artifacts.browser_console_log_path = log_path
async def browser_console_log(msg: ConsoleMessage) -> None:
@@ -55,7 +66,6 @@ def set_browser_console_log(browser_context: BrowserContext, browser_artifacts:
LOG.info("browser console log is saved", log_path=browser_artifacts.browser_console_log_path)
browser_context.on("console", browser_console_log)
return browser_artifacts.browser_console_log_path
class BrowserContextCreator(Protocol):
@@ -128,6 +138,7 @@ class BrowserContextFactory:
cls, playwright: Playwright, **kwargs: Any
) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]:
browser_type = SettingsManager.get_settings().BROWSER_TYPE
browser_context: BrowserContext | None = None
try:
creator = cls._creators.get(browser_type)
if not creator:
@@ -135,9 +146,15 @@ class BrowserContextFactory:
browser_context, browser_artifacts, cleanup_func = await creator(playwright, **kwargs)
set_browser_console_log(browser_context=browser_context, browser_artifacts=browser_artifacts)
return browser_context, browser_artifacts, cleanup_func
except UnknownBrowserType as e:
raise e
except Exception as e:
if browser_context is not None:
# FIXME: sometimes it can't close the browser context?
LOG.error("unexpected error happens after created browser context, going to close the context")
await browser_context.close()
if isinstance(e, UnknownBrowserType):
raise e
raise UnknownErrorWhileCreatingBrowserContext(browser_type, e) from e
@classmethod
@@ -217,7 +234,7 @@ class BrowserState:
def __init__(
self,
pw: Playwright | None = None,
pw: Playwright,
browser_context: BrowserContext | None = None,
page: Page | None = None,
browser_artifacts: BrowserArtifacts = BrowserArtifacts(),
@@ -253,10 +270,6 @@ class BrowserState:
workflow_run_id: str | None = None,
organization_id: str | None = None,
) -> None:
if self.pw is None:
LOG.info("Starting playwright")
self.pw = await async_playwright().start()
LOG.info("playwright is started")
if self.browser_context is None:
LOG.info("creating browser context")
(
@@ -276,8 +289,6 @@ class BrowserState:
self.browser_cleanup = browser_cleanup
LOG.info("browser context is created")
assert self.browser_context is not None
if await self.get_working_page() is None:
success = False
retries = 0

View File

@@ -53,17 +53,25 @@ class BrowserManager:
browser_cleanup=browser_cleanup,
)
async def get_or_create_for_task(self, task: Task) -> BrowserState:
if task.task_id in self.pages:
return self.pages[task.task_id]
elif task.workflow_run_id in self.pages:
def get_for_task(self, task_id: str, workflow_run_id: str | None = None) -> BrowserState | None:
if task_id in self.pages:
return self.pages[task_id]
if workflow_run_id and workflow_run_id in self.pages:
LOG.info(
"Browser state for task not found. Using browser state for workflow run",
task_id=task.task_id,
workflow_run_id=task.workflow_run_id,
task_id=task_id,
workflow_run_id=workflow_run_id,
)
self.pages[task.task_id] = self.pages[task.workflow_run_id]
return self.pages[task.task_id]
self.pages[task_id] = self.pages[workflow_run_id]
return self.pages[task_id]
return None
async def get_or_create_for_task(self, task: Task) -> BrowserState:
browser_state = self.get_for_task(task_id=task.task_id, workflow_run_id=task.workflow_run_id)
if browser_state is not None:
return browser_state
LOG.info("Creating browser state for task", task_id=task.task_id)
browser_state = await self._create_browser_state(
@@ -85,8 +93,10 @@ class BrowserManager:
return browser_state
async def get_or_create_for_workflow_run(self, workflow_run: WorkflowRun, url: str | None = None) -> BrowserState:
if workflow_run.workflow_run_id in self.pages:
return self.pages[workflow_run.workflow_run_id]
browser_state = self.get_for_workflow_run(workflow_run_id=workflow_run.workflow_run_id)
if browser_state is not None:
return browser_state
LOG.info(
"Creating browser state for workflow run",
workflow_run_id=workflow_run.workflow_run_id,
@@ -110,7 +120,7 @@ class BrowserManager:
self.pages[workflow_run.workflow_run_id] = browser_state
return browser_state
async def get_for_workflow_run(self, workflow_run_id: str) -> BrowserState | None:
def get_for_workflow_run(self, workflow_run_id: str) -> BrowserState | None:
if workflow_run_id in self.pages:
return self.pages[workflow_run_id]
return None
@@ -230,7 +240,18 @@ class BrowserManager:
await browser_state_to_close.close(close_browser_on_completion=close_browser_on_completion)
for task_id in task_ids:
self.pages.pop(task_id, None)
task_browser_state = self.pages.pop(task_id, None)
if task_browser_state is None:
continue
try:
await task_browser_state.close()
except Exception:
LOG.info(
"Failed to close the browser state from the task block, might because it's already closed.",
exc_info=True,
task_id=task_id,
workflow_run_id=workflow_run_id,
)
LOG.info("Workflow run is cleaned up")
return browser_state_to_close