clarifai 11.1.7__py3-none-any.whl → 11.1.7rc2__py3-none-any.whl

clarifai/runners/dockerfile_template/Dockerfile.nim

@@ -0,0 +1,71 @@
+FROM nvcr.io/nim/meta/llama-3.1-8b-instruct:1.1.2 as build
+
+FROM gcr.io/distroless/python3-debian12:debug
+
+
+COPY --from=build /bin/bash /bin/rbash
+COPY --from=build /bin/sh /bin/sh
+COPY --from=build /bin/rsh /bin/rsh
+
+# we have to overwrite the python3 binary that the distroless image uses
+COPY --from=build /opt/nim/llm/.venv/bin/python3.10 /usr/bin/python3
+COPY --from=build /opt/nim/llm/.venv/bin/python3.10 /usr/local/bin/python3.10
+
+# also copy in all the lib files for it.
+COPY --from=build /lib /lib
+COPY --from=build /lib64 /lib64
+COPY --from=build /usr/lib/ /usr/lib/
+COPY --from=build /usr/local/lib/ /usr/local/lib/
+# ldconfig is needed to update the shared library cache so system libraries (like CUDA) can be found
+COPY --from=build /usr/sbin/ldconfig /sbin/ldconfig
+COPY --from=build /usr/sbin/ldconfig.real /sbin/ldconfig.real
+COPY --from=build /etc/ld.so.conf /etc/ld.so.conf
+COPY --from=build /etc/ld.so.cache /etc/ld.so.cache
+COPY --from=build /etc/ld.so.conf.d/ /etc/ld.so.conf.d/
+
+# COPY NIM files
+COPY  --from=build /opt /opt
+COPY  --from=build /etc/nim /etc/nim
+
+# Set environment variables to use the nim libraries and python
+ENV PYTHONPATH=${PYTHONPATH}:/opt/nim/llm/.venv/lib/python3.10/site-packages:/opt/nim/llm
+ENV PATH="/opt/nim/llm/.venv/bin:/opt/hpcx/ucc/bin:/opt/hpcx/ucx/bin:/opt/hpcx/ompi/bin:$PATH"
+
+ENV LD_LIBRARY_PATH="/opt/hpcx/ucc/lib/ucc:/opt/hpcx/ucc/lib:/opt/hpcx/ucx/lib/ucx:/opt/hpcx/ucx/lib:/opt/hpcx/ompi/lib:/opt/hpcx/ompi/lib/openmpi:/opt/nim/llm/.venv/lib/python3.10/site-packages/tensorrt_llm/libs:/opt/nim/llm/.venv/lib/python3.10/site-packages/nvidia/cublas/lib:/opt/nim/llm/.venv/lib/python3.10/site-packages/tensorrt_libs:/opt/nim/llm/.venv/lib/python3.10/site-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH"
+
+ENV LIBRARY_PATH=/opt/hpcx/ucc/lib:/opt/hpcx/ucx/lib:/opt/hpcx/ompi/lib:$LIBRARY_PATH
+
+ENV CPATH=/opt/hpcx/ompi/include:/opt/hpcx/ucc/include:/opt/hpcx/ucx/include:$CPATH
+ENV LLM_PROJECT_DIR=/opt/nim/llm
+
+# Set environment variables for MPI
+ENV OMPI_HOME=/opt/hpcx/ompi
+ENV HPCX_MPI_DIR=/opt/hpcx/ompi
+ENV MPIf_HOME=/opt/hpcx/ompi
+ENV OPAL_PREFIX=/opt/hpcx/ompi
+
+# Set environment variables for UCC
+ENV UCC_DIR=/opt/hpcx/ucc/lib/cmake/ucc
+ENV UCC_HOME=/opt/hpcx/ucc
+ENV HPCX_UCC_DIR=/opt/hpcx/ucc
+ENV USE_UCC=1
+ENV USE_SYSTEM_UCC=1
+
+# Set environment variables for HPC-X
+ENV HPCX_DIR=/opt/hpcx
+ENV HPCX_UCX_DIR=/opt/hpcx/ucx
+ENV HPCX_MPI_DIR=/opt/hpcx/ompi
+
+# Set environment variables for UCX
+ENV UCX_DIR=/opt/hpcx/ucx/lib/cmake/ucx
+ENV UCX_HOME=/opt/hpcx/ucx
+
+ENV HOME=/opt/nim/llm
+
+# ln is needed to create symbolic links (needed by nvidia-container-runtime)
+COPY --from=build /usr/bin/ln /usr/bin/ln
+
+# Run ldconfig in the build stage to update the library cache else CUDA libraries won't be found
+RUN ldconfig -v
+
+SHELL ["/bin/rbash", "-c"]

