58
58
)
59
59
from ads .aqua .data import AquaResourceIdentifier
60
60
from ads .common .auth import AuthState , default_signer
61
+ from ads .common .decorator .threaded import threaded
61
62
from ads .common .extended_enum import ExtendedEnumMeta
62
63
from ads .common .object_storage_details import ObjectStorageDetails
63
64
from ads .common .oci_resource import SEARCH_TYPE , OCIResource
@@ -225,6 +226,7 @@ def read_file(file_path: str, **kwargs) -> str:
225
226
return UNKNOWN
226
227
227
228
229
+ @threaded ()
228
230
def load_config (file_path : str , config_file_name : str , ** kwargs ) -> dict :
229
231
artifact_path = f"{ file_path .rstrip ('/' )} /{ config_file_name } "
230
232
signer = default_signer () if artifact_path .startswith ("oci://" ) else {}
@@ -536,14 +538,14 @@ def _build_job_identifier(
536
538
return AquaResourceIdentifier ()
537
539
538
540
539
- def container_config_path ():
541
+ def service_config_path ():
540
542
return f"oci://{ AQUA_SERVICE_MODELS_BUCKET } @{ CONDA_BUCKET_NS } /service_models/config"
541
543
542
544
543
545
@cached (cache = TTLCache (maxsize = 1 , ttl = timedelta (hours = 5 ), timer = datetime .now ))
544
546
def get_container_config ():
545
547
config = load_config (
546
- file_path = container_config_path (),
548
+ file_path = service_config_path (),
547
549
config_file_name = CONTAINER_INDEX ,
548
550
)
549
551
@@ -568,7 +570,7 @@ def get_container_image(
568
570
"""
569
571
570
572
config = config_file_name or get_container_config ()
571
- config_file_name = container_config_path ()
573
+ config_file_name = service_config_path ()
572
574
573
575
if container_type not in config :
574
576
raise AquaValueError (
@@ -1062,3 +1064,18 @@ def get_hf_model_info(repo_id: str) -> ModelInfo:
1062
1064
return HfApi ().model_info (repo_id = repo_id )
1063
1065
except HfHubHTTPError as err :
1064
1066
raise format_hf_custom_error_message (err ) from err
1067
+
1068
+
1069
+ @cached (cache = TTLCache (maxsize = 1 , ttl = timedelta (hours = 5 ), timer = datetime .now ))
1070
+ def list_hf_models (query : str ) -> List [str ]:
1071
+ try :
1072
+ models = HfApi ().list_models (
1073
+ model_name = query ,
1074
+ task = "text-generation" ,
1075
+ sort = "downloads" ,
1076
+ direction = - 1 ,
1077
+ limit = 20 ,
1078
+ )
1079
+ return [model .id for model in models if model .disabled is None ]
1080
+ except HfHubHTTPError as err :
1081
+ raise format_hf_custom_error_message (err ) from err
0 commit comments