clarifai 11.8.0__py3-none-any.whl → 11.8.2__py3-none-any.whl

clarifai/client/search.py

@@ -91,14 +91,21 @@ class Search(Lister, BaseClient):
         )
         Lister.__init__(self, page_size=1000)
 
-    def _get_annot_proto(self, **kwargs):
+    def _get_annot_proto(self, **kwargs) -> resources_pb2.Annotation:
         """Get an Annotation proto message based on keyword arguments.
 
         Args:
-            **kwargs: Keyword arguments specifying the resource.
+            **kwargs: Keyword arguments specifying the annotation data.
+                     Supported keys:
+                     - image_bytes (bytes): Raw image bytes
+                     - image_url (str): URL to an image
+                     - concepts (List[Dict]): List of concept dictionaries
+                     - metadata (Dict): Metadata dictionary
+                     - geo_longitude (float): Geographic longitude
+                     - geo_latitude (float): Geographic latitude
 
         Returns:
-            resources_pb2.Annotation: An Annotation proto message.
+            resources_pb2.Annotation: An Annotation proto message with the specified data.
         """
         if not kwargs:
             return resources_pb2.Annotation()
@@ -139,14 +146,24 @@ class Search(Lister, BaseClient):
                 raise UserError(f"kwargs contain key that is not supported: {key}")
         return resources_pb2.Annotation(data=self.data_proto)
 
-    def _get_input_proto(self, **kwargs):
+    def _get_input_proto(self, **kwargs) -> resources_pb2.Input:
         """Get an Input proto message based on keyword arguments.
 
         Args:
-            **kwargs: Keyword arguments specifying the resource.
+            **kwargs: Keyword arguments specifying the input data.
+                     Supported keys:
+                     - input_types (List[str]): List of input types ('image', 'text', 'audio', 'video')
+                     - dataset_ids (List[str]): List of dataset IDs to filter by
+                     - image_bytes (bytes): Raw image bytes
+                     - image_url (str): URL to an image
+                     - text_raw (str): Raw text content
+                     - concepts (List[Dict]): List of concept dictionaries
+                     - metadata (Dict): Metadata dictionary
+                     - geo_longitude (float): Geographic longitude
+                     - geo_latitude (float): Geographic latitude
 
         Returns:
-            resources_pb2.Input: An Input proto message.
+            resources_pb2.Input: An Input proto message with the specified data.
         """
         if not kwargs:
             return resources_pb2.Input()
@@ -194,15 +211,22 @@ class Search(Lister, BaseClient):
     def _list_topk_generator(
         self, endpoint: Callable[..., Any], proto_message: Any, request_data: Dict[str, Any]
     ) -> Generator[Dict[str, Any], None, None]:
-        """Lists all pages of a resource.
+        """Lists top-k results with pagination support.
+
+        This method handles pagination for search results when top_k is specified,
+        automatically calculating the required number of pages and per-page limits.
 
         Args:
-            endpoint (Callable): The endpoint to call.
-            proto_message (Any): The proto message to use.
-            request_data (dict): The request data to use.
+            endpoint (Callable[..., Any]): The gRPC endpoint method to call for search.
+            proto_message (Any): The protobuf message class for the request.
+            request_data (Dict[str, Any]): The base request data dictionary.
 
         Yields:
-            response_dict: The next item in the listing.
+            Dict[str, Any]: Individual search result items from the API response.
+
+        Raises:
+            UserError: If pagination limits are exceeded or top_k is too large.
+            Exception: If the API request fails.
         """
         max_pages = ceil(self.top_k / self.default_page_size)
         total_hits = 0

clarifai/client/user.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Dict, Generator, List
+from typing import Any, Dict, Generator, List, Optional
 
 import yaml
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -53,14 +53,17 @@ class User(Lister, BaseClient):
         Lister.__init__(self)
 
     def list_apps(
-        self, filter_by: Dict[str, Any] = {}, page_no: int = None, per_page: int = None
+        self,
+        filter_by: Dict[str, Any] = {},
+        page_no: Optional[int] = None,
+        per_page: Optional[int] = None,
     ) -> Generator[App, None, None]:
         """Lists all the apps for the user.
 
         Args:
