clarifai 11.5.5__tar.gz → 11.6.0__tar.gz

{clarifai-11.5.5/clarifai.egg-info → clarifai-11.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: clarifai
-Version: 11.5.5
+Version: 11.6.0
 Home-page: https://github.com/Clarifai/clarifai-python
 Author: Clarifai
 Author-email: support@clarifai.com
@@ -20,7 +20,7 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: clarifai-grpc>=11.5.5
-Requires-Dist: clarifai-protocol>=0.0.24
+Requires-Dist: clarifai-protocol>=0.0.25
 Requires-Dist: numpy>=1.22.0
 Requires-Dist: tqdm>=4.65.0
 Requires-Dist: PyYAML>=6.0.1

clarifai-11.6.0/clarifai/__init__.py

@@ -0,0 +1 @@
+__version__ = "11.6.0"

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/cli/model.py

@@ -120,14 +120,22 @@ def init(model_path, model_type_id):
     is_flag=True,
     help='Flag to skip generating a dockerfile so that you can manually edit an already created dockerfile.',
 )
-def upload(model_path, stage, skip_dockerfile):
+@click.pass_context
+def upload(ctx, model_path, stage, skip_dockerfile):
     """Upload a model to Clarifai.
 
     MODEL_PATH: Path to the model directory. If not specified, the current directory is used by default.
     """
     from clarifai.runners.models.model_builder import upload_model
 
-    upload_model(model_path, stage, skip_dockerfile)
+    validate_context(ctx)
+    upload_model(
+        model_path,
+        stage,
+        skip_dockerfile,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base,
+    )
 
 
 @model.command()

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/cli/pipeline.py

@@ -27,6 +27,146 @@ def upload(path):
     upload_pipeline(path)
 
 
