clarifai 11.3.0rc2__py3-none-any.whl → 11.4.0__py3-none-any.whl

clarifai/runners/models/model_runner.py

@@ -2,180 +2,188 @@ from typing import Iterator
 
 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 ..utils.url_fetcher import ensure_urls_downloaded
 
+from ..utils.url_fetcher import ensure_urls_downloaded
 from .model_class import ModelClass
 
 
 class ModelRunner(BaseRunner, HealthProbeRequestHandler):
-  """
-  This is a subclass of the runner class which will handle only the work items relevant to models.
-  """
-
-  def __init__(
-      self,
-      model: ModelClass,
-      runner_id: str,
-      nodepool_id: str,
-      compute_cluster_id: str,
-      user_id: str = None,
-      check_runner_exists: bool = True,
-      base_url: str = "https://api.clarifai.com",
-      pat: str = None,
-      token: str = None,
-      num_parallel_polls: int = 4,
-      **kwargs,
-  ) -> None:
-    super().__init__(
-        runner_id,
-        nodepool_id,
-        compute_cluster_id,
-        user_id,
-        check_runner_exists,
-        base_url,
-        pat,
-        token,
-        num_parallel_polls,
-        **kwargs,
-    )
-    self.model = model
-
-    # After model load successfully set the health probe to ready and startup
-    HealthProbeRequestHandler.is_ready = True
-    HealthProbeRequestHandler.is_startup = True
-
-  def get_runner_item_output_for_status(self,
-                                        status: status_pb2.Status) -> service_pb2.RunnerItemOutput:
     """
-    Set the error message in the RunnerItemOutput message subfield, used during exception handling
-    where we may only have a status to return.
-
-    Args:
-      status: status_pb2.Status - the status to return
-
-    Returns:
-      service_pb2.RunnerItemOutput - the RunnerItemOutput message with the status set
-    """
-    rio = service_pb2.RunnerItemOutput(
-        multi_output_response=service_pb2.MultiOutputResponse(status=status))
-    return rio
-
-  def runner_item_predict(self,
-                          runner_item: service_pb2.RunnerItem) -> service_pb2.RunnerItemOutput:
+    This is a subclass of the runner class which will handle only the work items relevant to models.
     """
-    Run the model on the given request. You shouldn't need to override this method, see run_input
-    for the implementation to process each input in the request.
 
-    Args:
-      request: service_pb2.PostModelOutputsRequest - the request to run the model on
-
-    Returns:
-      service_pb2.MultiOutputResponse - the response from the model's run_input implementation.
-    """
-
-    if not runner_item.HasField('post_model_outputs_request'):
-      raise Exception("Unexpected work item type: {}".format(runner_item))
-    request = runner_item.post_model_outputs_request
-    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(
-          code=status_code_pb2.SUCCESS,
-          description="Success",
-      )
-    elif any(successes):
-      status = status_pb2.Status(
-          code=status_code_pb2.MIXED_STATUS,
-          description="Mixed Status",
-      )
-    else:
-      status = status_pb2.Status(
-          code=status_code_pb2.FAILURE,
-          description="Failed",
-      )
-
-    resp.status.CopyFrom(status)
-    return service_pb2.RunnerItemOutput(multi_output_response=resp)
-
-  def runner_item_generate(
-      self, runner_item: service_pb2.RunnerItem) -> Iterator[service_pb2.RunnerItemOutput]:
-    # Call the generate() method the underlying model implements.
-
-    if not runner_item.HasField('post_model_outputs_request'):
-      raise Exception("Unexpected work item type: {}".format(runner_item))
-    request = runner_item.post_model_outputs_request
-    ensure_urls_downloaded(request)
-
-    for resp in self.model.generate_wrapper(request):
-      if resp.status.code != status_code_pb2.SUCCESS:
-        yield service_pb2.RunnerItemOutput(multi_output_response=resp)
-        continue
-      successes = []
-      for output in resp.outputs:
-        if not output.HasField('status') or not output.status.code:
-          raise Exception("Output must have a status code, please check the model implementation.")
-        successes.append(output.status.code == status_code_pb2.SUCCESS)
-      if all(successes):
-        status = status_pb2.Status(
-            code=status_code_pb2.SUCCESS,
-            description="Success",
-        )
-      elif any(successes):
-        status = status_pb2.Status(
-            code=status_code_pb2.MIXED_STATUS,
-            description="Mixed Status",
-        )
-      else:
-        status = status_pb2.Status(
-            code=status_code_pb2.FAILURE,
-            description="Failed",
-        )
-      resp.status.CopyFrom(status)
-
-      yield service_pb2.RunnerItemOutput(multi_output_response=resp)
-
-  def runner_item_stream(self, runner_item_iterator: Iterator[service_pb2.RunnerItem]
-                        ) -> Iterator[service_pb2.RunnerItemOutput]:
-    # Call the generate() method the underlying model implements.
-    for resp in self.model.stream_wrapper(pmo_iterator(runner_item_iterator)):
-      if resp.status.code != status_code_pb2.SUCCESS:
-        yield service_pb2.RunnerItemOutput(multi_output_response=resp)
-        continue
-      successes = []
-      for output in resp.outputs:
-        if not output.HasField('status') or not output.status.code:
-          raise Exception("Output must have a status code, please check the model implementation.")
-        successes.append(output.status.code == status_code_pb2.SUCCESS)
-      if all(successes):
-        status = status_pb2.Status(
-            code=status_code_pb2.SUCCESS,
-            description="Success",
-        )
-      elif any(successes):
-        status = status_pb2.Status(
-            code=status_code_pb2.MIXED_STATUS,
-            description="Mixed Status",
+    def __init__(
+        self,
+        model: ModelClass,
+        runner_id: str,
+        nodepool_id: str,
+        compute_cluster_id: str,
+        user_id: str = None,
+        check_runner_exists: bool = True,
+        base_url: str = "https://api.clarifai.com",
+        pat: str = None,
+        token: str = None,
+        num_parallel_polls: int = 4,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            runner_id,
+            nodepool_id,
+            compute_cluster_id,
+            user_id,
+            check_runner_exists,
+            base_url,
+            pat,
+            token,
+            num_parallel_polls,
+            **kwargs,
         )
-      else:
-        status = status_pb2.Status(
-            code=status_code_pb2.FAILURE,
-            description="Failed",
+        self.model = model
+
+        # After model load successfully set the health probe to ready and startup
+        HealthProbeRequestHandler.is_ready = True
+        HealthProbeRequestHandler.is_startup = True
+
+    def get_runner_item_output_for_status(
+        self, status: status_pb2.Status
+    ) -> service_pb2.RunnerItemOutput:
+        """
+        Set the error message in the RunnerItemOutput message subfield, used during exception handling
+        where we may only have a status to return.
+
+        Args:
+          status: status_pb2.Status - the status to return
+
+        Returns:
+          service_pb2.RunnerItemOutput - the RunnerItemOutput message with the status set
+        """
+        rio = service_pb2.RunnerItemOutput(
+            multi_output_response=service_pb2.MultiOutputResponse(status=status)
         )
-      resp.status.CopyFrom(status)
-
-      yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+        return rio
+
+    def runner_item_predict(
+        self, runner_item: service_pb2.RunnerItem
+    ) -> service_pb2.RunnerItemOutput:
+        """
+        Run the model on the given request. You shouldn't need to override this method, see run_input
+        for the implementation to process each input in the request.
+
+        Args:
+          request: service_pb2.PostModelOutputsRequest - the request to run the model on
+
+        Returns:
+          service_pb2.MultiOutputResponse - the response from the model's run_input implementation.
+        """
+
+        if not runner_item.HasField('post_model_outputs_request'):
+            raise Exception("Unexpected work item type: {}".format(runner_item))
+        request = runner_item.post_model_outputs_request
+        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(
+                code=status_code_pb2.SUCCESS,
+                description="Success",
+            )
+        elif any(successes):
+            status = status_pb2.Status(
+                code=status_code_pb2.MIXED_STATUS,
+                description="Mixed Status",
+            )
+        else:
+            status = status_pb2.Status(
+                code=status_code_pb2.FAILURE,
+                description="Failed",
+            )
+
+        resp.status.CopyFrom(status)
+        return service_pb2.RunnerItemOutput(multi_output_response=resp)
+
+    def runner_item_generate(
+        self, runner_item: service_pb2.RunnerItem
+    ) -> Iterator[service_pb2.RunnerItemOutput]:
+        # Call the generate() method the underlying model implements.
+
+        if not runner_item.HasField('post_model_outputs_request'):
+            raise Exception("Unexpected work item type: {}".format(runner_item))
+        request = runner_item.post_model_outputs_request
+        ensure_urls_downloaded(request)
+
+        for resp in self.model.generate_wrapper(request):
+            if resp.status.code != status_code_pb2.SUCCESS:
+                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                continue
+            successes = []
+            for output in resp.outputs:
+                if not output.HasField('status') or not output.status.code:
+                    raise Exception(
+                        "Output must have a status code, please check the model implementation."
+                    )
+                successes.append(output.status.code == status_code_pb2.SUCCESS)
+            if all(successes):
+                status = status_pb2.Status(
+                    code=status_code_pb2.SUCCESS,
+                    description="Success",
+                )
+            elif any(successes):
+                status = status_pb2.Status(
+                    code=status_code_pb2.MIXED_STATUS,
+                    description="Mixed Status",
+                )
+            else:
+                status = status_pb2.Status(
+                    code=status_code_pb2.FAILURE,
+                    description="Failed",
+                )
+            resp.status.CopyFrom(status)
+
+            yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+
+    def runner_item_stream(
+        self, runner_item_iterator: Iterator[service_pb2.RunnerItem]
+    ) -> Iterator[service_pb2.RunnerItemOutput]:
+        # Call the generate() method the underlying model implements.
+        for resp in self.model.stream_wrapper(pmo_iterator(runner_item_iterator)):
+            if resp.status.code != status_code_pb2.SUCCESS:
+                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                continue
+            successes = []
+            for output in resp.outputs:
+                if not output.HasField('status') or not output.status.code:
+                    raise Exception(
+                        "Output must have a status code, please check the model implementation."
+                    )
+                successes.append(output.status.code == status_code_pb2.SUCCESS)
+            if all(successes):
+                status = status_pb2.Status(
+                    code=status_code_pb2.SUCCESS,
+                    description="Success",
+                )
+            elif any(successes):
+                status = status_pb2.Status(
+                    code=status_code_pb2.MIXED_STATUS,
+                    description="Mixed Status",
+                )
+            else:
+                status = status_pb2.Status(
+                    code=status_code_pb2.FAILURE,
+                    description="Failed",
+                )
+            resp.status.CopyFrom(status)
+
+            yield service_pb2.RunnerItemOutput(multi_output_response=resp)
 
 
 def pmo_iterator(runner_item_iterator):
-  for runner_item in runner_item_iterator:
-    if not runner_item.HasField('post_model_outputs_request'):
-      raise Exception("Unexpected work item type: {}".format(runner_item))
-    ensure_urls_downloaded(runner_item.post_model_outputs_request)
-    yield runner_item.post_model_outputs_request
+    for runner_item in runner_item_iterator:
+        if not runner_item.HasField('post_model_outputs_request'):
+            raise Exception("Unexpected work item type: {}".format(runner_item))
+        ensure_urls_downloaded(runner_item.post_model_outputs_request)
+        yield runner_item.post_model_outputs_request

clarifai/runners/models/model_servicer.py

@@ -11,83 +11,91 @@ _RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() in ("true", "
 
 
 class ModelServicer(service_pb2_grpc.V2Servicer):
-  """
-  This is the servicer that will handle the gRPC requests from either the dev server or runner loop.
-  """
-
-  def __init__(self, model):
     """
-    Args:
-        model: The class that will handle the model logic. Must implement predict(),
-    generate(), stream().
+    This is the servicer that will handle the gRPC requests from either the dev server or runner loop.
     """
-    self.model = model
 
-  def PostModelOutputs(self, request: service_pb2.PostModelOutputsRequest,
-                       context=None) -> service_pb2.MultiOutputResponse:
-    """
-    This is the method that will be called when the servicer is run. It takes in an input and
-    returns an output.
-    """
+    def __init__(self, model):
+        """
+        Args:
+            model: The class that will handle the model logic. Must implement predict(),
+        generate(), stream().
+        """
+        self.model = model
 
-    # Download any urls that are not already bytes.
-    ensure_urls_downloaded(request)
+    def PostModelOutputs(
+        self, request: service_pb2.PostModelOutputsRequest, context=None
+    ) -> service_pb2.MultiOutputResponse:
+        """
+        This is the method that will be called when the servicer is run. It takes in an input and
+        returns an output.
+        """
 
-    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",
-          details="",
-          internal_details=str(e),
-      ))
+        # Download any urls that are not already bytes.
+        ensure_urls_downloaded(request)
 
-  def GenerateModelOutputs(self, request: service_pb2.PostModelOutputsRequest,
-                           context=None) -> Iterator[service_pb2.MultiOutputResponse]:
-    """
-    This is the method that will be called when the servicer is run. It takes in an input and
-    returns an output.
-    """
-    # Download any urls that are not already bytes.
-    ensure_urls_downloaded(request)
+        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",
+                    details="",
+                    internal_details=str(e),
+                )
+            )
 
