Spaces:
Running
Running
| 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]) | |