clarifai 11.8.1__py3-none-any.whl → 11.8.3__py3-none-any.whl

clarifai/runners/models/model_builder.py

@@ -102,6 +102,7 @@ class ModelBuilder:
         self.folder = self._validate_folder(folder)
         self.config = self._load_config(os.path.join(self.folder, 'config.yaml'))
         self._validate_config()
+        self._validate_config_secrets()
         self._validate_stream_options()
         self.model_proto = self._get_model_proto()
         self.model_id = self.model_proto.id
@@ -465,6 +466,115 @@ class ModelBuilder:
                     "2) set_output_context"
                 )
 
+    def _validate_config_secrets(self):
+        """
+        Validate the secrets section in the config file.
+        """
+        if "secrets" not in self.config:
+            return
+
+        secrets = self.config.get("secrets", [])
+        if not isinstance(secrets, list):
+            raise ValueError("The 'secrets' field must be an array.")
+
+        for i, secret in enumerate(secrets):
+            if not isinstance(secret, dict):
+                raise ValueError(f"Secret at index {i} must be a dictionary.")
+
+            # Validate required fields
+            if "id" not in secret or not secret["id"]:
+                raise ValueError(f"Secret at index {i} must have a non-empty 'id' field.")
+
+            if "type" not in secret or not secret["type"]:
+                secret["type"] = "env"
+
+            if "env_var" not in secret or not secret["env_var"]:
+                raise ValueError(f"Secret at index {i} must have a non-empty 'env_var' field.")
+            # Validate secret type
+            if secret["type"] != "env":
+                raise ValueError(
+                    f"Secret at index {i} has invalid type '{secret['type']}'. Must be 'env'."
+                )
+
+        logger.info(f"Validated {len(secrets)} secrets in config file.")
+
+    def _process_secrets(self):
+        """
+        Process secrets from config file and create/validate them using the User client.
+        Returns the processed secrets array for inclusion in ModelVersion.OutputInfo.Params.
+        """
+        if "secrets" not in self.config:
+            return []
+
+        secrets = self.config.get("secrets", [])
+        if not secrets:
+            return []
+
+        # Get user client for secret operations
+        user = User(
+            user_id=self.config.get('model').get('user_id'),
+            pat=self.client.pat,
+            token=self.client.token,
+            base_url=self.client.base,
+        )
+
+        processed_secrets = []
+        secrets_to_create = []
+
+        for secret in secrets:
+            secret_id = secret["id"]
+            secret_type = secret.get("type", "env")
+            env_var = secret["env_var"]
+            secret_value = secret.get("value")  # Optional for existing secrets
+
+            # Check if secret already exists
+            try:
+                existing_secret = user.get_secret(secret_id)
+                logger.info(f"Secret '{secret_id}' already exists, using existing secret.")
+
+                # Add to processed secrets without the value
+                processed_secret = {
+                    "id": secret_id,
+                    "type": secret_type,
+                    "env_var": env_var,
+                }
+                processed_secrets.append(processed_secret)
+
+            except Exception:
+                # Secret doesn't exist, need to create it
+                if secret_value:
+                    logger.info(f"Secret '{secret_id}' does not exist, will create it.")
+                    secrets_to_create.append(
+                        {
+                            "id": secret_id,
+                            "value": secret_value,
+                            "description": secret.get("description", f"Secret for {env_var}"),
+                        }
+                    )
+
+                    # Add to processed secrets
+                    processed_secret = {
+                        "id": secret_id,
+                        "type": secret_type,
+                        "env_var": env_var,
+                    }
+                    processed_secrets.append(processed_secret)
+                else:
+                    raise ValueError(
+                        f"Secret '{secret_id}' does not exist and no value provided for creation."
+                    )
+
+        # Create new secrets if any
+        if secrets_to_create:
+            try:
+                created_secrets = user.create_secrets(secrets_to_create)
+                logger.info(f"Successfully created {len(created_secrets)} new secrets.")
+            except Exception as e:
+                logger.error(f"Failed to create secrets: {e}")
+                raise
+
+        return processed_secrets
+
     def _is_clarifai_internal(self):
         """
         Check if the current user is a Clarifai internal user based on email domain.
@@ -891,19 +1001,21 @@ class ModelBuilder:
                     )
                 torch_version = dependencies.get('torch', None)
                 if 'torch' in dependencies:
-                    if python_version != DEFAULT_PYTHON_VERSION:
-                        raise Exception(
-                            f"torch is not supported with Python version {python_version}, please use Python version {DEFAULT_PYTHON_VERSION} in your config.yaml"
-                        )
                     if not torch_version:
                         logger.info(
                             f"Setup: torch version not found in requirements.txt, using the default version {DEFAULT_AMD_TORCH_VERSION}"
                         )
                         torch_version = DEFAULT_AMD_TORCH_VERSION
-                    if torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
-                        raise Exception(
-                            f"torch version {torch_version} not supported, please use one of the following versions: {DEFAULT_AMD_TORCH_VERSION} in your requirements.txt"
-                        )
+                    elif torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
+                        # Currently, we have only one vLLM image built with the DEFAULT_AMD_TORCH_VERSION.
+                        # If the user requests a different PyTorch version, that specific version will be
+                        # installed during the requirements.txt installation step
+                        torch_version = DEFAULT_AMD_TORCH_VERSION
+                else:
+                    logger.info(
+                        f"`torch` not found in requirements.txt, using the default torch=={DEFAULT_AMD_TORCH_VERSION}"
+                    )
+                    torch_version = DEFAULT_AMD_TORCH_VERSION
                 python_version = DEFAULT_PYTHON_VERSION
                 gpu_version = DEFAULT_AMD_GPU_VERSION
                 final_image = AMD_VLLM_BASE_IMAGE.format(
@@ -912,21 +1024,17 @@ class ModelBuilder:
                     gpu_version=gpu_version,
                 )
                 logger.info("Setup: Using vLLM base image to build the Docker image")
-            elif 'torch' in dependencies:
+            elif (
+                'torch' in dependencies
+                and (dependencies['torch'] in [None, DEFAULT_AMD_TORCH_VERSION])
+                and python_version == DEFAULT_PYTHON_VERSION
+            ):
                 torch_version = dependencies['torch']
-                if python_version != DEFAULT_PYTHON_VERSION:
-                    raise Exception(
-                        f"torch is not supported with Python version {python_version}, please use Python version {DEFAULT_PYTHON_VERSION} in your config.yaml"
-                    )
                 if not torch_version:
                     logger.info(
                         f"torch version not found in requirements.txt, using the default version {DEFAULT_AMD_TORCH_VERSION}"
                     )
                     torch_version = DEFAULT_AMD_TORCH_VERSION
-                if torch_version not in [DEFAULT_AMD_TORCH_VERSION]:
-                    raise Exception(
-                        f"torch version {torch_version} not supported, please use one of the following versions: {DEFAULT_AMD_TORCH_VERSION} in your requirements.txt"
-                    )
                 python_version = DEFAULT_PYTHON_VERSION
                 gpu_version = DEFAULT_AMD_GPU_VERSION
                 final_image = AMD_TORCH_BASE_IMAGE.format(
@@ -1258,6 +1366,29 @@ class ModelBuilder:
             metadata_struct.update({'git_registry': git_info})
             model_version_proto.metadata.CopyFrom(metadata_struct)
 
+        # Process and add secrets to output_info.params
+        try:
+            processed_secrets = self._process_secrets()
+            if processed_secrets:
+                # Initialize output_info.params if not already present
+                if not model_version_proto.HasField("output_info"):
+                    model_version_proto.output_info.CopyFrom(resources_pb2.OutputInfo())
+
+                # Initialize params if not already present
+                if not model_version_proto.output_info.HasField("params"):
+                    from google.protobuf.struct_pb2 import Struct
+
+                    model_version_proto.output_info.params.CopyFrom(Struct())
+
+                # Add secrets to params
+                model_version_proto.output_info.params.update({"secrets": processed_secrets})
+                logger.info(
+                    f"Added {len(processed_secrets)} secrets to model version output_info.params"
+                )
+        except Exception as e:
+            logger.error(f"Failed to process secrets: {e}")
+            raise
+
         model_type_id = self.config.get('model').get('model_type_id')
         if model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
             if 'concepts' in self.config:
@@ -1382,6 +1513,7 @@ class ModelBuilder:
                     user_id=self.client.user_app_id.user_id,
                     app_id=self.client.user_app_id.app_id,
                     model_id=self.model_proto.id,
+                    colorize=True,
                 )
                 logger.info("""\n
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

