clarifai 11.2.3rc6__py3-none-any.whl → 11.2.3rc9__py3-none-any.whl

clarifai/cli/nodepool.py

@@ -1,32 +1,29 @@
 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
+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)
   if not compute_cluster_id:
     if 'compute_cluster' not in nodepool_config['nodepool']:
@@ -42,50 +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
 
-  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")
+  validate_context(ctx)
 
+  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/client/model_client.py

@@ -59,12 +59,41 @@ class ModelClient:
       Returns:
           Dict: The method signatures.
       '''
-    #request = resources_pb2.GetModelSignaturesRequest()
-    #response = self.stub.GetModelSignatures(request)
-    #self._method_signatures = json.loads(response.signatures)  # or define protos
-    # TODO this could use a new endpoint to get the signatures
-    # for local grpc models, we'll also have to add the endpoint to the model servicer
-    # for now we'll just use the predict endpoint with a special method name
+    try:
+      response = self.client.STUB.GetModelVersion(
+          service_pb2.GetModelVersionRequest(
+              user_app_id=self.request_template.user_app_id,
+              model_id= self.request_template.model_id,
+              version_id=self.request_template.model_version.id,
+          ))
+      method_signatures = None
+      if response.status.code == status_code_pb2.SUCCESS:
+        method_signatures = response.model_version.method_signatures
+      if response.status.code != status_code_pb2.SUCCESS:
+        raise Exception(f"Model failed with response {response!r}")
+      self._method_signatures= {}
+      for method_signature in method_signatures:
+        method_name = method_signature.name
+        # check for duplicate method names
+        if method_name in self._method_signatures:
+          raise ValueError(f"Duplicate method name {method_name}")
+        self._method_signatures[method_name] = method_signature
+      if not self._method_signatures: # if no method signatures, try to fetch from the model
+        self._fetch_signatures_backup()
+    except Exception as e:
+      logger.info(f"Failed to fetch method signatures from model: {e}")
+      # try to fetch from the model
+      self._fetch_signatures_backup()
+      if not self._method_signatures:
+        raise ValueError("Failed to fetch method signatures from model and backup method")
+
+  def _fetch_signatures_backup(self):
+    '''
+      This is a temporary method of fetching the method signatures from the model.
+
+      Returns:
+          Dict: The method signatures.
+      '''
 
     request = service_pb2.PostModelOutputsRequest()
     request.CopyFrom(self.request_template)

clarifai/runners/models/model_builder.py

@@ -192,17 +192,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(
@@ -221,15 +225,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):
@@ -250,9 +269,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):
@@ -266,7 +282,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)
@@ -282,7 +298,7 @@ class ModelBuilder:
           f"`num_threads` must be an integer greater than or equal to 1. Received type {type(num_threads)} with value {num_threads}."
       )
     else:
-      num_threads = int(os.environ.get("CLARIFAI_NUM_THREADS", 1))
+      num_threads = int(os.environ.get("CLARIFAI_NUM_THREADS", 16))
       self.config["num_threads"] = num_threads
 
   @staticmethod
@@ -354,8 +370,9 @@ class ModelBuilder:
 
     assert "model_type_id" in model, "model_type_id not found in the config file"
     assert "id" in model, "model_id not found in the config file"
-    assert "user_id" in model, "user_id not found in the config file"
-    assert "app_id" in model, "app_id not found in the config file"
+    if not self.download_validation_only:
+      assert "user_id" in model, "user_id not found in the config file"
+      assert "app_id" in model, "app_id not found in the config file"
 
     model_proto = json_format.ParseDict(model, resources_pb2.Model())
 
@@ -466,11 +483,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
@@ -547,8 +565,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:
@@ -557,14 +577,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:
@@ -598,53 +622,12 @@ class ModelBuilder:
     concepts = config.get('concepts')
     logger.info(f"Updated config.yaml with {len(concepts)} concepts.")
 
-  def filled_params_specs_with_inference_params(self, method_signatures: list[resources_pb2.MethodSignature]) -> list[resources_pb2.ModelTypeField]:
-    """
-    Fills the params_specs with the inference params.
-    """
-    inference_params = set()
-    for i, signature in enumerate(method_signatures):
-      for field in signature.input_fields:
-        if field.is_param:
-          if i==0:
-            inference_params.add(field.name)
-          else:
-            # if field.name not in inference_params then remove from inference_params
-            if field.name not in inference_params:
-              inference_params.remove(field.name)
-    output=[]
-    for signature in method_signatures:
-      for field in signature.input_fields:
-        if field.is_param and field.name in inference_params:
-          field.path = field.name
-          if field.type == resources_pb2.ModelTypeField.DataType.STR:
-            field.default_value= str(field.default)
-            field.field_type = resources_pb2.ModelTypeField.ModelTypeFieldType.STRING
-          elif field.type == resources_pb2.ModelTypeField.DataType.INT:
-            field.default_value= int(field.default)
-            field.field_type = resources_pb2.ModelTypeField.ModelTypeFieldType.NUMBER
-          elif field.type == resources_pb2.ModelTypeField.DataType.FLOAT:
-            field.default_value= float(field.default)
-            field.field_type = resources_pb2.ModelTypeField.ModelTypeFieldType.NUMBER
-          elif field.type == resources_pb2.ModelTypeField.DataType.BOOL:
-            field.default_value= bool(field.default)
-            field.field_type = resources_pb2.ModelTypeField.ModelTypeFieldType.BOOLEAN
-          else:
-            field.default_value= field.default
-            field.field_type = resources_pb2.ModelTypeField.ModelTypeFieldType.STRING
-          output.append(field)
-    return output
-
-
   def get_model_version_proto(self):
     signatures = self.get_method_signatures()
     model_version_proto = resources_pb2.ModelVersion(
         pretrained_model_config=resources_pb2.PretrainedModelConfig(),
         inference_compute_info=self.inference_compute_info,
         method_signatures=signatures,
-        # output_info= resources_pb2.OutputInfo(
-        #   params_specs=self.filled_params_specs_with_inference_params(signatures),
-        # )
     )
 
     model_type_id = self.config.get('model').get('model_type_id')
@@ -678,7 +661,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(
@@ -726,7 +709,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_class.py

@@ -30,7 +30,8 @@ class ModelClass(ABC):
   Example:
 
     from clarifai.runners.model_class import ModelClass
-    from clarifai.runners.utils.data_types import NamedFields, Stream
+    from clarifai.runners.utils.data_types import NamedFields
+    from typing import List, Iterator
 
     class MyModel(ModelClass):
 
@@ -39,12 +40,12 @@ class ModelClass(ABC):
         return [x] * y
 
       @ModelClass.method
-      def generate(self, x: str, y: int) -> Stream[str]:
+      def generate(self, x: str, y: int) -> Iterator[str]:
         for i in range(y):
           yield x + str(i)
 
       @ModelClass.method
-      def stream(self, input_stream: Stream[NamedFields(x=str, y=int)]) -> Stream[str]:
+      def stream(self, input_stream: Iterator[NamedFields(x=str, y=int)]) -> Iterator[str]:
         for item in input_stream:
           yield item.x + ' ' + str(item.y)
   '''
