[EAGLE-7214 ]: Add configuration options for health check port in local-runner

clarifai/cli/model.py

@@ -1,8 +1,10 @@
 import json
 import os
 import shutil
+import socket
 import subprocess
 import tempfile
+from contextlib import closing
 from typing import Any, Dict, Optional
 
 import click
@@ -54,6 +56,19 @@
 )
 
 
+def find_available_port(start_port=8080):
+    """Find the first available port starting from start_port."""
+    port = start_port
+    while port <= 65535:
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+            try:
+                sock.bind(('', port))
+                return port
+            except OSError:
+                port += 1
+    raise RuntimeError("No available port found")
+
+
 def _select_context(ctx_config: Config) -> Optional[Context]:
     contexts_map = getattr(ctx_config, "contexts", {}) or {}
     if not contexts_map:
@@ -1165,8 +1180,35 @@ def run_locally(ctx, model_path, port, mode, keep_env, keep_image, skip_dockerfi
     is_flag=True,
     help='Keep the Docker image after testing the model locally (applicable for container mode). Defaults to False.',
 )
+@click.option(
+    '--health-check-port',
+    type=int,
+    default=8080,
+    show_default=True,
+    help='The port to run the health check server on. Defaults to 8080.',
+)
+@click.option(
+    '--disable-health-check',
+    is_flag=True,
+    help='Disable the health check server.',
+)
+@click.option(
+    '--auto-find-health-check-port',
+    is_flag=True,
+    help='Automatically find an available port starting from --health-check-port.',
+)
 @click.pass_context
-def local_runner(ctx, model_path, pool_size, suppress_toolkit_logs, mode, keep_image):
+def local_runner(
+    ctx,
+    model_path,
+    pool_size,
+    suppress_toolkit_logs,
+    mode,
+    keep_image,
+    health_check_port,
+    disable_health_check,
+    auto_find_health_check_port,
+):
     """Run the model as a local runner to help debug your model connected to the API or to
       leverage local compute resources manually. This relies on many variables being present in the env
       of the currently selected context. If they are not present then default values will be used to
@@ -1214,6 +1256,11 @@ def local_runner(ctx, model_path, pool_size, suppress_toolkit_logs, mode, keep_i
     builder = ModelBuilder(model_path, download_validation_only=True)
     manager = ModelRunLocally(model_path, model_builder=builder)
 
+    if disable_health_check:
+        health_check_port = -1
+    elif auto_find_health_check_port:
+        health_check_port = find_available_port(health_check_port)
+
     port = 8080
     if mode == "env":
         manager.create_temp_venv()
@@ -1640,6 +1687,7 @@ def print_code_snippet():
         "base_url": ctx.obj.current.api_base,
         "pat": ctx.obj.current.pat,
         "context": ctx.obj.current,
+        "health_check_port": health_check_port,
     }
 
     if mode == "container":

clarifai/runners/models/model_run_locally.py

@@ -426,6 +426,8 @@ def run_docker_container(
                         str(kwargs.get("pat", None)),
                         "--num_threads",
                         str(kwargs.get("num_threads", 0)),
+                        "--health_check_port",
+                        str(kwargs.get("health_check_port", 8080)),
                     ]
                 )
             else:

clarifai/runners/models/model_runner.py

@@ -95,7 +95,7 @@ def __init__(
         HealthProbeRequestHandler.is_ready = True
         HealthProbeRequestHandler.is_startup = True
 
-        if health_check_port is not None:
+        if health_check_port is not None and health_check_port > 0:
             start_health_server_thread(port=health_check_port, address='')
 
     def get_runner_item_output_for_status(

clarifai/runners/server.py

@@ -108,6 +108,12 @@ def main():
         default=0,
         help='The number of threads for the runner to use (default: 0, which means read from config.yaml).',
     )
+    parser.add_argument(
+        '--health_check_port',
+        type=int,
+        default=8080,
+        help="The port to host the health check server at. Set to 0 or -1 to disable.",
+    )
 
     parsed_args = parser.parse_args()
     server = ModelServer(parsed_args.model_path)
@@ -137,6 +143,7 @@ def main():
         max_msg_length=parsed_args.max_msg_length,
         enable_tls=parsed_args.enable_tls,
         grpc=parsed_args.grpc,
+        health_check_port=parsed_args.health_check_port,
     )
 
 
@@ -273,6 +280,7 @@ def serve(
         max_msg_length=1024 * 1024 * 1024,
         enable_tls=False,
         grpc=False,
+        health_check_port=8080,
         user_id: Optional[str] = os.environ.get("CLARIFAI_USER_ID", None),
         compute_cluster_id: Optional[str] = os.environ.get("CLARIFAI_COMPUTE_CLUSTER_ID", None),
         nodepool_id: Optional[str] = os.environ.get("CLARIFAI_NODEPOOL_ID", None),
@@ -304,6 +312,7 @@ def serve(
                 base_url,
                 pat,
                 num_threads,
+                health_check_port,
             )
 
     def start_servicer(self, port, pool_size, max_queue_size, max_msg_length, enable_tls):
@@ -335,6 +344,7 @@ def start_runner(
         base_url,
         pat,
         num_threads,
+        health_check_port,
     ):
         # initialize the Runner class. This is what the user implements.
         assert compute_cluster_id is not None, "compute_cluster_id must be set for the runner."
@@ -350,6 +360,7 @@ def start_runner(
             base_url=base_url,
             pat=pat,
             num_parallel_polls=num_threads,
+            health_check_port=health_check_port,
         )
 
         self._runner.start()  # start the runner to fetch work from the API.