fix(cli): lightning-sdk installation process with uv (#640)

* fix(cli): Improve lightning-sdk installation process with uv support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add test for failure case in _ensure_lightning_installed function

* fix(cli): Enhance lightning-sdk installation process to support multiple installers

* updated tests

* re run

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bhimrazy

src/litserve/cli.py

@@ -1,13 +1,33 @@
+import importlib.util
+import shutil
 import subprocess
 import sys
 
 from litserve.utils import is_package_installed
 
 
 def _ensure_lightning_installed():
-    if not is_package_installed("lightning_sdk"):
-        print("Lightning CLI not found. Installing...")
-        subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"])
+    """Ensure lightning-sdk is installed, attempting auto-installation if needed."""
+    if is_package_installed("lightning_sdk"):
+        return
+
+    print("Lightning CLI not found. Installing lightning-sdk...")
+
+    # Build list of available installers (pip first as it respects the active environment)
+    installers = []
+    if importlib.util.find_spec("pip"):
+        installers.append([sys.executable, "-m", "pip"])
+    if shutil.which("uv"):
+        installers.append(["uv", "pip"])
+
+    for installer in installers:
+        try:
+            subprocess.run([*installer, "install", "-U", "lightning-sdk"], check=True)
+            return
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            continue
+
+    sys.exit("Failed to install lightning-sdk. Run: pip install lightning-sdk")
 
 
 def main():

tests/unit/test_cli.py

@@ -1,4 +1,5 @@
 import os
+import subprocess
 import sys
 from unittest.mock import MagicMock, patch
 
@@ -34,19 +35,98 @@ def test_dockerize_command(monkeypatch, capsys):
 
 
 @patch("litserve.cli.is_package_installed")
-@patch("subprocess.check_call")
-def test_ensure_lightning_installed(mock_check_call, mock_is_package_installed):
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_with_pip(mock_run, mock_which, mock_find_spec, mock_is_package_installed):
     mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = True  # pip available
+    mock_which.return_value = None  # uv not available
     _ensure_lightning_installed()
-    mock_check_call.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"])
+    mock_run.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"], check=True)
+
+
+@patch("litserve.cli.is_package_installed")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_pip_preferred(mock_run, mock_which, mock_find_spec, mock_is_package_installed):
+    """When both pip and uv are available, pip should be used first."""
+    mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = True  # pip available
+    mock_which.return_value = "/usr/bin/uv"  # uv also available
+    _ensure_lightning_installed()
+    mock_run.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"], check=True)
+
+
+@patch("litserve.cli.is_package_installed")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_with_uv(mock_run, mock_which, mock_find_spec, mock_is_package_installed):
+    mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = None  # pip not available
+    mock_which.return_value = "/usr/bin/uv"  # uv available
+    _ensure_lightning_installed()
+    mock_run.assert_called_once_with(["uv", "pip", "install", "-U", "lightning-sdk"], check=True)
+
+
+@patch("litserve.cli.is_package_installed")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_fallback_to_uv(mock_run, mock_which, mock_find_spec, mock_is_package_installed):
+    """When pip fails, should fall back to uv."""
+    mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = True  # pip available
+    mock_which.return_value = "/usr/bin/uv"  # uv also available
+    mock_run.side_effect = [subprocess.CalledProcessError(1, "pip"), None]  # pip fails, uv succeeds
+    _ensure_lightning_installed()
+    assert mock_run.call_count == 2
+    mock_run.assert_called_with(["uv", "pip", "install", "-U", "lightning-sdk"], check=True)
+
+
+@patch("litserve.cli.is_package_installed")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_failure(mock_run, mock_which, mock_find_spec, mock_is_package_installed):
+    """When all available installers fail, should exit with error."""
+    mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = True  # pip available
+    mock_which.return_value = "/usr/bin/uv"  # uv also available
+    mock_run.side_effect = subprocess.CalledProcessError(1, "install")  # both fail
+
+    with pytest.raises(SystemExit, match="Failed to install lightning-sdk"):
+        _ensure_lightning_installed()
+    assert mock_run.call_count == 2  # tried both pip and uv
+
+
+@patch("litserve.cli.is_package_installed")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
+def test_ensure_lightning_installed_no_installer_available(
+    mock_run, mock_which, mock_find_spec, mock_is_package_installed
+):
+    """When neither pip nor uv is available, should exit with error."""
+    mock_is_package_installed.return_value = False
+    mock_find_spec.return_value = None  # pip not available
+    mock_which.return_value = None  # uv not available
+
+    with pytest.raises(SystemExit, match="Failed to install lightning-sdk"):
+        _ensure_lightning_installed()
+    mock_run.assert_not_called()  # no installer was tried
 
 
 # TODO: Remove this once we have a fix for Python 3.10
 @pytest.mark.skipif(sys.version_info[:2] in [(3, 10)], reason="Test fails on Python 3.10")
 @patch("litserve.cli.is_package_installed")
-@patch("subprocess.check_call")
+@patch("litserve.cli.importlib.util.find_spec")
+@patch("litserve.cli.shutil.which")
+@patch("subprocess.run")
 @patch("builtins.__import__")
-def test_cli_main_lightning_not_installed(mock_import, mock_check_call, mock_is_package_installed):
+def test_cli_main_lightning_not_installed(mock_import, mock_run, mock_which, mock_find_spec, mock_is_package_installed):
     # Create a mock for the lightning_sdk module and its components
     mock_lightning_sdk = MagicMock()
     mock_lightning_sdk.cli.entrypoint.main_cli = MagicMock()
@@ -58,6 +138,8 @@ def side_effect(name, *args, **kwargs):
         return __import__(name, *args, **kwargs)
 
     mock_import.side_effect = side_effect
+    mock_find_spec.return_value = True  # pip available
+    mock_which.return_value = None  # uv not available
 
     # Test when lightning_sdk is not installed but gets installed dynamically
     mock_is_package_installed.side_effect = [False, True]  # First call returns False, second call returns True
@@ -66,7 +148,7 @@ def side_effect(name, *args, **kwargs):
     with patch.object(sys, "argv", test_args):
         cli_main()
 
-    mock_check_call.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"])
+    mock_run.assert_called_once_with([sys.executable, "-m", "pip", "install", "-U", "lightning-sdk"], check=True)
 
 
 @pytest.mark.skipif(sys.version_info[:2] in [(3, 10)], reason="Test fails on Python 3.10")