clarifai 11.4.1__py3-none-any.whl → 11.4.3rc1__py3-none-any.whl

clarifai/runners/models/mcp_class.py

@@ -0,0 +1,143 @@
+"""Base class for creating Model Context Protocol (MCP) servers."""
+
+import asyncio
+import json
+from typing import Any
+
+from fastmcp import Client, FastMCP  # use fastmcp v2 not the built in mcp
+from mcp import types
+from mcp.shared.exceptions import McpError
+
+from clarifai.runners.models.model_class import ModelClass
+
+# class MCPServerProvider(abc.ABC):
+#     """
+#     Base class for creating Model Context Protocol (MCP) servers.
+
+#     This class provides a base implementation of the MCP server, including
+#     methods for handling requests and responses, as well as error handling and
+#     logging.
+
+#     Attributes:
+#     _server: The FastMCP server instance.
+#     _tools: List of tools available in the server.
+#     _resources: List of resources available in the server.
+#     _prompts: List of prompts available in the server.
+
+#     Methods:
+#     get_server(): Returns the FastMCP server instance.
+#     mcp_transport(msg): Handles incoming messages and sends them to the FastMCP server.
+#     """
+
+#     @abc.abstractmethod
+#     def get_server(self) -> FastMCP:
+#         """Required method for each subclass to implement to return the FastMCP server to use."""
+#         if self._server is None:
+#             raise ValueError("Server not initialized")
+#         return self._server
+
+
+class MCPModelClass(ModelClass):
+    """Base class for wrapping FastMCP servers as a model running in Clarfai. This handles
+    all the transport between the API and the MCP server here. Simply subclass this and implement
+    the get_server() method to return the FastMCP server instance. The server is then used to
+    handle all the requests and responses.
+    """
+
+    def load_model(self):
+        # in memory transport provided in fastmcp v2 so we can easily use the client functions.
+        self.client = Client(self.get_server())
+
+    def get_server(self) -> FastMCP:
+        """Required method for each subclass to implement to return the FastMCP server to use."""
+        raise NotImplementedError("Subclasses must implement get_server() method")
+
+    @ModelClass.method
+    def mcp_transport(self, msg: str) -> str:
+        """The single model method to get the jsonrpc message and send it to the FastMCP server then
+        return it's response.
+
+        """
+
+        async def send_notification(client_message: types.ClientNotification) -> None:
+            async with self.client:
+                # Strip the jsonrpc field since send_notification will also pass it in for some reason.
+                client_message = types.ClientNotification.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc"}
+                    )
+                )
+                try:
+                    return await self.client.session.send_notification(client_message)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", error=e.error)
+
+        async def send_request(client_message: types.ClientRequest, id: str) -> Any:
+            async with self.client:
+                # Strip the jsonrpc and id fields as send_request sets them again too.
+                client_message = types.ClientRequest.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc", "id"}
+                    )
+                )
+
+                result_type = None
+                print("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
+                print(types.PingRequest)
+                if isinstance(client_message.root, types.PingRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.InitializeRequest):
+                    return await self.client.session.initialize()
+                elif isinstance(client_message.root, types.SetLevelRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListResourcesRequest):
+                    result_type = types.ListResourcesResult
+                elif isinstance(client_message.root, types.ListResourceTemplatesRequest):
+                    result_type = types.ListResourceTemplatesResult
+                elif isinstance(client_message.root, types.ReadResourceRequest):
+                    result_type = types.ReadResourceResult
+                elif isinstance(client_message.root, types.SubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.UnsubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListPromptsRequest):
+                    result_type = types.ListPromptsResult
+                elif isinstance(client_message.root, types.GetPromptRequest):
+                    result_type = types.GetPromptResult
+                elif isinstance(client_message.root, types.CompleteRequest):
+                    result_type = types.CompleteResult
+                elif isinstance(client_message.root, types.ListToolsRequest):
+                    result_type = types.ListToolsResult
+                elif isinstance(client_message.root, types.CallToolRequest):
+                    result_type = types.CallToolResult
+                else:
+                    # this is a special case where we need to return the list of tools.
+                    raise NotImplementedError(f"Method {client_message.method} not implemented")
+                # Call the mcp server using send_request() or send_notification() depending on the method.
+                try:
+                    return await self.client.session.send_request(client_message, result_type)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", id=id, error=e.error)
+
+        # The message coming here is the generic request. We look at it's .method
+        # to determine which client function to call and to further subparse the params.
+        # Note(zeiler): unfortunately the pydantic types in mcp/types.py are not consistent.
+        # The JSONRPCRequest are supposed to have an id but the InitializeRequest
+        # does not have it.
+        d = json.loads(msg)
+
+        # If we have an id it's a JSONRPCRequest
+        if not d.get('method', None).startswith("notifications/"):
+            # rpc_message2 = types.JSONRPCRequest.model_validate(rpc_message)
+            # underlying: types.JSONRPCRequest = rpc_message.root
+            client_message = types.ClientRequest.model_validate(d)
+            response = asyncio.run(
+                send_request(client_message, id=d.get('id', ""))
+            )  # underlying.id))
+        else:  # JSONRPCRequest
+            client_message = types.ClientNotification.model_validate(d)
+            response = asyncio.run(send_notification(client_message))
+        if response is None:
+            return ""
+        # return as a serialized json string
+        return response.model_dump_json(by_alias=True, exclude_none=True)

clarifai/runners/models/mcp_class.py~

@@ -0,0 +1,149 @@
+"""Base class for creating Model Context Protocol (MCP) servers."""
+
+import asyncio
+import json
+from typing import Any
+
+from fastmcp import Client, FastMCP  # use fastmcp v2 not the built in mcp
+from mcp import types
+from mcp.shared.exceptions import McpError
+
+from clarifai.runners.models.model_class import ModelClass
+
+# class MCPServerProvider(abc.ABC):
+#     """
+#     Base class for creating Model Context Protocol (MCP) servers.
+
+#     This class provides a base implementation of the MCP server, including
+#     methods for handling requests and responses, as well as error handling and
+#     logging.
+
+#     Attributes:
+#     _server: The FastMCP server instance.
+#     _tools: List of tools available in the server.
+#     _resources: List of resources available in the server.
+#     _prompts: List of prompts available in the server.
+
+#     Methods:
+#     get_server(): Returns the FastMCP server instance.
+#     mcp_transport(msg): Handles incoming messages and sends them to the FastMCP server.
+#     """
+
+#     @abc.abstractmethod
+#     def get_server(self) -> FastMCP:
+#         """Required method for each subclass to implement to return the FastMCP server to use."""
+#         if self._server is None:
+#             raise ValueError("Server not initialized")
+#         return self._server
+
+
+class MCPModelClass(ModelClass, MCPServerProvider):
+    """Base class for wrapping FastMCP servers as a model running in Clarfai. This handles
+    all the transport between the API and the MCP server here. Simply subclass this and implement
+    the get_server() method to return the FastMCP server instance. The server is then used to
+    handle all the requests and responses.
+    """
+
+    def load_model(self):
+        # in memory transport provided in fastmcp v2 so we can easily use the client functions.
+        self.client = Client(self.get_server())
+
+    def get_server(self) -> FastMCP:
+        """Required method for each subclass to implement to return the FastMCP server to use."""
+        if self._server is None:
+            raise ValueError("Server not initialized")
+        return self._server
+
+    @ModelClass.method
+    def mcp_transport(self, msg: str) -> str:
+        """The single model method to get the jsonrpc message and send it to the FastMCP server then
+        return it's response.
+
+        Arguments:
+          msg: The incoming message to be handled in serialized JSONRPC format from an MCP client.
+        Returns:
+          str: The response to the incoming message in serialized JSONRPC format
+        """
+
+        async def send_notification(client_message: types.ClientNotification) -> None:
+            async with self.client:
+                # Strip the jsonrpc field since send_notification will also pass it in for some reason.
+                client_message = types.ClientNotification.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc"}
+                    )
+                )
+                try:
+                    return await self.client.session.send_notification(client_message)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", error=e.error)
+
+        async def send_request(client_message: types.ClientRequest, id: str) -> Any:
+            async with self.client:
+                # Strip the jsonrpc and id fields as send_request sets them again too.
+                client_message = types.ClientRequest.model_validate(
+                    client_message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True, exclude={"jsonrpc", "id"}
+                    )
+                )
+
+                result_type = None
+                print("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA")
+                print(types.PingRequest)
+                if isinstance(client_message.root, types.PingRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.InitializeRequest):
+                    return await self.client.session.initialize()
+                elif isinstance(client_message.root, types.SetLevelRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListResourcesRequest):
+                    result_type = types.ListResourcesResult
+                elif isinstance(client_message.root, types.ListResourceTemplatesRequest):
+                    result_type = types.ListResourceTemplatesResult
+                elif isinstance(client_message.root, types.ReadResourceRequest):
+                    result_type = types.ReadResourceResult
+                elif isinstance(client_message.root, types.SubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.UnsubscribeRequest):
+                    result_type = types.EmptyResult
+                elif isinstance(client_message.root, types.ListPromptsRequest):
+                    result_type = types.ListPromptsResult
+                elif isinstance(client_message.root, types.GetPromptRequest):
+                    result_type = types.GetPromptResult
+                elif isinstance(client_message.root, types.CompleteRequest):
+                    result_type = types.CompleteResult
+                elif isinstance(client_message.root, types.ListToolsRequest):
+                    result_type = types.ListToolsResult
+                elif isinstance(client_message.root, types.CallToolRequest):
+                    result_type = types.CallToolResult
+                else:
+                    # this is a special case where we need to return the list of tools.
+                    raise NotImplementedError(f"Method {client_message.method} not implemented")
+                # Call the mcp server using send_request() or send_notification() depending on the method.
+                try:
+                    return await self.client.session.send_request(client_message, result_type)
+                except McpError as e:
+                    return types.JSONRPCError(jsonrpc="2.0", id=id, error=e.error)
+
+        # The message coming here is the generic request. We look at it's .method
+        # to determine which client function to call and to further subparse the params.
+        # Note(zeiler): unfortunately the pydantic types in mcp/types.py are not consistent.
+        # The JSONRPCRequest are supposed to have an id but the InitializeRequest
+        # does not have it.
+        d = json.loads(msg)
+
+        # If we have an id it's a JSONRPCRequest
+        if not d.get('method', None).startswith("notifications/"):
+            # rpc_message2 = types.JSONRPCRequest.model_validate(rpc_message)
+            # underlying: types.JSONRPCRequest = rpc_message.root
+            client_message = types.ClientRequest.model_validate(d)
+            response = asyncio.run(
+                send_request(client_message, id=d.get('id', ""))
+            )  # underlying.id))
+        else:  # JSONRPCRequest
+            client_message = types.ClientNotification.model_validate(d)
+            response = asyncio.run(send_notification(client_message))
+        if response is None:
+            return ""
+        # return as a serialized json string
+        return response.model_dump_json(by_alias=True, exclude_none=True)

