Spaces:
Running
Running
File size: 8,654 Bytes
c5b0bb7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
import platform
import gradio as gr
import subprocess
import psutil
import os
import signal
import logging
import threading
import shutil
# Configure Logging
# logging.basicConfig(
# level=logging.DEBUG, # Set to DEBUG to capture all levels of logs
# format='%(asctime)s - %(levelname)s - %(message)s',
# handlers=[
# logging.FileHandler("app.log"),
# logging.StreamHandler()
# ]
# )
def is_ollama_installed():
"""
Checks if the 'ollama' executable is available in the system's PATH.
Returns True if installed, False otherwise.
"""
return shutil.which('ollama') is not None
def get_ollama_models():
"""
Retrieves available Ollama models by executing 'ollama list'.
Returns a list of model names or an empty list if an error occurs.
"""
try:
result = subprocess.run(['ollama', 'list'], capture_output=True, text=True, check=True, timeout=10)
models = result.stdout.strip().split('\n')[1:] # Skip header
model_names = [model.split()[0] for model in models if model.strip()]
logging.debug(f"Available Ollama models: {model_names}")
return model_names
except FileNotFoundError:
logging.error("Ollama executable not found. Please ensure Ollama is installed and in your PATH.")
return []
except subprocess.TimeoutExpired:
logging.error("Ollama 'list' command timed out.")
return []
except subprocess.CalledProcessError as e:
logging.error(f"Error executing Ollama 'list': {e}")
return []
except Exception as e:
logging.error(f"Unexpected error in get_ollama_models: {e}")
return []
def pull_ollama_model(model_name):
"""
Pulls the specified Ollama model if Ollama is installed.
"""
if not is_ollama_installed():
logging.error("Ollama is not installed.")
return "Failed to pull model: Ollama is not installed or not in your PATH."
try:
subprocess.run(['ollama', 'pull', model_name], check=True, timeout=300) # Adjust timeout as needed
logging.info(f"Successfully pulled model: {model_name}")
return f"Successfully pulled model: {model_name}"
except subprocess.TimeoutExpired:
logging.error(f"Pulling model '{model_name}' timed out.")
return f"Failed to pull model '{model_name}': Operation timed out."
except subprocess.CalledProcessError as e:
logging.error(f"Failed to pull model '{model_name}': {e}")
return f"Failed to pull model '{model_name}': {e}"
except FileNotFoundError:
logging.error("Ollama executable not found. Please ensure Ollama is installed and in your PATH.")
return "Failed to pull model: Ollama executable not found."
except Exception as e:
logging.error(f"Unexpected error in pull_ollama_model: {e}")
return f"Failed to pull model '{model_name}': {e}"
def serve_ollama_model(model_name, port):
"""
Serves the specified Ollama model on the given port if Ollama is installed.
"""
if not is_ollama_installed():
logging.error("Ollama is not installed.")
return "Error: Ollama is not installed or not in your PATH."
try:
# Check if a server is already running on the specified port
for conn in psutil.net_connections():
if conn.laddr.port == int(port):
logging.warning(f"Port {port} is already in use.")
return f"Error: Port {port} is already in use. Please choose a different port."
# Start the Ollama server
cmd = ['ollama', 'serve', model_name, '--port', str(port)]
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
logging.info(f"Started Ollama server for model '{model_name}' on port {port}. PID: {process.pid}")
return f"Started Ollama server for model '{model_name}' on port {port}. Process ID: {process.pid}"
except FileNotFoundError:
logging.error("Ollama executable not found.")
return "Error: Ollama executable not found. Please ensure Ollama is installed and in your PATH."
except Exception as e:
logging.error(f"Error starting Ollama server: {e}")
return f"Error starting Ollama server: {e}"
def stop_ollama_server(pid):
"""
Stops the Ollama server with the specified process ID if Ollama is installed.
"""
if not is_ollama_installed():
logging.error("Ollama is not installed.")
return "Error: Ollama is not installed or not in your PATH."
try:
if platform.system() == "Windows":
subprocess.run(['taskkill', '/F', '/PID', str(pid)], check=True)
elif platform.system() in ["Linux", "Darwin"]:
os.kill(pid, signal.SIGTERM)
logging.info(f"Stopped Ollama server with PID {pid}")
return f"Stopped Ollama server with PID {pid}"
except ProcessLookupError:
logging.warning(f"No process found with PID {pid}")
return f"No process found with PID {pid}"
except Exception as e:
logging.error(f"Error stopping Ollama server: {e}")
return f"Error stopping Ollama server: {e}"
def create_ollama_tab():
"""
Creates the Ollama Model Serving tab in the Gradio interface with lazy loading.
"""
ollama_installed = is_ollama_installed()
with gr.Tab("Ollama Model Serving"):
if not ollama_installed:
gr.Markdown(
"# Ollama Model Serving\n\n"
"**Ollama is not installed or not found in your PATH. Please install Ollama to use this feature.**"
)
return # Exit early, no need to add further components
gr.Markdown("# Ollama Model Serving")
with gr.Row():
# Initialize Dropdowns with placeholders
model_list = gr.Dropdown(
label="Available Models",
choices=["Click 'Refresh Model List' to load models"],
value="Click 'Refresh Model List' to load models"
)
refresh_button = gr.Button("Refresh Model List")
with gr.Row():
new_model_name = gr.Textbox(label="Model to Pull", placeholder="Enter model name")
pull_button = gr.Button("Pull Model")
pull_output = gr.Textbox(label="Pull Status")
with gr.Row():
serve_model = gr.Dropdown(
label="Model to Serve",
choices=["Click 'Refresh Model List' to load models"],
value="Click 'Refresh Model List' to load models"
)
port = gr.Number(label="Port", value=11434, precision=0)
serve_button = gr.Button("Start Server")
serve_output = gr.Textbox(label="Server Status")
with gr.Row():
pid = gr.Number(label="Server Process ID (Enter the PID to stop)", precision=0)
stop_button = gr.Button("Stop Server")
stop_output = gr.Textbox(label="Stop Status")
def update_model_lists():
"""
Retrieves the list of available Ollama models and updates the dropdowns.
"""
models = get_ollama_models()
if models:
return gr.update(choices=models, value=models[0]), gr.update(choices=models, value=models[0])
else:
return gr.update(choices=["No models found"], value="No models found"), gr.update(choices=["No models found"], value="No models found")
def async_update_model_lists():
"""
Asynchronously updates the model lists to prevent blocking.
"""
def task():
choices1, choices2 = update_model_lists()
model_list.update(choices=choices1['choices'], value=choices1.get('value'))
serve_model.update(choices=choices2['choices'], value=choices2.get('value'))
threading.Thread(target=task).start()
# Bind the refresh button to the asynchronous update function
refresh_button.click(fn=async_update_model_lists, inputs=[], outputs=[])
# Bind the pull, serve, and stop buttons to their respective functions
pull_button.click(fn=pull_ollama_model, inputs=[new_model_name], outputs=[pull_output])
serve_button.click(fn=serve_ollama_model, inputs=[serve_model, port], outputs=[serve_output])
stop_button.click(fn=stop_ollama_server, inputs=[pid], outputs=[stop_output])
|