+@pipeline.command()
+@click.option(
+    '--config',
+    type=click.Path(exists=True),
+    required=False,
+    help='Path to the pipeline run config file.',
+)
+@click.option('--pipeline_id', required=False, help='Pipeline ID to run.')
+@click.option('--pipeline_version_id', required=False, help='Pipeline Version ID to run.')
+@click.option(
+    '--pipeline_version_run_id',
+    required=False,
+    help='Pipeline Version Run ID. If not provided, a UUID will be generated.',
+)
+@click.option('--user_id', required=False, help='User ID of the pipeline.')
+@click.option('--app_id', required=False, help='App ID that contains the pipeline.')
+@click.option('--nodepool_id', required=False, help='Nodepool ID to run the pipeline on.')
+@click.option(
+    '--compute_cluster_id', required=False, help='Compute Cluster ID to run the pipeline on.'
+)
+@click.option('--pipeline_url', required=False, help='Pipeline URL to run.')
+@click.option(
+    '--timeout',
+    type=int,
+    default=3600,
+    help='Maximum time to wait for completion in seconds. Default 3600 (1 hour).',
+)
+@click.option(
+    '--monitor_interval',
+    type=int,
+    default=10,
+    help='Interval between status checks in seconds. Default 10.',
+)
+@click.option(
+    '--log_file',
+    type=click.Path(),
+    required=False,
+    help='Path to file where logs should be written. If not provided, logs are displayed on console.',
+)
+@click.option(
+    '--monitor',
+    is_flag=True,
+    default=False,
+    help='Monitor an existing pipeline run instead of starting a new one. Requires pipeline_version_run_id.',
+)
+@click.pass_context
+def run(
+    ctx,
+    config,
+    pipeline_id,
+    pipeline_version_id,
+    pipeline_version_run_id,
+    user_id,
+    app_id,
+    nodepool_id,
+    compute_cluster_id,
+    pipeline_url,
+    timeout,
+    monitor_interval,
+    log_file,
+    monitor,
+):
+    """Run a pipeline and monitor its progress."""
+    import json
+
+    from clarifai.client.pipeline import Pipeline
+    from clarifai.utils.cli import from_yaml, validate_context
+
+    validate_context(ctx)
+
+    if config:
+        config_data = from_yaml(config)
+        pipeline_id = config_data.get('pipeline_id', pipeline_id)
+        pipeline_version_id = config_data.get('pipeline_version_id', pipeline_version_id)
+        pipeline_version_run_id = config_data.get(
+            'pipeline_version_run_id', pipeline_version_run_id
+        )
+        user_id = config_data.get('user_id', user_id)
+        app_id = config_data.get('app_id', app_id)
+        nodepool_id = config_data.get('nodepool_id', nodepool_id)
+        compute_cluster_id = config_data.get('compute_cluster_id', compute_cluster_id)
+        pipeline_url = config_data.get('pipeline_url', pipeline_url)
+        timeout = config_data.get('timeout', timeout)
+        monitor_interval = config_data.get('monitor_interval', monitor_interval)
+        log_file = config_data.get('log_file', log_file)
+        monitor = config_data.get('monitor', monitor)
+
+    # compute_cluster_id and nodepool_id are mandatory regardless of whether pipeline_url is provided
+    if not compute_cluster_id or not nodepool_id:
+        raise ValueError("--compute_cluster_id and --nodepool_id are mandatory parameters.")
+
+    # When monitor flag is used, pipeline_version_run_id is mandatory
+    if monitor and not pipeline_version_run_id:
+        raise ValueError("--pipeline_version_run_id is required when using --monitor flag.")
+
+    if pipeline_url:
+        # When using pipeline_url, other parameters are optional (will be parsed from URL)
+        required_params_provided = True
+    else:
+        # When not using pipeline_url, all individual parameters are required
+        required_params_provided = all([pipeline_id, user_id, app_id, pipeline_version_id])
+
+    if not required_params_provided:
+        raise ValueError(
+            "Either --user_id & --app_id & --pipeline_id & --pipeline_version_id or --pipeline_url must be provided."
+        )
+
+    if pipeline_url:
+        pipeline = Pipeline(
+            url=pipeline_url,
+            pat=ctx.obj.current.pat,
+            base_url=ctx.obj.current.api_base,
+            pipeline_version_run_id=pipeline_version_run_id,
+            nodepool_id=nodepool_id,
+            compute_cluster_id=compute_cluster_id,
+            log_file=log_file,
+        )
+    else:
+        pipeline = Pipeline(
+            pipeline_id=pipeline_id,
+            pipeline_version_id=pipeline_version_id,
+            pipeline_version_run_id=pipeline_version_run_id,
+            user_id=user_id,
+            app_id=app_id,
+            nodepool_id=nodepool_id,
+            compute_cluster_id=compute_cluster_id,
+            pat=ctx.obj.current.pat,
+            base_url=ctx.obj.current.api_base,
+            log_file=log_file,
+        )
+
+    if monitor:
+        # Monitor existing pipeline run instead of starting new one
+        result = pipeline.monitor_only(timeout=timeout, monitor_interval=monitor_interval)
+    else:
+        # Start new pipeline run and monitor it
+        result = pipeline.run(timeout=timeout, monitor_interval=monitor_interval)
+    click.echo(json.dumps(result, indent=2, default=str))
+
+
 @pipeline.command()
 @click.argument(
     "pipeline_path",

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/client/__init__.py

@@ -7,6 +7,7 @@ from clarifai.client.input import Inputs
 from clarifai.client.lister import Lister
 from clarifai.client.model import Model
 from clarifai.client.module import Module
+from clarifai.client.pipeline import Pipeline
 from clarifai.client.search import Search
 from clarifai.client.user import User
 from clarifai.client.workflow import Workflow
@@ -18,6 +19,7 @@ __all__ = [
     'App',
     'Model',
     'Workflow',
+    'Pipeline',
     'Module',
     'Lister',
     'Dataset',

clarifai-11.6.0/clarifai/client/pipeline.py

@@ -0,0 +1,312 @@
+import time
+import uuid
+from typing import Dict, List
+
+from clarifai_grpc.grpc.api import resources_pb2, service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2
+
+from clarifai.client.base import BaseClient
+from clarifai.client.lister import Lister
+from clarifai.errors import UserError
+from clarifai.urls.helper import ClarifaiUrlHelper
+from clarifai.utils.constants import DEFAULT_BASE
+from clarifai.utils.logging import logger
+
+
+def _get_status_name(status_code: int) -> str:
+    """Get the human-readable name for a status code."""
+    status_mapping = {
+        # Job status codes (these are the actual values based on the error message showing 64001)
+        64001: "JOB_QUEUED",
+        64002: "JOB_RUNNING",
+        64003: "JOB_COMPLETED",
+        64004: "JOB_FAILED",
+        64005: "JOB_UNEXPECTED_ERROR",
+        # Standard status codes
+        10000: "SUCCESS",
+        10010: "MIXED_STATUS",
+    }
+    return status_mapping.get(status_code, f"UNKNOWN_STATUS_{status_code}")
+
+
+class Pipeline(Lister, BaseClient):
+    """Pipeline is a class that provides access to Clarifai API endpoints related to Pipeline information."""
+
+    def __init__(
+        self,
+        url: str = None,
+        pipeline_id: str = None,
+        pipeline_version_id: str = None,
+        pipeline_version_run_id: str = None,
+        user_id: str = None,
+        app_id: str = None,
+        nodepool_id: str = None,
+        compute_cluster_id: str = None,
+        log_file: str = None,
+        base_url: str = DEFAULT_BASE,
+        pat: str = None,
+        token: str = None,
+        root_certificates_path: str = None,
+        **kwargs,
+    ):
+        """Initializes a Pipeline object.
+
+        Args:
+            url (str): The URL to initialize the pipeline object.
+            pipeline_id (str): The Pipeline ID to interact with.
+            pipeline_version_id (str): The Pipeline Version ID to interact with.
+            pipeline_version_run_id (str): The Pipeline Version Run ID. If not provided, a UUID will be generated.
+            user_id (str): The User ID that owns the pipeline.
+            app_id (str): The App ID that contains the pipeline.
+            nodepool_id (str): The Nodepool ID to run the pipeline on.
+            compute_cluster_id (str): The Compute Cluster ID to run the pipeline on.
+            log_file (str): Path to file where logs should be written. If not provided, logs are displayed on console.
+            base_url (str): Base API url. Default "https://api.clarifai.com"
+            pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
+            token (str): A session token for authentication. Accepts either a session token or a pat. Can be set as env var CLARIFAI_SESSION_TOKEN
+            root_certificates_path (str): Path to the SSL root certificates file, used to establish secure gRPC connections.
+            **kwargs: Additional keyword arguments to be passed to the Pipeline.
+        """
+        if url and pipeline_id:
+            raise UserError("You can only specify one of url or pipeline_id.")
+        if not url and not pipeline_id:
+            raise UserError("You must specify one of url or pipeline_id.")
+        if url:
+            parsed_user_id, parsed_app_id, _, parsed_pipeline_id, parsed_version_id = (
+                ClarifaiUrlHelper.split_clarifai_url(url)
+            )
+            user_id = user_id or parsed_user_id
+            app_id = app_id or parsed_app_id
+            pipeline_id = parsed_pipeline_id
+            pipeline_version_id = pipeline_version_id or parsed_version_id
+
+        self.pipeline_id = pipeline_id
+        self.pipeline_version_id = pipeline_version_id
+        self.pipeline_version_run_id = pipeline_version_run_id or str(uuid.uuid4())
+        self.user_id = user_id
+        self.app_id = app_id
+        self.nodepool_id = nodepool_id
+        self.compute_cluster_id = compute_cluster_id
+        self.log_file = log_file
+
+        BaseClient.__init__(
+            self,
+            user_id=user_id,
+            app_id=app_id,
+            base=base_url,
+            pat=pat,
+            token=token,
+            root_certificates_path=root_certificates_path,
+        )
+        Lister.__init__(self)
+
+        # Set up runner selector if compute cluster and nodepool are provided
+        self._runner_selector = None
+        if self.compute_cluster_id and self.nodepool_id:
+            from clarifai.client.nodepool import Nodepool
+
+            self._runner_selector = Nodepool.get_runner_selector(
+                user_id=self.user_id,
+                compute_cluster_id=self.compute_cluster_id,
+                nodepool_id=self.nodepool_id,
+            )
+
+    def run(self, inputs: List = None, timeout: int = 3600, monitor_interval: int = 10) -> Dict:
+        """Run the pipeline and monitor its progress.
+
+        Args:
+            inputs (List): List of inputs to run the pipeline with. If None, runs without inputs.
+            timeout (int): Maximum time to wait for completion in seconds. Default 3600 (1 hour).
+            monitor_interval (int): Interval between status checks in seconds. Default 10.
+
+        Returns:
+            Dict: The pipeline run result.
+        """
+        # Create a new pipeline version run
+        pipeline_version_run = resources_pb2.PipelineVersionRun()
+        pipeline_version_run.id = self.pipeline_version_run_id
+
+        # Set nodepools if nodepool information is available
+        if self.nodepool_id and self.compute_cluster_id:
+            nodepool = resources_pb2.Nodepool(
+                id=self.nodepool_id,
+                compute_cluster=resources_pb2.ComputeCluster(
+                    id=self.compute_cluster_id, user_id=self.user_id
+                ),
+            )
+            pipeline_version_run.nodepools.extend([nodepool])
+
+        run_request = service_pb2.PostPipelineVersionRunsRequest()
+        run_request.user_app_id.CopyFrom(self.user_app_id)
+        run_request.pipeline_id = self.pipeline_id
+        run_request.pipeline_version_id = self.pipeline_version_id or ""
+        run_request.pipeline_version_runs.append(pipeline_version_run)
+
+        # Add runner selector if available
+        if self._runner_selector:
+            run_request.runner_selector.CopyFrom(self._runner_selector)
+
+        logger.info(f"Starting pipeline run for pipeline {self.pipeline_id}")
+        response = self.STUB.PostPipelineVersionRuns(
+            run_request, metadata=self.auth_helper.metadata
+        )
+
+        if response.status.code != status_code_pb2.StatusCode.SUCCESS:
+            raise UserError(
+                f"Failed to start pipeline run: {response.status.description}. Details: {response.status.details}"
+            )
+
+        if not response.pipeline_version_runs:
+            raise UserError("No pipeline version run was created")
+
+        pipeline_version_run = response.pipeline_version_runs[0]
+        run_id = pipeline_version_run.id or self.pipeline_version_run_id
+
+        logger.info(f"Pipeline version run created with ID: {run_id}")
+
+        # Monitor the run
+        return self._monitor_pipeline_run(run_id, timeout, monitor_interval)
+
+    def monitor_only(self, timeout: int = 3600, monitor_interval: int = 10) -> Dict:
+        """Monitor an existing pipeline run without starting a new one.
+
+        Args:
+            timeout (int): Maximum time to wait for completion in seconds. Default 3600 (1 hour).
+            monitor_interval (int): Interval between status checks in seconds. Default 10.
+
+        Returns:
+            Dict: The pipeline run result.
+        """
+        if not self.pipeline_version_run_id:
+            raise UserError("pipeline_version_run_id is required for monitoring existing runs")
+
+        logger.info(f"Monitoring existing pipeline run with ID: {self.pipeline_version_run_id}")
+
+        # Monitor the existing run
+        return self._monitor_pipeline_run(self.pipeline_version_run_id, timeout, monitor_interval)
+
+    def _monitor_pipeline_run(self, run_id: str, timeout: int, monitor_interval: int) -> Dict:
+        """Monitor a pipeline version run until completion.
+
+        Args:
+            run_id (str): The pipeline version run ID to monitor.
+            timeout (int): Maximum time to wait for completion in seconds.
+            monitor_interval (int): Interval between status checks in seconds.
+
+        Returns:
+            Dict: The pipeline run result.
+        """
+        start_time = time.time()
+        seen_logs = set()
+
+        while time.time() - start_time < timeout:
+            # Get run status
+            get_run_request = service_pb2.GetPipelineVersionRunRequest()
+            get_run_request.user_app_id.CopyFrom(self.user_app_id)
+            get_run_request.pipeline_id = self.pipeline_id
+            get_run_request.pipeline_version_id = self.pipeline_version_id or ""
+            get_run_request.pipeline_version_run_id = run_id
+
+            try:
+                run_response = self.STUB.GetPipelineVersionRun(
+                    get_run_request, metadata=self.auth_helper.metadata
+                )
+
+                if run_response.status.code != status_code_pb2.StatusCode.SUCCESS:
+                    logger.error(f"Error getting run status: {run_response.status.description}")
+                    time.sleep(monitor_interval)
+                    continue
+
+                pipeline_run = run_response.pipeline_version_run
+
+                # Display new log entries
+                self._display_new_logs(run_id, seen_logs)
+
+                elapsed_time = time.time() - start_time
+                logger.info(f"Pipeline run monitoring... (elapsed {elapsed_time:.1f}s)")
+
+                # Check if we have orchestration status
+                if (
+                    hasattr(pipeline_run, 'orchestration_status')
+                    and pipeline_run.orchestration_status
+                ):
+                    orch_status = pipeline_run.orchestration_status
+                    if hasattr(orch_status, 'status') and orch_status.status:
+                        status_code = orch_status.status.code
+                        status_name = _get_status_name(status_code)
+                        logger.info(f"Pipeline run status: {status_code} ({status_name})")
+
+                        # Display orchestration status details if available
+                        if hasattr(orch_status, 'description') and orch_status.description:
+                            logger.info(f"Orchestration status: {orch_status.description}")
+
+                        # Success codes that allow continuation: JOB_RUNNING, JOB_QUEUED
+                        if status_code in [64001, 64002]:  # JOB_QUEUED, JOB_RUNNING
+                            logger.info(f"Pipeline run in progress: {status_code} ({status_name})")
+                            # Continue monitoring
+                        # Successful terminal state: JOB_COMPLETED
+                        elif status_code == 64003:  # JOB_COMPLETED
+                            logger.info("Pipeline run completed successfully!")
+                            return {"status": "success", "pipeline_version_run": pipeline_run}
+                        # Failure terminal states: JOB_UNEXPECTED_ERROR, JOB_FAILED
+                        elif status_code in [64004, 64005]:  # JOB_FAILED, JOB_UNEXPECTED_ERROR
+                            logger.error(
+                                f"Pipeline run failed with status: {status_code} ({status_name})"
+                            )
+                            return {"status": "failed", "pipeline_version_run": pipeline_run}
+                        # Handle legacy SUCCESS status for backward compatibility
+                        elif status_code == status_code_pb2.StatusCode.SUCCESS:
+                            logger.info("Pipeline run completed successfully!")
+                            return {"status": "success", "pipeline_version_run": pipeline_run}
+                        elif status_code != status_code_pb2.StatusCode.MIXED_STATUS:
+                            # Log other unexpected statuses but continue monitoring
+                            logger.warning(
+                                f"Unexpected pipeline run status: {status_code} ({status_name}). Continuing to monitor..."
+                            )
+
+            except Exception as e:
+                logger.error(f"Error monitoring pipeline run: {e}")
+
+            time.sleep(monitor_interval)
+
+        logger.error(f"Pipeline run timed out after {timeout} seconds")
+        return {"status": "timeout"}
+
+    def _display_new_logs(self, run_id: str, seen_logs: set):
+        """Display new log entries for a pipeline version run.
+
+        Args:
+            run_id (str): The pipeline version run ID.
+            seen_logs (set): Set of already seen log entry IDs.
+        """
+        try:
+            logs_request = service_pb2.ListLogEntriesRequest()
+            logs_request.user_app_id.CopyFrom(self.user_app_id)
+            logs_request.pipeline_id = self.pipeline_id
+            logs_request.pipeline_version_id = self.pipeline_version_id or ""
+            logs_request.pipeline_version_run_id = run_id
+            logs_request.log_type = "pipeline.version.run"  # Set required log type
+            logs_request.page = 1
+            logs_request.per_page = 50
+
+            logs_response = self.STUB.ListLogEntries(
+                logs_request, metadata=self.auth_helper.metadata
+            )
+
+            if logs_response.status.code == status_code_pb2.StatusCode.SUCCESS:
+                for log_entry in logs_response.log_entries:
+                    # Use log entry URL or timestamp as unique identifier
+                    log_id = log_entry.url or f"{log_entry.created_at.seconds}_{log_entry.message}"
+                    if log_id not in seen_logs:
+                        seen_logs.add(log_id)
+                        log_message = f"[LOG] {log_entry.message.strip()}"
+
+                        # Write to file if log_file is specified, otherwise log to console
+                        if self.log_file:
+                            with open(self.log_file, 'a', encoding='utf-8') as f:
+                                f.write(log_message + '\n')
+                        else:
+                            logger.info(log_message)
+
+        except Exception as e:
+            logger.debug(f"Error fetching logs: {e}")

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/runners/models/model_builder.py

@@ -73,6 +73,8 @@ class ModelBuilder:
         validate_api_ids: bool = True,
         download_validation_only: bool = False,
         app_not_found_action: Literal["auto_create", "prompt", "error"] = "error",
+        pat: str = None,
+        base_url: str = None,
     ):
         """
         :param folder: The folder containing the model.py, config.yaml, requirements.txt and
@@ -83,12 +85,16 @@ class ModelBuilder:
         just downloading a checkpoint.
         :param app_not_found_action: Defines how to handle the case when the app is not found.
         Options: 'auto_create' - create automatically, 'prompt' - ask user, 'error' - raise exception.
+        :param pat: Personal access token for authentication. If None, will use environment variables.
+        :param base_url: Base URL for the API. If None, will use environment variables.
         """
         assert app_not_found_action in ["auto_create", "prompt", "error"], ValueError(
             f"Expected one of {['auto_create', 'prompt', 'error']}, got {app_not_found_action=}"
         )
         self.app_not_found_action = app_not_found_action
         self._client = None
+        self._pat = pat
+        self._base_url = base_url
         if not validate_api_ids:  # for backwards compatibility
             download_validation_only = True
         self.download_validation_only = download_validation_only
@@ -487,8 +493,20 @@ class ModelBuilder:
             user_id = model.get('user_id')
             app_id = model.get('app_id')
 
-            self._base_api = os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
-            self._client = BaseClient(user_id=user_id, app_id=app_id, base=self._base_api)
+            # Use context parameters if provided, otherwise fall back to environment variables
+            self._base_api = (
+                self._base_url
+                if self._base_url
+                else os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
+            )
+
+            # Create BaseClient with explicit pat parameter if provided
+            if self._pat:
+                self._client = BaseClient(
+                    user_id=user_id, app_id=app_id, base=self._base_api, pat=self._pat
+                )
+            else:
+                self._client = BaseClient(user_id=user_id, app_id=app_id, base=self._base_api)
 
         return self._client
 
@@ -550,6 +568,11 @@ class ModelBuilder:
             "inference_compute_info not found in the config file"
         )
         inference_compute_info = self.config.get('inference_compute_info')