-            filter_by (dict): A dictionary of filters to be applied to the list of apps.
-            page_no (int): The page number to list.
-            per_page (int): The number of items per page.
+            filter_by (Dict[str, Any]): A dictionary of filters to be applied to the list of apps.
+            page_no (Optional[int]): The page number to list. If None, lists all pages.
+            per_page (Optional[int]): The number of items per page. If None, uses default.
 
         Yields:
             App: App objects for the user.
@@ -533,6 +536,178 @@ class User(Lister, BaseClient):
         ]
         return f"Clarifai User Details: \n{', '.join(attribute_strings)}\n"
 
+    def get_secret(self, secret_id: str) -> dict:
+        """Returns a secret object if exists.
+
+        Args:
+            secret_id (str): The secret ID to interact with
+
+        Returns:
+            Dict: A dictionary containing information about the existing secret ID.
+
+        Example:
+            >>> from clarifai.client.user import User
+            >>> client = User(user_id="user_id")
+            >>> secret_info = client.get_secret(secret_id="secret_id")
+        """
+        request = service_pb2.GetSecretRequest(user_app_id=self.user_app_id, id=secret_id)
+        response = self._grpc_request(self.STUB.GetSecret, request)
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(
+                f"""Error getting secret, are you sure this is a valid secret id {secret_id} at the user_id
+          {self.user_app_id.user_id}.
+          Error: {response.status.description}"""
+            )
+
+        dict_response = MessageToDict(response, preserving_proto_field_name=True)
+        kwargs = self.process_response_keys(dict_response["secret"], "secret")
+
+        return dict(auth=self.auth_helper, **kwargs)
+
+    def list_secrets(
+        self, page_no: int = None, per_page: int = None
+    ) -> Generator[dict, None, None]:
+        """List all secrets for the user
+
+        Args:
+            page_no (int): The page number to list.
+            per_page (int): The number of items per page.
+
+        Yields:
+            Dict: Dictionaries containing information about the secrets.
+
+        Example:
+            >>> from clarifai.client.user import User
+            >>> client = User(user_id="user_id")
+            >>> all_secrets = list(client.list_secrets())
+
+        Note:
+            Defaults to 16 per page if page_no is specified and per_page is not specified.
+            If both page_no and per_page are None, then lists all the resources.
+        """
+        request_data = dict(user_app_id=self.user_app_id)
+        all_secrets_info = self.list_pages_generator(
+            self.STUB.ListSecrets,
+            service_pb2.ListSecretsRequest,
+            request_data,
+            per_page=per_page,
+            page_no=page_no,
+        )
+        for secret_info in all_secrets_info:
+            yield dict(auth=self.auth_helper, **secret_info)
+
+    def create_secrets(self, secrets: List[Dict[str, Any]]) -> List[dict]:
+        """Creates secrets for the user.
+
+        Args:
+            secrets (List[Dict[str, Any]]): List of secret configurations to create.
+                Each secret dict can contain:
+                - id (str): The name/ID of the secret (required)
+                - value (str): The secret value (required)
+                - description (str): Optional description of the secret
+                - expires_at (str): Optional expiration timestamp
+
+        Returns:
+            List[Dict]: List of dictionaries containing information about the created secrets.
+
+        Example:
+            >>> from clarifai.client.user import User
+            >>> client = User(user_id="user_id")
+            >>> secrets = [{"id": "secret1", "value": "secret_value", "description": "My Secret"}]
+            >>> created_secrets = client.create_secrets(secrets)
+        """
+        assert isinstance(secrets, list), "secrets param should be a list"
+
+        # Convert dict secrets to protobuf Secret objects
+        secret_objects = []
+        for secret_config in secrets:
+            secret_objects.append(resources_pb2.Secret(**secret_config))
+
+        request = service_pb2.PostSecretsRequest(
+            user_app_id=self.user_app_id, secrets=secret_objects
+        )
+        response = self._grpc_request(self.STUB.PostSecrets, request)
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+
+        self.logger.info(f"Secrets created successfully:\n{response.status}")
+
+        # Convert response to list of dictionaries
+        dict_response = MessageToDict(response, preserving_proto_field_name=True)
+        created_secrets = []
+        for secret in dict_response.get("secrets", []):
+            kwargs = self.process_response_keys(secret, "secret")
+            created_secrets.append(dict(auth=self.auth_helper, **kwargs))
+
+        return created_secrets
+
+    def patch_secrets(
+        self, secrets: List[Dict[str, Any]], action: str = 'overwrite'
+    ) -> List[dict]:
+        """Patches secrets for the user.
+
+        Args:
+            secrets (List[Dict[str, Any]]): List of secret configurations to patch.
+                Each secret dict should contain:
+                - id (str): The name/ID of the secret to patch (required)
+                - value (str): Optional new secret value
+                - description (str): Optional new description
+                - expires_at (str): Optional new expiration timestamp
+            action (str): The action to perform on the secrets (overwrite/remove).
+
+        Returns:
+            List[Dict]: List of dictionaries containing information about the patched secrets.
+
+        Example:
+            >>> from clarifai.client.user import User
+            >>> client = User(user_id="user_id")
+            >>> secrets = [{"id": "secret1", "description": "Updated Secret Description"}]
+            >>> patched_secrets = client.patch_secrets(secrets, action="overwrite")
+        """
+        assert isinstance(secrets, list), "secrets param should be a list"
+
+        # Convert dict secrets to protobuf Secret objects
+        secret_objects = []
+        for secret_config in secrets:
+            secret_objects.append(resources_pb2.Secret(**secret_config))
+
+        request = service_pb2.PatchSecretsRequest(
+            user_app_id=self.user_app_id, secret=secret_objects, action=action
+        )
+        response = self._grpc_request(self.STUB.PatchSecrets, request)
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+
+        self.logger.info(f"Secrets patched successfully:\n{response.status}")
+
+        # Convert response to list of dictionaries
+        dict_response = MessageToDict(response, preserving_proto_field_name=True)
+        patched_secrets = []
+        for secret in dict_response.get("secrets", []):
+            kwargs = self.process_response_keys(secret, "secret")
+            patched_secrets.append(dict(auth=self.auth_helper, **kwargs))
+
+        return patched_secrets
+
+    def delete_secrets(self, secret_ids: List[str]) -> None:
+        """Deletes a list of secrets for the user.
+
+        Args:
+            secret_ids (List[str]): The secret IDs of the user to delete.
+
+        Example:
+            >>> from clarifai.client.user import User
+            >>> client = User(user_id="user_id")
+            >>> client.delete_secrets(secret_ids=["secret_id1", "secret_id2"])
+        """
+        assert isinstance(secret_ids, list), "secret_ids param should be a list"
+
+        request = service_pb2.DeleteSecretsRequest(user_app_id=self.user_app_id, ids=secret_ids)
+        response = self._grpc_request(self.STUB.DeleteSecrets, request)
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+        self.logger.info("\nSecrets Deleted\n%s", response.status)
+
     def list_models(
         self,
         user_id: str = None,

clarifai/client/workflow.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import Dict, Generator, List
+from typing import Any, Dict, Generator, List, Optional
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.resources_pb2 import Input
@@ -25,29 +25,30 @@ class Workflow(Lister, BaseClient):
 
     def __init__(
         self,
-        url: str = None,
-        workflow_id: str = None,
-        workflow_version: Dict = {'id': ""},
-        output_config: Dict = {'min_value': 0},
+        url: Optional[str] = None,
+        workflow_id: Optional[str] = None,
+        workflow_version: Dict[str, str] = {'id': ""},
+        output_config: Dict[str, Any] = {'min_value': 0},
         base_url: str = DEFAULT_BASE,
-        pat: str = None,
-        token: str = None,
-        root_certificates_path: str = None,
+        pat: Optional[str] = None,
+        token: Optional[str] = None,
+        root_certificates_path: Optional[str] = None,
         **kwargs,
     ):
         """Initializes a Workflow object.
 
         Args:
-            url (str): The URL to initialize the workflow object.
-            workflow_id (str): The Workflow ID to interact with.
-            workflow_version (dict): The Workflow Version to interact with.
-            output_config (dict): The output config to interact with.
-              min_value (float): The minimum value of the prediction confidence to filter.
-              max_concepts (int): The maximum number of concepts to return.
-              select_concepts (list[Concept]): The concepts to select.
-              sample_ms (int): The number of milliseconds to sample.
+            url (Optional[str]): The URL to initialize the workflow object.
+            workflow_id (Optional[str]): The Workflow ID to interact with.
+            workflow_version (Dict[str, str]): The Workflow Version to interact with.
+                                              Defaults to {'id': ""} for latest version.
+            output_config (Dict[str, Any]): The output config to interact with.
+              - min_value (float): The minimum value of the prediction confidence to filter.
+              - max_concepts (int): The maximum number of concepts to return.
+              - select_concepts (List[Concept]): The concepts to select.
+              - sample_ms (int): The number of milliseconds to sample.
             base_url (str): Base API url. Default "https://api.clarifai.com"
-            pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
+            pat (Optional[str]): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
             token (str): A session token for authentication. Accepts either a session token or a pat. Can be set as env var CLARIFAI_SESSION_TOKEN
             root_certificates_path (str): Path to the SSL root certificates file, used to establish secure gRPC connections.
             **kwargs: Additional keyword arguments to be passed to the Workflow.

clarifai/runners/models/model_builder.py

@@ -102,6 +102,7 @@ class ModelBuilder:
         self.folder = self._validate_folder(folder)
         self.config = self._load_config(os.path.join(self.folder, 'config.yaml'))
         self._validate_config()
+        self._validate_config_secrets()
         self._validate_stream_options()
         self.model_proto = self._get_model_proto()
         self.model_id = self.model_proto.id
@@ -465,6 +466,115 @@ class ModelBuilder:
                     "2) set_output_context"
                 )
 
+    def _validate_config_secrets(self):
+        """
+        Validate the secrets section in the config file.
+        """
+        if "secrets" not in self.config:
+            return
+
+        secrets = self.config.get("secrets", [])
+        if not isinstance(secrets, list):
+            raise ValueError("The 'secrets' field must be an array.")
+
+        for i, secret in enumerate(secrets):
+            if not isinstance(secret, dict):
+                raise ValueError(f"Secret at index {i} must be a dictionary.")
+
+            # Validate required fields
+            if "id" not in secret or not secret["id"]:
+                raise ValueError(f"Secret at index {i} must have a non-empty 'id' field.")
+
+            if "type" not in secret or not secret["type"]:
+                secret["type"] = "env"
+
+            if "env_var" not in secret or not secret["env_var"]:
+                raise ValueError(f"Secret at index {i} must have a non-empty 'env_var' field.")
+            # Validate secret type
+            if secret["type"] != "env":
+                raise ValueError(
+                    f"Secret at index {i} has invalid type '{secret['type']}'. Must be 'env'."
+                )
+
+        logger.info(f"Validated {len(secrets)} secrets in config file.")
+
+    def _process_secrets(self):
+        """
+        Process secrets from config file and create/validate them using the User client.
+        Returns the processed secrets array for inclusion in ModelVersion.OutputInfo.Params.
+        """
+        if "secrets" not in self.config:
+            return []
+
+        secrets = self.config.get("secrets", [])
+        if not secrets:
+            return []
+
+        # Get user client for secret operations
+        user = User(
+            user_id=self.config.get('model').get('user_id'),
+            pat=self.client.pat,
+            token=self.client.token,
+            base_url=self.client.base,
+        )
+
+        processed_secrets = []
+        secrets_to_create = []
+
+        for secret in secrets:
+            secret_id = secret["id"]
+            secret_type = secret.get("type", "env")
+            env_var = secret["env_var"]
+            secret_value = secret.get("value")  # Optional for existing secrets
+
+            # Check if secret already exists
+            try:
+                existing_secret = user.get_secret(secret_id)
+                logger.info(f"Secret '{secret_id}' already exists, using existing secret.")
+
+                # Add to processed secrets without the value
+                processed_secret = {
+                    "id": secret_id,
+                    "type": secret_type,
+                    "env_var": env_var,
+                }
+                processed_secrets.append(processed_secret)
+
+            except Exception:
+                # Secret doesn't exist, need to create it
+                if secret_value:
+                    logger.info(f"Secret '{secret_id}' does not exist, will create it.")
+                    secrets_to_create.append(
+                        {
+                            "id": secret_id,
+                            "value": secret_value,
+                            "description": secret.get("description", f"Secret for {env_var}"),
+                        }
+                    )
+
+                    # Add to processed secrets
+                    processed_secret = {
+                        "id": secret_id,
+                        "type": secret_type,
+                        "env_var": env_var,
+                    }
+                    processed_secrets.append(processed_secret)
+                else:
+                    raise ValueError(
+                        f"Secret '{secret_id}' does not exist and no value provided for creation."
+                    )
+
+        # Create new secrets if any
+        if secrets_to_create:
+            try:
+                created_secrets = user.create_secrets(secrets_to_create)
+                logger.info(f"Successfully created {len(created_secrets)} new secrets.")
+            except Exception as e:
+                logger.error(f"Failed to create secrets: {e}")
+                raise
+
+        return processed_secrets
+
     def _is_clarifai_internal(self):
         """
         Check if the current user is a Clarifai internal user based on email domain.
@@ -891,19 +1001,21 @@ class ModelBuilder:
                     )
                 torch_version = dependencies.get('torch', None)
                 if 'torch' in dependencies:
