sitecustomize.py based autoinstrumentation

Author: piotrekno1

tests/test_cli.py

@@ -2,7 +2,6 @@
 
 import json
 import os
-from types import SimpleNamespace
 
 import pytest
 
@@ -12,16 +11,20 @@
 from wildedge.runtime import runner as runtime_runner
 
 
-def test_cli_run_script_invokes_runner(monkeypatch):
-    captured = {}
+def _fake_execle(captured: dict):
+    def _execle(path, *args):  # type: ignore[no-untyped-def]
+        # Last positional arg is the env dict (os.execle convention).
+        captured["path"] = path
+        captured["argv"] = list(args[:-1])
+        captured["env"] = args[-1]
 
-    def fake_run(cmd, env, check):  # type: ignore[no-untyped-def]
-        captured["cmd"] = cmd
-        captured["env"] = env
-        captured["check"] = check
-        return SimpleNamespace(returncode=0)
+    return _execle
 
-    monkeypatch.setattr(cli.subprocess, "run", fake_run)
+
+def test_cli_run_execs_command_with_env(monkeypatch):
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
 
     rc = cli.main(
         [
@@ -32,24 +35,16 @@ def fake_run(cmd, env, check):  # type: ignore[no-untyped-def]
             "1.2.3",
             "--debug",
             "--",
-            "python",
-            "app.py",
-            "--foo",
-            "bar",
+            "gunicorn",
+            "myapp.wsgi:app",
+            "--workers",
+            "4",
         ]
     )
 
     assert rc == 0
-    assert captured["cmd"][1:8] == [
-        "-m",
-        "wildedge.runtime.runner",
-        "--mode",
-        "script",
-        "--target",
-        "app.py",
-        "--",
-    ]
-    assert captured["cmd"][8:] == ["--foo", "bar"]
+    assert captured["path"] == "/usr/bin/gunicorn"
+    assert captured["argv"] == ["/usr/bin/gunicorn", "myapp.wsgi:app", "--workers", "4"]
     assert (
         captured["env"][bootstrap.RUN_DSN_ENV]
         == "https://secret@ingest.wildedge.dev/key"
@@ -64,62 +59,86 @@ def fake_run(cmd, env, check):  # type: ignore[no-untyped-def]
     )
 
 
-def test_cli_run_sets_no_propagate_and_strict(monkeypatch):
-    captured = {}
+def test_cli_run_prepends_autoload_to_pythonpath(monkeypatch):
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
+    monkeypatch.delenv("PYTHONPATH", raising=False)
+
+    cli.main(["run", "--", "gunicorn", "myapp.wsgi:app"])
+
+    autoload_dir = str(cli.Path(__file__).parent.parent / "wildedge" / "autoload")
+    assert captured["env"]["PYTHONPATH"] == autoload_dir
+
+
+def test_cli_run_preserves_existing_pythonpath(monkeypatch):
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
+    monkeypatch.setenv("PYTHONPATH", "/existing/path")
 
-    def fake_run(cmd, env, check):  # type: ignore[no-untyped-def]
-        captured["cmd"] = cmd
-        captured["env"] = env
-        captured["check"] = check
-        return SimpleNamespace(returncode=0)
+    cli.main(["run", "--", "gunicorn", "myapp.wsgi:app"])
 
-    monkeypatch.setattr(cli.subprocess, "run", fake_run)
+    pythonpath = captured["env"]["PYTHONPATH"]
+    assert pythonpath.endswith(os.pathsep + "/existing/path")
+
+
+def test_cli_run_sets_no_propagate_and_strict(monkeypatch):
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
 
     rc = cli.main(
         [
             "run",
             "--strict-integrations",
             "--no-propagate",
             "--",
-            "python",
-            "-m",
-            "pkg.main",
-            "--foo",
+            "gunicorn",
+            "myapp.wsgi:app",
         ]
     )
 
     assert rc == 0
-    assert captured["cmd"][1:8] == [
-        "-m",
-        "wildedge.runtime.runner",
-        "--mode",
-        "module",
-        "--target",
-        "pkg.main",
-        "--",
-    ]
     assert captured["env"][bootstrap.RUN_PROPAGATE_ENV] == "0"
     assert captured["env"][bootstrap.RUN_STRICT_INTEGRATIONS_ENV] == "1"
 
 
 def test_cli_run_sets_print_startup_report(monkeypatch):
-    captured = {}
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: f"/usr/bin/{cmd}")
 
-    def fake_run(cmd, env, check):  # type: ignore[no-untyped-def]
-        captured["env"] = env
-        return SimpleNamespace(returncode=0)
+    rc = cli.main(["run", "--print-startup-report", "--", "gunicorn", "myapp.wsgi:app"])
 
