Backend-defined Dependencies (chaiNNer-org#1834)

* start of defining dependencies on backend

* ncnn dep

* nvidia helper

* define rest of deps in backend

* some fixes

* import testing & dep installing

* fixed install stuff

* this dependency stuff is a mess

* ...

* WIP core dependency workaround

* Fixes

* fix scoping issue

* Fixes

* cleanup

* fix workflow

* some fixes

* lint

* make it not run the b uild workflow

* update requirements.txt

* renaming

* reorganize stuff

* Correct order of operations

* SEE for backend ready event

* wip more

* getting closer to the finish line

* move a few required deps over

* remove a few prints

* float progress

* Install multiple dependencies at once + other pr suggestions

* Other minor PR suggestions

* fix dumb mistake

* fixes & debugging

* lint

* attempting to use a separate sse

* setup-sse from main process

* use progress slice

* use task

* sleepy time

* reuse event loop, because we probably should

* put_and_wait

* replace console.log with comment

* backend linting

* lint + remove logs

* remove comment

* retry getting python info

* Don't --upgrade, allow None versions

* clean up requirements.txt

* Revert "clean up requirements.txt"

This reverts commit 4c0911a.

* More declarative deps

* remove unnecessary global

* remove comments

* other cleanup

* rename some things

* add error listener

* add back loading status

* Pr suggestions

* Added `nodes_available` function

* Removed my stupid log.debug

* Update src/renderer/splash.tsx

* Update src/renderer/main.tsx

---------

Co-authored-by: RunDevelopment <mitchi5000.ms@googlemail.com>

.github/workflows/make.yml

@@ -16,7 +16,7 @@ on:
 jobs:
   test-build-release:
     # commit the word "build" to the commit message to enable this job
-    if: contains(${{ github.event.head_commit.message }}, 'build')
+    if: contains(github.event.head_commit.message, 'build')
     runs-on: ${{ matrix.os }}
 
     strategy:

.github/workflows/test-backend.yml

@@ -49,7 +49,6 @@ jobs:
             - uses: actions/setup-python@v4
               with:
                   python-version: ${{ matrix.python-version }}
-            - run: npx ts-node ./scripts/install-required-deps.ts
             - run: python ./backend/src/run.py --no-run
               env:
                   TYPE_CHECK_LEVEL: 'error'

backend/src/api.py

@@ -19,6 +19,10 @@
     typeValidateSchema,
 )
 
+KB = 1024**1
+MB = 1024**2
+GB = 1024**3
+
 
 def _process_inputs(base_inputs: Iterable[Union[BaseInput, NestedGroup]]):
     inputs: List[BaseInput] = []
@@ -174,11 +178,34 @@ def toDict(self):
         }
 
 
+@dataclass
+class Dependency:
+    display_name: str
+    pypi_name: str
+    version: str
+    size_estimate: int | float
+    auto_update: bool = False
+    extra_index_url: str | None = None
+
+    import_name: str | None = None
+
+    def toDict(self):
+        return {
+            "displayName": self.display_name,
+            "pypiName": self.pypi_name,
+            "version": self.version,
+            "sizeEstimate": int(self.size_estimate),
+            "autoUpdate": self.auto_update,
+            "findLink": self.extra_index_url,
+        }
+
+
 @dataclass
 class Package:
     where: str
     name: str