clarifai/runners/pipeline_steps/pipeline_step_builder.py

@@ -4,6 +4,7 @@ import sys
 import tarfile
 import time
 from string import Template
+from typing import List, Optional
 
 import yaml
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -11,6 +12,7 @@ from clarifai_grpc.grpc.api.status import status_code_pb2
 from google.protobuf import json_format
 
 from clarifai.client.base import BaseClient
+from clarifai.utils.hashing import hash_directory
 from clarifai.utils.logging import logger
 from clarifai.utils.misc import get_uuid
 from clarifai.versions import CLIENT_VERSION
@@ -22,12 +24,13 @@ UPLOAD_CHUNK_SIZE = 14 * 1024 * 1024
 class PipelineStepBuilder:
     """Pipeline Step Builder class for managing pipeline step upload to Clarifai."""
 
-    def __init__(self, folder: str):
+    def __init__(self, folder: str, hash_exclusions: Optional[List[str]] = None):
         """
         Initialize PipelineStepBuilder.
 
         :param folder: The folder containing the pipeline step files (config.yaml, requirements.txt,
                       dockerfile, and pipeline_step.py in 1/ subdirectory)
+        :param hash_exclusions: List of file names to exclude from hash calculation (defaults to ['config-lock.yaml'])
         """
         self._client = None
         self.folder = self._validate_folder(folder)
