Merge branch 'master' of github.com:Clarifai/clarifai-python

Author: David Eigen

clarifai/cli/model.py

@@ -171,6 +171,18 @@ def run_locally(model_path, port, mode, keep_env, keep_image):
     click.echo(f"Failed to starts model server locally: {e}", err=True)
 
 
+@model.command()
+@click.option(
+    '--model_path',
+    type=click.Path(exists=True),
+    required=True,
+    help='Path to the model directory.')
+def local_dev(model_path):
+  """Run the model as a local dev runner to help debug your model connected to the API. You must set several envvars such as CLARIFAI_PAT, CLARIFAI_RUNNER_ID, CLARIFAI_NODEPOOL_ID, CLARIFAI_COMPUTE_CLUSTER_ID. """
+  from clarifai.runners.server import serve
+  serve(model_path)
+
+
 @model.command()
 @click.option(
     '--config',

clarifai/runners/models/model_run_locally.py

@@ -482,8 +482,10 @@ def main(model_path,
     sys.exit(1)
   manager = ModelRunLocally(model_path)
   # get whatever stage is in config.yaml to force download now
+  # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
   _, _, _, when = manager.builder._validate_config_checkpoints()
-  manager.builder.download_checkpoints(stage=when)
+  manager.builder.download_checkpoints(
+      stage=when, checkpoint_path_override=manager.builder.checkpoint_path)
   if inside_container:
     if not manager.is_docker_installed():
       sys.exit(1)

clarifai/runners/server.py

@@ -68,30 +68,43 @@ def main():
 
   parsed_args = parser.parse_args()
 
-  builder = ModelBuilder(parsed_args.model_path, download_validation_only=True)
+  serve(parsed_args.model_path, parsed_args.port, parsed_args.pool_size,
+        parsed_args.max_queue_size, parsed_args.max_msg_length, parsed_args.enable_tls,
+        parsed_args.grpc)
+
+
+def serve(model_path,
+          port=8000,
+          pool_size=32,
+          max_queue_size=10,
+          max_msg_length=1024 * 1024 * 1024,
+          enable_tls=False,
+          grpc=False):
+
+  builder = ModelBuilder(model_path, download_validation_only=True)
 
   model = builder.create_model_instance()
 
   # Setup the grpc server for local development.
-  if parsed_args.grpc:
+  if grpc:
 
     # initialize the servicer with the runner so that it gets the predict(), generate(), stream() classes.
     servicer = ModelServicer(model)
 
     server = GRPCServer(
         futures.ThreadPoolExecutor(
-            max_workers=parsed_args.pool_size,
+            max_workers=pool_size,
             thread_name_prefix="ServeCalls",
         ),
-        parsed_args.max_msg_length,
-        parsed_args.max_queue_size,
+        max_msg_length,
+        max_queue_size,
     )
-    server.add_port_to_server('[::]:%s' % parsed_args.port, parsed_args.enable_tls)
+    server.add_port_to_server('[::]:%s' % port, enable_tls)
 
     service_pb2_grpc.add_V2Servicer_to_server(servicer, server)
     server.start()
-    logger.info("Started server on port %s", parsed_args.port)
-    logger.info(f"Access the model at http://localhost:{parsed_args.port}")
+    logger.info("Started server on port %s", port)
+    logger.info(f"Access the model at http://localhost:{port}")
     server.wait_for_termination()
   else:  # start the runner with the proper env variables and as a runner protocol.
 

tests/runners/test_download_checkpoints.py

@@ -5,7 +5,6 @@
 import pytest
 
 from clarifai.runners.models.model_builder import ModelBuilder
-from clarifai.runners.utils.const import DEFAULT_RUNTIME_DOWNLOAD_PATH
 from clarifai.runners.utils.loader import HuggingFaceLoader
 
 MODEL_ID = "timm/mobilenetv3_small_100.lamb_in1k"
@@ -62,16 +61,20 @@ def test_download_checkpoints(dummy_runner_models_dir):
   model_builder = ModelBuilder(model_folder_path, download_validation_only=True)
   # defaults to runtime stage which matches config.yaml not having a when field.
   # get whatever stage is in config.yaml to force download now
+  # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
   _, _, _, when = model_builder._validate_config_checkpoints()
-  checkpoint_dir = model_builder.download_checkpoints(stage=when)
-  assert checkpoint_dir == DEFAULT_RUNTIME_DOWNLOAD_PATH
+  checkpoint_dir = model_builder.download_checkpoints(
+      stage=when, checkpoint_path_override=model_builder.checkpoint_path)
+  assert checkpoint_dir == model_builder.checkpoint_path
 
   # This doesn't have when in it's config.yaml so build.
   model_folder_path = os.path.join(os.path.dirname(__file__), "hf_mbart_model")
   model_builder = ModelBuilder(model_folder_path, download_validation_only=True)
   # defaults to runtime stage which matches config.yaml not having a when field.
   # get whatever stage is in config.yaml to force download now
+  # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
   _, _, _, when = model_builder._validate_config_checkpoints()
-  checkpoint_dir = model_builder.download_checkpoints(stage=when)
+  checkpoint_dir = model_builder.download_checkpoints(
+      stage=when, checkpoint_path_override=model_builder.checkpoint_path)
   assert checkpoint_dir == os.path.join(
       os.path.dirname(__file__), "hf_mbart_model", "1", "checkpoints")

tests/runners/test_model_run_locally-container.py

@@ -81,8 +81,12 @@ def test_hf_docker_build_and_test_container(hf_model_run_locally):
   This test will be skipped if Docker is not installed.
   """
 
+  # get whatever stage is in config.yaml to force download now
+  # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
   # Download the checkpoints for the model
-  hf_model_run_locally.builder.download_checkpoints()
+  _, _, _, when = hf_model_run_locally.builder._validate_config_checkpoints()
+  hf_model_run_locally.builder.download_checkpoints(
+      stage=when, checkpoint_path_override=hf_model_run_locally.builder.checkpoint_path)
 
   # Test if Docker is installed
   assert hf_model_run_locally.is_docker_installed(), "Docker not installed, skipping."

tests/runners/test_model_run_locally.py

@@ -171,8 +171,10 @@ def test_hf_test_model_success(hf_model_run_locally):
   This calls the script's test_model method, which runs a subprocess.
   """
   # get whatever stage is in config.yaml to force download now
+  # also always write to where upload/build wants to, not the /tmp folder that runtime stage uses
   _, _, _, when = hf_model_run_locally.builder._validate_config_checkpoints()
-  hf_model_run_locally.builder.download_checkpoints(stage=when)
+  hf_model_run_locally.builder.download_checkpoints(
+      stage=when, checkpoint_path_override=hf_model_run_locally.builder.checkpoint_path)
   hf_model_run_locally.create_temp_venv()
   hf_model_run_locally.install_requirements()