start_jupyter.py

#!/usr/bin/env python3
"""
PyTorch Jupyter Lab Starter Script
A cross-platform Python script to set up and start Jupyter Lab with PyTorch
"""

import os
import sys
import subprocess
import platform

def run_command(cmd, shell=True):
    """Run a command and return the result"""
    try:
        result = subprocess.run(cmd, shell=shell, capture_output=True, text=True)
        return result.returncode == 0, result.stdout, result.stderr
    except Exception as e:
        return False, "", str(e)

def check_conda():
    """Check if conda is available"""
    success, _, _ = run_command("conda --version")
    return success

def check_environment_exists(env_name):
    """Check if conda environment exists"""
    success, output, _ = run_command("conda env list")
    if success:
        return env_name in output
    return False

def create_environment(env_name):
    """Create and setup the conda environment"""
    print("📦 Creating environment...")
    
    # Create environment
    success, _, error = run_command(f"conda create -n {env_name} python=3.11 -y")
    if not success:
        print(f"❌ Failed to create environment: {error}")
        return False
    
    # Install PyTorch
    print("📥 Installing PyTorch...")
    if platform.system() == "Darwin":  # macOS
        pytorch_cmd = f"conda install -n {env_name} pytorch torchvision torchaudio -c pytorch -y"
    else:  # Windows/Linux
        pytorch_cmd = f"conda install -n {env_name} pytorch torchvision torchaudio cpuonly -c pytorch -y"
    
    success, _, error = run_command(pytorch_cmd)
    if not success:
        print(f"❌ Failed to install PyTorch: {error}")
        return False
    
    # Install Java
    print("📥 Installing Java (OpenJDK) for PySpark...")
    java_cmd = f"conda install -n {env_name} openjdk=17 -c conda-forge -y"
    success, _, error = run_command(java_cmd)
    if not success:
        print(f"❌ Failed to install Java: {error}")
        return False
    
    # Install PySpark
    print("📥 Installing PySpark...")
    pyspark_cmd = f"conda install -n {env_name} pyspark -c conda-forge -y"
    success, _, error = run_command(pyspark_cmd)
    if not success:
        print(f"❌ Failed to install PySpark: {error}")
        return False
    
    # Install Jupyter and other packages
    print("📥 Installing Jupyter and data science packages...")
    jupyter_cmd = f"conda install -n {env_name} jupyterlab matplotlib seaborn pandas numpy scipy tensorboard tqdm -c conda-forge -y"
    success, _, error = run_command(jupyter_cmd)
    if not success:
        print(f"❌ Failed to install Jupyter: {error}")
        return False
    
    return True

def verify_pytorch(env_name):
    """Verify PyTorch installation"""
    print("🔍 Checking PyTorch installation...")
    
    if platform.system() == "Windows":
        activate_cmd = f"conda activate {env_name} && python -c \"import torch; print(f'PyTorch version: {{torch.__version__}}'); print('✅ PyTorch is ready!')\""
    else:
        activate_cmd = f"source $(conda info --base)/etc/profile.d/conda.sh && conda activate {env_name} && python -c \"import torch; print(f'PyTorch version: {{torch.__version__}}'); print('✅ PyTorch is ready!')\""
    
    success, output, error = run_command(activate_cmd)
    if success:
        print(output)
        return True
    else:
        print(f"⚠️ PyTorch verification failed: {error}")
        return False

