81 lines
3 KiB
Python
81 lines
3 KiB
Python
"""File cleanup operations for the quantisation workflow.
|
|
|
|
Manages removal of temporary files, model cleanup after processing,
|
|
and disk space recovery during quantisation operations.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from shutil import rmtree as shutil_rmtree
|
|
from typing import TYPE_CHECKING
|
|
|
|
from helpers.logger import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
from helpers.models.quantisation import ModelSource
|
|
|
|
|
|
class FileCleanup:
|
|
"""Handles cleanup of temporary and intermediate files.
|
|
|
|
Provides methods for removing processed model files, temporary
|
|
conversions, and other artifacts to manage disk space efficiently
|
|
during quantisation workflows.
|
|
"""
|
|
|
|
@staticmethod
|
|
def cleanup_files(f16_model_path: Path, model_source: ModelSource, models_dir: Path) -> None:
|
|
"""Clean up temporary files after processing.
|
|
|
|
Removes F16 model and original format files to save disk space
|
|
after successful quantisation and upload. Processes both F16
|
|
GGUF files and original model formats to maximise storage recovery.
|
|
"""
|
|
if f16_model_path.exists():
|
|
logger.info(f"Removing F16 model {f16_model_path.name} to save disk space...")
|
|
f16_model_path.unlink()
|
|
|
|
if not model_source.is_gguf_repo:
|
|
FileCleanup.cleanup_original_model(model_source, models_dir)
|
|
|
|
@staticmethod
|
|
def cleanup_original_model(model_source: ModelSource, models_dir: Path) -> None:
|
|
"""Clean up original model files after successful conversion.
|
|
|
|
Removes SafeTensors files to save disk space whilst preserving
|
|
configuration, tokeniser, and metadata files for reference. The
|
|
design prioritises space efficiency over re-conversion capability.
|
|
"""
|
|
model_dir = models_dir / model_source.model_name
|
|
|
|
safetensor_files = list(model_dir.glob("*.safetensors"))
|
|
if safetensor_files:
|
|
logger.info(f"Removing {len(safetensor_files)} SafeTensors files...")
|
|
for file in safetensor_files:
|
|
file.unlink()
|
|
|
|
logger.info("Keeping config files, tokeniser, and metadata for reference")
|
|
|
|
@staticmethod
|
|
def cleanup_quantisation_file(file_path: Path) -> None:
|
|
"""Remove a single quantisation file.
|
|
|
|
Safely removes the specified file with existence checking and
|
|
logging for disk space management during quantisation workflows.
|
|
"""
|
|
if file_path.exists():
|
|
logger.info(f"Removing {file_path.name} to save disk space...")
|
|
file_path.unlink()
|
|
|
|
@staticmethod
|
|
def cleanup_temp_directory(temp_dir: Path) -> None:
|
|
"""Clean up a temporary directory and all its contents.
|
|
|
|
Recursively removes the directory and all subdirectories with
|
|
error tolerance to handle locked or missing files gracefully.
|
|
"""
|
|
if temp_dir.exists() and temp_dir.is_dir():
|
|
logger.debug(f"Cleaning up temporary directory: {temp_dir}")
|
|
shutil_rmtree(temp_dir, ignore_errors=True)
|