clarifai 11.1.5__py3-none-any.whl → 11.1.5rc6__py3-none-any.whl

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 == ".":
@@ -226,6 +234,15 @@ class ModelBuilder:
           )
           logger.info("Continuing without Hugging Face token")
 
+    num_threads = self.config.get("num_threads")
+    if num_threads or num_threads == 0:
+      assert isinstance(num_threads, int) and num_threads >= 1, ValueError(
+          f"`num_threads` must be an integer greater than or equal to 1. Received type {type(num_threads)} with value {num_threads}."
+      )
+    else:
+      num_threads = int(os.environ.get("CLARIFAI_NUM_THREADS", 1))
+      self.config["num_threads"] = num_threads
+
   @staticmethod
   def _get_tar_file_content_size(tar_file_path):
     """
@@ -244,6 +261,15 @@ 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.items()}
+    return signatures_to_yaml(signatures)
+
   @property
   def client(self):
     if self._client is None:

clarifai/runners/models/model_class.py

@@ -1,41 +1,286 @@
-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() == "true"
+
+
+class methods:
+  '''
+  Decorators to mark methods as predict, generate, or stream methods.
+  '''
+
+  @staticmethod
+  def predict(method):
+    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'predict'))
+    return method
+
+  @staticmethod
+  def generate(method):
+    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'generate'))
+    return method
+
+  @staticmethod
+  def stream(method):
+    setattr(method, _METHOD_INFO_ATTR, _MethodInfo(method, 'stream'))
+    return method
 
 
 class ModelClass(ABC):
+  '''
+  Base class for model classes that can be run as a service.
+
+  Define methods as predict, generate, or stream methods using the @methods decorators.
+
+  Example:
+
+    from clarifai.runners.model_class import ModelClass, methods
+    from clarifai.runners.utils.data_types import Input, Stream
+
+    class MyModel(ModelClass):
+
+      @methods.predict
+      def predict(self, x: str, y: int) -> List[str]:
+        return [x] * y
+
+      @methods.generate
+      def generate(self, x: str, y: int) -> Stream[str]:
+        for i in range(y):
+          yield x + str(i)
+
+      @methods.stream
+      def stream(self, input_stream: Stream[Input(x=str, y=int)]) -> Stream[str]:
+        for item in input_stream:
+          yield item.x + ' ' + str(item.y)
+  '''
+
+  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 = None
+      if len(request.inputs) > 0:
+        method_name = request.inputs[0].data.metadata.get('_method_name', None)
+      # call_params = dict(request.model.model_version.output_info.params)
+      # method_name = call_params.get('_method_name', 'predict')
+      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.inputs,
+                                                    python_param_types)
+      if len(inputs) == 1:
+        inputs = inputs[0]
+        output = method(**inputs)
+        outputs.append(self._convert_output_to_proto(output, signature.outputs))
+      else:
+        outputs = self.batch_predict(method, inputs)
+        outputs = [self._convert_output_to_proto(output, signature.outputs) 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:
+      call_params = dict(request.model.model_version.output_info.params)
+      method_name = call_params.get('_method_name', 'generate')
+      method = getattr(self, method_name)
+      method_info = method._cf_method_info
+      signature = method_info.signature
+      python_param_types = method_info.python_param_types
 
-  def stream_wrapper(self, request: service_pb2.PostModelOutputsRequest
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.inputs,
+                                                    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.outputs, 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.outputs, 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_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")
+      call_params = dict(request.model.model_version.output_info.params)
+      method_name = call_params.get('_method_name', 'stream')
+      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_argname, streaming_var_signatures = get_stream_from_signature(signature.inputs)
+
+      # convert all inputs for the first request, including the first stream value
+      inputs = self._convert_input_protos_to_python(request.inputs, signature.inputs,
+                                                    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, streaming_var_signatures,
+                                                      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.outputs, 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,
+                                      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,
+                               proto=None) -> resources_pb2.Output:
+    if proto is None:
+      proto = resources_pb2.Output()
+    if isinstance(output, tuple):
+      output = {f'return.{i}': item for i, item in enumerate(output)}
+    if not isinstance(output, dict):  # TODO Output type, not just dict
+      output = {'return': output}
+    serialize(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, method_type=name)
+    # 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, method_type):
+    self.name = method.__name__
+    self.signature = build_function_signature(method, method_type)
+    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_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."""

clarifai/runners/models/model_runner.py

@@ -82,6 +82,8 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
     ensure_urls_downloaded(request)
 
     resp = self.model.predict_wrapper(request)
+    if resp.status.code != status_code_pb2.SUCCESS:
+      return service_pb2.RunnerItemOutput(multi_output_response=resp)
     successes = [o.status.code == status_code_pb2.SUCCESS for o in resp.outputs]
     if all(successes):
       status = status_pb2.Status(

clarifai/runners/models/model_servicer.py

@@ -1,3 +1,4 @@
+import os
 from itertools import tee
 from typing import Iterator
 
@@ -6,6 +7,8 @@ from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 
 from ..utils.url_fetcher import ensure_urls_downloaded
 
+_RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() in ("true", "1")
+
 
 class ModelServicer(service_pb2_grpc.V2Servicer):
   """
@@ -33,6 +36,8 @@ class ModelServicer(service_pb2_grpc.V2Servicer):
     try:
       return self.model.predict_wrapper(request)
     except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
       return service_pb2.MultiOutputResponse(status=status_pb2.Status(
           code=status_code_pb2.MODEL_PREDICTION_FAILED,
           description="Failed",
@@ -50,8 +55,10 @@ class ModelServicer(service_pb2_grpc.V2Servicer):
     ensure_urls_downloaded(request)
 
     try:
-      return self.model.generate_wrapper(request)
+      yield from self.model.generate_wrapper(request)
     except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
       yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
           code=status_code_pb2.MODEL_PREDICTION_FAILED,
           description="Failed",
@@ -74,8 +81,10 @@ class ModelServicer(service_pb2_grpc.V2Servicer):
       ensure_urls_downloaded(req)
 
     try:
-      return self.model.stream_wrapper(request_copy)
+      yield from self.model.stream_wrapper(request_copy)
     except Exception as e:
+      if _RAISE_EXCEPTIONS:
+        raise
       yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
           code=status_code_pb2.MODEL_PREDICTION_FAILED,
           description="Failed",

clarifai/runners/server.py

@@ -85,6 +85,10 @@ def serve(model_path,
 
   model = builder.create_model_instance()
 
+  # `num_threads` can be set in config.yaml or via the environment variable CLARIFAI_NUM_THREADS="<integer>".
+  # Note: The value in config.yaml takes precedence over the environment variable.
+  num_threads = builder.config.get("num_threads")
+
   # Setup the grpc server for local development.
   if grpc:
 
@@ -115,7 +119,7 @@ def serve(model_path,
         nodepool_id=os.environ["CLARIFAI_NODEPOOL_ID"],
         compute_cluster_id=os.environ["CLARIFAI_COMPUTE_CLUSTER_ID"],
         base_url=os.environ.get("CLARIFAI_API_BASE", "https://api.clarifai.com"),
-        num_parallel_polls=int(os.environ.get("CLARIFAI_NUM_THREADS", 1)),
+        num_parallel_polls=num_threads,
     )
     runner.start()  # start the runner to fetch work from the API.