diff --git a/core/services/kraken/api/v2/routers/jobs.py b/core/services/kraken/api/v2/routers/jobs.py index af20a2b329..7a760fbfef 100644 --- a/core/services/kraken/api/v2/routers/jobs.py +++ b/core/services/kraken/api/v2/routers/jobs.py @@ -3,9 +3,10 @@ from typing import Any, Callable, List, Tuple from fastapi import APIRouter, Body, HTTPException, status +from fastapi.responses import StreamingResponse from fastapi_versioning import versioned_api_route from jobs import JobsManager -from jobs.exceptions import JobNotFound +from jobs.exceptions import JobIsRunning, JobNotFound from jobs.models import Job, JobMethod jobs_router_v2 = APIRouter( @@ -21,6 +22,8 @@ def jobs_to_http_exception(endpoint: Callable[..., Any]) -> Callable[..., Any]: async def wrapper(*args: Tuple[Any], **kwargs: dict[str, Any]) -> Any: try: return await endpoint(*args, **kwargs) + except JobIsRunning as error: + raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(error)) from error except JobNotFound as error: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(error)) from error except Exception as error: @@ -45,6 +48,12 @@ async def fetch() -> List[Job]: return JobsManager.get() +@jobs_router_v2.get("/{identifier}/stream", status_code=status.HTTP_200_OK) +@jobs_to_http_exception +async def stream(identifier: str) -> StreamingResponse: + return StreamingResponse(JobsManager.stream(identifier), media_type="text/plain") + + @jobs_router_v2.get("/{identifier}", status_code=status.HTTP_200_OK) @jobs_to_http_exception async def fetch_by_identifier(identifier: str) -> Job: diff --git a/core/services/kraken/jobs/exceptions.py b/core/services/kraken/jobs/exceptions.py index 32fe30fd80..a174a0a799 100644 --- a/core/services/kraken/jobs/exceptions.py +++ b/core/services/kraken/jobs/exceptions.py @@ -1,2 +1,6 @@ class JobNotFound(Exception): pass + + +class JobIsRunning(Exception): + pass diff --git a/core/services/kraken/jobs/jobs.py b/core/services/kraken/jobs/jobs.py index d3ff8d11eb..268e68d36c 100644 --- a/core/services/kraken/jobs/jobs.py +++ b/core/services/kraken/jobs/jobs.py @@ -1,15 +1,17 @@ import asyncio -from typing import List, Optional +from typing import AsyncGenerator, Dict, List, Optional import aiohttp -from jobs.exceptions import JobNotFound -from jobs.models import Job +from jobs.exceptions import JobIsRunning, JobNotFound +from jobs.models import Job, JobStatus +from jobs.stream import JobStream from loguru import logger class JobsManager: _jobs: List[Job] = [] _executing_job: Optional[Job] = None + _streams: Dict[str, JobStream] = {} def __init__(self) -> None: self.is_running = True @@ -18,6 +20,7 @@ def __init__(self) -> None: async def execute_job(self, job: Job) -> None: job_name = f"{job.method.value} - {job.route}" logger.info(f"Executing job {job_name}") + stream = self._streams[job.id] async with aiohttp.ClientSession() as session: for i in range(job.retries): try: @@ -25,20 +28,30 @@ async def execute_job(self, job: Job) -> None: method=job.method, url=f"{self.base_host}/{job.route}", json=job.body ) as response: response.raise_for_status() - await response.read() + async for chunk in response.content.iter_any(): + stream.publish(chunk) + job.status = JobStatus.SUCCESS return except Exception: logger.warning(f"Failed job {job_name} attempt {i + 1}/{job.retries}") + stream.reset() await asyncio.sleep(5) + job.status = JobStatus.ERROR logger.error(f"Job {job_name} failed to be executed") async def start(self) -> None: while self.is_running: await asyncio.sleep(1) - if self._jobs: - self._executing_job = self._jobs.pop(0) - await self.execute_job(self._executing_job) - self._executing_job = None + if JobsManager._jobs: + JobsManager._executing_job = JobsManager._jobs.pop(0) + JobsManager._executing_job.status = JobStatus.RUNNING + stream = JobsManager._streams[JobsManager._executing_job.id] + try: + await self.execute_job(JobsManager._executing_job) + finally: + stream.close() + JobsManager._streams.pop(JobsManager._executing_job.id, None) + JobsManager._executing_job = None async def stop(self) -> None: self.is_running = False @@ -48,6 +61,7 @@ def set_base_host(self, host: str) -> None: @classmethod def add(cls, job: Job) -> None: + cls._streams[job.id] = JobStream() cls._jobs.append(job) @classmethod @@ -61,7 +75,19 @@ def get_by_identifier(cls, identifier: str) -> Job: raise JobNotFound(f"Job with id {identifier} not found") return job + @classmethod + def stream(cls, identifier: str) -> AsyncGenerator[bytes, None]: + cls.get_by_identifier(identifier) + return cls._streams[identifier].subscribe() + @classmethod def delete(cls, identifier: str) -> None: - job = cls.get_by_identifier(identifier) - cls._jobs.remove(job) + cls.get_by_identifier(identifier) + job = next((job for job in cls._jobs if job.id == identifier), None) + if job: + cls._jobs.remove(job) + stream = cls._streams.pop(identifier, None) + if stream: + stream.close() + elif cls._executing_job and cls._executing_job.id == identifier: + raise JobIsRunning(f"Cannot delete job {identifier} while it is running") diff --git a/core/services/kraken/jobs/models.py b/core/services/kraken/jobs/models.py index f5d72d46bf..d655f52b55 100644 --- a/core/services/kraken/jobs/models.py +++ b/core/services/kraken/jobs/models.py @@ -13,9 +13,17 @@ class JobMethod(str, Enum): DELETE = "DELETE" +class JobStatus(str, Enum): + PENDING = "PENDING" + RUNNING = "RUNNING" + SUCCESS = "SUCCESS" + ERROR = "ERROR" + + class Job(BaseModel): id: str route: str method: JobMethod body: Any retries: int = 5 + status: JobStatus = JobStatus.PENDING diff --git a/core/services/kraken/jobs/stream.py b/core/services/kraken/jobs/stream.py new file mode 100644 index 0000000000..fc7c688c9f --- /dev/null +++ b/core/services/kraken/jobs/stream.py @@ -0,0 +1,38 @@ +import asyncio +from typing import AsyncGenerator, List + + +class JobStream: + def __init__(self) -> None: + self._chunks: List[bytes] = [] + self._updated = asyncio.Event() + self._done = False + self._retries = 0 + + def publish(self, chunk: bytes) -> None: + self._chunks.append(chunk) + self._updated.set() + + def reset(self) -> None: + self._chunks = [] + self._retries += 1 + self._updated.set() + + def close(self) -> None: + self._done = True + self._updated.set() + + async def subscribe(self) -> AsyncGenerator[bytes, None]: + index = 0 + retries = self._retries + while True: + self._updated.clear() + if retries != self._retries: + retries = self._retries + index = 0 + while index < len(self._chunks): + yield self._chunks[index] + index += 1 + if self._done: + return + await self._updated.wait()