58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
# 代码内容与原文件基本一致，确保它是通用的
|
||
import httpx
|
||
from typing import Optional
|
||
from app.core.logging import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
# xuhuanqing
|
||
class RuntimeClient:
|
||
"""HTTP client for communicating with deepdata-runtime service"""
|
||
|
||
def __init__(self, base_url: str = "http://deepdata-runtime:8081"):
|
||
self.base_url = base_url
|
||
self.client = httpx.AsyncClient(timeout=60.0)
|
||
|
||
async def submit_task(self, task_id: str, retry_count: int = 0) -> bool:
|
||
"""Submit task to runtime executor (通用接口)"""
|
||
try:
|
||
url = f"{self.base_url}/api/task/{task_id}/submit"
|
||
# 注意：虽然 json payload 目前是 {retry_count}，
|
||
# 但 Runtime 接收的是一个通用的 JSON，未来可以扩展增加 executor_type 等字段
|
||
response = await self.client.post(url, json={"retry_count": retry_count})
|
||
response.raise_for_status()
|
||
logger.info(f"Task {task_id} submitted to Runtime successfully.")
|
||
return True
|
||
except httpx.HTTPError as e:
|
||
logger.error(f"Failed to submit task {task_id} to Runtime: {e}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error submitting task {task_id}: {e}")
|
||
return False
|
||
|
||
async def stop_task(self, task_id: str) -> bool:
|
||
"""Stop running task"""
|
||
try:
|
||
url = f"{self.base_url}/api/task/{task_id}/stop"
|
||
response = await self.client.post(url)
|
||
response.raise_for_status()
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Failed to stop task {task_id}: {e}")
|
||
return False
|
||
|
||
async def get_task_status(self, task_id: str) -> Optional[dict]:
|
||
"""Get task execution status from runtime"""
|
||
try:
|
||
url = f"{self.base_url}/api/task/{task_id}/status"
|
||
response = await self.client.get(url)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
except Exception as e:
|
||
logger.error(f"Failed to get task status {task_id}: {e}")
|
||
return None
|
||
|
||
async def close(self):
|
||
"""Close HTTP client"""
|
||
await self.client.aclose()
|