+        # Ensure cpu_limit is a string if it exists and is an int
+        if 'cpu_limit' in inference_compute_info and isinstance(
+            inference_compute_info['cpu_limit'], int
+        ):
+            inference_compute_info['cpu_limit'] = str(inference_compute_info['cpu_limit'])
         return json_format.ParseDict(inference_compute_info, resources_pb2.ComputeInfo())
 
     def check_model_exists(self):
@@ -1221,15 +1244,17 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                 return False
 
 
-def upload_model(folder, stage, skip_dockerfile):
+def upload_model(folder, stage, skip_dockerfile, pat=None, base_url=None):
     """
     Uploads a model to Clarifai.
 
     :param folder: The folder containing the model files.
     :param stage: The stage we are calling download checkpoints from. Typically this would "upload" and will download checkpoints if config.yaml checkpoints section has when set to "upload". Other options include "runtime" to be used in load_model or "upload" to be used during model upload. Set this stage to whatever you have in config.yaml to force downloading now.
     :param skip_dockerfile: If True, will not create a Dockerfile.
+    :param pat: Personal access token for authentication. If None, will use environment variables.
+    :param base_url: Base URL for the API. If None, will use environment variables.
     """
-    builder = ModelBuilder(folder, app_not_found_action="prompt")
+    builder = ModelBuilder(folder, app_not_found_action="prompt", pat=pat, base_url=base_url)
     builder.download_checkpoints(stage=stage)
     if not skip_dockerfile:
         builder.create_dockerfile()

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/runners/models/model_class.py

