diff --git a/setup.py b/setup.py index 9721ee7f..f4bbc58d 100755 --- a/setup.py +++ b/setup.py @@ -110,4 +110,9 @@ def _prepare_extras(requirements_dir: str = _PATH_REQUIRES, skip_files: tuple = "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", ], + entry_points={ + "console_scripts": [ + "litserve=litserve.__main__:main", + ], + }, ) diff --git a/src/litserve/__main__.py b/src/litserve/__main__.py new file mode 100644 index 00000000..6f8aa6ce --- /dev/null +++ b/src/litserve/__main__.py @@ -0,0 +1,27 @@ +# Copyright The Lightning AI team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 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) + + +if __name__ == "__main__": + main() diff --git a/src/litserve/docker_builder.py b/src/litserve/docker_builder.py new file mode 100644 index 00000000..c1b39e45 --- /dev/null +++ b/src/litserve/docker_builder.py @@ -0,0 +1,120 @@ +# Copyright The Lightning AI team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +import logging +import os +from pathlib import Path + +import warnings +import litserve as ls + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.propagate = False +console_handler = logging.StreamHandler() +console_handler.setLevel(logging.INFO) +formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") +console_handler.setFormatter(formatter) +logger.addHandler(console_handler) + +# COLOR CODES +RESET = "\u001b[0m" +RED = "\u001b[31m" +GREEN = "\u001b[32m" +BLUE = "\u001b[34m" +MAGENTA = "\u001b[35m" +BG_MAGENTA = "\u001b[45m" + +# ACTION CODES +BOLD = "\u001b[1m" +UNDERLINE = "\u001b[4m" +INFO = f"{BOLD}{BLUE}[INFO]" +WARNING = f"{BOLD}{RED}[WARNING]" + + +def color(text, color_code, action_code=None): + if action_code: + return f"{action_code} {color_code}{text}{RESET}" + return f"{color_code}{text}{RESET}" + + +REQUIREMENTS_FILE = "requirements.txt" +DOCKERFILE_TEMPLATE = """FROM python:3.10-slim + +####### Add your own installation commands here ####### +# RUN pip install some-package +# RUN wget https://path/to/some/data/or/weights +# RUN apt-get update && apt-get install -y + +WORKDIR /app +COPY . /app + +# Install litserve and requirements +RUN pip install --no-cache-dir litserve=={version} {requirements} +EXPOSE {port} +CMD ["python", "/app/{server_filename}"] +""" + +# Link our documentation as the bottom of this msg +SUCCESS_MSG = """{BOLD}{MAGENTA}Dockerfile created successfully at{RESET} {UNDERLINE}{dockerfile_path}{RESET} +Build the container with: +> {UNDERLINE}docker build -t litserve-model .{RESET} + +To run the Docker container on the machine: +> {UNDERLINE}docker run -p 8000:8000 litserve-model{RESET} + +To push the container to a registry: +> {UNDERLINE}docker push litserve-model{RESET} +""" + + +def dockerize(server_filename: str, port: int = 8000): + """Generate a Dockerfile for the given server code. + + Args: + server_filename (str): The path to the server file. Example sever.py or app.py. + port (int, optional): The port to expose in the Docker container. Defaults to 8000. + + """ + requirements = "" + if os.path.exists(REQUIREMENTS_FILE): + requirements = f"-r {REQUIREMENTS_FILE}" + else: + warnings.warn( + f"requirements.txt not found at {os.getcwd()}. " + f"Make sure to install the required packages in the Dockerfile.", + UserWarning, + ) + + current_dir = Path.cwd() + if not (current_dir / server_filename).is_file(): + raise FileNotFoundError(f"Server file `{server_filename}` must be in the current directory: {os.getcwd()}") + + version = ls.__version__ + dockerfile_content = DOCKERFILE_TEMPLATE.format( + server_filename=server_filename, port=port, version=version, requirements=requirements + ) + with open("Dockerfile", "w") as f: + f.write(dockerfile_content) + success_msg = SUCCESS_MSG.format( + dockerfile_path=os.path.abspath("Dockerfile"), + port=port, + BOLD=BOLD, + MAGENTA=MAGENTA, + GREEN=GREEN, + BLUE=BLUE, + UNDERLINE=UNDERLINE, + BG_MAGENTA=BG_MAGENTA, + RESET=RESET, + ) + print(success_msg) diff --git a/tests/test_docker_builder.py b/tests/test_docker_builder.py new file mode 100644 index 00000000..4eda4a9a --- /dev/null +++ b/tests/test_docker_builder.py @@ -0,0 +1,62 @@ +# Copyright The Lightning AI team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +import pytest + +import litserve as ls +from litserve import docker_builder + + +def test_color(): + assert docker_builder.color("hi", docker_builder.RED) == f"{docker_builder.RED}hi{docker_builder.RESET}" + + expected = f"{docker_builder.INFO} {docker_builder.RED}hi{docker_builder.RESET}" + assert docker_builder.color("hi", docker_builder.RED, docker_builder.INFO) == expected + + +EXPECTED_CONENT = f"""FROM python:3.10-slim + +####### Add your own installation commands here ####### +# RUN pip install some-package +# RUN wget https://path/to/some/data/or/weights +# RUN apt-get update && apt-get install -y + +WORKDIR /app +COPY . /app + +# Install litserve and requirements +RUN pip install --no-cache-dir litserve=={ls.__version__} -r requirements.txt +EXPOSE 8000 +CMD ["python", "/app/app.py"] +""" + + +def test_build(tmp_path, monkeypatch): + with open(tmp_path / "app.py", "w") as f: + f.write("print('hello')") + + # Temporarily change the current working directory to tmp_path + monkeypatch.chdir(tmp_path) + + with pytest.warns(UserWarning, match="Make sure to install the required packages in the Dockerfile."): + docker_builder.dockerize("app.py", 8000) + + with open(tmp_path / "requirements.txt", "w") as f: + f.write("lightning") + docker_builder.dockerize("app.py", 8000) + with open("Dockerfile") as f: + content = f.read() + assert content == EXPECTED_CONENT + + with pytest.raises(FileNotFoundError, match="must be in the current directory"): + docker_builder.dockerize("random_file_name.py", 8000)