clarifai 11.0.7rc2__py3-none-any.whl → 11.1.0__py3-none-any.whl

clarifai/__init__.py

@@ -1 +1 @@
-__version__ = "11.0.7rc2"
+__version__ = "11.1.0"

clarifai/cli/model.py

@@ -28,9 +28,8 @@ def model():
 )
 def upload(model_path, download_checkpoints, skip_dockerfile):
   """Upload a model to Clarifai."""
-  from clarifai.runners.models import model_upload
-
-  model_upload.main(model_path, download_checkpoints, skip_dockerfile)
+  from clarifai.runners.models.model_builder import upload_model
+  upload_model(model_path, download_checkpoints, skip_dockerfile)
 
 
 @model.command()
@@ -50,9 +49,9 @@ def upload(model_path, download_checkpoints, skip_dockerfile):
 def download_checkpoints(model_path, out_path):
   """Download checkpoints from external source to local model_path"""
 
-  from clarifai.runners.models.model_upload import ModelUploader
-  uploader = ModelUploader(model_path, download_validation_only=True)
-  uploader.download_checkpoints(out_path)
+  from clarifai.runners.models.model_builder import ModelBuilder
+  builder = ModelBuilder(model_path, download_validation_only=True)
+  builder.download_checkpoints(out_path)
 
 
 @model.command()
@@ -173,8 +172,6 @@ def run_locally(model_path, port, mode, keep_env, keep_image):
 @click.option('--file_path', required=False, help='File path of file for the model to predict')
 @click.option('--url', required=False, help='URL to the file for the model to predict')
 @click.option('--bytes', required=False, help='Bytes to the file for the model to predict')
-@click.option(
-    '--input_id', required=False, help='Existing input id in the app for the model to predict')
 @click.option('--input_type', required=False, help='Type of input')
 @click.option(
     '-cc_id',
@@ -188,36 +185,28 @@ def run_locally(model_path, port, mode, keep_env, keep_image):
     '--inference_params', required=False, default='{}', help='Inference parameters to override')
 @click.option('--output_config', required=False, default='{}', help='Output config to override')
 @click.pass_context
-def predict(ctx, config, model_id, user_id, app_id, model_url, file_path, url, bytes, input_id,
-            input_type, compute_cluster_id, nodepool_id, deployment_id, inference_params,
-            output_config):
+def predict(ctx, config, model_id, user_id, app_id, model_url, file_path, url, bytes, input_type,
+            compute_cluster_id, nodepool_id, deployment_id, inference_params, output_config):
   """Predict using the given model"""
   import json
 
-  from clarifai.client.deployment import Deployment
-  from clarifai.client.input import Input
   from clarifai.client.model import Model
-  from clarifai.client.nodepool import Nodepool
   from clarifai.utils.cli import from_yaml
   if config:
     config = from_yaml(config)
-    model_id, user_id, app_id, model_url, file_path, url, bytes, input_id, input_type, compute_cluster_id, nodepool_id, deployment_id, inference_params, output_config = (
+    model_id, user_id, app_id, model_url, file_path, url, bytes, input_type, compute_cluster_id, nodepool_id, deployment_id, inference_params, output_config = (
         config.get(k, v)
         for k, v in [('model_id', model_id), ('user_id', user_id), ('app_id', app_id), (
             'model_url', model_url), ('file_path', file_path), ('url', url), ('bytes', bytes), (
-                'input_id',
-                input_id), ('input_type',
-                            input_type), ('compute_cluster_id',
-                                          compute_cluster_id), ('nodepool_id', nodepool_id), (
-                                              'deployment_id',
-                                              deployment_id), ('inference_params',
-                                                               inference_params), ('output_config',
-                                                                                   output_config)])
+                'input_type', input_type), ('compute_cluster_id', compute_cluster_id), (
+                    'nodepool_id',
+                    nodepool_id), ('deployment_id',
+                                   deployment_id), ('inference_params',
+                                                    inference_params), ('output_config',
+                                                                        output_config)])
   if sum([opt[1] for opt in [(model_id, 1), (user_id, 1), (app_id, 1), (model_url, 3)]
           if opt[0]]) != 3:
     raise ValueError("Either --model_id & --user_id & --app_id or --model_url must be provided.")