@@ -37,6 +40,10 @@ class PipelineStepBuilder:
         self.pipeline_step_id = self.pipeline_step_proto.id
         self.pipeline_step_version_id = None
         self.pipeline_step_compute_info = self._get_pipeline_step_compute_info()
+        # Configure files to exclude from hash calculation
+        self.hash_exclusions = (
+            hash_exclusions if hash_exclusions is not None else ['config-lock.yaml']
+        )
 
     @property
     def client(self):
@@ -490,6 +497,95 @@ COPY --link=true requirements.txt config.yaml /home/nonroot/main/
 
         raise TimeoutError("Pipeline step build did not finish in time")
 
+    def load_config_lock(self):
+        """
+        Load existing config-lock.yaml if it exists.
+
+        :return: Dictionary with config-lock data or None if file doesn't exist
+        """
+        config_lock_path = os.path.join(self.folder, "config-lock.yaml")
+        if os.path.exists(config_lock_path):
+            try:
+                with open(config_lock_path, 'r', encoding='utf-8') as f:
+                    return yaml.safe_load(f)
+            except Exception as e:
+                logger.warning(f"Failed to load config-lock.yaml: {e}")
+                return None
+        return None
+
+    def should_upload_step(self, algo="md5"):
+        """
+        Check if the pipeline step should be uploaded based on hash comparison.
+
+        :param algo: Hash algorithm to use
+        :return: True if step should be uploaded, False otherwise
+        """
+        config_lock = self.load_config_lock()
+
+        # If no config-lock.yaml exists, upload the step (first time upload)
+        if config_lock is None:
+            logger.info("No config-lock.yaml found, will upload pipeline step")
+            return True
+
+        # Compare stored hash with freshly computed one
+        current_hash = hash_directory(self.folder, algo=algo, exclude_files=self.hash_exclusions)
+        stored_hash_info = config_lock.get("hash", {})
+        stored_hash = stored_hash_info.get("value", "")
+        stored_algo = stored_hash_info.get("algo", "md5")
+
+        # If algorithm changed, re-upload to update hash
+        if stored_algo != algo:
+            logger.info(
+                f"Hash algorithm changed from {stored_algo} to {algo}, will upload pipeline step"
+            )
+            return True
+
+        # If hash changed, upload
+        if current_hash != stored_hash:
+            logger.info(
+                f"Hash changed (was: {stored_hash}, now: {current_hash}), will upload pipeline step"
+            )
+            return True
+
+        logger.info(f"Hash unchanged ({current_hash}), skipping pipeline step upload")
+        return False
+
+    def generate_config_lock(self, version_id, algo="md5"):
+        """
+        Generate config-lock.yaml content for the pipeline step.
+
+        :param version_id: Pipeline step version ID
+        :param algo: Hash algorithm used
+        :return: Dictionary with config-lock data
+        """
+        # Compute hash
+        hash_value = hash_directory(self.folder, algo=algo, exclude_files=self.hash_exclusions)
+
+        # Create config-lock structure
+        config_lock = {"id": version_id, "hash": {"algo": algo, "value": hash_value}}
+
+        # Append the original config.yaml contents
+        config_lock.update(self.config)
+
+        return config_lock
+
+    def save_config_lock(self, version_id, algo="md5"):
+        """
+        Save config-lock.yaml file with pipeline step metadata.
+
+        :param version_id: Pipeline step version ID
+        :param algo: Hash algorithm used
+        """
+        config_lock_data = self.generate_config_lock(version_id, algo)
+        config_lock_path = os.path.join(self.folder, "config-lock.yaml")
+
+        try:
+            with open(config_lock_path, 'w', encoding='utf-8') as f:
+                yaml.dump(config_lock_data, f, default_flow_style=False, allow_unicode=True)
+            logger.info(f"Generated config-lock.yaml at {config_lock_path}")
+        except Exception as e:
+            logger.error(f"Failed to save config-lock.yaml: {e}")
+
 
 def upload_pipeline_step(folder, skip_dockerfile=False):
     """