-                    if python_version != DEFAULT_PYTHON_VERSION:
-                        raise Exception(
-                            f"torch is not supported with Python version {python_version}, please use Python version {DEFAULT_PYTHON_VERSION} in your config.yaml"
-                        )
                     if not torch_version:
                         logger.info(
                             f"Setup: torch version not found in requirements.txt, using the default version {DEFAULT_AMD_TORCH_VERSION}"
                         )
                         torch_version = DEFAULT_AMD_TORCH_VERSION
-                    if torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
-                        raise Exception(
-                            f"torch version {torch_version} not supported, please use one of the following versions: {DEFAULT_AMD_TORCH_VERSION} in your requirements.txt"
-                        )
+                    elif torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
+                        # Currently, we have only one vLLM image built with the DEFAULT_AMD_TORCH_VERSION.
+                        # If the user requests a different PyTorch version, that specific version will be
+                        # installed during the requirements.txt installation step
+                        torch_version = DEFAULT_AMD_TORCH_VERSION
+                else:
+                    logger.info(
+                        f"`torch` not found in requirements.txt, using the default torch=={DEFAULT_AMD_TORCH_VERSION}"
+                    )
+                    torch_version = DEFAULT_AMD_TORCH_VERSION
                 python_version = DEFAULT_PYTHON_VERSION
                 gpu_version = DEFAULT_AMD_GPU_VERSION
                 final_image = AMD_VLLM_BASE_IMAGE.format(
@@ -912,21 +1024,17 @@ class ModelBuilder:
                     gpu_version=gpu_version,
                 )
                 logger.info("Setup: Using vLLM base image to build the Docker image")
