Reload job functions when --watch detects a file change

arq/cli.py

@@ -1,16 +1,17 @@
 import asyncio
+import importlib
 import logging.config
 import os
 import sys
 from signal import Signals
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import click
 
 from .logs import default_log_config
 from .utils import import_string
 from .version import VERSION
-from .worker import check_health, create_worker, run_worker
+from .worker import Worker, check_health, create_worker, run_worker
 
 if TYPE_CHECKING:
     from .typing import WorkerSettingsType
@@ -54,6 +55,48 @@ def cli(*, worker_settings: str, burst: bool, check: bool, watch: str, verbose:
             run_worker(worker_settings_, **kwargs)
 
 
+def _reload_worker_settings(worker_settings: 'WorkerSettingsType') -> 'WorkerSettingsType':
+    """Re-import worker settings so --watch runs the edited job functions."""
+    module_name = getattr(worker_settings, '__module__', None)
+    qualname = getattr(worker_settings, '__qualname__', None)
+    if not module_name or not qualname:
+        return worker_settings
+    module = sys.modules.get(module_name)
+    if module is None:
+        return worker_settings
+    # A save in the same mtime second still has to recompile; otherwise reload
+    # keeps the bytecode from startup and --watch runs the old function.
+    cached = getattr(module, '__cached__', None)
+    if isinstance(cached, str):
+        try:
+            os.remove(cached)
+        except OSError:
+            pass
+    importlib.invalidate_caches()
+    reloaded = importlib.reload(module)
+    obj: Any = reloaded
+    for part in qualname.split('.'):
+        obj = getattr(obj, part)
+    return cast('WorkerSettingsType', obj)
+
+
+def _stop_worker_for_reload(worker: Worker) -> None:
+    """Cancel the running worker without leaving --watch.
+
+    Unix uses SIGUSR1 so ``worker_on_stop`` does not set the watch stop event.
+    Windows has no SIGUSR1; cancel the tasks directly.
+    """
+    sigusr1 = getattr(Signals, 'SIGUSR1', None)
+    if sigusr1 is not None:
+        worker.handle_sig(sigusr1)
+        return
+    for task in list(worker.tasks.values()):
+        if not task.done():
+            task.cancel()
+    if worker.main_task is not None and not worker.main_task.done():
+        worker.main_task.cancel()
+
+
 async def watch_reload(path: str, worker_settings: 'WorkerSettingsType') -> None:
     try:
         from watchfiles import awatch
@@ -64,7 +107,7 @@ async def watch_reload(path: str, worker_settings: 'WorkerSettingsType') -> None
     stop_event = asyncio.Event()
 
     def worker_on_stop(s: Signals) -> None:
-        if s != Signals.SIGUSR1:  # pragma: no cover
+        if s != getattr(Signals, 'SIGUSR1', None):  # pragma: no cover
             stop_event.set()
 
     worker = create_worker(worker_settings)
@@ -73,8 +116,11 @@ def worker_on_stop(s: Signals) -> None:
         loop.create_task(worker.async_run())
         async for _ in awatch(path, stop_event=stop_event):
             print('\nfiles changed, reloading arq worker...')
-            worker.handle_sig(Signals.SIGUSR1)
+            _stop_worker_for_reload(worker)
             await worker.close()
+            worker_settings = _reload_worker_settings(worker_settings)
+            worker = create_worker(worker_settings)
+            worker.on_stop = worker_on_stop
             loop.create_task(worker.async_run())
     finally:
         await worker.close()

tests/test_cli.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 from click.testing import CliRunner
 
@@ -10,14 +12,29 @@ async def foobar(ctx):
     return 42
 
 
+def _redis_settings_from_env():
+    # --watch reloads this module. The reloaded class must still target the test Redis.
+    host = os.environ.get('ARQ_TEST_REDIS_HOST')
+    port = os.environ.get('ARQ_TEST_REDIS_PORT')
+    if not host or not port:
+        return None
+    return RedisSettings(host=host, port=int(port))
+
+
 class WorkerSettings:
     burst = True
     functions = [foobar]
+    redis_settings = _redis_settings_from_env()
 
 
 @pytest.fixture(scope='module', autouse=True)
 def setup_worker_connection(test_redis_host: str, test_redis_port: int):
+    os.environ['ARQ_TEST_REDIS_HOST'] = test_redis_host
+    os.environ['ARQ_TEST_REDIS_PORT'] = str(test_redis_port)
     WorkerSettings.redis_settings = RedisSettings(host=test_redis_host, port=test_redis_port)
+    yield
+    os.environ.pop('ARQ_TEST_REDIS_HOST', None)
+    os.environ.pop('ARQ_TEST_REDIS_PORT', None)
 
 
 def test_help():

tests/test_watch_reload.py

@@ -0,0 +1,56 @@
+import importlib
+import sys
+from types import ModuleType
+
+from arq.cli import watch_reload
+
+
+async def test_watch_reload_runs_edited_job(tmp_path, monkeypatch):
+    """--watch must import job functions again after the settings module changes."""
+    module_name = 'arq_watch_demo_429'
+    module_path = tmp_path / f'{module_name}.py'
+    source = 'async def my_task(ctx):\n    return "HELLO-1"\n\nclass WorkerSettings:\n    functions = [my_task]\n'
+    module_path.write_text(source, encoding='utf-8')
+    monkeypatch.syspath_prepend(str(tmp_path))
+    module = importlib.import_module(module_name)
+
+    async def awatch(path, stop_event=None):
+        module_path.write_text(source.replace('HELLO-1', 'HELLO-2'), encoding='utf-8')
+        yield {'changed'}
+
+    watchfiles = ModuleType('watchfiles')
+    watchfiles.awatch = awatch
+    monkeypatch.setitem(sys.modules, 'watchfiles', watchfiles)
+
+    created = []
+
+    class _FakeWorker:
+        def __init__(self, settings):
+            self.settings = settings
+            self.on_stop = None
+            self.tasks = {}
+            self.main_task = None
+
+        def handle_sig(self, sig):
+            return None
+
+        async def close(self):
+            return None
+
+        async def async_run(self):
+            return None
+
+    def fake_create(settings, **kwargs):
+        worker = _FakeWorker(settings)
+        created.append(worker)
+        return worker
+
+    monkeypatch.setattr('arq.cli.create_worker', fake_create)
+    try:
+        await watch_reload(str(tmp_path), module.WorkerSettings)
+    finally:
+        sys.modules.pop(module_name, None)
+
+    assert len(created) == 2
+    assert await created[0].settings.functions[0]({}) == 'HELLO-1'
+    assert await created[1].settings.functions[0]({}) == 'HELLO-2'