-    dependencies: List[str] = field(default_factory=list)
+    description: str
+    dependencies: List[Dependency] = field(default_factory=list)
     categories: List[Category] = field(default_factory=list)
 
     def add_category(
@@ -200,6 +227,12 @@ def add_category(
         self.categories.append(result)
         return result
 
+    def add_dependency(
+        self,
+        dependency: Dependency,
+    ):
+        self.dependencies.append(dependency)
+
 
 def _iter_py_files(directory: str):
     for root, _, files in os.walk(directory):
@@ -271,5 +304,7 @@ def _refresh_nodes(self):
 registry = PackageRegistry()
 
 
-def add_package(where: str, name: str, dependencies: List[str]) -> Package:
-    return registry.add(Package(where, name, dependencies))
+def add_package(
+    where: str, name: str, description: str, dependencies: List[Dependency]
+) -> Package:
+    return registry.add(Package(where, name, description, dependencies))

backend/src/dependencies/install_core_deps.py

@@ -0,0 +1,54 @@
+from typing import List
+
+from system import is_arm_mac, is_windows
+
+from .store import DependencyInfo
+from .versioned_dependency_helpers import install_version_checked_dependencies
+
+# I'm leaving this here in case I can use the Dependency class in the future, so I don't lose the extra info
+
+# dependencies=[
+#         Dependency("OpenCV", "opencv-python", "4.7.0.68", 30 * MB, import_name="cv2"),
+#         Dependency("NumPy", "numpy", "1.23.2", 15 * MB),
+#         Dependency("Pillow (PIL)", "Pillow", "9.2.0", 3 * MB, import_name="PIL"),
+#         Dependency("appdirs", "appdirs", "1.4.4", 13.5 * KB),
+#         Dependency("FFMPEG", "ffmpeg-python", "0.2.0", 25 * KB, import_name="ffmpeg"),
+#         Dependency("Requests", "requests", "2.28.2", 452 * KB),
+#         Dependency("re2", "google-re2", "1.0", 275 * KB, import_name="re2"),
+#         Dependency("scipy", "scipy", "1.9.3", 42 * MB),
+#     ],
+
+# if is_arm_mac:
+#     package.add_dependency(Dependency("Pasteboard", "pasteboard", "0.3.3", 19 * KB))
+# elif is_windows:
+#     package.add_dependency(
+#         Dependency("Pywin32", "pywin32", "304", 12 * MB, import_name="win32clipboard")
+#     )
+
+deps: List[DependencyInfo] = [
+    {
+        "package_name": "appdirs",
+        "version": "1.4.4",
+    },
+    {
+        "package_name": "ffmpeg-python",
+        "version": "0.2.0",
+    },
+    {
+        "package_name": "requests",
+        "version": "2.28.2",
+    },
+    {
+        "package_name": "scipy",
+        "version": "1.9.3",
+    },
+    {"package_name": "pynvml", "version": "11.5.0"},
+    {"package_name": "typing_extensions", "version": "4.6.2"},
+]
+
+if is_arm_mac:
+    deps.append({"package_name": "pasteboard", "version": "0.3.3"})
+elif is_windows:
+    deps.append({"package_name": "pywin32", "version": None})
+
+install_version_checked_dependencies(deps)

backend/src/dependencies/install_essential_deps.py

@@ -0,0 +1,15 @@
+from typing import List
+
+from .store import DependencyInfo, install_dependencies
+
+# These are the dependencies we _absolutely need_ to install before we can do anything else
+deps: List[DependencyInfo] = [
+    {
+        "package_name": "semver",
+        "version": "3.0.0",
+    },
+]
+
+
+# Note: We can't be sure we have semver yet so we can't use it to compare versions
+install_dependencies(deps)

backend/src/dependencies/install_server_deps.py

@@ -0,0 +1,43 @@
+import subprocess
+from json import loads as json_parse
+from typing import List
+
+from .store import DependencyInfo, installed_packages, python_path
+from .versioned_dependency_helpers import install_version_checked_dependencies
+
+# Get the list of installed packages
+# We can't rely on using the package's __version__ attribute because not all packages actually have it
+try:
+    pip_list = subprocess.check_output(
+        [python_path, "-m", "pip", "list", "--format=json"]
+    )
+    for p in json_parse(pip_list):
+        installed_packages[p["name"]] = p["version"]
+except Exception as e:
+    print(f"Failed to get installed packages: {e}")
+
+
+deps: List[DependencyInfo] = [
+    {
+        "package_name": "sanic",
+        "version": "23.3.0",
+    },
+    {
+        "package_name": "Sanic-Cors",
+        "version": "2.2.0",
+    },
+    {
+        "package_name": "numpy",
+        "version": "1.23.2",
+    },
+    {
+        "package_name": "opencv-python",
+        "version": "4.7.0.68",
+    },
+    {
+        "package_name": "Pillow",
+        "version": "9.2.0",
+    },
+]
+
+install_version_checked_dependencies(deps)

backend/src/dependencies/store.py

@@ -0,0 +1,47 @@
+import subprocess
+import sys
+from typing import Dict, List, TypedDict, Union
+
+python_path = sys.executable
+
+installed_packages: Dict[str, Union[str, None]] = {}
+
+
+class DependencyInfo(TypedDict):
+    package_name: str
+    version: Union[str, None]
+
+
+def pin(package_name: str, version: Union[str, None]) -> str:
+    if version is None:
+        return package_name
+    return f"{package_name}=={version}"
+
+
+def install_dependencies(dependency_info_array: List[DependencyInfo]):
+    subprocess.check_call(
+        [
+            python_path,
+            "-m",
+            "pip",
+            "install",
+            *[
+                pin(dep_info["package_name"], dep_info["version"])
+                for dep_info in dependency_info_array
+            ],
+            "--disable-pip-version-check",
+            "--no-warn-script-location",
+        ]
+    )
+    for dep_info in dependency_info_array:
+        package_name = dep_info["package_name"]
+        version = dep_info["version"]
+        installed_packages[package_name] = version
+
+
+__all__ = [
+    "DependencyInfo",
+    "python_path",
+    "install_dependencies",
+    "installed_packages",
+]

backend/src/dependencies/versioned_dependency_helpers.py

@@ -0,0 +1,36 @@
+import re
+from typing import List
+
+from semver.version import Version
+
+from .store import DependencyInfo, install_dependencies, installed_packages
+
+
+def coerce_semver(version: str) -> Version:
+    try:
+        return Version.parse(version, True)
+    except Exception:
+        regex = r"(\d+)\.(\d+)\.(\d+)"
+        match = re.search(regex, version)
+        if match:
+            return Version(
+                int(match.group(1)),
+                int(match.group(2)),
+                int(match.group(3)),
+            )
+        return Version(0, 0, 0)
+
+
+def install_version_checked_dependencies(dependencies: List[DependencyInfo]):
+    dependencies_to_install = []
+    for dependency in dependencies:
+        version = installed_packages.get(dependency["package_name"], None)
+        if dependency["version"] and version:
+            installed_version = coerce_semver(version)
+            dep_version = coerce_semver(dependency["version"])
+            if installed_version < dep_version:
+                dependencies_to_install.append(dependency)
+        elif not version:
+            dependencies_to_install.append(dependency)
+    if len(dependencies_to_install) > 0:
+        install_dependencies(dependencies_to_install)

backend/src/events.py

@@ -47,6 +47,11 @@ class IteratorProgressUpdateData(TypedDict):
     running: Optional[List[NodeId]]
 
 
+class BackendStatusData(TypedDict):
+    message: str
+    percent: float
+
+
 class FinishEvent(TypedDict):
     event: Literal["finish"]
     data: FinishData
@@ -67,11 +72,23 @@ class IteratorProgressUpdateEvent(TypedDict):
     data: IteratorProgressUpdateData
 
 
+class BackendStatusEvent(TypedDict):
+    event: Literal["backend-status"]
+    data: BackendStatusData
+
+
+class BackendStateEvent(TypedDict):
+    event: Union[Literal["backend-ready"], Literal["backend-started"]]
+    data: None
+
+
 Event = Union[
     FinishEvent,
     ExecutionErrorEvent,
     NodeFinishEvent,
     IteratorProgressUpdateEvent,
+    BackendStatusEvent,
+    BackendStateEvent,
 ]
 
 
@@ -84,3 +101,14 @@ async def get(self) -> Event:
 
     async def put(self, event: Event) -> None:
         await self.queue.put(event)
+
+    async def wait_until_empty(self, timeout: float) -> None:
+        while timeout > 0:
+            if self.queue.empty():
+                return
+            await asyncio.sleep(0.01)
+            timeout -= 0.01
+
+    async def put_and_wait(self, event: Event, timeout: float = float("inf")) -> None:
+        await self.queue.put(event)
+        await self.wait_until_empty(timeout)

backend/src/gpu.py

@@ -0,0 +1,49 @@
+import pynvml as nv
+from sanic.log import logger
+
+nvidia_is_available = False
+
+try:
+    nv.nvmlInit()
+    nvidia_is_available = True
+    nv.nvmlShutdown()
+except nv.NVMLError as e:
+    logger.info("No Nvidia GPU found, or invalid driver installed.")
+except Exception as e:
+    logger.info(f"Unknown error occurred when trying to initialize Nvidia GPU: {e}")
+
+
+class NvidiaHelper:
+    def __init__(self):
+        nv.nvmlInit()
+
+        self.__num_gpus = nv.nvmlDeviceGetCount()
+
+        self.__gpus = []
+        for i in range(self.__num_gpus):
+            handle = nv.nvmlDeviceGetHandleByIndex(i)
+            self.__gpus.append(
+                {
+                    "name": nv.nvmlDeviceGetName(handle),
+                    "uuid": nv.nvmlDeviceGetUUID(handle),
+                    "index": i,
+                    "handle": handle,
+                }
+            )
+
+    def __del__(self):
+        nv.nvmlShutdown()
+
+    def list_gpus(self):
+        return self.__gpus
+
+    def get_current_vram_usage(self, gpu_index=0):
+        info = nv.nvmlDeviceGetMemoryInfo(self.__gpus[gpu_index]["handle"])
+
+        return info.total, info.used, info.free
+
+
+__all__ = [
+    "nvidia_is_available",
+    "NvidiaHelper",
+]

backend/src/nodes/impl/pytorch/architecture/SPSR.py

@@ -60,7 +60,6 @@ def __init__(
         self.out_nc: int = self.state["f_HR_conv1.0.bias"].shape[0]
 
         self.scale = self.get_scale(4)
-        print(self.scale)
         self.num_filters: int = self.state["model.0.weight"].shape[0]
 
         self.supports_fp16 = True