-            elif 'torch' in dependencies:
+            elif (
+                'torch' in dependencies
+                and (dependencies['torch'] in [None, DEFAULT_AMD_TORCH_VERSION])
+                and python_version == DEFAULT_PYTHON_VERSION
+            ):
                 torch_version = dependencies['torch']
-                if python_version != DEFAULT_PYTHON_VERSION:
-                    raise Exception(
-                        f"torch is not supported with Python version {python_version}, please use Python version {DEFAULT_PYTHON_VERSION} in your config.yaml"
-                    )
                 if not torch_version:
                     logger.info(
                         f"torch version not found in requirements.txt, using the default version {DEFAULT_AMD_TORCH_VERSION}"
                     )
                     torch_version = DEFAULT_AMD_TORCH_VERSION
-                if torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
-                    raise Exception(
-                        f"torch version {torch_version} not supported, please use one of the following versions: {DEFAULT_AMD_TORCH_VERSION} in your requirements.txt"
-                    )
                 python_version = DEFAULT_PYTHON_VERSION
                 gpu_version = DEFAULT_AMD_GPU_VERSION
                 final_image = AMD_TORCH_BASE_IMAGE.format(
@@ -1258,6 +1366,29 @@ class ModelBuilder:
             metadata_struct.update({'git_registry': git_info})
             model_version_proto.metadata.CopyFrom(metadata_struct)
 
+        # Process and add secrets to output_info.params
+        try:
+            processed_secrets = self._process_secrets()
+            if processed_secrets:
+                # Initialize output_info.params if not already present
+                if not model_version_proto.HasField("output_info"):
+                    model_version_proto.output_info.CopyFrom(resources_pb2.OutputInfo())
+
+                # Initialize params if not already present
+                if not model_version_proto.output_info.HasField("params"):
+                    from google.protobuf.struct_pb2 import Struct
+
+                    model_version_proto.output_info.params.CopyFrom(Struct())
+
+                # Add secrets to params
+                model_version_proto.output_info.params.update({"secrets": processed_secrets})
+                logger.info(
+                    f"Added {len(processed_secrets)} secrets to model version output_info.params"
+                )
+        except Exception as e:
+            logger.error(f"Failed to process secrets: {e}")
+            raise
+
         model_type_id = self.config.get('model').get('model_type_id')
         if model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
             if 'concepts' in self.config:
@@ -1382,6 +1513,7 @@ class ModelBuilder:
                     user_id=self.client.user_app_id.user_id,
                     app_id=self.client.user_app_id.app_id,
                     model_id=self.model_proto.id,
+                    colorize=True,
                 )
                 logger.info("""\n
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

clarifai/runners/models/model_runner.py

@@ -1,10 +1,10 @@
 import time
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Union
 
 from clarifai_grpc.grpc.api import service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 from clarifai_protocol import BaseRunner
-from clarifai_protocol.utils.health import HealthProbeRequestHandler
+from clarifai_protocol.utils.health import HealthProbeRequestHandler, start_health_server_thread
 
 from clarifai.client.auth.helper import ClarifaiAuthHelper
 from clarifai.utils.constants import STATUS_FAIL, STATUS_MIXED, STATUS_OK, STATUS_UNKNOWN
@@ -15,7 +15,7 @@ from ..utils.url_fetcher import ensure_urls_downloaded
 from .model_class import ModelClass
 
 
-class ModelRunner(BaseRunner, HealthProbeRequestHandler):
+class ModelRunner(BaseRunner):
     """
     This is a subclass of the runner class which will handle only the work items relevant to models.
     """
@@ -32,6 +32,7 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
         pat: Optional[str] = None,
         token: Optional[str] = None,
         num_parallel_polls: int = 4,
+        health_check_port: Union[int, None] = 8080,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -70,10 +71,27 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
                 # If auth helper creation fails, proceed without authentication
                 self._auth_helper = None
 
+        # if the model has a handle_liveness_probe method, call it to determine liveness
+        # otherwise rely on self.is_alive from the protocol
+        if hasattr(self.model, 'handle_liveness_probe'):
+            HealthProbeRequestHandler.handle_liveness_probe = self.model.handle_liveness_probe
+
+        # if the model has a handle_readiness_probe method, call it to determine readiness
+        # otherwise rely on self.is_ready from the protocol
+        if hasattr(self.model, 'handle_readiness_probe'):
+            HealthProbeRequestHandler.handle_readiness_probe = self.model.handle_readiness_probe
+
+        # if the model has a handle_startup_probe method, call it to determine startup
+        # otherwise rely on self.is_startup from the protocol
+        if hasattr(self.model, 'handle_startup_probe'):
+            HealthProbeRequestHandler.handle_startup_probe = self.model.handle_startup_probe
+
         # After model load successfully set the health probe to ready and startup
         HealthProbeRequestHandler.is_ready = True
         HealthProbeRequestHandler.is_startup = True
 
+        start_health_server_thread(port=health_check_port, address='')
+
     def get_runner_item_output_for_status(
         self, status: status_pb2.Status
     ) -> service_pb2.RunnerItemOutput:
@@ -301,27 +319,6 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
         """Set the model for this runner."""
         self.model = model
 
-    def handle_liveness_probe(self):
-        # if the model has a handle_liveness_probe method, call it to determine liveness
-        # otherwise rely on HealthProbeRequestHandler.is_alive from the protocol
-        if hasattr(self.model, 'handle_liveness_probe'):
-            HealthProbeRequestHandler.is_alive = self.model.handle_liveness_probe()
-        return super().handle_liveness_probe()
-
-    def handle_readiness_probe(self):
-        # if the model has a handle_readiness_probe method, call it to determine readiness
-        # otherwise rely on HealthProbeRequestHandler.is_ready from the protocol
-        if hasattr(self.model, 'handle_readiness_probe'):
-            HealthProbeRequestHandler.is_ready = self.model.handle_readiness_probe()
-        return super().handle_readiness_probe()
-
-    def handle_startup_probe(self):
-        # if the model has a handle_startup_probe method, call it to determine startup
-        # otherwise rely on HealthProbeRequestHandler.is_startup from the protocol
-        if hasattr(self.model, 'handle_startup_probe'):
-            HealthProbeRequestHandler.is_startup = self.model.handle_startup_probe()
-        return super().handle_startup_probe()
-
 
 def pmo_iterator(runner_item_iterator, auth_helper=None):
     for runner_item in runner_item_iterator:

clarifai/runners/server.py

@@ -290,6 +290,7 @@ class ModelServer:
                 model_id=context.model_id,
                 deployment_id=context.deployment_id,
                 base_url=context.api_base,
+                colorize=True,
             )
             logger.info(
                 "✅ Your model is running locally and is ready for requests from the API...\n"