Improve typing and reuse variable for timeout

zencoder/huggingface/hf_deployment.py

@@ -3,28 +3,34 @@
 from typing import Tuple, Optional
 from zenml.services import ServiceType, ServiceState
 from zenml.services.service import BaseDeploymentService
-from huggingface_hub import InferenceClient, InferenceEndpointError
+from huggingface_hub import (
+    InferenceClient,
+    InferenceEndpointError,
+    InferenceEndpoint,
+)
 
 logger = get_logger(__name__)
 
 
 class HuggingFaceModelService(BaseDeploymentService):
+    """HuggingFace model deployment service."""
+
     SERVICE_TYPE = ServiceType(
         name="hf-endpoint-deployment",
         type="model-serving",
         flavor="hfendpoint",
         description="Huggingface inference endpoint service",
     )
 
-    def __init__(self, endpoint):
-        """_summary_
+    def __init__(self, endpoint: InferenceEndpoint):
+        """_summary_.
 
         Args:
-            endpoint (_type_): _description_
+            endpoint (InferenceEndpoint): _description_
         """
         self.endpoint = endpoint
 
-    def wait_for_startup(self, timeout=300) -> bool:
+    def wait_for_startup(self, timeout: int = 300) -> bool:
         """_summary_
 
         Args:

zencoder/huggingface/hf_model_deployer.py

@@ -9,24 +9,32 @@
 from zenml.services import BaseService, ServiceConfig
 from huggingface_hub import create_inference_endpoint
 from huggingface.hf_deployment import HuggingFaceModelService
-from zenml.model_deployers.base_model_deployer import BaseModelDeployerFlavor
+from zenml.model_deployers.base_model_deployer import (
+    DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    BaseModelDeployerFlavor,
+)
 
 logger = get_logger(__name__)
 
 
 class HFEndpointModelDeployer(BaseModelDeployer):
+    """Huggingface endpoint model deployer."""
+
     NAME: ClassVar[str] = "HFEndpoint"
     FLAVOR: ClassVar[Type[BaseModelDeployerFlavor]] = HFModelDeployerFlavor
 
     def deploy_model(
-        self, config: ServiceConfig, replace: bool = False, timeout: int = 300
+        self,
+        config: ServiceConfig,
+        replace: bool = False,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
     ) -> BaseService:
         """_summary_.
 
         Args:
             config (ServiceConfig): _description_
             replace (bool, optional): _description_. Defaults to False.
-            timeout (int, optional): _description_. Defaults to 300.
+            timeout (int, optional): _description_. Defaults to DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT.
 
         Raises:
             ValueError: _description_

zencoder/huggingface/hf_model_deployer_flavor.py

@@ -22,7 +22,7 @@ class HFInferenceEndpointConfig(BaseModelDeployerConfig):
 
 
 class HFModelDeployerFlavor(BaseModelDeployerFlavor):
-    """Seldon Core model deployer flavor."""
+    """Huggingface Endpoint model deployer flavor."""
 
     @property
     def name(self) -> str:
@@ -35,7 +35,7 @@ def name(self) -> str:
 
     @property
     def config_class(self):
-        """Returns `SeldonModelDeployerConfig` config class.
+        """Returns `HFInferenceEndpointConfig` config class.
 
         Returns:
                 The config class.