fix(ml): limit load retries (#10494)

This commit is contained in:
Mert 2024-06-20 14:13:18 -04:00 committed by GitHub
parent 79a8ab71ef
commit a42af06889
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 26 additions and 11 deletions

View file

@ -31,6 +31,7 @@ class InferenceModel(ABC):
**model_kwargs: Any,
) -> None:
self.loaded = False
self.load_attempts = 0
self.model_name = clean_name(model_name)
self.cache_dir = Path(cache_dir) if cache_dir is not None else self.cache_dir_default
self.providers = providers if providers is not None else self.providers_default
@ -48,9 +49,11 @@ class InferenceModel(ABC):
def load(self) -> None:
if self.loaded:
return
self.load_attempts += 1
self.download()
log.info(f"Loading {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
attempt = f"Attempt #{self.load_attempts + 1} to load" if self.load_attempts else "Loading"
log.info(f"{attempt} {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
self.session = self._load()
self.loaded = True