move ensure_urls_downloaded into model class with enable flag

Author: David Eigen

clarifai/runners/models/base_typed_model.py

@@ -6,6 +6,9 @@
 from clarifai_grpc.grpc.api.service_pb2 import PostModelOutputsRequest
 from google.protobuf import json_format
 
+from clarifai.runners.utils.stream_utils import readahead
+from clarifai.runners.utils.url_fetcher import ensure_urls_downloaded
+
 from ..utils.data_handler import InputDataHandler, OutputDataHandler
 from .model_class import ModelClass
 
@@ -46,13 +49,17 @@ def convert_output_to_proto(self, outputs: list):
 
   def predict_wrapper(
       self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
+    if self.download_request_urls:
+      ensure_urls_downloaded(request)
     list_dict_input, inference_params = self.parse_input_request(request)
     outputs = self.predict(list_dict_input, inference_parameters=inference_params)
     return self.convert_output_to_proto(outputs)
 
   def generate_wrapper(
       self, request: PostModelOutputsRequest) -> Iterator[service_pb2.MultiOutputResponse]:
     list_dict_input, inference_params = self.parse_input_request(request)
+    if self.download_request_urls:
+      ensure_urls_downloaded(request)
     outputs = self.generate(list_dict_input, inference_parameters=inference_params)
     for output in outputs:
       yield self.convert_output_to_proto(output)
@@ -69,6 +76,8 @@ def stream_wrapper(self, request: Iterator[PostModelOutputsRequest]
     first_request = next(request)
     _, inference_params = self.parse_input_request(first_request)
     request_iterator = itertools.chain([first_request], request)
+    if self.download_request_urls:
+      request_iterator = readahead(map(ensure_urls_downloaded, request_iterator))
     outputs = self.stream(self._preprocess_stream(request_iterator), inference_params)
     for output in outputs:
       yield self.convert_output_to_proto(output)

clarifai/runners/models/model_class.py

@@ -3,23 +3,41 @@
 
 from clarifai_grpc.grpc.api import service_pb2
 
+from clarifai.runners.utils.stream_utils import readahead
+from clarifai.runners.utils.url_fetcher import ensure_urls_downloaded
+
 
 class ModelClass(ABC):
 
+  download_request_urls = True
+
   def predict_wrapper(
       self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
     """This method is used for input/output proto data conversion"""
+    # Download any urls that are not already bytes.
+    if self.download_request_urls:
+      ensure_urls_downloaded(request)
+
     return self.predict(request)
 
   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"""
+    # Download any urls that are not already bytes.
+    if self.download_request_urls:
+      ensure_urls_downloaded(request)
+
     return self.generate(request)
 
-  def stream_wrapper(self, request: service_pb2.PostModelOutputsRequest
+  def stream_wrapper(self, request_stream: 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)
+
+    # Download any urls that are not already bytes.
+    if self.download_request_urls:
+      request_stream = readahead(map(ensure_urls_downloaded, request_stream))
+
+    return self.stream(request_stream)
 
   @abstractmethod
   def load_model(self):

clarifai/runners/models/model_runner.py

@@ -5,7 +5,6 @@
 
 from clarifai_protocol import BaseRunner
 from clarifai_protocol.utils.health import HealthProbeRequestHandler
-from ..utils.url_fetcher import ensure_urls_downloaded
 
 from .model_class import ModelClass
 
@@ -79,7 +78,6 @@ def runner_item_predict(self,
     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)
     successes = [o.status.code == status_code_pb2.SUCCESS for o in resp.outputs]
@@ -109,7 +107,6 @@ def runner_item_generate(
     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):
       successes = []
@@ -169,5 +166,4 @@ 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

clarifai/runners/models/model_servicer.py

@@ -3,9 +3,6 @@
 from clarifai_grpc.grpc.api import service_pb2, service_pb2_grpc
 from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 
-from ..utils.stream_utils import readahead
-from ..utils.url_fetcher import ensure_urls_downloaded
-
 
 class ModelServicer(service_pb2_grpc.V2Servicer):
   """
@@ -27,9 +24,6 @@ def PostModelOutputs(self, request: service_pb2.PostModelOutputsRequest,
     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:
@@ -46,9 +40,6 @@ def GenerateModelOutputs(self, request: service_pb2.PostModelOutputsRequest,
     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.generate_wrapper(request)
     except Exception as e:
@@ -66,13 +57,8 @@ def StreamModelOutputs(self,
     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.
-    def _download_urls_stream(requests):
-      return readahead(map(ensure_urls_downloaded, requests))
-
     try:
-      return self.model_class.stream(_download_urls_stream(request))
+      return self.model_class.stream(request)
     except Exception as e:
       yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
           code=status_code_pb2.MODEL_PREDICTION_FAILED,