Server-side package installation (#2683)

* Add uninstalling functions to store.py

* Add new routes

* slightly change api

* make it kinda work

* get more of the way there

* Some cleanup

* Fix error modal

* remove "direct pip" install mode

* PR suggestions

* revert back to put_and_wait

* Slight refactor + improvement

* use a small timeout

backend/src/dependencies/store.py

@@ -16,6 +16,7 @@
 installed_packages: dict[str, str | None] = {}
 
 COLLECTING_REGEX = re.compile(r"Collecting ([a-zA-Z0-9-_]+)")
+UNINSTALLING_REGEX = re.compile(r"Uninstalling ([a-zA-Z0-9-_]+)-+")
 
 DEP_MAX_PROGRESS = 0.8
 
@@ -159,6 +160,7 @@ def get_progress_amount():
             "--disable-chainner_pip-version-check",
             "--no-warn-script-location",
             "--progress-bar=json",
+            "--no-cache-dir",
             *extra_index_args,
         ],
         stdout=subprocess.PIPE,
@@ -233,6 +235,116 @@ def get_progress_amount():
     return len(dependencies_to_install)
 
 
+def uninstall_dependencies_sync(
+    dependencies: list[DependencyInfo],
+):
+    if len(dependencies) == 0:
+        return
+
+    exit_code = subprocess.check_call(
+        [
+            python_path,
+            "-m",
+            "pip",
+            "uninstall",
+            *[d.package_name for d in dependencies],
+            "-y",
+        ],
+    )
+    if exit_code != 0:
+        raise ValueError("An error occurred while uninstalling dependencies.")
+
+    for dep_info in dependencies:
+        installed_packages[dep_info.package_name] = dep_info.version
+
+
+async def uninstall_dependencies(
+    dependencies: list[DependencyInfo],
+    update_progress_cb: UpdateProgressFn | None = None,
+    logger: Logger | None = None,
+):
+    # If there's no progress callback, just uninstall the dependencies synchronously
+    if update_progress_cb is None:
+        return uninstall_dependencies_sync(dependencies)
+
+    if len(dependencies) == 0:
+        return
+
+    dependency_name_map = {
+        dep_info.package_name: dep_info.display_name or dep_info.package_name
+        for dep_info in dependencies
+    }
+    deps_count = len(dependencies)
+    deps_counter = 0
+    transitive_deps_counter = 0
+
+    def get_progress_amount():
+        transitive_progress = 1 - 1 / (2**transitive_deps_counter)
+        progress = (deps_counter + transitive_progress) / (deps_count + 1)
+        return min(max(0, progress), 1)
+
+    # Used to increment by a small amount between collect and download
+    dep_small_incr = (1 / deps_count) / 2
+
+    process = subprocess.Popen(
+        [
+            python_path,
+            "-m",
+            # TODO: Change this back to "pip" once pip updates with my changes
+            "chainner_pip",
+            "uninstall",
+            *[d.package_name for d in dependencies],
+            "-y",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+    uninstalling_name = "Unknown"
+    while True:
+        nextline = process.stdout.readline()  # type: ignore
+        if nextline == b"" and process.poll() is not None:
+            break
+        line = nextline.decode("utf-8").strip()
+        if not line:
+            continue
+
+        if logger is not None and not line.startswith("Progress:"):
+            logger.info(line)
+
+        # The Uninstalling step of pip. It tells us what package is being UNinstalled.
+        if "Uninstalling" in line:
+            match = UNINSTALLING_REGEX.search(line)
+            if match:
+                package_name = match.group(1)
+                uninstalling_name = dependency_name_map.get(package_name, None)
+                if uninstalling_name is None:
+                    uninstalling_name = package_name
+                    transitive_deps_counter += 1
+                else:
+                    deps_counter += 1
+                await update_progress_cb(
+                    f"Uninstalling {uninstalling_name}...", get_progress_amount(), None
+                )
+        # The Downloading step of pip. It tells us what package is currently being downloaded.
+        # Later, we can use this to get the progress of the download.
+        # For now, we just tell the user that it's happening.
+        elif "Successfully uninstalled" in line:
+            await update_progress_cb(
+                f"Uninstalled {uninstalling_name}.",
+                get_progress_amount() + dep_small_incr,
+                None,
+            )
+
+    exit_code = process.wait()
+    if exit_code != 0:
+        raise ValueError("An error occurred while installing dependencies.")
+
+    await update_progress_cb("Finished installing dependencies...", 1, None)
+
+    for dep_info in dependencies:
+        del installed_packages[dep_info.package_name]
+
+
 __all__ = [
     "DependencyInfo",
     "python_path",

backend/src/events.py

@@ -17,7 +17,7 @@ class BackendStatusData(TypedDict):
 
 
 class BackendStatusEvent(TypedDict):
-    event: Literal["backend-status"]
+    event: Literal["backend-status", "package-install-status"]
     data: BackendStatusData
 
 

backend/src/server_host.py

@@ -16,7 +16,12 @@
 
 import api
 from custom_types import UpdateProgressFn
-from dependencies.store import DependencyInfo, install_dependencies, installed_packages
+from dependencies.store import (
+    DependencyInfo,
+    install_dependencies,
+    installed_packages,
+    uninstall_dependencies,
+)
 from events import EventQueue
 from gpu import get_nvidia_helper
 from response import error_response, success_response
@@ -171,6 +176,104 @@ async def get_features(request: Request):
     return await worker.proxy_request(request)
 
 
+def deps_to_dep_info(deps: list[api.Dependency]) -> list[DependencyInfo]:
+    return [
+        DependencyInfo(
+            package_name=dep.pypi_name,
+            display_name=dep.display_name,
+            version=dep.version,
+            extra_index_url=dep.extra_index_url,
+        )
+        for dep in deps
+    ]
+
+
+@app.route("/packages/uninstall", methods=["POST"])
+async def uninstall_dependencies_request(request: Request):
+    full_data = dict(request.json)  # type: ignore
+    package_to_uninstall = full_data["package"]
+
+    packages = await worker.get_packages()
+
+    package = next((x for x in packages if x.id == package_to_uninstall), None)
+
+    if package is None:
+        return json(
+            {"status": "error", "message": f"Package {package_to_uninstall} not found"},
+            status=404,
+        )
+
+    def update_progress(
+        message: str, progress: float, status_progress: float | None = None
+    ):
+        return AppContext.get(request.app).setup_queue.put_and_wait(
+            {
+                "event": "package-install-status",
+                "data": {
+                    "message": message,
+                    "progress": progress,
+                    "statusProgress": status_progress,
+                },
+            },
+            timeout=0.01,
+        )
+
+    try:
+        await worker.stop()
+        try:
+            await uninstall_dependencies(
+                deps_to_dep_info(package.dependencies), update_progress, logger
+            )
+        finally:
+            await worker.start()
+        return json({"status": "ok"})
+    except Exception as ex:
+        logger.error(f"Error uninstalling dependencies: {ex}", exc_info=True)
+        return json({"status": "error", "message": str(ex)}, status=500)
+
+
+@app.route("/packages/install", methods=["POST"])
+async def install_dependencies_request(request: Request):
+    full_data = dict(request.json)  # type: ignore
+    package_to_install = full_data["package"]
+
+    def update_progress(
+        message: str, progress: float, status_progress: float | None = None
+    ):
+        logger.info(f"Progress: {message} {progress} {status_progress}")
+        return AppContext.get(request.app).setup_queue.put_and_wait(
+            {
+                "event": "package-install-status",
+                "data": {
+                    "message": message,
+                    "progress": progress,
+                    "statusProgress": status_progress,
+                },
+            },
+            timeout=0.01,
+        )
+
+    packages = await worker.get_packages()
+    package = next((x for x in packages if x.id == package_to_install), None)
+
+    if package is None:
+        return json(
+            {"status": "error", "message": f"Package {package_to_install} not found"},
+            status=404,
+        )
+
+    try:
+        await worker.stop()
+        await install_dependencies(
+            deps_to_dep_info(package.dependencies), update_progress, logger
+        )
+        await worker.start()
+        return json({"status": "ok"})
+    except Exception as ex:
+        logger.error(f"Error installing dependencies: {ex}", exc_info=True)
+        return json({"status": "error", "message": str(ex)}, status=500)
+
+
 @app.get("/sse")
 async def sse(request: Request):
     headers = {"Cache-Control": "no-cache"}

src/common/Backend.ts

@@ -237,6 +237,24 @@ export class Backend {
     features(): Promise<FeatureState[]> {
         return this.fetchJson('/features', 'GET');
     }
+
+    installPackage(pkg: Package): Promise<void> {
+        return this.fetchJson('/packages/install', 'POST', {
+            package: pkg.id,
+        });
+    }
+
+    uninstallPackage(pkg: Package): Promise<void> {
+        return this.fetchJson('/packages/uninstall', 'POST', {
+            package: pkg.id,
+        });
+    }
+
+    updatePackage(pkg: Package): Promise<void> {
+        return this.fetchJson('/packages/install', 'POST', {
+            package: pkg.id,
+        });
+    }
 }
 
 const backendCache = new Map<string, Backend>();
@@ -293,4 +311,9 @@ export interface BackendEventMap {
         statusProgress?: number | null;
     };
     'backend-ready': null;
+    'package-install-status': {
+        message: string;
+        progress: number;
+        statusProgress?: number | null;
+    };
 }