Spaces:

kota101
/

ta_ml

Sleeping

ta_ml / api /backup.py

muhammadjalaludin

Implement service layers for backup, model management, file handling, tokenization, and training

1424118 over 1 year ago

4.98 kB

# ta_ml/api/backup.py
"""
API endpoints for backup and cache management.
"""
import logging

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import FileResponse

from models.api import CacheFileListResponse, BackupUploadResponse
from services.backup_service import BackupService
from api.dependencies import get_backup_service

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/backup", tags=["Cache Backup"])


@router.get(
    "/list-cache",
    response_model=CacheFileListResponse,
    summary="List all files in HF cache",
    description="Returns a JSON array of all cached file paths relative to the root of the Hugging Face cache directory.",
    responses={
        200: {"description": "A list of all files in the cache."},
        404: {"description": "The Hugging Face cache directory was not found."},
    },
)
def list_cache_files(
    backup_service: BackupService = Depends(get_backup_service)
):
    try:
        files = backup_service.list_cache_files()
        return CacheFileListResponse(files=files)
    except FileNotFoundError as e:
        logger.warning(f"List cache failed: {e}")
        raise HTTPException(status_code=404, detail=str(e))


@router.get(
    "/download-log",
    summary="Download training logs",
    description="Serves the `training_logs.pkl` file, which contains a history of all structured logs from training runs, directly for download.",
    response_class=FileResponse,
    responses={
        200: {
            "description": "The log file is successfully streamed as a download.",
            "content": {"application/octet-stream": {}},
        },
        404: {"description": "The log file `training_logs.pkl` was not found in the cache."},
    },
)
def download_log(backup_service: BackupService = Depends(get_backup_service)):
    try:
        log_file_path = backup_service.get_log_file_for_download()
        return FileResponse(
            path=log_file_path,
            filename="training_logs.pkl",
            media_type="application/octet-stream",
        )
    except FileNotFoundError as e:
        logger.warning(f"Log file download failed: {e}")
        raise HTTPException(status_code=404, detail=str(e))


@router.get(
    "/download-cache",
    summary="Download HF cache as ZIP",
    description="Creates a ZIP archive of the entire Hugging Face cache directory on the fly and serves it for download. This can be useful for creating a full backup of all downloaded models and datasets. **Warning**: This may be a slow and resource-intensive operation if the cache is large.",
    response_class=FileResponse,
    responses={
        200: {
            "description": "The ZIP archive is successfully streamed as a download.",
            "content": {"application/zip": {}},
        },
        404: {"description": "The Hugging Face cache directory was not found."},
        500: {"description": "An error occurred during the ZIP archive creation."},
    },
)
def download_cache(backup_service: BackupService = Depends(get_backup_service)):
    try:
        archive_path = backup_service.get_cache_archive_for_download()
        return FileResponse(
            path=archive_path,
            filename="hf_cache_backup.zip",
            media_type="application/zip",
        )
    except FileNotFoundError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Cache archiving failed: {e}", exc_info=True)
        raise HTTPException(
            status_code=500, detail="Failed to create cache archive.")


