| from enum import Enum |
| import json |
| import logging |
| from pathlib import Path as PyPath |
| import traceback |
| import yaml |
| import sys |
| import shutil |
| import zipfile |
| import httpx |
| import os |
| import tempfile |
| import time |
| import subprocess |
|
|
| from fastapi import ( |
| FastAPI, |
| HTTPException, |
| UploadFile, |
| File, |
| BackgroundTasks, |
| Path as FastApiPath, |
| Query, |
| Body, |
| Request, |
| ) |
| from fastapi.responses import FileResponse, StreamingResponse, JSONResponse, RedirectResponse |
| from fastapi.staticfiles import StaticFiles |
| from typing import List, Optional, Dict, Any |
| from pydantic import BaseModel, Field |
|
|
| from fastapi.middleware.cors import CORSMiddleware |
| from contextlib import asynccontextmanager |
| from ocr.config import Settings as OCRSettings |
| from ocr.worker_api import create_app as create_ocr_app |
|
|
|
|
| from huey_config import huey |
| from tasks import ( |
| run_pipeline_task, |
| set_final_status, |
| resolve_job_output_dir, |
| ) |
| from huey.exceptions import ( |
| HueyException, |
| TaskException, |
| ) |
| from redis_config import ( |
| get_progress, |
| redis_client, |
| set_progress, |
| ) |
| import signal |
|
|
| from resolve_path import resolve_path |
|
|
| from file_operation_helpers import ( |
| get_safe_path, |
| zip_directory, |
| get_dir_structure, |
| FileStructure, |
| MoveItemRequest, |
| handle_get_structure, |
| handle_download_item, |
| handle_delete_item, |
| handle_move_item, |
| handle_create_directory, |
| ) |
|
|
| from augmentoolkit.hf_persistence import ( |
| HFHubManager, |
| CheckpointManager, |
| JobManifestManager, |
| IncompatibleResumeError, |
| JobManifest, |
| ) |
|
|
|
|
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
|
|
| SUPER_CONFIG_PATH = PyPath("super_config.yaml") |
| PATH_ALIASES = {} |
|
|
| try: |
| with open(SUPER_CONFIG_PATH, "r", encoding="utf-8") as f: |
| super_config = yaml.safe_load(f) |
| PATH_ALIASES = super_config.get("path_aliases", {}) |
| except FileNotFoundError: |
| print(f"ERROR: Super config file not found at {SUPER_CONFIG_PATH}.") |
| except yaml.YAMLError as e: |
| print(f"ERROR: Error parsing super config file {SUPER_CONFIG_PATH}: {e}") |
|
|
| DATA_ROOT = PyPath(os.environ.get("ATK_DATA_ROOT", "/tmp/augmentoolkit")).resolve() |
| INPUTS_DIR = (DATA_ROOT / "inputs").resolve() |
| OUTPUTS_DIR = (DATA_ROOT / "outputs").resolve() |
| GENERATION_DIR = (DATA_ROOT / "generation").resolve() |
| CONFIGS_DIR = (DATA_ROOT / "external_configs").resolve() |
| MAX_FILE_SIZE = 1024 * 1024 * 1000 |
|
|
| LOGS_DIR = (DATA_ROOT / "logs").resolve() |
|
|
| INPUTS_DIR.mkdir(parents=True, exist_ok=True) |
| OUTPUTS_DIR.mkdir(parents=True, exist_ok=True) |
| GENERATION_DIR.mkdir(parents=True, exist_ok=True) |
| CONFIGS_DIR.mkdir(parents=True, exist_ok=True) |
| LOGS_DIR.mkdir(parents=True, exist_ok=True) |
|
|
|
|
| class PipelineRunRequest(BaseModel): |
| node_path: str |
| config_path: Optional[str] = None |
| parameters: Optional[Dict[str, Any]] = None |
|
|
|
|
| class PipelineRunResponse(BaseModel): |
| pipeline_id: str |
| message: str |
|
|
|
|
| class PipelineStatus(str, Enum): |
| PENDING = "PENDING" |
| RUNNING = "RUNNING" |
| COMPLETED = "COMPLETED" |
| FAILED = "FAILED" |
| REVOKED = "REVOKED" |
|
|
|
|
| class PipelineStatusResponse(BaseModel): |
| task_id: str |
| status: PipelineStatus |
| message: Optional[str] = None |
| progress: Optional[float] = Field(None, ge=0.0, le=1.0) |
| details: Optional[Dict[str, Any]] = None |
|
|
|
|
| class QueueStatusResponse(BaseModel): |
| pending_tasks: List[str] = Field(..., description="List of task IDs currently pending execution.") |
| scheduled_tasks: List[str] = Field(..., description="List of task IDs scheduled for future execution.") |
| message: str |
|
|
|
|
| class CreateDirectoryRequest(BaseModel): |
| relative_path: str = Field(..., description="The relative path within base directory.") |
|
|
|
|
| class DuplicateConfigRequest(BaseModel): |
| source_alias: str = Field(..., description="Alias from super_config.yaml.") |
| destination_relative_path: str = Field(..., description="Desired relative path.") |
|
|
|
|
| class TaskParametersResponse(BaseModel): |
| task_id: str |
| parameters: Dict[str, Any] |
|
|
|
|
|
|
| ocr_settings = OCRSettings.from_env() |
| ocr_app = create_ocr_app(settings=ocr_settings) |
|
|
| @asynccontextmanager |
| async def app_lifespan(main_app: FastAPI): |
| async with ocr_app.router.lifespan_context(ocr_app): |
| |
| logger.info("Augmentoolkit API starting up...") |
| work_repo = os.environ.get("HF_WORK_REPO") |
| if work_repo: |
| try: |
| ckpt_mgr = CheckpointManager() |
| incomplete_jobs = ckpt_mgr.discover_incomplete_jobs(work_repo=work_repo) |
| logger.info(f"Discovered {len(incomplete_jobs)} incomplete job(s) in HF_WORK_REPO '{work_repo}'.") |
| auto_resume = os.environ.get("HF_AUTO_RESUME", "false").lower() in ("true", "1", "yes") |
| if auto_resume: |
| for manifest in incomplete_jobs: |
| logger.info(f"Auto-resuming durable job {manifest.job_id}...") |
| run_pipeline_task( |
| node_path=manifest.pipeline, |
| parameters={ |
| "job_id": manifest.job_id, |
| "hf_source_repo": manifest.source_repo, |
| "hf_work_repo": manifest.work_repo, |
| "hf_output_repo": manifest.output_repo, |
| "source_revision": manifest.source_revision, |
| }, |
| ) |
| except Exception as e: |
| logger.error(f"Error during startup HF recovery discovery: {e}") |
| yield |
|
|
| app = FastAPI( |
| title="Augmentoolkit Resumable Dataset Factory API", |
| description="API for managing and running Augmentoolkit dataset pipelines with Hugging Face Hub remote persistence.", |
| version="2.0", |
| lifespan=app_lifespan, |
| ) |
|
|
| app.mount("/ocr", ocr_app) |
|
|
|
|
| origins = [ |
| "http://localhost:5173", |
| "http://127.0.0.1:5173", |
| "http://localhost:5174", |
| "http://127.0.0.1:5174", |
| "http://localhost:3000", |
| ] |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=origins, |
| allow_credentials=True, |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
|
|
|
|
|
|
| @app.get("/health", summary="Health Check") |
| async def health_check(): |
| """Basic health check endpoint for Hugging Face.""" |
| return {"status": "ok", "message": "Augmentoolkit API is running."} |
|
|
|
|
| @app.get("/readiness", summary="Readiness Check") |
| async def readiness_check(): |
| """ |
| Comprehensive readiness check reporting Redis/Huey connection, Hugging Face authentication, |
| repository accessibility (HF_SOURCE_REPO, HF_WORK_REPO, HF_OUTPUT_REPO), and worker status. |
| """ |
| redis_ok = False |
| redis_error = None |
| if redis_client: |
| try: |
| redis_client.ping() |
| redis_ok = True |
| except Exception as e: |
| redis_error = str(e) |
| else: |
| redis_error = "Redis client not initialized" |
|
|
| hf_mgr = HFHubManager() |
| hf_health = hf_mgr.check_all_health() |
|
|
| worker_ok = False |
| if huey: |
| try: |
| huey.pending() |
| worker_ok = True |
| except Exception: |
| pass |
|
|
| overall_ready = redis_ok and worker_ok and hf_health.get("healthy", False) |
|
|
| return { |
| "status": "ready" if overall_ready else "degraded", |
| "redis": {"status": "ok" if redis_ok else "error", "error": redis_error}, |
| "huey_queue": {"status": "ok" if worker_ok else "error"}, |
| "huggingface": hf_health, |
| } |
|
|
|
|
| @app.get("/api-docs", include_in_schema=False) |
| async def api_docs_redirect(): |
| return RedirectResponse(url="/docs") |
|
|
|
|
| @app.post( |
| "/pipelines/run", |
| response_model=PipelineRunResponse, |
| status_code=202, |
| summary="Queue a dataset generation pipeline for execution.", |
| ) |
| def queue_pipeline_run(request: PipelineRunRequest): |
| try: |
| if huey is None: |
| raise HTTPException(status_code=503, detail="Job functionality disabled: Huey/Redis unavailable.") |
| task = run_pipeline_task( |
| node_path=request.node_path, |
| config_path=request.config_path, |
| parameters=request.parameters, |
| ) |
| return PipelineRunResponse( |
| pipeline_id=task.id, |
| message="Pipeline run queued successfully.", |
| ) |
| except Exception as e: |
| logger.error(f"ERROR during pipeline queueing: {e}", exc_info=True) |
| raise HTTPException( |
| status_code=500, |
| detail=f"Failed to enqueue pipeline task: {e}", |
| ) |
|
|
|
|
| @app.get("/jobs", summary="List jobs from HF_WORK_REPO") |
| def list_jobs(work_repo: Optional[str] = Query(None)): |
| target_repo = work_repo or os.environ.get("HF_WORK_REPO", "") |
| if not target_repo: |
| return {"work_repo": "", "jobs": []} |
| hf_mgr = HFHubManager() |
| hf_jobs = hf_mgr.list_jobs(work_repo=target_repo) |
| return {"work_repo": target_repo, "jobs": hf_jobs} |
|
|
|
|
| @app.get("/jobs/{job_id}/manifest", summary="Inspect job manifest") |
| def inspect_job_manifest(job_id: str, work_repo: Optional[str] = Query(None)): |
| target_repo = work_repo or os.environ.get("HF_WORK_REPO", "") |
| hf_mgr = HFHubManager() |
| manifest = hf_mgr.fetch_manifest(job_id=job_id, work_repo=target_repo) |
| if not manifest: |
| raise HTTPException(status_code=404, detail=f"Manifest for job '{job_id}' not found.") |
| return manifest |
|
|
|
|
| @app.get("/jobs/{job_id}/checkpoints", summary="Inspect job checkpoints") |
| def inspect_job_checkpoints(job_id: str, work_repo: Optional[str] = Query(None)): |
| target_repo = work_repo or os.environ.get("HF_WORK_REPO", "") |
| if not target_repo: |
| raise HTTPException(status_code=400, detail="HF_WORK_REPO is not configured.") |
| hf_mgr = HFHubManager() |
| try: |
| files = hf_mgr.api.list_repo_files(repo_id=target_repo, repo_type="dataset", token=hf_mgr.token) |
| ckpts = [f for f in files if f.startswith(f"jobs/{job_id}/checkpoints/")] |
| return {"job_id": job_id, "checkpoints": ckpts} |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error checking checkpoints: {e}") |
|
|
|
|
| @app.get("/jobs/{job_id}/outputs", summary="Inspect job output references") |
| def inspect_job_outputs(job_id: str, output_repo: Optional[str] = Query(None)): |
| target_repo = output_repo or os.environ.get("HF_OUTPUT_REPO", "") |
| if not target_repo: |
| raise HTTPException(status_code=400, detail="HF_OUTPUT_REPO is not configured.") |
| hf_mgr = HFHubManager() |
| try: |
| files = hf_mgr.api.list_repo_files(repo_id=target_repo, repo_type="dataset", token=hf_mgr.token) |
| outputs = [f for f in files if f.startswith(f"datasets/{job_id}/")] |
| return {"job_id": job_id, "output_repo": target_repo, "outputs": outputs} |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error checking outputs: {e}") |
|
|
|
|
| @app.post("/jobs/{job_id}/resume", summary="Explicitly resume a job from remote state") |
| def resume_job(job_id: str, parameters: Optional[Dict[str, Any]] = Body(None)): |
| work_repo = os.environ.get("HF_WORK_REPO", "") |
| ckpt_mgr = CheckpointManager() |
| manifest = ckpt_mgr.load_remote_manifest(job_id=job_id, work_repo=work_repo) |
| if not manifest: |
| raise HTTPException(status_code=404, detail=f"Job '{job_id}' not found in {work_repo}.") |
|
|
| target_output_dir = resolve_job_output_dir( |
| node_path=manifest.pipeline, |
| config_path=parameters.get("config_path") if parameters else None, |
| parameters=parameters, |
| durable_job_id=job_id, |
| ) |
|
|
| live_prompt_hash = JobManifestManager.compute_live_prompt_hash( |
| prompt_folder=parameters.get("prompt_folder") if parameters else None, |
| default_prompt_folder=parameters.get("default_prompt_folder") if parameters else None, |
| ) |
|
|
| try: |
| reconciled_manifest, completed_shards, restored_data = ckpt_mgr.reconcile_and_resume_job( |
| job_id=job_id, |
| target_pipeline=manifest.pipeline, |
| target_config_dict=parameters or {}, |
| target_prompt_hash=live_prompt_hash, |
| target_source_repo=manifest.source_repo, |
| target_source_revision=manifest.source_revision, |
| target_source_path=manifest.source_path, |
| target_source_split=manifest.source_split, |
| target_output_schema_version=manifest.output_schema_version, |
| work_repo=work_repo, |
| output_dir=target_output_dir, |
| ) |
| except IncompatibleResumeError as e: |
| raise HTTPException(status_code=400, detail=str(e)) |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Failed to reconcile job state: {e}") |
|
|
| run_params = parameters or {} |
| run_params.update({ |
| "job_id": job_id, |
| "output_dir": target_output_dir, |
| "hf_source_repo": manifest.source_repo, |
| "hf_work_repo": manifest.work_repo, |
| "hf_output_repo": manifest.output_repo, |
| "source_revision": manifest.source_revision, |
| "restored_shards": list(completed_shards), |
| }) |
|
|
| if huey is None: |
| raise HTTPException(status_code=503, detail="Job functionality disabled: Huey/Redis unavailable.") |
| task = run_pipeline_task( |
| node_path=manifest.pipeline, |
| parameters=run_params, |
| ) |
| return { |
| "job_id": job_id, |
| "task_id": task.id, |
| "output_dir": target_output_dir, |
| "message": f"Resume task for durable job '{job_id}' enqueued successfully.", |
| "completed_shards": list(completed_shards), |
| } |
|
|
|
|
| @app.post("/jobs/{job_id}/fork", summary="Fork an existing job run into a new durable job") |
| def fork_job( |
| job_id: str, |
| new_job_id: Optional[str] = Query(None), |
| parameters: Optional[Dict[str, Any]] = Body(None), |
| ): |
| work_repo = os.environ.get("HF_WORK_REPO", "") |
| ckpt_mgr = CheckpointManager() |
| parent_manifest = ckpt_mgr.load_remote_manifest(job_id=job_id, work_repo=work_repo) |
| if not parent_manifest: |
| raise HTTPException(status_code=404, detail=f"Parent job '{job_id}' not found in {work_repo}.") |
|
|
| forked_id = new_job_id or f"{job_id}_fork_{int(time.time())}" |
|
|
| fork_params = parameters or {} |
| run_params = { |
| "job_id": forked_id, |
| "forked_from_job_id": job_id, |
| "hf_source_repo": fork_params.get("hf_source_repo", parent_manifest.source_repo), |
| "hf_work_repo": work_repo, |
| "hf_output_repo": fork_params.get("hf_output_repo", parent_manifest.output_repo), |
| "source_revision": fork_params.get("source_revision", parent_manifest.source_revision), |
| "source_path": fork_params.get("source_path", parent_manifest.source_path), |
| "source_split": fork_params.get("source_split", parent_manifest.source_split), |
| } |
| run_params.update(fork_params) |
|
|
| if huey is None: |
| raise HTTPException(status_code=503, detail="Job functionality disabled: Huey/Redis unavailable.") |
| task = run_pipeline_task( |
| node_path=parent_manifest.pipeline, |
| parameters=run_params, |
| ) |
|
|
| return { |
| "forked_job_id": forked_id, |
| "parent_job_id": job_id, |
| "task_id": task.id, |
| "message": "Forked job queued successfully.", |
| } |
|
|
|
|
| @app.get( |
| "/pipelines/available", |
| response_model=List[str], |
| summary="Get available pipeline aliases from super_config.yaml.", |
| ) |
| def get_available_pipelines(): |
| available_pipelines = [] |
| if not SUPER_CONFIG_PATH.exists(): |
| return [] |
|
|
| try: |
| with open(SUPER_CONFIG_PATH, "r", encoding="utf-8") as f: |
| super_config = yaml.safe_load(f) or {} |
|
|
| aliases = super_config.get("path_aliases", {}) |
| if not isinstance(aliases, dict): |
| return [] |
|
|
| for alias, path_value in aliases.items(): |
| if isinstance(path_value, str) and not path_value.strip().lower().endswith(".yaml"): |
| available_pipelines.append(alias) |
|
|
| return sorted(available_pipelines) |
|
|
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error parsing super config: {e}") |
|
|
|
|
| @app.get( |
| "/configs/aliases", |
| response_model=List[str], |
| summary="Get available config file aliases from super_config.yaml.", |
| ) |
| def get_available_config_aliases(): |
| config_aliases = [] |
| if not SUPER_CONFIG_PATH.exists(): |
| return [] |
|
|
| try: |
| with open(SUPER_CONFIG_PATH, "r", encoding="utf-8") as f: |
| super_config = yaml.safe_load(f) or {} |
|
|
| aliases = super_config.get("path_aliases", {}) |
| if not isinstance(aliases, dict): |
| return [] |
|
|
| for alias, path_value in aliases.items(): |
| if isinstance(path_value, str) and path_value.strip().lower().endswith(".yaml"): |
| config_aliases.append(alias) |
|
|
| return sorted(config_aliases) |
|
|
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error parsing super config: {e}") |
|
|
|
|
| @app.get( |
| "/tasks/queue", |
| response_model=QueueStatusResponse, |
| summary="Get lists of pending and scheduled tasks.", |
| ) |
| def get_queue_status(): |
| try: |
| pending_tasks = [task.id for task in huey.pending()] if huey else [] |
| scheduled_tasks = [task.id for task in huey.scheduled()] if huey else [] |
| return QueueStatusResponse( |
| pending_tasks=pending_tasks, |
| scheduled_tasks=scheduled_tasks, |
| message=f"Found {len(pending_tasks)} pending and {len(scheduled_tasks)} scheduled tasks.", |
| ) |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Failed to retrieve queue status: {e}") |
|
|
|
|
| @app.get( |
| "/tasks/{task_id}/status", |
| response_model=PipelineStatusResponse, |
| summary="Get the status of a pipeline run.", |
| ) |
| def get_pipeline_status(task_id: str): |
| redis_status_key = f"status_for_task:{task_id}" |
| try: |
| final_status_json = redis_client.get(redis_status_key) |
| if final_status_json: |
| try: |
| status_data = json.loads(final_status_json) |
| final_status_enum = PipelineStatus(status_data["status"]) |
| return PipelineStatusResponse( |
| task_id=task_id, |
| status=final_status_enum, |
| message=status_data.get("message"), |
| progress=1.0, |
| details=status_data.get("details"), |
| ) |
| except Exception: |
| pass |
|
|
| pipeline_progress = get_progress(task_id) |
| if pipeline_progress: |
| return PipelineStatusResponse( |
| task_id=task_id, |
| status=PipelineStatus.RUNNING, |
| message=pipeline_progress.get("message", "Task is running."), |
| progress=min(pipeline_progress.get("progress", 0.0), 1.0), |
| details={"source": "progress_tracker"}, |
| ) |
|
|
| pending_ids = {task.id for task in huey.pending()} if huey else set() |
| scheduled_ids = {task.id for task in huey.scheduled()} if huey else set() |
| if task_id in pending_ids or task_id in scheduled_ids: |
| return PipelineStatusResponse( |
| task_id=task_id, |
| status=PipelineStatus.PENDING, |
| message="Task is pending in the queue.", |
| progress=0.0, |
| details={"source": "huey_queue"}, |
| ) |
|
|
| if huey and huey.is_revoked(task_id): |
| return PipelineStatusResponse( |
| task_id=task_id, |
| status=PipelineStatus.REVOKED, |
| message="Task was revoked.", |
| progress=0.0, |
| details={"source": "huey_revoked"}, |
| ) |
|
|
| raise HTTPException(status_code=404, detail=f"Task with ID '{task_id}' not found.") |
| except HTTPException: |
| raise |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error retrieving status: {e}") |
|
|
|
|
| @app.get( |
| "/tasks/{task_id}/parameters", |
| response_model=TaskParametersResponse, |
| summary="Get the parameters a task was executed with.", |
| ) |
| def get_task_parameters(task_id: str): |
| redis_key = f"parameters_for_task:{task_id}" |
| try: |
| params_json = redis_client.get(redis_key) |
| if params_json is None: |
| raise HTTPException(status_code=404, detail=f"Parameters for task '{task_id}' not found.") |
| parameters = json.loads(params_json) |
| return TaskParametersResponse(task_id=task_id, parameters=parameters) |
| except HTTPException: |
| raise |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error retrieving parameters: {e}") |
|
|
|
|
| def _find_and_kill_run_augmentoolkit_process() -> Optional[Dict[str, Any]]: |
| try: |
| result = subprocess.run(["ps", "-ef"], capture_output=True, text=True, check=True) |
| for line in result.stdout.splitlines(): |
| if "run_augmentoolkit.py" in line and "grep" not in line: |
| fields = line.split() |
| if len(fields) >= 2: |
| try: |
| pid = int(fields[1]) |
| os.kill(pid, signal.SIGINT) |
| time.sleep(2) |
| return {"pid": pid, "killed": True, "method": "SIGINT"} |
| except Exception: |
| continue |
| return None |
| except Exception: |
| return None |
|
|
|
|
| @app.post( |
| "/tasks/{task_id}/interrupt", |
| status_code=200, |
| summary="Interrupt a running pipeline task subprocess or revoke a pending task.", |
| ) |
| def interrupt_or_revoke_task(task_id: str): |
| redis_pid_key = f"worker_pid_for_task:{task_id}" |
| redis_status_key = f"status_for_task:{task_id}" |
|
|
| try: |
| final_status_json = redis_client.get(redis_status_key) |
| if final_status_json: |
| try: |
| status_data = json.loads(final_status_json) |
| existing_status = status_data.get("status", "UNKNOWN").upper() |
| if existing_status in ["COMPLETED", "FAILED", "REVOKED"]: |
| raise HTTPException( |
| status_code=409, |
| detail=f"Task {task_id} has already finished with status: {existing_status}.", |
| ) |
| except HTTPException: |
| raise |
| except Exception: |
| pass |
|
|
| pid_bytes = redis_client.get(redis_pid_key) |
| if pid_bytes: |
| try: |
| pid = int(pid_bytes) |
| os.kill(pid, signal.SIGINT) |
| time.sleep(1) |
| set_final_status(task_id, "REVOKED", "Task interrupted via SIGINT.") |
| return {"message": f"Task {task_id} interrupted via SIGINT."} |
| except Exception: |
| pass |
|
|
| was_revoked = huey.revoke(task_id, revoke_once=True) if huey else False |
| if was_revoked: |
| set_final_status(task_id, "REVOKED", "Task revoked while pending.") |
| return {"message": f"Task {task_id} was pending and has been revoked."} |
|
|
| raise HTTPException(status_code=404, detail=f"Task {task_id} not running or not found.") |
| except HTTPException: |
| raise |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Error interrupting task: {e}") |
|
|
|
|
| @app.get("/tasks/{task_id}/logs", summary="Get logs for a specific task.") |
| def get_task_logs(task_id: str, tail: Optional[int] = Query(None, ge=1)): |
| log_file_path = LOGS_DIR / f"{task_id}.log" |
| if not log_file_path.exists(): |
| raise HTTPException(status_code=404, detail=f"Log file for task {task_id} not found.") |
|
|
| try: |
| with open(log_file_path, "r", encoding="utf-8") as f: |
| lines = f.readlines() |
| log_content = "".join(lines[-tail:]) if tail else "".join(lines) |
| return JSONResponse(content={"task_id": task_id, "logs": log_content}) |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Failed to read log file: {e}") |
|
|
|
|
| @app.delete("/logs", status_code=200, summary="Clear all task log files.") |
| def clear_all_logs(): |
| deleted_count = 0 |
| try: |
| for item in LOGS_DIR.iterdir(): |
| if item.is_file() and item.suffix == ".log": |
| item.unlink() |
| deleted_count += 1 |
| return {"message": f"Successfully deleted {deleted_count} log file(s)."} |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Failed to clear logs: {e}") |
|
|
|
|
| @app.get( |
| "/tasks/{task_id}/outputs/download", summary="Download output directory for a task." |
| ) |
| def download_task_output(task_id: str): |
| redis_key = f"output_dir_for_task:{task_id}" |
| output_dir_str = redis_client.get(redis_key) |
| if not output_dir_str: |
| raise HTTPException(status_code=404, detail=f"Output directory not found for task {task_id}.") |
|
|
| task_output_dir = PyPath(output_dir_str).resolve() |
| if not task_output_dir.is_dir(): |
| raise HTTPException(status_code=404, detail="Output directory does not exist.") |
|
|
| zip_filename = f"task_{task_id}_output.zip" |
| temp_zip_path = PyPath(tempfile.gettempdir()) / zip_filename |
| zip_directory(task_output_dir, temp_zip_path) |
| return FileResponse(temp_zip_path, media_type="application/zip", filename=zip_filename) |
|
|
|
|
| @app.get( |
| "/outputs/structure/{relative_path:path}", |
| response_model=List[FileStructure], |
| summary="Get structure of a path within the outputs directory.", |
| ) |
| def get_output_structure(relative_path: str = "."): |
| return handle_get_structure(OUTPUTS_DIR, relative_path) |
|
|
|
|
| @app.get( |
| "/inputs/structure/{relative_path:path}", |
| response_model=List[FileStructure], |
| summary="Get structure of a path within the inputs directory.", |
| ) |
| def get_input_structure(relative_path: str = "."): |
| return handle_get_structure(INPUTS_DIR, relative_path) |
|
|
|
|
| @app.get( |
| "/configs/structure/{relative_path:path}", |
| response_model=List[FileStructure], |
| summary="Get structure of a path within the configs directory.", |
| ) |
| def get_config_structure(relative_path: str = "."): |
| return handle_get_structure(CONFIGS_DIR, relative_path) |
|
|
|
|
| FRONTEND_DIST_DIR = PyPath("atk-interface/dist") |
|
|
| @app.get("/{full_path:path}", include_in_schema=False) |
| async def serve_frontend(full_path: str): |
| file_path = FRONTEND_DIST_DIR / full_path |
| if file_path.is_file(): |
| return FileResponse(file_path) |
|
|
| assets_path = FRONTEND_DIST_DIR / "assets" / full_path |
| if assets_path.is_file(): |
| return FileResponse(assets_path) |
|
|
| index_html = FRONTEND_DIST_DIR / "index.html" |
| if index_html.is_file(): |
| return FileResponse(index_html) |
|
|
| if not full_path or full_path == "/": |
| return {"status": "ok", "message": "Augmentoolkit API is running (frontend not built)."} |
|
|
| raise HTTPException(status_code=404, detail="Not found") |
|
|