clarifai/runners/models/model_builder.py

@@ -14,15 +14,17 @@ import yaml
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2
 from google.protobuf import json_format
-from rich import print
-from rich.markup import escape
 
 from clarifai.client.base import BaseClient
 from clarifai.runners.models.model_class import ModelClass
 from clarifai.runners.utils.const import (
+    AMD_PYTHON_BASE_IMAGE,
+    AMD_VLLM_BASE_IMAGE,
     AVAILABLE_PYTHON_IMAGES,
     AVAILABLE_TORCH_IMAGES,
     CONCEPTS_REQUIRED_MODEL_TYPE,
+    DEFAULT_AMD_GPU_VERSION,
+    DEFAULT_AMD_TORCH_VERSION,
     DEFAULT_DOWNLOAD_CHECKPOINT_WHEN,
     DEFAULT_PYTHON_VERSION,
     DEFAULT_RUNTIME_DOWNLOAD_PATH,
@@ -43,13 +45,6 @@ dependencies = [
 ]
 
 
-def _clear_line(n: int = 1) -> None:
-    LINE_UP = '\033[1A'  # Move cursor up one line
-    LINE_CLEAR = '\x1b[2K'  # Clear the entire line
-    for _ in range(n):
-        print(LINE_UP, end=LINE_CLEAR, flush=True)
-
-
 def is_related(object_class, main_class):
     # Check if the object_class is a subclass of main_class
     if issubclass(object_class, main_class):
@@ -361,13 +356,23 @@ class ModelBuilder:
         if self.config.get("checkpoints"):
             loader_type, _, hf_token, _, _, _ = self._validate_config_checkpoints()
 
-            if loader_type == "huggingface" and hf_token:
-                is_valid_token = HuggingFaceLoader.validate_hftoken(hf_token)
-                if not is_valid_token:
+            if loader_type == "huggingface":
+                is_valid_token = hf_token and HuggingFaceLoader.validate_hftoken(hf_token)
+                if not is_valid_token and hf_token:
+                    logger.info(
+                        "Continuing without Hugging Face token for validating config in model builder."
+                    )
+
+                has_repo_access = HuggingFaceLoader.validate_hf_repo_access(
+                    repo_id=self.config.get("checkpoints", {}).get("repo_id"),
+                    token=hf_token if is_valid_token else None,
+                )
+
+                if not has_repo_access:
                     logger.error(
-                        "Invalid Hugging Face token provided in the config file, this might cause issues with downloading the restricted model checkpoints."
+                        f"Invalid Hugging Face repo access for repo {self.config.get('checkpoints').get('repo_id')}. Please check your repo and try again."
                     )
-                    logger.info("Continuing without Hugging Face token")
+                    sys.exit("Token does not have access to HuggingFace repo , exiting.")
 
         num_threads = self.config.get("num_threads")
         if num_threads or num_threads == 0:
@@ -405,11 +410,11 @@ class ModelBuilder:
         signatures = {method.name: method.signature for method in method_info.values()}
         return signatures_to_yaml(signatures)
 
-    def get_method_signatures(self):
+    def get_method_signatures(self, mocking=True):
         """
         Returns the method signatures for the model class.
         """
-        model_class = self.load_model_class(mocking=True)
+        model_class = self.load_model_class(mocking=mocking)
         method_info = model_class._get_method_info()
         signatures = [method.signature for method in method_info.values()]
         return signatures
@@ -532,6 +537,30 @@ class ModelBuilder:
                 dependencies_version[dependency] = version if version else None
         return dependencies_version
 
+    def _is_amd(self):
+        """
+        Check if the model is AMD or not.
+        """
+        is_amd_gpu = False
+        is_nvidia_gpu = False
+        if "inference_compute_info" in self.config:
+            inference_compute_info = self.config.get('inference_compute_info')
+            if 'accelerator_type' in inference_compute_info:
+                for accelerator in inference_compute_info['accelerator_type']:
+                    if 'amd' in accelerator.lower():
+                        is_amd_gpu = True
+                    elif 'nvidia' in accelerator.lower():
+                        is_nvidia_gpu = True
+        if is_amd_gpu and is_nvidia_gpu:
+            raise Exception(
+                "Both AMD and NVIDIA GPUs are specified in the config file, please use only one type of GPU."
+            )
+        if is_amd_gpu:
+            logger.info("Using AMD base image to build the Docker image and upload the model")
+        elif is_nvidia_gpu:
+            logger.info("Using NVIDIA base image to build the Docker image and upload the model")
+        return is_amd_gpu
+
     def create_dockerfile(self):
         dockerfile_template = os.path.join(
             os.path.dirname(os.path.dirname(__file__)),
@@ -562,30 +591,85 @@ class ModelBuilder:
             )
             python_version = DEFAULT_PYTHON_VERSION
 
-        # This is always the final image used for runtime.
-        final_image = PYTHON_BASE_IMAGE.format(python_version=python_version)
-        downloader_image = PYTHON_BASE_IMAGE.format(python_version=python_version)
-
         # Parse the requirements.txt file to determine the base image
         dependencies = self._parse_requirements()
-        if 'torch' in dependencies and dependencies['torch']:
-            torch_version = dependencies['torch']
-
-            # Sort in reverse so that newer cuda versions come first and are preferred.
-            for image in sorted(AVAILABLE_TORCH_IMAGES, reverse=True):
-                if torch_version in image and f'py{python_version}' in image:
-                    # like cu124, rocm6.3, etc.
-                    gpu_version = image.split('-')[-1]
-                    final_image = TORCH_BASE_IMAGE.format(
-                        torch_version=torch_version,
-                        python_version=python_version,
-                        gpu_version=gpu_version,
+
+        is_amd_gpu = self._is_amd()
+        if is_amd_gpu:
+            final_image = AMD_PYTHON_BASE_IMAGE.format(python_version=python_version)
+            downloader_image = AMD_PYTHON_BASE_IMAGE.format(python_version=python_version)
+            if 'vllm' in dependencies:
+                if python_version != DEFAULT_PYTHON_VERSION:
+                    raise Exception(
+                        f"vLLM is not supported with Python version {python_version}, please use Python version {DEFAULT_PYTHON_VERSION} in your config.yaml"
                     )
+                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"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_VLLM_BASE_IMAGE.format(
+                    torch_version=torch_version,
+                    python_version=python_version,
+                    gpu_version=gpu_version,
+                )
+                logger.info("Using vLLM base image to build the Docker image")
+            elif 'torch' in dependencies:
+                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"Using Torch version {torch_version} base image to build the Docker image"
+                        f"torch version not found in requirements.txt, using the default version {DEFAULT_AMD_TORCH_VERSION}"
                     )