@@ -107,7 +108,8 @@ class ModelClass(ABC):
         is_convert = DataConverter.is_old_format(input.data)
         if is_convert:
           # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data, signature.input_fields)
+          new_data = DataConverter.convert_input_data_to_new_format(input.data,
+                                                                    signature.input_fields)
           input.data.CopyFrom(new_data)
       # convert inputs to python types
       inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
@@ -153,7 +155,8 @@ class ModelClass(ABC):
         is_convert = DataConverter.is_old_format(input.data)
         if is_convert:
           # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data, signature.input_fields)
+          new_data = DataConverter.convert_input_data_to_new_format(input.data,
+                                                                    signature.input_fields)
           input.data.CopyFrom(new_data)
       inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
                                                     signature.input_fields, python_param_types)
@@ -214,7 +217,8 @@ class ModelClass(ABC):
         is_convert = DataConverter.is_old_format(input.data)
         if is_convert:
           # convert to new format
-          new_data = DataConverter.convert_input_data_to_new_format(input.data, signature.input_fields)
+          new_data = DataConverter.convert_input_data_to_new_format(input.data,
+                                                                    signature.input_fields)
           input.data.CopyFrom(new_data)
       # convert all inputs for the first request, including the first stream value
       inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
@@ -267,8 +271,8 @@ class ModelClass(ABC):
         if k not in python_param_types:
           continue
 
-        if hasattr(python_param_types[k], "__args__") and getattr(
-            python_param_types[k], "__origin__", None) == data_types.Stream:
+        if hasattr(python_param_types[k], "__args__") and getattr(python_param_types[k],
+                                                                  "__origin__", None) == Iterator:
           # get the type of the items in the stream
           stream_type = python_param_types[k].__args__[0]
 

clarifai/runners/models/model_run_locally.py

@@ -11,7 +11,6 @@ import traceback
 import venv
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
-
 from clarifai.runners.models.model_builder import ModelBuilder
 from clarifai.utils.logging import logger
 
@@ -104,11 +103,6 @@ class ModelRunLocally:
         ],
     )
 
-  def _build_stream_request(self):
-    request = self._build_request()
-    for i in range(1):
-      yield request
-
   def _run_test(self):
     """Test the model locally by making a prediction."""
     # Create the model
@@ -404,39 +398,33 @@ def main(model_path,
          inside_container=False,
          port=8080,
          keep_env=False,
-         keep_image=False):
+         keep_image=False,
+         skip_dockerfile: bool = False):
 
-  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)
   # 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:
     if not manager.is_docker_installed():
       sys.exit(1)
-    manager.builder.create_dockerfile()
+    if not skip_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']
+    model_id = manager.config['model']['id'].lower()
+    # must be in lowercase
+    image_name = f"{model_id}:{image_tag}"
+    container_name = model_id
     if not manager.docker_image_exists(image_name):
       manager.build_docker_image(image_name=image_name)
     try:
-      envs = {
-          'CLARIFAI_PAT': os.environ['CLARIFAI_PAT'],
-          'CLARIFAI_API_BASE': os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
-      }
       if run_model_server:
         manager.run_docker_container(
-            image_name=image_name, container_name=container_name, port=port, env_vars=envs)
+            image_name=image_name, container_name=container_name, port=port)
       else:
-        manager.test_model_container(
-            image_name=image_name, container_name=container_name, env_vars=envs)
+        manager.test_model_container(image_name=image_name, container_name=container_name)
     finally:
       if manager.container_exists(container_name):
         manager.stop_docker_container(container_name)

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'