clarifai/runners/dockerfile_template/Dockerfile.template

@@ -44,6 +44,9 @@ ENV PYTHONPATH=${PYTHONPATH}:/home/nonroot/main \
     CLARIFAI_COMPUTE_CLUSTER_ID=${CLARIFAI_COMPUTE_CLUSTER_ID} \
     CLARIFAI_API_BASE=${CLARIFAI_API_BASE:-https://api.clarifai.com}
 
+# # Write out the model function signatures
+# RUN ["python", "-m", "clarifai.cli", "model", "signatures", "--model_path", "/home/nonroot/main", "--out_path", "/home/nonroot/main/signatures.yaml"]
+
 # Finally run the clarifai entrypoint to start the runner loop and local dev server.
 # Note(zeiler): we may want to make this a clarifai CLI call.
 ENTRYPOINT ["python", "-m", "clarifai.runners.server"]

clarifai/runners/models/model_builder.py

@@ -14,13 +14,14 @@ from google.protobuf import json_format
 from rich import print
 from rich.markup import escape
 
-from clarifai.client import BaseClient
+from clarifai.client.base import BaseClient
 from clarifai.runners.models.model_class import ModelClass
 from clarifai.runners.utils.const import (
     AVAILABLE_PYTHON_IMAGES, AVAILABLE_TORCH_IMAGES, CONCEPTS_REQUIRED_MODEL_TYPE,
     DEFAULT_DOWNLOAD_CHECKPOINT_WHEN, DEFAULT_PYTHON_VERSION, DEFAULT_RUNTIME_DOWNLOAD_PATH,
     PYTHON_BASE_IMAGE, TORCH_BASE_IMAGE)
 from clarifai.runners.utils.loader import HuggingFaceLoader
+from clarifai.runners.utils.method_signatures import signatures_to_yaml
 from clarifai.urls.helper import ClarifaiUrlHelper
 from clarifai.utils.logging import logger
 from clarifai.versions import CLIENT_VERSION
@@ -69,6 +70,18 @@ class ModelBuilder:
     """
     Create an instance of the model class, as specified in the config file.
     """
+    model_class = self.load_model_class()
+
+    # initialize the model
+    model = model_class()
+    if load_model:
+      model.load_model()
+    return model
+
+  def load_model_class(self):
+    """
+    Import the model class from the model.py file.
+    """
     # look for default model.py file location
     for loc in ["model.py", "1/model.py"]:
       model_file = os.path.join(self.folder, loc)
@@ -107,12 +120,7 @@ class ModelBuilder:
           "Could not determine model class. There should be exactly one model inheriting from ModelClass defined in the model.py"
       )
     model_class = classes[0]
-
-    # initialize the model
-    model = model_class()
-    if load_model:
-      model.load_model()
-    return model
+    return model_class
 
   def _validate_folder(self, folder):
     if folder == ".":
@@ -253,6 +261,24 @@ class ModelBuilder:
           total_size += member.size
     return total_size
 
+  def method_signatures_yaml(self):
+    """
+    Returns the method signatures for the model class in YAML format.
+    """
+    model_class = self.load_model_class()
+    method_info = model_class._get_method_info()
+    signatures = {name: m.signature for name, m in method_info.values()}
+    return signatures_to_yaml(signatures)
+
+  def get_method_signatures(self):
+    """
+    Returns the method signatures for the model class.
+    """
+    model_class = self.load_model_class()
+    method_info = model_class._get_method_info()
+    signatures = [method.signature for method in method_info.values()]
+    return signatures
+
   @property
   def client(self):
     if self._client is None:
@@ -536,6 +562,8 @@ class ModelBuilder:
         pretrained_model_config=resources_pb2.PretrainedModelConfig(),
         inference_compute_info=self.inference_compute_info,
     )
