grosse
/
jupyterlab_ollama_launcher
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

wrap in main

This commit is contained in:
Malte Grosse 2024-06-30 10:57:32 +00:00
parent 56b128f5c4
commit 609ca76405
2 changed files with 130 additions and 127 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

253
ollama_launcher/launcher.py
View File

@ -6,163 +6,166 @@ import threading
import itertools import itertools
import os import os
import argparse import argparse
print("##### Init Ollama #####") def main():
# --- Global list to store running processes --- print("##### Init Ollama #####")
running_processes = [] # --- Global list to store running processes ---
running_processes = []
# --- Function to run a command and wait for a specific string in its output --- # --- Function to run a command and wait for a specific string in its output ---
def run_and_wait_for_string(command, target_string,printstr, debug=False): def run_and_wait_for_string(command, target_string,printstr, debug=False):
"""Runs a subprocess, waits for a string, and hides output (with spinner).""" """Runs a subprocess, waits for a string, and hides output (with spinner)."""
if debug: if debug:
print(f"DEBUG: Executing command: {command}") print(f"DEBUG: Executing command: {command}")
print("##### "+printstr+" #####") print("##### "+printstr+" #####")
def spinning_cursor(): def spinning_cursor():
while True: while True:
for cursor in itertools.cycle(['-', '/', '|', '\\']): for cursor in itertools.cycle(['-', '/', '|', '\\']):
yield cursor yield cursor
spinner = spinning_cursor() spinner = spinning_cursor()
stop_event = threading.Event() stop_event = threading.Event()
def print_spinner(): def print_spinner():
while not stop_event.is_set(): while not stop_event.is_set():
sys.stdout.write(next(spinner)) sys.stdout.write(next(spinner))
sys.stdout.flush() sys.stdout.flush()
sys.stdout.write('\b') sys.stdout.write('\b')
time.sleep(0.1) time.sleep(0.1)
spinner_thread = threading.Thread(target=print_spinner) spinner_thread = threading.Thread(target=print_spinner)
spinner_thread.start() spinner_thread.start()
process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, text=True) stderr=subprocess.STDOUT, text=True)
running_processes.append(process) # Add the process to the list
while True:
output = process.stdout.readline()
if output == '' and process.poll() is not None:
break
if target_string in output:
stop_event.set()
spinner_thread.join()
print("done")
return True
stop_event.set()
spinner_thread.join()
return False
# --- Function to run a command until it succeeds ---
def run_until_success(command,printstr, debug=False):
"""Runs a command in a loop until it exits with a 0 return code (hidden)."""
if debug:
print(f"DEBUG: Executing command until successful: {command}")
print("##### "+printstr+" #####")
while True:
process = subprocess.Popen(command, shell=True, stdout=subprocess.DEVNULL,
stderr=subprocess.STDOUT)
running_processes.append(process) # Add the process to the list running_processes.append(process) # Add the process to the list
process.communicate()
if process.returncode == 0: while True:
print("done") output = process.stdout.readline()
break if output == '' and process.poll() is not None:
break
if target_string in output:
stop_event.set()
spinner_thread.join()
print("done")
return True
stop_event.set()
spinner_thread.join()
return False
# --- Function to gracefully terminate running processes --- # --- Function to run a command until it succeeds ---
def terminate_processes(): def run_until_success(command,printstr, debug=False):
"""Sends SIGINT to all processes in the global list.""" """Runs a command in a loop until it exits with a 0 return code (hidden)."""
for process in running_processes: if debug:
if process.poll() is None: print(f"DEBUG: Executing command until successful: {command}")
try:
process.send_signal(signal.SIGINT) print("##### "+printstr+" #####")
process.wait() while True:
except: process = subprocess.Popen(command, shell=True, stdout=subprocess.DEVNULL,
print(f"Error terminating process: {process}") stderr=subprocess.STDOUT)
running_processes.clear() running_processes.append(process) # Add the process to the list
process.communicate()
if process.returncode == 0:
print("done")
break
# --- Parse command line arguments --- # --- Function to gracefully terminate running processes ---
parser = argparse.ArgumentParser(description="Start Ollama and related services.") def terminate_processes():
parser.add_argument("--debug", action="store_true", help="Enable debug mode (print executed commands)") """Sends SIGINT to all processes in the global list."""
args = parser.parse_args() for process in running_processes:
if process.poll() is None:
try:
process.send_signal(signal.SIGINT)
process.wait()
except:
print(f"Error terminating process: {process}")
running_processes.clear()
# --- Main script ---
# --- Subprocess 1: ollama serve --- # --- Parse command line arguments ---
command1 = "ollama serve" parser = argparse.ArgumentParser(description="Start Ollama and related services.")
target_string1 = "inference compute" parser.add_argument("--debug", action="store_true", help="Enable debug mode (print executed commands)")
args = parser.parse_args()
# --- Subprocess 2: webui serve --- # --- Main script ---
command2 = "open-webui serve"
target_string2 = "Uvicorn running"
# --- Subprocess 3: ollama pull --- # --- Subprocess 1: ollama serve ---
command3 = "ollama pull qwen2:0.5b" command1 = "ollama serve"
target_string1 = "inference compute"
# --- Get token from user --- # --- Subprocess 2: webui serve ---
token = input("Please visit https://tun.iuk.hdm-stuttgart.de to obtain a token and paste it here: ") command2 = "open-webui serve"
jupyterhub_user = os.environ.get("JUPYTERHUB_USER") target_string2 = "Uvicorn running"
if not jupyterhub_user:
print("Error: Environment variable 'JUPYTERHUB_USER' not found.")
sys.exit(1)
# --- Subprocess 4: pgrok init --- # --- Subprocess 3: ollama pull ---
command4 = f"pgrok init --remote-addr tun.iuk.hdm-stuttgart.de:80 --forward-addr https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de --token {token}" command3 = "ollama pull qwen2:0.5b"
url = f"https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de"
# --- Subprocess 5: pgrok http ---
command5 = "pgrok http 8080"
target_string5 = "You're ready to go live"
# --- Run subprocesses sequentially --- # --- Get token from user ---
if run_and_wait_for_string(command1, target_string1, "Starting Ollama Server", debug=args.debug): token = input("Please visit https://tun.iuk.hdm-stuttgart.de to obtain a token and paste it here: ")
if run_and_wait_for_string(command2, target_string2, "Starting WebUI", debug=args.debug): jupyterhub_user = os.environ.get("JUPYTERHUB_USER")
run_until_success(command3,"Loading default model", debug=args.debug) if not jupyterhub_user:
print("Error: Environment variable 'JUPYTERHUB_USER' not found.")
sys.exit(1)
if subprocess.call(command4, shell=True, stdout=subprocess.DEVNULL) == 0: # --- Subprocess 4: pgrok init ---
print("##### Init Tunnel #####") command4 = f"pgrok init --remote-addr tun.iuk.hdm-stuttgart.de:80 --forward-addr https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de --token {token}"
if run_and_wait_for_string(command5, target_string5,"Create Tunnel", debug=args.debug): url = f"https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de"
# --- Subprocess 5: pgrok http ---
command5 = "pgrok http 8080"
target_string5 = "You're ready to go live"
# --- Ollama shell --- # --- Run subprocesses sequentially ---
if run_and_wait_for_string(command1, target_string1, "Starting Ollama Server", debug=args.debug):
if run_and_wait_for_string(command2, target_string2, "Starting WebUI", debug=args.debug):
run_until_success(command3,"Loading default model", debug=args.debug)
print("##### All services started, Please visit " + url + " for the WebUI! #####") if subprocess.call(command4, shell=True, stdout=subprocess.DEVNULL) == 0:
print("##### Launching Ollama shell! #####") print("##### Init Tunnel #####")
while True: if run_and_wait_for_string(command5, target_string5,"Create Tunnel", debug=args.debug):
command = input("ollama> ")
if command.lower() == "exit":
break
try: # --- Ollama shell ---
full_command = f"ollama {command}"
if args.debug:
print(f"DEBUG: Executing command: {full_command}")
process = subprocess.Popen(full_command, shell=True)
process.communicate()
except FileNotFoundError: print("##### All services started, Please visit " + url + " for the WebUI! #####")
print("Command not found. Please make sure 'ollama' is installed and in your PATH.") print("##### Launching Ollama shell! #####")
except Exception as e: while True:
print(f"An error occurred: {e}") command = input("ollama> ")
if command.lower() == "exit":
break
try:
full_command = f"ollama {command}"
if args.debug:
print(f"DEBUG: Executing command: {full_command}")
process = subprocess.Popen(full_command, shell=True)
process.communicate()
except FileNotFoundError:
print("Command not found. Please make sure 'ollama' is installed and in your PATH.")
except Exception as e:
print(f"An error occurred: {e}")
else:
print("Error: 'pgrok http' failed to start.")
terminate_processes()
sys.exit(1)
else: else:
print("Error: 'pgrok http' failed to start.") print("Error: 'pgrok init' failed.")
terminate_processes() terminate_processes()
sys.exit(1) sys.exit(1)
else: else:
print("Error: 'pgrok init' failed.") print("Error: 'webui serve' failed to start.")
terminate_processes() terminate_processes()
sys.exit(1) sys.exit(1)
else: else:
print("Error: 'webui serve' failed to start.") print("Error: 'ollama serve' failed to start.")
terminate_processes() terminate_processes()
sys.exit(1) sys.exit(1)
else:
print("Error: 'ollama serve' failed to start.")
terminate_processes()
sys.exit(1)
# Gracefully terminate running processes on script exit or error # Gracefully terminate running processes on script exit or error
terminate_processes() terminate_processes()
if __name__ == "__main__":
main()

2
setup.py
View File

@ -10,7 +10,7 @@ setup(
], ],
entry_points={ entry_points={
'console_scripts': [ 'console_scripts': [
'ollama_start=ollama_launcher.launcher' 'ollama_start=ollama_launcher.launcher:main'
], ],
}, },
) )