-                    break
-
+                    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 = TORCH_BASE_IMAGE.format(
+                    torch_version=torch_version,
+                    python_version=python_version,
+                    gpu_version=gpu_version,
+                )
+                logger.info(
+                    f"Using Torch version {torch_version} base image to build the Docker image"
+                )
+        else:
+            final_image = PYTHON_BASE_IMAGE.format(python_version=python_version)
+            downloader_image = PYTHON_BASE_IMAGE.format(python_version=python_version)
+            if 'torch' in dependencies and dependencies['torch']:
+                torch_version = dependencies['torch']
+                # Sort in reverse so that newer cuda versions come first and are preferred.
+                for image in sorted(AVAILABLE_TORCH_IMAGES, reverse=True):
+                    if torch_version in image and f'py{python_version}' in image:
+                        # like cu124, rocm6.3, etc.
+                        gpu_version = image.split('-')[-1]
+                        final_image = TORCH_BASE_IMAGE.format(
+                            torch_version=torch_version,
+                            python_version=python_version,
+                            gpu_version=gpu_version,
+                        )
+                        logger.info(
+                            f"Using Torch version {torch_version} base image to build the Docker image"
+                        )
+                        break
         if 'clarifai' not in dependencies:
             raise Exception(
                 f"clarifai not found in requirements.txt, please add clarifai to the requirements.txt file with a fixed version. Current version is clarifai=={CLIENT_VERSION}"
@@ -835,7 +919,6 @@ class ModelBuilder:
                 percent_completed = response.status.percent_completed
             details = response.status.details
 
-            _clear_line()
             print(
                 f"Status: {response.status.description}, Progress: {percent_completed}% - {details} ",
                 f"request_id: {response.status.req_id}",
@@ -849,7 +932,48 @@ class ModelBuilder:
         logger.info(f"Created Model Version ID: {self.model_version_id}")
         logger.info(f"Full url to that version is: {self.model_url}")
         try:
-            self.monitor_model_build()
+            is_uploaded = self.monitor_model_build()
+            if is_uploaded:
+                # Provide an mcp client config
+                if model_type_id == "mcp":
+                    snippet = (
+                        """
+import asyncio
+import os
+from fastmcp import Client
+from fastmcp.client.transports import StreamableHttpTransport
+
+transport = StreamableHttpTransport(url="%s/mcp",
+                                    headers={"Authorization": "Bearer " + os.environ["CLARIFAI_PAT"]})
+
+async def main():
+  async with Client(transport) as client:
+    tools = await client.list_tools()
+    print(f"Available tools: {tools}")
+    result = await client.call_tool(tools[0].name, {"a": 5, "b": 3})
+    print(f"Result: {result[0].text}")
+
+if __name__ == "__main__":
+  asyncio.run(main())
+"""
+                        % self.model_url
+                    )
+                else:  # python code to run the model.
+                    from clarifai.runners.utils import code_script
+
+                    method_signatures = self.get_method_signatures()
+                    snippet = code_script.generate_client_script(
+                        method_signatures,
+                        user_id=self.client.user_app_id.user_id,
+                        app_id=self.client.user_app_id.app_id,
+                        model_id=self.model_proto.id,
+                    )
+                logger.info("""\n
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+# Here is a code snippet to use this model:
+XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+                """)
+                logger.info(snippet)
         finally:
             if os.path.exists(self.tar_file):
                 logger.debug(f"Cleaning up upload file: {self.tar_file}")
@@ -933,7 +1057,12 @@ class ModelBuilder:
             for log_entry in logs.log_entries:
                 if log_entry.url not in seen_logs:
                     seen_logs.add(log_entry.url)
-                    logger.info(f"{escape(log_entry.message.strip())}")
+                    log_entry_msg = re.sub(
+                        r"(\\*)(\[[a-z#/@][^[]*?])",
+                        lambda m: f"{m.group(1)}{m.group(1)}\\{m.group(2)}",
+                        log_entry.message.strip(),
+                    )
+                    logger.info(log_entry_msg)
             if status_code == status_code_pb2.MODEL_BUILDING:
                 print(
                     f"Model is building... (elapsed {time.time() - st:.1f}s)", end='\r', flush=True

clarifai/runners/models/model_class.py

@@ -9,7 +9,6 @@ from typing import Any, Dict, Iterator, List
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
-from google.protobuf import json_format
 
 from clarifai.runners.utils import data_types
 from clarifai.runners.utils.data_utils import DataConverter
@@ -100,7 +99,6 @@ class ModelClass(ABC):
         try:
             # TODO add method name field to proto
             method_name = 'predict'
-            inference_params = get_inference_params(request)
             if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
                 method_name = request.inputs[0].data.metadata['_method_name']
             if (
@@ -124,7 +122,7 @@ class ModelClass(ABC):
                     input.data.CopyFrom(new_data)
             # convert inputs to python types
             inputs = self._convert_input_protos_to_python(
-                request.inputs, inference_params, signature.input_fields, python_param_types
+                request.inputs, signature.input_fields, python_param_types
             )
             if len(inputs) == 1:
                 inputs = inputs[0]
@@ -163,7 +161,6 @@ class ModelClass(ABC):
     ) -> Iterator[service_pb2.MultiOutputResponse]:
         try:
             method_name = 'generate'
-            inference_params = get_inference_params(request)
             if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
                 method_name = request.inputs[0].data.metadata['_method_name']
             method = getattr(self, method_name)
@@ -180,7 +177,7 @@ class ModelClass(ABC):
                     )
                     input.data.CopyFrom(new_data)
             inputs = self._convert_input_protos_to_python(
-                request.inputs, inference_params, signature.input_fields, python_param_types
+                request.inputs, signature.input_fields, python_param_types
             )
             if len(inputs) == 1:
                 inputs = inputs[0]
@@ -226,7 +223,6 @@ class ModelClass(ABC):
             assert len(request.inputs) == 1, "Streaming requires exactly one input"
 
             method_name = 'stream'
-            inference_params = get_inference_params(request)
             if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
                 method_name = request.inputs[0].data.metadata['_method_name']
             method = getattr(self, method_name)
@@ -251,7 +247,7 @@ class ModelClass(ABC):
                     input.data.CopyFrom(new_data)
             # convert all inputs for the first request, including the first stream value
             inputs = self._convert_input_protos_to_python(
-                request.inputs, inference_params, signature.input_fields, python_param_types
+                request.inputs, signature.input_fields, python_param_types
             )
             kwargs = inputs[0]
 
@@ -264,7 +260,7 @@ class ModelClass(ABC):
                 # subsequent streaming items contain only the streaming input
                 for request in request_iterator:
                     item = self._convert_input_protos_to_python(
-                        request.inputs, inference_params, [stream_sig], python_param_types
+                        request.inputs, [stream_sig], python_param_types
                     )
                     item = item[0][stream_argname]
                     yield item
@@ -297,13 +293,12 @@ class ModelClass(ABC):
     def _convert_input_protos_to_python(
         self,
         inputs: List[resources_pb2.Input],
-        inference_params: dict,
         variables_signature: List[resources_pb2.ModelTypeField],
         python_param_types,
     ) -> List[Dict[str, Any]]:
         result = []
         for input in inputs:
-            kwargs = deserialize(input.data, variables_signature, inference_params)
+            kwargs = deserialize(input.data, variables_signature)
             # dynamic cast to annotated types
             for k, v in kwargs.items():
                 if k not in python_param_types:
@@ -374,18 +369,6 @@ class ModelClass(ABC):
         return method_info
 
 
-# Helper function to get the inference params
-def get_inference_params(request) -> dict:
-    """Get the inference params from the request."""
-    inference_params = {}
-    if request.model.model_version.id != "":
-        output_info = request.model.model_version.output_info
-        output_info = json_format.MessageToDict(output_info, preserving_proto_field_name=True)
-        if "params" in output_info:
-            inference_params = output_info["params"]
-    return inference_params
-
-
 class _MethodInfo:
     def __init__(self, method):
         self.name = method.__name__