-    monkeypatch.setattr(cli.subprocess, "run", fake_run)
-    rc = cli.main(["run", "--print-startup-report", "--", "python", "app.py"])
     assert rc == 0
     assert captured["env"][bootstrap.RUN_PRINT_STARTUP_REPORT_ENV] == "1"
 
 
-def test_cli_rejects_non_python_command(capsys):
-    rc = cli.main(["run", "--", "bash", "script.sh"])
-    captured = capsys.readouterr()
-    assert rc == 2
-    assert "unsupported command format" in captured.err
+def test_cli_run_returns_127_for_missing_command(capsys, monkeypatch):
+    monkeypatch.setattr(cli.shutil, "which", lambda cmd: None)
+    monkeypatch.setattr(cli.Path, "is_file", lambda self: False)
+    rc = cli.main(["run", "--", "nonexistent-command"])
+    assert rc == 127
+    assert "command not found" in capsys.readouterr().err
+
+
+def test_cli_run_wraps_python_script_with_interpreter(monkeypatch, tmp_path):
+    script = tmp_path / "app.py"
+    script.write_text("pass")
+    captured: dict = {}
+    monkeypatch.setattr(cli.os, "execle", _fake_execle(captured))
+    monkeypatch.setattr(
+        cli.shutil,
+        "which",
+        lambda cmd: None if cmd.endswith(".py") else f"/usr/bin/{cmd}",
+    )
+
+    rc = cli.main(["run", "--", str(script)])
+
+    assert rc == 0
+    assert captured["argv"][1] == str(script)
+    assert "python" in captured["path"].lower()
 
 
 def test_install_runtime_requires_dsn(monkeypatch):
@@ -145,6 +164,9 @@ def flush(self, timeout):  # type: ignore[no-untyped-def]
         def close(self):  # type: ignore[no-untyped-def]
             pass
 
+        def _register_at_fork(self):  # type: ignore[no-untyped-def]
+            pass
+
     monkeypatch.setattr(bootstrap, "WildEdge", FakeWildEdge)
     monkeypatch.setenv(bootstrap.RUN_DSN_ENV, "https://secret@ingest.wildedge.dev/key")
     monkeypatch.delenv(bootstrap.RUN_FLUSH_TIMEOUT_ENV, raising=False)
@@ -177,6 +199,9 @@ def flush(self, timeout):  # type: ignore[no-untyped-def]
         def close(self):  # type: ignore[no-untyped-def]
             events.append(("close", ""))
 
+        def _register_at_fork(self):  # type: ignore[no-untyped-def]
+            pass
+
     monkeypatch.setattr(bootstrap, "WildEdge", FakeWildEdge)
     monkeypatch.setenv(bootstrap.RUN_DSN_ENV, "https://secret@ingest.wildedge.dev/key")
     monkeypatch.setenv(bootstrap.RUN_APP_VERSION_ENV, "2.0.0")
@@ -205,6 +230,9 @@ def __init__(self, *, dsn, app_version, debug):  # type: ignore[no-untyped-def]
         def instrument(self, name):  # type: ignore[no-untyped-def]
             raise RuntimeError("boom")
 
+        def _register_at_fork(self):  # type: ignore[no-untyped-def]
+            pass
+
     monkeypatch.setattr(bootstrap, "WildEdge", FakeWildEdge)
     monkeypatch.setenv(bootstrap.RUN_DSN_ENV, "https://secret@ingest.wildedge.dev/key")
     monkeypatch.setenv(bootstrap.RUN_INTEGRATIONS_ENV, "onnx")
@@ -361,6 +389,9 @@ def flush(self, timeout):  # type: ignore[no-untyped-def]
         def close(self):  # type: ignore[no-untyped-def]
             pass
 
