#!/usr/bin/env python3 """Automated notebook execution, PDF conversion, and Slack notification.""" import logging import os import subprocess import sys from datetime import datetime from pathlib import Path from typing import Dict, Optional import requests import yaml def setup_logging(log_dir: Path) -> logging.Logger: """Set up logging with file and console handlers.""" log_dir.mkdir(parents=True, exist_ok=True) log_file = log_dir / f"notebook_automation_{datetime.now().strftime('%Y%m%d')}.log" logger = logging.getLogger("notebook_automation") logger.setLevel(logging.INFO) # File handler file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.INFO) # Console handler console_handler = logging.StreamHandler() console_handler.setLevel(logging.INFO) # Formatter formatter = logging.Formatter( "%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) file_handler.setFormatter(formatter) console_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.addHandler(console_handler) return logger def load_config(config_path: Path) -> Dict: """Load configuration from YAML file.""" with open(config_path, "r") as f: return yaml.safe_load(f) def execute_notebook(notebook_path: Path, logger: logging.Logger) -> bool: """Execute notebook in-place.""" try: logger.info(f"Starting notebook execution: {notebook_path}") # Execute notebook without PDF conversion cmd = [ "jupyter", "nbconvert", "--execute", "--to", "notebook", "--inplace", "--ExecutePreprocessor.timeout=3600", # 1 hour timeout "--allow-errors", # Continue execution even if there are errors str(notebook_path), ] # Set environment to disable tqdm progress bars env = os.environ.copy() env["TQDM_DISABLE"] = "1" result = subprocess.run( cmd, capture_output=True, text=True, check=False, env=env ) if result.returncode != 0: logger.error(f"Notebook execution failed: {result.stderr}") return False logger.info("Notebook executed successfully") return True except Exception as e: logger.error(f"Error during notebook execution: {e}") return False def extract_plots_from_notebook( notebook_path: Path, reports_dir: Path, logger: logging.Logger ) -> list[Path]: """Extract plots from executed notebook and save as JPEGs.""" import json import base64 from PIL import Image import io try: logger.info(f"Extracting plots from notebook: {notebook_path}") # Read the executed notebook with open(notebook_path, "r") as f: nb = json.load(f) # Define the cells we want to extract plots from target_cells = [ "get_preference_counts(interesting_clips)", "get_preference_counts(user_intersting_clips_3p5)", ] extracted_plots = [] timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M") # Find and extract plots from target cells for cell in nb["cells"]: if cell["cell_type"] == "code": source = "".join(cell.get("source", [])) for idx, target in enumerate(target_cells): if target in source: # Look for image in outputs for output in cell.get("outputs", []): if output.get("output_type") == "display_data": data = output.get("data", {}) if "image/png" in data: plot_data = data["image/png"] logger.info(f"Found plot for: {target}") # Decode and save the image plot_name = ( "interesting_clips" if idx == 0 else "user_intersting_clips_3p5" ) output_path = ( reports_dir / f"preference_plot_{plot_name}_{timestamp}.jpg" ) # Decode base64 and save as JPEG image_data = base64.b64decode(plot_data) # Convert PNG to JPEG using PIL img = Image.open(io.BytesIO(image_data)) # Convert RGBA to RGB if necessary if img.mode in ("RGBA", "LA", "P"): background = Image.new( "RGB", img.size, (255, 255, 255) ) if img.mode == "P": img = img.convert("RGBA") background.paste( img, mask=img.split()[-1] if img.mode == "RGBA" else None, ) img = background img.save(output_path, "JPEG", quality=95) logger.info(f"Plot saved to: {output_path}") extracted_plots.append(output_path) break if not extracted_plots: logger.error("No plots found in notebook outputs") return [] logger.info(f"Successfully extracted {len(extracted_plots)} plot(s)") return extracted_plots except Exception as e: logger.error(f"Error extracting plots from notebook: {e}") return [] def send_slack_webhook(webhook_url: str, message: str, logger: logging.Logger) -> bool: """Send a message to Slack via webhook.""" try: payload = {"text": message} response = requests.post(webhook_url, json=payload, timeout=30) response.raise_for_status() logger.info("Slack message sent successfully via webhook") return True except Exception as e: logger.error(f"Failed to send Slack webhook: {e}") return False def get_channel_id( bot_token: str, channel_name: str, logger: logging.Logger ) -> Optional[str]: """ Get Slack channel ID from channel name. Args: bot_token: Slack bot token channel_name: Channel name (with or without #) logger: Logger instance Returns: Channel ID if found, None otherwise """ try: # Remove # if present clean_name = channel_name.lstrip("#") headers = {"Authorization": f"Bearer {bot_token}"} response = requests.get( "https://slack.com/api/conversations.list", headers=headers, params={"types": "public_channel,private_channel", "limit": 1000}, timeout=30, ) result = response.json() if result.get("ok"): for channel in result.get("channels", []): if channel.get("name") == clean_name: return channel.get("id") logger.error(f"Channel '{channel_name}' not found") return None except Exception as e: logger.error(f"Failed to get channel ID: {e}") return None def upload_file_to_slack_webhook( webhook_url: str, bot_token: Optional[str], channel: str, file_path: Path, message: str, logger: logging.Logger, ) -> bool: """ Upload a file to Slack channel. Args: webhook_url: Slack webhook URL (for text message) bot_token: Slack bot token for file upload channel: Slack channel name or ID file_path: Path to file to upload message: Message to accompany the file logger: Logger instance Returns: True if successful, False otherwise """ try: # Use bot token if available for file upload (better for files) if bot_token: headers = {"Authorization": f"Bearer {bot_token}"} # Check if channel is already an ID (starts with C, G, D) or a name if channel.startswith(("C", "G", "D")): # Already a channel ID channel_id = channel logger.info(f"Using provided channel ID: {channel_id}") else: # Get channel ID from channel name channel_id = get_channel_id(bot_token, channel, logger) if not channel_id: logger.error(f"Could not find channel ID for '{channel}'") send_slack_webhook(webhook_url, message, logger) return False # Step 1: Get upload URL file_size = file_path.stat().st_size upload_url_response = requests.post( "https://slack.com/api/files.getUploadURLExternal", headers=headers, data={ "filename": file_path.name, "length": file_size, }, timeout=30, ) upload_data = upload_url_response.json() if not upload_data.get("ok"): error = upload_data.get("error") logger.error(f"Failed to get upload URL: {error}") send_slack_webhook(webhook_url, message, logger) return False upload_url = upload_data["upload_url"] file_id = upload_data["file_id"] # Step 2: Upload file to the URL with open(file_path, "rb") as f: upload_response = requests.post( upload_url, files={"file": f}, timeout=120, ) if upload_response.status_code != 200: logger.error( f"File upload to URL failed: {upload_response.status_code}" ) send_slack_webhook(webhook_url, message, logger) return False # Step 3: Complete the upload and share to channel import json as json_module complete_response = requests.post( "https://slack.com/api/files.completeUploadExternal", headers=headers, data={ "files": json_module.dumps( [{"id": file_id, "title": file_path.name}] ), "channel_id": channel_id, "initial_comment": message, }, timeout=30, ) result = complete_response.json() if result.get("ok"): logger.info(f"File uploaded to Slack successfully: {file_path.name}") return True else: error = result.get("error") logger.error(f"Slack file upload failed: {error}") logger.error(f"Full response: {result}") # Fall back to just sending a message send_slack_webhook(webhook_url, message, logger) return False else: # Without bot token, just send a message via webhook logger.warning("No bot token provided, sending message only (no file)") return send_slack_webhook(webhook_url, message, logger) except Exception as e: logger.error(f"Failed to upload file to Slack: {e}") # Try to send error message send_slack_webhook( webhook_url, f"{message}\n\n⚠️ Failed to upload file: {e}", logger ) return False def cleanup_old_files( directory: Path, days_to_keep: int, logger: logging.Logger ) -> None: """ Remove files older than specified days. Args: directory: Directory to clean days_to_keep: Number of days to keep files logger: Logger instance """ try: cutoff_time = datetime.now().timestamp() - (days_to_keep * 86400) removed_count = 0 for file_path in directory.glob("*"): if file_path.is_file() and file_path.stat().st_mtime < cutoff_time: file_path.unlink() removed_count += 1 if removed_count > 0: logger.info(f"Cleaned up {removed_count} old files from {directory}") except Exception as e: logger.warning(f"Failed to cleanup old files: {e}") def main() -> int: """ Main execution function. Returns: Exit code (0 for success, 1 for failure) """ import argparse # Parse command line arguments parser = argparse.ArgumentParser( description="Execute notebook and send plot to Slack" ) parser.add_argument( "--skip-notebook", action="store_true", help="Skip notebook execution and upload most recent plot", ) args = parser.parse_args() # Determine script directory and config path script_dir = Path(__file__).parent config_path = script_dir / "config.yaml" # Setup logging first log_dir = script_dir / "logs" logger = setup_logging(log_dir) try: # Load configuration logger.info("Loading configuration") config = load_config(config_path) # Setup paths notebook_path = script_dir / config["notebook_path"] reports_dir = script_dir / config["reports_dir"] reports_dir.mkdir(parents=True, exist_ok=True) # Generate timestamp timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M") # Execute notebook if not skipped if not args.skip_notebook: logger.info("Executing notebook...") if not execute_notebook(notebook_path, logger): error_msg = f"❌ Notebook execution failed at {timestamp}\nCheck logs for details." send_slack_webhook(config["slack_webhook_url"], error_msg, logger) return 1 else: logger.info("Skipping notebook execution") # Extract plots from notebook output_jpegs = extract_plots_from_notebook(notebook_path, reports_dir, logger) if not output_jpegs: # If extraction fails, try to find existing plots logger.warning("Plot extraction failed, looking for existing plots...") jpeg_files = sorted( reports_dir.glob("preference_plot_*.jpg"), key=lambda p: p.stat().st_mtime, reverse=True, ) if not jpeg_files: error_msg = "❌ No plots found in notebook output or reports directory\nCheck logs for details." send_slack_webhook(config["slack_webhook_url"], error_msg, logger) return 1 # Use the two most recent plots output_jpegs = jpeg_files[:2] logger.info(f"Using {len(output_jpegs)} existing plot(s)") # Send plots to Slack for idx, output_jpeg in enumerate(output_jpegs, 1): plot_type = ( "interesting_clips" if "interesting_clips" in output_jpeg.name else "user_intersting_clips_3p5" ) success_msg = ( f"✅ Preference Plot Generated ({idx}/{len(output_jpegs)})\n" f"Dataset: {plot_type}\n" f"Timestamp: {timestamp}\n" f"File: {output_jpeg.name}" ) upload_success = upload_file_to_slack_webhook( webhook_url=config["slack_webhook_url"], bot_token=config.get("slack_bot_token"), channel=config["slack_channel"], file_path=output_jpeg, message=success_msg, logger=logger, ) if not upload_success: logger.warning( f"Upload of {output_jpeg.name} had issues, continuing..." ) # Cleanup old files (keep last 1 day for reports, 7 days for logs) cleanup_old_files(reports_dir, 1, logger) cleanup_old_files(log_dir, 7, logger) logger.info("Automation completed successfully") return 0 except Exception as e: logger.error(f"Fatal error in main execution: {e}", exc_info=True) try: config = load_config(config_path) error_msg = f"❌ Critical error in notebook automation:\n{str(e)}" send_slack_webhook(config["slack_webhook_url"], error_msg, logger) except Exception: pass return 1 if __name__ == "__main__": sys.exit(main())