-    try:
-      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",
-          details="",
-          internal_details=str(e),
-      ))
+    def GenerateModelOutputs(
+        self, request: service_pb2.PostModelOutputsRequest, context=None
+    ) -> Iterator[service_pb2.MultiOutputResponse]:
+        """
+        This is the method that will be called when the servicer is run. It takes in an input and
+        returns an output.
+        """
+        # Download any urls that are not already bytes.
+        ensure_urls_downloaded(request)
 
-  def StreamModelOutputs(self,
-                         request: Iterator[service_pb2.PostModelOutputsRequest],
-                         context=None) -> Iterator[service_pb2.MultiOutputResponse]:
-    """
-    This is the method that will be called when the servicer is run. It takes in an input and
-    returns an output.
-    """
-    # Duplicate the iterator
-    request, request_copy = tee(request)
+        try:
+            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",
+                    details="",
+                    internal_details=str(e),
+                )
+            )
+
+    def StreamModelOutputs(
+        self, request: Iterator[service_pb2.PostModelOutputsRequest], context=None
+    ) -> Iterator[service_pb2.MultiOutputResponse]:
+        """
+        This is the method that will be called when the servicer is run. It takes in an input and
+        returns an output.
+        """
+        # Duplicate the iterator
+        request, request_copy = tee(request)
 
-    # Download any urls that are not already bytes.
-    for req in request:
-      ensure_urls_downloaded(req)
+        # Download any urls that are not already bytes.
+        for req in request:
+            ensure_urls_downloaded(req)
 
-    try:
-      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",
-          details="",
-          internal_details=str(e),
-      ))
+        try:
+            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",
+                    details="",
+                    internal_details=str(e),
+                )
+            )