+    # TODO: update this to `model_signatures` field when it's available in the API
+    model_version_proto.model_signature.extend(self.get_method_signatures())
 
     model_type_id = self.config.get('model').get('model_type_id')
     if model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:

clarifai/runners/models/model_class.py

@@ -1,41 +1,275 @@
-from abc import ABC, abstractmethod
-from typing import Iterator
+import inspect
+import itertools
+import logging
+import os
+import traceback
+from abc import ABC
+from typing import Any, Dict, Iterator, List
 
-from clarifai_grpc.grpc.api import service_pb2
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
+
+from clarifai.runners.utils import data_types
+from clarifai.runners.utils.method_signatures import (build_function_signature, deserialize,
+                                                      get_stream_from_signature, serialize,
+                                                      signatures_to_json)
+
+_METHOD_INFO_ATTR = '_cf_method_info'
+
+_RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() in ("true", "1")
 
 
 class ModelClass(ABC):
+  '''
+  Base class for model classes that can be run as a service.
+
+  Define predict, generate, or stream methods using the @ModelClass.method decorator.
+
+  Example:
+
+    from clarifai.runners.model_class import ModelClass
+    from clarifai.runners.utils.data_types import NamedFields, Stream
+
+    class MyModel(ModelClass):
+
+      @ModelClass.method
+      def predict(self, x: str, y: int) -> List[str]:
+        return [x] * y
+
+      @ModelClass.method
+      def generate(self, x: str, y: int) -> Stream[str]:
+        for i in range(y):
+          yield x + str(i)
+
+      @ModelClass.method
+      def stream(self, input_stream: Stream[NamedFields(x=str, y=int)]) -> Stream[str]:
+        for item in input_stream:
+          yield item.x + ' ' + str(item.y)
+  '''
+
+  @staticmethod
+  def method(func):
+    setattr(func, _METHOD_INFO_ATTR, _MethodInfo(func))
+    return func
+
+  def load_model(self):
+    """Load the model."""
+
+  def _handle_get_signatures_request(self) -> service_pb2.MultiOutputResponse:
+    methods = self._get_method_info()
+    signatures = {method.name: method.signature for method in methods.values()}
+    resp = service_pb2.MultiOutputResponse(status=status_pb2.Status(code=status_code_pb2.SUCCESS))
+    output = resp.outputs.add()
+    output.status.code = status_code_pb2.SUCCESS
+    output.data.text.raw = signatures_to_json(signatures)
+    return resp
+
+  def _batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
+    """Batch predict method for multiple inputs."""
+    outputs = []
+    for input in inputs:
+      output = method(**input)
+      outputs.append(output)
+    return outputs
+
+  def _batch_generate(self, method, inputs: List[Dict[str, Any]]) -> Iterator[List[Any]]:
+    """Batch generate method for multiple inputs."""
+    generators = [method(**input) for input in inputs]
+    for outputs in itertools.zip_longest(*generators):
+      yield outputs
 
   def predict_wrapper(
       self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
-    """This method is used for input/output proto data conversion"""
-    return self.predict(request)
+    outputs = []
+    try:
+      # TODO add method name field to proto
+      method_name = 'predict'
+      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
+        method_name = request.inputs[0].data.metadata['_method_name']
+      if method_name == '_GET_SIGNATURES':  # special case to fetch signatures, TODO add endpoint for this
+        return self._handle_get_signatures_request()
+      if method_name not in self._get_method_info():
+        raise ValueError(f"Method {method_name} not found in model class")
+      method = getattr(self, method_name)
+      method_info = method._cf_method_info
+      signature = method_info.signature
+      python_param_types = method_info.python_param_types
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
+                                                    python_param_types)
+      if len(inputs) == 1:
+        inputs = inputs[0]
+        output = method(**inputs)
+        outputs.append(self._convert_output_to_proto(output, signature.output_fields))
+      else:
+        outputs = self._batch_predict(method, inputs)
+        outputs = [
+            self._convert_output_to_proto(output, signature.output_fields) for output in outputs
+        ]
+
+      return service_pb2.MultiOutputResponse(
+          outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS))
+    except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
+      logging.exception("Error in predict")
+      return service_pb2.MultiOutputResponse(status=status_pb2.Status(
+          code=status_code_pb2.FAILURE,
+          details=str(e),
+          stack_trace=traceback.format_exc().split('\n')))
 
   def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
                       ) -> Iterator[service_pb2.MultiOutputResponse]:
-    """This method is used for input/output proto data conversion and yield outcome"""
-    return self.generate(request)
+    try:
+      method_name = 'generate'
+      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)
+      method_info = method._cf_method_info
+      signature = method_info.signature
+      python_param_types = method_info.python_param_types
+
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
+                                                    python_param_types)
+      if len(inputs) == 1:
+        inputs = inputs[0]
+        for output in method(**inputs):
+          resp = service_pb2.MultiOutputResponse()
+          self._convert_output_to_proto(output, signature.output_fields, proto=resp.outputs.add())
+          resp.status.code = status_code_pb2.SUCCESS
+          yield resp
+      else:
+        for outputs in self._batch_generate(method, inputs):
+          resp = service_pb2.MultiOutputResponse()
+          for output in outputs:
+            self._convert_output_to_proto(
+                output, signature.output_fields, proto=resp.outputs.add())
+          resp.status.code = status_code_pb2.SUCCESS
+          yield resp
+    except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
+      logging.exception("Error in generate")
+      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
+          code=status_code_pb2.FAILURE,
+          details=str(e),
+          stack_trace=traceback.format_exc().split('\n')))
 
-  def stream_wrapper(self, request: service_pb2.PostModelOutputsRequest
+  def stream_wrapper(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
                     ) -> Iterator[service_pb2.MultiOutputResponse]:
-    """This method is used for input/output proto data conversion and yield outcome"""
-    return self.stream(request)
+    try:
+      request = next(request_iterator)  # get first request to determine method
+      assert len(request.inputs) == 1, "Streaming requires exactly one input"
 
-  @abstractmethod
-  def load_model(self):
-    raise NotImplementedError("load_model() not implemented")
-
-  @abstractmethod
-  def predict(self,
-              request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
-    raise NotImplementedError("run_input() not implemented")
-
-  @abstractmethod
-  def generate(self, request: service_pb2.PostModelOutputsRequest
-              ) -> Iterator[service_pb2.MultiOutputResponse]:
-    raise NotImplementedError("generate() not implemented")
-
-  @abstractmethod
-  def stream(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
-            ) -> Iterator[service_pb2.MultiOutputResponse]:
-    raise NotImplementedError("stream() not implemented")
+      method_name = 'generate'
+      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)
+      method_info = method._cf_method_info
+      signature = method_info.signature
+      python_param_types = method_info.python_param_types
+
+      # find the streaming vars in the signature
+      stream_sig = get_stream_from_signature(signature.input_fields)
+      if stream_sig is None:
+        raise ValueError("Streaming method must have a Stream input")
+      stream_argname = stream_sig.name
+
+      # convert all inputs for the first request, including the first stream value
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.input_fields,
+                                                    python_param_types)
+      kwargs = inputs[0]
+
+      # first streaming item
+      first_item = kwargs.pop(stream_argname)
+
+      # streaming generator
+      def InputStream():
+        yield first_item
+        # subsequent streaming items contain only the streaming input
+        for request in request_iterator:
+          item = self._convert_input_protos_to_python(request.inputs, [stream_sig],
+                                                      python_param_types)
+          item = item[0][stream_argname]
+          yield item
+
+      # add stream generator back to the input kwargs
+      kwargs[stream_argname] = InputStream()
+
+      for output in method(**kwargs):
+        resp = service_pb2.MultiOutputResponse()
+        self._convert_output_to_proto(output, signature.output_fields, proto=resp.outputs.add())
+        resp.status.code = status_code_pb2.SUCCESS
+        yield resp
+    except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
+      logging.exception("Error in stream")
+      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
+          code=status_code_pb2.FAILURE,
+          details=str(e),
+          stack_trace=traceback.format_exc().split('\n')))
+
+  def _convert_input_protos_to_python(self, inputs: List[resources_pb2.Input],
+                                      variables_signature: List[resources_pb2.ModelTypeField],
+                                      python_param_types) -> List[Dict[str, Any]]:
+    result = []
+    for input in inputs:
+      kwargs = deserialize(input.data, variables_signature)
+      # dynamic cast to annotated types
+      for k, v in kwargs.items():
+        if k not in python_param_types:
+          continue
+        kwargs[k] = data_types.cast(v, python_param_types[k])
+      result.append(kwargs)
+    return result
+
+  def _convert_output_to_proto(self,
+                               output: Any,
+                               variables_signature: List[resources_pb2.ModelTypeField],
+                               proto=None) -> resources_pb2.Output:
+    if proto is None:
+      proto = resources_pb2.Output()
+    serialize({'return': output}, variables_signature, proto.data, is_output=True)
+    proto.status.code = status_code_pb2.SUCCESS
+    return proto
+
+  @classmethod
+  def _register_model_methods(cls):
+    # go up the class hierarchy to find all decorated methods, and add to registry of current class
+    methods = {}
+    for base in reversed(cls.__mro__):
+      for name, method in base.__dict__.items():
+        method_info = getattr(method, _METHOD_INFO_ATTR, None)
+        if not method_info:  # regular function, not a model method
+          continue
+        methods[name] = method_info
+    # check for generic predict(request) -> response, etc. methods
+    #for name in ('predict', 'generate', 'stream'):
+    #  if hasattr(cls, name):
+    #    method = getattr(cls, name)
+    #    if not hasattr(method, _METHOD_INFO_ATTR):  # not already put in registry
+    #      methods[name] = _MethodInfo(method)
+    # set method table for this class in the registry
+    return methods
+
+  @classmethod
+  def _get_method_info(cls, func_name=None):
+    if not hasattr(cls, _METHOD_INFO_ATTR):
+      setattr(cls, _METHOD_INFO_ATTR, cls._register_model_methods())
+    method_info = getattr(cls, _METHOD_INFO_ATTR)
+    if func_name:
+      return method_info[func_name]
+    return method_info
+
+
+class _MethodInfo:
+
+  def __init__(self, method):
+    self.name = method.__name__
+    self.signature = build_function_signature(method)
+    self.python_param_types = {
+        p.name: p.annotation
+        for p in inspect.signature(method).parameters.values()
+        if p.annotation != inspect.Parameter.empty
+    }
+    self.python_param_types.pop('self', None)