-  if sum([1 for opt in [file_path, url, bytes, input_id] if opt]) != 1:
-    raise ValueError("Exactly one of --file_path, --url, --bytes or --input_id must be provided.")
   if compute_cluster_id or nodepool_id or deployment_id:
     if sum([
         opt[1] for opt in [(compute_cluster_id, 0.5), (nodepool_id, 0.5), (deployment_id, 1)]
@@ -267,21 +256,5 @@ def predict(ctx, config, model_id, user_id, app_id, model_url, file_path, url, b
         nodepool_id=nodepool_id,
         deployment_id=deployment_id,
         inference_params=inference_params,
-        output_config=output_config)
-  elif input_id:
-    inputs = [Input.get_input(input_id)]
-    runner_selector = None
-    if deployment_id:
-      runner_selector = Deployment.get_runner_selector(
-          user_id=ctx.obj['user_id'], deployment_id=deployment_id)
-    elif compute_cluster_id and nodepool_id:
-      runner_selector = Nodepool.get_runner_selector(
-          user_id=ctx.obj['user_id'],
-          compute_cluster_id=compute_cluster_id,
-          nodepool_id=nodepool_id)
-    model_prediction = model.predict(
-        inputs=inputs,
-        runner_selector=runner_selector,
-        inference_params=inference_params,
-        output_config=output_config)
+        output_config=output_config)  ## TO DO: Add support for input_id
   click.echo(model_prediction)

clarifai/client/input.py

@@ -22,7 +22,7 @@ from clarifai.constants.dataset import MAX_RETRIES
 from clarifai.constants.input import MAX_UPLOAD_BATCH_SIZE
 from clarifai.errors import UserError
 from clarifai.utils.logging import logger
-from clarifai.utils.misc import BackoffIterator, Chunker
+from clarifai.utils.misc import BackoffIterator, Chunker, clean_input_id
 
 
 class Inputs(Lister, BaseClient):
@@ -282,7 +282,7 @@ class Inputs(Lister, BaseClient):
     for filename in os.listdir(folder_path):
       if filename.split('.')[-1] not in ['jpg', 'jpeg', 'png', 'tiff', 'webp']:
         continue
-      input_id = filename.split('.')[0]
+      input_id = clean_input_id(filename.split('.')[0])
       image_pb = resources_pb2.Image(base64=open(os.path.join(folder_path, filename), 'rb').read())
       input_protos.append(
           Inputs._get_proto(
@@ -473,7 +473,7 @@ class Inputs(Lister, BaseClient):
     for filename in os.listdir(folder_path):
       if filename.split('.')[-1] != 'txt':
         continue
-      input_id = filename.split('.')[0]
+      input_id = clean_input_id(filename.split('.')[0])
       text_pb = resources_pb2.Text(raw=open(os.path.join(folder_path, filename), 'rb').read())
       input_protos.append(
           Inputs._get_proto(
@@ -572,6 +572,24 @@ class Inputs(Lister, BaseClient):
 
     return input_mask_proto
 
+  def get_input(self, input_id: str) -> Input:
+    """Get Input object of input with input_id provided from the app.
+
+    Args:
+        input_id (str): The input ID for the annotation to get.
+
+    Returns:
+        Input: An Input object for the specified input ID.
+
+    Example:
+        >>> from clarifai.client.input import Inputs
+        >>> input_obj = Inputs(user_id = 'user_id', app_id = 'demo_app')
+        >>> input_obj.get_input(input_id='demo')
+    """
+    request = service_pb2.GetInputRequest(user_app_id=self.user_app_id, input_id=input_id)
+    response = self._grpc_request(self.STUB.GetInput, request)
+    return response.input
+
   def upload_from_url(self,
                       input_id: str,
                       image_url: str = None,

clarifai/runners/__init__.py

@@ -1,11 +1,11 @@
 from .models.base_typed_model import AnyAnyModel, TextInputModel, VisualInputModel
+from .models.model_builder import ModelBuilder
 from .models.model_runner import ModelRunner
-from .models.model_upload import ModelUploader
 from .utils.data_handler import InputDataHandler, OutputDataHandler
 
 __all__ = [
     "ModelRunner",
-    "ModelUploader",
+    "ModelBuilder",
     "InputDataHandler",
     "OutputDataHandler",
     "AnyAnyModel",

clarifai/runners/dockerfile_template/Dockerfile.template

@@ -17,7 +17,7 @@ RUN pip install --no-cache-dir -r /home/nonroot/requirements.txt && \
 FROM --platform=$TARGETPLATFORM ${DOWNLOADER_IMAGE} as downloader
 
 # make sure we have the latest clarifai package.
-RUN (pip install --upgrade --upgrade-strategy only-if-needed --no-cache-dir clarifai==11.0.7rc1 clarifai-grpc clarifai-protocol || true)
+RUN (pip install --upgrade --upgrade-strategy only-if-needed --no-cache-dir clarifai clarifai-grpc clarifai-protocol || true)
 #####
 
 
@@ -46,10 +46,6 @@ COPY --link=true --from=downloader /venv /venv
 # This creates the directory that HF downloader will populate and with nonroot:nonroot permissions up.
 COPY --chown=nonroot:nonroot downloader/unused.yaml /home/nonroot/main/1/checkpoints/.cache/unused.yaml
 
-# Set entryponit to our CLI
-# ENTRYPOINT ["python", "-m", "clarifai.cli"]
-# ENTRYPOINT ["python", "-m", "clarifai"]
-
 #####
 # Download checkpoints
 COPY --link=true config.yaml /home/nonroot/main/
@@ -66,7 +62,7 @@ COPY --link=true --from=pip_packages /venv /venv
 # for the actual model.
 # If checkpoints aren't downloaded since a checkpoints: block is not provided, then they will
 # be in the build context and copied here as well.
-COPY --link=true 1/model.py /home/nonroot/main/1/model.py
+COPY --link=true 1 /home/nonroot/main/1
 # At this point we only need these for validation in the SDK.
 COPY --link=true requirements.txt config.yaml /home/nonroot/main/
 

clarifai/runners/models/base_typed_model.py

@@ -7,10 +7,10 @@ from clarifai_grpc.grpc.api.service_pb2 import PostModelOutputsRequest
 from google.protobuf import json_format
 
 from ..utils.data_handler import InputDataHandler, OutputDataHandler
-from .model_runner import ModelRunner
+from .model_class import ModelClass
 
 
-class AnyAnyModel(ModelRunner):
+class AnyAnyModel(ModelClass):
 
   def load_model(self):
     """

clarifai/runners/models/{model_upload.py → model_builder.py}

@@ -1,3 +1,5 @@
+import importlib
+import inspect
 import os
 import re
 import sys
@@ -13,6 +15,7 @@ from rich import print
 from rich.markup import escape
 
 from clarifai.client 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_PYTHON_VERSION, PYTHON_BUILDER_IMAGE, PYTHON_RUNTIME_IMAGE, TORCH_BASE_IMAGE)
@@ -28,7 +31,8 @@ def _clear_line(n: int = 1) -> None:
     print(LINE_UP, end=LINE_CLEAR, flush=True)
 
 
-class ModelUploader:
+class ModelBuilder:
+  DEFAULT_CHECKPOINT_SIZE = 50 * 1024**3  # 50 GiB
 
   def __init__(self, folder: str, validate_api_ids: bool = True, download_validation_only=False):
     """
@@ -52,6 +56,55 @@ class ModelUploader:
     self.inference_compute_info = self._get_inference_compute_info()
     self.is_v3 = True  # Do model build for v3
 
+  def create_model_instance(self, load_model=True):
+    """
+    Create an instance of the model class, as specified in the config file.
+    """
+    # look for default model.py file location
+    for loc in ["model.py", "1/model.py"]:
+      model_file = os.path.join(self.folder, loc)
+      if os.path.exists(model_file):
+        break
+    if not os.path.exists(model_file):
+      raise Exception("Model file not found.")
+
+    module_name = os.path.basename(model_file).replace(".py", "")
+
+    spec = importlib.util.spec_from_file_location(module_name, model_file)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+
+    # Find all classes in the model.py file that are subclasses of ModelClass
+    classes = [
+        cls for _, cls in inspect.getmembers(module, inspect.isclass)
+        if issubclass(cls, ModelClass) and cls.__module__ == module.__name__
+    ]
+    #  Ensure there is exactly one subclass of BaseRunner in the model.py file
+    if len(classes) != 1:
+      # check for old inheritence structure, ModelRunner used to be a ModelClass
+      runner_classes = [
+          cls for _, cls in inspect.getmembers(module, inspect.isclass)
+          if cls.__module__ == module.__name__ and any(c.__name__ == 'ModelRunner'
+                                                       for c in cls.__bases__)
+      ]
+      if runner_classes and len(runner_classes) == 1:
+        raise Exception(
+            f'Could not determine model class.'
+            f' Models should now inherit from {ModelClass.__module__}.ModelClass, not ModelRunner.'
+            f' Please update your model "{runner_classes[0].__name__}" to inherit from ModelClass.'
+        )
+      raise Exception(
+          "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
+
   def _validate_folder(self, folder):
     if folder == ".":
       folder = ""  # will getcwd() next which ends with /
@@ -102,6 +155,9 @@ class ModelUploader:
     resp = self.client.STUB.GetApp(service_pb2.GetAppRequest(user_app_id=self.client.user_app_id))
     if resp.status.code == status_code_pb2.SUCCESS:
       return True
+    logger.error(
+        f"Error checking API {self._base_api} for user app {self.client.user_app_id.user_id}/{self.client.user_app_id.app_id}. Error code: {resp.status.code}"
+    )
     return False
 
   def _validate_config_model(self):
@@ -148,6 +204,24 @@ class ModelUploader:
           )
           logger.info("Continuing without Hugging Face token")
 
+  @staticmethod
+  def _get_tar_file_content_size(tar_file_path):
+    """
+    Calculates the total size of the contents of a tar file.
+
+    Args:
+      tar_file_path (str): The path to the tar file.
+
+    Returns:
+      int: The total size of the contents in bytes.
+    """
+    total_size = 0
+    with tarfile.open(tar_file_path, 'r') as tar:
+      for member in tar:
+        if member.isfile():
+          total_size += member.size
+    return total_size
+
   @property
   def client(self):
     if self._client is None:
@@ -159,9 +233,8 @@ class ModelUploader:
       user_id = model.get('user_id')
       app_id = model.get('app_id')
 
-      base = os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
-
-      self._client = BaseClient(user_id=user_id, app_id=app_id, base=base)
+      self._base_api = os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
+      self._client = BaseClient(user_id=user_id, app_id=app_id, base=self._base_api)
 
     return self._client
 
@@ -468,6 +541,18 @@ class ModelUploader:
     file_size = os.path.getsize(self.tar_file)
     logger.info(f"Size of the tar is: {file_size} bytes")
 
+    self.storage_request_size = self._get_tar_file_content_size(file_path)
+    if not download_checkpoints and self.config.get("checkpoints"):
+      # Get the checkpoint size to add to the storage request.
+      # First check for the env variable, then try querying huggingface. If all else fails, use the default.
+      checkpoint_size = os.environ.get('CHECKPOINT_SIZE_BYTES', 0)
+      if not checkpoint_size:
+        _, repo_id, _ = self._validate_config_checkpoints()
+        checkpoint_size = HuggingFaceLoader.get_huggingface_checkpoint_total_size(repo_id)
+      if not checkpoint_size:
+        checkpoint_size = self.DEFAULT_CHECKPOINT_SIZE
+      self.storage_request_size += checkpoint_size
+
     self.maybe_create_model()
     if not self.check_model_exists():
       logger.error(f"Failed to create model: {self.model_proto.id}")
@@ -542,6 +627,7 @@ class ModelUploader:
             model_id=self.model_proto.id,
             model_version=model_version_proto,
             total_size=file_size,
+            storage_request_size=self.storage_request_size,
             is_v3=self.is_v3,
         ))
     return result
@@ -589,19 +675,19 @@ class ModelUploader:
         return False
 
 
-def main(folder, download_checkpoints, skip_dockerfile):
-  uploader = ModelUploader(folder)
+def upload_model(folder, download_checkpoints, skip_dockerfile):
+  builder = ModelBuilder(folder)
   if download_checkpoints:
-    uploader.download_checkpoints()
+    builder.download_checkpoints()
   if not skip_dockerfile:
-    uploader.create_dockerfile()
-  exists = uploader.check_model_exists()
+    builder.create_dockerfile()
+  exists = builder.check_model_exists()
   if exists:
     logger.info(
-        f"Model already exists at {uploader.model_url}, this upload will create a new version for it."
+        f"Model already exists at {builder.model_url}, this upload will create a new version for it."
     )
   else:
-    logger.info(f"New model will be created at {uploader.model_url} with it's first version.")
+    logger.info(f"New model will be created at {builder.model_url} with it's first version.")
 
   input("Press Enter to continue...")
-  uploader.upload_model_version(download_checkpoints)
+  builder.upload_model_version(download_checkpoints)

clarifai/runners/models/model_run_locally.py

@@ -1,6 +1,4 @@
 import hashlib
-import importlib.util
-import inspect
 import os
 import platform
 import shutil
@@ -14,9 +12,8 @@ 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_protocol import BaseRunner
 
-from clarifai.runners.models.model_upload import ModelUploader
+from clarifai.runners.models.model_builder import ModelBuilder
 from clarifai.runners.utils.url_fetcher import ensure_urls_downloaded
 from clarifai.utils.logging import logger
 
@@ -27,9 +24,9 @@ class ModelRunLocally:
     self.model_path = model_path
     self.requirements_file = os.path.join(self.model_path, "requirements.txt")
 
-    # ModelUploader contains multiple useful methods to interact with the model
-    self.uploader = ModelUploader(self.model_path, download_validation_only=True)
-    self.config = self.uploader.config
+    # ModelBuilder contains multiple useful methods to interact with the model
+    self.builder = ModelBuilder(self.model_path, download_validation_only=True)
+    self.config = self.builder.config
 
   def _requirements_hash(self):
     """Generate a hash of the requirements file."""
@@ -91,38 +88,10 @@ class ModelRunLocally:
       self.clean_up()
       sys.exit(1)
 
-  def _get_model_runner(self):
-    """Dynamically import the runner class from the model file."""
-
-    # import the runner class that to be implement by the user
-    runner_path = os.path.join(self.model_path, "1", "model.py")
-
-    # arbitrary name given to the module to be imported
-    module = "runner_module"
-
-    spec = importlib.util.spec_from_file_location(module, runner_path)
-    runner_module = importlib.util.module_from_spec(spec)
-    sys.modules[module] = runner_module
-    spec.loader.exec_module(runner_module)
-
-    # Find all classes in the model.py file that are subclasses of BaseRunner
-    classes = [
-        cls for _, cls in inspect.getmembers(runner_module, inspect.isclass)
-        if issubclass(cls, BaseRunner) and cls.__module__ == runner_module.__name__
-    ]
-
-    #  Ensure there is exactly one subclass of BaseRunner in the model.py file
-    if len(classes) != 1:
-      raise Exception("Expected exactly one subclass of BaseRunner, found: {}".format(
-          len(classes)))
-
-    MyRunner = classes[0]
-    return MyRunner
-
   def _build_request(self):
     """Create a mock inference request for testing the model."""
 
-    model_version_proto = self.uploader.get_model_version_proto()
+    model_version_proto = self.builder.get_model_version_proto()
     model_version_proto.id = "model_version"
 
     return service_pb2.PostModelOutputsRequest(
@@ -142,8 +111,8 @@ class ModelRunLocally:
     for i in range(1):
       yield request
 
-  def _run_model_inference(self, runner):
-    """Perform inference using the runner."""
+  def _run_model_inference(self, model):
+    """Perform inference using the model."""
     request = self._build_request()
     stream_request = self._build_stream_request()
 
@@ -152,7 +121,7 @@ class ModelRunLocally:
     generate_response = None
     stream_response = None
     try:
-      predict_response = runner.predict(request)
+      predict_response = model.predict(request)
     except NotImplementedError:
       logger.info("Model does not implement predict() method.")
     except Exception as e:
@@ -172,7 +141,7 @@ class ModelRunLocally:
         logger.info(f"Model Prediction succeeded: {predict_response}")
 
     try:
-      generate_response = runner.generate(request)
+      generate_response = model.generate(request)
     except NotImplementedError:
       logger.info("Model does not implement generate() method.")
     except Exception as e:
@@ -194,7 +163,7 @@ class ModelRunLocally:
             f"Model Prediction succeeded for generate and first response: {generate_first_res}")
 
     try:
-      stream_response = runner.stream(stream_request)
+      stream_response = model.stream(stream_request)
     except NotImplementedError:
       logger.info("Model does not implement stream() method.")
     except Exception as e:
@@ -217,16 +186,10 @@ class ModelRunLocally:
 
   def _run_test(self):
     """Test the model locally by making a prediction."""
-    # construct MyRunner which will call load_model()
-    MyRunner = self._get_model_runner()
-    runner = MyRunner(
-        runner_id="n/a",
-        nodepool_id="n/a",
-        compute_cluster_id="n/a",
-        user_id="n/a",
-    )
+    # Create the model
+    model = self.builder.create_model_instance()
     # send an inference.
-    self._run_model_inference(runner)
+    self._run_model_inference(model)
 
   def test_model(self):
     """Test the model by running it locally in the virtual environment."""
@@ -274,7 +237,7 @@ class ModelRunLocally:
 
     command = [
         self.python_executable, "-m", "clarifai.runners.server", "--model_path", self.model_path,
-        "--start_dev_server", "--port",
+        "--grpc", "--port",
         str(port)
     ]
     try:
@@ -324,7 +287,9 @@ class ModelRunLocally:
       # Comment out the COPY instruction that copies the current folder
       modified_lines = []
       for line in lines:
-        if 'COPY .' in line and '/app/model_dir/main' in line:
+        if 'COPY' in line and '/home/nonroot/main' in line:
+          modified_lines.append(f'# {line}')
+        elif 'download-checkpoints' in line and '/home/nonroot/main' in line:
           modified_lines.append(f'# {line}')
         else:
           modified_lines.append(line)
@@ -375,7 +340,7 @@ class ModelRunLocally:
       if self._gpu_is_available():
         cmd.extend(["--gpus", "all"])
       # Add volume mappings
-      cmd.extend(["-v", f"{self.model_path}:/app/model_dir/main"])
+      cmd.extend(["-v", f"{self.model_path}:/home/nonroot/main"])
       # Add environment variables
       if env_vars:
         for key, value in env_vars.items():
@@ -383,9 +348,7 @@ class ModelRunLocally:
       # Add the image name
       cmd.append(image_name)
       # update the CMD to run the server
-      cmd.extend(
-          ["--model_path", "/app/model_dir/main", "--start_dev_server", "--port",
-           str(port)])
+      cmd.extend(["--model_path", "/home/nonroot/main", "--grpc", "--port", str(port)])
       # Run the container
       process = subprocess.Popen(cmd,)
       logger.info(
@@ -426,7 +389,7 @@ class ModelRunLocally:
       # update the entrypoint for testing the model
       cmd.extend(["--entrypoint", "python"])
       # Add volume mappings
-      cmd.extend(["-v", f"{self.model_path}:/app/model_dir/main"])
+      cmd.extend(["-v", f"{self.model_path}:/home/nonroot/main"])
       # Add environment variables
       if env_vars:
         for key, value in env_vars.items():
@@ -436,7 +399,7 @@ class ModelRunLocally:
       # update the CMD to test the model inside the container
       cmd.extend([
           "-c",
-          "from clarifai.runners.models.model_run_locally import ModelRunLocally; ModelRunLocally('/app/model_dir/main')._run_test()"
+          "from clarifai.runners.models.model_run_locally import ModelRunLocally; ModelRunLocally('/home/nonroot/main')._run_test()"
       ])
       # Run the container
       subprocess.check_call(cmd)
@@ -518,11 +481,11 @@ def main(model_path,
     )
     sys.exit(1)
   manager = ModelRunLocally(model_path)
-  manager.uploader.download_checkpoints()
+  manager.builder.download_checkpoints()
   if inside_container:
     if not manager.is_docker_installed():
       sys.exit(1)
-    manager.uploader.create_dockerfile()
+    manager.builder.create_dockerfile()
     image_tag = manager._docker_hash()
     image_name = f"{manager.config['model']['id']}:{image_tag}"
     container_name = manager.config['model']['id']

clarifai/runners/models/model_runner.py

@@ -10,16 +10,14 @@ from ..utils.url_fetcher import ensure_urls_downloaded
 from .model_class import ModelClass
 
 
-class ModelRunner(BaseRunner, ModelClass, HealthProbeRequestHandler):
+class ModelRunner(BaseRunner, HealthProbeRequestHandler):
   """
   This is a subclass of the runner class which will handle only the work items relevant to models.
-
-  It is also a subclass of ModelClass so that any subclass of ModelRunner will need to just
-  implement predict(), generate() and stream() methods and load_model() if needed.
   """
 
   def __init__(
       self,
+      model: ModelClass,
       runner_id: str,
       nodepool_id: str,
       compute_cluster_id: str,
@@ -43,7 +41,7 @@ class ModelRunner(BaseRunner, ModelClass, HealthProbeRequestHandler):
         num_parallel_polls,
         **kwargs,
     )
-    self.load_model()
+    self.model = model
 
     # After model load successfully set the health probe to ready and startup
     HealthProbeRequestHandler.is_ready = True
@@ -83,7 +81,7 @@ class ModelRunner(BaseRunner, ModelClass, HealthProbeRequestHandler):
     request = runner_item.post_model_outputs_request
     ensure_urls_downloaded(request)
 
-    resp = self.predict_wrapper(request)
+    resp = self.model.predict_wrapper(request)
     successes = [o.status.code == status_code_pb2.SUCCESS for o in resp.outputs]
     if all(successes):
       status = status_pb2.Status(
@@ -113,7 +111,7 @@ class ModelRunner(BaseRunner, ModelClass, HealthProbeRequestHandler):
     request = runner_item.post_model_outputs_request
     ensure_urls_downloaded(request)
 
-    for resp in self.generate_wrapper(request):
+    for resp in self.model.generate_wrapper(request):
       successes = []
       for output in resp.outputs:
         if not output.HasField('status') or not output.status.code:
@@ -141,7 +139,7 @@ class ModelRunner(BaseRunner, ModelClass, HealthProbeRequestHandler):
   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.stream_wrapper(pmo_iterator(runner_item_iterator)):
+    for resp in self.model.stream_wrapper(pmo_iterator(runner_item_iterator)):
       successes = []
       for output in resp.outputs:
         if not output.HasField('status') or not output.status.code: