catch asyncio CancelledError (#337)

This commit is contained in:
Shuchang Zheng 2024-05-17 10:02:03 -07:00 committed by GitHub
parent 9a6ec13af6
commit 215c74b123
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,6 +1,7 @@
import asyncio import asyncio
import json import json
import random import random
from asyncio.exceptions import CancelledError
from datetime import datetime from datetime import datetime
from typing import Any, Tuple from typing import Any, Tuple
@ -538,6 +539,20 @@ class ForgeAgent:
output=detailed_agent_step_output.to_agent_step_output(), output=detailed_agent_step_output.to_agent_step_output(),
) )
return completed_step, detailed_agent_step_output return completed_step, detailed_agent_step_output
except CancelledError:
LOG.exception(
"CancelledError in agent_step, marking step as failed",
task_id=task.task_id,
step_id=step.step_id,
step_order=step.order,
step_retry=step.retry_index,
)
failed_step = await self.update_step(
step=step,
status=StepStatus.failed,
output=detailed_agent_step_output.to_agent_step_output(),
)
return failed_step, detailed_agent_step_output
except Exception: except Exception:
LOG.exception( LOG.exception(
"Unexpected exception in agent_step, marking step as failed", "Unexpected exception in agent_step, marking step as failed",