clarifai 11.2.3rc6__tar.gz → 11.2.3rc8__tar.gz

{clarifai-11.2.3rc6/clarifai.egg-info → clarifai-11.2.3rc8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: clarifai
-Version: 11.2.3rc6
+Version: 11.2.3rc8
 Summary: Clarifai Python SDK
 Home-page: https://github.com/Clarifai/clarifai-python
 Author: Clarifai

clarifai-11.2.3rc8/clarifai/__init__.py

@@ -0,0 +1 @@
+__version__ = "11.2.3rc8"

clarifai-11.2.3rc8/clarifai/cli/base.py

@@ -0,0 +1,260 @@
+import json
+import os
+import sys
+
+import click
+import yaml
+
+from clarifai.utils.cli import AliasedGroup, TableFormatter, load_command_modules
+from clarifai.utils.config import Config, Context
+from clarifai.utils.constants import DEFAULT_CONFIG
+from clarifai.utils.logging import logger
+
+
+#@click.group(cls=CustomMultiGroup)
+@click.group(cls=AliasedGroup)
+@click.option('--config', default=DEFAULT_CONFIG)
+@click.pass_context
+def cli(ctx, config):
+  """Clarifai CLI"""
+  ctx.ensure_object(dict)
+  if os.path.exists(config):
+    cfg = Config.from_yaml(filename=config)
+    ctx.obj = cfg
+  else:
+    cfg = Config(
+        filename=config,
+        current_context='default',
+        contexts={
+            'default':
+                Context(
+                    'default',
+                    CLARIFAI_PAT=os.environ.get('CLARIFAI_PAT', ''),
+                    CLARIFAI_USER_ID=os.environ.get('CLARIFAI_USER_ID', ''),
+                    CLARIFAI_API_BASE=os.environ.get('CLARIFAI_API_BASE', 'api.clarifai.com'),
+                )
+        })
+    try:
+      cfg.to_yaml(config)
+    except Exception:
+      logger.warning("Could not write configuration to disk. Could be a read only file system.")
+    ctx.obj = cfg  # still have the default config even if couldn't write.
+
+
+@cli.command()
+@click.argument('shell', type=click.Choice(['bash', 'zsh']))
+def shell_completion(shell):
+  """Shell completion script"""
+  os.system(f"_CLARIFAI_COMPLETE={shell}_source clarifai")
+
+
+@cli.group(['cfg'], cls=AliasedGroup)
+def config():
+  """Manage CLI configuration"""
+
+
+@config.command(['e'])
+@click.pass_context
+def edit(ctx):
+  """Edit the configuration file"""
+  os.system(f'{os.environ.get("EDITOR", "vi")} {ctx.obj.filename}')
+
+
+@config.command(['current'])
+@click.option('-o', '--output-format', default='name', type=click.Choice(['name', 'json', 'yaml']))
+@click.pass_context
+def current_context(ctx, output_format):
+  """Get the current context"""
+  if output_format == 'name':
+    print(ctx.obj.current_context)
+  else:
+    if output_format == 'json':
+      print(json.dumps(ctx.obj.contexts[ctx.obj.current_context].to_serializable_dict()))
+    else:
+      print(yaml.safe_dump(ctx.obj.contexts[ctx.obj.current_context].to_serializable_dict()))
+
+
+@config.command(['list', 'ls'])
+@click.option(
+    '-o', '--output-format', default='wide', type=click.Choice(['wide', 'name', 'json', 'yaml']))
+@click.pass_context
+def get_contexts(ctx, output_format):
+  """Get all contexts"""
+  if output_format == 'wide':
+    formatter = TableFormatter(custom_columns={
+        '': lambda c: '*' if c.name == ctx.obj.current_context else '',
+        'NAME': lambda c: c.name,
+        'USER_ID': lambda c: c.user_id,
+        'API_BASE': lambda c: c.api_base,
+        'PAT': lambda c: pat_display(c.pat),
+    })
+    print(formatter.format(ctx.obj.contexts.values(), fmt="plain"))
+  elif output_format == 'name':
+    print('\n'.join(ctx.obj.contexts))
+  elif output_format in ('json', 'yaml'):
+    dicts = [v.__dict__ for c, v in ctx.obj.contexts.items()]
+    for d in dicts:
+      d.pop('pat')
+    if output_format == 'json':
+      print(json.dumps(dicts))
+    elif output_format == 'yaml':
+      print(yaml.safe_dump(dicts))
+
+
+@config.command(['use'])
+@click.argument('context-name', type=str)
+@click.pass_context
+def use_context(ctx, context_name):
+  """Set the current context"""
+  if context_name not in ctx.obj.contexts:
+    raise click.UsageError('Context not found')
+  ctx.obj.current_context = context_name
+  ctx.obj.to_yaml()
+  print(f'Set {context_name} as the current context')
+
+
+@config.command(['cat'])
+@click.option('-o', '--output-format', default='yaml', type=click.Choice(['yaml', 'json']))
+@click.pass_obj
+def dump(ctx_obj, output_format):
+  """Dump the configuration to stdout"""
+  if output_format == 'yaml':
+    yaml.safe_dump(ctx_obj.to_dict(), sys.stdout)
+  else:
+    json.dump(ctx_obj.to_dict(), sys.stdout, indent=2)
+
+
+@cli.command()
+@click.argument('api_url', default="https://api.clarifai.com")
+@click.option('--user_id', required=False, help='User ID')
+@click.pass_context
+def login(ctx, api_url, user_id):
+  """Login command to set PAT and other configurations."""
+
+  name = input('context name (default: "default"): ')
+  user_id = user_id if user_id is not None else input('user id: ')
+  pat = input_or_default(
+      'personal access token value (default: "ENVVAR" to get our of env var rather than config): ',
+      'ENVVAR')
+
+  context = Context(
+      name,
+      CLARIFAI_API_BASE=api_url,
+      CLARIFAI_USER_ID=user_id,
+      CLARIFAI_PAT=pat,
+  )
+
+  if context.name == '':
+    context.name = 'default'
+
+  ctx.obj.contexts[context.name] = context
+  ctx.obj.current_context = context.name
+
+  ctx.obj.to_yaml()
+
+
+@cli.group(cls=AliasedGroup)
+def context():
+  """Manage contexts"""
+
+
+def pat_display(pat):
+  return pat[:5] + "****"
+
+
+@context.command(['ls'])
+@click.pass_context
+def list(ctx):
+  """List available contexts"""
+  formatter = TableFormatter(custom_columns={
+      '': lambda c: '*' if c.name == ctx.obj.current_context else '',
+      'NAME': lambda c: c.name,
+      'USER_ID': lambda c: c.user_id,
+      'API_BASE': lambda c: c.api_base,
+      'PAT': lambda c: pat_display(c.pat)
+  })
+  print(formatter.format(ctx.obj.contexts.values(), fmt="plain"))
+
+
+def input_or_default(prompt, default):
+  value = input(prompt)
+  return value if value else default
+
+
+@context.command()
+@click.argument('name')
+@click.option('--user-id', required=False, help='User ID')
+@click.option('--base-url', required=False, help='Base URL')
+@click.option('--pat', required=False, help='Personal access token')
+@click.pass_context
+def create(
+    ctx,
+    name,
+    user_id=None,
+    base_url=None,
+    pat=None,
+):
+  """Create a new context"""
+  if name in ctx.obj.contexts:
+    print(f'{name} already exists')
+    exit(1)
+  if not user_id:
+    user_id = input('user id: ')
+  if not base_url:
+    base_url = input_or_default('base url (default: https://api.clarifai.com): ',
+                                'https://api.clarifai.com')
+  if not pat:
+    pat = input_or_default(
+        'personal access token value (default: "ENVVAR" to get our of env var rather than config): ',
+        'ENVVAR')
+
+  context = Context(name, CLARIFAI_USER_ID=user_id, CLARIFAI_API_BASE=base_url, CLARIFAI_PAT=pat)
+  ctx.obj.contexts[context.name] = context
+  ctx.obj.to_yaml()
+
+
+# write a click command to delete a context
+@context.command(['rm'])
+@click.argument('name')
+@click.pass_context
+def delete(ctx, name):
+  """Delete a context"""
+  if name not in ctx.obj.contexts:
+    print(f'{name} is not a valid context')
+    exit(1)
+  ctx.obj.contexts.pop(name)
+  ctx.obj.to_yaml()
+  print(f'{name} deleted')
+
+
+@context.command()
+@click.argument('name', type=str)
+@click.pass_context
+def use(ctx, name):
+  """Set the current context"""
+  if name not in ctx.obj.contexts:
+    raise click.UsageError('Context not found')
+  ctx.obj.current_context = name
+  ctx.obj.to_yaml()
+  print(f'Set {name} as the current context')
+
+
+@cli.command()
+@click.argument('script', type=str)
+@click.option('--context', type=str, help='Context to use')
+@click.pass_context
+def run(ctx, script, context=None):
+  """Execute a script with the current context's environment"""
+  context = ctx.obj.current if not context else context
+  cmd = f'CLARIFAI_USER_ID={context.user_id} CLARIFAI_API_BASE={context.base_url} CLARIFAI_PAT={context.pat} '
+  cmd += ' '.join([f'{k}={v}' for k, v in context.env.items()])
+  cmd += f' {script}'
+  os.system(cmd)
+
+
+# Import the CLI commands to register them
+load_command_modules()
+
+
+def main():
+  cli()

clarifai-11.2.3rc8/clarifai/cli/compute_cluster.py

@@ -0,0 +1,61 @@
+import click
+
+from clarifai.cli.base import cli
+from clarifai.utils.cli import AliasedGroup, display_co_resources, validate_context
+
+
+@cli.group(['computecluster', 'cc'], cls=AliasedGroup)
+def computecluster():
+  """Manage Compute Clusters: create, delete, list"""
+
+
+@computecluster.command(['c'])
+@click.argument('compute_cluster_id')
+@click.option(
+    '--config',
+    type=click.Path(exists=True),
+    required=True,
+    help='Path to the compute cluster config file.')
+@click.pass_context
+def create(ctx, compute_cluster_id, config):
+  """Create a new Compute Cluster with the given config file."""
+  from clarifai.client.user import User
+  validate_context(ctx)
+  user = User(
+      user_id=ctx.obj.current.user_id, pat=ctx.obj.current.pat, base_url=ctx.obj.current.api_base)
+  if compute_cluster_id:
+    user.create_compute_cluster(config, compute_cluster_id=compute_cluster_id)
+  else:
+    user.create_compute_cluster(config)
+
+
+@computecluster.command(['ls'])
+@click.option('--page_no', required=False, help='Page number to list.', default=1)
+@click.option('--per_page', required=False, help='Number of items per page.', default=16)
+@click.pass_context
+def list(ctx, page_no, per_page):
+  """List all compute clusters for the user."""
+  from clarifai.client.user import User
+  validate_context(ctx)
+  user = User(
+      user_id=ctx.obj.current.user_id, pat=ctx.obj.current.pat, base_url=ctx.obj.current.api_base)
+  response = user.list_compute_clusters(page_no, per_page)
+  display_co_resources(
+      response,
+      custom_columns={
+          'ID': lambda c: c.id,
+          'USER_ID': lambda c: c.user_id,
+          'DESCRIPTION': lambda c: c.description,
+      })
+
+
+@computecluster.command(['rm'])
+@click.argument('compute_cluster_id')
+@click.pass_context
+def delete(ctx, compute_cluster_id):
+  """Deletes a compute cluster for the user."""
+  from clarifai.client.user import User
+  validate_context(ctx)
+  user = User(
+      user_id=ctx.obj.current.user_id, pat=ctx.obj.current.pat, base_url=ctx.obj.current.api_base)
+  user.delete_compute_clusters([compute_cluster_id])

clarifai-11.2.3rc8/clarifai/cli/deployment.py

@@ -0,0 +1,112 @@
+import click
+
+from clarifai.cli.base import cli
+from clarifai.utils.cli import AliasedGroup, display_co_resources, from_yaml, validate_context
+
+
+@cli.group(['deployment', 'dp'], cls=AliasedGroup)
+def deployment():
+  """Manage Deployments: create, delete, list"""
+
+
+@deployment.command(['c'])
+@click.argument('nodepool_id')
+@click.argument('deployment_id')
+@click.option(
+    '--config',
+    type=click.Path(exists=True),
+    required=True,
+    help='Path to the deployment config file.')
+@click.pass_context
+def create(ctx, nodepool_id, deployment_id, config):
+  """Create a new Deployment with the given config file."""
+
+  from clarifai.client.nodepool import Nodepool
+  validate_context(ctx)
+  if not nodepool_id:
+    deployment_config = from_yaml(config)
+    nodepool_id = deployment_config['deployment']['nodepools'][0]['id']
+
+  nodepool = Nodepool(
+      nodepool_id=nodepool_id,
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
+  if deployment_id:
+    nodepool.create_deployment(config, deployment_id=deployment_id)
+  else:
+    nodepool.create_deployment(config)
+
+
+@deployment.command(['ls'])
+@click.argument('nodepool_id', default="")
+@click.option('--page_no', required=False, help='Page number to list.', default=1)
+@click.option('--per_page', required=False, help='Number of items per page.', default=16)
+@click.pass_context
+def list(ctx, nodepool_id, page_no, per_page):
+  """List all deployments for the nodepool."""
+  from clarifai.client.compute_cluster import ComputeCluster
+  from clarifai.client.nodepool import Nodepool
+  from clarifai.client.user import User
+
+  validate_context(ctx)
+  if nodepool_id:
+    nodepool = Nodepool(
+        nodepool_id=nodepool_id,
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    response = nodepool.list_deployments(page_no=page_no, per_page=per_page)
+  else:
+    user = User(
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    ccs = user.list_compute_clusters(page_no, per_page)
+    nps = []
+    for cc in ccs:
+      compute_cluster = ComputeCluster(
+          compute_cluster_id=cc.id,
+          user_id=ctx.obj.current.user_id,
+          pat=ctx.obj.current.pat,
+          base_url=ctx.obj.current.api_base)
+      nps.extend([i for i in compute_cluster.list_nodepools(page_no, per_page)])
+    response = []
+    for np in nps:
+      nodepool = Nodepool(
+          nodepool_id=np.id,
+          user_id=ctx.obj.current.user_id,
+          pat=ctx.obj.current.pat,
+          base_url=ctx.obj.current.api_base)
+      response.extend([i for i in nodepool.list_deployments(page_no=page_no, per_page=per_page)])
+
+  display_co_resources(
+      response,
+      custom_columns={
+          'ID': lambda c: c.id,
+          'USER_ID': lambda c: c.user_id,
+          'COMPUTE_CLUSTER_ID': lambda c: c.nodepools[0].compute_cluster.id,
+          'NODEPOOL_ID': lambda c: c.nodepools[0].id,
+          'MODEL_USER_ID': lambda c: c.worker.model.user_id,
+          'MODEL_APP_ID': lambda c: c.worker.model.app_id,
+          'MODEL_ID': lambda c: c.worker.model.id,
+          'MODEL_VERSION_ID': lambda c: c.worker.model.model_version.id,
+          'DESCRIPTION': lambda c: c.description,
+      })
+
+
+@deployment.command(['rm'])
+@click.argument('nodepool_id')
+@click.argument('deployment_id')
+@click.pass_context
+def delete(ctx, nodepool_id, deployment_id):
+  """Deletes a deployment for the nodepool."""
+  from clarifai.client.nodepool import Nodepool
+
+  validate_context(ctx)
+  nodepool = Nodepool(
+      nodepool_id=nodepool_id,
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
+  nodepool.delete_deployments([deployment_id])

{clarifai-11.2.3rc6 → clarifai-11.2.3rc8}/clarifai/cli/model.py

@@ -5,7 +5,7 @@ from clarifai.cli.base import cli
 
 @cli.group(['model'])
 def model():
-  """Manage models: upload, test locally, run locally, predict, and more"""
+  """Manage models: upload, test, local dev, predict, etc"""
 
 
 @model.command()
@@ -121,11 +121,14 @@ def signatures(model_path, out_path):
     help=
     'Keep the Docker image after testing the model locally (applicable for container mode). Defaults to False.'
 )
-def test_locally(model_path, keep_env=False, keep_image=False, mode='env'):
-  """Test model locally.
-
-  MODEL_PATH: Path to the model directory. If not specified, the current directory is used by default.
-  """
+@click.option(
+    '--skip_dockerfile',
+    is_flag=True,
+    help=
+    'Flag to skip generating a dockerfile so that you can manually edit an already created dockerfile. Apply for `--mode conatainer`.',
+)
+def test_locally(model_path, keep_env=False, keep_image=False, mode='env', skip_dockerfile=False):
+  """Test model locally."""
   try:
     from clarifai.runners.models import model_run_locally
     if mode == 'env' and keep_image:
@@ -139,7 +142,11 @@ def test_locally(model_path, keep_env=False, keep_image=False, mode='env'):
     elif mode == "container":
       click.echo("Testing model locally inside a container...")
       model_run_locally.main(
-          model_path, inside_container=True, run_model_server=False, keep_image=keep_image)
+          model_path,
+          inside_container=True,
+          run_model_server=False,
+          keep_image=keep_image,
+          skip_dockerfile=skip_dockerfile)
     click.echo("Model tested successfully.")
   except Exception as e:
     click.echo(f"Failed to test model locally: {e}", err=True)
@@ -179,11 +186,14 @@ def test_locally(model_path, keep_env=False, keep_image=False, mode='env'):
     help=
     'Keep the Docker image after testing the model locally (applicable for container mode). Defaults to False.'
 )
-def run_locally(model_path, port, mode, keep_env, keep_image):
-  """Run the model locally and start a gRPC server to serve the model.
-
-  MODEL_PATH: Path to the model directory. If not specified, the current directory is used by default.
-  """
+@click.option(
+    '--skip_dockerfile',
+    is_flag=True,
+    help=
+    'Flag to skip generating a dockerfile so that you can manually edit an already created dockerfile. Apply for `--mode conatainer`.',
+)
+def run_locally(model_path, port, mode, keep_env, keep_image, skip_dockerfile=False):
+  """Run the model locally and start a gRPC server to serve the model."""
   try:
     from clarifai.runners.models import model_run_locally
     if mode == 'env' and keep_image:
@@ -201,7 +211,8 @@ def run_locally(model_path, port, mode, keep_env, keep_image):
           inside_container=True,
           run_model_server=True,
           port=port,
-          keep_image=keep_image)
+          keep_image=keep_image,
+          skip_dockerfile=skip_dockerfile)
     click.echo(f"Model server started locally from {model_path} in {mode} mode.")
   except Exception as e:
     click.echo(f"Failed to starts model server locally: {e}", err=True)
@@ -255,7 +266,8 @@ def predict(ctx, config, model_id, user_id, app_id, model_url, file_path, url, b
   import json
 
   from clarifai.client.model import Model
-  from clarifai.utils.cli import from_yaml
+  from clarifai.utils.cli import from_yaml, validate_context
+  validate_context(ctx)
   if config:
     config = from_yaml(config)
     model_id, user_id, app_id, model_url, file_path, url, bytes, input_type, compute_cluster_id, nodepool_id, deployment_id, inference_params, output_config = (

clarifai-11.2.3rc8/clarifai/cli/nodepool.py

@@ -0,0 +1,112 @@
+import click
+
+from clarifai.cli.base import cli
+from clarifai.utils.cli import (AliasedGroup, display_co_resources, dump_yaml, from_yaml,
+                                validate_context)
+
+
+@cli.group(['nodepool', 'np'], cls=AliasedGroup)
+def nodepool():
+  """Manage Nodepools: create, delete, list"""
+
+
+@nodepool.command(['c'])
+@click.argument('compute_cluster_id')
+@click.argument('nodepool_id')
+@click.option(
+    '--config',
+    type=click.Path(exists=True),
+    required=True,
+    help='Path to the nodepool config file.')
+@click.pass_context
+def create(ctx, compute_cluster_id, nodepool_id, config):
+  """Create a new Nodepool with the given config file."""
+  from clarifai.client.compute_cluster import ComputeCluster
+
+  validate_context(ctx)
+  nodepool_config = from_yaml(config)
+  if not compute_cluster_id:
+    if 'compute_cluster' not in nodepool_config['nodepool']:
+      click.echo(
+          "Please provide a compute cluster ID either in the config file or using --compute_cluster_id flag",
+          err=True)
+      return
+    compute_cluster_id = nodepool_config['nodepool']['compute_cluster']['id']
+  else:
+    if 'compute_cluster' not in nodepool_config['nodepool']:
+      nodepool_config['nodepool']['compute_cluster']['id'] = compute_cluster_id
+      dump_yaml(config, nodepool_config)
+
+  compute_cluster = ComputeCluster(
+      compute_cluster_id=compute_cluster_id,
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
+  if nodepool_id:
+    compute_cluster.create_nodepool(config, nodepool_id=nodepool_id)
+  else:
+    compute_cluster.create_nodepool(config)
+
+
+@nodepool.command(['ls'])
+@click.argument('compute_cluster_id', default="")
+@click.option('--page_no', required=False, help='Page number to list.', default=1)
+@click.option('--per_page', required=False, help='Number of items per page.', default=128)
+@click.pass_context
+def list(ctx, compute_cluster_id, page_no, per_page):
+  """List all nodepools for the user across all compute clusters. If compute_cluster_id is provided
+  it will list only within that compute cluster. """
+  from clarifai.client.compute_cluster import ComputeCluster
+  from clarifai.client.user import User
+
+  validate_context(ctx)
+
+  cc_id = compute_cluster_id
+
+  if cc_id:
+    compute_cluster = ComputeCluster(
+        compute_cluster_id=cc_id,
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    response = compute_cluster.list_nodepools(page_no, per_page)
+  else:
+    user = User(
+        user_id=ctx.obj.current.user_id,
+        pat=ctx.obj.current.pat,
+        base_url=ctx.obj.current.api_base)
+    ccs = user.list_compute_clusters(page_no, per_page)
+    response = []
+    for cc in ccs:
+      compute_cluster = ComputeCluster(
+          compute_cluster_id=cc.id,
+          user_id=ctx.obj.current.user_id,
+          pat=ctx.obj.current.pat,
+          base_url=ctx.obj.current.api_base)
+      response.extend([i for i in compute_cluster.list_nodepools(page_no, per_page)])
+
+  display_co_resources(
+      response,
+      custom_columns={
+          'ID': lambda c: c.id,
+          'USER_ID': lambda c: c.compute_cluster.user_id,
+          'COMPUTE_CLUSTER_ID': lambda c: c.compute_cluster.id,
+          'DESCRIPTION': lambda c: c.description,
+      })
+
+
+@nodepool.command(['rm'])
+@click.argument('compute_cluster_id')
+@click.argument('nodepool_id')
+@click.pass_context
+def delete(ctx, compute_cluster_id, nodepool_id):
+  """Deletes a nodepool for the user."""
+  from clarifai.client.compute_cluster import ComputeCluster
+
+  validate_context(ctx)
+  compute_cluster = ComputeCluster(
+      compute_cluster_id=compute_cluster_id,
+      user_id=ctx.obj.current.user_id,
+      pat=ctx.obj.current.pat,
+      base_url=ctx.obj.current.api_base)
+  compute_cluster.delete_nodepools([nodepool_id])

{clarifai-11.2.3rc6 → clarifai-11.2.3rc8}/clarifai/client/app.py

@@ -629,7 +629,7 @@ class App(Lister, BaseClient):
 
     Args:
         model_id (str): The model ID for the model to interact with.
-        model_version_id (str): The model version ID for the model version to interact with.
+        model_version (Dict): The model version ID for the model version to interact with.
 
     Returns:
         Model: A Model object for the existing model ID.

{clarifai-11.2.3rc6 → clarifai-11.2.3rc8}/clarifai/client/auth/stub.py

@@ -1,5 +1,4 @@
 import itertools
-import logging
 import time
 from concurrent.futures import ThreadPoolExecutor
 
@@ -8,7 +7,7 @@ from clarifai_grpc.grpc.api.status import status_code_pb2
 
 from clarifai.client.auth.helper import ClarifaiAuthHelper
 from clarifai.client.auth.register import RpcCallable, V2Stub
-
+from clarifai.utils.logging import logger
 throttle_status_codes = {
     status_code_pb2.CONN_THROTTLED,
     status_code_pb2.CONN_EXCEED_HOURLY_LIMIT,
@@ -26,7 +25,7 @@ def validate_response(response, attempt, max_attempts):
   def handle_simple_response(response):
     if hasattr(response, 'status') and hasattr(response.status, 'code'):
       if (response.status.code in throttle_status_codes) and attempt < max_attempts:
-        logging.debug('Retrying with status %s' % str(response.status))
+        logger.debug('Retrying with status %s' % str(response.status))
         return None  # Indicates a retry is needed
       else:
         return response
@@ -42,7 +41,7 @@ def validate_response(response, attempt, max_attempts):
         return itertools.chain([validated_response], response)
       return None  # Indicates a retry is needed
     except grpc.RpcError as e:
-      logging.error('Error processing streaming response: %s' % str(e))
+      logger.error('Error processing streaming response: %s' % str(e))
       return None  # Indicates an error
   else:
     # Handle simple response validation
@@ -143,7 +142,7 @@ class _RetryRpcCallable(RpcCallable):
           return v
       except grpc.RpcError as e:
         if (e.code() in retry_codes_grpc) and attempt < self.max_attempts:
-          logging.debug('Retrying with status %s' % e.code())
+          logger.debug('Retrying with status %s' % e.code())
         else:
           raise