+        def _register_at_fork(self):  # type: ignore[no-untyped-def]
+            pass
+
     monkeypatch.setattr(bootstrap, "WildEdge", FakeWildEdge)
     monkeypatch.setattr(
         bootstrap,

tests/test_consumer.py

@@ -335,3 +335,81 @@ def fake_register(fn, *args, **kwargs):
         consumer = self._make_consumer(queue, mock_transmitter)
         assert registered["fn"] == consumer.flush
         assert registered["args"] == (constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC,)
+
+
+class TestConsumerForkSafety:
+    def setup_method(self):
+        self._consumers: list[Consumer] = []
+
+    def teardown_method(self):
+        for c in self._consumers:
+            c.stop()
+
+    def _make_consumer(self, queue, transmitter, **kwargs) -> Consumer:
+        c = Consumer(
+            queue=queue,
+            transmitter=transmitter,
+            device=DeviceInfo(app_version="1.0", device_id="d", device_type="linux"),
+            get_models=lambda: {},
+            session_id="sess-fork",
+            **kwargs,
+        )
+        self._consumers.append(c)
+        return c
+
+    def test_before_fork_stops_thread(self, monkeypatch):
+        monkeypatch.setattr(Consumer, "run", lambda self: None)
+        queue = EventQueue(max_size=100)
+        mock_tx = MagicMock(spec=Transmitter)
+        consumer = self._make_consumer(queue, mock_tx)
+
+        original_thread = consumer.thread
+        consumer._before_fork()
+
+        assert consumer.stop_event.is_set()
+        assert not original_thread.is_alive()
+        # stopped is reset to False so _restart() can start a new thread
+        assert consumer.stopped is False
+
+    def test_restart_creates_fresh_thread(self, monkeypatch):
+        monkeypatch.setattr(Consumer, "run", lambda self: None)
+        queue = EventQueue(max_size=100)
+        mock_tx = MagicMock(spec=Transmitter)
+        consumer = self._make_consumer(queue, mock_tx)
+
+        original_thread = consumer.thread
+        consumer._before_fork()
+        consumer._restart()
+
+        assert consumer.thread is not original_thread
+        # thread was started (ident is set) even though the no-op run() exits fast
+        assert consumer.thread.ident is not None
+        assert not consumer.stop_event.is_set()
+
+    def test_restart_resets_backoff_and_snapshot(self, monkeypatch):
+        monkeypatch.setattr(Consumer, "run", lambda self: None)
+        queue = EventQueue(max_size=100)
+        mock_tx = MagicMock(spec=Transmitter)
+        consumer = self._make_consumer(queue, mock_tx)
+
+        consumer.backoff = constants.BACKOFF_MAX
+        consumer._held_snapshot = ([], {})
+
+        consumer._before_fork()
+        consumer._restart()
+
+        assert consumer.backoff == constants.BACKOFF_MIN
+        assert consumer._held_snapshot is None
+
+    def test_flush_is_noop_after_before_fork_before_restart(self, monkeypatch):
+        """flush() on a pre-fork-stopped consumer with empty queue returns immediately."""
+        monkeypatch.setattr(Consumer, "run", lambda self: None)
+        queue = EventQueue(max_size=100)
+        mock_tx = MagicMock(spec=Transmitter)
+        consumer = self._make_consumer(queue, mock_tx)
+
+        consumer._before_fork()
+        # stopped is False (reset by _before_fork) and queue is empty, so flush
+        # calls drain_once which returns False immediately — no transmit calls.
+        consumer.flush(timeout=0.1)
+        mock_tx.send.assert_not_called()

wildedge/autoload/__init__.py

@@ -0,0 +1,67 @@
+"""WildEdge autoload sitecustomize.
+
+Loaded automatically by the Python interpreter when wildedge/autoload/ is
+prepended to PYTHONPATH (by `wildedge run`). Calls install_runtime() before
+any user code runs, enabling framework instrumentation and fork-safe operation
+for gunicorn, celery, and other pre-fork servers.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+_GUARD = "WILDEDGE_AUTOLOAD_ACTIVE"
+_RUN_DSN = "WILDEDGE_RUN_DSN"
+_DSN = "WILDEDGE_DSN"
+
+
+def _bootstrap() -> None:
+    # Idempotent: skip if already initialized (e.g. in a forked worker that
+    # exec'd a subprocess, or if sitecustomize.py is imported twice).
+    if os.environ.get(_GUARD):
+        return
+
+    # Require a DSN. Silently skip if missing so that processes which
+    # inherit PYTHONPATH without wildedge config don't crash.
+    if not (os.environ.get(_RUN_DSN) or os.environ.get(_DSN)):
+        return
+
+    # Set the guard before importing wildedge to prevent re-entry.
+    os.environ[_GUARD] = "1"
+
+    try:
+        from wildedge.runtime.bootstrap import install_runtime  # noqa: PLC0415
+
+        # Don't install signal handlers: the host process (gunicorn, celery,
+        # etc.) manages SIGTERM/SIGINT itself.
+        install_runtime(install_signal_handlers=False)
+    except Exception as exc:  # pragma: no cover
+        print(f"wildedge: bootstrap failed: {exc}", file=sys.stderr)
+
+
+_bootstrap()
+
+
+# Chain any pre-existing sitecustomize that would otherwise be shadowed.
+# Use importlib to find and exec it directly — avoids sys.modules manipulation
+# which can trigger CPython's module GC and clear globals mid-execution.
+def _chain_sitecustomize() -> None:
+    import importlib.util as _iutil
+
+    _autoload_dir = os.path.dirname(os.path.abspath(__file__))
+    _saved = sys.path[:]
+    sys.path = [p for p in sys.path if p != _autoload_dir]
+    try:
+        _spec = _iutil.find_spec("sitecustomize")
+        if _spec is not None and _spec.origin is not None:
+            _mod = _iutil.module_from_spec(_spec)
+            sys.modules["sitecustomize"] = _mod
+            _spec.loader.exec_module(_mod)  # type: ignore[union-attr]
+    except Exception:
+        pass
+    finally:
+        sys.path = _saved
+
+
+_chain_sitecustomize()