clarifai 11.1.3__py3-none-any.whl → 11.1.4rc2__py3-none-any.whl

clarifai/__init__.py

@@ -1 +1 @@
-__version__ = "11.1.3"
+__version__ = "11.1.4rc2"

clarifai/cli/__main__.py~

@@ -0,0 +1,4 @@
+from my_cli_tool.cli import main
+
+if __name__ == "__main__":
+  main()

clarifai/cli/model.py

@@ -18,7 +18,7 @@ def model():
     '--download_checkpoints',
     is_flag=True,
     help=
-    'Flag to download checkpoints before uploading and including them in the tar file that is uploaded. Defaults to False, which will attempt to download them at docker build time.',
+    'Flag to download checkpoints before uploading and including them in the tar file that is uploaded. Defaults to False, which will use the "when" field in config.yaml checkpoints section which can be "upload", "build", or "runtime".',
 )
 @click.option(
     '--skip_dockerfile',
@@ -29,7 +29,11 @@ def model():
 def upload(model_path, download_checkpoints, skip_dockerfile):
   """Upload a model to Clarifai."""
   from clarifai.runners.models.model_builder import upload_model
-  upload_model(model_path, download_checkpoints, skip_dockerfile)
+  if download_checkpoints:
+    stage = "any"  # ignore config.yaml and force download now.
+  else:
+    stage = "upload"
+  upload_model(model_path, stage, skip_dockerfile)
 
 
 @model.command()
@@ -44,14 +48,23 @@ def upload(model_path, download_checkpoints, skip_dockerfile):
     required=False,
     default=None,
     help=
-    'Option path to write the checkpoints to. This will place them in {out_path}/ If not provided it will default to {model_path}/1/checkpoints where the config.yaml is read..'
+    'Option path to write the checkpoints to. This will place them in {out_path}/1/checkpoints If not provided it will default to {model_path}/1/checkpoints where the config.yaml is read.'
+)
+@click.option(
+    '--stage',
+    type=str,
+    required=False,
+    default="build",
+    show_default=True,
+    help=
+    'The stage we are calling download checkpoints from. Typically this would be in the build stage which is the default. Other options include "runtime" to be used in load_model, "upload" to be used during model upload or "any" which will force download now regardless of config.yaml'
 )
-def download_checkpoints(model_path, out_path):
+def download_checkpoints(model_path, out_path, stage):
   """Download checkpoints from external source to local model_path"""
 
   from clarifai.runners.models.model_builder import ModelBuilder
   builder = ModelBuilder(model_path, download_validation_only=True)
-  builder.download_checkpoints(out_path)
+  builder.download_checkpoints(stage=stage, checkpoint_path_override=out_path)
 
 
 @model.command()

clarifai/runners/dockerfile_template/Dockerfile.debug

