wehub-resource-sync
136 lines
3.6 KiB
Python
136 lines
3.6 KiB
Python
"""
|
|
Internal job APIs for UI invocation
|
|
"""
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
from mlflow.entities._job import Job as JobEntity
|
|
from mlflow.entities._job_status import JobStatus
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
job_api_router = APIRouter(prefix="/ajax-api/3.0/jobs", tags=["Job"])
|
|
|
|
|
|
class Job(BaseModel):
|
|
"""
|
|
Pydantic model for job query response.
|
|
"""
|
|
|
|
job_id: str
|
|
creation_time: int
|
|
job_name: str
|
|
params: dict[str, Any]
|
|
timeout: float | None
|
|
status: JobStatus
|
|
result: Any
|
|
retry_count: int
|
|
last_update_time: int
|
|
status_details: dict[str, Any] | None = None
|
|
|
|
@classmethod
|
|
def from_job_entity(cls, job: JobEntity) -> "Job":
|
|
return cls(
|
|
job_id=job.job_id,
|
|
creation_time=job.creation_time,
|
|
job_name=job.job_name,
|
|
params=json.loads(job.params),
|
|
timeout=job.timeout,
|
|
status=job.status,
|
|
result=job.parsed_result,
|
|
retry_count=job.retry_count,
|
|
last_update_time=job.last_update_time,
|
|
status_details=job.status_details,
|
|
)
|
|
|
|
|
|
@job_api_router.get("/{job_id}", response_model=Job)
|
|
def get_job(job_id: str) -> Job:
|
|
from mlflow.server.jobs import get_job
|
|
|
|
try:
|
|
job = get_job(job_id)
|
|
return Job.from_job_entity(job)
|
|
except MlflowException as e:
|
|
raise HTTPException(
|
|
status_code=e.get_http_status_code(),
|
|
detail=e.message,
|
|
)
|
|
|
|
|
|
class SubmitJobPayload(BaseModel):
|
|
job_name: str
|
|
params: dict[str, Any]
|
|
timeout: float | None = None
|
|
|
|
|
|
@job_api_router.post("/", response_model=Job)
|
|
def submit_job(payload: SubmitJobPayload) -> Job:
|
|
from mlflow.server.jobs import submit_job
|
|
from mlflow.server.jobs.utils import _load_function, get_job_fn_fullname
|
|
|
|
job_name = payload.job_name
|
|
try:
|
|
function_fullname = get_job_fn_fullname(job_name)
|
|
function = _load_function(function_fullname)
|
|
job = submit_job(function, payload.params, payload.timeout)
|
|
return Job.from_job_entity(job)
|
|
except MlflowException as e:
|
|
raise HTTPException(
|
|
status_code=e.get_http_status_code(),
|
|
detail=e.message,
|
|
)
|
|
|
|
|
|
@job_api_router.patch("/cancel/{job_id}", response_model=Job)
|
|
def cancel_job(job_id: str) -> Job:
|
|
from mlflow.server.jobs import cancel_job
|
|
|
|
try:
|
|
job = cancel_job(job_id)
|
|
return Job.from_job_entity(job)
|
|
except MlflowException as e:
|
|
raise HTTPException(
|
|
status_code=e.get_http_status_code(),
|
|
detail=e.message,
|
|
)
|
|
|
|
|
|
class SearchJobPayload(BaseModel):
|
|
job_name: str | None = None
|
|
params: dict[str, Any] | None = None
|
|
statuses: list[JobStatus] | None = None
|
|
|
|
|
|
class SearchJobsResponse(BaseModel):
|
|
"""
|
|
Pydantic model for job searching response.
|
|
"""
|
|
|
|
jobs: list[Job]
|
|
|
|
|
|
@job_api_router.post("/search", response_model=SearchJobsResponse)
|
|
def search_jobs(payload: SearchJobPayload) -> SearchJobsResponse:
|
|
from mlflow.server.handlers import _get_job_store
|
|
|
|
try:
|
|
store = _get_job_store()
|
|
job_results = [
|
|
Job.from_job_entity(job)
|
|
for job in store.list_jobs(
|
|
job_name=payload.job_name,
|
|
statuses=payload.statuses,
|
|
params=payload.params,
|
|
)
|
|
]
|
|
return SearchJobsResponse(jobs=job_results)
|
|
except MlflowException as e:
|
|
raise HTTPException(
|
|
status_code=e.get_http_status_code(),
|
|
detail=e.message,
|
|
)
|