clarifai 11.5.2__py3-none-any.whl → 11.5.3__py3-none-any.whl

clarifai/client/model_client.py

@@ -1,3 +1,4 @@
+import asyncio
 import inspect
 import time
 from typing import Any, Dict, Iterator, List
@@ -5,6 +6,7 @@ from typing import Any, Dict, Iterator, List
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2
 
+from clarifai.client.auth.register import V2Stub
 from clarifai.constants.model import MAX_MODEL_PREDICT_INPUTS
 from clarifai.errors import UserError
 from clarifai.runners.utils import code_script, method_signatures
@@ -20,12 +22,31 @@ from clarifai.utils.logging import logger
 from clarifai.utils.misc import BackoffIterator, status_is_retryable
 
 
+def is_async_context():
+    """Check if code is running in an async context."""
+    try:
+        asyncio.get_running_loop()
+        import sys
+
+        # In Jupyter, to check if we're actually in an async cell. Becaue by default jupyter considers it as async.
+        if 'ipykernel' in sys.modules:
+            return False
+        return True
+    except RuntimeError:
+        return False
+
+
 class ModelClient:
     '''
     Client for calling model predict, generate, and stream methods.
     '''
 
-    def __init__(self, stub, request_template: service_pb2.PostModelOutputsRequest = None):
+    def __init__(
+        self,
+        stub,
+        async_stub: V2Stub = None,
+        request_template: service_pb2.PostModelOutputsRequest = None,
+    ):
         '''
         Initialize the model client.
 
@@ -35,6 +56,7 @@ class ModelClient:
             common fields like model_id, model_version, cluster, etc.
         '''
         self.STUB = stub
+        self.async_stub = async_stub
         self.request_template = request_template or service_pb2.PostModelOutputsRequest()
         self._method_signatures = None
         self._defined = False
@@ -137,16 +159,19 @@ class ModelClient:
             # define the function in this client instance
             if resources_pb2.RunnerMethodType.Name(method_signature.method_type) == 'UNARY_UNARY':
                 call_func = self._predict
+                async_call_func = self._async_predict
             elif (
                 resources_pb2.RunnerMethodType.Name(method_signature.method_type)
                 == 'UNARY_STREAMING'
             ):
                 call_func = self._generate
+                async_call_func = self._async_generate
             elif (
                 resources_pb2.RunnerMethodType.Name(method_signature.method_type)
                 == 'STREAMING_STREAMING'
             ):
                 call_func = self._stream
+                async_call_func = self._async_stream
             else:
                 raise ValueError(f"Unknown method type {method_signature.method_type}")
 
@@ -158,8 +183,8 @@ class ModelClient:
                     continue
                 method_argnames.append(outer)
 
-            def bind_f(method_name, method_argnames, call_func):
-                def f(*args, **kwargs):
+            def bind_f(method_name, method_argnames, call_func, async_call_func):
+                def sync_f(*args, **kwargs):
                     if len(args) > len(method_argnames):
                         raise TypeError(
                             f"{method_name}() takes {len(method_argnames)} positional arguments but {len(args)} were given"
@@ -187,10 +212,45 @@ class ModelClient:
                         kwargs[name] = arg
                     return call_func(kwargs, method_name)
 
-                return f
+                async def async_f(*args, **kwargs):
+                    # Async version to call the async function
+                    if len(args) > len(method_argnames):
+                        raise TypeError(
+                            f"{method_name}() takes {len(method_argnames)} positional arguments but {len(args)} were given"
+                        )
+                    if len(args) + len(kwargs) > len(method_argnames):
+                        raise TypeError(
+                            f"{method_name}() got an unexpected keyword argument {next(iter(kwargs))}"
+                        )
+                    if len(args) == 1 and (not kwargs) and isinstance(args[0], list):
+                        batch_inputs = args[0]
+                        # Validate each input is a dictionary
+                        is_batch_input_valid = all(
+                            isinstance(input, dict) for input in batch_inputs
+                        )
+                        if is_batch_input_valid and (not is_openai_chat_format(batch_inputs)):
+                            # If the batch input is valid, call the function with the batch inputs and the method name
+                            return async_call_func(batch_inputs, method_name)
+
+                    for name, arg in zip(
+                        method_argnames, args
+                    ):  # handle positional with zip shortest
+                        if name in kwargs:
+                            raise TypeError(f"Multiple values for argument {name}")
+                        kwargs[name] = arg
+
+                    return async_call_func(kwargs, method_name)
+
+                class MethodWrapper:
+                    def __call__(self, *args, **kwargs):
+                        if is_async_context():
+                            return async_f(*args, **kwargs)
+                        return sync_f(*args, **kwargs)
+
+                return MethodWrapper()
 
             # need to bind method_name to the value, not the mutating loop variable
-            f = bind_f(method_name, method_argnames, call_func)
+            f = bind_f(method_name, method_argnames, call_func, async_call_func)
 
             # set names, annotations and docstrings
             f.__name__ = method_name
@@ -231,7 +291,11 @@ class ModelClient:
             self.fetch()
         return method_signatures.get_method_signature(self._method_signatures[method_name])
 
-    def generate_client_script(self) -> str:
+    def generate_client_script(
+        self,
+        base_url: str = None,
+        use_ctx: bool = False,
+    ) -> str:
         """Generate a client script for this model.
 
         Returns:
@@ -247,6 +311,11 @@ class ModelClient:
             user_id=self.request_template.user_app_id.user_id,
             app_id=self.request_template.user_app_id.app_id,
             model_id=self.request_template.model_id,
+            base_url=base_url,
+            deployment_id=self.request_template.runner_selector.deployment.id,
+            compute_cluster_id=self.request_template.runner_selector.nodepool.compute_cluster.id,
+            nodepool_id=self.request_template.runner_selector.nodepool.id,
+            use_ctx=use_ctx,
         )
 
     def _define_compatability_functions(self):
@@ -354,6 +423,109 @@ class ModelClient:
             break
         return response
 
+    async def _async_predict(
+        self,
+        inputs,
+        method_name: str = 'predict',
+    ) -> Any:
+        """Asynchronously process inputs and make predictions.
+
+        Args:
+            inputs: Input data to process
+            method_name (str): Name of the method to call
+
+        Returns:
+            Processed prediction results
+        """
+        # method_name is set to 'predict' by default, this is because to replicate the input and output signature behaviour of sync to async predict.
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        batch_input = True
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+            batch_input = False
+
+        proto_inputs = []
+        for input in inputs:
+            proto = resources_pb2.Input()
+            serialize(input, input_signature, proto.data)
+            proto_inputs.append(proto)
+        response = await self._async_predict_by_proto(proto_inputs, method_name)
+        outputs = []
+        for output in response.outputs:
+            outputs.append(deserialize(output.data, output_signature, is_output=True))
+
+        return outputs if batch_input else outputs[0]
+
+    async def _async_predict_by_proto(
+        self,
+        inputs: List[resources_pb2.Input],
+        method_name: str = None,
+        inference_params: Dict = None,
+        output_config: Dict = None,
+    ) -> service_pb2.MultiOutputResponse:
+        """Asynchronously predicts the model based on the given inputs.
+
+        Args:
+            inputs (List[resources_pb2.Input]): The inputs to predict.
+            method_name (str): The remote method name to call.
+            inference_params (Dict): Inference parameters to override.
+            output_config (Dict): Output configuration to override.
+
+        Returns:
+            service_pb2.MultiOutputResponse: The prediction response(s).
+
+        Raises:
+            UserError: If inputs are invalid or exceed maximum limit.
+            Exception: If the model prediction fails.
+        """
+        if not isinstance(inputs, list):
+            raise UserError('Invalid inputs, inputs must be a list of Input objects.')
+        if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
+            raise UserError(f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}.")
+
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+        request.inputs.extend(inputs)
+
+        if method_name:
+            for inp in request.inputs:
+                inp.data.metadata['_method_name'] = method_name
+        if inference_params:
+            request.model.model_version.output_info.params.update(inference_params)
+        if output_config:
+            request.model.model_version.output_info.output_config.MergeFrom(
+                resources_pb2.OutputConfig(**output_config)
+            )
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+
+        while True:
+            try:
+                response = await self.async_stub.PostModelOutputs(request)
+
+                if (
+                    status_is_retryable(response.status.code)
+                    and time.time() - start_time < 60 * 10
+                ):  # 10 minutes
+                    logger.info("Model is still deploying, please wait...")
+                    await asyncio.sleep(next(backoff_iterator))
+                    continue
+
+                if response.status.code != status_code_pb2.SUCCESS:
+                    raise Exception(f"Model predict failed with response {response!r}")
+
+                return response
+
+            except Exception as e:
+                if time.time() - start_time >= 10 * 1:  # 10 minutes timeout
+                    raise Exception("Model prediction timed out after 10 minutes") from e
+                logger.error(f"Error during prediction: {e}")
+                await asyncio.sleep(next(backoff_iterator))
+                continue
+
     def _generate(
         self,
         inputs,  # TODO set up functions according to fetched signatures?
@@ -444,6 +616,99 @@ class ModelClient:
                 raise Exception(f"Model Generate failed with response {response.status!r}")
             yield response
 
+    async def _async_generate(
+        self,
+        inputs,
+        method_name: str = 'generate',
+    ) -> Any:
+        # method_name is set to 'generate' by default, this is because to replicate the input and output signature behaviour of sync to async generate.
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        batch_input = True
+        if isinstance(inputs, dict):
+            inputs = [inputs]
+            batch_input = False
+
+        proto_inputs = []
+        for input in inputs:
+            proto = resources_pb2.Input()
+        serialize(input, input_signature, proto.data)
+        proto_inputs.append(proto)
+        response_stream = self._async_generate_by_proto(proto_inputs, method_name)
+
+        async for response in response_stream:
+            outputs = []
+            for output in response.outputs:
+                outputs.append(deserialize(output.data, output_signature, is_output=True))
+            if batch_input:
+                yield outputs
+            else:
+                yield outputs[0]
+
+    async def _async_generate_by_proto(
+        self,
+        inputs: List[resources_pb2.Input],
+        method_name: str = None,
+        inference_params: Dict = {},
+        output_config: Dict = {},
+    ):
+        """Generate the async stream output on model based on the given inputs.
+
+        Args:
+            inputs (list[Input]): The inputs to generate, must be less than 128.
+            method_name (str): The remote method name to call.
+            inference_params (dict): The inference params to override.
+            output_config (dict): The output config to override.
+        """
+        if not isinstance(inputs, list):
+            raise UserError('Invalid inputs, inputs must be a list of Input objects.')
+        if len(inputs) > MAX_MODEL_PREDICT_INPUTS:
+            raise UserError(
+                f"Too many inputs. Max is {MAX_MODEL_PREDICT_INPUTS}."
+            )  # TODO Use Chunker for inputs len > 128
+
+        request = service_pb2.PostModelOutputsRequest()
+        request.CopyFrom(self.request_template)
+
+        request.inputs.extend(inputs)
+
+        if method_name:
+            # TODO put in new proto field?
+            for inp in request.inputs:
+                inp.data.metadata['_method_name'] = method_name
+        if inference_params:
+            request.model.model_version.output_info.params.update(inference_params)
+        if output_config:
+            request.model.model_version.output_info.output_config.MergeFromDict(output_config)
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        started = False
+        while not started:
+            # stream response returns gRPC async iterable - UnaryStreamCall
+            stream_response = self.async_stub.GenerateModelOutputs(request)
+            stream_resp = await stream_response  # get the async iterable
+            iterator = stream_resp.__aiter__()  # get the async iterator for the response
+            try:
+                response = await iterator.__anext__()  # getting the first response
+            except StopAsyncIteration:
+                raise Exception("Model Generate failed with no response")
+            if status_is_retryable(response.status.code) and time.time() - start_time < 60 * 10:
+                logger.info("Model is still deploying, please wait...")
+                await asyncio.sleep(next(backoff_iterator))
+                continue
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model Generate failed with response {response.status!r}")
+            started = True
+
+        yield response  # yield the first response
+
+        async for response in iterator:
+            if response.status.code != status_code_pb2.SUCCESS:
+                raise Exception(f"Model Generate failed with response {response.status!r}")
+            yield response
+
     def _stream(
         self,
         inputs,
@@ -551,3 +816,87 @@ class ModelClient:
                     if not generation_started:
                         generation_started = True
                     yield response
+
+    # TODO: Test async streaming.
+    async def _async_stream(
+        self,
+        inputs,
+        method_name: str = 'stream',
+    ) -> Any:
+        # method_name is set to 'stream' by default, this is because to replicate the input and output signature behaviour of sync to async stream.
+        input_signature = self._method_signatures[method_name].input_fields
+        output_signature = self._method_signatures[method_name].output_fields
+
+        if isinstance(inputs, list):
+            assert len(inputs) == 1, 'streaming methods do not support batched calls'
+            inputs = inputs[0]
+        assert isinstance(inputs, dict)
+        kwargs = inputs
+
+        # find the streaming vars in the input signature, and the streaming input python param
+        stream_sig = get_stream_from_signature(input_signature)
+        if stream_sig is None:
+            raise ValueError("Streaming method must have a Stream input")
+        stream_argname = stream_sig.name
+
+        # get the streaming input generator from the user-provided function arg values
+        user_inputs_generator = kwargs.pop(stream_argname)
+
+        async def _input_proto_stream():
+            # first item contains all the inputs and the first stream item
+            proto = resources_pb2.Input()
+            try:
+                item = await user_inputs_generator.__anext__()
+            except StopAsyncIteration:
+                return  # no items to stream
+            kwargs[stream_argname] = item
+            serialize(kwargs, input_signature, proto.data)
+
+            yield proto
+
+        # subsequent items are just the stream items
+        async for item in user_inputs_generator:
+            proto = resources_pb2.Input()
+            serialize({stream_argname: item}, [stream_sig], proto.data)
+            yield proto
+
+        response_stream = await self._async_stream_by_proto(_input_proto_stream(), method_name)
+
+        async for response in response_stream:
+            assert len(response.outputs) == 1, 'streaming methods must have exactly one output'
+            yield deserialize(response.outputs[0].data, output_signature, is_output=True)
+
+    async def _async_stream_by_proto(
+        self,
+        inputs: Iterator[List[resources_pb2.Input]],
+        method_name: str = None,
+        inference_params: Dict = {},
+        output_config: Dict = {},
+    ):
+        """Generate the async stream output on model based on the given stream of inputs."""
+        # if not isinstance(inputs, Iterator[List[Input]]):
+        #   raise UserError('Invalid inputs, inputs must be a iterator of list of Input objects.')
+
+        request = self._req_iterator(inputs, method_name, inference_params, output_config)
+
+        start_time = time.time()
+        backoff_iterator = BackoffIterator(10)
+        generation_started = False
+        while True:
+            if generation_started:
+                break
+            stream_response = await self.async_stub.StreamModelOutputs(request)
+            async for response in stream_response:
+                if (
+                    status_is_retryable(response.status.code)
+                    and time.time() - start_time < 60 * 10
+                ):
+                    logger.info("Model is still deploying, please wait...")
+                    await asyncio.sleep(next(backoff_iterator))
+                    break
+                if response.status.code != status_code_pb2.SUCCESS:
+                    raise Exception(f"Model Predict failed with response {response.status!r}")
+                else:
+                    if not generation_started:
+                        generation_started = True
+                    yield response

clarifai/client/user.py

@@ -456,3 +456,84 @@ class User(Lister, BaseClient):
             if hasattr(self.user_info, param)
         ]
         return f"Clarifai User Details: \n{', '.join(attribute_strings)}\n"
+
+    def list_models(
+        self,
+        user_id: str = None,
+        app_id: str = None,
+        show: bool = True,
+        return_clarifai_model: bool = False,
+        **kwargs,
+    ):
+        if user_id == "all":
+            params = {}
+        elif user_id:
+            user_app_id = resources_pb2.UserAppIDSet(user_id=user_id, app_id=app_id)
+            params = {"user_app_id": user_app_id}
+        elif not user_id:
+            user_app_id = resources_pb2.UserAppIDSet(
+                user_id=self.user_app_id.user_id, app_id=app_id
+            )
+            params = {"user_app_id": user_app_id}
+
+        params.update(**kwargs)
+        models = self.list_pages_generator(
+            self.STUB.ListModels, service_pb2.ListModelsRequest, request_data=params
+        )
+        all_data = []
+        for model in models:
+            url = (
+                f"https://clarifai.com/{model['user_id']}/{model['app_id']}/models/{model['name']}"
+            )
+            data = dict(
+                user_id=model["user_id"],
+                app_id=model["app_id"],
+                id=model["model_id"],
+                model_type=model["model_type_id"],
+                url=url,
+            )
+            method_types_data = dict(
+                supported_openai_client=False,
+                UNARY_UNARY={"predict"},
+                UNARY_STREAMING=set(),
+                STREAMING_STREAMING=set(),
+            )
+            for each_method in model.get("model_version", {}).get("method_signatures", []):
+                name = each_method["name"]
+                method_type = each_method["method_type"]
+                method_types_data[method_type].add(name)
+            if (
+                "openai_transport" in method_types_data["UNARY_UNARY"]
+                and "openai_stream_transport" in method_types_data["UNARY_STREAMING"]
+            ):
+                method_types_data["supported_openai_client"] = True
+                method_types_data["UNARY_UNARY"].remove("openai_transport")
+                method_types_data["UNARY_STREAMING"].remove("openai_stream_transport")
+            for k, v in method_types_data.items():
+                if k != "supported_openai_client":
+                    if not v:
+                        method_types_data[k] = None
+                    else:
+                        method_types_data[k] = list(v)
+
+            data.update(method_types_data)
+            all_data.append(data)
+
+        if show:
+            from tabulate import tabulate
+
+            print(tabulate(all_data, headers="keys"))
+
+        if return_clarifai_model:
+            from clarifai.client import Model
+
+            models = []
+            for each_data in all_data:
+                model = Model.from_auth_helper(
+                    self.auth_helper,
+                    url=each_data["url"],
+                )
+                models.append(model)
+            return models
+        else:
+            return all_data

clarifai/runners/models/model_builder.py

@@ -10,6 +10,7 @@ import tarfile
 import time
 import webbrowser
 from string import Template
+from typing import Literal
 from unittest.mock import MagicMock
 
 import yaml
@@ -22,6 +23,7 @@ from clarifai.client.user import User
 from clarifai.runners.models.model_class import ModelClass
 from clarifai.runners.utils.const import (
     AMD_PYTHON_BASE_IMAGE,
+    AMD_TORCH_BASE_IMAGE,
     AMD_VLLM_BASE_IMAGE,
     AVAILABLE_PYTHON_IMAGES,
     AVAILABLE_TORCH_IMAGES,
@@ -45,6 +47,7 @@ from clarifai.versions import CLIENT_VERSION
 dependencies = [
     'torch',
     'clarifai',
+    'vllm',
 ]
 
 
@@ -69,6 +72,7 @@ class ModelBuilder:
         folder: str,
         validate_api_ids: bool = True,
         download_validation_only: bool = False,
+        app_not_found_action: Literal["auto_create", "prompt", "error"] = "error",
     ):
         """
         :param folder: The folder containing the model.py, config.yaml, requirements.txt and
@@ -77,7 +81,13 @@ class ModelBuilder:
         deprecate in favor of download_validation_only.
         :param download_validation_only: Whether to skip the API config validation. Set to True if
         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.
         """
+        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
         if not validate_api_ids:  # for backwards compatibility
             download_validation_only = True
@@ -304,13 +314,46 @@ class ModelBuilder:
                 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
+
+        user_id = self.client.user_app_id.user_id
+        app_id = self.client.user_app_id.app_id
+
+        if self.app_not_found_action == "error":
+            logger.error(
+                f"Error checking API {self._base_api} for user app `{user_id}/{app_id}`. Error code: {resp.status.code}"
+            )
+            logger.error(
+                f"App `{app_id}` not found for user `{user_id}`. Please create the app first and try again."
+            )
+            return False
+        else:
+            user = User(
+                user_id=user_id,
+                pat=self.client.pat,
+                token=self.client.token,
+                base_url=self.client.base,
+            )
+
+            def create_app():
+                logger.info(f"Creating App `{app_id}` user `{user_id}`.")
+                user.create_app(app_id=app_id)
+
+            logger.info(f"App {app_id} not found for user {user_id}.")
+
+            if self.app_not_found_action == "prompt":
+                create_app_prompt = input(f"Do you want to create App `{app_id}`? (y/n): ")
+                if create_app_prompt.lower() == 'y':
+                    create_app()
+                    return True
+                else:
+                    logger.error(
+                        f"App `{app_id}` has not been created for user `{user_id}`. Please create the app first or switch to an existing one, then try again."
+                    )
+                    return False
+
+            elif self.app_not_found_action == "auto_create":
+                create_app()
+                return True
 
     def _validate_config_model(self):
         assert "model" in self.config, "model section not found in the config file"
@@ -759,7 +802,7 @@ class ModelBuilder:
                     )
                 python_version = DEFAULT_PYTHON_VERSION
                 gpu_version = DEFAULT_AMD_GPU_VERSION
-                final_image = TORCH_BASE_IMAGE.format(
+                final_image = AMD_TORCH_BASE_IMAGE.format(
                     torch_version=torch_version,
                     python_version=python_version,
                     gpu_version=gpu_version,
@@ -1186,7 +1229,7 @@ def upload_model(folder, stage, skip_dockerfile):
     :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.
     """
-    builder = ModelBuilder(folder)
+    builder = ModelBuilder(folder, app_not_found_action="prompt")
     builder.download_checkpoints(stage=stage)
     if not skip_dockerfile:
         builder.create_dockerfile()