clarifai/runners/models/model_class_refract.py

@@ -0,0 +1,80 @@
+import inspect
+from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Iterator, List, get_type_hints
+
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+
+from clarifai.runners.utils.data_handler import Output, kwargs_to_proto, proto_to_kwargs
+
+
+class ModelClass(ABC):
+
+  @abstractmethod
+  def load_model(self):
+    raise NotImplementedError("load_model() not implemented")
+
+  @abstractmethod
+  def predict(self, *args, **kwargs) -> Output:
+    raise NotImplementedError("predict() not implemented")
+
+  @abstractmethod
+  def generate(self, *args, **kwargs) -> Iterator[Output]:
+    raise NotImplementedError("generate() not implemented")
+
+  @abstractmethod
+  def stream(self, *args, **kwargs) -> Iterator[Output]:
+    raise NotImplementedError("stream() not implemented")
+
+  def batch_predict(self, inputs: List[Dict[str, Any]]) -> List[Output]:
+    with ThreadPoolExecutor() as executor:
+      return list(executor.map(lambda x: self.predict(**x), inputs))
+
+  def _process_request(self, request, process_func, is_stream=False):
+    inputs = self._convert_proto_to_python(request.inputs)
+    if len(inputs) == 1:
+      result = process_func(**inputs[0])
+      if is_stream:
+        return (self._convert_output_to_proto(output) for output in result)
+      else:
+        return [self._convert_output_to_proto(result)]
+    else:
+      results = self.batch_predict(inputs) if not is_stream else []
+      return [self._convert_output_to_proto(output) for output in results]
+
+  def predict_wrapper(
+      self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
+    outputs = self._process_request(request, self.predict)
+    return service_pb2.MultiOutputResponse(outputs=outputs)
+
+  def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
+                      ) -> Iterator[service_pb2.MultiOutputResponse]:
+    outputs = self._process_request(request, self.generate, is_stream=True)
+    for output in outputs:
+      yield service_pb2.MultiOutputResponse(outputs=[output])
+
+  def stream_wrapper(self, requests: Iterator[service_pb2.PostModelOutputsRequest]
+                    ) -> Iterator[service_pb2.MultiOutputResponse]:
+    for request in requests:
+      outputs = self._process_request(request, self.stream, is_stream=True)
+      yield service_pb2.MultiOutputResponse(outputs=outputs)
+
+  def _convert_proto_to_python(self, inputs: List[resources_pb2.Input]) -> List[Dict[str, Any]]:
+    get_type_hints(self.predict)
+    required_params = [
+        name for name, param in inspect.signature(self.predict).parameters.items()
+        if param.default == inspect.Parameter.empty
+    ]
+    kwargs_list = []
+    for input_proto in inputs:
+      kwargs = proto_to_kwargs(input_proto.data)
+      missing = [name for name in required_params if name not in kwargs]
+      if missing:
+        raise ValueError(f"Missing required parameters: {missing}")
+      kwargs_list.append(kwargs)
+    return kwargs_list
+
+  def _convert_output_to_proto(self, output: Any) -> resources_pb2.Output:
+    if isinstance(output, Output):
+      return output.to_proto()
+    return kwargs_to_proto(**output).outputs.add()

clarifai/runners/models/model_run_locally.py

@@ -7,14 +7,11 @@ import subprocess
 import sys
 import tempfile
 import time
-import traceback
 import venv
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
-from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 
 from clarifai.runners.models.model_builder import ModelBuilder
-from clarifai.runners.utils.url_fetcher import ensure_urls_downloaded
 from clarifai.utils.logging import logger
 
 
@@ -111,85 +108,13 @@ class ModelRunLocally:
     for i in range(1):
       yield request
 
-  def _run_model_inference(self, model):
-    """Perform inference using the model."""
-    request = self._build_request()
-    stream_request = self._build_stream_request()
-
-    ensure_urls_downloaded(request)
-    predict_response = None
-    generate_response = None
-    stream_response = None
-    try:
-      predict_response = model.predict(request)
-    except NotImplementedError:
-      logger.info("Model does not implement predict() method.")
-    except Exception as e:
-      logger.error(f"Model Prediction failed: {e}")
-      traceback.print_exc()
-      predict_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.MODEL_PREDICTION_FAILED,
-          description="Prediction failed",
-          details="",
-          internal_details=str(e),
-      ))
-
-    if predict_response:
-      if predict_response.outputs[0].status.code != status_code_pb2.SUCCESS:
-        logger.error(f"Moddel Prediction failed: {predict_response}")
-      else:
-        logger.info(f"Model Prediction succeeded: {predict_response}")
-
-    try:
-      generate_response = model.generate(request)
-    except NotImplementedError:
-      logger.info("Model does not implement generate() method.")
-    except Exception as e:
-      logger.error(f"Model Generation failed: {e}")
-      traceback.print_exc()
-      generate_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.MODEL_GENERATION_FAILED,
-          description="Generation failed",
-          details="",
-          internal_details=str(e),
-      ))
-
-    if generate_response:
-      generate_first_res = next(generate_response)
-      if generate_first_res.outputs[0].status.code != status_code_pb2.SUCCESS:
-        logger.error(f"Moddel Prediction failed: {generate_first_res}")
-      else:
-        logger.info(
-            f"Model Prediction succeeded for generate and first response: {generate_first_res}")
-
-    try:
-      stream_response = model.stream(stream_request)
-    except NotImplementedError:
-      logger.info("Model does not implement stream() method.")
-    except Exception as e:
-      logger.error(f"Model Stream failed: {e}")
-      traceback.print_exc()
-      stream_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
-          code=status_code_pb2.MODEL_STREAM_FAILED,
-          description="Stream failed",
-          details="",
-          internal_details=str(e),
-      ))
-
-    if stream_response:
-      stream_first_res = next(stream_response)
-      if stream_first_res.outputs[0].status.code != status_code_pb2.SUCCESS:
-        logger.error(f"Moddel Prediction failed: {stream_first_res}")
-      else:
-        logger.info(
-            f"Model Prediction succeeded for stream and first response: {stream_first_res}")
-
   def _run_test(self):
     """Test the model locally by making a prediction."""
     # Create the model
     model = self.builder.create_model_instance()
-    # send an inference.
-    self._run_model_inference(model)
+    # call its test method, if it has one
+    if hasattr(model, "test"):
+      model.test()
 
   def test_model(self):
     """Test the model by running it locally in the virtual environment."""