clarifai 10.10.1__py3-none-any.whl → 10.11.1__py3-none-any.whl

clarifai/runners/models/model_run_locally.py

@@ -1,10 +1,13 @@
+import hashlib
 import importlib.util
 import inspect
 import os
 import shutil
+import signal
 import subprocess
 import sys
 import tempfile
+import time
 import traceback
 import venv
 
@@ -13,6 +16,7 @@ from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 from clarifai_protocol import BaseRunner
 
 from clarifai.runners.models.model_upload import ModelUploader
+from clarifai.runners.utils.url_fetcher import ensure_urls_downloaded
 from clarifai.utils.logging import logger
 
 
@@ -22,19 +26,36 @@ class ModelRunLocally:
     self.model_path = model_path
     self.requirements_file = os.path.join(self.model_path, "requirements.txt")
 
+    # ModelUploader contains multiple useful methods to interact with the model
+    self.uploader = ModelUploader(self.model_path)
+    self.config = self.uploader.config
+
+  def _requirements_hash(self):
+    """Generate a hash of the requirements file."""
+    with open(self.requirements_file, "r") as f:
+      return hashlib.md5(f.read().encode('utf-8')).hexdigest()
+
   def create_temp_venv(self):
     """Create a temporary virtual environment."""
-    logger.info("Creating temporary virtual environment...")
-    temp_dir = tempfile.mkdtemp()
+    requirements_hash = self._requirements_hash()
+
+    temp_dir = os.path.join(tempfile.gettempdir(), str(requirements_hash))
     venv_dir = os.path.join(temp_dir, "venv")
-    venv.create(venv_dir, with_pip=True)
+
+    if os.path.exists(temp_dir):
+      logger.info(f"Using previous virtual environment at {temp_dir}")
+      use_existing_venv = True
+    else:
+      logger.info("Creating temporary virtual environment...")
+      use_existing_venv = False
+      venv.create(venv_dir, with_pip=True)
+      logger.info(f"Created temporary virtual environment at {venv_dir}")
 
     self.venv_dir = venv_dir
     self.temp_dir = temp_dir
     self.python_executable = os.path.join(venv_dir, "bin", "python")
 
-    logger.info(f"Created temporary virtual environment at {venv_dir}")
-    return venv_dir, temp_dir
+    return use_existing_venv
 
   def install_requirements(self):
     """Install the dependencies from requirements.txt and Clarifai."""
@@ -94,31 +115,91 @@ class ModelRunLocally:
                 text=resources_pb2.Text(raw="How many people live in new york?"),
                 image=resources_pb2.Image(url="https://samples.clarifai.com/metro-north.jpg"),
                 audio=resources_pb2.Audio(url="https://samples.clarifai.com/GoodMorning.wav"),
+                video=resources_pb2.Video(url="https://samples.clarifai.com/beer.mp4"),
             ))
         ],
     )
 
+  def _build_stream_request(self):
+    request = self._build_request()
+    for i in range(1):
+      yield request
+
   def _run_model_inference(self, runner):
     """Perform inference using the runner."""
     request = self._build_request()
+    stream_request = self._build_stream_request()
 
+    ensure_urls_downloaded(request)
+    predict_response = None
+    generate_response = None
+    stream_response = None
     try:
-      return runner.predict(request)
+      predict_response = runner.predict(request)
+    except NotImplementedError:
+      logger.info("Model does not implement predict() method.")
     except Exception as e:
       logger.error(f"Model Prediction failed: {e}")
       traceback.print_exc()
-      return service_pb2.MultiOutputResponse(status=status_pb2.Status(
+      predict_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
           code=status_code_pb2.MODEL_PREDICTION_FAILED,
           description="Prediction failed",
           details="",
           internal_details=str(e),
       ))
 
+    if predict_response:
+      if predict_response.outputs[0].status.code != status_code_pb2.SUCCESS:
+        logger.error(f"Moddel Prediction failed: {predict_response}")
+      else:
+        logger.info(f"Model Prediction succeeded: {predict_response}")
+
+    try:
+      generate_response = runner.generate(request)
+    except NotImplementedError:
+      logger.info("Model does not implement generate() method.")
+    except Exception as e:
+      logger.error(f"Model Generation failed: {e}")
+      traceback.print_exc()
+      generate_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
+          code=status_code_pb2.MODEL_GENERATION_FAILED,
+          description="Generation failed",
+          details="",
+          internal_details=str(e),
+      ))
+
+    if generate_response:
+      generate_first_res = next(generate_response)
+      if generate_first_res.outputs[0].status.code != status_code_pb2.SUCCESS:
+        logger.error(f"Moddel Prediction failed: {generate_first_res}")
+      else:
+        logger.info(
+            f"Model Prediction succeeded for generate and first response: {generate_first_res}")
+
+    try:
+      stream_response = runner.stream(stream_request)
+    except NotImplementedError:
+      logger.info("Model does not implement stream() method.")
+    except Exception as e:
+      logger.error(f"Model Stream failed: {e}")
+      traceback.print_exc()
+      stream_response = service_pb2.MultiOutputResponse(status=status_pb2.Status(
+          code=status_code_pb2.MODEL_STREAM_FAILED,
+          description="Stream failed",
+          details="",
+          internal_details=str(e),
+      ))
+
+    if stream_response:
+      stream_first_res = next(stream_response)
+      if stream_first_res.outputs[0].status.code != status_code_pb2.SUCCESS:
+        logger.error(f"Moddel Prediction failed: {stream_first_res}")
+      else:
+        logger.info(
+            f"Model Prediction succeeded for stream and first response: {stream_first_res}")
+
   def _run_test(self):
     """Test the model locally by making a prediction."""
-    # validate that we have checkpoints downloaded before constructing MyRunner
-    uploader = ModelUploader(self.model_path)
-    uploader.download_checkpoints()
     # construct MyRunner which will call load_model()
     MyRunner = self._get_model_runner()
     runner = MyRunner(
@@ -127,13 +208,8 @@ class ModelRunLocally:
         compute_cluster_id="n/a",
         user_id="n/a",
     )
-
     # send an inference.
-    response = self._run_model_inference(runner)
-    if response.outputs[0].status.code != status_code_pb2.SUCCESS:
-      logger.error(f"Moddel Prediction failed: {response}")
-    else:
-      logger.info(f"Model Prediction succeeded: {response}")
+    self._run_model_inference(runner)
 
   def test_model(self):
     """Test the model by running it locally in the virtual environment."""
@@ -143,31 +219,258 @@ class ModelRunLocally:
         f"import sys; sys.path.append('{os.path.dirname(os.path.abspath(__file__))}'); "
         f"from model_run_locally import ModelRunLocally; ModelRunLocally('{self.model_path}')._run_test()",
     ]
+    process = None
     try:
       logger.info("Testing the model locally...")
-      subprocess.check_call(command)
-      logger.info("Model tested successfully!")
+      process = subprocess.Popen(command)
+      # Wait for the process to complete
+      process.wait()
+      if process.returncode == 0:
+        logger.info("Model tested successfully!")
+      if process.returncode != 0:
+        raise subprocess.CalledProcessError(process.returncode, command)
     except subprocess.CalledProcessError as e:
       logger.error(f"Error testing the model: {e}")
       sys.exit(1)
+    except Exception as e:
+      logger.error(f"Unexpected error: {e}")
+      sys.exit(1)
+    finally:
+      # After the function runs, check if the process is still running
+      if process and process.poll() is None:
+        logger.info("Process is still running. Terminating process.")
+        process.terminate()
+        try:
+          process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+          logger.info("Process did not terminate gracefully. Killing process.")
+          # Kill the process if it doesn't terminate after 5 seconds
+          process.kill()
 
   # run the model server
-  def run_model_server(self):
+  def run_model_server(self, port=8080):
     """Run the Clarifai Runners's model server."""
 
     command = [
         self.python_executable, "-m", "clarifai.runners.server", "--model_path", self.model_path,
-        "--start_dev_server"
+        "--start_dev_server", "--port",
+        str(port)
     ]
     try:
-      logger.info(f"Starting model server with model at {self.model_path}...")
+      logger.info(
+          f"Starting model server at localhost:{port} with the model at {self.model_path}...")
       subprocess.check_call(command)
-      logger.info("Model server started successfully!")
+      logger.info("Model server started successfully and running at localhost:{port}")
     except subprocess.CalledProcessError as e:
       logger.error(f"Error running model server: {e}")
       self.clean_up()
       sys.exit(1)
 
+  def _docker_hash(self):
+    """Generate a hash of the combined requirements file and Dockefile"""
+    with open(self.requirements_file, "r") as f:
+      requirements_hash = hashlib.md5(f.read().encode('utf-8')).hexdigest()
+    with open(os.path.join(self.model_path, "Dockerfile"), "r") as f:
+      dockerfile_hash = hashlib.md5(f.read().encode('utf-8')).hexdigest()
+
+    return hashlib.md5(f"{requirements_hash}{dockerfile_hash}".encode('utf-8')).hexdigest()
+
+  def is_docker_installed(self):
+    """Checks if Docker is installed on the system."""
+    try:
+      logger.info("Checking if Docker is installed...")
+      subprocess.run(["docker", "--version"], check=True)
+      return True
+    except subprocess.CalledProcessError:
+      logger.error(
+          "Docker is not installed! Please install Docker to run the model in a container.")
+      return False
+
+  def build_docker_image(
+      self,
+      image_name="model_image",
+  ):
+    """Build the docker image using the Dockerfile in the model directory."""
+    try:
+      logger.info(f"Building docker image from Dockerfile in {self.model_path}...")
+
+      # since we don't want to copy the model directory into the container, we need to modify the Dockerfile and comment out the COPY instruction
+      dockerfile_path = os.path.join(self.model_path, "Dockerfile")
+      # Read the Dockerfile
+      with open(dockerfile_path, 'r') as file:
+        lines = file.readlines()
+
+      # Comment out the COPY instruction that copies the current folder
+      modified_lines = []
+      for line in lines:
+        if 'COPY .' in line and '/app/model_dir/main' in line:
+          modified_lines.append(f'# {line}')
+        else:
+          modified_lines.append(line)
+
+      # Create a temporary directory to store the modified Dockerfile
+      with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dockerfile_path = os.path.join(temp_dir, "Dockerfile.temp")
+
+        # Write the modified Dockerfile to the temporary file
+        with open(temp_dockerfile_path, 'w') as file:
+          file.writelines(modified_lines)
+
+        # Build the Docker image using the temporary Dockerfile
+        subprocess.check_call(
+            ['docker', 'build', '-t', image_name, '-f', temp_dockerfile_path, self.model_path])
+      logger.info(f"Docker image '{image_name}' built successfully!")
+    except subprocess.CalledProcessError as e:
+      logger.info(f"Error occurred while building the Docker image: {e}")
+      sys.exit(1)
+
+  def docker_image_exists(self, image_name):
+    """Check if the Docker image exists."""
+    try:
+      logger.info(f"Checking if Docker image '{image_name}' exists...")
+      subprocess.run(["docker", "inspect", image_name], check=True)
+      logger.info(f"Docker image '{image_name}' exists!")
+      return True
+    except subprocess.CalledProcessError:
+      logger.info(f"Docker image '{image_name}' does not exist!")
+      return False
+
+  def run_docker_container(self,
+                           image_name,
+                           container_name="clarifai-model-container",
+                           port=8080,
+                           env_vars=None):
+    """Runs a Docker container from the specified image."""
+    try:
+      logger.info(f"Running Docker container '{container_name}' from image '{image_name}'...")
+      # Base docker run command
+      cmd = [
+          "docker", "run", "--name", container_name, '--rm', "--gpus", "all", "--network", "host"
+      ]
+      # Add volume mappings
+      cmd.extend(["-v", f"{self.model_path}:/app/model_dir/main"])
+      # Add environment variables
+      if env_vars:
+        for key, value in env_vars.items():
+          cmd.extend(["-e", f"{key}={value}"])
+      # Add the image name
+      cmd.append(image_name)
+      # update the CMD to run the server
+      cmd.extend(
+          ["--model_path", "/app/model_dir/main", "--start_dev_server", "--port",
+           str(port)])
+      # Run the container
+      process = subprocess.Popen(cmd,)
+      logger.info(
+          f"Docker container '{container_name}' is running successfully! access the model at http://localhost:{port}"
+      )
+
+      # Function to handle Ctrl+C (SIGINT) gracefully
+      def signal_handler(sig, frame):
+        logger.info(f"Stopping Docker container '{container_name}'...")
+        subprocess.run(["docker", "stop", container_name], check=True)
+        process.terminate()
+        logger.info(f"Docker container '{container_name}' stopped successfully!")
+        time.sleep(1)
+        sys.exit(0)
+
+      # Register the signal handler for SIGINT (Ctrl+C)
+      signal.signal(signal.SIGINT, signal_handler)
+      # Wait for the process to finish (keeps the container running until it's stopped)
+      process.wait()
+    except subprocess.CalledProcessError as e:
+      logger.info(f"Error occurred while running the Docker container: {e}")
+      sys.exit(1)
+    except Exception as e:
+      logger.info(f"Error occurred while running the Docker container: {e}")
+      sys.exit(1)
+
+  def test_model_container(self,
+                           image_name,
+                           container_name="clarifai-model-container",
+                           env_vars=None):
+    """Test the model inside the Docker container."""
+    try:
+      logger.info("Testing the model inside the Docker container...")
+      # Base docker run command
+      cmd = [
+          "docker", "run", "--name", container_name, '--rm', "--gpus", "all", "--network", "host"
+      ]
+      # update the entrypoint for testing the model
+      cmd.extend(["--entrypoint", "python"])
+      # Add volume mappings
+      cmd.extend(["-v", f"{self.model_path}:/app/model_dir/main"])
+      # Add environment variables
+      if env_vars:
+        for key, value in env_vars.items():
+          cmd.extend(["-e", f"{key}={value}"])
+      # Add the image name
+      cmd.append(image_name)
+      # update the CMD to test the model inside the container
+      cmd.extend([
+          "-c",
+          "from clarifai.runners.models.model_run_locally import ModelRunLocally; ModelRunLocally('/app/model_dir/main')._run_test()"
+      ])
+      # Run the container
+      subprocess.check_call(cmd)
+      logger.info("Model tested successfully!")
+    except subprocess.CalledProcessError as e:
+      logger.error(f"Error testing the model inside the Docker container: {e}")
+      sys.exit(1)
+
+  def container_exists(self, container_name="clarifai-model-container"):
+    """Check if the Docker container exists."""
+    try:
+      # Run docker ps -a to list all containers (running and stopped)
+      result = subprocess.run(
+          ["docker", "ps", "-a", "--filter", f"name={container_name}", "--format", "{{.Names}}"],
+          check=True,
+          capture_output=True,
+          text=True)
+      # If the container name is returned, it exists
+      if result.stdout.strip() == container_name:
+        logger.info(f"Docker container '{container_name}' exists.")
+        return True
+      else:
+        return False
+    except subprocess.CalledProcessError as e:
+      logger.error(f"Error occurred while checking if container exists: {e}")
+      return False
+
+  def stop_docker_container(self, container_name="clarifai-model-container"):
+    """Stop the Docker container if it's running."""
+    try:
+      # Check if the container is running
+      result = subprocess.run(
+          ["docker", "ps", "--filter", f"name={container_name}", "--format", "{{.Names}}"],
+          check=True,
+          capture_output=True,
+          text=True)
+      if result.stdout.strip() == container_name:
+        logger.info(f"Docker container '{container_name}' is running. Stopping it...")
+        subprocess.run(["docker", "stop", container_name], check=True)
+        logger.info(f"Docker container '{container_name}' stopped successfully!")
+    except subprocess.CalledProcessError as e:
+      logger.error(f"Error occurred while stopping the Docker container: {e}")
+
+  def remove_docker_container(self, container_name="clarifai-model-container"):
+    """Remove the Docker container."""
+    try:
+      logger.info(f"Removing Docker container '{container_name}'...")
+      subprocess.run(["docker", "rm", container_name], check=True)
+      logger.info(f"Docker container '{container_name}' removed successfully!")
+    except subprocess.CalledProcessError as e:
+      logger.error(f"Error occurred while removing the Docker container: {e}")
+
+  def remove_docker_image(self, image_name):
+    """Remove the Docker image."""
+    try:
+      logger.info(f"Removing Docker image '{image_name}'...")
+      subprocess.run(["docker", "rmi", image_name], check=True)
+      logger.info(f"Docker image '{image_name}' removed successfully!")
+    except subprocess.CalledProcessError as e:
+      logger.error(f"Error occurred while removing the Docker image: {e}")
+
   def clean_up(self):
     """Clean up the temporary virtual environment."""
     if os.path.exists(self.temp_dir):
@@ -175,16 +478,53 @@ class ModelRunLocally:
       shutil.rmtree(self.temp_dir)
 
 
-def main(model_path, run_model_server=False):
+def main(model_path,
+         run_model_server=False,
+         inside_container=False,
+         port=8080,
+         keep_env=False,
+         keep_image=False):
 
+  if not os.environ['CLARIFAI_PAT']:
+    logger.error(
+        "CLARIFAI_PAT environment variable is not set! Please set your PAT in the 'CLARIFAI_PAT' environment variable."
+    )
+    sys.exit(1)
   manager = ModelRunLocally(model_path)
-  manager.create_temp_venv()
+  manager.uploader.download_checkpoints()
+  if inside_container:
+    if not manager.is_docker_installed():
+      sys.exit(1)
+    manager.uploader.create_dockerfile()
+    image_tag = manager._docker_hash()
+    image_name = f"{manager.config['model']['id']}:{image_tag}"
+    container_name = manager.config['model']['id']
+    if not manager.docker_image_exists(image_name):
+      manager.build_docker_image(image_name=image_name)
+    try:
+      envs = {'CLARIFAI_PAT': os.environ['CLARIFAI_PAT'], 'CLARIFAI_USER_ID': 'n/a'}
+      if run_model_server:
+        manager.run_docker_container(
+            image_name=image_name, container_name=container_name, port=port, env_vars=envs)
+      else:
+        manager.test_model_container(
+            image_name=image_name, container_name=container_name, env_vars=envs)
+    finally:
+      if manager.container_exists(container_name):
+        manager.stop_docker_container(container_name)
+        manager.remove_docker_container(container_name=container_name)
+      if not keep_image:
+        manager.remove_docker_image(image_name)
 
-  try:
-    manager.install_requirements()
-    if run_model_server:
-      manager.run_model_server()
-    else:
-      manager.test_model()
-  finally:
-    manager.clean_up()
+  else:
+    try:
+      use_existing_env = manager.create_temp_venv()
+      if not use_existing_env:
+        manager.install_requirements()
+      if run_model_server:
+        manager.run_model_server(port)
+      else:
+        manager.test_model()
+    finally:
+      if not keep_env:
+        manager.clean_up()

clarifai/runners/models/model_upload.py

@@ -10,6 +10,9 @@ from google.protobuf import json_format
 from rich import print
 
 from clarifai.client import BaseClient
+from clarifai.runners.utils.const import (AVAILABLE_PYTHON_IMAGES, AVAILABLE_TORCH_IMAGES,
+                                          CONCEPTS_REQUIRED_MODEL_TYPE, DEFAULT_PYTHON_VERSION,
+                                          PYTHON_BASE_IMAGE, TORCH_BASE_IMAGE)
 from clarifai.runners.utils.loader import HuggingFaceLoader
 from clarifai.urls.helper import ClarifaiUrlHelper
 from clarifai.utils.logging import logger
@@ -23,48 +26,6 @@ def _clear_line(n: int = 1) -> None:
 
 
 class ModelUploader:
-  DEFAULT_PYTHON_VERSION = 3.11
-  DEFAULT_TORCH_VERSION = '2.4.0'
-  DEFAULT_CUDA_VERSION = '124'
-  # List of available torch images for matrix
-  '''
-        python_version: ['3.8', '3.9', '3.10', '3.11']
-        torch_version: ['2.0.0', '2.1.0', '2.2.0', '2.3.0', '2.4.0', '2.4.1', '2.5.0']
-        cuda_version: ['124']
-  '''
-  AVAILABLE_TORCH_IMAGES = [
-      '2.0.0-py3.8-cuda124',
-      '2.0.0-py3.9-cuda124',
-      '2.0.0-py3.10-cuda124',
-      '2.0.0-py3.11-cuda124',
-      '2.1.0-py3.8-cuda124',
-      '2.1.0-py3.9-cuda124',
-      '2.1.0-py3.10-cuda124',
-      '2.1.0-py3.11-cuda124',
-      '2.2.0-py3.8-cuda124',
-      '2.2.0-py3.9-cuda124',
-      '2.2.0-py3.10-cuda124',
-      '2.2.0-py3.11-cuda124',
-      '2.3.0-py3.8-cuda124',
-      '2.3.0-py3.9-cuda124',
-      '2.3.0-py3.10-cuda124',
-      '2.3.0-py3.11-cuda124',
-      '2.4.0-py3.8-cuda124',
-      '2.4.0-py3.9-cuda124',
-      '2.4.0-py3.10-cuda124',
-      '2.4.0-py3.11-cuda124',
-      '2.4.1-py3.8-cuda124',
-      '2.4.1-py3.9-cuda124',
-      '2.4.1-py3.10-cuda124',
-      '2.4.1-py3.11-cuda124',
-  ]
-  AVAILABLE_PYTHON_IMAGES = ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
-  PYTHON_BASE_IMAGE = 'public.ecr.aws/clarifai-models/python-base:{python_version}'
-  TORCH_BASE_IMAGE = 'public.ecr.aws/clarifai-models/torch:{torch_version}-py{python_version}-cuda{cuda_version}'
-
-  CONCEPTS_REQUIRED_MODEL_TYPE = [
-      'visual-classifier', 'visual-detector', 'visual-segmenter', 'text-classifier'
-  ]
 
   def __init__(self, folder: str):
     self._client = None
@@ -126,7 +87,7 @@ class ModelUploader:
       user_id = model.get('user_id')
       app_id = model.get('app_id')
 
-      base = os.environ.get('CLARIFAI_API_BASE', 'https://api-dev.clarifai.com')
+      base = os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
 
       self._client = BaseClient(user_id=user_id, app_id=app_id, base=base)
 
@@ -233,33 +194,35 @@ class ModelUploader:
     build_info = self.config.get('build_info', {})
     if 'python_version' in build_info:
       python_version = build_info['python_version']
-      if python_version not in self.AVAILABLE_PYTHON_IMAGES:
+      if python_version not in AVAILABLE_PYTHON_IMAGES:
         logger.error(
-            f"Python version {python_version} not supported, please use one of the following versions: {self.AVAILABLE_PYTHON_IMAGES}"
+            f"Python version {python_version} not supported, please use one of the following versions: {AVAILABLE_PYTHON_IMAGES}"
         )
         return
       logger.info(
           f"Using Python version {python_version} from the config file to build the Dockerfile")
     else:
       logger.info(
-          f"Python version not found in the config file, using default Python version: {self.DEFAULT_PYTHON_VERSION}"
+          f"Python version not found in the config file, using default Python version: {DEFAULT_PYTHON_VERSION}"
       )
-      python_version = self.DEFAULT_PYTHON_VERSION
+      python_version = DEFAULT_PYTHON_VERSION
 
-    base_image = self.PYTHON_BASE_IMAGE.format(python_version=python_version)
+    base_image = PYTHON_BASE_IMAGE.format(python_version=python_version)
 
     # Parse the requirements.txt file to determine the base image
     dependencies = self._parse_requirements()
     if 'torch' in dependencies and dependencies['torch']:
       torch_version = dependencies['torch']
 
