[EAGLE-5342] Added Tests for Model run locally

clarifai/runners/models/model_run_locally.py

@@ -2,6 +2,7 @@
 import importlib.util
 import inspect
 import os
+import platform
 import shutil
 import signal
 import subprocess
@@ -27,14 +28,31 @@ def __init__(self, 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.uploader = ModelUploader(self.model_path, download_validation_only=True)
     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 _get_env_executable(self):
+    """Get the python executable from the virtual environment."""
+    # Depending on the platform, venv scripts are placed in either "Scripts" (Windows) or "bin" (Linux/Mac)
+    if platform.system().lower().startswith("win"):
+      scripts_folder = "Scripts"
+      python_exe = "python.exe"
+      pip_exe = "pip.exe"
+    else:
+      scripts_folder = "bin"
+      python_exe = "python"
+      pip_exe = "pip"
+
+    self.python_executable = os.path.join(self.venv_dir, scripts_folder, python_exe)
+    self.pip_executable = os.path.join(self.venv_dir, scripts_folder, pip_exe)
+
+    return self.python_executable, self.pip_executable
+
   def create_temp_venv(self):
     """Create a temporary virtual environment."""
     requirements_hash = self._requirements_hash()
@@ -53,13 +71,13 @@ def create_temp_venv(self):
 
     self.venv_dir = venv_dir
     self.temp_dir = temp_dir
-    self.python_executable = os.path.join(venv_dir, "bin", "python")
+    self.python_executable, self.pip_executable = self._get_env_executable()
 
     return use_existing_venv
 
   def install_requirements(self):
     """Install the dependencies from requirements.txt and Clarifai."""
-    pip_executable = os.path.join(self.venv_dir, "bin", "pip")
+    _, pip_executable = self._get_env_executable()
     try:
       logger.info(
           f"Installing requirements from {self.requirements_file}... in the virtual environment {self.venv_dir}"
@@ -104,8 +122,7 @@ def _get_model_runner(self):
   def _build_request(self):
     """Create a mock inference request for testing the model."""
 
-    uploader = ModelUploader(self.model_path)
-    model_version_proto = uploader.get_model_version_proto()
+    model_version_proto = self.uploader.get_model_version_proto()
     model_version_proto.id = "model_version"
 
     return service_pb2.PostModelOutputsRequest(
@@ -213,12 +230,16 @@ def _run_test(self):
 
   def test_model(self):
     """Test the model by running it locally in the virtual environment."""
-    command = [
-        self.python_executable,
-        "-c",
-        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()",
-    ]
+
+    import_path = repr(os.path.dirname(os.path.abspath(__file__)))
+    model_path = repr(self.model_path)
+
+    command_string = (f"import sys; "
+                      f"sys.path.append({import_path}); "
+                      f"from model_run_locally import ModelRunLocally; "
+                      f"ModelRunLocally({model_path})._run_test()")
+
+    command = [self.python_executable, "-c", command_string]
     process = None
     try:
       logger.info("Testing the model locally...")
@@ -335,6 +356,12 @@ def docker_image_exists(self, image_name):
       logger.info(f"Docker image '{image_name}' does not exist!")
       return False
 
+  def _gpu_is_available(self):
+    """
+    Checks if nvidia-smi is available, indicating a GPU is likely accessible.
+    """
+    return shutil.which("nvidia-smi") is not None
+
   def run_docker_container(self,
                            image_name,
                            container_name="clarifai-model-container",
@@ -344,9 +371,9 @@ def run_docker_container(self,
     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"
-      ]
+      cmd = ["docker", "run", "--name", container_name, '--rm', "--network", "host"]
+      if self._gpu_is_available():
+        cmd.extend(["--gpus", "all"])
       # Add volume mappings
       cmd.extend(["-v", f"{self.model_path}:/app/model_dir/main"])
       # Add environment variables
@@ -393,9 +420,9 @@ def test_model_container(self,
     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"
-      ]
+      cmd = ["docker", "run", "--name", container_name, '--rm', "--network", "host"]
+      if self._gpu_is_available():
+        cmd.extend(["--gpus", "all"])
       # update the entrypoint for testing the model
       cmd.extend(["--entrypoint", "python"])
       # Add volume mappings

tests/runners/dummy_runner_models/1/model.py

@@ -1,4 +1,3 @@
-import hashlib
 from typing import Iterator
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -44,11 +43,9 @@ def predict(self,
       if data.image.url != "":
         output.data.text.raw = data.image.url.replace(
             "samples.clarifai.com",
-            "newdomain.com" + params_dict.get("domain",),
+            "newdomain.com" + params_dict.get("domain", ""),
         )
-      if data.image.base64 != b"":
-        output.data.image.base64 = (
-            hashlib.md5(data.image.base64).hexdigest() + "Hello World run_input")
+
       output.status.code = status_code_pb2.SUCCESS
       outputs.append(output)
     return service_pb2.MultiOutputResponse(outputs=outputs,)
@@ -74,9 +71,7 @@ def generate(self, request: service_pb2.PostModelOutputsRequest
         if inp.data.text.raw != "":
           output.data.text.raw = f"{inp.data.text.raw}Generate Hello World {i}" + params_dict.get(
               "hello", "")
-        if inp.data.image.base64 != b"":
-          output.data.text.raw = (
-              hashlib.md5(inp.data.image.base64).hexdigest() + f"Generate Hello World {i}")
+
         output.status.code = status_code_pb2.SUCCESS
         outputs.append(output)
       yield service_pb2.MultiOutputResponse(outputs=outputs,)
@@ -102,8 +97,6 @@ def stream(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
           output = resources_pb2.Output()
           if inp.data.text.raw != "":
             out_text = inp.data.text.raw + f"Stream Hello World {i}" + params_dict.get("hello", "")
-          if inp.data.image.base64 != b"":
-            out_text = hashlib.md5(inp.data.image.base64).hexdigest() + f"Stream Hello World {i}"
           output.status.code = status_code_pb2.SUCCESS
           print(out_text)
           output.data.text.raw = out_text

tests/runners/dummy_runner_models/config.yaml

@@ -7,7 +7,7 @@ model:
   model_type_id: "multimodal-to-text"
 
 build_info:
-  python_version: "3.11"
+  python_version: "3.12"
 
 inference_compute_info:
   cpu_limit: "1"

tests/runners/dummy_runner_models/requirements.txt

@@ -0,0 +1,2 @@
+aiohttp
+requests

tests/runners/test_model_run_locally.py

@@ -0,0 +1,140 @@
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+
+from clarifai.runners.models.model_run_locally import ModelRunLocally
+
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "dummy_runner_models")
+CLARIFAI_USER_ID = os.environ["CLARIFAI_USER_ID"]
+CLARIFAI_PAT = os.environ["CLARIFAI_PAT"]
+
+
+@pytest.fixture
+def model_run_locally():
+  """
+  Fixture that instantiates the ModelRunLocally class
+  with the dummy model_path that already exists.
+  """
+  return ModelRunLocally(MODEL_PATH)
+
+
+def test_get_model_runner(model_run_locally):
+  """
+  Test that _get_model_runner successfully retrieves exactly one subclass of BaseRunner
+  from the dummy_runner_model's model.py
+  """
+  runner_cls = model_run_locally._get_model_runner()
+  assert runner_cls is not None, "Expected a runner class to be returned."
+  # Verify it's truly a subclass of clarifai_protocol.BaseRunner
+  from clarifai_protocol import BaseRunner
+  assert issubclass(runner_cls, BaseRunner), "Retrieved class must inherit from BaseRunner."
+
+
+def test_build_request(model_run_locally):
+  """
+  Test that _build_request returns a well-formed PostModelOutputsRequest
+  """
+  request = model_run_locally._build_request()
+  assert request is not None
+  assert len(request.inputs) == 1, "Expected exactly one input in constructed request."
+
+
+def test_create_temp_venv(model_run_locally):
+  """
+  Test whether create_temp_venv correctly initializes a virtual environment directory.
+  """
+  use_existing_venv = model_run_locally.create_temp_venv()
+  # Confirm we get back a boolean
+  assert isinstance(use_existing_venv, bool)
+  # Check that the venv_dir was set
+  venv_dir = Path(model_run_locally.venv_dir)
+  assert venv_dir.exists()
+  # Clean up
+  model_run_locally.clean_up()
+  assert not venv_dir.exists(), "Temporary virtual environment was not cleaned up"
+
+
+def test_install_requirements(model_run_locally):
+  """
+  Test installing requirements into the virtual environment.
+  Note: This actually installs from the dummy requirements.txt.
+  """
+  # Create the environment
+  model_run_locally.create_temp_venv()
+  # Attempt to install requirements
+  try:
+    model_run_locally.install_requirements()
+  except SystemExit:
+    pytest.fail("install_requirements() failed and exited.")
+  # You might want to verify the presence of installed packages by checking
+  # the venv's site-packages or something similar. For simplicity, we'll only
+  # verify that no exception was raised.
+  # Clean up
+  model_run_locally.clean_up()
+
+
+def test_test_model_success(model_run_locally):
+  """
+  Test that test_model succeeds with the dummy model.
+  This calls the script's test_model method, which runs a subprocess.
+  """
+  model_run_locally.create_temp_venv()
+  model_run_locally.install_requirements()
+
+  # Catch the subprocess call. If the dummy model is correct, exit code should be 0.
+  try:
+    model_run_locally.test_model()
+  except SystemExit:
+    pytest.fail("test_model() triggered a system exit with non-zero code.")
+  except subprocess.CalledProcessError:
+    # If the process didn't return code 0, fail the test
+    pytest.fail("The model test did not complete successfully in the subprocess.")
+  finally:
+    # Clean up
+    model_run_locally.clean_up()
+
+
+# @pytest.mark.skipif(shutil.which("docker") is None, reason="Docker not installed or not in PATH.")
+@pytest.mark.skip(reason="Will add later after new clarifai package is released")
+def test_docker_build_and_test_container(model_run_locally):
+  """
+  Test building a Docker image and running a container test using the dummy model.
+  This test will be skipped if Docker is not installed.
+  """
+  # Setup
+  # download_checkpoints & createDockerfile are called in the main()
+  # but we can do it here if needed. The code calls them automatically
+  # in main if inside_container is True, we directly test the method:
+
+  # Test if Docker is installed
+  assert model_run_locally.is_docker_installed(), "Docker not installed, skipping."
+
+  # Build or re-build the Docker image
+  model_run_locally.uploader.create_dockerfile()
+  image_tag = model_run_locally._docker_hash()
+  image_name = f"{model_run_locally.config['model']['id']}:{image_tag}"
+
+  if not model_run_locally.docker_image_exists(image_name):
+    model_run_locally.build_docker_image(image_name=image_name)
+
+  # Run tests inside the container
+  try:
+    model_run_locally.test_model_container(
+        image_name=image_name,
+        container_name="test_clarifai_model_container",
+        env_vars={
+            'CLARIFAI_PAT': CLARIFAI_PAT,
+            'CLARIFAI_API_BASE': os.environ.get('CLARIFAI_API_BASE', 'https://api.clarifai.com')
+        })
+  except subprocess.CalledProcessError:
+    pytest.fail("Failed to test the model inside the docker container.")
+  finally:
+    # Clean up the container if it still exists
+    if model_run_locally.container_exists("test_clarifai_model_container"):
+      model_run_locally.stop_docker_container("test_clarifai_model_container")
+      model_run_locally.remove_docker_container("test_clarifai_model_container")
+
+    # Remove the image
+    model_run_locally.remove_docker_image(image_name)