add lightning cli (#446)

Author: aniketmaurya

setup.py

@@ -111,8 +111,6 @@ def _prepare_extras(requirements_dir: str = _PATH_REQUIRES, skip_files: tuple =
         "Programming Language :: Python :: 3.11",
     ],
     entry_points={
-        "console_scripts": [
-            "litserve=litserve.__main__:main",
-        ],
+        "console_scripts": ["litserve=litserve.__main__:main", "lightning=litserve.cli:main"],
     },
 )

src/litserve/cli.py

@@ -0,0 +1,17 @@
+import importlib.util
+import subprocess
+import sys
+
+
+def _ensure_lightning_installed():
+    if not importlib.util.find_spec("lightning_sdk"):
+        print("Lightning CLI not found. Installing...")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"])
+
+
+def main():
+    _ensure_lightning_installed()
+
+    # Forward CLI arguments to the real lightning command
+    cli_args = sys.argv[1:]
+    subprocess.run(["lightning"] + cli_args)

tests/minimal_run.py

@@ -21,10 +21,10 @@
 
 def main():
     process = subprocess.Popen(
-        ["python", "tests/simple_server.py"],
+        ["lightning", "serve", "api", "tests/simple_server.py"],
     )
     print("Waiting for server to start...")
-    time.sleep(5)
+    time.sleep(10)
     try:
         url = "http://127.0.0.1:8000/predict"
         data = json.dumps({"input": 4.0}).encode("utf-8")

tests/test_cli.py

@@ -1,8 +1,11 @@
 import os
+import sys
+from unittest.mock import patch
 
 import pytest
 
 from litserve.__main__ import main
+from litserve.cli import _ensure_lightning_installed
 
 
 def test_dockerize_help(monkeypatch, capsys):
@@ -27,3 +30,11 @@ def test_dockerize_command(monkeypatch, capsys):
     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"
+
+
+@patch("importlib.util.find_spec")
+@patch("subprocess.check_call")
+def test_ensure_lightning_installed(mock_check_call, mock_find_spec):
+    mock_find_spec.return_value = False
+    _ensure_lightning_installed()
+    mock_check_call.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"])