Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Server-side package installation #2683

Merged
merged 17 commits into from
Mar 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 112 additions & 0 deletions backend/src/dependencies/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion backend/src/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ class BackendStatusData(TypedDict):


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


Expand Down
105 changes: 104 additions & 1 deletion backend/src/server_host.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
joeyballentine marked this conversation as resolved.
Show resolved Hide resolved
{
"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"}
Expand Down
23 changes: 23 additions & 0 deletions src/common/Backend.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>();
Expand Down Expand Up @@ -293,4 +311,9 @@ export interface BackendEventMap {
statusProgress?: number | null;
};
'backend-ready': null;
'package-install-status': {
message: string;
progress: number;
statusProgress?: number | null;
};
}
Loading
Loading