CoolFace
Apppublic

uxoxo/eb2ab

sourceHugging Faceapache-2.0updated 11mo agoView on Hugging Face
0likes
storage.py267 linesDownload Raw Back to api
1"""2File storage and cleanup management for API-generated audio files.3"""4 5import os6import shutil7import threading8import time9from datetime import datetime, timedelta10from pathlib import Path11from typing import Optional12 13 14# Storage configuration15def get_api_output_dir() -> Path:16    """Get the output directory for API-generated files."""17    output_dir = os.environ.get("API_OUTPUT_DIR", "/app/audiobooks/api")18    path = Path(output_dir)19    path.mkdir(parents=True, exist_ok=True)20    return path21 22 23def get_retention_hours() -> int:24    """Get file retention period in hours."""25    return int(os.environ.get("AUDIO_RETENTION_HOURS", "24"))26 27 28def get_job_directory(job_id: str) -> Path:29    """30    Get directory for a specific job.31 32    Args:33        job_id: Unique job identifier34 35    Returns:36        Path to job directory37    """38    job_dir = get_api_output_dir() / job_id39    job_dir.mkdir(parents=True, exist_ok=True)40    return job_dir41 42 43def get_audio_file_path(job_id: str, output_format: str) -> Path:44    """45    Get expected path for generated audio file.46 47    Args:48        job_id: Unique job identifier49        output_format: Audio format extension50 51    Returns:52        Path to audio file53    """54    job_dir = get_job_directory(job_id)55    return job_dir / f"output.{output_format}"56 57 58def get_text_file_path(job_id: str) -> Path:59    """60    Get path for input text file.61 62    Args:63        job_id: Unique job identifier64 65    Returns:66        Path to text file67    """68    job_dir = get_job_directory(job_id)69    return job_dir / "input.txt"70 71 72def save_text_input(job_id: str, text: str) -> Path:73    """74    Save input text to file for TTS processing.75 76    Args:77        job_id: Unique job identifier78        text: Text content to save79 80    Returns:81        Path to saved text file82    """83    text_path = get_text_file_path(job_id)84    text_path.write_text(text, encoding="utf-8")85    return text_path86 87 88def get_job_metadata_path(job_id: str) -> Path:89    """90    Get path for job metadata file.91 92    Args:93        job_id: Unique job identifier94 95    Returns:96        Path to metadata file97    """98    job_dir = get_job_directory(job_id)99    return job_dir / "metadata.txt"100 101 102def save_job_metadata(job_id: str, created_at: datetime):103    """104    Save job metadata for cleanup tracking.105 106    Args:107        job_id: Unique job identifier108        created_at: Job creation timestamp109    """110    metadata_path = get_job_metadata_path(job_id)111    metadata_path.write_text(created_at.isoformat(), encoding="utf-8")112 113 114def get_job_creation_time(job_id: str) -> Optional[datetime]:115    """116    Get job creation time from metadata.117 118    Args:119        job_id: Unique job identifier120 121    Returns:122        Job creation datetime or None if not found123    """124    metadata_path = get_job_metadata_path(job_id)125    if not metadata_path.exists():126        return None127 128    try:129        iso_time = metadata_path.read_text(encoding="utf-8").strip()130        return datetime.fromisoformat(iso_time)131    except Exception:132        return None133 134 135def get_expiration_time(job_id: str) -> Optional[datetime]:136    """137    Get job expiration time.138 139    Args:140        job_id: Unique job identifier141 142    Returns:143        Expiration datetime or None if not found144    """145    created_at = get_job_creation_time(job_id)146    if not created_at:147        return None148 149    retention_hours = get_retention_hours()150    return created_at + timedelta(hours=retention_hours)151 152 153def delete_job_files(job_id: str) -> bool:154    """155    Delete all files for a job.156 157    Args:158        job_id: Unique job identifier159 160    Returns:161        True if deleted successfully, False otherwise162    """163    job_dir = get_job_directory(job_id)164 165    if not job_dir.exists():166        return False167 168    try:169        shutil.rmtree(job_dir)170        return True171    except Exception as e:172        print(f"Error deleting job {job_id}: {e}")173        return False174 175 176def cleanup_expired_jobs():177    """178    Cleanup expired job files.179    Should be called periodically by background thread.180    """181    output_dir = get_api_output_dir()182    current_time = datetime.now()183    retention_hours = get_retention_hours()184 185    deleted_count = 0186 187    # Iterate through all job directories188    for job_dir in output_dir.iterdir():189        if not job_dir.is_dir():190            continue191 192        job_id = job_dir.name193 194        # Get job creation time195        created_at = get_job_creation_time(job_id)196 197        # If no metadata, use directory modification time198        if not created_at:199            try:200                mtime = datetime.fromtimestamp(job_dir.stat().st_mtime)201                created_at = mtime202            except Exception:203                continue204 205        # Check if expired206        age_hours = (current_time - created_at).total_seconds() / 3600207 208        if age_hours > retention_hours:209            if delete_job_files(job_id):210                deleted_count += 1211                print(f"Cleaned up expired job: {job_id} (age: {age_hours:.1f}h)")212 213    if deleted_count > 0:214        print(f"Cleanup completed: {deleted_count} jobs deleted")215 216 217def start_cleanup_thread(interval_seconds: int = 3600):218    """219    Start background thread for periodic cleanup.220 221    Args:222        interval_seconds: Cleanup interval in seconds (default: 1 hour)223    """224    def cleanup_loop():225        while True:226            try:227                cleanup_expired_jobs()228            except Exception as e:229                print(f"Error in cleanup thread: {e}")230 231            time.sleep(interval_seconds)232 233    thread = threading.Thread(target=cleanup_loop, daemon=True, name="CleanupThread")234    thread.start()235    print(f"Cleanup thread started (interval: {interval_seconds}s)")236 237 238def get_storage_stats() -> dict:239    """240    Get storage statistics.241 242    Returns:243        Dictionary with storage stats244    """245    output_dir = get_api_output_dir()246 247    total_jobs = 0248    total_size = 0249 250    for job_dir in output_dir.iterdir():251        if not job_dir.is_dir():252            continue253 254        total_jobs += 1255 256        # Calculate directory size257        for file_path in job_dir.rglob("*"):258            if file_path.is_file():259                total_size += file_path.stat().st_size260 261    return {262        "total_jobs": total_jobs,263        "total_size_bytes": total_size,264        "total_size_mb": round(total_size / (1024 * 1024), 2),265        "output_directory": str(output_dir)266    }267