@@ -1,6 +1,7 @@
 import inspect
 import itertools
 import os
+import threading
 import traceback
 from abc import ABC
 from collections import abc
@@ -59,15 +60,22 @@ class ModelClass(ABC):
             yield item.x + ' ' + str(item.y)
     '''
 
+    def __init__(self):
+        super().__init__()
+        self._thread_local = threading.local()
+
     @staticmethod
     def method(func):
         setattr(func, _METHOD_INFO_ATTR, _MethodInfo(func))
         return func
 
     def set_output_context(self, prompt_tokens=None, completion_tokens=None):
-        """This is used to set the prompt and completion tokens in the Output proto"""
-        self._prompt_tokens = prompt_tokens
-        self._completion_tokens = completion_tokens
+        """Set the prompt and completion tokens for the Output proto.
+        In batch mode, call this once per output, in order, before returning each output.
+        """
+        if not hasattr(self._thread_local, 'token_contexts'):
+            self._thread_local.token_contexts = []
+        self._thread_local.token_contexts.append((prompt_tokens, completion_tokens))
 
     def load_model(self):
         """Load the model."""
@@ -213,6 +221,7 @@ class ModelClass(ABC):
         self, request: service_pb2.PostModelOutputsRequest
     ) -> Iterator[service_pb2.MultiOutputResponse]:
         try:
+            assert len(request.inputs) == 1, "Generate requires exactly one input"
             method_name = 'generate'
             if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
                 method_name = request.inputs[0].data.metadata['_method_name']
@@ -385,12 +394,18 @@ class ModelClass(ABC):
             data = DataConverter.convert_output_data_to_old_format(proto.data)
             proto.data.CopyFrom(data)
         proto.status.code = status_code_pb2.SUCCESS
-        if hasattr(self, "_prompt_tokens") and self._prompt_tokens is not None:
-            proto.prompt_tokens = self._prompt_tokens
-        if hasattr(self, "_completion_tokens") and self._completion_tokens is not None:
-            proto.completion_tokens = self._completion_tokens
-        self._prompt_tokens = None
-        self._completion_tokens = None
+        # Per-output token context support
+        token_contexts = getattr(self._thread_local, 'token_contexts', None)
+        prompt_tokens = completion_tokens = None
+        if token_contexts and len(token_contexts) > 0:
+            prompt_tokens, completion_tokens = token_contexts.pop(0)
+            # If this was the last, clean up
+            if len(token_contexts) == 0:
+                del self._thread_local.token_contexts
+        if prompt_tokens is not None:
+            proto.prompt_tokens = prompt_tokens
+        if completion_tokens is not None:
+            proto.completion_tokens = completion_tokens
         return proto
 
     @classmethod

{clarifai-11.5.5 → clarifai-11.6.0}/clarifai/runners/models/openai_class.py

@@ -34,6 +34,7 @@ class OpenAIModelClass(ModelClass):
     model = None
 
     def __init__(self) -> None:
+        super().__init__()
         if self.client is None:
             raise NotImplementedError("Subclasses must set the 'client' class attribute")
         if self.model is None: