generate Dockerfile with CLI (#298)

* docker v0

* update

* handle filenotfound

* handle filenotfound

* update

* expose port

* update

* fix

* default level info

* concurrent copy

* add tests

* fix tests

* add license info

* update tests

* update

* add colors

* update tests

* clean up

* update dockerfile

* fix tests

* update msg

* update copy

* Update src/litserve/docker_builder.py

Co-authored-by: William Falcon <[email protected]>

* fix tests

* increase test coverage

* apply feedbacks

* rename cli

---------

Co-authored-by: William Falcon <[email protected]>

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",
+        ],
+    },
 )

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()

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 <package-name>
+
+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)

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 <package-name>
+
+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)