@@ -0,0 +1,11 @@
+FROM --platform=$TARGETPLATFORM public.ecr.aws/docker/library/python:3.12-slim-bookworm as pybase
+
+#############################
+# Final runtime image
+#############################
+FROM --platform=$TARGETPLATFORM ${RUNTIME_IMAGE} as final
+
+COPY --from=pybase --link=true /usr/bin/ls /usr/bin/cat  /usr/bin/which /usr/bin/bash /usr/bin/sort /usr/bin/du /usr/bin/
+COPY --from=pybase --link=true /bin/rbash /bin/sh /bin/rm /bin/
+COPY --from=pybase --link=true /lib/*-linux-gnu/libselinux.so.1 /lib/*-linux-gnu/libpcre2-8.so.0 /lib/x86_64-linux-gnu/
+COPY --from=pybase --link=true /lib/*-linux-gnu/libselinux.so.1 /lib/*-linux-gnu/libpcre2-8.so.0 /lib/aarch64-linux-gnu/

clarifai/runners/dockerfile_template/Dockerfile.debug~

@@ -0,0 +1,9 @@
+FROM --platform=$TARGETPLATFORM public.ecr.aws/docker/library/python:3.12-slim-bookworm as pybase
+
+#############################
+# Final runtime image
+#############################
+FROM --platform=$TARGETPLATFORM ${RUNTIME_IMAGE} as final
+
+COPY --from=pybase --link=true /usr/bin/ls /usr/bin/cat  /usr/bin/which /usr/bin/bash /usr/bin/sort /usr/bin/du /usr/bin/
+COPY --from=pybase --link=true /bin/rbash /bin/sh /bin/rm /bin/

clarifai/runners/dockerfile_template/Dockerfile.template

@@ -47,9 +47,9 @@ COPY --link=true --from=downloader /venv /venv
 COPY --chown=nonroot:nonroot downloader/unused.yaml /home/nonroot/main/1/checkpoints/.cache/unused.yaml
 
 #####
-# Download checkpoints
+# Download checkpoints if config.yaml has checkpoints.when = "build"
 COPY --link=true config.yaml /home/nonroot/main/
-RUN ["python", "-m", "clarifai.cli", "model", "download-checkpoints", "--model_path", "/home/nonroot/main", "--out_path", "/home/nonroot/main"]
+RUN ["python", "-m", "clarifai.cli", "model", "download-checkpoints", "--model_path", "/home/nonroot/main", "--out_path", "/home/nonroot/main/1/checkpoints", "--stage", "build"]
 #####
 
 

clarifai/runners/models/model_builder.py

@@ -16,9 +16,10 @@ 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_BASE_IMAGE, TORCH_BASE_IMAGE)
+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.urls.helper import ClarifaiUrlHelper
 from clarifai.utils.logging import logger
@@ -145,11 +146,25 @@ class ModelBuilder:
     :return: repo_id location of checkpoint.
     :return: hf_token token to access checkpoint.
     """
+    if "checkpoints" not in self.config:
+      return None, None, None, DEFAULT_DOWNLOAD_CHECKPOINT_WHEN
     assert "type" in self.config.get("checkpoints"), "No loader type specified in the config file"
     loader_type = self.config.get("checkpoints").get("type")
     if not loader_type:
       logger.info("No loader type specified in the config file for checkpoints")
       return None, None, None
+    checkpoints = self.config.get("checkpoints")
+    if 'when' not in checkpoints:
+      logger.warn(
+          f"No 'when' specified in the config file for checkpoints, defaulting to download at {DEFAULT_DOWNLOAD_CHECKPOINT_WHEN}"
+      )
+    when = checkpoints.get("when", DEFAULT_DOWNLOAD_CHECKPOINT_WHEN)
+    # In the config.yaml we don't allow "any", that's only used in download_checkpoints to force download.
+    assert when in [
+        "upload",
+        "build",
+        "runtime",
+    ], "Invalid value for when in the checkpoint loader when, needs to be one of ['upload', 'build', 'runtime']"
     assert loader_type == "huggingface", "Only huggingface loader supported for now"
     if loader_type == "huggingface":
       assert "repo_id" in self.config.get("checkpoints"), "No repo_id specified in the config file"
@@ -157,7 +172,7 @@ class ModelBuilder:
 
       # get from config.yaml otherwise fall back to HF_TOKEN env var.
       hf_token = self.config.get("checkpoints").get("hf_token", os.environ.get("HF_TOKEN", None))
-      return loader_type, repo_id, hf_token
+      return loader_type, repo_id, hf_token, when
 
   def _check_app_exists(self):
     resp = self.client.STUB.GetApp(service_pb2.GetAppRequest(user_app_id=self.client.user_app_id))
@@ -202,7 +217,7 @@ class ModelBuilder:
         assert model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE, f"Model type {model_type_id} not supported for concepts"
 
     if self.config.get("checkpoints"):
-      loader_type, _, hf_token = self._validate_config_checkpoints()
+      loader_type, _, hf_token, _ = self._validate_config_checkpoints()
 
       if loader_type == "huggingface" and hf_token:
         is_valid_token = HuggingFaceLoader.validate_hftoken(hf_token)
@@ -434,25 +449,43 @@ class ModelBuilder:
   def tar_file(self):
     return f"{self.folder}.tar.gz"
 
-  def download_checkpoints(self, checkpoint_path_override: str = None):
+  def default_runtime_checkpoint_path(self):
+    return DEFAULT_RUNTIME_DOWNLOAD_PATH
+
+  def download_checkpoints(self, stage: str, checkpoint_path_override: str = None):
     """
     Downloads the checkpoints specified in the config file.
 
-    :param checkpoint_path_override: The path to download the checkpoints to. If not provided, the
-    default path is used based on the folder ModelUploader was initialized with. The
-    checkpoint_suffix will be appended to the path.
+    :param stage: The stage of the build process. This is used to determine when to download the
+    checkpoints. The stage can be one of ['build', 'upload', 'runtime', 'any']. If "any" it will always try to download
+    regardless of what is specified in config.yaml. Otherwise it must match what is in config.yaml
+    :param checkpoint_path_override: The path to download the checkpoints to (with 1/checkpoints added as suffix). If not provided, the
+    default path is used based on the folder ModelUploader was initialized with. The checkpoint_suffix will be appended to the path.
+    If stage is 'runtime' and checkpoint_path_override is None, the default runtime path will be used.
+
+    :return: The path to the downloaded checkpoints. Even if it doesn't download anything, it will return the default path.
     """
+    path = self.checkpoint_path  # default checkpoint path.
     if not self.config.get("checkpoints"):
       logger.info("No checkpoints specified in the config file")
-      return True
+      return path
 
-    loader_type, repo_id, hf_token = self._validate_config_checkpoints()
+    loader_type, repo_id, hf_token, when = self._validate_config_checkpoints()
+    if stage not in ["build", "upload", "runtime", "any"]:
+      raise Exception("Invalid stage provided, must be one of ['build', 'upload', 'runtime']")
+    if when != stage and stage != "any":
+      logger.info(
+          f"Skipping downloading checkpoints for stage {stage} since config.yaml says to download them at stage {when}"
+      )
+      return path
 
     success = True
     if loader_type == "huggingface":
       loader = HuggingFaceLoader(repo_id=repo_id, token=hf_token)
-      path = self._checkpoint_path(
-          checkpoint_path_override) if checkpoint_path_override else self.checkpoint_path
+      # for runtime default to /tmp path
+      if stage == "runtime" and checkpoint_path_override is None:
+        checkpoint_path_override = self.default_runtime_checkpoint_path()
+      path = checkpoint_path_override if checkpoint_path_override else self.checkpoint_path
       success = loader.download_checkpoints(path)
 
     if loader_type:
@@ -461,7 +494,7 @@ class ModelBuilder:
         sys.exit(1)
       else:
         logger.info(f"Downloaded checkpoints for model {repo_id}")
-    return success
+    return path
 
   def _concepts_protos_from_concepts(self, concepts):
     concept_protos = []
@@ -520,11 +553,12 @@ class ModelBuilder:
             self._concepts_protos_from_concepts(labels))
     return model_version_proto
 
-  def upload_model_version(self, download_checkpoints):
+  def upload_model_version(self):
     file_path = f"{self.folder}.tar.gz"
     logger.debug(f"Will tar it into file: {file_path}")
 
     model_type_id = self.config.get('model').get('model_type_id')
+    loader_type, repo_id, hf_token, when = self._validate_config_checkpoints()
 
     if (model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE) and 'concepts' not in self.config:
       logger.info(
@@ -534,15 +568,13 @@ class ModelBuilder:
         logger.info(
             "Checkpoints specified in the config.yaml file, will download the HF model's config.json file to infer the concepts."
         )
-
-        if not download_checkpoints and not HuggingFaceLoader.validate_config(
-            self.checkpoint_path):
-
-          input(
-              "Press Enter to download the HuggingFace model's config.json file to infer the concepts and continue..."
-          )
-          loader_type, repo_id, hf_token = self._validate_config_checkpoints()
-          if loader_type == "huggingface":
+        # If we don't already have the concepts, download the config.json file from HuggingFace
+        if loader_type == "huggingface":
+          # If the config.yaml says we'll download in the future (build time or runtime) then we need to get this config now.
+          if when != "upload" and not HuggingFaceLoader.validate_config(self.checkpoint_path):
+            input(
+                "Press Enter to download the HuggingFace model's config.json file to infer the concepts and continue..."
+            )
             loader = HuggingFaceLoader(repo_id=repo_id, token=hf_token)
             loader.download_config(self.checkpoint_path)
 
@@ -557,7 +589,7 @@ class ModelBuilder:
     def filter_func(tarinfo):
       name = tarinfo.name
       exclude = [self.tar_file, "*~"]
-      if not download_checkpoints:
+      if when != "upload":
         exclude.append(self.checkpoint_suffix)
       return None if any(name.endswith(ex) for ex in exclude) else tarinfo
 
@@ -569,12 +601,12 @@ class ModelBuilder:
     logger.debug(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"):
+    if when != "upload" 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()
+        _, 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
@@ -702,10 +734,16 @@ class ModelBuilder:
         return False
 
 
-def upload_model(folder, download_checkpoints, skip_dockerfile):
+def upload_model(folder, stage, skip_dockerfile):
+  """
+  Uploads a model to Clarifai.
+
+  :param folder: The folder containing the model files.
+  :param stage: The stage of when you're uploading this model. This is used to determine when to download the checkpoints based on a match with the "when" field in the config.yaml checkpoints section or if you set stage to "any" it will always download the checkpoints.
+  :param skip_dockerfile: If True, skips creating the Dockerfile so you can re-use the local one.
+  """
   builder = ModelBuilder(folder)
-  if download_checkpoints:
-    builder.download_checkpoints()
+  builder.download_checkpoints(stage=stage)
   if not skip_dockerfile:
     builder.create_dockerfile()
   exists = builder.check_model_exists()
@@ -717,4 +755,4 @@ def upload_model(folder, download_checkpoints, skip_dockerfile):
     logger.info(f"New model will be created at {builder.model_url} with it's first version.")
 
   input("Press Enter to continue...")
-  builder.upload_model_version(download_checkpoints)
+  builder.upload_model_version()

clarifai/runners/models/model_run_locally.py

@@ -481,7 +481,8 @@ def main(model_path,
     )
     sys.exit(1)
   manager = ModelRunLocally(model_path)
-  manager.builder.download_checkpoints()
+  # stage="any" forces downloaded now regardless of config.yaml
+  manager.builder.download_checkpoints(stage="any")
   if inside_container:
     if not manager.is_docker_installed():
       sys.exit(1)