ETF_Suite_Portal/cli_manager.py

265 lines
10 KiB
Python
Executable File

#!/usr/bin/env python3
import os
import sys
import time
import signal
import logging
import argparse
import subprocess
from pathlib import Path
from typing import Optional, Dict, Any
import psutil
class CLIManager:
"""CLI Manager for ETF Portal application.
Provides commands to:
- start: Launch the application
- stop: Gracefully stop the application
- restart: Stop and start the application
- status: Check if the application is running
"""
def __init__(self):
self.app_name = "ETF_Portal"
self.pid_file = Path("/var/run/etf-portal/etf_portal.pid")
self.log_file = Path("/var/log/etf-portal/cli_manager.log")
self._setup_logging()
def _setup_logging(self):
"""Configure logging for the CLI manager."""
self.log_file.parent.mkdir(parents=True, exist_ok=True)
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler(self.log_file),
logging.StreamHandler()
]
)
self.logger = logging.getLogger("CLIManager")
def _get_pid(self) -> Optional[int]:
"""Get the PID from the PID file if it exists."""
try:
if self.pid_file.exists():
with open(self.pid_file, 'r') as f:
pid = int(f.read().strip())
return pid
except (ValueError, IOError) as e:
self.logger.error(f"Error reading PID file: {e}")
return None
def _save_pid(self, pid: int):
"""Save the PID to the PID file."""
try:
self.pid_file.parent.mkdir(parents=True, exist_ok=True)
with open(self.pid_file, 'w') as f:
f.write(str(pid))
except IOError as e:
self.logger.error(f"Error saving PID file: {e}")
def _is_process_running(self, pid: int) -> bool:
"""Check if a process with the given PID is running."""
try:
process = psutil.Process(pid)
# Check if it's a Python process running Streamlit
is_running = process.is_running()
is_streamlit = any('streamlit' in cmd.lower() for cmd in process.cmdline())
is_root = process.uids().real == 0 # Check if process is running as root
if is_running and is_streamlit:
if is_root:
self.logger.warning(f"Process {pid} is running as root")
return True
return False
except psutil.NoSuchProcess:
return False
except psutil.AccessDenied:
self.logger.warning(f"Access denied to process {pid}")
return False
def _find_streamlit_process(self) -> Optional[int]:
"""Find the main Streamlit process."""
for proc in psutil.process_iter(['pid', 'cmdline', 'uids']):
try:
if any('streamlit' in cmd.lower() for cmd in proc.info['cmdline']):
if proc.info['uids'].real == 0:
self.logger.warning(f"Found Streamlit process {proc.info['pid']} running as root")
return proc.info['pid']
except (psutil.NoSuchProcess, psutil.AccessDenied):
continue
return None
def start(self):
"""Start the ETF Portal application."""
# Check if Streamlit is already running
existing_pid = self._find_streamlit_process()
if existing_pid:
self.logger.info(f"{self.app_name} is already running with PID {existing_pid}")
self._save_pid(existing_pid)
return
try:
# Get the full path to streamlit in the virtual environment
venv_streamlit = os.path.join(os.path.dirname(sys.executable), "streamlit")
if not os.path.exists(venv_streamlit):
raise Exception(f"Streamlit not found at {venv_streamlit}")
# Ensure we're using the virtual environment's Python
env = os.environ.copy()
env["PATH"] = os.path.dirname(sys.executable) + ":" + env["PATH"]
# Start the application using streamlit
process = subprocess.Popen(
[venv_streamlit, "run", "ETF_Suite_Launcher.py", "--server.port=8500"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
start_new_session=True,
env=env,
preexec_fn=os.setsid # Create a new session
)
# Wait a moment for Streamlit to start
time.sleep(2)
# Find the actual Streamlit process
streamlit_pid = self._find_streamlit_process()
if streamlit_pid:
# Verify the process is running with correct permissions
try:
proc = psutil.Process(streamlit_pid)
if proc.uids().real == 0:
self.logger.error("Process started as root. Stopping...")
self.stop()
raise Exception("Process started with incorrect permissions")
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
self._save_pid(streamlit_pid)
self.logger.info(f"{self.app_name} started with PID {streamlit_pid}")
else:
raise Exception("Streamlit process not found after startup")
except Exception as e:
self.logger.error(f"Failed to start {self.app_name}: {e}")
sys.exit(1)
def stop(self):
"""Stop the ETF Portal application."""
# First check the PID file
pid = self._get_pid()
if not pid:
# If no PID file, try to find running Streamlit process
pid = self._find_streamlit_process()
if not pid:
self.logger.info(f"{self.app_name} is not running")
return
try:
if self._is_process_running(pid):
# Get the process using psutil
process = psutil.Process(pid)
# Check if process is running as root
if process.uids().real == 0:
self.logger.error(f"Cannot stop process {pid} - it is running as root")
self.logger.info("Please use sudo to stop the process:")
self.logger.info(f"sudo kill {pid}")
return
# Try to terminate the process and its children
try:
# Get all child processes
children = process.children(recursive=True)
# Terminate children first
for child in children:
try:
child.terminate()
except psutil.NoSuchProcess:
pass
# Terminate the main process
process.terminate()
# Wait for processes to terminate
gone, alive = psutil.wait_procs([process] + children, timeout=10)
# If any processes are still alive, force kill them
for p in alive:
try:
p.kill()
except psutil.NoSuchProcess:
pass
self.logger.info(f"{self.app_name} stopped")
except psutil.NoSuchProcess:
self.logger.info(f"{self.app_name} is not running")
else:
self.logger.info(f"{self.app_name} is not running")
# Remove PID file
if self.pid_file.exists():
self.pid_file.unlink()
except Exception as e:
self.logger.error(f"Error stopping {self.app_name}: {e}")
# Don't exit with error, just log it
self.logger.info("Attempting to clean up PID file...")
if self.pid_file.exists():
self.pid_file.unlink()
def restart(self):
"""Restart the ETF Portal application."""
self.logger.info("Restarting ETF Portal...")
self.stop()
time.sleep(2) # Wait for processes to fully stop
self.start()
def status(self):
"""Check the status of the ETF Portal application."""
# First check the PID file
pid = self._get_pid()
if not pid:
# If no PID file, try to find running Streamlit process
pid = self._find_streamlit_process()
if not pid:
print(f"{self.app_name} is not running")
return
if self._is_process_running(pid):
print(f"{self.app_name} is running with PID {pid}")
else:
print(f"{self.app_name} is not running (stale PID file)")
if self.pid_file.exists():
self.pid_file.unlink()
def main():
"""Main entry point for the CLI."""
parser = argparse.ArgumentParser(
description="ETF Portal CLI Manager",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
etf-portal start # Start the ETF Portal application
etf-portal stop # Stop the ETF Portal application
etf-portal restart # Restart the ETF Portal application
etf-portal status # Check the status of the ETF Portal application
"""
)
parser.add_argument(
"command",
choices=["start", "stop", "restart", "status"],
help="Command to execute"
)
args = parser.parse_args()
manager = CLIManager()
# Execute the requested command
getattr(manager, args.command)()
if __name__ == "__main__":
main()