Start: Rewrite start scripts

Start scripts now don't update dependencies by default due to mishandling
caches from pip. Also add dedicated update scripts and save options
to a JSON file instead of a text one.

Signed-off-by: kingbri <bdashore3@proton.me>

.gitignore

@@ -200,5 +200,8 @@ sampler_overrides/*
 # Gpu lib preferences file
 gpu_lib.txt
 
+# Start options file
+start_options.json
+
 # OpenAPI JSON
 openapi.json

start.bat

@@ -19,3 +19,5 @@ if exist "%CONDA_PREFIX%" (
 
 :: Call the python script with batch args
 call python start.py %*
+
+pause

start.py

@@ -1,6 +1,7 @@
 """Utility to automatically upgrade and start the API"""
 
 import argparse
+import json
 import os
 import pathlib
 import platform
@@ -11,6 +12,9 @@ from shutil import copyfile
 from common.args import convert_args_to_dict, init_argparser
 
 
+start_options = {}
+
+
 def get_user_choice(question: str, options_dict: dict):
     """
     Gets user input in a commandline script.
@@ -39,36 +43,24 @@ def get_install_features(lib_name: str = None):
     """Fetches the appropriate requirements file depending on the GPU"""
     install_features = None
     possible_features = ["cu121", "cu118", "amd"]
-    saved_lib_path = pathlib.Path("gpu_lib.txt")
 
-    if lib_name:
+    if not lib_name:
-        print("Overriding GPU lib name from args.")
+        # Ask the user for the GPU lib
-    else:
+        gpu_lib_choices = {
-        # Try getting the GPU lib from file
+            "A": {"pretty": "NVIDIA Cuda 12.x", "internal": "cu121"},
-        if saved_lib_path.exists():
+            "B": {"pretty": "NVIDIA Cuda 11.8", "internal": "cu118"},
-            with open(saved_lib_path.resolve(), "r") as f:
+            "C": {"pretty": "AMD", "internal": "amd"},
-                lib_name = f.readline().strip()
+        }
-        else:
+        user_input = get_user_choice(
-            # Ask the user for the GPU lib
+            "Select your GPU. If you don't know, select Cuda 12.x (A)",
-            gpu_lib_choices = {
+            gpu_lib_choices,
-                "A": {"pretty": "NVIDIA Cuda 12.x", "internal": "cu121"},
+        )
-                "B": {"pretty": "NVIDIA Cuda 11.8", "internal": "cu118"},
-                "C": {"pretty": "AMD", "internal": "amd"},
-            }
-            user_input = get_user_choice(
-                "Select your GPU. If you don't know, select Cuda 12.x (A)",
-                gpu_lib_choices,
-            )
 
-            lib_name = gpu_lib_choices.get(user_input, {}).get("internal")
+        lib_name = gpu_lib_choices.get(user_input, {}).get("internal")
 
-            # Write to a file for subsequent runs
+        # Write to start options
-            with open(saved_lib_path.resolve(), "w") as f:
+        start_options["gpu_lib"] = lib_name
-                f.write(lib_name)
+        print("Saving your choice to start options.")
-                print(
-                    "Saving your choice to gpu_lib.txt. "
-                    "Delete this file and restart if you want to change your selection."
-                )
 
     # Assume default if the file is invalid
     if lib_name and lib_name in possible_features:
@@ -104,10 +96,16 @@ def add_start_args(parser: argparse.ArgumentParser):
     """Add start script args to the provided parser"""
     start_group = parser.add_argument_group("start")
     start_group.add_argument(
-        "-iu",
+        "-ur",
-        "--ignore-upgrade",
+        "--update-repository",
         action="store_true",
-        help="Ignore requirements upgrade",
+        help="Update local git repository to latest",
+    )
+    start_group.add_argument(
+        "-ud",
+        "--update-deps",
+        action="store_true",
+        help="Update all pip dependencies",
     )
     start_group.add_argument(
         "-nw",
@@ -122,6 +120,25 @@ def add_start_args(parser: argparse.ArgumentParser):
     )
 
 
+def migrate_gpu_lib():
+    gpu_lib_path = pathlib.Path("gpu_lib.txt")
+
+    if not gpu_lib_path.exists():
+        return
+
+    print("Migrating gpu_lib.txt to the new start_options.json")
+    with open("gpu_lib.txt", "r") as gpu_lib_file:
+        start_options["gpu_lib"] = gpu_lib_file.readline().strip()
+        start_options["first_run_done"] = True
+
+    # Remove the old file
+    gpu_lib_path.unlink()
+
+    print(
+        "Successfully migrated gpu lib options to start_options. "
+        "The old file has been deleted."
+    )
+
 if __name__ == "__main__":
     subprocess.run(["pip", "-V"])
 
@@ -129,6 +146,34 @@ if __name__ == "__main__":
     parser = init_argparser()
     add_start_args(parser)
     args = parser.parse_args()
+    script_ext = "bat" if platform.system() == "Windows" else "sh"
+
+    start_options_path = pathlib.Path("start_options.json")
+    if start_options_path.exists():
+        with open(start_options_path) as start_options_file:
+            start_options = json.load(start_options_file)
+
+        if start_options.get("first_run_done"):
+            first_run = False
+    else:
+        print(
+            "It looks like you're running TabbyAPI for the first time. "
+            "Getting things ready..."
+        )
+
+    # Migrate from old setting storage
+    migrate_gpu_lib()
+
+    # Set variables that rely on start options
+    first_run = not start_options.get("first_run_done")
+
+    if args.gpu_lib:
+        print("Overriding GPU lib name from args.")
+        gpu_lib = args.gpu_lib
+    elif "gpu_lib" in start_options:
+        gpu_lib = start_options.get("gpu_lib")
+    else:
+        gpu_lib = None
 
     # Create a config if it doesn't exist
     # This is not necessary to run TabbyAPI, but is new user proof
@@ -144,10 +189,13 @@ if __name__ == "__main__":
             f"Created one at {str(config_path.resolve())}"
         )
 
-    if args.ignore_upgrade:
+    if args.update_repository:
-        print("Ignoring pip dependency upgrade due to user request.")
+        print("Pulling latest changes from Github.")
-    else:
+        pull_command = "git pull"
-        install_features = None if args.nowheel else get_install_features(args.gpu_lib)
+        subprocess.run(pull_command.split(" "))
+
+    if first_run or args.update_deps:
+        install_features = None if args.nowheel else get_install_features(gpu_lib)
         features = f"[{install_features}]" if install_features else ""
 
         # pip install .[features]
@@ -155,8 +203,35 @@ if __name__ == "__main__":
         print(f"Running install command: {install_command}")
         subprocess.run(install_command.split(" "))
 
+        if args.update_deps:
+            print(
+                f"Dependencies updated. Please run TabbyAPI with `start.{script_ext}`. "
+                "Exiting."
+            )
+            sys.exit(0)
+        else:
+            print(
+                f"Dependencies installed. Update them with `update_deps.{script_ext}` "
+                "inside the `update_scripts` folder."
+            )
+
+    if first_run:
+        start_options["first_run_done"] = True
+
+    # Save start options
+    with open("start_options.json", "w") as start_file:
+        start_file.write(json.dumps(start_options))
+
+        print(
+            "Successfully wrote your start script options to `start_options.json`. \n"
+            "If something goes wrong, editing or deleting the file will reinstall TabbyAPI "
+            "as a first-time user."
+        )
+
     # Import entrypoint after installing all requirements
     from main import entrypoint
 
     converted_args = convert_args_to_dict(args, parser)
+
+    print("Starting TabbyAPI...")
     entrypoint(converted_args)

update_scripts/update_deps.bat

@@ -0,0 +1,24 @@
+@echo off
+
+:: Creates a venv if it doesn't exist and runs the start script for requirements upgrades
+:: This is intended for users who want to start the API and have everything upgraded and installed
+
+:: cd to the parent directory
+cd "%~dp0.."
+
+:: Don't create a venv if a conda environment is active
+if exist "%CONDA_PREFIX%" (
+    echo It looks like you're in a conda environment. Skipping venv check.
+) else (
+    if not exist "venv\" (
+        echo Venv doesn't exist! Please run start.bat instead.
+        exit 0
+    )
+
+    call .\venv\Scripts\activate.bat
+)
+
+:: Call the python script with batch args
+call python start.py --update-deps %*
+
+pause

update_scripts/update_deps.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+cd "$(dirname "$0")/.." || exit
+
+if [ -n "$CONDA_PREFIX" ]; then
+    echo "It looks like you're in a conda environment. Skipping venv check."
+else
+    if [ ! -d "venv" ]; then
+        echo "Venv doesn't exist! Please run start.sh instead."
+        exit 0
+    fi
+
+    echo "Activating venv"
+
+    # shellcheck source=/dev/null
+    source venv/bin/activate
+fi
+
+python3 start.py --update-deps "$@"

update_scripts/update_deps_and_pull.bat

@@ -0,0 +1,24 @@
+@echo off
+
+:: Creates a venv if it doesn't exist and runs the start script for requirements upgrades
+:: This is intended for users who want to start the API and have everything upgraded and installed
+
+:: cd to the parent directory
+cd "%~dp0.."
+
+:: Don't create a venv if a conda environment is active
+if exist "%CONDA_PREFIX%" (
+    echo It looks like you're in a conda environment. Skipping venv check.
+) else (
+    if not exist "venv\" (
+        echo Venv doesn't exist! Please run start.bat instead.
+        exit 0
+    )
+
+    call .\venv\Scripts\activate.bat
+)
+
+:: Call the python script with batch args
+call python start.py --update-deps --update-repository %*
+
+pause

update_scripts/update_deps_and_pull.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+cd "$(dirname "$0")/.." || exit
+
+if [ -n "$CONDA_PREFIX" ]; then
+    echo "It looks like you're in a conda environment. Skipping venv check."
+else
+    if [ ! -d "venv" ]; then
+        echo "Venv doesn't exist! Please run start.sh instead."
+        exit 0
+    fi
+
+    echo "Activating venv"
+
+    # shellcheck source=/dev/null
+    source venv/bin/activate
+fi
+
+python3 start.py --update-deps --update-repository "$@"