switch to argparse for CLI (#309)

* switch to argparse

* format

* remove file first

* update

* remove comments

requirements.txt

@@ -1,4 +1,3 @@
 fastapi >=0.100
 httpx
 uvicorn[standard] >=0.29.0
-jsonargparse[signatures]>=4.29.0

src/litserve/__main__.py

@@ -11,16 +11,44 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from jsonargparse import set_config_read_mode, set_docstring_parse_options, CLI
-
+from argparse import ArgumentParser
 from litserve.docker_builder import dockerize
 
 
 def main():
-    cli_components = {"dockerize": dockerize}
-    set_docstring_parse_options(attribute_docstrings=True)
-    set_config_read_mode(urls_enabled=True)
-    CLI(cli_components)
+    parser = ArgumentParser(description="CLI for LitServe")
+    subparsers = parser.add_subparsers(dest="command", help="Sub-command help")
+
+    # dockerize sub-command
+    dockerize_parser = subparsers.add_parser(
+        "dockerize",
+        help=dockerize.__doc__,
+        description="Generate a Dockerfile for the given server code.",
+    )
+    dockerize_parser.add_argument(
+        "server_filename",
+        type=str,
+        help="The path to the server file. Example: server.py or app.py.",
+    )
+    dockerize_parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="The port to expose in the Docker container. Defaults to 8000.",
+    )
+    dockerize_parser.add_argument(
+        "--gpu",
+        default=False,
+        action="store_true",
+        help="Whether to use a GPU-enabled Docker image. Defaults to false.",
+    )
+    dockerize_parser.set_defaults(func=lambda args: dockerize(args.server_filename, args.port, args.gpu))
+    args = parser.parse_args()
+
+    if hasattr(args, "func"):
+        args.func(args)
+    else:
+        parser.print_help()
 
 
 if __name__ == "__main__":

src/litserve/docker_builder.py

@@ -145,7 +145,7 @@ def dockerize(server_filename: str, port: int = 8000, gpu: bool = False):
         run_cmd = f"docker run --gpus all -p {port}:{port} litserve-model:latest"
         docker_template = CUDA_DOCKER_TEMPLATE
     else:
-        run_cmd = "docker run -p {port}:{port} litserve-model:latest"
+        run_cmd = f"docker run -p {port}:{port} litserve-model:latest"
         docker_template = DOCKERFILE_TEMPLATE
     dockerfile_content = docker_template.format(
         server_filename=server_filename,

tests/test_cli.py

@@ -0,0 +1,27 @@
+import pytest
+from litserve.__main__ import main
+import os
+
+
+def test_dockerize_help(monkeypatch, capsys):
+    monkeypatch.setattr("sys.argv", ["litserve", "dockerize", "--help"])
+    # argparse calls sys.exit() after printing help
+    with pytest.raises(SystemExit):
+        main()
+    captured = capsys.readouterr()
+    assert "usage:" in captured.out, "CLI did not print help message"
+    assert "The path to the server file." in captured.out, "CLI did not print help message"
+
+
+def test_dockerize_command(monkeypatch, capsys):
+    # Assuming you have a dummy server file for testing
+    dummy_server_file = "dummy_server.py"
+    with open(dummy_server_file, "w") as f:
+        f.write("# Dummy server file for testing\n")
+
+    monkeypatch.setattr("sys.argv", ["litserve", "dockerize", dummy_server_file])
+    main()
+    captured = capsys.readouterr()
+    os.remove(dummy_server_file)
+    assert "Dockerfile created successfully" in captured.out, "CLI did not create Dockerfile"
+    assert os.path.exists("Dockerfile"), "CLI did not create Dockerfile"