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): # Original startup logic 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")