refactoring

Author: christineyu123

clarifai/cli/pipeline.py

@@ -93,6 +93,61 @@ def pipeline():
     """Create and manage pipelines."""
 
 
+@pipeline.command('local-run')
+@click.argument("step_path", type=click.Path(exists=True), required=False, default=".")
+@click.option(
+    '--mode',
+    type=click.Choice(['container']),
+    default='container',
+    help='Execution mode. Currently only "container" is supported.',
+)
+@click.option(
+    '--keep-image',
+    is_flag=True,
+    default=False,
+    help='Keep the Docker image after the pipeline step finishes.',
+)
+def local_run(step_path, mode, keep_image):
+    """Run a pipeline step locally in a Docker container.
+
+    STEP_PATH: Path to the pipeline step directory (containing config.yaml,
+    requirements.txt, and 1/pipeline_step.py). Defaults to current directory.
+
+    This reuses the same Docker build infrastructure as ``clarifai model serve
+    --mode container`` but executes pipeline_step.py once and exits.
+    """
+    from clarifai.runners.pipeline_steps.pipeline_run_locally import PipelineStepRunLocally
+
+    manager = PipelineStepRunLocally(step_path)
+
+    if not manager.is_docker_installed():
+        raise click.ClickException("Docker is not installed.")
+
+    # Generate Dockerfile if missing
+    manager.builder.create_dockerfile()
+
+    image_tag = manager._docker_hash()
+    step_id = manager.config['pipeline_step']['id'].lower()
+    image_name = f"{step_id}:{image_tag}"
+    container_name = f"{step_id}-local-run"
+
+    if not manager.docker_image_exists(image_name):
+        logger.info("Building Docker image...")
+        manager.build_docker_image(image_name=image_name)
+
+    try:
+        manager.run_pipeline_step_container(
+            image_name=image_name,
+            container_name=container_name,
+        )
+    finally:
+        if manager.container_exists(container_name):
+            manager.stop_docker_container(container_name)
+            manager.remove_docker_container(container_name)
+        if not keep_image:
+            manager.remove_docker_image(image_name)
+
+
 @pipeline.command()
 @click.argument("path", type=click.Path(exists=True), required=False, default=".")
 @click.option(

clarifai/runners/pipeline_steps/pipeline_run_locally.py

@@ -0,0 +1,138 @@
+"""Run a pipeline step locally in a Docker container.
+
+Reuses Docker infrastructure from ModelRunLocally for building and managing
+containers, but runs pipeline_step.py instead of the model server.
+"""
+
+import os
+import signal
+import subprocess
+import sys
+import time
+
+from clarifai.runners.models.model_run_locally import ModelRunLocally
+from clarifai.runners.pipeline_steps.pipeline_step_builder import PipelineStepBuilder
+from clarifai.utils.logging import logger
+
+
+class PipelineStepRunLocally:
+    """Run a single pipeline step locally in a Docker container.
+
+    Reuses ModelRunLocally for Docker build/image/container management,
+    but overrides the container command to run pipeline_step.py.
+    """
+
+    def __init__(self, step_path):
+        self.step_path = os.path.abspath(step_path)
+        self.builder = PipelineStepBuilder(self.step_path)
+        self.config = self.builder.config
+
+        # Create a ModelRunLocally instance for Docker utilities.
+        # We bypass __init__ since it expects a ModelBuilder, and set the
+        # attributes that the Docker methods rely on directly.
+        self._docker = ModelRunLocally.__new__(ModelRunLocally)
+        self._docker.model_path = self.step_path
+        self._docker.requirements_file = os.path.join(self.step_path, "requirements.txt")
+
+    # ── Delegated Docker utilities ──────────────────────────────────────
+
+    def is_docker_installed(self):
+        return self._docker.is_docker_installed()
+
+    def docker_image_exists(self, image_name):
+        return self._docker.docker_image_exists(image_name)
+
+    def build_docker_image(self, image_name):
+        return self._docker.build_docker_image(image_name=image_name)
+
+    def container_exists(self, container_name):
+        return self._docker.container_exists(container_name)
+
+    def stop_docker_container(self, container_name):
+        return self._docker.stop_docker_container(container_name)
+
+    def remove_docker_container(self, container_name):
+        return self._docker.remove_docker_container(container_name)
+
+    def remove_docker_image(self, image_name):
+        return self._docker.remove_docker_image(image_name)
+
+    def _docker_hash(self):
+        return self._docker._docker_hash()
+
+    def _gpu_is_available(self):
+        return self._docker._gpu_is_available()
+
+    # ── Pipeline-step-specific container run ────────────────────────────
+
+    def run_pipeline_step_container(
+        self,
+        image_name,
+        container_name="clarifai-pipeline-step-container",
+        env_vars=None,
+    ):
+        """Run pipeline_step.py inside a Docker container and wait for it to finish.
+
+        Unlike ``ModelRunLocally.run_docker_container`` which starts a long-running
+        server, this method executes the pipeline step script once and exits.
+        """
+        try:
+            cmd = ["docker", "run", "--name", container_name, "--rm", "--network", "host"]
+
+            if self._gpu_is_available():
+                cmd.extend(["--gpus", "all"])
+
+            # Mount step directory into container (same target as model serve)
+            cmd.extend(
+                [
+                    "--mount",
+                    f"type=bind,source={self.step_path},target=/home/nonroot/main",
+                ]
+            )
+
+            if env_vars:
+                for key, value in env_vars.items():
+                    cmd.extend(["-e", f"{key}={value}"])
+
+            cmd.extend(["-e", "PYTHONDONTWRITEBYTECODE=1"])
+
+            # Override entrypoint to run pipeline_step.py directly
+            cmd.extend(["--entrypoint", "python"])
+            cmd.append(image_name)
+            cmd.extend(["/home/nonroot/main/1/pipeline_step.py"])
+
+            logger.info(f"Running pipeline step in container '{container_name}'...")
+            logger.info(f"Docker command: {cmd}")
+
+            process = subprocess.Popen(cmd)
+
+            # Graceful Ctrl+C handling
+            original_sigint = signal.getsignal(signal.SIGINT)
+
+            def signal_handler(sig, frame):
+                logger.info(f"Stopping container '{container_name}'...")
+                subprocess.run(["docker", "stop", container_name], check=False)
+                process.terminate()
+                signal.signal(signal.SIGINT, original_sigint)
+                time.sleep(1)
+                sys.exit(0)
+
+            signal.signal(signal.SIGINT, signal_handler)
+
+            process.wait()
+
+            # Restore original handler
+            signal.signal(signal.SIGINT, original_sigint)
+
+            if process.returncode != 0:
+                logger.error(f"Pipeline step failed with exit code {process.returncode}")
+                sys.exit(process.returncode)
+
+            logger.info("Pipeline step completed successfully!")
+
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Error running pipeline step container: {e}")
+            sys.exit(1)
+        except Exception as e:
+            logger.error(f"Error running pipeline step container: {e}")
+            sys.exit(1)

tests/runners/dummy_pipeline_step/1/pipeline_step.py

@@ -0,0 +1,4 @@
+"""Dummy pipeline step for testing local-run."""
+
+print("Pipeline step started")
+print("Pipeline step completed successfully")

tests/runners/dummy_pipeline_step/config.yaml

@@ -0,0 +1,12 @@
+pipeline_step:
+  id: "dummy-pipeline-step"
+  user_id: "test_user"
+  app_id: "test_app"
+
+build_info:
+  python_version: "3.12"
+
+pipeline_step_compute_info:
+  cpu_limit: "500m"
+  cpu_memory: "500Mi"
+  num_accelerators: 0

tests/runners/dummy_pipeline_step/requirements.txt

@@ -0,0 +1 @@
+clarifai

tests/runners/test_pipeline_run_locally_container.py

@@ -0,0 +1,58 @@
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from clarifai.runners.pipeline_steps.pipeline_run_locally import PipelineStepRunLocally
+
+
+@pytest.fixture
+def dummy_pipeline_step_path(tmp_path):
+    """Copy the dummy_pipeline_step folder to a temp directory."""
+    tests_dir = Path(__file__).parent.resolve()
+    original_path = tests_dir / "dummy_pipeline_step"
+    if not original_path.exists():
+        raise FileNotFoundError(f"Could not find dummy_pipeline_step at {original_path}.")
+    target_folder = tmp_path / "dummy_pipeline_step"
+    shutil.copytree(original_path, target_folder)
+    return str(target_folder)
+
+
+@pytest.fixture
+def pipeline_run_locally(dummy_pipeline_step_path):
+    """Instantiate PipelineRunLocally with the dummy pipeline step."""
+    return PipelineStepRunLocally(dummy_pipeline_step_path)
+
+
+@pytest.mark.skipif(shutil.which("docker") is None, reason="Docker not installed or not in PATH.")
+@pytest.mark.skipif(
+    sys.platform not in ["linux", "darwin"],
+    reason="Test only runs on Linux and macOS.",
+)
+def test_pipeline_step_docker_build_and_run(pipeline_run_locally):
+    """Test building a Docker image and running a pipeline step in a container."""
+    assert pipeline_run_locally.is_docker_installed(), "Docker not installed."
+
+    pipeline_run_locally.builder.create_dockerfile()
+    image_tag = pipeline_run_locally._docker_hash()
+    step_id = pipeline_run_locally.config['pipeline_step']['id'].lower()
+    image_name = f"{step_id}:{image_tag}"
+    container_name = "test-pipeline-step-container"
+
+    if not pipeline_run_locally.docker_image_exists(image_name):
+        pipeline_run_locally.build_docker_image(image_name=image_name)
+
+    try:
+        pipeline_run_locally.run_pipeline_step_container(
+            image_name=image_name,
+            container_name=container_name,
+        )
+    except subprocess.CalledProcessError:
+        pytest.fail("Failed to run pipeline step inside the docker container.")
+    finally:
+        if pipeline_run_locally.container_exists(container_name):
+            pipeline_run_locally.stop_docker_container(container_name)
+            pipeline_run_locally.remove_docker_container(container_name)
+        pipeline_run_locally.remove_docker_image(image_name)