Add helper function to clean up artifacts

Author: VipulMascarenhas

ads/aqua/common/utils.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 """AQUA utils and constants."""
 
@@ -11,6 +11,7 @@
 import random
 import re
 import shlex
+import shutil
 import subprocess
 from datetime import datetime, timedelta
 from functools import wraps
@@ -21,6 +22,8 @@
 import fsspec
 import oci
 from cachetools import TTLCache, cached
+from huggingface_hub.constants import HF_HUB_CACHE
+from huggingface_hub.file_download import repo_folder_name
 from huggingface_hub.hf_api import HfApi, ModelInfo
 from huggingface_hub.utils import (
     GatedRepoError,
@@ -788,7 +791,9 @@ def get_ocid_substring(ocid: str, key_len: int) -> str:
     return ocid[-key_len:] if ocid and len(ocid) > key_len else ""
 
 
-def upload_folder(os_path: str, local_dir: str, model_name: str, exclude_pattern: str = None) -> str:
+def upload_folder(
+    os_path: str, local_dir: str, model_name: str, exclude_pattern: str = None
+) -> str:
     """Upload the local folder to the object storage
 
     Args:
@@ -818,6 +823,38 @@ def upload_folder(os_path: str, local_dir: str, model_name: str, exclude_pattern
     return f"oci://{os_details.bucket}@{os_details.namespace}" + "/" + object_path
 
 
+def cleanup_local_hf_model_artifact(
+    model_name: str,
+    local_dir: str = None,
+):
+    """
+    Helper function that deletes local artifacts downloaded from Hugging Face to free up disk space.
+    Parameters
+    ----------
+    model_name (str): Name of the huggingface model
+    local_dir (str): Local directory where the object is downloaded
+
+    """
+    if local_dir and os.path.exists(local_dir):
+        model_dir = os.path.join(local_dir, model_name)
+        if os.path.exists(model_dir):
+            shutil.rmtree(model_dir)
+            logger.debug(f"Deleted local model artifact directory: {model_dir}")
+
+        if not os.listdir(local_dir):
+            shutil.rmtree(local_dir)
+            logger.debug(f"Deleted local directory {model_dir} as it is empty.")
+
+    hf_local_path = os.path.join(
+        HF_HUB_CACHE, repo_folder_name(repo_id=model_name, repo_type="model")
+    )
+    if os.path.exists(hf_local_path):
+        shutil.rmtree(hf_local_path)
+        logger.debug(
+            f"Deleted local Hugging Face cache directory {hf_local_path} for the model {model_name} "
+        )
+
+
 def is_service_managed_container(container):
     return container and container.startswith(SERVICE_MANAGED_CONTAINER_URI_SCHEME)
 

ads/aqua/model/model.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import os
 import pathlib
@@ -23,6 +23,7 @@
 from ads.aqua.common.utils import (
     LifecycleStatus,
     _build_resource_identifier,
+    cleanup_local_hf_model_artifact,
     copy_model_config,
     create_word_icon,
     generate_tei_cmd_var,
@@ -1322,20 +1323,18 @@ def _download_model_from_hf(
         Returns
         -------
         model_artifact_path (str): Location where the model artifacts are downloaded.
-
         """
         # Download the model from hub
-        if not local_dir:
-            local_dir = os.path.join(os.path.expanduser("~"), "cached-model")
-        local_dir = os.path.join(local_dir, model_name)
-        os.makedirs(local_dir, exist_ok=True)
+        if local_dir:
+            local_dir = os.path.join(local_dir, model_name)
+            os.makedirs(local_dir, exist_ok=True)
         snapshot_download(
             repo_id=model_name,
             local_dir=local_dir,
             allow_patterns=allow_patterns,
             ignore_patterns=ignore_patterns,
         )
-        # Upload to object storage and skip .cache/huggingface/ folder
+        # Upload to object storage
         model_artifact_path = upload_folder(
             os_path=os_path,
             local_dir=local_dir,
@@ -1365,6 +1364,8 @@ def register(
                 ignore_patterns (list): Model files matching any of the patterns are not downloaded.
                     Example: ["*.json"] will ignore all .json files. ["folder/*"] will ignore all files under `folder`.
                     Patterns are Standard Wildcards (globbing patterns) and rules can be found here: https://docs.python.org/3/library/fnmatch.html
+                delete_from_local (bool): Deletes downloaded files from local machine after model is successfully
+                registered. Set to True by default.
 
         Returns:
             AquaModel:
@@ -1474,6 +1475,14 @@ def register(
             detail=validation_result.telemetry_model_name,
         )
 
+        if (
+            import_model_details.download_from_hf
+            and import_model_details.delete_from_local
+        ):
+            cleanup_local_hf_model_artifact(
+                model_name=model_name, local_dir=import_model_details.local_dir
+            )
+
         return AquaModel(**aqua_model_attributes)
 
     def _if_show(self, model: DataScienceModel) -> bool: