Skip to content

Commit

Permalink
switch to argparse for CLI (#309)
Browse files Browse the repository at this point in the history
* switch to argparse

* format

* remove file first

* update

* remove comments
  • Loading branch information
aniketmaurya authored Oct 1, 2024
1 parent 6b5d43d commit 92e19e5
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 8 deletions.
1 change: 0 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
fastapi >=0.100
httpx
uvicorn[standard] >=0.29.0
jsonargparse[signatures]>=4.29.0
40 changes: 34 additions & 6 deletions src/litserve/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__":
Expand Down
2 changes: 1 addition & 1 deletion src/litserve/docker_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
27 changes: 27 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
@@ -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"

0 comments on commit 92e19e5

Please sign in to comment.