add step args

Author: christineyu123

clarifai/cli/pipeline_step.py

@@ -125,15 +125,30 @@ def init(pipeline_step_path):
     default=False,
     help='Keep the Docker image after the pipeline step finishes.',
 )
-def local_run(pipeline_step_path, mode, keep_image):
+@click.option(
+    '--step-args',
+    default=None,
+    help='Arguments to pass to pipeline_step.py (e.g., "--param_a hello --param_b world").',
+)
+def local_run(pipeline_step_path, mode, keep_image, step_args):
     """Run a pipeline step locally in a Docker container.
 
+    \b
     PIPELINE_STEP_PATH: Path to the pipeline step directory (containing config.yaml,
     requirements.txt, and 1/pipeline_step.py). Defaults to current directory.
 
+    \b
+    Pass arguments to the step script via --step-args:
+      clarifai pipelinestep local-run ./my-step --step-args "--param_a hello --param_b world"
+
     This reuses the same Docker build infrastructure as ``clarifai model serve
     --mode container`` but executes pipeline_step.py once and exits.
     """
+    # Parse step-args string into a list
+    if step_args:
+        import shlex
+
+        step_args = shlex.split(step_args)
     from clarifai.runners.pipeline_steps.pipeline_run_locally import PipelineStepRunLocally
 
     manager = PipelineStepRunLocally(pipeline_step_path)
@@ -157,6 +172,7 @@ def local_run(pipeline_step_path, mode, keep_image):
         manager.run_pipeline_step_container(
             image_name=image_name,
             container_name=container_name,
+            step_args=list(step_args) if step_args else None,
         )
     finally:
         if manager.container_exists(container_name):

clarifai/runners/pipeline_steps/pipeline_run_locally.py

@@ -70,11 +70,15 @@ def run_pipeline_step_container(
         image_name,
         container_name="clarifai-pipeline-step-container",
         env_vars=None,
+        step_args=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.
+
+        Args:
+            step_args: Optional list of arguments to pass to pipeline_step.py.
         """
         try:
             cmd = ["docker", "run", "--name", container_name, "--rm", "--network", "host"]
@@ -101,6 +105,10 @@ def run_pipeline_step_container(
             cmd.append(image_name)
             cmd.extend(["/home/nonroot/main/1/pipeline_step.py"])
 
+            # Append any extra arguments for the pipeline step script
+            if step_args:
+                cmd.extend(step_args)
+
             logger.info(f"Running pipeline step in container '{container_name}'...")
             logger.info(f"Docker command: {cmd}")
 

tests/cli/test_pipeline_step.py

@@ -464,6 +464,7 @@ def test_local_run_command_help(self):
         assert '--mode' in result.output
         assert '--keep-image' in result.output
         assert 'PIPELINE_STEP_PATH' in result.output
+        assert '--step-args' in result.output
 
 
 class TestPipelineStepLocalRunCommand:
@@ -504,3 +505,20 @@ def test_local_run_no_docker(self, mock_docker):
 
             assert result.exit_code != 0
             assert 'Docker is not installed' in result.output
+
+    def test_local_run_step_args_parsing(self):
+        """Test that --step-args string is correctly parsed into a list."""
+        import shlex
+
+        # This is the parsing logic used inside the local_run command
+        step_args_str = "--param_a hello --param_b world"
+        parsed = shlex.split(step_args_str)
+        assert parsed == ['--param_a', 'hello', '--param_b', 'world']
+
+        # Quoted values
+        step_args_str = '--param_a "hello world" --param_b value'
+        parsed = shlex.split(step_args_str)
+        assert parsed == ['--param_a', 'hello world', '--param_b', 'value']
+
+        # None input means no args
+        assert shlex.split("") == []

tests/runners/dummy_pipeline_step/1/pipeline_step.py

@@ -1,4 +1,18 @@
-"""Dummy pipeline step for testing local-run."""
+"""Dummy pipeline step matching the default clarifai pipelinestep init template."""
 
-print("Pipeline step started")
-print("Pipeline step completed successfully")
+import argparse
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Dummy pipeline step.')
+    parser.add_argument('--param_a', type=str, required=True, help='First parameter')
+    parser.add_argument('--param_b', type=str, default='default_b', help='Second parameter')
+
+    args = parser.parse_args()
+
+    print(f"Pipeline step started: param_a={args.param_a}, param_b={args.param_b}")
+    print("Pipeline step completed successfully")
+
+
+if __name__ == "__main__":
+    main()

tests/runners/test_pipeline_step_run_locally_container.py

@@ -48,6 +48,7 @@ def test_pipeline_step_docker_build_and_run(pipeline_step_run_locally):
         pipeline_step_run_locally.run_pipeline_step_container(
             image_name=image_name,
             container_name=container_name,
+            step_args=["--param_a", "hello", "--param_b", "world"],
         )
     except subprocess.CalledProcessError:
         pytest.fail("Failed to run pipeline step inside the docker container.")