diff --git a/README.md b/README.md index 4921dfc..aa2f94e 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ commit to git (they're gitignored and volume-mounted from the host): - **`aragorn_omnicorp`** → `./omnicorp_lmdb/` (`curies.lmdb`, `shared_counts.lmdb`) - **`score_paths`** → `./pathfinder_embeddings/` (a directory-style LMDB) -- **`arax_pathfinder`** → `./arax_pathfinder_dbs/` (`curie_ngd_v1.0_.sqlite`, `tier0-info-for-overlay_v1.0_.sqlite`) +- **`arax_pathfinder`** → `./arax_pathfinder_dbs/` (`curie_ngd_v1.0_.sqlite`, `tier0-info-for-overlay_v1.0_.sqlite`, `general_concepts.json`) So a new developer doesn't have to source these by hand, each worker can fetch its dataset on first startup. Two download mechanisms are supported, depending on where the dataset lives: @@ -37,7 +37,7 @@ PATHFINDER_EMBEDDINGS_URL=https://example.org/path/pathfinder_embeddings.tar.gz The archive for each dataset should contain the expected files at its top level: `curies.lmdb` and `shared_counts.lmdb` for omnicorp, `data.mdb` (and `lock.mdb`) for the embeddings. -**arax_pathfinder's sqlite databases** are served as plain files over HTTPS. The +**arax_pathfinder's sqlite databases** are served as plain files over HTTPS, no credentials needed. The filenames embed a Knowledge Graph version that changes periodically, so only one variable needs updating when a new Knowledge Graph ships: @@ -45,14 +45,43 @@ when a new Knowledge Graph ships: ARAX_PATHFINDER_TIER_VERSION=tier0-20260621 ``` -This requires an SSH key with access to that host, mounted read-only into the container -(`~/.ssh:/home/nru/.ssh:ro` in compose.yml). +The ARAX blocked-concept list (`general_concepts.json`, fetched from GitHub) lands in this same +directory, so it is downloaded once and then persists with the databases rather than being re-fetched +by every new container. It is only fetched when absent — delete it from the volume to pick up an +updated upstream list. On startup, each worker checks whether its files already exist in the volume-mounted directory. If -they're missing and a source is configured (URL or scp path), it fetches them into that directory — -which lives on the host, so the data persists across restarts and is only downloaded once. If the -files are already present, or no source is configured, the download is skipped (production mounts -this data out of band, so it's unaffected). +they're missing and a URL is configured, it fetches them into that directory — which lives on the +host, so the data persists across restarts and is only downloaded once. If the files are already +present, or no source is configured, the download is skipped (production mounts this data out of +band, so it's unaffected). + +Downloads are bounded by `DATASET_DOWNLOAD_TIMEOUT_SEC` (default 60), which applies per socket +operation rather than to the whole transfer — a large file downloads for as long as it needs, but a +connection that opens and then stalls fails loudly instead of hanging worker startup. + +#### Deploying these workers + +The presence check is an exact match on the configured directory **and** the tier-versioned filenames. +A deployment that mounts the data somewhere else, or whose `ARAX_PATHFINDER_TIER_VERSION` doesn't match +the filenames on the volume, will not use the mounted copies — it will decide the dataset is missing and +download it again, into whatever path the settings do point at. If that path isn't the mount, the files +land on the container's writable layer and the pod is eventually evicted for exceeding its ephemeral +storage. So when deploying, set the directory explicitly and confirm it resolves to your mount: + +```dotenv +ARAX_PATHFINDER_DBS_DIR=/data/arax_pathfinder_dbs # default is relative: resolved against /app +``` + +```console +$ kubectl exec deploy/arax-pathfinder -- python -c \ + "from shepherd_utils.data_download import arax_pathfinder_sqlite_paths as p; print(*p(), sep='\n')" +``` + +Because `general_concepts.json` now shares that directory, the volume needs to be writable for the +first startup that fetches it — or the file can be preloaded alongside the sqlite databases, after +which the worker only ever reads it. A read-only mount with no preloaded copy fails at startup with a +permission error rather than silently continuing. ### Worker diff --git a/compose.yml b/compose.yml index 0e1a39a..bd6d77d 100644 --- a/compose.yml +++ b/compose.yml @@ -359,9 +359,14 @@ services: - ./logs:/app/logs - ./.env:/app/.env # First run? The worker downloads its two sqlite dbs from - # kg2webhost.rtx.ai on startup, no credentials needed. Set + # kg2webhost.rtx.ai, plus the ARAX blocked-concept list from GitHub, on + # startup -- no credentials needed. All three live here, so they persist + # across restarts instead of being re-fetched by each new container. Set # ARAX_PATHFINDER_TIER_VERSION in your .env if you need a tier other than # the default (see README "Worker data (LMDB / sqlite) downloads"). + # Deploying elsewhere? ARAX_PATHFINDER_DBS_DIR must point at the mount -- + # it defaults to a path relative to /app, and a mismatch silently + # re-downloads everything instead of using the mounted copies. - ./arax_pathfinder_dbs:/app/arax_pathfinder_dbs arax_rank: diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 7f7e65c..78e9907 100644 --- a/shepherd_utils/config.py +++ b/shepherd_utils/config.py @@ -122,6 +122,19 @@ class Settings(BaseSettings): omnicorp_lmdb_url: str = "" pathfinder_embeddings_url: str = "" + # Socket timeout (seconds) for those startup dataset downloads. This bounds + # each individual blocking socket operation -- the connect and every + # subsequent read -- not the transfer as a whole, so a genuinely large file + # still downloads for as long as it needs provided it keeps making progress. + # Without it, urllib inherits Python's default of no timeout: a connection + # that opens and then stalls (an egress proxy that swallows the request, a + # server that accepts and never responds) hangs the worker's startup + # forever, before it ever reaches the poll loop. That is silent -- the + # worker never registers a heartbeat and never picks up a task, but it also + # never crashes. With a timeout the download fails loudly instead. Set to 0 + # to restore the unbounded behavior. + dataset_download_timeout_sec: float = 60.0 + otel_enabled: bool = True jaeger_host: str = "http://jaeger" jaeger_port: int = 4317 diff --git a/shepherd_utils/data_download.py b/shepherd_utils/data_download.py index 610a948..e10998f 100644 --- a/shepherd_utils/data_download.py +++ b/shepherd_utils/data_download.py @@ -15,7 +15,8 @@ ``aragorn_omnicorp`` and ``score_paths`` below). * **Per-file** -- individual files fetched directly, no archive/extract step (``arax_pathfinder`` below, whose two sqlite databases are served as plain - files rather than bundled into one archive). + files rather than bundled into one archive, plus the ARAX blocked-concept + list that shares their directory). When a download source is configured (read via :mod:`shepherd_utils.config`), each worker calls its matching ``ensure_*`` helper at startup: @@ -50,14 +51,25 @@ def _missing_files(target_dir: str, required_files: List[str]) -> List[str]: def _download(url: str, dest_path: str, logger: logging.Logger) -> None: - """Stream ``url`` to ``dest_path``, logging progress periodically.""" + """Stream ``url`` to ``dest_path``, logging progress periodically. + + Bounded by ``dataset_download_timeout_sec`` (see config). The timeout is + per socket operation rather than for the whole transfer, so a multi-GB file + downloads for as long as it needs while a stalled connection raises instead + of hanging worker startup indefinitely. + """ logger.info(f"Downloading dataset from {url} ...") + timeout = float(settings.dataset_download_timeout_sec) + # urlopen treats timeout=None as "no timeout"; 0 is the documented opt-out. + kwargs = {"timeout": timeout} if timeout > 0 else {} + # Hoisted so the timeout handler below can report how far the transfer got + # even when it never made it past opening the connection. + read = 0 try: # nosec B310: the URL is operator-configured (an env var), not user input. - with urllib.request.urlopen(url) as resp: # noqa: S310 + with urllib.request.urlopen(url, **kwargs) as resp: # noqa: S310 header = resp.headers.get("Content-Length") total = int(header) if header and header.isdigit() else None - read = 0 step = 50 * 1024 * 1024 # log roughly every 50 MB next_log = step with open(dest_path, "wb") as out: @@ -82,10 +94,22 @@ def _download(url: str, dest_path: str, logger: logging.Logger) -> None: f"URL is correct and reachable." ) from e except urllib.error.URLError as e: + # A connect timeout arrives here, wrapped, with e.reason set to the + # underlying socket.timeout. raise RuntimeError( f"download failed for {url}: {e.reason}. Confirm the URL is correct " f"and reachable from inside the container." ) from e + except TimeoutError as e: + # A stall part-way through the body raises straight out of resp.read() + # rather than being wrapped in URLError, so it needs its own handler -- + # otherwise this surfaces as a bare TimeoutError with no context. + raise RuntimeError( + f"download failed for {url}: no data for {timeout:.0f}s " + f"({read / 1e6:.0f} MB transferred). The host is reachable but the " + f"transfer stalled -- check for an egress proxy or a rate limit, or " + f"raise DATASET_DOWNLOAD_TIMEOUT_SEC." + ) from e logger.info(f"Download complete: {read / 1e6:.0f} MB") @@ -123,7 +147,7 @@ def ensure_lmdb_dataset( Idempotent: once the files are present this returns immediately, so it's safe to call unconditionally on every worker startup. """ - logger = logger or logging.getLogger(__name__) + logger = logger or logging.getLogger("shepherd.data_download") missing = _missing_files(target_dir, required_files) if not missing: @@ -191,7 +215,7 @@ def ensure_http_files_dataset( Idempotent: once a file is present it's left alone, so it's safe to call unconditionally on every worker startup. """ - logger = logger or logging.getLogger(__name__) + logger = logger or logging.getLogger("shepherd.data_download") required_files = list(file_sources.keys()) missing = _missing_files(target_dir, required_files) @@ -344,3 +368,44 @@ def ensure_arax_pathfinder_dbs(logger: Optional[logging.Logger] = None) -> None: }, logger=logger, ) + + +# The ARAX blocked-concept list lives alongside the pathfinder sqlite databases +# so it lands on the same mounted volume. It previously went to the worker's +# working directory (``/app``), which is the container's writable layer and +# therefore thrown away on every restart -- meaning each new pod re-fetched it +# from GitHub during startup, before the poll loop, on a code path whose logs +# were being discarded. On the volume it is fetched once and persists. +ARAX_BLOCKED_LIST_FILENAME = "general_concepts.json" + + +def arax_blocked_list_path() -> str: + """Return the on-disk path of the ARAX blocked-concept list. + + Single source of truth, in the same spirit as + ``arax_pathfinder_sqlite_paths``: ``ensure_arax_blocked_list`` uses it to + know where to download, and worker.py uses it to know what to open. + """ + return os.path.join(settings.arax_pathfinder_dbs_dir, ARAX_BLOCKED_LIST_FILENAME) + + +def ensure_arax_blocked_list(logger: Optional[logging.Logger] = None) -> None: + """Ensure the ARAX blocked-concept list is present next to the sqlite dbs. + + Fetched via the shared downloader so it lands through a temp file + atomic + rename (a direct write let concurrent tasks race on a half-written file) + and inherits the download timeout. Idempotent, so it is safe to call at + startup and again lazily from a pool child. + + Note the flip side of persisting this on the volume: it is now only fetched + when absent, so a refreshed upstream list is not picked up until the file is + deleted. Delete it from the volume to force a re-fetch on the next restart. + """ + ensure_http_files_dataset( + name="arax_blocked_list", + # ``or "."`` so an unset/blank dbs dir degrades to the working directory + # rather than handing makedirs an empty path. + target_dir=os.path.dirname(arax_blocked_list_path()) or ".", + file_sources={ARAX_BLOCKED_LIST_FILENAME: settings.arax_blocked_list_url}, + logger=logger, + ) diff --git a/shepherd_utils/db.py b/shepherd_utils/db.py index 0380903..dfdff43 100644 --- a/shepherd_utils/db.py +++ b/shepherd_utils/db.py @@ -193,7 +193,7 @@ async def initialize_db() -> None: except Exception: # A failed upgrade must never keep a worker from starting: the schema # additions are performance aids, and the janitor/next boot retries. - logging.getLogger(__name__).warning( + logging.getLogger("shepherd.db").warning( "Failed to apply startup schema upgrades", exc_info=True ) diff --git a/shepherd_utils/logger.py b/shepherd_utils/logger.py index 791de4c..e43957b 100644 --- a/shepherd_utils/logger.py +++ b/shepherd_utils/logger.py @@ -128,10 +128,32 @@ def get_logging_config(): "default": {"format": "[%(asctime)s: %(levelname)s/%(name)s]: %(message)s"} }, "handlers": handlers, + # The output handlers live on root, and root alone. Previously they were + # attached to ``shepherd`` and root was left unconfigured, so any record + # logged outside that namespace -- a stray ``logging.info``, a + # third-party library, a logger named for its module or its Redis stream + # -- reached a handler-less root and was dropped by logging's + # ``lastResort`` fallback, which only emits WARNING+ and ignores our + # formatter. That silently swallowed every worker's entire startup + # phase; see ``get_worker_logger``. + # + # Root's level applies only to records logged directly on root, not to + # ones propagated up from a child (those are level-checked at the + # originating logger). So WARNING here means third-party libraries are + # quiet below WARNING -- httpx logs a line per request at INFO, which + # would bury our own output -- while ``shepherd.*`` still emits at DEBUG + # via the entry below. Everything then reaches these handlers by + # propagation, which also keeps pytest's ``caplog`` working. + "root": { + "level": "WARNING", + "handlers": logger_handlers, + }, "loggers": { + # No handlers: records propagate to root's. Only the level is set + # here, which is what lets our own logging through at DEBUG while + # leaving third-party loggers at root's WARNING. "shepherd": { "level": "DEBUG", - "handlers": logger_handlers, }, # psycopg's pool retries to keep min_size connections warm and logs # a WARNING on every failed attempt. When the DB is down that floods @@ -150,6 +172,25 @@ def get_logging_config(): return logging_config +def get_worker_logger(name: str) -> logging.Logger: + """Return a logger under the configured ``shepherd`` namespace. + + ``setup_logging`` only attaches handlers to ``shepherd`` (and, as a + WARNING-level backstop, root), so a logger named for the stream alone + (``logging.getLogger("arax.pathfinder")``) inherits no handler at INFO and + its records vanish. Every worker's startup phase -- dataset downloads, pool + sizing, poll-loop errors -- logged through exactly such a logger, so a + worker that hung before reaching ``get_tasks`` produced no output at all + and looked identical to a healthy idle one. + + Passing a name that is already namespaced is a no-op, so this is safe to + apply to existing ``shepherd.``-prefixed names. + """ + if name == "shepherd" or name.startswith("shepherd."): + return logging.getLogger(name) + return logging.getLogger(f"shepherd.{name}") + + def setup_logging(): """Set up logging.""" config = get_logging_config() diff --git a/tests/unit/test_data_download.py b/tests/unit/test_data_download.py index 851cdae..1b32170 100644 --- a/tests/unit/test_data_download.py +++ b/tests/unit/test_data_download.py @@ -5,7 +5,9 @@ """ import logging +import os import tarfile +import urllib.request # noqa: F401 (patched by name in the timeout tests) import pytest @@ -155,3 +157,91 @@ def test_ensure_pathfinder_embeddings_wires_settings(tmp_path, mocker): assert (target / "data.mdb").exists() assert (target / "lock.mdb").exists() + + +# --- download timeout -------------------------------------------------------- +# +# urllib defaults to no timeout, so a connection that opens and then stalls hung +# worker startup forever -- before the poll loop, before any heartbeat, and with +# no crash to point at. + + +def test_download_passes_configured_timeout(tmp_path, mocker): + mocker.patch.object(data_download.settings, "dataset_download_timeout_sec", 12.5) + urlopen = mocker.patch.object(data_download.urllib.request, "urlopen") + urlopen.return_value.__enter__.return_value.headers.get.return_value = None + urlopen.return_value.__enter__.return_value.read.return_value = b"" + + data_download._download("http://example.invalid/x", str(tmp_path / "out"), logger) + + assert urlopen.call_args.kwargs["timeout"] == 12.5 + + +def test_download_timeout_zero_restores_unbounded_behavior(tmp_path, mocker): + mocker.patch.object(data_download.settings, "dataset_download_timeout_sec", 0) + urlopen = mocker.patch.object(data_download.urllib.request, "urlopen") + urlopen.return_value.__enter__.return_value.headers.get.return_value = None + urlopen.return_value.__enter__.return_value.read.return_value = b"" + + data_download._download("http://example.invalid/x", str(tmp_path / "out"), logger) + + assert "timeout" not in urlopen.call_args.kwargs + + +def test_download_reports_a_mid_transfer_stall(tmp_path, mocker): + """A stall inside resp.read() raises TimeoutError directly rather than being + wrapped in URLError, so it needs its own handler to get a useful message.""" + mocker.patch.object(data_download.settings, "dataset_download_timeout_sec", 30) + urlopen = mocker.patch.object(data_download.urllib.request, "urlopen") + resp = urlopen.return_value.__enter__.return_value + resp.headers.get.return_value = None + resp.read.side_effect = [b"partial", TimeoutError("timed out")] + + with pytest.raises(RuntimeError, match="transfer stalled"): + data_download._download( + "http://example.invalid/x", str(tmp_path / "out"), logger + ) + + +# --- ARAX blocked list ------------------------------------------------------- + + +def test_blocked_list_lives_with_the_pathfinder_sqlite_dbs(tmp_path, mocker): + """It used to be written to the working directory -- the container's + writable layer -- so every new pod re-fetched it from GitHub at startup.""" + mocker.patch.object( + data_download.settings, "arax_pathfinder_dbs_dir", str(tmp_path / "dbs") + ) + + path = data_download.arax_blocked_list_path() + + assert path == str(tmp_path / "dbs" / "general_concepts.json") + curie_ngd, _ = data_download.arax_pathfinder_sqlite_paths() + # Same volume as the sqlite databases, which is the whole point. + assert os.path.dirname(path) == os.path.dirname(curie_ngd) + + +def test_ensure_arax_blocked_list_downloads_into_the_volume(tmp_path, mocker): + target = tmp_path / "dbs" + source = tmp_path / "general_concepts.json" + source.write_text('{"curies": ["CHEBI:1"], "synonyms": ["Water"]}') + mocker.patch.object(data_download.settings, "arax_pathfinder_dbs_dir", str(target)) + mocker.patch.object( + data_download.settings, "arax_blocked_list_url", source.as_uri() + ) + + data_download.ensure_arax_blocked_list(logger) + + assert (target / "general_concepts.json").exists() + + +def test_ensure_arax_blocked_list_is_a_noop_once_present(tmp_path, mocker): + target = tmp_path / "dbs" + target.mkdir() + (target / "general_concepts.json").write_text('{"curies": [], "synonyms": []}') + mocker.patch.object(data_download.settings, "arax_pathfinder_dbs_dir", str(target)) + spy = mocker.patch.object(data_download, "_download") + + data_download.ensure_arax_blocked_list(logger) + + spy.assert_not_called() diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py index c0d37c2..df97bfd 100644 --- a/tests/unit/test_logger.py +++ b/tests/unit/test_logger.py @@ -7,6 +7,7 @@ QueryLogger, ReasonerLogEntryFormatter, get_logging_config, + get_worker_logger, ) @@ -91,7 +92,7 @@ def test_get_logging_config_local_includes_file_handler(monkeypatch, tmp_path): config = get_logging_config() assert "file" in config["handlers"] assert "console" in config["handlers"] - assert set(config["loggers"]["shepherd"]["handlers"]) == {"console", "file"} + assert set(config["root"]["handlers"]) == {"console", "file"} # The function eagerly creates the logs/ dir for file output. assert os.path.isdir(tmp_path / "logs") @@ -100,7 +101,7 @@ def test_get_logging_config_kubernetes_skips_file_handler(monkeypatch): monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") config = get_logging_config() assert "file" not in config["handlers"] - assert config["loggers"]["shepherd"]["handlers"] == ["console"] + assert config["root"]["handlers"] == ["console"] def test_get_logging_config_pool_child_skips_file_handler(monkeypatch, tmp_path): @@ -119,4 +120,64 @@ def test_get_logging_config_pool_child_skips_file_handler(monkeypatch, tmp_path) ) config = get_logging_config() assert "file" not in config["handlers"] - assert config["loggers"]["shepherd"]["handlers"] == ["console"] + assert config["root"]["handlers"] == ["console"] + + +# --- worker logger namespacing ---------------------------------------------- +# +# Regression cover for a silent-startup bug: workers logged through +# ``logging.getLogger(STREAM)`` (e.g. "arax.pathfinder"), which sits outside the +# only namespace ``setup_logging`` attaches handlers to. Those records reached a +# handler-less root and were dropped by logging's lastResort fallback, so a +# worker that hung during startup produced no output whatsoever. + + +def test_get_worker_logger_namespaces_stream_names(): + assert get_worker_logger("arax.pathfinder").name == "shepherd.arax.pathfinder" + + +def test_get_worker_logger_does_not_double_prefix(): + """Safe to apply to names that are already namespaced.""" + assert get_worker_logger("shepherd.monitor").name == "shepherd.monitor" + assert get_worker_logger("shepherd").name == "shepherd" + + +def test_worker_logger_inherits_configured_handlers(monkeypatch): + """The whole point: a worker logger must resolve to a real handler at INFO.""" + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") + logging.config.dictConfig(get_logging_config()) + + worker_logger = get_worker_logger("arax.pathfinder") + # Walk the ancestry the way logging does when emitting a record. + effective = [] + node = worker_logger + while node: + effective.extend(node.handlers) + node = node.parent if node.propagate else None + + assert effective, "worker logger resolved to no handler" + assert worker_logger.getEffectiveLevel() <= logging.INFO + + +def test_root_logger_configured_as_warning_backstop(monkeypatch): + """Stray records outside ``shepherd.*`` still land somewhere, formatted. + + WARNING rather than INFO on purpose -- at INFO, libraries such as httpx emit + a line per request and drown out our own logs. + """ + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") + config = get_logging_config() + assert config["root"]["level"] == "WARNING" + assert config["root"]["handlers"] == ["console"] + + +def test_handlers_are_attached_only_once(monkeypatch): + """Handlers live on root only; shepherd sets a level and propagates to them. + + Attaching them in both places would emit every shepherd record twice. + """ + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") + config = get_logging_config() + assert config["root"]["handlers"] == ["console"] + assert "handlers" not in config["loggers"]["shepherd"] + assert config["loggers"]["shepherd"]["level"] == "DEBUG" diff --git a/workers/aragorn/worker.py b/workers/aragorn/worker.py index 4bdde0d..e58d8c7 100644 --- a/workers/aragorn/worker.py +++ b/workers/aragorn/worker.py @@ -6,6 +6,7 @@ import uuid from shepherd_utils.db import get_message +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import ( examine_query, @@ -20,6 +21,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def aragorn(task, logger: logging.Logger): @@ -76,9 +78,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/aragorn_lookup/worker.py b/workers/aragorn_lookup/worker.py index e529717..50beea9 100644 --- a/workers/aragorn_lookup/worker.py +++ b/workers/aragorn_lookup/worker.py @@ -23,6 +23,7 @@ remove_callback_id, save_message, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -33,6 +34,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def examine_query(message): @@ -367,9 +369,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/aragorn_omnicorp/worker.py b/workers/aragorn_omnicorp/worker.py index f72dd23..cff432b 100644 --- a/workers/aragorn_omnicorp/worker.py +++ b/workers/aragorn_omnicorp/worker.py @@ -29,6 +29,7 @@ from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.data_download import ensure_omnicorp_lmdb from shepherd_utils.db import get_message_sync, save_message_sync +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -43,6 +44,7 @@ # many run in parallel; keep it modest since each message can be large. TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) # Matches the upstream `redis_batch_size`. LMDB_BATCH_SIZE = 1000 @@ -560,7 +562,7 @@ async def poll_for_tasks(): # to open them lazily (a first-run local `docker compose up` starts with the # volume-mounted directory empty). No-op once present or when no download URL # is configured (e.g. production, where the data is mounted out of band). - ensure_omnicorp_lmdb(logging.getLogger(STREAM)) + ensure_omnicorp_lmdb(LOGGER) loop = asyncio.get_running_loop() # The overlay is CPU-bound, so cap real parallelism at the number of cores. @@ -571,8 +573,8 @@ async def poll_for_tasks(): # messages at once and OOM-killing the pod). resolve_pool_workers reads the # cgroup CPU limit and honours a POOL_MAX_WORKERS override for memory-tight # deployments. Extra in-flight tasks queue against the pool without blocking. - max_workers = resolve_pool_workers(TASK_LIMIT, logging.getLogger(STREAM)) - logging.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") pool = ProcessPoolManager( max_workers, max_tasks_per_child=settings.pool_max_tasks_per_child, @@ -588,9 +590,9 @@ async def poll_for_tasks(): process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/aragorn_pathfinder/worker.py b/workers/aragorn_pathfinder/worker.py index 9cc0f3b..dc003d5 100644 --- a/workers/aragorn_pathfinder/worker.py +++ b/workers/aragorn_pathfinder/worker.py @@ -16,6 +16,7 @@ get_running_callbacks, save_message, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import ( get_tasks, @@ -29,6 +30,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def shadowfax(task, logger: logging.Logger) -> str: @@ -278,9 +280,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/aragorn_score/worker.py b/workers/aragorn_score/worker.py index c4f7c77..d5db81b 100644 --- a/workers/aragorn_score/worker.py +++ b/workers/aragorn_score/worker.py @@ -14,6 +14,7 @@ from shepherd_utils.config import settings from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.db import get_message_sync, save_message_sync +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -25,6 +26,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) DEFAULT_WEIGHT = 1e-2 @@ -1255,8 +1257,8 @@ async def poll_for_tasks(): # Size the pool by the pod's actual CPU allocation (cgroup limit), not # os.cpu_count() -- see aragorn_omnicorp.poll_for_tasks. Each child loads a # full message, so this also bounds peak memory. POOL_MAX_WORKERS overrides. - max_workers = resolve_pool_workers(TASK_LIMIT, logging.getLogger(STREAM)) - logging.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") pool = ProcessPoolManager( max_workers, max_tasks_per_child=settings.pool_max_tasks_per_child, @@ -1272,9 +1274,9 @@ async def poll_for_tasks(): process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/arax/worker.py b/workers/arax/worker.py index 84ab9e4..34a9b46 100644 --- a/workers/arax/worker.py +++ b/workers/arax/worker.py @@ -11,6 +11,7 @@ from shepherd_utils.config import settings from shepherd_utils.db import get_message, save_message +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -20,6 +21,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def is_pathfinder_query(message): @@ -82,9 +84,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/arax_pathfinder/worker.py b/workers/arax_pathfinder/worker.py index e1e8e9f..bbe2d0b 100644 --- a/workers/arax_pathfinder/worker.py +++ b/workers/arax_pathfinder/worker.py @@ -14,9 +14,10 @@ from shepherd_utils.config import settings from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.data_download import ( + arax_blocked_list_path, arax_pathfinder_sqlite_paths, + ensure_arax_blocked_list, ensure_arax_pathfinder_dbs, - ensure_http_files_dataset, ) from shepherd_utils.db import ( get_message_sync, @@ -25,6 +26,7 @@ from shepherd_utils.inject_shepherd_arax_provenance import ( add_shepherd_arax_to_edge_sources, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -36,6 +38,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) NUM_TOTAL_HOPS = 4 MAX_HOPS_TO_EXPLORE = 4 @@ -43,13 +46,6 @@ PRUNE_TOP_K = 75 NODE_DEGREE_THRESHOLD = 10000 -# The ARAX blocked-concept list, fetched once at worker startup (see -# poll_for_tasks) and read by each pool child. Kept in the working directory -# (/app in the image) so it needs no extra volume mount. -BLOCKED_LIST_DIR = "." -BLOCKED_LIST_FILENAME = "general_concepts.json" -BLOCKED_LIST_PATH = Path(BLOCKED_LIST_DIR) / BLOCKED_LIST_FILENAME - BIOLINK_CACHE_DIR = "/tmp/biolink" REHYDRATE_TIMEOUT_SEC = 30.0 @@ -65,30 +61,16 @@ _descendants_cache: dict = {} -def ensure_blocked_list(logger: logging.Logger) -> None: - """Fetch the ARAX blocked-concept list if it isn't on disk yet. - - Uses the shared downloader so the file lands via a temp file + atomic - rename; the previous per-task ``requests.get`` wrote the destination - directly, so concurrent tasks could race on a half-written file. - Idempotent, so it is safe to call at startup and again lazily in a child. - """ - ensure_http_files_dataset( - name="arax_blocked_list", - target_dir=BLOCKED_LIST_DIR, - file_sources={BLOCKED_LIST_FILENAME: settings.arax_blocked_list_url}, - logger=logger, - ) - - def get_blocked_list(logger: logging.Logger): """``(blocked_curies, blocked_synonyms)``, parsed once per pool child.""" global _blocked_list_cache if _blocked_list_cache is None: - if not BLOCKED_LIST_PATH.exists(): - # Startup fetch failed or this child outlived a wiped working dir. - ensure_blocked_list(logger) - with open(BLOCKED_LIST_PATH, "r") as file: + blocked_list_path = Path(arax_blocked_list_path()) + if not blocked_list_path.exists(): + # Startup fetch failed, or the volume was replaced under a + # long-lived child. + ensure_arax_blocked_list(logger) + with open(blocked_list_path, "r") as file: json_block_list = json.load(file) synonyms = set(s.lower() for s in json_block_list["synonyms"]) _blocked_list_cache = (set(json_block_list["curies"]), synonyms) @@ -326,21 +308,25 @@ async def _run(task, logger): async def poll_for_tasks(): """On initialization, poll indefinitely for available tasks.""" - startup_logger = logging.getLogger(STREAM) # Ensure the two sqlite databases exist before any task tries to open them # (a first-run local `docker compose up` starts with the volume-mounted - # directory empty). No-op once present or when no scp source is configured - # (e.g. production, where the data is mounted out of band). - ensure_arax_pathfinder_dbs(startup_logger) + # directory empty). No-op once present -- which is the case in production, + # where the data is mounted out of band. Note the "already present" check is + # an exact match on ARAX_PATHFINDER_DBS_DIR plus the tier-versioned + # filenames, so a deployment whose mount path or ARAX_PATHFINDER_TIER_VERSION + # disagrees with those settings downloads the databases again rather than + # using the mounted copies. + ensure_arax_pathfinder_dbs(LOGGER) # Fetch the blocked-concept list once here rather than per task, so the pool - # children only ever read it. - ensure_blocked_list(startup_logger) + # children only ever read it. It lives on the same volume as the sqlite dbs, + # so this is a no-op after the first run. + ensure_arax_blocked_list(LOGGER) loop = asyncio.get_running_loop() # Size the pool by the pod's actual CPU allocation (cgroup limit), not # os.cpu_count() -- see aragorn_score.poll_for_tasks. Each child loads a full # message, so this also bounds peak memory. POOL_MAX_WORKERS overrides. - max_workers = resolve_pool_workers(TASK_LIMIT, startup_logger) - logging.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") pool = ProcessPoolManager( max_workers, max_tasks_per_child=settings.pool_max_tasks_per_child, @@ -356,9 +342,9 @@ async def poll_for_tasks(): process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/arax_rank/worker.py b/workers/arax_rank/worker.py index af9cb06..197dc12 100644 --- a/workers/arax_rank/worker.py +++ b/workers/arax_rank/worker.py @@ -18,6 +18,7 @@ from shepherd_utils.config import settings from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.db import get_message_sync, save_message_sync +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -31,6 +32,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def rank_message(in_message: dict, logger: logging.Logger) -> dict: @@ -124,8 +126,8 @@ async def poll_for_tasks() -> None: # Size the pool by the pod's actual CPU allocation (cgroup limit), not # os.cpu_count() -- see aragorn_omnicorp.poll_for_tasks. Each child loads a # full message, so this also bounds peak memory. POOL_MAX_WORKERS overrides. - max_workers = resolve_pool_workers(TASK_LIMIT, logging.getLogger(STREAM)) - logging.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") pool = ProcessPoolManager( max_workers, max_tasks_per_child=settings.pool_max_tasks_per_child, @@ -142,9 +144,9 @@ async def poll_for_tasks() -> None: process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/bte/worker.py b/workers/bte/worker.py index c882263..f1aee0d 100644 --- a/workers/bte/worker.py +++ b/workers/bte/worker.py @@ -6,6 +6,7 @@ import uuid from shepherd_utils.db import get_message +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import ( examine_query, @@ -20,6 +21,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def bte(task, logger: logging.Logger): @@ -69,9 +71,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/bte_lookup/worker.py b/workers/bte_lookup/worker.py index 21f7391..5e3c1e4 100644 --- a/workers/bte_lookup/worker.py +++ b/workers/bte_lookup/worker.py @@ -23,6 +23,7 @@ remove_callback_id, save_message, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -33,6 +34,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def examine_query(message): @@ -443,9 +445,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/example_ara/worker.py b/workers/example_ara/worker.py index 12f0449..db8069d 100644 --- a/workers/example_ara/worker.py +++ b/workers/example_ara/worker.py @@ -7,6 +7,7 @@ import uuid from shepherd_utils.db import get_message from shepherd_utils.shared import get_tasks, run_task_lifecycle +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -15,6 +16,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def example_ara(task, logger: logging.Logger): @@ -50,9 +52,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/example_lookup/worker.py b/workers/example_lookup/worker.py index d95b458..0865d03 100644 --- a/workers/example_lookup/worker.py +++ b/workers/example_lookup/worker.py @@ -17,6 +17,7 @@ get_running_callbacks, save_message, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -26,6 +27,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 100 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def example_lookup(task, logger: logging.Logger): @@ -135,9 +137,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/example_score/worker.py b/workers/example_score/worker.py index 8236f53..0b935ce 100644 --- a/workers/example_score/worker.py +++ b/workers/example_score/worker.py @@ -7,6 +7,7 @@ import uuid from shepherd_utils.db import get_message, save_message from shepherd_utils.shared import get_tasks, run_task_lifecycle +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -15,6 +16,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 3 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def example_score(task, logger: logging.Logger): @@ -47,9 +49,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/filter_analyses_top_n/worker.py b/workers/filter_analyses_top_n/worker.py index 5db85de..d8402c1 100644 --- a/workers/filter_analyses_top_n/worker.py +++ b/workers/filter_analyses_top_n/worker.py @@ -6,6 +6,7 @@ import uuid from shepherd_utils.db import get_message, save_message, get_query_state from shepherd_utils.shared import get_tasks, run_task_lifecycle +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -14,6 +15,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 5 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def filter_analyses_top_n(task, logger: logging.Logger): @@ -56,9 +58,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/filter_kgraph_orphans/worker.py b/workers/filter_kgraph_orphans/worker.py index 7cffeda..5627589 100644 --- a/workers/filter_kgraph_orphans/worker.py +++ b/workers/filter_kgraph_orphans/worker.py @@ -6,6 +6,7 @@ import uuid from shepherd_utils.db import get_message, save_message +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import ( filter_kgraph_orphans, @@ -19,6 +20,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def do_filter_kgraph_orphans(task, logger: logging.Logger): @@ -51,9 +53,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/filter_results_top_n/worker.py b/workers/filter_results_top_n/worker.py index 4652669..d535e2a 100644 --- a/workers/filter_results_top_n/worker.py +++ b/workers/filter_results_top_n/worker.py @@ -6,6 +6,7 @@ import uuid from shepherd_utils.db import get_message, save_message, get_query_state from shepherd_utils.shared import get_tasks, run_task_lifecycle +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -14,6 +15,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 5 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def filter_results_top_n(task, logger: logging.Logger): @@ -59,9 +61,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/finish_query/worker.py b/workers/finish_query/worker.py index 3546f2e..2cbed87 100644 --- a/workers/finish_query/worker.py +++ b/workers/finish_query/worker.py @@ -19,6 +19,7 @@ set_query_completed, ) from shepherd_utils.shared import get_tasks +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -27,6 +28,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) CALLBACK_RETRIES = 3 @@ -130,9 +132,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/merge_message/worker.py b/workers/merge_message/worker.py index 6b52901..bdfa0ba 100644 --- a/workers/merge_message/worker.py +++ b/workers/merge_message/worker.py @@ -30,7 +30,7 @@ save_logs, save_message_sync, ) -from shepherd_utils.logger import QueryLogger +from shepherd_utils.logger import QueryLogger, get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import filter_kgraph_orphans, get_tasks, merge_kgraph @@ -41,6 +41,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 10 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def get_edgeset(result): @@ -690,7 +691,7 @@ def merge_messages_by_ids( # otherwise handlers would accumulate across the child's successive tasks # and leak one query's logs into the next. query_log_handler = QueryLogger().log_handler - worker_logger = logging.getLogger(f"merge_message.worker.{os.getpid()}") + worker_logger = get_worker_logger(f"merge_message.worker.{os.getpid()}") worker_logger.setLevel(log_level) worker_logger.addHandler(query_log_handler) try: @@ -766,8 +767,8 @@ async def poll_for_tasks(): # and the in-flight task limit below: each merge runs a child that loads the # growing response blob, so pool size == concurrency bounds peak memory. # POOL_MAX_WORKERS overrides. - max_workers = resolve_pool_workers(TASK_LIMIT, logging.getLogger(STREAM)) - logging.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") # Shared self-healing pool: spawn-context executor that replaces itself in # place on a BrokenProcessPool (same implementation the aragorn.omnicorp / # aragorn.score / arax.rank workers use). run() swaps the dead pool before @@ -957,13 +958,13 @@ async def process_query(task, parent_ctx, logger, limiter): inflight.add(t) t.add_done_callback(inflight.discard) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") for t in inflight: t.cancel() pool.shutdown() return except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index 635264e..1c84e63 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -1,7 +1,6 @@ """Path scoring module""" import asyncio -import logging import time import uuid from concurrent.futures import ThreadPoolExecutor @@ -17,6 +16,7 @@ from shepherd_utils.config import settings from shepherd_utils.data_download import ensure_pathfinder_embeddings from shepherd_utils.db import get_message_sync, save_message_sync +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -26,6 +26,7 @@ TASK_LIMIT = 4 EMBEDDING_DIR = settings.pathfinder_embeddings_dir tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) def convert_path_to_components(source, target, path, knowledge_graph, logger): @@ -275,7 +276,7 @@ async def poll_for_tasks(): # local `docker compose up` starts with the volume-mounted directory empty). # No-op once present or when no download URL is configured (e.g. production, # where the data is mounted out of band). - ensure_pathfinder_embeddings(logging.getLogger(STREAM)) + ensure_pathfinder_embeddings(LOGGER) clf = XGBClassifier() clf.load_model("model_weights/squashbert_classifier_weights.json") bmt = Toolkit() @@ -283,7 +284,7 @@ async def poll_for_tasks(): EMBEDDING_DIR, readonly=True, lock=False, readahead=False, subdir=True ) count, sample = _probe_cache(embedding_env) - logging.info(f"embeddings cache: {count} entries (sample key: {sample!r})") + LOGGER.info(f"embeddings cache: {count} entries (sample key: {sample!r})") mlp = nn.Sequential( nn.Linear(11 * 768, 1536), nn.GELU(), @@ -304,9 +305,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) diff --git a/workers/sipr/worker.py b/workers/sipr/worker.py index 83d2610..6e472a7 100644 --- a/workers/sipr/worker.py +++ b/workers/sipr/worker.py @@ -14,6 +14,7 @@ get_message, save_message, ) +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer from shepherd_utils.shared import get_tasks, run_task_lifecycle @@ -24,6 +25,7 @@ TASK_LIMIT = 10 MAX_QUERY_TIME = 2400 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def get_neighborhood(id_list: list[str], depth: int, logger): @@ -349,9 +351,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying diff --git a/workers/sort_results_score/worker.py b/workers/sort_results_score/worker.py index e3ac1d9..e0790fc 100644 --- a/workers/sort_results_score/worker.py +++ b/workers/sort_results_score/worker.py @@ -11,6 +11,7 @@ get_query_state, ) from shepherd_utils.shared import get_tasks, run_task_lifecycle +from shepherd_utils.logger import get_worker_logger from shepherd_utils.otel import setup_tracer # Queue name @@ -19,6 +20,7 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 3 tracer = setup_tracer(STREAM) +LOGGER = get_worker_logger(STREAM) async def sort_results_score(task, logger: logging.Logger): @@ -82,9 +84,9 @@ async def poll_for_tasks(): ): asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) except asyncio.CancelledError: - logging.info("Poll loop cancelled, shutting down.") + LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: - logging.error(f"Error in task polling loop: {e}", exc_info=True) + LOGGER.error(f"Error in task polling loop: {e}", exc_info=True) await asyncio.sleep(5) # back off before retrying