diff --git a/procrastinate/cli.py b/procrastinate/cli.py index 7fb43364f..1b0d6bcb8 100644 --- a/procrastinate/cli.py +++ b/procrastinate/cli.py @@ -336,6 +336,49 @@ def configure_worker_parser(subparsers: argparse._SubParsersAction): envvar_help="use 0/n/f or 1/y/t", envvar_type=env_bool, ) + + # Reload options (similar to uvicorn) + add_argument( + worker_parser, + "--reload", + action="store_true", + help="Enable auto-reload for development", + envvar="WORKER_RELOAD", + envvar_help="use 0/n/f or 1/y/t", + envvar_type=env_bool, + ) + add_argument( + worker_parser, + "--reload-dir", + dest="reload_dirs", + action="append", + help="Set reload directories explicitly, instead of using current working directory. May be used multiple times", + envvar="WORKER_RELOAD_DIRS", + ) + add_argument( + worker_parser, + "--reload-include", + dest="reload_includes", + action="append", + help="Set glob patterns to include while watching for files. Includes '*.py' by default. May be used multiple times", + envvar="WORKER_RELOAD_INCLUDES", + ) + add_argument( + worker_parser, + "--reload-exclude", + dest="reload_excludes", + action="append", + help="Set glob patterns to exclude while watching for files. May be used multiple times", + envvar="WORKER_RELOAD_EXCLUDES", + ) + add_argument( + worker_parser, + "--reload-delay", + type=float, + default=0.25, + help="Delay between previous and next check if application needs to be reloaded", + envvar="WORKER_RELOAD_DELAY", + ) add_argument( worker_parser, "--delete-jobs", @@ -551,6 +594,11 @@ async def worker_( Launch a worker, listening on the given queues (or all queues). Values default to App.worker_defaults and then App.run_worker() defaults values. """ + # Handle reload mode + if kwargs.get("reload"): + return run_worker_with_reload(app, **kwargs) + + # Standard worker mode queues = kwargs.get("queues") print_stderr( f"Launching a worker on {'all queues' if not queues else ', '.join(queues)}" @@ -674,6 +722,61 @@ async def shell_(app: procrastinate.App, shell_command: list[str]): await utils.sync_to_async(shell_obj.cmdloop) +def run_worker_with_reload(app: procrastinate.App, **kwargs) -> int: + """ + Run worker with auto-reload functionality. + + This function switches from async context to sync and uses subprocess + to enable file watching and process restarting. + """ + # Import here to avoid import errors if watchfiles not installed + from .reloader import run_with_reload + + # Extract reload configuration + reload_config = { + "reload_dirs": kwargs.get("reload_dirs"), + "reload_includes": kwargs.get("reload_includes"), + "reload_excludes": kwargs.get("reload_excludes"), + "reload_delay": kwargs.get("reload_delay", 0.25), + } + + # Build command to run worker without reload + # We need to reconstruct the CLI command + + # Start with base command + cmd = [sys.executable, "-m", "procrastinate"] + + # Add app parameter if available + app_import_string = getattr(app, "_import_string", None) + if app_import_string: + cmd.extend(["--app", app_import_string]) + + # Add worker subcommand + cmd.append("worker") + + # Add worker-specific arguments (excluding reload options) + worker_args = { + k: v + for k, v in kwargs.items() + if not k.startswith("reload") and k != "reload" and v is not None + } + + for key, value in worker_args.items(): + cli_key = key.replace("_", "-") + if isinstance(value, bool): + if value: + cmd.append(f"--{cli_key}") + elif isinstance(value, list): + for item in value: + cmd.extend([f"--{cli_key}", str(item)]) + else: + cmd.extend([f"--{cli_key}", str(value)]) + + # Run with reload + print_stderr("๐Ÿ”„ Starting Procrastinate worker with auto-reload enabled") + return run_with_reload(cmd, **reload_config) + + def main(): if os.name == "nt": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) diff --git a/procrastinate/reloader.py b/procrastinate/reloader.py new file mode 100644 index 000000000..b15acd680 --- /dev/null +++ b/procrastinate/reloader.py @@ -0,0 +1,303 @@ +""" +Auto-reload functionality for Procrastinate CLI worker command. + +This module provides file watching and worker restarting capabilities +similar to uvicorn's --reload functionality, using watchfiles with inotify +for efficient file system monitoring. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import signal +import sys +from pathlib import Path +from typing import Any + +try: + from watchfiles import awatch # type: ignore[import-untyped] + + HAS_WATCHFILES = True +except ImportError: + HAS_WATCHFILES = False + awatch = None # type: ignore[assignment] + +logger = logging.getLogger(__name__) + + +class ChangeReloader: + """ + Auto-reload supervisor for Procrastinate worker CLI. + + Monitors file changes and restarts worker processes, + similar to uvicorn's reload functionality. + """ + + def __init__( + self, + target_cmd: list[str], + *, + reload_dirs: list[str] | None = None, + reload_includes: list[str] | None = None, + reload_excludes: list[str] | None = None, + reload_delay: float = 0.25, + ) -> None: + """ + Initialize the change reloader. + + Args: + target_cmd: Command and arguments to run as subprocess + reload_dirs: Directories to watch (default: current directory) + reload_includes: Glob patterns to include (default: ['*.py']) + reload_excludes: Glob patterns to exclude + reload_delay: Delay between file change and restart + """ + if not HAS_WATCHFILES: + raise RuntimeError( + "watchfiles is required for --reload functionality. " + "Install with: pip install 'procrastinate[reload]'" + ) + + self.target_cmd = target_cmd + self.reload_dirs = reload_dirs or [os.getcwd()] + self.reload_includes = reload_includes or ["*.py"] + self.reload_excludes = reload_excludes or [] + self.reload_delay = reload_delay + + self.should_exit = asyncio.Event() + self.process: asyncio.subprocess.Process | None = None + + def startup_message(self) -> None: + """Print startup information about file watching.""" + dirs = ", ".join(self.reload_dirs) + includes = ( + ", ".join(self.reload_includes) if self.reload_includes else "all files" + ) + + logger.info("๐Ÿ”„ Started reloader process using watchfiles") + logger.info(f"๐Ÿ‘๏ธ Watching for file changes in: {dirs}") + logger.info(f"๐Ÿ“„ Including patterns: {includes}") + + if self.reload_excludes: + excludes = ", ".join(self.reload_excludes) + logger.info(f"๐Ÿšซ Excluding patterns: {excludes}") + + def should_reload(self, path: str) -> bool: + """ + Check if a file change should trigger a reload. + + Args: + path: Path to the changed file + + Returns: + True if the change should trigger a reload + """ + from fnmatch import fnmatch + + path_obj = Path(path) + + # Check excludes first (more restrictive) + for exclude_pattern in self.reload_excludes: + if fnmatch(str(path_obj), exclude_pattern): + return False + + # Check includes + for include_pattern in self.reload_includes: + if fnmatch(str(path_obj), include_pattern): + return True + + # Default: don't reload if no patterns matched + return False + + async def start_process(self) -> None: + """Start the target process.""" + logger.info(f"๐Ÿš€ Starting process: {' '.join(self.target_cmd)}") + + self.process = await asyncio.create_subprocess_exec( + *self.target_cmd, + stdout=sys.stdout, + stderr=sys.stderr, + stdin=sys.stdin, + ) + + logger.info(f"โœ… Process started with PID {self.process.pid}") + + async def restart_process(self) -> None: + """Restart the target process.""" + if self.process: + logger.info("๐Ÿ”„ Stopping current process...") + + # Try graceful shutdown first + self.process.terminate() + + try: + await asyncio.wait_for(self.process.wait(), timeout=10.0) + logger.info("โœ… Process stopped gracefully") + except asyncio.TimeoutError: + logger.warning("โš ๏ธ Process didn't stop gracefully, killing...") + self.process.kill() + await self.process.wait() + logger.info("๐Ÿ’€ Process killed") + + # Start new process + await self.start_process() + + async def watch_files(self) -> None: + """Watch for file changes and trigger restarts.""" + if awatch is None: + raise RuntimeError("watchfiles not available") + + # Convert string paths to Path objects for validation + watch_paths = [ + str(Path(d).resolve()) for d in self.reload_dirs if Path(d).exists() + ] + + if not watch_paths: + logger.warning("โš ๏ธ No valid watch directories found") + return + + logger.info(f"๐Ÿ‘๏ธ Watching for changes in: {watch_paths}") + + try: + async for changes in awatch(*watch_paths): + if self.should_exit.is_set(): + break + + # Filter changes based on patterns + relevant_changes = [ + (change_type, path) + for change_type, path in changes + if self.should_reload(path) + ] + + if not relevant_changes: + continue + + logger.info( + f"๐Ÿ“ Detected {len(relevant_changes)} relevant file changes:" + ) + for change_type, path in relevant_changes: + logger.info(f" {change_type.name}: {path}") + + # Debounce rapid changes + if self.reload_delay > 0: + logger.info(f"โณ Waiting {self.reload_delay}s before restart...") + await asyncio.sleep(self.reload_delay) + + # Check if we should still exit (in case signal came during delay) + if not self.should_exit.is_set(): + await self.restart_process() + + except Exception as e: + logger.error(f"โŒ Error watching files: {e}") + + def setup_signal_handlers(self) -> None: + """Setup signal handlers for graceful shutdown.""" + + def signal_handler(signum: int) -> None: + logger.info(f"๐Ÿ“ก Received signal {signum}, shutting down...") + self.should_exit.set() + + # Handle common termination signals + signal.signal(signal.SIGTERM, lambda s, f: signal_handler(s)) + signal.signal(signal.SIGINT, lambda s, f: signal_handler(s)) + + # Handle additional signals on Unix + if hasattr(signal, "SIGHUP"): + signal.signal(signal.SIGHUP, lambda s, f: signal_handler(s)) + + async def run(self) -> int: + """ + Main run loop for the reloader. + + Returns: + Exit code from the target process + """ + self.setup_signal_handlers() + self.startup_message() + + try: + # Start initial process + await self.start_process() + + # Start file watching + watch_task = asyncio.create_task(self.watch_files()) + + # Wait for either process completion or shutdown signal + if self.process: + process_task = asyncio.create_task(self.process.wait()) + + done, pending = await asyncio.wait( + [process_task, asyncio.create_task(self.should_exit.wait())], + return_when=asyncio.FIRST_COMPLETED, + ) + + # Cancel pending tasks + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Get exit code if process completed + if process_task in done: + exit_code = process_task.result() + logger.info(f"๐Ÿ›‘ Process exited with code {exit_code}") + else: + exit_code = 0 # Shutdown signal received + else: + exit_code = 1 # Process failed to start + + # Cancel file watcher + watch_task.cancel() + try: + await watch_task + except asyncio.CancelledError: + pass + + # Cleanup process if still running + if self.process and self.process.returncode is None: + logger.info("๐Ÿงน Cleaning up process...") + self.process.terminate() + try: + await asyncio.wait_for(self.process.wait(), timeout=5.0) + except asyncio.TimeoutError: + self.process.kill() + await self.process.wait() + + return exit_code + + except KeyboardInterrupt: + logger.info("โŒจ๏ธ Keyboard interrupt received") + return 0 + except Exception as e: + logger.error(f"๐Ÿ’ฅ Unexpected error in reloader: {e}") + return 1 + + +def run_with_reload( + target_cmd: list[str], + **reload_config: Any, +) -> int: + """ + Run a command with auto-reload functionality. + + Args: + target_cmd: Command and arguments to execute + **reload_config: Configuration for reload behavior + + Returns: + Exit code from the target process + """ + reloader = ChangeReloader(target_cmd, **reload_config) + + try: + return asyncio.run(reloader.run()) + except KeyboardInterrupt: + return 0 + except Exception as e: + logger.error(f"๐Ÿ’ฅ Failed to run with reload: {e}") + return 1 diff --git a/pyproject.toml b/pyproject.toml index 2cf38e5cf..fa5554578 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ sqlalchemy = ["sqlalchemy~=2.0"] aiopg = ["aiopg", "psycopg2-binary"] psycopg2 = ["psycopg2-binary"] sphinx = ["sphinx"] +reload = ["watchfiles>=0.21.0"] [project.urls] homepage = "https://procrastinate.readthedocs.io/"