maintain venv in user cache dir and add pdb flags

Author: David Eigen

clarifai/cli/model.py

@@ -68,9 +68,11 @@ def download_checkpoints(model_path, out_path):
     help=
     'Specify how to test the model locally: "env" for virtual environment or "container" for Docker container. Defaults to "env".'
 )
+@click.option('--pdb', is_flag=True, help='Enable PDB debugging when testing the model locally.')
 @click.option(
     '--keep_env',
     is_flag=True,
+    default=True,
     help=
     'Keep the virtual environment after testing the model locally (applicable for virtualenv mode). Defaults to False.'
 )
@@ -80,7 +82,7 @@ def download_checkpoints(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'):
+def test_locally(model_path, pdb=False, keep_env=False, keep_image=False, mode='env'):
   """Test model locally."""
   try:
     from clarifai.runners.models import model_run_locally
@@ -89,9 +91,12 @@ def test_locally(model_path, keep_env=False, keep_image=False, mode='env'):
     if mode == 'container' and keep_env:
       raise ValueError("'keep_env' is applicable only for 'env' mode")
 
+    if pdb and mode == "container":
+      raise ValueError("PDB debugging is not supported in container mode.")
+
     if mode == "env":
       click.echo("Testing model locally in a virtual environment...")
-      model_run_locally.main(model_path, run_model_server=False, keep_env=keep_env)
+      model_run_locally.main(model_path, run_model_server=False, keep_env=keep_env, use_pdb=pdb)
     elif mode == "container":
       click.echo("Testing model locally inside a container...")
       model_run_locally.main(
@@ -125,6 +130,7 @@ def test_locally(model_path, keep_env=False, keep_image=False, mode='env'):
 @click.option(
     '--keep_env',
     is_flag=True,
+    default=True,
     help=
     'Keep the virtual environment after testing the model locally (applicable for virtualenv mode). Defaults to False.'
 )

clarifai/runners/models/model_run_locally.py

@@ -10,6 +10,7 @@
 import traceback
 import venv
 
+import platformdirs
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
 from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
 
@@ -52,37 +53,42 @@ def _get_env_executable(self):
 
   def create_temp_venv(self):
     """Create a temporary virtual environment."""
-    requirements_hash = self._requirements_hash()
-
-    temp_dir = os.path.join(tempfile.gettempdir(), str(requirements_hash))
+    venv_key = hashlib.md5(self.model_path.encode('utf-8')).hexdigest()
+    temp_dir = os.path.join(platformdirs.user_cache_dir('clarifai', 'clarifai'), 'venvs', venv_key)
     venv_dir = os.path.join(temp_dir, "venv")
 
     if os.path.exists(temp_dir):
-      logger.info(f"Using previous virtual environment at {temp_dir}")
-      use_existing_venv = True
+      logger.info(f"Using previous virtual environment at {venv_dir}")
     else:
-      logger.info("Creating temporary virtual environment...")
-      use_existing_venv = False
+      logger.info("Creating virtual environment...")
       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, self.pip_executable = self._get_env_executable()
 
-    return use_existing_venv
-
   def install_requirements(self):
     """Install the dependencies from requirements.txt and Clarifai."""
+    if os.path.exists(os.path.join(self.temp_dir, "requirements.txt")):
+      requirements = open(self.requirements_file).read()
+      installed_requirements = open(os.path.join(self.temp_dir, "requirements.txt")).read()
+      if requirements == installed_requirements:
+        logger.info("Requirements already installed.")
+        return
     _, pip_executable = self._get_env_executable()
     try:
       logger.info(
           f"Installing requirements from {self.requirements_file}... in the virtual environment {self.venv_dir}"
       )
-      subprocess.check_call([pip_executable, "install", "-r", self.requirements_file])
       logger.info("Installing Clarifai package...")
-      subprocess.check_call([pip_executable, "install", "clarifai"])
+      subprocess.check_call(
+          [pip_executable, "install", "clarifai==" + __import__("clarifai").__version__])
+      logger.info("Installing model requirements...")
+      subprocess.check_call([pip_executable, "install", "-r", self.requirements_file])
       logger.info("Requirements installed successfully!")
+      with open(os.path.join(self.temp_dir, "requirements.txt"), "w") as f:
+        f.write(open(self.requirements_file).read())
     except subprocess.CalledProcessError as e:
       logger.error(f"Error installing requirements: {e}")
       self.clean_up()
@@ -177,9 +183,12 @@ def _run_model_inference(self, model):
       ))
 
     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}")
+      try:
+        stream_first_res = next(stream_response)
+      except StopIteration:
+        stream_first_res = None
+      if stream_first_res is None or stream_first_res.outputs[0].status.code != status_code_pb2.SUCCESS:
+        logger.error(f"Model stream failed: {stream_first_res}")
       else:
         logger.info(
             f"Model Prediction succeeded for stream and first response: {stream_first_res}")
@@ -191,7 +200,7 @@ def _run_test(self):
     # send an inference.
     self._run_model_inference(model)
 
-  def test_model(self):
+  def test_model(self, use_pdb=False):
     """Test the model by running it locally in the virtual environment."""
 
     import_path = repr(os.path.dirname(os.path.abspath(__file__)))
@@ -201,8 +210,12 @@ def test_model(self):
                       f"sys.path.append({import_path}); "
                       f"from model_run_locally import ModelRunLocally; "
                       f"ModelRunLocally({model_path})._run_test()")
+    main_file = tempfile.NamedTemporaryFile(mode="w")
+    with open(main_file.name, "w") as f:
+      f.write(command_string)
 
-    command = [self.python_executable, "-c", command_string]
+    pdb_args = ["-m", "pdb"] if use_pdb else []
+    command = [self.python_executable, *pdb_args, main_file.name]
     process = None
     try:
       logger.info("Testing the model locally...")
@@ -232,12 +245,13 @@ def test_model(self):
           process.kill()
 
   # run the model server
-  def run_model_server(self, port=8080):
+  def run_model_server(self, port=8080, use_pdb=False):
     """Run the Clarifai Runners's model server."""
 
+    pdb_args = ["-m", "pdb"] if use_pdb else []
     command = [
-        self.python_executable, "-m", "clarifai.runners.server", "--model_path", self.model_path,
-        "--grpc", "--port",
+        self.python_executable, *pdb_args, "-m", "clarifai.runners.server", "--model_path",
+        self.model_path, "--grpc", "--port",
         str(port)
     ]
     try:
@@ -462,16 +476,19 @@ def remove_docker_image(self, image_name):
   def clean_up(self):
     """Clean up the temporary virtual environment."""
     if os.path.exists(self.temp_dir):
-      logger.info("Cleaning up temporary virtual environment...")
+      logger.info("Cleaning up virtual environment...")
       shutil.rmtree(self.temp_dir)
 
 
-def main(model_path,
-         run_model_server=False,
-         inside_container=False,
-         port=8080,
-         keep_env=False,
-         keep_image=False):
+def main(
+    model_path,
+    run_model_server=False,
+    inside_container=False,
+    port=8080,
+    keep_env=True,
+    keep_image=False,
+    use_pdb=False,
+):
 
   if not os.environ['CLARIFAI_PAT']:
     logger.error(
@@ -513,9 +530,9 @@ def main(model_path,
       if not use_existing_env:
         manager.install_requirements()
       if run_model_server:
-        manager.run_model_server(port)
+        manager.run_model_server(port, use_pdb=use_pdb)
       else:
-        manager.test_model()
+        manager.test_model(use_pdb=use_pdb)
     finally:
       if not keep_env:
         manager.clean_up()