clarifai 11.1.2__py3-none-any.whl → 11.1.4rc1__py3-none-any.whl

clarifai/__init__.py

@@ -1 +1 @@
-__version__ = "11.1.2"
+__version__ = "11.1.4rc1"

clarifai/cli/__main__.py~

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

clarifai/cli/model.py

@@ -14,22 +14,16 @@ def model():
     type=click.Path(exists=True),
     required=True,
     help='Path to the model directory.')
-@click.option(
-    '--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.',
-)
 @click.option(
     '--skip_dockerfile',
     is_flag=True,
     help=
     'Flag to skip generating a dockerfile so that you can manually edit an already created dockerfile.',
 )
-def upload(model_path, download_checkpoints, skip_dockerfile):
+def upload(model_path, skip_dockerfile):
   """Upload a model to Clarifai."""
   from clarifai.runners.models.model_builder import upload_model
-  upload_model(model_path, download_checkpoints, skip_dockerfile)
+  upload_model(model_path, "upload", skip_dockerfile)
 
 
 @model.command()
@@ -46,12 +40,21 @@ def upload(model_path, download_checkpoints, skip_dockerfile):
     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..'
 )
-def download_checkpoints(model_path, out_path):
+@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, 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", "--stage", "build"]
 #####
 
 

clarifai/runners/models/model_builder.py

@@ -16,9 +16,9 @@ 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, 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
@@ -30,19 +30,6 @@ dependencies = [
     'torch',
     'clarifai',
 ]
-# Escape dependency names for regex
-dep_pattern = '|'.join(map(re.escape, dependencies))
-# All possible version specifiers
-version_specifiers = '==|>=|<=|!=|~=|>|<'
-# Compile a regex pattern with verbose mode for readability
-pattern = re.compile(r"""
-      ^\s*                                   # Start of line, optional whitespace
-      (?P<dependency>""" + dep_pattern + r""")   # Dependency name
-      \s*                                   # Optional whitespace
-      (?P<specifier>""" + version_specifiers + r""")?  # Optional version specifier
-      \s*                                   # Optional whitespace
-      (?P<version>[^\s;]+)?                 # Optional version (up to space or semicolon)
-      """, re.VERBOSE)
 
 
 def _clear_line(n: int = 1) -> None:
@@ -158,11 +145,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"
@@ -170,7 +171,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))
@@ -215,7 +216,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)
@@ -310,25 +311,41 @@ class ModelBuilder:
     )
     return self.client.STUB.PostModels(request)
 
+  def _match_req_line(self, line):
+    line = line.strip()
+    if not line or line.startswith('#'):
+      return None, None
+    # split on whitespace followed by #
+    line = re.split(r'\s+#', line)[0]
+    if "==" in line:
+      pkg, version = line.split("==")
+    elif ">=" in line:
+      pkg, version = line.split(">=")
+    elif ">" in line:
+      pkg, version = line.split(">")
+    elif "<=" in line:
+      pkg, version = line.split("<=")
+    elif "<" in line:
+      pkg, version = line.split("<")
+    else:
+      pkg, version = line, None  # No version specified
+    for dep in dependencies:
+      if dep == pkg:
+        if dep == 'torch' and line.find(
+            'whl/cpu') > 0:  # Ignore torch-cpu whl files, use base mage.
+          return None, None
+        return dep.strip(), version.strip() if version else None
+    return None, None
+
   def _parse_requirements(self):
     dependencies_version = {}
     with open(os.path.join(self.folder, 'requirements.txt'), 'r') as file:
       for line in file:
         # Skip empty lines and comments
-        line = line.strip()
-        if not line or line.startswith('#'):
+        dependency, version = self._match_req_line(line)
+        if dependency is None:
           continue
-        # split on whitespace followed by #
-        line = re.split(r'\s+#', line)[0]
-        match = pattern.match(line)
-        if match:
-          dependency = match.group('dependency')
-          version = match.group('version')
-          if dependency == "torch" and line.find(
-              'whl/cpu') > 0:  # Ignore torch-cpu whl files, use base mage.
-            continue
-
-          dependencies_version[dependency] = version if version else None
+        dependencies_version[dependency] = version if version else None
     return dependencies_version
 
   def create_dockerfile(self):
@@ -395,10 +412,8 @@ class ModelBuilder:
       with open(os.path.join(self.folder, 'requirements.txt'), 'r') as file:
         for line in file:
           # if the line without whitespace is "clarifai"
-          # split on whitespace followed by #
-          matchline = re.split(r'\s+#', line)[0]
-          match = pattern.match(matchline)
-          if match and match.group('dependency') == "clarifai":
+          dependency, version = self._match_req_line(line)
+          if dependency and dependency == "clarifai":
             lines.append(line.replace("clarifai", f"clarifai=={CLIENT_VERSION}"))
           else:
             lines.append(line)
@@ -433,10 +448,13 @@ class ModelBuilder:
   def tar_file(self):
     return f"{self.folder}.tar.gz"
 
-  def download_checkpoints(self, checkpoint_path_override: str = None):
+  def download_checkpoints(self, stage: str, checkpoint_path_override: str = None):
     """
     Downloads the checkpoints specified in the config file.
 
+    :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. 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.
@@ -445,7 +463,14 @@ class ModelBuilder:
       logger.info("No checkpoints specified in the config file")
       return True
 
-    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 True  # success seems fine here.
 
     success = True
     if loader_type == "huggingface":
@@ -519,11 +544,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(
@@ -533,15 +559,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)
 
@@ -556,7 +580,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
 
@@ -568,12 +592,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
@@ -701,10 +725,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()
@@ -716,4 +746,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

@@ -475,13 +475,14 @@ def main(model_path,
          keep_env=False,
          keep_image=False):
 
-  if not os.environ['CLARIFAI_PAT']:
+  if not os.environ.get("CLARIFAI_PAT", None):
     logger.error(
         "CLARIFAI_PAT environment variable is not set! Please set your PAT in the 'CLARIFAI_PAT' environment variable."
     )
     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)