clarifai 11.7.4__py3-none-any.whl → 11.7.5rc1__py3-none-any.whl

clarifai/runners/models/model_builder.py

@@ -10,7 +10,7 @@ import tarfile
 import time
 import webbrowser
 from string import Template
-from typing import Literal
+from typing import Any, Dict, Literal, Optional
 from unittest.mock import MagicMock
 
 import yaml
@@ -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.
@@ -1107,7 +1217,7 @@ class ModelBuilder:
                 logger.error(f"Failed to download checkpoints for model {repo_id}")
                 sys.exit(1)
             else:
-                logger.info(f"Downloaded checkpoints for model {repo_id}")
+                logger.info(f"Downloaded checkpoints for model {repo_id} successfully to {path}")
         return path
 
     def _concepts_protos_from_concepts(self, concepts):
@@ -1140,7 +1250,109 @@ class ModelBuilder:
         concepts = config.get('concepts')
         logger.info(f"Updated config.yaml with {len(concepts)} concepts.")
 
-    def get_model_version_proto(self):
+    def _get_git_info(self) -> Optional[Dict[str, Any]]:
+        """
+        Get git repository information for the model path.
+
+        Returns:
+            Dict with git info (url, commit, branch) or None if not a git repository
+        """
+        try:
+            # Check if the folder is within a git repository
+            result = subprocess.run(
+                ['git', 'rev-parse', '--git-dir'],
+                cwd=self.folder,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+
+            # Get git remote URL
+            remote_result = subprocess.run(
+                ['git', 'config', '--get', 'remote.origin.url'],
+                cwd=self.folder,
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+
+            # Get current commit hash
+            commit_result = subprocess.run(
+                ['git', 'rev-parse', 'HEAD'],
+                cwd=self.folder,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+
+            # Get current branch
+            branch_result = subprocess.run(
+                ['git', 'branch', '--show-current'],
+                cwd=self.folder,
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+
+            git_info = {
+                'commit': commit_result.stdout.strip(),
+                'branch': branch_result.stdout.strip()
+                if branch_result.returncode == 0
+                else 'HEAD',
+            }
+
+            if remote_result.returncode == 0:
+                git_info['url'] = remote_result.stdout.strip()
+
+            return git_info
+
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            # Not a git repository or git not available
+            return None
+
+    def _check_git_status_and_prompt(self) -> bool:
+        """
+        Check for uncommitted changes in git repository within the model path and prompt user.
+
+        Returns:
+            True if should continue with upload, False if should abort
+        """
+        try:
+            # Check for uncommitted changes within the model path only
+            status_result = subprocess.run(
+                ['git', 'status', '--porcelain', '.'],
+                cwd=self.folder,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+
+            if status_result.stdout.strip():
+                logger.warning("Uncommitted changes detected in model path:")
+                logger.warning(status_result.stdout)
+
+                response = input(
+                    "\nDo you want to continue upload with uncommitted changes? (y/N): "
+                )
+                return response.lower() in ['y', 'yes']
+            else:
+                logger.info("Model path has no uncommitted changes.")
+                return True
+
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            # Error checking git status, but we already know it's a git repo
+            logger.warning("Could not check git status, continuing with upload.")
+            return True
+
+    def get_model_version_proto(self, git_info: Optional[Dict[str, Any]] = None):
+        """
+        Create a ModelVersion protobuf message for the model.
+        Args:
+          git_info (Optional[Dict[str, Any]]): Git repository information to include in metadata.
+        Returns:
+          resources_pb2.ModelVersion: The ModelVersion protobuf message.
+        """
+
         signatures = self.get_method_signatures()
         model_version_proto = resources_pb2.ModelVersion(
             pretrained_model_config=resources_pb2.PretrainedModelConfig(),
@@ -1148,6 +1360,37 @@ class ModelBuilder:
             method_signatures=signatures,
         )
 
+        # Add git information to metadata if available
+        if git_info:
+            from google.protobuf.struct_pb2 import Struct
+
+            metadata_struct = Struct()
+            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:
@@ -1175,7 +1418,7 @@ class ModelBuilder:
                 )
         return model_version_proto
 
-    def upload_model_version(self):
+    def upload_model_version(self, git_info=None):
         file_path = f"{self.folder}.tar.gz"
         logger.debug(f"Will tar it into file: {file_path}")
 
@@ -1208,7 +1451,7 @@ class ModelBuilder:
                 )
                 return
 
-        model_version_proto = self.get_model_version_proto()
+        model_version_proto = self.get_model_version_proto(git_info)
 
         def filter_func(tarinfo):
             name = tarinfo.name
@@ -1228,7 +1471,7 @@ class ModelBuilder:
         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)
+            checkpoint_size = int(os.environ.get('CHECKPOINT_SIZE_BYTES', 0))
             if not checkpoint_size:
                 _, repo_id, _, _, _, _ = self._validate_config_checkpoints()
                 checkpoint_size = HuggingFaceLoader.get_huggingface_checkpoint_total_size(repo_id)
@@ -1332,13 +1575,13 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
         )
         return result
 
-    def get_model_build_logs(self):
+    def get_model_build_logs(self, current_page=1):
         logs_request = service_pb2.ListLogEntriesRequest(
             log_type="builder",
             user_app_id=self.client.user_app_id,
             model_id=self.model_proto.id,
             model_version_id=self.model_version_id,
-            page=1,
+            page=current_page,
             per_page=50,
         )
         response = self.client.STUB.ListLogEntries(logs_request)
@@ -1347,6 +1590,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     def monitor_model_build(self):
         st = time.time()
         seen_logs = set()  # To avoid duplicate log messages
+        current_page = 1  # Track current page for log pagination
         while True:
             resp = self.client.STUB.GetModelVersion(
                 service_pb2.GetModelVersionRequest(
@@ -1357,8 +1601,10 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
             )
 
             status_code = resp.model_version.status.code
-            logs = self.get_model_build_logs()
+            logs = self.get_model_build_logs(current_page)
+            entries_count = 0
             for log_entry in logs.log_entries:
+                entries_count += 1
                 if log_entry.url not in seen_logs:
                     seen_logs.add(log_entry.url)
                     log_entry_msg = re.sub(
@@ -1367,6 +1613,12 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                         log_entry.message.strip(),
                     )
                     logger.info(log_entry_msg)
+
+            # If we got a full page (50 entries), there might be more logs on the next page
+            # If we got fewer than 50 entries, we've reached the end and should stay on current page
+            if entries_count == 50:
+                current_page += 1
+            # else: stay on current_page
             if status_code == status_code_pb2.MODEL_BUILDING:
                 print(
                     f"Model is building... (elapsed {time.time() - st:.1f}s)", end='\r', flush=True
@@ -1414,8 +1666,20 @@ def upload_model(folder, stage, skip_dockerfile, pat=None, base_url=None):
             f"New model will be created at {builder.model_ui_url} with it's first version."
         )
 
+    # Check for git repository information
+    git_info = builder._get_git_info()
+    if git_info:
+        logger.info(f"Detected git repository: {git_info.get('url', 'local repository')}")
+        logger.info(f"Current commit: {git_info['commit']}")
+        logger.info(f"Current branch: {git_info['branch']}")
+
+        # Check for uncommitted changes and prompt user
+        if not builder._check_git_status_and_prompt():
+            logger.info("Upload cancelled by user due to uncommitted changes.")
+            return
     input("Press Enter to continue...")
-    model_version = builder.upload_model_version()
+
+    model_version = builder.upload_model_version(git_info)
 
     # Ask user if they want to deploy the model
     if model_version is not None:  # if it comes back None then it failed.

clarifai/runners/models/model_runner.py

@@ -9,7 +9,7 @@ from clarifai_protocol.utils.health import HealthProbeRequestHandler
 from clarifai.client.auth.helper import ClarifaiAuthHelper
 from clarifai.utils.constants import STATUS_FAIL, STATUS_MIXED, STATUS_OK, STATUS_UNKNOWN
 from clarifai.utils.logging import get_req_id_from_context, logger
-from clarifai.utils.secrets import inject_secrets
+from clarifai.utils.secrets import inject_secrets, req_secrets_context
 
 from ..utils.url_fetcher import ensure_urls_downloaded
 from .model_class import ModelClass
@@ -126,7 +126,10 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
         if method_name == '_GET_SIGNATURES':
             logging = False
 
-        resp = self.model.predict_wrapper(request)
+        # Use req_secrets_context to temporarily set request-type secrets as environment variables
+        with req_secrets_context(request):
+            resp = self.model.predict_wrapper(request)
+
         # if we have any non-successful code already it's an error we can return.
         if (
             resp.status.code != status_code_pb2.SUCCESS
@@ -185,45 +188,49 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
         status_str = STATUS_UNKNOWN
         endpoint = "POST /v2/.../outputs/generate"
 
-        for resp in self.model.generate_wrapper(request):
-            # if we have any non-successful code already it's an error we can return.
-            if (
-                resp.status.code != status_code_pb2.SUCCESS
-                and resp.status.code != status_code_pb2.ZERO
-            ):
-                status_str = f"{resp.status.code} ERROR"
-                duration_ms = (time.time() - start_time) * 1000
-                logger.info(f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}")
-                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
-                continue
-            successes = []
-            for output in resp.outputs:
-                if not output.HasField('status') or not output.status.code:
-                    raise Exception(
-                        "Output must have a status code, please check the model implementation."
+        # Use req_secrets_context to temporarily set request-type secrets as environment variables
+        with req_secrets_context(request):
+            for resp in self.model.generate_wrapper(request):
+                # if we have any non-successful code already it's an error we can return.
+                if (
+                    resp.status.code != status_code_pb2.SUCCESS
+                    and resp.status.code != status_code_pb2.ZERO
+                ):
+                    status_str = f"{resp.status.code} ERROR"
+                    duration_ms = (time.time() - start_time) * 1000
+                    logger.info(
+                        f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}"
                     )
-                successes.append(output.status.code == status_code_pb2.SUCCESS)
-            if all(successes):
-                status = status_pb2.Status(
-                    code=status_code_pb2.SUCCESS,
-                    description="Success",
-                )
-                status_str = STATUS_OK
-            elif any(successes):
-                status = status_pb2.Status(
-                    code=status_code_pb2.MIXED_STATUS,
-                    description="Mixed Status",
-                )
-                status_str = STATUS_MIXED
-            else:
-                status = status_pb2.Status(
-                    code=status_code_pb2.FAILURE,
-                    description="Failed",
-                )
-                status_str = STATUS_FAIL
-            resp.status.CopyFrom(status)
+                    yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                    continue
+                successes = []
+                for output in resp.outputs:
+                    if not output.HasField('status') or not output.status.code:
+                        raise Exception(
+                            "Output must have a status code, please check the model implementation."
+                        )
+                    successes.append(output.status.code == status_code_pb2.SUCCESS)
+                if all(successes):
+                    status = status_pb2.Status(
+                        code=status_code_pb2.SUCCESS,
+                        description="Success",
+                    )
+                    status_str = STATUS_OK
+                elif any(successes):
+                    status = status_pb2.Status(
+                        code=status_code_pb2.MIXED_STATUS,
+                        description="Mixed Status",
+                    )
+                    status_str = STATUS_MIXED
+                else:
+                    status = status_pb2.Status(
+                        code=status_code_pb2.FAILURE,
+                        description="Failed",
+                    )
+                    status_str = STATUS_FAIL
+                resp.status.CopyFrom(status)
 
-            yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
 
         duration_ms = (time.time() - start_time) * 1000
         logger.info(f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}")
@@ -237,45 +244,55 @@ class ModelRunner(BaseRunner, HealthProbeRequestHandler):
         status_str = STATUS_UNKNOWN
         endpoint = "POST /v2/.../outputs/stream  "
 
-        for resp in self.model.stream_wrapper(pmo_iterator(runner_item_iterator)):
-            # if we have any non-successful code already it's an error we can return.
-            if (
-                resp.status.code != status_code_pb2.SUCCESS
-                and resp.status.code != status_code_pb2.ZERO
-            ):
-                status_str = f"{resp.status.code} ERROR"
-                duration_ms = (time.time() - start_time) * 1000
-                logger.info(f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}")
-                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
-                continue
-            successes = []
-            for output in resp.outputs:
-                if not output.HasField('status') or not output.status.code:
-                    raise Exception(
-                        "Output must have a status code, please check the model implementation."
+        # Get the first request to establish secrets context
+        first_request = None
+        runner_items = list(runner_item_iterator)  # Convert to list to avoid consuming iterator
+        if runner_items:
+            first_request = runner_items[0].post_model_outputs_request
+
+        # Use req_secrets_context based on the first request (secrets should be consistent across stream)
+        with req_secrets_context(first_request):
+            for resp in self.model.stream_wrapper(pmo_iterator(iter(runner_items))):
+                # if we have any non-successful code already it's an error we can return.
+                if (
+                    resp.status.code != status_code_pb2.SUCCESS
+                    and resp.status.code != status_code_pb2.ZERO
+                ):
+                    status_str = f"{resp.status.code} ERROR"
+                    duration_ms = (time.time() - start_time) * 1000
+                    logger.info(
+                        f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}"
                     )
-                successes.append(output.status.code == status_code_pb2.SUCCESS)
-            if all(successes):
-                status = status_pb2.Status(
-                    code=status_code_pb2.SUCCESS,
-                    description="Success",
-                )
-                status_str = STATUS_OK
-            elif any(successes):
-                status = status_pb2.Status(
-                    code=status_code_pb2.MIXED_STATUS,
-                    description="Mixed Status",
-                )
-                status_str = STATUS_MIXED
-            else:
-                status = status_pb2.Status(
-                    code=status_code_pb2.FAILURE,
-                    description="Failed",
-                )
-                status_str = STATUS_FAIL
-            resp.status.CopyFrom(status)
+                    yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                    continue
+                successes = []
+                for output in resp.outputs:
+                    if not output.HasField('status') or not output.status.code:
+                        raise Exception(
+                            "Output must have a status code, please check the model implementation."
+                        )
+                    successes.append(output.status.code == status_code_pb2.SUCCESS)
+                if all(successes):
+                    status = status_pb2.Status(
+                        code=status_code_pb2.SUCCESS,
+                        description="Success",
+                    )
+                    status_str = STATUS_OK
+                elif any(successes):
+                    status = status_pb2.Status(
+                        code=status_code_pb2.MIXED_STATUS,
+                        description="Mixed Status",
+                    )
+                    status_str = STATUS_MIXED
+                else:
+                    status = status_pb2.Status(
+                        code=status_code_pb2.FAILURE,
+                        description="Failed",
+                    )
+                    status_str = STATUS_FAIL
+                resp.status.CopyFrom(status)
 
-            yield service_pb2.RunnerItemOutput(multi_output_response=resp)
+                yield service_pb2.RunnerItemOutput(multi_output_response=resp)
 
         duration_ms = (time.time() - start_time) * 1000
         logger.info(f"{endpoint} | {status_str} | {duration_ms:.2f}ms | req_id={req_id}")