-      for image in self.AVAILABLE_TORCH_IMAGES:
+      for image in AVAILABLE_TORCH_IMAGES:
         if torch_version in image and f'py{python_version}' in image:
-          base_image = self.TORCH_BASE_IMAGE.format(
+          cuda_version = image.split('-')[-1].replace('cuda', '')
+          base_image = TORCH_BASE_IMAGE.format(
               torch_version=torch_version,
               python_version=python_version,
-              cuda_version=self.DEFAULT_CUDA_VERSION)
-          logger.info(f"Using Torch version {torch_version} base image  to build the Docker image")
+              cuda_version=cuda_version,
+          )
+          logger.info(f"Using Torch version {torch_version} base image to build the Docker image")
           break
 
     # Replace placeholders with actual values
@@ -314,7 +277,7 @@ class ModelUploader:
       config = yaml.safe_load(file)
     model = config.get('model')
     model_type_id = model.get('model_type_id')
-    assert model_type_id in self.CONCEPTS_REQUIRED_MODEL_TYPE, f"Model type {model_type_id} not supported for concepts"
+    assert model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE, f"Model type {model_type_id} not supported for concepts"
     concept_protos = self._concepts_protos_from_concepts(labels)
 
     config['concepts'] = [{'id': concept.id, 'name': concept.name} for concept in concept_protos]
@@ -332,7 +295,7 @@ class ModelUploader:
     )
 
     model_type_id = self.config.get('model').get('model_type_id')
-    if model_type_id in self.CONCEPTS_REQUIRED_MODEL_TYPE:
+    if model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE:
 
       if 'concepts' in self.config:
         labels = self.config.get('concepts')
@@ -347,7 +310,10 @@ class ModelUploader:
         labels = sorted(labels.items(), key=lambda x: int(x[0]))
 
         config_file = os.path.join(self.folder, 'config.yaml')
-        self.hf_labels_to_config(labels, config_file)
+        try:
+          self.hf_labels_to_config(labels, config_file)
+        except Exception as e:
+          logger.error(f"Failed to update the config.yaml file with the concepts: {e}")
 
         model_version_proto.output_info.data.concepts.extend(
             self._concepts_protos_from_concepts(labels))
@@ -359,7 +325,7 @@ class ModelUploader:
 
     model_type_id = self.config.get('model').get('model_type_id')
 
-    if (model_type_id in self.CONCEPTS_REQUIRED_MODEL_TYPE) and 'concepts' not in self.config:
+    if (model_type_id in CONCEPTS_REQUIRED_MODEL_TYPE) and 'concepts' not in self.config:
       logger.info(
           f"Model type {model_type_id} requires concepts to be specified in the config.yaml file.."
       )
@@ -473,8 +439,21 @@ class ModelUploader:
             is_v3=self.is_v3,
         ))
 
+  def get_model_build_logs(self):
+    logs_request = service_pb2.ListLogEntriesRequest(
+        log_type="builder",
+        user_app_id=self.client.user_app_id,
+        model_id=self.model_proto.id,
+        model_version_id=self.model_version_id,
+        page=1,
+        per_page=50)
+    response = self.client.STUB.ListLogEntries(logs_request)
+
+    return response
+
   def monitor_model_build(self):
     st = time.time()
+    seen_logs = set()  # To avoid duplicate log messages
     while True:
       resp = self.client.STUB.GetModelVersion(
           service_pb2.GetModelVersionRequest(
@@ -485,6 +464,13 @@ class ModelUploader:
       status_code = resp.model_version.status.code
       if status_code == status_code_pb2.MODEL_BUILDING:
         print(f"Model is building... (elapsed {time.time() - st:.1f}s)", end='\r', flush=True)
+
+        # Fetch and display the logs
+        logs = self.get_model_build_logs()
+        for log_entry in logs.log_entries:
+          if log_entry.url not in seen_logs:
+            seen_logs.add(log_entry.url)
+            print(f"Model Building Logs...: {log_entry.message.strip()}")
         time.sleep(1)
       elif status_code == status_code_pb2.MODEL_TRAINED:
         logger.info(f"\nModel build complete! (elapsed {time.time() - st:.1f}s)")