1+
2+ import os
3+ from pathlib import Path
4+ from huggingface_hub import HfApi
5+ from huggingface_hub .utils import RepositoryNotFoundError
6+ from datetime import datetime
7+ import time
8+ import shutil
9+ from fastapi import HTTPException
10+
11+ from inferadmin .config .loader import config_manager
12+
13+ def check_hf_model_exists (model_path ):
14+ """
15+ Check if a complete Hugging Face model exists at the given path.
16+
17+ Returns:
18+ bool: True if valid model, False otherwise
19+ """
20+ path = Path (model_path )
21+
22+ # Check if the directory exists
23+ if not path .exists () or not path .is_dir ():
24+ return False
25+ # Check if the directory is empty
26+ if not any (path .iterdir ()):
27+ return False
28+
29+ # Required files to check
30+ required_files = [
31+ "*.safetensors" , # At least one safetensors file
32+ "generation_config.json" ,
33+ "config.json" ,
34+ "tokenizer.json" ,
35+ "special_tokens_map.json" ,
36+ "tokenizer_config.json" ,
37+ "vocab.json"
38+ ]
39+
40+ # Check each required file
41+ for file_pattern in required_files :
42+ if "*" in file_pattern :
43+ # Handle wildcard pattern (for safetensors)
44+ if not list (path .glob (file_pattern )):
45+ return False
46+ else :
47+ # Handle exact filename
48+ if not (path / file_pattern ).exists ():
49+ return False
50+
51+ return True
52+
53+ def get_folder_size_gb (folder_path ):
54+ """
55+ Calculate the total size of all files in a folder and its subfolders in GB.
56+
57+ Args:
58+ folder_path (str): Path to the folder
59+
60+ Returns:
61+ float: Size in gigabytes
62+ """
63+ path = Path (folder_path )
64+
65+ if not path .exists () or not path .is_dir ():
66+ return 0.0
67+
68+ total_size = 0
69+
70+ # Walk through all files in the directory and subdirectories
71+ for dirpath , dirnames , filenames in os .walk (path ):
72+ for filename in filenames :
73+ file_path = os .path .join (dirpath , filename )
74+ # Skip if it's a symlink
75+ if not os .path .islink (file_path ):
76+ total_size += os .path .getsize (file_path )
77+
78+ # Convert bytes to GB
79+ size_gb = total_size / (1024 ** 3 )
80+
81+ return round (size_gb , 2 )
82+
83+ def get_most_recent_modified_date (folder_path ):
84+ """
85+ Get the most recent modification date of any file in the directory.
86+
87+ Args:
88+ folder_path (str): Path to the folder
89+
90+ Returns:
91+ datetime
92+ """
93+ most_recent_time = 0
94+
95+ # Walk through all files in the directory and subdirectories
96+ for dirpath , dirnames , filenames in os .walk (folder_path ):
97+ for filename in filenames :
98+ file_path = os .path .join (dirpath , filename )
99+ if os .path .isfile (file_path ):
100+ file_mtime = os .path .getmtime (file_path )
101+ if file_mtime > most_recent_time :
102+ most_recent_time = file_mtime
103+
104+ # If no files were found, return current time
105+ if most_recent_time == 0 :
106+ most_recent_time = time .time ()
107+
108+ # Convert to datetime
109+ modified_date = datetime .fromtimestamp (most_recent_time )
110+
111+ return modified_date
112+
113+ def scan_hf_models_directory ():
114+ """
115+ Scan a directory for Hugging Face models and collect info about each valid model.
116+
117+ Returns:
118+ list: List of dictionaries with model info (repo_id, path, size_gb, last_updated)
119+ """
120+ base_path = Path (config_manager .get_config ().model_storage_path )
121+
122+ if not base_path .exists () or not base_path .is_dir ():
123+ return []
124+
125+ model_info_list = []
126+
127+ # Get all subdirectories in the base path
128+ subdirs = [d for d in base_path .iterdir () if d .is_dir ()]
129+
130+ for folder in subdirs :
131+ # Get folder name
132+ folder_name = folder .name
133+ folder_path = str (folder )
134+
135+ # Check if it's a valid HF model
136+ is_valid_model = check_hf_model_exists (folder_path )
137+
138+ # Get folder size
139+ size_gb = get_folder_size_gb (folder_path )
140+
141+ # Get the most recent modification date of any file in the directory
142+ last_updated = get_most_recent_modified_date (folder_path )
143+
144+ # Create model info dictionary
145+
146+ model_info = {
147+ "repo_id" : folder_name .replace ('_' , '/' ),
148+ "path" : folder_path ,
149+ "size_gb" : size_gb ,
150+ "last_updated" : last_updated
151+ }
152+
153+ if is_valid_model :
154+ model_info_list .append (model_info )
155+
156+ return model_info_list
157+
158+ def delete_model (repo_id :str ):
159+ """
160+ Delete a model from the storage path and all its contents.
161+
162+ Args:
163+ model_name (str): hf_model name
164+ """
165+ folder_name = repo_id .replace ("/" , "_" )
166+ folder_path = f"{ config_manager .get_config ().model_storage_path } { folder_name }
167+
168+ path = Path (folder_path )
169+
170+ # Check if the path exists and is a directory
171+ if not path .exists ():
172+ raise HTTPException (
173+ status_code = 404 ,
174+ detail = f"Path does not exist: { folder_path }
175+ )
176+
177+ if not path .is_dir ():
178+ raise HTTPException (
179+ status_code = 404 ,
180+ detail = f"Path is not a directory: { folder_path }
181+ )
182+
183+ try :
184+ # Remove the directory and all its contents
185+ shutil .rmtree (folder_path )
186+ except Exception as e :
187+ raise HTTPException (
188+ status_code = 500 ,
189+ detail = f"Error deleting model: { str (e )}
190+ )
191+
192+ def download_hf_model (repo_id :str ):
193+ """
194+ Downloads the given model to the local directory from huggingface.
195+
196+ Args:
197+ repo_id (str): the HF repo name of thing to download
198+ """
199+ hf = HfApi (
200+ token = config_manager .get_config ().hf_token
201+ )
202+
203+ base_path = config_manager .get_config ().model_storage_path
204+ file_name = repo_id .replace ('/' , '_' )
205+ target_path = f"{ base_path } { file_name }
206+
207+ try :
208+ hf .snapshot_download (
209+ repo_id = repo_id ,
210+ repo_type = 'model' ,
211+ local_dir = target_path
212+ )
213+
214+ except RepositoryNotFoundError as e :
215+ raise HTTPException (
216+ status_code = 404 ,
217+ detail = f"HuggingFace Repository not found: { repo_id }
218+ )
219+
220+ except Exception as e :
221+ raise HTTPException (
222+ status_code = 500 ,
223+ detail = f"Error downloading model: { str (e )}
224+ )
0 commit comments