clarifai 11.2.1__py3-none-any.whl → 11.2.3__py3-none-any.whl

clarifai/cli/nodepool.py

@@ -1,31 +1,27 @@
 import click
+
 from clarifai.cli.base import cli
-from clarifai.client.compute_cluster import ComputeCluster
-from clarifai.utils.cli import display_co_resources, dump_yaml, from_yaml, validate_context
+from clarifai.utils.cli import (AliasedGroup, display_co_resources, dump_yaml, from_yaml,
+                                validate_context)
 
 
-@cli.group(['nodepool', 'np'])
+@cli.group(['nodepool', 'np'], cls=AliasedGroup)
 def nodepool():
   """Manage Nodepools: create, delete, list"""
-  pass
 
 
-@nodepool.command()
-@click.option(
-    '-cc_id',
-    '--compute_cluster_id',
-    required=False,
-    help='Compute Cluster ID for the compute cluster to interact with.')
+@nodepool.command(['c'])
+@click.argument('compute_cluster_id')
+@click.argument('nodepool_id')
 @click.option(
     '--config',
     type=click.Path(exists=True),
     required=True,
     help='Path to the nodepool config file.')
-@click.option(
-    '-np_id', '--nodepool_id', required=False, help='New Nodepool ID for the nodepool to create.')
 @click.pass_context
-def create(ctx, compute_cluster_id, config, nodepool_id):
+def create(ctx, compute_cluster_id, nodepool_id, config):
   """Create a new Nodepool with the given config file."""
+  from clarifai.client.compute_cluster import ComputeCluster
 
   validate_context(ctx)
   nodepool_config = from_yaml(config)
@@ -43,52 +39,74 @@ def create(ctx, compute_cluster_id, config, nodepool_id):
 
   compute_cluster = ComputeCluster(
       compute_cluster_id=compute_cluster_id,
-      user_id=ctx.obj['user_id'],
-      pat=ctx.obj['pat'],
-      base_url=ctx.obj['base_url'])
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
   if nodepool_id:
     compute_cluster.create_nodepool(config, nodepool_id=nodepool_id)
   else:
     compute_cluster.create_nodepool(config)
 
 
-@nodepool.command()
-@click.option(
-    '-cc_id',
-    '--compute_cluster_id',
-    required=True,
-    help='Compute Cluster ID for the compute cluster to interact with.')
+@nodepool.command(['ls'])
+@click.argument('compute_cluster_id', default="")
 @click.option('--page_no', required=False, help='Page number to list.', default=1)
-@click.option('--per_page', required=False, help='Number of items per page.', default=16)
+@click.option('--per_page', required=False, help='Number of items per page.', default=128)
 @click.pass_context
 def list(ctx, compute_cluster_id, page_no, per_page):
-  """List all nodepools for the user."""
+  """List all nodepools for the user across all compute clusters. If compute_cluster_id is provided
+  it will list only within that compute cluster. """
+  from clarifai.client.compute_cluster import ComputeCluster
+  from clarifai.client.user import User
 
   validate_context(ctx)
-  compute_cluster = ComputeCluster(
-      compute_cluster_id=compute_cluster_id,
-      user_id=ctx.obj['user_id'],
-      pat=ctx.obj['pat'],
-      base_url=ctx.obj['base_url'])
-  response = compute_cluster.list_nodepools(page_no, per_page)
-  display_co_resources(response, "Nodepool")
 
+  cc_id = compute_cluster_id
 
