diff --git a/requirements.txt b/requirements.txt index 817aed9..2c359e4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,7 @@ git+https://github.com/falkordb/falkordb-bulk-loader.git@v1.0.6 setuptools>=66 pytest PyYAML -git+https://github.com/helxplatform/dug@develop +git+https://github.com/helxplatform/dug@tranql-crawl-concurrency orjson>=3.11 git+https://github.com/helxplatform/kg_utils.git@v0.0.10.1 # kg_utils/merging.py hashes an f-string; xxhash 4.0 dropped the diff --git a/scripts/gzip_artifacts.py b/scripts/gzip_artifacts.py new file mode 100644 index 0000000..07b6f39 --- /dev/null +++ b/scripts/gzip_artifacts.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python +"""Gzip-compress annotate/crawl artifacts committed before storage.py's +write_object started gzipping them (see roger.core.storage). + +A dataset annotated before that fix has its elements.txt/concepts.txt sitting +in lakefs uncompressed -- ~3MB/dir raw for a typical dbGaP data dict pair. +crawl's own output is gzipped going forward, but crawl still has to +*download* that old, uncompressed annotation output as input first: for +bdc-parent's 61,597 dirs that is ~186GB of local disk before crawl writes a +single byte, which is most of what blew the crawl task's 200G PVC. + +Unlike migrate_pickled_classes.py this does no jsonpickle decode/re-encode -- +pure byte-level gzip, so the serialized content is untouched, just smaller. + + lakectl local clone lakefs:////annotate_and_index// ./out + python scripts/gzip_artifacts.py ./out # compress in place + lakectl local commit ./out -m "gzip annotation output" + + python scripts/gzip_artifacts.py --dry-run ./out # report only + python scripts/gzip_artifacts.py --self-check # no dir needed +""" + +import argparse +import gzip +import sys +from pathlib import Path + +ARTIFACTS = ('elements.txt', 'concepts.txt', 'expanded_concepts.txt') +GZIP_MAGIC = b'\x1f\x8b' +# level 6 matches storage.py's write_object -- level 9 (gzip's default) +# burned 3-4x the cpu for the same ratio on this repetitive JSON. +COMPRESSLEVEL = 6 + + +def artifact_files(root): + return sorted(p for p in Path(root).rglob('*.txt') if p.name in ARTIFACTS) + + +def compress_file(path, dry_run=False): + """Returns (raw_size, compressed_size) if compressed, None if skipped + (already gzip or empty).""" + raw = path.read_bytes() + if not raw or raw[:2] == GZIP_MAGIC: + return None + compressed = gzip.compress(raw, compresslevel=COMPRESSLEVEL) + if not dry_run: + tmp = path.with_name(path.name + '.gzip-tmp') + tmp.write_bytes(compressed) + tmp.replace(path) + return len(raw), len(compressed) + + +def run(root, dry_run=False): + files = artifact_files(root) + total_raw = total_compressed = 0 + changed = skipped = 0 + for i, path in enumerate(files, 1): + result = compress_file(path, dry_run=dry_run) + if result is None: + skipped += 1 + continue + raw_size, compressed_size = result + total_raw += raw_size + total_compressed += compressed_size + changed += 1 + if changed % 5000 == 0: + print(f" {changed} compressed ({i}/{len(files)} scanned)") + verb = "would compress" if dry_run else "compressed" + print(f"{verb} {changed} of {len(files)} file(s), " + f"{skipped} already gzip or empty") + if total_raw: + print(f"{total_raw / 1e9:.2f}GB -> {total_compressed / 1e9:.2f}GB " + f"({total_raw / total_compressed:.1f}x)") + return changed + + +def self_check(): + "Round-trip a small payload; fails loudly if compress_file regresses." + import tempfile + payload = b'{"id": "UMLS:C1"}' * 500 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / 'concepts.txt' + path.write_bytes(payload) + + result = compress_file(path, dry_run=True) + assert result is not None, "dry-run should report a would-be change" + assert path.read_bytes() == payload, "dry-run must not touch the file" + + result = compress_file(path) + assert result is not None + raw_size, compressed_size = result + assert raw_size == len(payload) + assert compressed_size < raw_size, "should have actually shrunk" + + with open(path, 'rb') as f: + assert f.read(2) == GZIP_MAGIC + assert gzip.decompress(path.read_bytes()) == payload + + # idempotent: running again on an already-gzipped file is a no-op + assert compress_file(path) is None + assert gzip.decompress(path.read_bytes()) == payload + + print("self-check ok") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('root', nargs='?', help='directory to compress in place') + parser.add_argument('--dry-run', action='store_true', + help='report what would change without writing') + parser.add_argument('--self-check', action='store_true', + help='round-trip a synthetic payload, no dir needed') + args = parser.parse_args() + + if args.self_check: + self_check() + return + if not args.root: + parser.error("root directory required unless --self-check") + run(args.root, dry_run=args.dry_run) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/scripts/migrate_pickled_classes.py b/scripts/migrate_pickled_classes.py index a6e06b1..6ace171 100644 --- a/scripts/migrate_pickled_classes.py +++ b/scripts/migrate_pickled_classes.py @@ -16,6 +16,7 @@ """ import argparse +import gzip import importlib import os import json @@ -39,6 +40,23 @@ PY_OBJECT = re.compile(r'"py/object":\s*"([^"]+)"') ARTIFACTS = ('elements.txt', 'concepts.txt', 'expanded_concepts.txt') +GZIP_MAGIC = b'\x1f\x8b' + + +def read_artifact_text(path): + """roger's storage.write_object gzips these now; older artifacts + committed before that are still plain text, so detect and handle both.""" + raw = path.read_bytes() + if raw[:2] == GZIP_MAGIC: + return gzip.decompress(raw).decode('utf-8') + return raw.decode('utf-8') + + +def write_artifact_text(path, text): + """Always write gzip -- migration is also the chance to bring an old + plain-text artifact up to the current on-disk format.""" + path.write_bytes(gzip.compress(text.encode('utf-8'))) + def install_alias(module_path): """Stand in for a legacy module, resolving class names against CURRENT.""" @@ -110,7 +128,7 @@ def field_drift_paths(paths): """ drift = {} for path in paths: - obj = jsonpickle.decode(path.read_text()) + obj = jsonpickle.decode(read_artifact_text(path)) stack, seen = [obj], set() while stack: item = stack.pop() @@ -137,7 +155,7 @@ def scan(root): print(f"{len(files)} artifact file(s) under {root}") found = set() for path in files: - found |= classes_in(path.read_text()) + found |= classes_in(read_artifact_text(path)) broken = broken_modules(found) for cls in sorted(found): module_path = cls.rpartition('.')[0] @@ -180,7 +198,7 @@ def dead_modules(sample_paths): """ dead, found = set(), set() for path in sample_paths: - found |= classes_in(path.read_text()) + found |= classes_in(read_artifact_text(path)) for cls in found: module_path = cls.rpartition('.')[0] if module_path in dead or module_path in sys.modules: @@ -271,7 +289,7 @@ def restamp(root, sample=20, dry_run=False): changed = 0 for i, path in enumerate(files, 1): - text = path.read_text() + text = read_artifact_text(path) new_text, unmapped = restamp_text(text, dead) if unmapped: raise SystemExit( @@ -284,7 +302,7 @@ def restamp(root, sample=20, dry_run=False): # write-then-rename: a pod killed mid-write must not leave a # truncated artifact behind, and there are 150k of them tmp = path.with_name(path.name + '.restamp-tmp') - tmp.write_text(new_text) + write_artifact_text(tmp, new_text) os.replace(tmp, path) if changed % 5000 == 0: print(f" {changed} rewritten ({i}/{len(files)} scanned)") @@ -296,13 +314,13 @@ def restamp(root, sample=20, dry_run=False): def fix(root, dry_run=False): files = artifact_files(root) for module_path in broken_modules( - {c for p in files for c in classes_in(p.read_text())}): + {c for p in files for c in classes_in(read_artifact_text(p))}): print(f"aliasing legacy module {module_path}") install_alias(module_path) changed = 0 for path in files: - text = path.read_text() + text = read_artifact_text(path) obj = jsonpickle.decode(text) fill_defaults(obj) rewritten = jsonpickle.encode(obj, indent=2) @@ -311,7 +329,7 @@ def fix(root, dry_run=False): changed += 1 print(f"{'would rewrite' if dry_run else 'rewrote'} {path}") if not dry_run: - path.write_text(rewritten) + write_artifact_text(path, rewritten) print(f"{changed} of {len(files)} file(s) needed migration") return changed diff --git a/src/roger/config/__init__.py b/src/roger/config/__init__.py index a7732b0..ddfe1d2 100644 --- a/src/roger/config/__init__.py +++ b/src/roger/config/__init__.py @@ -232,6 +232,19 @@ class IndexingConfig(DictLike): "anat_to_pheno": ["anatomical_entity", "phenotypic_feature"], }) tranql_endpoint: str = "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + # Concurrency for the crawl. `crawl_workers` threads the TranQL fetches + # within one concept; `crawl_file_workers` threads whole input files. + # They multiply, and the product should not exceed what the TranQL + # service can serve at once (its gunicorn worker count) -- past that, + # requests only queue. + crawl_workers: int = 4 + crawl_file_workers: int = 4 + # crawl_file_workers threads doing real work (TranQL fetches, jsonpickle + # encode, gzip) on the chart's default cpu limit throttled a crawl pod + # 70% of its scheduling periods, cutting throughput to a third. Match + # crawl_file_workers 1:1 with cores so each worker thread gets its own, + # instead of 4 threads fighting over a fraction of one. + crawl_cpu: str = "4" # by default skips node to element queries node_to_element_queries: dict = field(default_factory=lambda: {}) element_mapping: str = "" diff --git a/src/roger/config/config.yaml b/src/roger/config/config.yaml index 3d883d1..4d802ae 100644 --- a/src/roger/config/config.yaml +++ b/src/roger/config/config.yaml @@ -120,6 +120,9 @@ indexing: "chemical_mixture_to_disease": ["chemical_mixture", "disease"] "phen_to_anat": ["phenotypic_feature", "anatomical_entity"] tranql_endpoint: "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + crawl_workers: 4 + crawl_file_workers: 4 + crawl_cpu: "4" node_to_element_queries: enabled: false cde: diff --git a/src/roger/core/bulkload.py b/src/roger/core/bulkload.py index 4dc2736..a2dd0db 100644 --- a/src/roger/core/bulkload.py +++ b/src/roger/core/bulkload.py @@ -54,12 +54,19 @@ def create_nodes_csv_file(self, input_data_path=None, output_data_path=None): merged_nodes_file = storage.merged_objects('nodes', input_data_path) counter = 1 for node in storage.json_line_iter(merged_nodes_file): + # \r alone (no paired \n) still reads as a line break under + # universal newlines -- both our own row count and the bulk + # loader's CSV parser split on it, turning one row into two + # and corrupting the column count. Seen in dbGaP codebook text + # pasted in with stray CRs (e.g. a Study "activitybk" field). if node.get('description'): - node['description'] = node['description'].replace('\n', - ' ') + node['description'] = ( + node['description'].replace('\r\n', ' ') + .replace('\r', ' ').replace('\n', ' ')) if node.get('name'): - node['name'] = node['name'].replace('\n', - ' ') + node['name'] = ( + node['name'].replace('\r\n', ' ') + .replace('\r', ' ').replace('\n', ' ')) if not node.get('category'): category_error_nodes.add(node['id']) node['category'] = [BiolinkModel.root_type] diff --git a/src/roger/core/storage.py b/src/roger/core/storage.py index d10d0d6..0f9f1cc 100644 --- a/src/roger/core/storage.py +++ b/src/roger/core/storage.py @@ -5,6 +5,7 @@ import os import glob +import gzip import time import pathlib import pickle @@ -86,9 +87,25 @@ def read_object(path, key=None): with open(file=path, mode="rb") as stream: obj = pickle.load(stream) elif path.endswith(".jsonl") or path.endswith('.txt'): - obj = read_data(path) + obj = read_gzip_or_plain_text(path) if not is_web(path) \ + else read_data(path) return obj +# gzip magic number -- distinguishes a compressed artifact from the plain +# text ones already committed to lakefs before this was added, so both +# read transparently and nothing downstream needs to change. +GZIP_MAGIC = b'\x1f\x8b' + +def read_gzip_or_plain_text(path): + """ Read a local .txt/.jsonl artifact, decompressing it if it was + gzip-written by write_object; falls back to plain text for artifacts + written before compression was added. """ + with open(path, 'rb') as stream: + raw = stream.read() + if raw[:2] == GZIP_MAGIC: + return gzip.decompress(raw).decode('utf-8') + return raw.decode('utf-8') + def is_web (uri): """ The URI is a web URI (starts with http or https). :param uri: A URI """ @@ -122,8 +139,19 @@ def write_object (obj, path, key=None): with open (path, "wb") as stream: pickle.dump(obj, file=stream) elif path.endswith(".jsonl") or path.endswith('.txt'): - with open (path, "w", encoding="utf-8") as stream: - stream.write(obj) + # gzip -- these are the crawl-stage KG-answer artifacts, and the + # same repeated CURIEs/biolink categories/JSON keys compress + # 5-10x; that's the difference between fitting a large dataset's + # crawl output on the PVC and hitting ENOSPC mid-run. Same + # filename/extension as before so every glob pattern that finds + # these files by name still matches; read_object detects the + # gzip magic number so old uncompressed artifacts still read. + with open(path, "wb") as stream: + # level 9 (gzip.compress's default) burned 3-4x the CPU of + # level 6 for the same ratio on this repetitive JSON -- deadly + # under a crawl task's thin CPU limit (measured 70% of periods + # throttled on a 250m limit with 4 crawl workers). + stream.write(gzip.compress(obj.encode('utf-8'), compresslevel=6)) else: # Raise an exception if invalid. raise ValueError (f"Unrecognized extension: {path}") diff --git a/src/roger/pipelines/base.py b/src/roger/pipelines/base.py index f3af15f..fc83edf 100644 --- a/src/roger/pipelines/base.py +++ b/src/roger/pipelines/base.py @@ -388,6 +388,30 @@ def annotation_is_complete(cls, parse_file, output_data_path): for path in cls.annotation_output_paths(parse_file, output_data_path)) + @staticmethod + def crawl_output_path(concept_file, output_data_path=None): + "The expanded_concepts.txt file crawl_one_file writes for a concept_file" + data_set_name = os.path.split(os.path.dirname(concept_file))[-1] + output_file_name = os.path.join(data_set_name, 'expanded_concepts.txt') + if not output_data_path: + return storage.dug_expanded_concepts_path(output_file_name) + return os.path.join(output_data_path, output_file_name) + + @classmethod + def crawl_is_complete(cls, concept_file, output_data_path): + """True if this file's crawl output is already fully written. + + crawl_one_file has no partial-write hazard like annotation's two + files do -- expanded_concepts.txt is written in one call -- so + existence and non-empty is the whole check. Without this, a resumed + try re-crawls every file from scratch: TranQL's response cache makes + the already-done ones cheap, but not free, and a large dataset's + never-before-seen files still queue behind all the free-but-not- + instant redone work before any real progress resumes. + """ + path = cls.crawl_output_path(concept_file, output_data_path) + return os.path.isfile(path) and os.path.getsize(path) > 0 + def annotate_one_file(self, parse_file, parser, output_data_path, index=0, total=0): "Parse and annotate a single input file, writing pickles for it" @@ -778,10 +802,6 @@ def crawl_concepts(self, concepts, data_set_name, output_path=None): :param data_set_name: :return: """ - # TODO crawl dir seems to be storaing crawling info to avoid - # re-crawling, but is that consting us much? , it was when tranql was - # slow, but might right to consider getting rid of it. - crawl_dir = storage.dug_crawl_path('crawl_output') output_file_name = os.path.join(data_set_name, 'expanded_concepts.txt') extracted_dug_elements_file_name = os.path.join( @@ -796,7 +816,6 @@ def crawl_concepts(self, concepts, data_set_name, output_path=None): extracted_output_file = os.path.join( output_path, extracted_dug_elements_file_name) - Path(crawl_dir).mkdir(parents=True, exist_ok=True) extracted_dug_elements = [] log.debug("Creating Dug Crawler object") crawler = Crawler( @@ -806,8 +825,8 @@ def crawl_concepts(self, concepts, data_set_name, output_path=None): tranqlizer=self.tranqlizer, tranql_queries=self.tranql_queries, http_session=self.cached_session, + crawl_workers=self.config.indexing.crawl_workers, ) - crawler.crawlspace = crawl_dir counter = 0 total = len(concepts) for concept in concepts.values(): @@ -1220,78 +1239,106 @@ def crawl_tranql(self, to_string=False, concept_files=None, input_data_path, format='txt') if output_data_path: - crawl_dir = os.path.join(output_data_path, 'crawl_output') expanded_concepts_dir = os.path.join(output_data_path, 'expanded_concepts') else: - crawl_dir = storage.dug_crawl_path('crawl_output') expanded_concepts_dir = storage.dug_expanded_concepts_path("") - log.info("Clearing crawl output dir %s", crawl_dir) - storage.clear_dir(crawl_dir) log.info("Clearing expanded concepts dir: %s", expanded_concepts_dir) storage.clear_dir(expanded_concepts_dir) - log.info("Crawling Dug Concepts, found %d file(s).", - len(concept_files)) - for file_ in concept_files: - objects = storage.read_object(file_) - objects = objects or {} - if not objects: - log.info(f'no concepts in {file_}') - data_set = jsonpickle.decode(objects) - original_variables_dataset_name = os.path.split( - os.path.dirname(file_))[-1] - self.crawl_concepts(concepts=data_set, - data_set_name=original_variables_dataset_name, - output_path= output_data_path) - - # After expanding concepts with KG answers, update the - # corresponding elements' optional_terms so that KG-derived - # search terms are present when elements are later indexed. - # This mirrors what Crawler.crawl() does after concept expansion. - # The updated elements are written to the expanded concepts - # directory (alongside expanded_concepts.txt) rather than - # mutating the annotate step's output. - annotation_elements_file = os.path.join( - os.path.dirname(file_), 'elements.txt') - expanded_elements_file_name = os.path.join( - original_variables_dataset_name, 'elements.txt') - if not output_data_path: - expanded_elements_file = ( - storage.dug_expanded_concepts_path( - expanded_elements_file_name)) - else: - expanded_elements_file = os.path.join( - output_data_path, expanded_elements_file_name) - if os.path.exists(annotation_elements_file): - log.info("Updating element optional terms from expanded " - "concepts for %s", original_variables_dataset_name) - elements = jsonpickle.decode( - storage.read_object(annotation_elements_file)) - for element in elements: - if isinstance(element, DugConcept): - continue - # Replace each element's concept references with - # the expanded versions that now carry kg_answers. - for concept_id in list(element.concepts.keys()): - if concept_id in data_set: - element.concepts[concept_id] = data_set[ - concept_id] - element.set_optional_terms() - storage.write_object( - jsonpickle.encode(elements, indent=2), - expanded_elements_file) - log.info("Updated elements serialized to %s", - expanded_elements_file) - else: - log.warning("Elements file not found at %s, skipping " - "optional terms update", - annotation_elements_file) + pending = [f for f in concept_files + if not self.crawl_is_complete(f, output_data_path)] + skipped = len(concept_files) - len(pending) + if skipped: + log.info("Resuming: %d of %d files already crawled, %d to go", + skipped, len(concept_files), len(pending)) + + if not pending: + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + workers = max(1, int(self.config.indexing.crawl_file_workers)) + workers = min(workers, len(pending)) or 1 + log.info("Crawling Dug Concepts, found %d file(s) with %d worker(s).", + len(pending), workers) + if workers == 1: + for file_ in pending: + self.crawl_one_file(file_, output_data_path) + else: + # Files are independent: each decodes its own concepts, expands + # them and writes its own output dir. The work is nearly all + # http wait on TranQL, so threads scale it despite the GIL. + with ThreadPoolExecutor(max_workers=workers, + thread_name_prefix='crawl') as pool: + futures = [pool.submit(self.crawl_one_file, file_, + output_data_path) + for file_ in pending] + # surface the first failure rather than letting the pool + # swallow it; a raised exception means that file produced + # nothing and the task must not report success + for future in futures: + future.result() output_log = self.log_stream.getvalue() if to_string else '' return output_log + def crawl_one_file(self, file_, output_data_path=None): + "Expand one annotate output file's concepts and write its outputs" + objects = storage.read_object(file_) + objects = objects or {} + if not objects: + log.info(f'no concepts in {file_}') + data_set = jsonpickle.decode(objects) + original_variables_dataset_name = os.path.split( + os.path.dirname(file_))[-1] + self.crawl_concepts(concepts=data_set, + data_set_name=original_variables_dataset_name, + output_path= output_data_path) + + # After expanding concepts with KG answers, update the + # corresponding elements' optional_terms so that KG-derived + # search terms are present when elements are later indexed. + # This mirrors what Crawler.crawl() does after concept expansion. + # The updated elements are written to the expanded concepts + # directory (alongside expanded_concepts.txt) rather than + # mutating the annotate step's output. + annotation_elements_file = os.path.join( + os.path.dirname(file_), 'elements.txt') + expanded_elements_file_name = os.path.join( + original_variables_dataset_name, 'elements.txt') + if not output_data_path: + expanded_elements_file = ( + storage.dug_expanded_concepts_path( + expanded_elements_file_name)) + else: + expanded_elements_file = os.path.join( + output_data_path, expanded_elements_file_name) + if os.path.exists(annotation_elements_file): + log.info("Updating element optional terms from expanded " + "concepts for %s", original_variables_dataset_name) + elements = jsonpickle.decode( + storage.read_object(annotation_elements_file)) + for element in elements: + if isinstance(element, DugConcept): + continue + # Replace each element's concept references with + # the expanded versions that now carry kg_answers. + for concept_id in list(element.concepts.keys()): + if concept_id in data_set: + element.concepts[concept_id] = data_set[ + concept_id] + element.set_optional_terms() + storage.write_object( + jsonpickle.encode(elements, indent=2), + expanded_elements_file) + log.info("Updated elements serialized to %s", + expanded_elements_file) + else: + log.warning("Elements file not found at %s, skipping " + "optional terms update", + annotation_elements_file) + def index_concepts(self, to_string=False, input_data_path=None, output_data_path=None): "Index concepts from expanded concept files" diff --git a/src/roger/tasks.py b/src/roger/tasks.py index 37f60f6..248e742 100755 --- a/src/roger/tasks.py +++ b/src/roger/tasks.py @@ -127,21 +127,35 @@ def get_executor_config(data_path='/opt/airflow/share/data'): def memory_override(limit: str, request: str = None) -> dict: - """executor_config bumping only this task's memory. + """executor_config bumping only this task's memory. See + resource_override -- this is kept as a thin wrapper since it is the + common case and already used at several call sites.""" + return resource_override(memory_limit=limit, memory_request=request) + + +def resource_override(memory_limit: str = None, memory_request: str = None, + cpu_limit: str = None, cpu_request: str = None) -> dict: + """executor_config bumping only this task's cpu and/or memory. Everything else (image, volumes, env, service account) is inherited from the chart's worker pod template; this patches the 'base' container so one heavy task does not force the default up for every task. Keep request - well under limit: the namespace quota counts requests.memory and - limits.memory separately. + well under limit: the namespace quota counts requests.memory/cpu and + limits.memory/cpu separately. """ from kubernetes.client import models as k8s + requests, limits = {}, {} + if memory_limit: + limits["memory"] = memory_limit + requests["memory"] = memory_request or "1Gi" + if cpu_limit: + limits["cpu"] = cpu_limit + requests["cpu"] = cpu_request or cpu_limit return {"pod_override": k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ k8s.V1Container( name="base", resources=k8s.V1ResourceRequirements( - requests={"memory": request or "1Gi"}, - limits={"memory": limit}))]))} + requests=requests, limits=limits))]))} def init_lakefs_client(config: RogerConfig) -> LakeFsWrapper: @@ -772,7 +786,7 @@ def create_python_task(dag, name, a_callable, func_kwargs=None, external_repos=None, pass_conf=True, no_output_files=False, no_input_files=False, incremental_pull=True, clear_output_prefix=False, - memory=None, resumable=False): + memory=None, cpu=None, resumable=False): """ Create a python task. :param func_kwargs: additional arguments for callable. :param dag: dag to add task to. @@ -793,6 +807,8 @@ def create_python_task(dag, name, a_callable, func_kwargs=None, vary run to run (the bulk-load CSVs) and would otherwise accumulate. :param memory: memory limit for this task's pod, e.g. '15Gi'. Omit to take the chart's worker default. + :param cpu: cpu limit for this task's pod, e.g. '1'. Omit to take the + chart's worker default. """ if external_repos is None: @@ -814,8 +830,9 @@ def create_python_task(dag, name, a_callable, func_kwargs=None, # executor_config example left commented; fill if needed "dag": dag, } - if memory: - python_operator_args["executor_config"] = memory_override(memory) + if memory or cpu: + python_operator_args["executor_config"] = resource_override( + memory_limit=memory, cpu_limit=cpu) if config.lakefs_config.enabled: pre_exec_conf = { @@ -962,6 +979,11 @@ def create_pipeline_taskgroup( crawl_callable, # expands every concept through tranql, accumulating answers memory=configparam.annotation.annotate_memory, + # crawl_file_workers threads doing real CPU work (TranQL + # fetches, jsonpickle encode, gzip) on the chart's thin default + # cpu limit throttled a crawl pod 70% of its scheduling + # periods, cutting throughput to a third. + cpu=configparam.indexing.crawl_cpu, pass_conf=False) crawl_task.set_upstream(annotate_task) diff --git a/tests/unit/test_bulkload_newline_sanitization.py b/tests/unit/test_bulkload_newline_sanitization.py new file mode 100644 index 0000000..c50ca3f --- /dev/null +++ b/tests/unit/test_bulkload_newline_sanitization.py @@ -0,0 +1,50 @@ +"""A bare \\r (no paired \\n) in a node's description/name still reads as a +line break under universal newlines. Both our own tooling and the +falkordb_bulk_loader CSV parser split on it, turning one row into two and +corrupting the column count -- seen in dbGaP codebook text pasted in with +stray CRs (a Study "activitybk" field). +""" +from roger.core.bulkload import BulkLoad +from roger.config import config + + +def test_bare_cr_in_description_does_not_split_the_row(tmp_path, monkeypatch): + from roger.core import bulkload as bulkload_mod + + node = { + 'id': 'HDP00066:activitybk', + 'category': ['biolink:Study'], + 'name': 'activitybk', + 'description': 'SECTION H. READ Scale\r Ask: do you read to your child?', + } + leaf_class = 'biolink:Study' + schema = { + leaf_class: { + 'category': 'list', 'name': 'str', + 'description': 'str', 'id': 'str', + } + } + + monkeypatch.setattr(bulkload_mod.storage, 'merged_objects', + lambda kind, path=None: 'nodes.jsonl') + monkeypatch.setattr(bulkload_mod.storage, 'json_line_iter', + lambda path: iter([node])) + monkeypatch.setattr(bulkload_mod.storage, 'read_schema', + lambda schema_type, path=None: schema) + monkeypatch.setattr(bulkload_mod.storage, 'bulk_path', + lambda name, path=None: str(tmp_path / name)) + + class FakeBiolink: + def get_leaf_class(self, names): + return leaf_class + + bulk = BulkLoad(FakeBiolink(), config=config) + bulk.create_nodes_csv_file(input_data_path=None, output_data_path=None) + + out_file = tmp_path / 'nodes' / f"{leaf_class.replace(':', '~')}.csv-0-1" + lines = out_file.read_bytes().split(b'\n') + lines = [l for l in lines if l] + assert len(lines) == 2, lines # header + one data row, not split in two + header, row = (l.decode().split('\x1e') for l in lines) + assert len(header) == len(row) == 4 + assert b'\r' not in out_file.read_bytes() diff --git a/tests/unit/test_crawl_file_workers.py b/tests/unit/test_crawl_file_workers.py new file mode 100644 index 0000000..0e2914f --- /dev/null +++ b/tests/unit/test_crawl_file_workers.py @@ -0,0 +1,128 @@ +"""crawl_tranql threads whole input files; failures must not be swallowed. + +The dir loop is where the bulk of crawl concurrency comes from -- one file +per annotated input, tens of thousands of them for a dbGaP dataset. Two +things must hold: every file gets processed exactly once regardless of +worker count, and a file that raises fails the task rather than quietly +producing no output. +""" +import os +import threading +import types + +import pytest + + +def make_pipeline(workers, crawl_one): + """A stand-in carrying only what crawl_tranql touches.""" + from roger.pipelines.base import DugPipeline + + pipeline = types.SimpleNamespace() + pipeline.config = types.SimpleNamespace( + indexing=types.SimpleNamespace(crawl_file_workers=workers)) + pipeline.log_stream = types.SimpleNamespace(getvalue=lambda: '') + pipeline.crawl_one_file = crawl_one + # nothing exists yet in these tests -- every file is pending, same as + # the pre-skip-check behavior this test suite was written against + pipeline.crawl_is_complete = lambda f, output_data_path: False + pipeline.crawl_tranql = types.MethodType( + DugPipeline.crawl_tranql.__wrapped__ + if hasattr(DugPipeline.crawl_tranql, '__wrapped__') + else DugPipeline.crawl_tranql, pipeline) + return pipeline + + +FILES = [f"/in/file{i}/concepts.txt" for i in range(12)] + + +@pytest.mark.parametrize("workers", [1, 4, 8]) +def test_every_file_processed_once(workers, monkeypatch, tmp_path): + import roger.pipelines.base as base + + seen = [] + lock = threading.Lock() + + def crawl_one(file_, output_data_path=None): + with lock: + seen.append(file_) + + monkeypatch.setattr(base.storage, 'clear_dir', lambda *a, **k: None) + pipeline = make_pipeline(workers, crawl_one) + pipeline.crawl_tranql(concept_files=list(FILES), + output_data_path=str(tmp_path)) + + assert sorted(seen) == sorted(FILES) + assert len(seen) == len(FILES) + + +def test_one_bad_file_fails_the_task(monkeypatch, tmp_path): + import roger.pipelines.base as base + + def crawl_one(file_, output_data_path=None): + if file_.endswith("file7/concepts.txt"): + raise ValueError("bad pickle") + + monkeypatch.setattr(base.storage, 'clear_dir', lambda *a, **k: None) + pipeline = make_pipeline(4, crawl_one) + with pytest.raises(ValueError, match="bad pickle"): + pipeline.crawl_tranql(concept_files=list(FILES), + output_data_path=str(tmp_path)) + + +def test_worker_count_is_capped_by_file_count(monkeypatch, tmp_path): + """Two files must not spin up eight threads.""" + import roger.pipelines.base as base + + threads = set() + lock = threading.Lock() + + def crawl_one(file_, output_data_path=None): + with lock: + threads.add(threading.current_thread().name) + + monkeypatch.setattr(base.storage, 'clear_dir', lambda *a, **k: None) + pipeline = make_pipeline(8, crawl_one) + pipeline.crawl_tranql(concept_files=FILES[:2], + output_data_path=str(tmp_path)) + assert len(threads) <= 2 + + +def test_already_crawled_files_are_skipped(monkeypatch, tmp_path): + """A resumed try must not redo files an earlier try already crawled -- + TranQL's response cache makes that cheap, not free, and it stands + between a large dataset and any real new progress.""" + import roger.pipelines.base as base + + seen = [] + + def crawl_one(file_, output_data_path=None): + seen.append(file_) + + monkeypatch.setattr(base.storage, 'clear_dir', lambda *a, **k: None) + pipeline = make_pipeline(4, crawl_one) + already_done = set(FILES[:5]) + pipeline.crawl_is_complete = lambda f, output_data_path: f in already_done + + pipeline.crawl_tranql(concept_files=list(FILES), + output_data_path=str(tmp_path)) + + assert sorted(seen) == sorted(set(FILES) - already_done) + + +def test_crawl_is_complete_checks_the_real_pipeline(tmp_path): + """Exercise the actual DugPipeline method, not just the stand-in used + above -- this is what would have caught crawl_output_path drifting + from where crawl_concepts actually writes.""" + from roger.pipelines.base import DugPipeline + + concept_file = "/in/phs000123.v1.data_dict/concepts.txt" + output_data_path = str(tmp_path) + + assert DugPipeline.crawl_is_complete(concept_file, output_data_path) is False + + path = DugPipeline.crawl_output_path(concept_file, output_data_path) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, 'w') as f: + f.write('{}') + + assert DugPipeline.crawl_is_complete(concept_file, output_data_path) is True diff --git a/tests/unit/test_storage_gzip_artifacts.py b/tests/unit/test_storage_gzip_artifacts.py new file mode 100644 index 0000000..b0f37ac --- /dev/null +++ b/tests/unit/test_storage_gzip_artifacts.py @@ -0,0 +1,27 @@ +"""write_object gzips .txt/.jsonl artifacts now -- the repeated CURIEs and +biolink categories in crawl-stage KG-answer JSON compress 5-10x, which is +the difference between a large dataset's crawl fitting on the PVC and +hitting ENOSPC mid-run. read_object must still read artifacts committed +before this was added. +""" +import gzip + +from roger.core import storage + + +def test_txt_artifact_round_trips_through_gzip(tmp_path): + path = str(tmp_path / 'concepts.txt') + text = '{"id": "UMLS:C1"}' * 100 + + storage.write_object(text, path) + + assert open(path, 'rb').read(2) == b'\x1f\x8b' + assert storage.read_object(path) == text + + +def test_txt_artifact_backward_compat_with_plain_text(tmp_path): + path = tmp_path / 'concepts.txt' + text = '{"id": "UMLS:C1"}' + path.write_text(text, encoding='utf-8') + + assert storage.read_object(str(path)) == text diff --git a/tests/unit/test_tasks_incremental.py b/tests/unit/test_tasks_incremental.py index 7b6bbc3..9d92fce 100644 --- a/tests/unit/test_tasks_incremental.py +++ b/tests/unit/test_tasks_incremental.py @@ -427,6 +427,16 @@ def test_memory_override_patches_base_container(): assert container.resources.requests == {"memory": "1Gi"} +def test_resource_override_patches_cpu_and_memory_together(): + pytest.importorskip("kubernetes") + cfg = tasks.resource_override(memory_limit="15Gi", cpu_limit="1") + container = cfg["pod_override"].spec.containers[0] + assert container.resources.limits == {"memory": "15Gi", "cpu": "1"} + # cpu request defaults to the limit (no separate cheap-request case, + # unlike memory) since cpu limits are compressible, not a quota risk + assert container.resources.requests == {"memory": "1Gi", "cpu": "1"} + + def test_es_taskgroup_pulls_crawl_outputs_only(monkeypatch, lakefs_env): """index_variables must read crawl's expanded elements.txt, not annotate's: only the crawl copy carries KG-derived optional_terms, and