wrap in main

ollama_launcher/launcher.py

@@ -6,12 +6,13 @@ import threading
 import itertools
 import os
 import argparse
-print("##### Init Ollama #####")
-# --- Global list to store running processes ---
-running_processes = []
+def main():
+    print("##### Init Ollama #####")
+    # --- Global list to store running processes ---
+    running_processes = []
 
-# --- 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):
+    # --- 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):
         """Runs a subprocess, waits for a string, and hides output (with spinner)."""
         if debug:
             print(f"DEBUG: Executing command: {command}")
@@ -54,8 +55,8 @@ def run_and_wait_for_string(command, target_string,printstr, debug=False):
         return False
 
 
-# --- Function to run a command until it succeeds ---
-def run_until_success(command,printstr, debug=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}")
@@ -71,8 +72,8 @@ def run_until_success(command,printstr, debug=False):
                 break
 
 
-# --- Function to gracefully terminate running processes ---
-def terminate_processes():
+    # --- Function to gracefully terminate running processes ---
+    def terminate_processes():
         """Sends SIGINT to all processes in the global list."""
         for process in running_processes:
             if process.poll() is None:
@@ -84,40 +85,40 @@ def terminate_processes():
         running_processes.clear()
 
 
-# --- Parse command line arguments ---
-parser = argparse.ArgumentParser(description="Start Ollama and related services.")
-parser.add_argument("--debug", action="store_true", help="Enable debug mode (print executed commands)")
-args = parser.parse_args()
+    # --- Parse command line arguments ---
+    parser = argparse.ArgumentParser(description="Start Ollama and related services.")
+    parser.add_argument("--debug", action="store_true", help="Enable debug mode (print executed commands)")
+    args = parser.parse_args()
 
-# --- Main script ---
+    # --- Main script ---
 
-# --- Subprocess 1: ollama serve ---
-command1 = "ollama serve"
-target_string1 = "inference compute"
+    # --- Subprocess 1: ollama serve ---
+    command1 = "ollama serve"
+    target_string1 = "inference compute"
 
-# --- Subprocess 2: webui serve ---
-command2 = "open-webui serve"
-target_string2 = "Uvicorn running"
+    # --- Subprocess 2: webui serve ---
+    command2 = "open-webui serve"
+    target_string2 = "Uvicorn running"
 
-# --- Subprocess 3: ollama pull ---
-command3 = "ollama pull qwen2:0.5b"
+    # --- Subprocess 3: ollama pull ---
+    command3 = "ollama pull qwen2:0.5b"
 
-# --- Get token from user ---
-token = input("Please visit https://tun.iuk.hdm-stuttgart.de to obtain a token and paste it here: ")
-jupyterhub_user = os.environ.get("JUPYTERHUB_USER")
-if not jupyterhub_user:
+    # --- Get token from user ---
+    token = input("Please visit https://tun.iuk.hdm-stuttgart.de to obtain a token and paste it here: ")
+    jupyterhub_user = os.environ.get("JUPYTERHUB_USER")
+    if not jupyterhub_user:
         print("Error: Environment variable 'JUPYTERHUB_USER' not found.")
         sys.exit(1)
 
-# --- Subprocess 4: pgrok init ---
-command4 = f"pgrok init --remote-addr tun.iuk.hdm-stuttgart.de:80 --forward-addr https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de --token {token}"
-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"
+    # --- Subprocess 4: pgrok init ---
+    command4 = f"pgrok init --remote-addr tun.iuk.hdm-stuttgart.de:80 --forward-addr https://{jupyterhub_user}.tun.iuk.hdm-stuttgart.de --token {token}"
+    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 ---
-if run_and_wait_for_string(command1, target_string1, "Starting Ollama Server", debug=args.debug):
+    # --- 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)
 
@@ -159,10 +160,12 @@ if run_and_wait_for_string(command1, target_string1, "Starting Ollama Server", d
             print("Error: 'webui serve' failed to start.")
             terminate_processes()
             sys.exit(1)
-else:
+    else:
         print("Error: 'ollama serve' failed to start.")
         terminate_processes()
         sys.exit(1)
 
-# Gracefully terminate running processes on script exit or error
-terminate_processes()
+    # Gracefully terminate running processes on script exit or error
+    terminate_processes()
+if __name__ == "__main__":
+    main()

setup.py

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