@router.post(
    "/upload-cache",
    response_model=BackupUploadResponse,
    summary="Upload HF cache to Model repo",
    description=(
        "Uploads the entire contents of the local Hugging Face cache directory to a designated subfolder (`cache_backup/`) in a Hugging Face Model repository. This is useful for persisting a cache between different environments. "
        "**Requires environment variables**: `HF_TOKEN` (with write access) and `MODEL_REPO_ID` must be set."
    ),
    responses={
        200: {"description": "The cache was uploaded successfully."},
        400: {"description": "A required environment variable (`HF_TOKEN` or `MODEL_REPO_ID`) is missing, or another client-side error occurred."},
        500: {"description": "The upload process failed due to an API error or other issue. Check logs for details."},
    },
)
def upload_cache_to_hub(
    backup_service: BackupService = Depends(get_backup_service)
):
    try:
        repo_id = backup_service.upload_cache_to_hub()
        return BackupUploadResponse(
            status="Cache uploaded successfully",
            repo=repo_id
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Cache upload to Hub failed: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Upload failed: {e}")

	# ta_ml/api/backup.py
	"""
	API endpoints for backup and cache management.
	"""
	import logging

	from fastapi import APIRouter, Depends, HTTPException
	from fastapi.responses import FileResponse

	from models.api import CacheFileListResponse, BackupUploadResponse
	from services.backup_service import BackupService
	from api.dependencies import get_backup_service

	logger = logging.getLogger(__name__)
	router = APIRouter(prefix="/backup", tags=["Cache Backup"])


	@router.get(
	"/list-cache",
	response_model=CacheFileListResponse,
	summary="List all files in HF cache",
	description="Returns a JSON array of all cached file paths relative to the root of the Hugging Face cache directory.",
	responses={
	200: {"description": "A list of all files in the cache."},
	404: {"description": "The Hugging Face cache directory was not found."},
	},
	)
	def list_cache_files(
	backup_service: BackupService = Depends(get_backup_service)
	):
	try:
	files = backup_service.list_cache_files()
	return CacheFileListResponse(files=files)
	except FileNotFoundError as e:
	logger.warning(f"List cache failed: {e}")
	raise HTTPException(status_code=404, detail=str(e))


	@router.get(
	"/download-log",
	summary="Download training logs",
	description="Serves the `training_logs.pkl` file, which contains a history of all structured logs from training runs, directly for download.",
	response_class=FileResponse,
	responses={
	200: {
	"description": "The log file is successfully streamed as a download.",
	"content": {"application/octet-stream": {}},
	},
	404: {"description": "The log file `training_logs.pkl` was not found in the cache."},
	},
	)
	def download_log(backup_service: BackupService = Depends(get_backup_service)):
	try:
	log_file_path = backup_service.get_log_file_for_download()
	return FileResponse(
	path=log_file_path,
	filename="training_logs.pkl",
	media_type="application/octet-stream",
	)
	except FileNotFoundError as e:
	logger.warning(f"Log file download failed: {e}")
	raise HTTPException(status_code=404, detail=str(e))


	@router.get(
	"/download-cache",
	summary="Download HF cache as ZIP",
	description="Creates a ZIP archive of the entire Hugging Face cache directory on the fly and serves it for download. This can be useful for creating a full backup of all downloaded models and datasets. Warning: This may be a slow and resource-intensive operation if the cache is large.",
	response_class=FileResponse,
	responses={
	200: {
	"description": "The ZIP archive is successfully streamed as a download.",
	"content": {"application/zip": {}},
	},
	404: {"description": "The Hugging Face cache directory was not found."},
	500: {"description": "An error occurred during the ZIP archive creation."},
	},
	)
	def download_cache(backup_service: BackupService = Depends(get_backup_service)):
	try:
	archive_path = backup_service.get_cache_archive_for_download()
	return FileResponse(
	path=archive_path,
	filename="hf_cache_backup.zip",
	media_type="application/zip",
	)
	except FileNotFoundError as e:
	raise HTTPException(status_code=404, detail=str(e))
	except Exception as e:
	logger.error(f"Cache archiving failed: {e}", exc_info=True)
	raise HTTPException(
	status_code=500, detail="Failed to create cache archive.")


	@router.post(
	"/upload-cache",
	response_model=BackupUploadResponse,
	summary="Upload HF cache to Model repo",
	description=(
	"Uploads the entire contents of the local Hugging Face cache directory to a designated subfolder (`cache_backup/`) in a Hugging Face Model repository. This is useful for persisting a cache between different environments. "
	"Requires environment variables: `HF_TOKEN` (with write access) and `MODEL_REPO_ID` must be set."
	),
	responses={
	200: {"description": "The cache was uploaded successfully."},
	400: {"description": "A required environment variable (`HF_TOKEN` or `MODEL_REPO_ID`) is missing, or another client-side error occurred."},
	500: {"description": "The upload process failed due to an API error or other issue. Check logs for details."},
	},
	)
	def upload_cache_to_hub(
	backup_service: BackupService = Depends(get_backup_service)
	):
	try:
	repo_id = backup_service.upload_cache_to_hub()
	return BackupUploadResponse(
	status="Cache uploaded successfully",
	repo=repo_id
	)
	except ValueError as e:
	raise HTTPException(status_code=400, detail=str(e))
	except Exception as e:
	logger.error(f"Cache upload to Hub failed: {e}", exc_info=True)
	raise HTTPException(status_code=500, detail=f"Upload failed: {e}")