def verify_java_pyspark(env_name):
    """Verify Java and PySpark installation"""
    print("🔍 Checking Java and PySpark installation...")
    
    if platform.system() == "Windows":
        activate_cmd = f"conda activate {env_name} && set JAVA_HOME=%CONDA_PREFIX%\\lib\\jvm && set PATH=%JAVA_HOME%\\bin;%PATH% && set PYSPARK_PYTHON=%CONDA_PREFIX%\\python.exe && set PYSPARK_DRIVER_PYTHON=%CONDA_PREFIX%\\python.exe && python -c \"import os; print(f'JAVA_HOME: {{os.environ.get(\"JAVA_HOME\", \"Not set\")}}'); print(f'Java binary: {{os.popen(\"where java\").read().strip()}}'); import pyspark; print(f'PySpark version: {{pyspark.__version__}}'); print('✅ Java and PySpark are ready!')\""
    else:
        activate_cmd = f"source $(conda info --base)/etc/profile.d/conda.sh && conda activate {env_name} && export JAVA_HOME=$CONDA_PREFIX/lib/jvm && export PATH=$JAVA_HOME/bin:$PATH && export PYSPARK_PYTHON=$CONDA_PREFIX/bin/python && export PYSPARK_DRIVER_PYTHON=$CONDA_PREFIX/bin/python && python -c \"import os; print(f'JAVA_HOME: {{os.environ.get(\"JAVA_HOME\", \"Not set\")}}'); print(f'Java binary: {{os.popen(\"which java\").read().strip()}}'); import pyspark; print(f'PySpark version: {{pyspark.__version__}}'); print('✅ Java and PySpark are ready!')\""
    
    success, output, error = run_command(activate_cmd)
    if success:
        print(output)
        return True
    else:
        print(f"⚠️ Java/PySpark verification failed: {error}")
        return False

def start_jupyter(env_name):
    """Start Jupyter Lab"""
    print("🌐 Starting Jupyter Lab...")
    
    # Get current directory
    current_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(current_dir)
    
    print(f"📁 Working directory: {current_dir}")
    print("🔗 Jupyter Lab will open in your default browser")
    print("\n📝 To stop Jupyter Lab: Press Ctrl+C")
    print("=" * 50)
    
    if platform.system() == "Windows":
        jupyter_cmd = f"conda activate {env_name} && set JAVA_HOME=%CONDA_PREFIX%\\lib\\jvm && set PATH=%JAVA_HOME%\\bin;%PATH% && set PYSPARK_PYTHON=%CONDA_PREFIX%\\python.exe && set PYSPARK_DRIVER_PYTHON=%CONDA_PREFIX%\\python.exe && jupyter lab --no-browser --port=8888"
    else:
        jupyter_cmd = f"source $(conda info --base)/etc/profile.d/conda.sh && conda activate {env_name} && export JAVA_HOME=$CONDA_PREFIX/lib/jvm && export PATH=$JAVA_HOME/bin:$PATH && export PYSPARK_PYTHON=$CONDA_PREFIX/bin/python && export PYSPARK_DRIVER_PYTHON=$CONDA_PREFIX/bin/python && jupyter lab --no-browser --port=8888"
    
    try:
        subprocess.run(jupyter_cmd, shell=True)
    except KeyboardInterrupt:
        print("\n👋 Jupyter Lab stopped. Have a great day!")

def main():
    """Main function"""
    print("🚀 PyTorch Jupyter Lab Starter")
    print("=" * 50)
    
    # Check conda
    if not check_conda():
        print("❌ Conda not found. Please install Miniconda or Anaconda first.")
        sys.exit(1)
    
    env_name = "pytorch-mps"
    
    # Check if environment exists
    if not check_environment_exists(env_name):
        print(f"📦 Environment '{env_name}' not found.")
        if not create_environment(env_name):
            print("❌ Failed to create environment.")
            sys.exit(1)
        print("✅ Environment setup complete!")
    else:
        print(f"✅ Environment '{env_name}' found.")
    
    # Verify PyTorch
    if not verify_pytorch(env_name):
        print("❌ PyTorch verification failed.")
        sys.exit(1)
    
    # Verify Java and PySpark
    if not verify_java_pyspark(env_name):
        print("❌ Java/PySpark verification failed.")
        sys.exit(1)
    
    # Start Jupyter Lab
    start_jupyter(env_name)

if __name__ == "__main__":
    main()