-@nodepool.command()
-@click.option(
-    '-cc_id',
-    '--compute_cluster_id',
-    required=True,
-    help='Compute Cluster ID for the compute cluster to interact with.')
-@click.option('-np_id', '--nodepool_id', help='Nodepool ID of the user to delete.')
+  if cc_id:
+    compute_cluster = ComputeCluster(
+        compute_cluster_id=cc_id,
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    response = compute_cluster.list_nodepools(page_no, per_page)
+  else:
+    user = User(
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    ccs = user.list_compute_clusters(page_no, per_page)
+    response = []
+    for cc in ccs:
+      compute_cluster = ComputeCluster(
+          compute_cluster_id=cc.id,
+          user_id=ctx.obj.current.user_id,
+          pat=ctx.obj.current.pat,
+          base_url=ctx.obj.current.api_base)
+      response.extend([i for i in compute_cluster.list_nodepools(page_no, per_page)])
+
+  display_co_resources(
+      response,
+      custom_columns={
+          'ID': lambda c: c.id,
+          'USER_ID': lambda c: c.compute_cluster.user_id,
+          'COMPUTE_CLUSTER_ID': lambda c: c.compute_cluster.id,
+          'DESCRIPTION': lambda c: c.description,
+      })
+
+
+@nodepool.command(['rm'])
+@click.argument('compute_cluster_id')
+@click.argument('nodepool_id')
 @click.pass_context
 def delete(ctx, compute_cluster_id, nodepool_id):
   """Deletes a nodepool for the user."""
+  from clarifai.client.compute_cluster import ComputeCluster
 
   validate_context(ctx)
   compute_cluster = ComputeCluster(
       compute_cluster_id=compute_cluster_id,
-      user_id=ctx.obj['user_id'],
-      pat=ctx.obj['pat'],
-      base_url=ctx.obj['base_url'])
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
   compute_cluster.delete_nodepools([nodepool_id])

clarifai/client/app.py

@@ -629,7 +629,7 @@ class App(Lister, BaseClient):
 
     Args:
         model_id (str): The model ID for the model to interact with.
-        model_version_id (str): The model version ID for the model version to interact with.
+        model_version (Dict): The model version ID for the model version to interact with.
 
     Returns:
         Model: A Model object for the existing model ID.

clarifai/client/auth/stub.py

@@ -1,5 +1,4 @@
 import itertools
-import logging
 import time
 from concurrent.futures import ThreadPoolExecutor
 
@@ -8,7 +7,7 @@ from clarifai_grpc.grpc.api.status import status_code_pb2
 
 from clarifai.client.auth.helper import ClarifaiAuthHelper
 from clarifai.client.auth.register import RpcCallable, V2Stub
-
+from clarifai.utils.logging import logger
 throttle_status_codes = {
     status_code_pb2.CONN_THROTTLED,
     status_code_pb2.CONN_EXCEED_HOURLY_LIMIT,
@@ -26,7 +25,7 @@ def validate_response(response, attempt, max_attempts):
   def handle_simple_response(response):
     if hasattr(response, 'status') and hasattr(response.status, 'code'):
       if (response.status.code in throttle_status_codes) and attempt < max_attempts:
-        logging.debug('Retrying with status %s' % str(response.status))
+        logger.debug('Retrying with status %s' % str(response.status))
         return None  # Indicates a retry is needed
       else:
         return response
@@ -42,7 +41,7 @@ def validate_response(response, attempt, max_attempts):
         return itertools.chain([validated_response], response)
       return None  # Indicates a retry is needed
     except grpc.RpcError as e:
-      logging.error('Error processing streaming response: %s' % str(e))
+      logger.error('Error processing streaming response: %s' % str(e))
       return None  # Indicates an error
   else:
     # Handle simple response validation
@@ -143,7 +142,7 @@ class _RetryRpcCallable(RpcCallable):
           return v
       except grpc.RpcError as e:
         if (e.code() in retry_codes_grpc) and attempt < self.max_attempts:
-          logging.debug('Retrying with status %s' % e.code())
+          logger.debug('Retrying with status %s' % e.code())
         else:
           raise
 

clarifai/client/dataset.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import time
 import uuid
@@ -354,7 +353,7 @@ class Dataset(Lister, BaseClient):
         break
       if failed_input_ids:
         retry_input_ids = [dataset_obj.all_input_ids[id] for id in failed_input_ids]
-        logging.warning(
+        logger.warning(
             f"Retrying upload for {len(failed_input_ids)} inputs in current batch: {retry_input_ids}\n"
         )
         failed_retrying_inputs, _, retry_response = self._upload_inputs_annotations(
@@ -494,7 +493,7 @@ class Dataset(Lister, BaseClient):
       add_file_handler(self.logger, f"Dataset_Upload{str(int(datetime.now().timestamp()))}.log")
 
     if retry_duplicates and duplicate_input_ids:
-      logging.warning(f"Retrying upload for {len(duplicate_input_ids)} duplicate inputs...\n")
+      logger.warning(f"Retrying upload for {len(duplicate_input_ids)} duplicate inputs...\n")
       duplicate_inputs_indexes = [input["Index"] for input in duplicate_input_ids]
       self.upload_dataset(
           dataloader=dataloader,
@@ -505,7 +504,7 @@ class Dataset(Lister, BaseClient):
 
     if failed_input_ids:
       #failed_inputs= ([input["Input_ID"] for input in failed_input_ids])
-      logging.warning(f"Retrying upload for {len(failed_input_ids)} failed inputs...\n")
+      logger.warning(f"Retrying upload for {len(failed_input_ids)} failed inputs...\n")
       failed_input_indexes = [input["Index"] for input in failed_input_ids]
       self.upload_dataset(
           dataloader=dataloader, log_retry_ids=failed_input_indexes, is_log_retry=True, **kwargs)

clarifai/runners/models/model_builder.py

@@ -142,17 +142,21 @@ class ModelBuilder:
   def _validate_config_checkpoints(self):
     """
     Validates the checkpoints section in the config file.
+      return loader_type, repo_id, hf_token, when, allowed_file_patterns, ignore_file_patterns
     :return: loader_type the type of loader or None if no checkpoints.
     :return: repo_id location of checkpoint.
     :return: hf_token token to access checkpoint.
+    :return: when one of ['upload', 'build', 'runtime'] to download checkpoint
+    :return: allowed_file_patterns patterns to allow in downloaded checkpoint
+    :return: ignore_file_patterns patterns to ignore in downloaded checkpoint
     """
     if "checkpoints" not in self.config:
-      return None, None, None, DEFAULT_DOWNLOAD_CHECKPOINT_WHEN
+      return None, None, None, DEFAULT_DOWNLOAD_CHECKPOINT_WHEN, None, None
     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
+      return None, None, None, DEFAULT_DOWNLOAD_CHECKPOINT_WHEN, None, None
     checkpoints = self.config.get("checkpoints")
     if 'when' not in checkpoints:
       logger.warn(
@@ -171,15 +175,30 @@ 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, when
+
+      allowed_file_patterns = self.config.get("checkpoints").get('allowed_file_patterns', None)
+      if isinstance(allowed_file_patterns, str):
+        allowed_file_patterns = [allowed_file_patterns]
+      ignore_file_patterns = self.config.get("checkpoints").get('ignore_file_patterns', None)
+      if isinstance(ignore_file_patterns, str):
+        ignore_file_patterns = [ignore_file_patterns]
+      return loader_type, repo_id, hf_token, when, allowed_file_patterns, ignore_file_patterns
 
   def _check_app_exists(self):
     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
+    if resp.status.code == status_code_pb2.CONN_KEY_INVALID:
+      logger.error(
+          f"Invalid PAT provided for user {self.client.user_app_id.user_id}. Please check your PAT and try again."
+      )
+      return False
     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}"
     )
+    logger.error(
+        f"App {self.client.user_app_id.app_id} not found for user {self.client.user_app_id.user_id}. Please create the app first and try again."
+    )
     return False
 
   def _validate_config_model(self):
@@ -200,9 +219,6 @@ class ModelBuilder:
     assert model.get('id') != "", "model_id cannot be empty in the config file"
 
     if not self._check_app_exists():
-      logger.error(
-          f"App {self.client.user_app_id.app_id} not found for user {self.client.user_app_id.user_id}"
-      )
       sys.exit(1)
 
   def _validate_config(self):
@@ -216,7 +232,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)
@@ -399,11 +415,12 @@ class ModelBuilder:
       # Sort in reverse so that newer cuda versions come first and are preferred.
       for image in sorted(AVAILABLE_TORCH_IMAGES, reverse=True):
         if torch_version in image and f'py{python_version}' in image:
-          cuda_version = image.split('-')[-1].replace('cuda', '')
+          # like cu124, rocm6.3, etc.
+          gpu_version = image.split('-')[-1]
           final_image = TORCH_BASE_IMAGE.format(
               torch_version=torch_version,
               python_version=python_version,
-              cuda_version=cuda_version,
+              gpu_version=gpu_version,
           )
           logger.info(f"Using Torch version {torch_version} base image to build the Docker image")
           break
@@ -480,8 +497,10 @@ class ModelBuilder:
     if not self.config.get("checkpoints"):
       logger.info("No checkpoints specified in the config file")
       return path
+    clarifai_model_type_id = self.config.get('model').get('model_type_id')
 
-    loader_type, repo_id, hf_token, when = self._validate_config_checkpoints()
+    loader_type, repo_id, hf_token, when, allowed_file_patterns, ignore_file_patterns = self._validate_config_checkpoints(
+    )
     if stage not in ["build", "upload", "runtime"]:
       raise Exception("Invalid stage provided, must be one of ['build', 'upload', 'runtime']")
     if when != stage:
@@ -490,14 +509,18 @@ class ModelBuilder:
       )
       return path
 
-    success = True
+    success = False
     if loader_type == "huggingface":
-      loader = HuggingFaceLoader(repo_id=repo_id, token=hf_token)
+      loader = HuggingFaceLoader(
+          repo_id=repo_id, token=hf_token, model_type_id=clarifai_model_type_id)
       # 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)
+      success = loader.download_checkpoints(
+          path,
+          allowed_file_patterns=allowed_file_patterns,
+          ignore_file_patterns=ignore_file_patterns)
 
     if loader_type:
       if not success:
@@ -569,7 +592,7 @@ class ModelBuilder:
     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()
+    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(
@@ -617,7 +640,7 @@ class ModelBuilder:
       # 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

clarifai/runners/models/model_run_locally.py

@@ -108,6 +108,7 @@ class ModelRunLocally:
 
   def _build_stream_request(self):
     request = self._build_request()
+    ensure_urls_downloaded(request)
     for i in range(1):
       yield request
 
@@ -479,7 +480,7 @@ def main(model_path,
   manager = ModelRunLocally(model_path)
   # get whatever stage is in config.yaml to force download now
   # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
-  _, _, _, when = manager.builder._validate_config_checkpoints()
+  _, _, _, when, _, _ = manager.builder._validate_config_checkpoints()
   manager.builder.download_checkpoints(
       stage=when, checkpoint_path_override=manager.builder.checkpoint_path)
   if inside_container:

clarifai/runners/utils/const.py

@@ -2,10 +2,10 @@ import os
 
 registry = os.environ.get('CLARIFAI_BASE_IMAGE_REGISTRY', 'public.ecr.aws/clarifai-models')
 
-GIT_SHA = "df565436eea93efb3e8d1eb558a0a46df29523ec"
+GIT_SHA = "b8ae56bf3b7c95e686ca002b07ca83d259c716eb"
 
 PYTHON_BASE_IMAGE = registry + '/python-base:{python_version}-' + GIT_SHA
-TORCH_BASE_IMAGE = registry + '/torch:{torch_version}-py{python_version}-cuda{cuda_version}-' + GIT_SHA
+TORCH_BASE_IMAGE = registry + '/torch:{torch_version}-py{python_version}-{gpu_version}-' + GIT_SHA
 
 # List of available python base images
 AVAILABLE_PYTHON_IMAGES = ['3.11', '3.12']
@@ -21,12 +21,13 @@ DEFAULT_RUNTIME_DOWNLOAD_PATH = os.path.join(os.sep, "tmp", ".cache")
 # List of available torch images
 # Keep sorted by most recent cuda version.
 AVAILABLE_TORCH_IMAGES = [
-    '2.4.1-py3.11-cuda124',
-    '2.5.1-py3.11-cuda124',
-    '2.4.1-py3.12-cuda124',
-    '2.5.1-py3.12-cuda124',
-    # '2.4.1-py3.13-cuda124',
-    # '2.5.1-py3.13-cuda124',
+    '2.4.1-py3.11-cu124',
+    '2.5.1-py3.11-cu124',
+    '2.4.1-py3.12-cu124',
+    '2.5.1-py3.12-cu124',
+    '2.6.0-py3.12-cu126',
+    '2.7.0-py3.12-cu128',
+    '2.7.0-py3.12-rocm6.3',
 ]
 CONCEPTS_REQUIRED_MODEL_TYPE = [
     'visual-classifier', 'visual-detector', 'visual-segmenter', 'text-classifier'

clarifai/runners/utils/loader.py

@@ -6,6 +6,7 @@ import shutil
 
 import requests
 
+from clarifai.runners.utils.const import CONCEPTS_REQUIRED_MODEL_TYPE
 from clarifai.utils.logging import logger
 
 
@@ -13,9 +14,10 @@ class HuggingFaceLoader:
 
   HF_DOWNLOAD_TEXT = "The 'huggingface_hub' package is not installed. Please install it using 'pip install huggingface_hub'."
 
-  def __init__(self, repo_id=None, token=None):
+  def __init__(self, repo_id=None, token=None, model_type_id=None):
     self.repo_id = repo_id
     self.token = token
+    self.clarifai_model_type_id = model_type_id
     if token:
       if self.validate_hftoken(token):
         try:
@@ -43,13 +45,17 @@ class HuggingFaceLoader:
           f"Error setting up Hugging Face token, please make sure you have the correct token: {e}")
       return False
 
-  def download_checkpoints(self, checkpoint_path: str):
+  def download_checkpoints(self,
+                           checkpoint_path: str,
+                           allowed_file_patterns=None,
+                           ignore_file_patterns=None):
     # throw error if huggingface_hub wasn't installed
     try:
       from huggingface_hub import snapshot_download
     except ImportError:
       raise ImportError(self.HF_DOWNLOAD_TEXT)
-    if os.path.exists(checkpoint_path) and self.validate_download(checkpoint_path):
+    if os.path.exists(checkpoint_path) and self.validate_download(
+        checkpoint_path, allowed_file_patterns, ignore_file_patterns):
       logger.info("Checkpoints already exist")
       return True
     else:
@@ -61,10 +67,16 @@ class HuggingFaceLoader:
           return False
 
         self.ignore_patterns = self._get_ignore_patterns()
+        if ignore_file_patterns:
+          if self.ignore_patterns:
+            self.ignore_patterns.extend(ignore_file_patterns)
+          else:
+            self.ignore_patterns = ignore_file_patterns
         snapshot_download(
             repo_id=self.repo_id,
             local_dir=checkpoint_path,
             local_dir_use_symlinks=False,
+            allow_patterns=allowed_file_patterns,
             ignore_patterns=self.ignore_patterns)
         # Remove the `.cache` folder if it exists
         cache_path = os.path.join(checkpoint_path, ".cache")
@@ -75,7 +87,8 @@ class HuggingFaceLoader:
         logger.error(f"Error downloading model checkpoints {e}")
         return False
       finally:
-        is_downloaded = self.validate_download(checkpoint_path)
+        is_downloaded = self.validate_download(checkpoint_path, allowed_file_patterns,
+                                               ignore_file_patterns)
         if not is_downloaded:
           logger.error("Error validating downloaded model checkpoints")
           return False
@@ -109,9 +122,13 @@ class HuggingFaceLoader:
       from huggingface_hub import file_exists, repo_exists
     except ImportError:
       raise ImportError(self.HF_DOWNLOAD_TEXT)
-    return repo_exists(self.repo_id) and file_exists(self.repo_id, 'config.json')
+    if self.clarifai_model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
+      return repo_exists(self.repo_id) and file_exists(self.repo_id, 'config.json')
+    else:
+      return repo_exists(self.repo_id)
 
-  def validate_download(self, checkpoint_path: str):
+  def validate_download(self, checkpoint_path: str, allowed_file_patterns: list,
+                        ignore_file_patterns: list):
     # check if model exists on HF
     try:
       from huggingface_hub import list_repo_files
@@ -120,7 +137,20 @@ class HuggingFaceLoader:
     # Get the list of files on the repo
     repo_files = list_repo_files(self.repo_id, token=self.token)
 
+    # Get the list of files on the repo that are allowed
+    if allowed_file_patterns:
+
+      def should_allow(file_path):
+        return any(fnmatch.fnmatch(file_path, pattern) for pattern in allowed_file_patterns)
+
+      repo_files = [f for f in repo_files if should_allow(f)]
+
     self.ignore_patterns = self._get_ignore_patterns()
+    if ignore_file_patterns:
+      if self.ignore_patterns:
+        self.ignore_patterns.extend(ignore_file_patterns)
+      else:
+        self.ignore_patterns = ignore_file_patterns
     # Get the list of files on the repo that are not ignored
     if getattr(self, "ignore_patterns", None):
       patterns = self.ignore_patterns