From fb37209d73dd7ec69be6837e002c466ade6dddef Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Mon, 20 Jul 2026 11:57:34 -0700 Subject: [PATCH 1/2] Separate benchmark pipeline into dedicated top-level directory --- benchmark-tests/run_evaluation.py | 291 ------------------ .../BENCHMARKING.md => benchmarks/README.md | 31 +- benchmarks/__init__.py | 4 + benchmarks/datasets/__init__.py | 4 + .../benchmark.concurrentqa.prototype | 0 .../datasets}/benchmark.cuad.prototype | 0 .../datasets}/benchmark.pga | 0 .../datasets}/benchmark.wikihow | 0 benchmarks/env.template | 26 ++ benchmarks/requirements.txt | 14 + benchmarks/scripts/__init__.py | 4 + .../scripts}/benchmark_build.py | 6 +- .../scripts}/benchmark_evaluate.py | 6 +- .../scripts}/benchmark_extract.py | 6 +- .../scripts}/benchmark_query.py | 16 +- benchmarks/scripts/integration_test_base.py | 30 ++ .../scripts/integration_test_handler.py | 223 ++++++++++++++ benchmarks/tests/__init__.py | 4 + .../tests}/test_agentic_retriever.py | 2 +- .../tests}/test_benchmark_query.py | 4 +- .../tests}/test_comparison_report.py | 2 +- .../tests}/test_hop_classifier.py | 2 +- .../tests}/test_metrics_summary.py | 4 +- .../tests}/test_retriever_factory.py | 6 +- .../tests}/test_token_tracker.py | 4 +- .../utils}/__init__.py | 0 .../utils}/agentic_retriever.py | 0 .../utils}/comparison_report.py | 0 .../utils}/hop_classifier.py | 0 .../utils}/metrics_summary.py | 0 .../utils}/retriever_factory.py | 2 +- .../utils}/run_evaluation.py | 0 .../utils}/s3_utils.py | 0 .../utils}/token_tracker.py | 0 integration-tests/build-tests.sh | 13 +- integration-tests/test-scripts/test_suite.py | 7 +- 36 files changed, 376 insertions(+), 335 deletions(-) delete mode 100644 benchmark-tests/run_evaluation.py rename integration-tests/BENCHMARKING.md => benchmarks/README.md (92%) create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/datasets/__init__.py rename {integration-tests => benchmarks/datasets}/benchmark.concurrentqa.prototype (100%) rename {integration-tests => benchmarks/datasets}/benchmark.cuad.prototype (100%) rename {integration-tests => benchmarks/datasets}/benchmark.pga (100%) rename {integration-tests => benchmarks/datasets}/benchmark.wikihow (100%) create mode 100644 benchmarks/env.template create mode 100644 benchmarks/requirements.txt create mode 100644 benchmarks/scripts/__init__.py rename {integration-tests/test-scripts/graphrag_toolkit_tests => benchmarks/scripts}/benchmark_build.py (96%) rename {integration-tests/test-scripts/graphrag_toolkit_tests => benchmarks/scripts}/benchmark_evaluate.py (96%) rename {integration-tests/test-scripts/graphrag_toolkit_tests => benchmarks/scripts}/benchmark_extract.py (95%) rename {integration-tests/test-scripts/graphrag_toolkit_tests => benchmarks/scripts}/benchmark_query.py (95%) create mode 100644 benchmarks/scripts/integration_test_base.py create mode 100644 benchmarks/scripts/integration_test_handler.py create mode 100644 benchmarks/tests/__init__.py rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_agentic_retriever.py (95%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_benchmark_query.py (98%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_comparison_report.py (99%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_hop_classifier.py (97%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_metrics_summary.py (98%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_retriever_factory.py (96%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/tests}/test_token_tracker.py (99%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/__init__.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/agentic_retriever.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/comparison_report.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/hop_classifier.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/metrics_summary.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/retriever_factory.py (99%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/run_evaluation.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/s3_utils.py (100%) rename {integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils => benchmarks/utils}/token_tracker.py (100%) diff --git a/benchmark-tests/run_evaluation.py b/benchmark-tests/run_evaluation.py deleted file mode 100644 index 3be421b19..000000000 --- a/benchmark-tests/run_evaluation.py +++ /dev/null @@ -1,291 +0,0 @@ -from boto3 import Session -from botocore.config import Config -import logging -import json -import argparse -from tqdm import tqdm -import os -import time - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - -def call_bedrock_invoke_model(prompt, bedrock, model_id, is_json_output=True, max_retries=10): - last_exception = None - for _attempt in range(max_retries): - try: - accept = 'application/json' - contentType = 'application/json' - - if 'anthropic.claude-3' in model_id: - payload_body = { - "anthropic_version": "bedrock-2023-05-31", - "max_tokens": 2048, - "temperature": 0.0, - "top_p": 1, - "top_k": 50, - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": prompt - } - ] - } - ] - } - body = json.dumps(payload_body) - else: - body = json.dumps({"prompt": prompt, "max_tokens_to_sample": 2048, "temperature": 0.0, "top_p": 1, "top_k": 50, - "stop_sequences": ["\\n\\nHuman:"]}) - - response = bedrock.invoke_model(body=body, modelId=model_id, accept=accept, contentType=contentType) - - if 'anthropic.claude-3' in model_id: - response = response['body'].read().decode('utf-8') - response = json.loads(response) - response_text = response['content'][0]['text'] - - else: - response_obj = response.get('body').read().decode('utf-8') - print("RESPONSE OBJ: " + str(response_obj)) - json_response_obj = json.loads(response_obj) - response_text = json_response_obj['completion'] - if is_json_output: - try: - start_idx = response_text.find("{") - end_idx = response_text.find("}") - parsed_completion = response_text[start_idx:end_idx + 1] - parsed_json = json.loads(parsed_completion) - parsed_json['llm_response'] = response_text - return parsed_json - except: - logger.error(response_text) - return { - 'grade': "incorrect", - 'justification': "LLM failed grading", - 'llm_response': response_text - } - else: - return response_text - except Exception as e: - last_exception = e - logger.error(str(e)) - time.sleep(3) - raise last_exception - - -BKB_CORRECTNESS_GRADING = """ -Human: -You are a teacher grading a quiz. -You are given a question, the student's answer, and the true answer, and are asked to score the student answer as either Correct or Incorrect. - -Example Format: -QUESTION: question here -STUDENT ANSWER: student's answer here -TRUE ANSWER: true answer here -GRADE: Correct or Incorrect here - -Grade the student answers based ONLY on their factual accuracy. Ignore differences in punctuation and phrasing between the student answer and true answer. It is OK if the student answer contains more information than the true answer, as long as it does not contain any conflicting statements. If the student fails to answers or claims that the search results do not mention the answer then mark as incorrect. Begin! - -QUESTION: {query} -STUDENT ANSWER: {answer} -TRUE ANSWER: {expected_answer} -GRADE: - -Your response should be in json format as follows: -{{ - "grade": (correct or incorrect), - "justification": (Without mentioning the student/teacher framing of this prompt, explain why the STUDENT ANSWER is Correct or Incorrect. Use one or two sentences maximum. Keep the answer as concise as possible.) -}} - - -Assistant: -""" - -IDK_DETECTION = """You are a teacher grading a quiz. Based on students' response, you are asked to determine if the students think they can not answer the question because some information are missing. -Response: {response} -Please output "Unanswerable" if the students identify that they can not answer the question. Otherwise, output "Answerable". -""" -_REGION = os.environ.get('AWS_REGION', os.environ.get('REGION_NAME', 'us-west-2')) - - -class GenerationEvaluator: - bedrock = Session().client( - service_name='bedrock-runtime', - region_name=_REGION, - config=Config( - max_pool_connections=50, - retries={"max_attempts": 10, "mode": "standard"}, - connect_timeout=500, - read_timeout=500, - region_name=_REGION - )) - - - def __init__(self, model_id): - self.model_id = model_id - -class CorrectnessEvaluator(GenerationEvaluator): - def __init__(self, model_id): - super().__init__(model_id) - - def evaluate(self, question, answer, response): - grading = {} - grading.update(self._llm_evaluate(question, answer, response)) - return grading - - def _llm_evaluate(self, question, answer, response): - prompt = BKB_CORRECTNESS_GRADING.format( - query=question, - answer=response, - expected_answer=answer - ) - completion = call_bedrock_invoke_model(prompt, self.bedrock, model_id=self.model_id) - if answer == "": - completion['grade'] = "incorrect" - completion['justification'] = "No answer was provided" - - if not completion or not completion['grade'] or not completion['justification']: - logger.error("Failed to grade") - logger.error(str(completion)) - return { - 'question': question, - 'llmCorrectnessGrade': "incorrect", - 'llmCorrectnessGradeJustification': "LLM failed grading", - 'llm_response': completion.get('llm_response', str(completion)) # Store the raw response - } - - try: - grading = { - 'question': question, - 'llmCorrectnessGrade': completion['grade'].lower(), - 'llmCorrectnessGradeJustification': completion['justification'].replace("\"", "\\\""), - 'llm_response': completion.get('llm_response', str(completion)) # Store the raw response - } - return grading - except Exception as e: - logger.info(str(e)) - return { - 'question': question, - 'llmCorrectnessGrade': "incorrect", - 'llmCorrectnessGradeJustification': "LLM failed grading", - 'llm_response': completion.get('llm_response', str(completion)) # Store the raw response - } - - -class IDKEvaluator(GenerationEvaluator): - def __init__(self, model_id): - super().__init__(model_id) - - def evaluate(self, question, answer, response): - grading = {} - grading.update(self._llm_evaluate(question, answer, response)) - return grading - - def _llm_evaluate(self, question, answer, response): - prompt = IDK_DETECTION.format( - question=question, - answer=answer, - response=response - ) - completion = call_bedrock_invoke_model(prompt, self.bedrock, model_id=self.model_id, is_json_output=False) - if "Unanswerable" in completion: - return { - "label": "unanswerable" - } - else: - return { - "label": "answerable" - } - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--input-file-path", type=str) - parser.add_argument("--metrics-output-path", type=str) - parser.add_argument("--eval-artifacts", type=str) - parser.add_argument("--metric", type=str, default="correctness", choices=["correctness", "idk", "correctness_on_answerable"]) - args = parser.parse_args() - - if args.metric == "correctness_on_answerable": - eval_correctness_artifact, eval_idk_artifact = args.eval_artifacts.split(",") - with open(eval_correctness_artifact) as fin: - eval_correctness_data = json.load(fin) - with open(eval_idk_artifact) as fin: - eval_idk_data = json.load(fin) - assert len(eval_correctness_data) == len(eval_idk_data) - total, count = 0, 0 - for correctness_eval, idk_eval in zip(eval_correctness_data, eval_idk_data): - if idk_eval["label"] == "answerable": - total += 1 - if correctness_eval["llmCorrectnessGrade"] == "correct": - count += 1 - logger.info("{}: {}".format(args.metric, count / total)) - else: - if args.metric == "correctness": - evaluator = CorrectnessEvaluator(model_id="anthropic.claude-3-sonnet-20240229-v1:0") - elif args.metric == "idk": - evaluator = IDKEvaluator(model_id="anthropic.claude-3-sonnet-20240229-v1:0") - - - data = [] - with open(args.input_file_path) as fin: - for line in fin: - data.append(json.loads(line)) - evaluation_outputs = [] - for example in tqdm(data): - answer = example["raw_example"]["answer"] - response = example["response"] - question = example["raw_example"]["question"] - evaluation_outputs.append(evaluator.evaluate(question, answer, response)) - count, total = 0, 0 - for evaluation in evaluation_outputs: - total += 1 - if args.metric == "correctness": - if evaluation["llmCorrectnessGrade"] == "correct": - count += 1 - if args.metric == "idk": - if evaluation["label"] == "unanswerable": - count += 1 - - logger.info("{}: {}".format(args.metric, count / total)) - os.makedirs(os.path.dirname(args.metrics_output_path), exist_ok=True) - with open(args.metrics_output_path, "w") as fout: - json.dump(evaluation_outputs, fout, indent=4) - with open(os.path.join(os.path.dirname(args.metrics_output_path), "{}.json".format(args.metric)), "w") as fout: - json.dump({ - args.metric: count / total - }, fout, indent=4) - -import multiprocessing as mp -import yaml - -class SafeCounter(): - # constructor - def __init__(self): - # initialize counter - self._counter = mp.Value('i', 0) - # initialize lock - self._lock = mp.Lock() - - # increment the counter - def increment(self): - # get the lock - with self._lock: - self._counter.value += 1 - - # get the counter value - def value(self): - with self._lock: - return self._counter.value - -def read_yaml(file_path): - if file_path: - with open(file_path, 'r') as file: - data = yaml.safe_load(file) - return data - else: - return {} \ No newline at end of file diff --git a/integration-tests/BENCHMARKING.md b/benchmarks/README.md similarity index 92% rename from integration-tests/BENCHMARKING.md rename to benchmarks/README.md index 1d3270a7b..dfde9e22d 100644 --- a/integration-tests/BENCHMARKING.md +++ b/benchmarks/README.md @@ -21,14 +21,22 @@ The benchmarking pipeline has up to four stages: ## Configuration -Create a `.env` file in the `integration-tests/` directory based on `env.template`: +Create a `.env` file based on the benchmark-specific template: ```bash +cd benchmarks cp env.template .env -# Edit .env with your bucket name, region, and model preferences +# Edit .env with your bucket name, region, and benchmark preferences ``` -See `env.template` for all available configuration options. +The benchmark deployment still uses `integration-tests/build-tests.sh`, so source the `.env` file before running: + +```bash +cd ../integration-tests +source ../benchmarks/.env +``` + +See `benchmarks/env.template` for all available benchmark configuration options. Note: benchmark variables are a subset of those in `integration-tests/env.template` — only variables relevant to benchmarking are included. ## S3 Data Layout @@ -152,7 +160,7 @@ sh build-tests.sh \ --neptune-instance-type db.r8g.2xlarge \ --notebook-instance-type ml.m5.4xlarge \ --benchmark-data-s3-uri s3://my-benchmarking-bucket/benchmark-data/ \ - --test-file benchmark.wikihow + --test-file benchmarks/datasets/benchmark.wikihow ``` **Resource requirements:** @@ -251,9 +259,8 @@ aws s3 sync s3:///wikihow-benchmark-results/ benchmark-results/wiki aws s3 sync s3:///pga-benchmark-results/ benchmark-results/pga/ --region us-west-2 # Generate comparison reports -cd integration-tests/test-scripts python -c " -from graphrag_toolkit_tests.benchmark_utils.comparison_report import generate_comparison_report +from benchmarks.utils.comparison_report import generate_comparison_report generate_comparison_report('cuad', '/path/to/benchmark-results') generate_comparison_report('concurrentqa', '/path/to/benchmark-results') generate_comparison_report('wikihow', '/path/to/benchmark-results') @@ -272,13 +279,13 @@ Use prototype mode to validate the pipeline with a small subset of data (2 docum sh build-tests.sh \ --benchmark-data-s3-uri s3://my-benchmarking-bucket/benchmark-data/ \ --benchmark-prototype \ - --test-file benchmark.cuad.prototype + --test-file benchmarks/datasets/benchmark.cuad.prototype # ConcurrentQA prototype (includes extraction) sh build-tests.sh \ --benchmark-data-s3-uri s3://my-benchmarking-bucket/benchmark-data/ \ --benchmark-prototype \ - --test-file benchmark.concurrentqa.prototype + --test-file benchmarks/datasets/benchmark.concurrentqa.prototype ``` ## Monitoring @@ -336,7 +343,7 @@ The build stage uses conservative settings to avoid OOM on large datasets: | `build_batch_size` | 10 | Documents per batch | | `build_batch_write_size` | 25 | Graph write batch size | -These are set in `benchmark_build.py` and apply to all benchmark datasets. For smaller datasets (like CUAD), these are more conservative than necessary but ensure stability. +These are set in `benchmarks/scripts/benchmark_build.py` and apply to all benchmark datasets. For smaller datasets (like CUAD), these are more conservative than necessary but ensure stability. ## Dataset Details @@ -379,7 +386,7 @@ After running all retrievers, generate `comparison_report.json` at `benchmark-re ## Adding a New Benchmark Dataset -1. Add a dataset entry to `DATASET_CONFIG` in `benchmark_build.py`: +1. Add a dataset entry to `DATASET_CONFIG` in `benchmarks/scripts/benchmark_build.py`: ```python 'my-dataset': { 'num_docs': 100, @@ -388,7 +395,7 @@ After running all retrievers, generate `comparison_report.json` at `benchmark-re } ``` -2. Add QA file mapping in `benchmark_query.py`: +2. Add QA file mapping in `benchmarks/scripts/benchmark_query.py`: ```python QA_FILE_MAP = { ... @@ -400,4 +407,4 @@ After running all retrievers, generate `comparison_report.json` at `benchmark-re 4. Create test classes for Build, Query, and Evaluate (follow the CUAD/ConcurrentQA patterns). -5. Optionally create a test file (e.g., `benchmark.my-dataset`) listing the test classes. +5. Optionally create a suite file (e.g., `benchmarks/datasets/benchmark.my-dataset`) listing the test classes. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..de658a563 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmarks package for graphrag-toolkit retriever evaluation pipeline diff --git a/benchmarks/datasets/__init__.py b/benchmarks/datasets/__init__.py new file mode 100644 index 000000000..1cd0e1430 --- /dev/null +++ b/benchmarks/datasets/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmark dataset suite configurations diff --git a/integration-tests/benchmark.concurrentqa.prototype b/benchmarks/datasets/benchmark.concurrentqa.prototype similarity index 100% rename from integration-tests/benchmark.concurrentqa.prototype rename to benchmarks/datasets/benchmark.concurrentqa.prototype diff --git a/integration-tests/benchmark.cuad.prototype b/benchmarks/datasets/benchmark.cuad.prototype similarity index 100% rename from integration-tests/benchmark.cuad.prototype rename to benchmarks/datasets/benchmark.cuad.prototype diff --git a/integration-tests/benchmark.pga b/benchmarks/datasets/benchmark.pga similarity index 100% rename from integration-tests/benchmark.pga rename to benchmarks/datasets/benchmark.pga diff --git a/integration-tests/benchmark.wikihow b/benchmarks/datasets/benchmark.wikihow similarity index 100% rename from integration-tests/benchmark.wikihow rename to benchmarks/datasets/benchmark.wikihow diff --git a/benchmarks/env.template b/benchmarks/env.template new file mode 100644 index 000000000..a8a50d8bf --- /dev/null +++ b/benchmarks/env.template @@ -0,0 +1,26 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ============================================================ +# Required Variables +# ============================================================ + +export GRAPHRAG_TOOLKIT_DIR= +export BUCKET_NAME= +export REGION_NAME= + +# ============================================================ +# Optional Variables — Benchmark Configuration +# ============================================================ + +export BENCHMARK_DATA_DIR= +export BENCHMARK_DATA_S3_URI= +export BENCHMARK_QA_LIMIT= +export BENCHMARK_IS_PROTOTYPE= + +# ============================================================ +# Optional Variables — Infrastructure +# ============================================================ + +export EXISTING_VPC_ID= +export EXISTING_SUBNET_IDS= diff --git a/benchmarks/requirements.txt b/benchmarks/requirements.txt new file mode 100644 index 000000000..38b278f17 --- /dev/null +++ b/benchmarks/requirements.txt @@ -0,0 +1,14 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmark-specific Python dependencies +# These are packages required by the benchmark pipeline that are NOT already +# listed in integration-tests/requirements-integ-test.txt. +# +# Note: graphrag-toolkit (and its transitive deps like llama-index-core, +# llama-index-llms-bedrock-converse, boto3, botocore) is assumed to be +# installed separately via `pip install -e ./lexical-graph`. + +python-dotenv>=1.0.0 +boto3>=1.28.0 +botocore>=1.31.0 diff --git a/benchmarks/scripts/__init__.py b/benchmarks/scripts/__init__.py new file mode 100644 index 000000000..c28d8dd96 --- /dev/null +++ b/benchmarks/scripts/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmark pipeline scripts diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_build.py b/benchmarks/scripts/benchmark_build.py similarity index 96% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_build.py rename to benchmarks/scripts/benchmark_build.py index a74ba8ba9..943574b9b 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_build.py +++ b/benchmarks/scripts/benchmark_build.py @@ -6,9 +6,9 @@ from typing import Dict, Any, Optional import logging -from graphrag_toolkit_tests.integration_test_base import IntegrationTestBase -from graphrag_toolkit_tests.integration_test_handler import IntegrationTestHandler -from graphrag_toolkit_tests.benchmark_utils.s3_utils import sync_benchmark_data_from_s3 +from benchmarks.scripts.integration_test_base import IntegrationTestBase +from benchmarks.scripts.integration_test_handler import IntegrationTestHandler +from benchmarks.utils.s3_utils import sync_benchmark_data_from_s3 from graphrag_toolkit.lexical_graph import LexicalGraphIndex from graphrag_toolkit.lexical_graph import GraphRAGConfig diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_evaluate.py b/benchmarks/scripts/benchmark_evaluate.py similarity index 96% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_evaluate.py rename to benchmarks/scripts/benchmark_evaluate.py index 3b2c0b37a..e27f271fd 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_evaluate.py +++ b/benchmarks/scripts/benchmark_evaluate.py @@ -17,9 +17,9 @@ import unittest from typing import Dict, Any, Optional, List -from graphrag_toolkit_tests.integration_test_base import IntegrationTestBase -from graphrag_toolkit_tests.integration_test_handler import IntegrationTestHandler -from graphrag_toolkit_tests.benchmark_utils.run_evaluation import evaluate_responses +from benchmarks.scripts.integration_test_base import IntegrationTestBase +from benchmarks.scripts.integration_test_handler import IntegrationTestHandler +from benchmarks.utils.run_evaluation import evaluate_responses logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py b/benchmarks/scripts/benchmark_extract.py similarity index 95% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py rename to benchmarks/scripts/benchmark_extract.py index 510dd0545..f767d52de 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_extract.py +++ b/benchmarks/scripts/benchmark_extract.py @@ -5,9 +5,9 @@ import unittest from typing import Dict, Any, Optional -from graphrag_toolkit_tests.integration_test_base import IntegrationTestBase -from graphrag_toolkit_tests.integration_test_handler import IntegrationTestHandler -from graphrag_toolkit_tests.benchmark_utils.s3_utils import sync_benchmark_data_from_s3 +from benchmarks.scripts.integration_test_base import IntegrationTestBase +from benchmarks.scripts.integration_test_handler import IntegrationTestHandler +from benchmarks.utils.s3_utils import sync_benchmark_data_from_s3 from graphrag_toolkit.lexical_graph import LexicalGraphIndex from graphrag_toolkit.lexical_graph import GraphRAGConfig, IndexingConfig diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_query.py b/benchmarks/scripts/benchmark_query.py similarity index 95% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_query.py rename to benchmarks/scripts/benchmark_query.py index 547c380e1..db546d156 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_query.py +++ b/benchmarks/scripts/benchmark_query.py @@ -8,14 +8,14 @@ from typing import Dict, Any, Optional, List import logging -from graphrag_toolkit_tests.integration_test_base import IntegrationTestBase -from graphrag_toolkit_tests.integration_test_handler import IntegrationTestHandler -from graphrag_toolkit_tests.benchmark_utils.s3_utils import sync_benchmark_data_from_s3 -from graphrag_toolkit_tests.benchmark_utils.retriever_factory import create_query_engine, get_retriever_config, ByoKGQueryEngineWrapper -from graphrag_toolkit_tests.benchmark_utils.token_tracker import TokenTrackingLLMCache, extract_token_usage -from graphrag_toolkit_tests.benchmark_utils.metrics_summary import compute_metrics_summary -from graphrag_toolkit_tests.benchmark_utils.hop_classifier import classify_hop -from graphrag_toolkit_tests.benchmark_utils.agentic_retriever import AgenticRetriever, AgenticQueryResult +from benchmarks.scripts.integration_test_base import IntegrationTestBase +from benchmarks.scripts.integration_test_handler import IntegrationTestHandler +from benchmarks.utils.s3_utils import sync_benchmark_data_from_s3 +from benchmarks.utils.retriever_factory import create_query_engine, get_retriever_config, ByoKGQueryEngineWrapper +from benchmarks.utils.token_tracker import TokenTrackingLLMCache, extract_token_usage +from benchmarks.utils.metrics_summary import compute_metrics_summary +from benchmarks.utils.hop_classifier import classify_hop +from benchmarks.utils.agentic_retriever import AgenticRetriever, AgenticQueryResult from graphrag_toolkit.lexical_graph import GraphRAGConfig from graphrag_toolkit.lexical_graph.storage import GraphStoreFactory, VectorStoreFactory diff --git a/benchmarks/scripts/integration_test_base.py b/benchmarks/scripts/integration_test_base.py new file mode 100644 index 000000000..647eb175c --- /dev/null +++ b/benchmarks/scripts/integration_test_base.py @@ -0,0 +1,30 @@ +# CANONICAL SOURCE: integration-tests/test-scripts/graphrag_toolkit_tests/integration_test_base.py +# This is a copy for benchmark deployment independence. If modifying shared behavior, +# update the canonical source first, then sync this copy. + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +import abc +from typing import Dict, Any +from benchmarks.scripts.integration_test_handler import IntegrationTestHandler + +class IntegrationTestBase(): + + @property + @abc.abstractmethod + def description(self): + pass + + def wait(self) -> bool: + return False + + @abc.abstractmethod + def _run_test(self, handler:IntegrationTestHandler, params:Dict[str, Any]): + pass + + def init_test_details(self, handler:IntegrationTestHandler, params:Dict[str, Any]): + handler.init_with_test_details(self.__class__.__name__, self.description, params) + + def run_test(self, handler:IntegrationTestHandler, params:Dict[str, Any]): + handler.start_test() + self._run_test(handler, params) \ No newline at end of file diff --git a/benchmarks/scripts/integration_test_handler.py b/benchmarks/scripts/integration_test_handler.py new file mode 100644 index 000000000..1d55989f3 --- /dev/null +++ b/benchmarks/scripts/integration_test_handler.py @@ -0,0 +1,223 @@ +# CANONICAL SOURCE: integration-tests/test-scripts/graphrag_toolkit_tests/integration_test_handler.py +# This is a copy for benchmark deployment independence. If modifying shared behavior, +# update the canonical source first, then sync this copy. + +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +import time +import json +import boto3 +import logging +import os +import unittest +import traceback +from typing import Any, Dict, List +from copy import deepcopy + +from graphrag_toolkit.lexical_graph import set_logging_config + +logger = logging.getLogger(__name__) + + +class IntegrationTestHandler(): + def __init__(self, test_run:int, index:int, s3_results_bucket:str, s3_results_prefix:str, results:List[Dict[str, Any]]): + self.test_run = test_run + self.index = index + self.s3_results_bucket = s3_results_bucket + self.s3_results_prefix = s3_results_prefix + self.results = results + self.skipped = False + self.props = { + 'name': None, + 'description': None, + 'init_time': 0, + 'start_time': 0, + 'end_time': 0, + 'duration_seconds': None, + 'result': 'FAIL', + 'input_params': {}, + 'exceptions': [], + 'test_results': {}, + 'output': [] + } + + def init_with_test_details(self, name:str, description:str, params:Dict[str, Any]): + self.props['name'] = name + self.props['description'] = description + self.props['input_params'] = deepcopy(params) + + file_path = f"test-logs/{self.index:02d}-{self.props['name']}.log" + if os.path.exists(file_path): + os.remove(file_path) + + set_logging_config( + 'DEBUG', + ['graphrag_toolkit'], + filename=file_path + ) + + self.props['init_time'] = int(time.time()) + + def start_test(self): + self.props['start_time'] = int(time.time()) + + def add_output(self, key:str, value:Any): + self.props['output'].append({ + 'name': key, + 'timestamp': int(time.time()), + 'value': value + }) + + def run_assertions(self, assertions:unittest.TestCase): + + def format_test_name(s): + return s.split(' ')[0] + + def format_test_case(f): + test_case = f[0] + msg = f[1] + return { + 'name': format_test_name(str(test_case)), + 'description': test_case.shortDescription(), + 'message': msg + } + + def format_param(s): + if len(s) > 100: + return f'{s[:100]} ... <{len(s) - 100} more chars>' + else: + return s + + runner = unittest.TextTestRunner() + suite = unittest.makeSuite(assertions) + + + + all_tests = { + str(test_case) : { + 'name': format_test_name(str(test_case)), + 'description': test_case.shortDescription() + } + for test_case in suite + } + + + + print(f'all_tests: {all_tests}') + + r = runner.run(suite) + + suite_param_keys = [ + k + for k in assertions.__dict__.keys() + if k.startswith('_') and k not in [ + '__module__', + 'setUpClass', + '__doc__', + '_classSetupFailed', + '_class_cleanups', + 'tearDown_exceptions' + ] + ] + + suite_params = { + k:format_param(str(v)) for k, v in assertions.__dict__.items() if k in suite_param_keys + } + + test_results = { + 'passed': r.wasSuccessful(), + 'num_tests': r.testsRun, + 'num_failures': len(r.failures), + 'num_errors': len(r.errors), + 'successes': [], + 'failures': [format_test_case(f) for f in r.failures], + 'errors': [format_test_case(e) for e in r.errors], + 'suite_params': suite_params + } + + for f in r.failures: + failure_key = str(f[0]) + print(f'failure_key: {failure_key}') + if failure_key in all_tests: + del all_tests[failure_key] + + for e in r.errors: + error_key = str(e[0]) + print(f'error_key: {error_key}') + if error_key in all_tests: + del all_tests[error_key] + + test_results['successes'] = list(all_tests.values()) + + self.props['test_results'] = test_results + + def skip(self): + self.skipped = True + + def succeeded(self): + result = 'PASS' if self.props['test_results'].get('passed', False) else 'FAIL' + result = 'SKIPPED' if self.skipped else result + self.props['result'] = result + + def add_exception(self, ex): + self.props['exceptions'].append(''.join(traceback.TracebackException.from_exception(ex).format())) + logger.exception(ex, stack_info=True) + + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, exception_traceback): + + self.props['end_time'] = int(time.time()) + if self.props['start_time'] > 0: + self.props['duration_seconds'] = self.props['end_time'] - self.props['start_time'] + else: + self.props['duration_seconds'] = self.props['end_time'] - self.props['init_time'] + + results_file_name = f"{self.index:02d}-{self.props['name']}.{self.props['result']}.json" + results_file_path = f'test-results/{results_file_name}' + s3_results_key = f'{self.s3_results_prefix}/test-runs/{self.test_run}/test-results/{results_file_name}' + + log_file_name = f"{self.index:02d}-{self.props['name']}.log" + log_file_path = f'test-logs/{log_file_name}' + s3_log_key = f'{self.s3_results_prefix}/test-runs/{self.test_run}/test-logs/{log_file_name}' + + with open(results_file_path, 'w') as out_file: + json.dump(self.props, out_file, indent=2, ensure_ascii=False) + + client = boto3.client('s3') + + client.upload_file( + results_file_path, + self.s3_results_bucket, + s3_results_key, + { + 'ServerSideEncryption':'AES256', + 'ContentType': 'application/json' + } + ) + + result = { + 'name': self.props['name'], + 'result': self.props['result'], + 'duration_seconds': self.props['duration_seconds'], + 'details': f's3://{self.s3_results_bucket}/{s3_results_key}' + } + + if os.path.exists(log_file_path): + + result['logs'] = f's3://{self.s3_results_bucket}/{s3_log_key}' + + client.upload_file( + log_file_path, + self.s3_results_bucket, + s3_log_key, + { + 'ServerSideEncryption':'AES256', + 'ContentType': 'text/plain' + } + ) + + self.results.append(result) + + \ No newline at end of file diff --git a/benchmarks/tests/__init__.py b/benchmarks/tests/__init__.py new file mode 100644 index 000000000..caa463cf7 --- /dev/null +++ b/benchmarks/tests/__init__.py @@ -0,0 +1,4 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmark unit tests diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_agentic_retriever.py b/benchmarks/tests/test_agentic_retriever.py similarity index 95% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_agentic_retriever.py rename to benchmarks/tests/test_agentic_retriever.py index 5e5086cde..4636aad98 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_agentic_retriever.py +++ b/benchmarks/tests/test_agentic_retriever.py @@ -9,7 +9,7 @@ from hypothesis import given, settings, assume from hypothesis.strategies import integers -from graphrag_toolkit_tests.benchmark_utils.agentic_retriever import _validate_max_iterations +from benchmarks.utils.agentic_retriever import _validate_max_iterations class TestAgenticIterationBoundProperty: diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_benchmark_query.py b/benchmarks/tests/test_benchmark_query.py similarity index 98% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_benchmark_query.py rename to benchmarks/tests/test_benchmark_query.py index 6ff372ad6..5b9d7d94a 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_benchmark_query.py +++ b/benchmarks/tests/test_benchmark_query.py @@ -251,12 +251,12 @@ def test_all_required_fields_present( from hypothesis.strategies import sampled_from, tuples from hypothesis import assume -from graphrag_toolkit_tests.benchmark_utils.retriever_factory import ( +from benchmarks.utils.retriever_factory import ( VALID_RETRIEVER_IDS, SUB_RETRIEVER_IDS, get_retriever_config, ) -from graphrag_toolkit_tests.benchmark_utils.metrics_summary import compute_metrics_summary +from benchmarks.utils.metrics_summary import compute_metrics_summary # Strategy for dataset names: alphanumeric + hyphens, min_size=1 diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_comparison_report.py b/benchmarks/tests/test_comparison_report.py similarity index 99% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_comparison_report.py rename to benchmarks/tests/test_comparison_report.py index d9abbb341..16bc14cc6 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_comparison_report.py +++ b/benchmarks/tests/test_comparison_report.py @@ -12,7 +12,7 @@ from hypothesis import given, settings, assume from hypothesis import strategies as st -from graphrag_toolkit_tests.benchmark_utils.comparison_report import ( +from benchmarks.utils.comparison_report import ( _compute_cost_per_query, _compute_cost_efficiency, _compute_latency_efficiency, diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_hop_classifier.py b/benchmarks/tests/test_hop_classifier.py similarity index 97% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_hop_classifier.py rename to benchmarks/tests/test_hop_classifier.py index ed9122bfe..750cee22f 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_hop_classifier.py +++ b/benchmarks/tests/test_hop_classifier.py @@ -8,7 +8,7 @@ from hypothesis import given, settings from hypothesis.strategies import text, dictionaries, one_of, integers, none, floats, just -from graphrag_toolkit_tests.benchmark_utils.hop_classifier import ( +from benchmarks.utils.hop_classifier import ( classify_hop, VALID_CLASSIFICATIONS, ) diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_metrics_summary.py b/benchmarks/tests/test_metrics_summary.py similarity index 98% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_metrics_summary.py rename to benchmarks/tests/test_metrics_summary.py index 2e9c7e1d6..6bc642761 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_metrics_summary.py +++ b/benchmarks/tests/test_metrics_summary.py @@ -11,7 +11,7 @@ from hypothesis import strategies as st from hypothesis.strategies import lists, integers, floats, fixed_dictionaries -from graphrag_toolkit_tests.benchmark_utils.metrics_summary import ( +from benchmarks.utils.metrics_summary import ( _compute_latency_stats, compute_metrics_summary, BEDROCK_PRICING, @@ -262,7 +262,7 @@ def test_cost_equals_formula(self, per_query_tokens, input_price, output_price): } with patch( - 'graphrag_toolkit_tests.benchmark_utils.metrics_summary.BEDROCK_PRICING', + 'benchmarks.utils.metrics_summary.BEDROCK_PRICING', patched_pricing, ): result = compute_metrics_summary( diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_retriever_factory.py b/benchmarks/tests/test_retriever_factory.py similarity index 96% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_retriever_factory.py rename to benchmarks/tests/test_retriever_factory.py index b67e69da2..eefb9b082 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_retriever_factory.py +++ b/benchmarks/tests/test_retriever_factory.py @@ -10,7 +10,7 @@ from hypothesis import given, settings from hypothesis.strategies import sampled_from -from graphrag_toolkit_tests.benchmark_utils.retriever_factory import ( +from benchmarks.utils.retriever_factory import ( create_query_engine, _SUB_RETRIEVER_MAP, _SUB_RETRIEVER_PROCESSOR_ARGS, @@ -66,7 +66,7 @@ def test_sub_retriever_factory_creates_correct_configuration(self, retriever_id) mock_query_engine = MagicMock() with patch( - 'graphrag_toolkit_tests.benchmark_utils.retriever_factory.LexicalGraphQueryEngine.for_traversal_based_search', + 'benchmarks.utils.retriever_factory.LexicalGraphQueryEngine.for_traversal_based_search', return_value=mock_query_engine, ) as mock_for_traversal: result = create_query_engine( @@ -144,7 +144,7 @@ def test_sub_retriever_factory_creates_correct_configuration(self, retriever_id) from hypothesis import given, settings, assume from hypothesis.strategies import text -from graphrag_toolkit_tests.benchmark_utils.retriever_factory import ( +from benchmarks.utils.retriever_factory import ( VALID_RETRIEVER_IDS, ) diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_token_tracker.py b/benchmarks/tests/test_token_tracker.py similarity index 99% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_token_tracker.py rename to benchmarks/tests/test_token_tracker.py index a568b49a3..e28cb5a39 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/test_token_tracker.py +++ b/benchmarks/tests/test_token_tracker.py @@ -12,7 +12,7 @@ import pytest -from graphrag_toolkit_tests.benchmark_utils.token_tracker import ( +from benchmarks.utils.token_tracker import ( TokenTrackingLLMCache, extract_token_usage, ) @@ -314,7 +314,7 @@ def llm_predict(prompt, **kwargs): from hypothesis import given, settings from hypothesis.strategies import integers, text -from graphrag_toolkit_tests.benchmark_utils.token_tracker import estimate_token_count +from benchmarks.utils.token_tracker import estimate_token_count class TestEstimateTokenCount: diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/__init__.py b/benchmarks/utils/__init__.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/__init__.py rename to benchmarks/utils/__init__.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/agentic_retriever.py b/benchmarks/utils/agentic_retriever.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/agentic_retriever.py rename to benchmarks/utils/agentic_retriever.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/comparison_report.py b/benchmarks/utils/comparison_report.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/comparison_report.py rename to benchmarks/utils/comparison_report.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/hop_classifier.py b/benchmarks/utils/hop_classifier.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/hop_classifier.py rename to benchmarks/utils/hop_classifier.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/metrics_summary.py b/benchmarks/utils/metrics_summary.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/metrics_summary.py rename to benchmarks/utils/metrics_summary.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/retriever_factory.py b/benchmarks/utils/retriever_factory.py similarity index 99% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/retriever_factory.py rename to benchmarks/utils/retriever_factory.py index 95de0b10b..00124cae3 100644 --- a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/retriever_factory.py +++ b/benchmarks/utils/retriever_factory.py @@ -14,7 +14,7 @@ from graphrag_toolkit.lexical_graph import LexicalGraphQueryEngine -from graphrag_toolkit_tests.benchmark_utils.agentic_retriever import AgenticRetriever +from benchmarks.utils.agentic_retriever import AgenticRetriever from graphrag_toolkit.lexical_graph.retrieval.retrievers import ( ChunkBasedSearch, ChunkBasedSemanticSearch, diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/run_evaluation.py b/benchmarks/utils/run_evaluation.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/run_evaluation.py rename to benchmarks/utils/run_evaluation.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/s3_utils.py b/benchmarks/utils/s3_utils.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/s3_utils.py rename to benchmarks/utils/s3_utils.py diff --git a/integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/token_tracker.py b/benchmarks/utils/token_tracker.py similarity index 100% rename from integration-tests/test-scripts/graphrag_toolkit_tests/benchmark_utils/token_tracker.py rename to benchmarks/utils/token_tracker.py diff --git a/integration-tests/build-tests.sh b/integration-tests/build-tests.sh index a17cca6a8..459cb59d4 100644 --- a/integration-tests/build-tests.sh +++ b/integration-tests/build-tests.sh @@ -219,13 +219,23 @@ if [[ "$TEST_FILE" ]]; then files=$(echo $TEST_FILE | tr "," "\n") for f in $files do + suite_file="$f" + if [[ ! -f "$suite_file" ]]; then + alt_path="../benchmarks/datasets/$f" + if [[ -f "$alt_path" ]]; then + suite_file="$alt_path" + else + echo "ERROR: Suite file not found: $f (searched CWD and ../benchmarks/datasets/)" + exit 1 + fi + fi while IFS= read -r line || [[ -n "$line" ]]; do if [[ "$TESTS" ]]; then TESTS+=" $line" else TESTS+="$line" fi - done < $f + done < $suite_file done fi @@ -299,6 +309,7 @@ cp -r $GRAPHRAG_TOOLKIT_DIR/examples/lexical-graph/notebooks/* lexical-graph-exa cp -r $GRAPHRAG_TOOLKIT_DIR/examples/byokg-rag/* lexical-graph-examples cp -r ./../test-scripts/* lexical-graph-examples cp -r ./../source-data lexical-graph-examples/source-data +cp -r $GRAPHRAG_TOOLKIT_DIR/benchmarks lexical-graph-examples/benchmarks # Include benchmark data if local dir is specified and no S3 URI is provided # Only copies dataset subdirectories that match the tests being run diff --git a/integration-tests/test-scripts/test_suite.py b/integration-tests/test-scripts/test_suite.py index 3b42cb11a..08b075d9b 100644 --- a/integration-tests/test-scripts/test_suite.py +++ b/integration-tests/test-scripts/test_suite.py @@ -30,8 +30,13 @@ def get_lexical_graph_version(): def to_test_class(t): parts = t.split('.') - module_name = f'graphrag_toolkit_tests.{parts[0]}' class_name = parts[1] + + if parts[0].startswith('benchmark_'): + module_name = f'benchmarks.scripts.{parts[0]}' + else: + module_name = f'graphrag_toolkit_tests.{parts[0]}' + m = importlib.import_module(module_name) c = getattr(m, class_name) return c From 91b37142933df9c7ad2c14c73266aaaed3d8bb5b Mon Sep 17 00:00:00 2001 From: Oussama Hansal Date: Mon, 20 Jul 2026 14:49:43 -0700 Subject: [PATCH 2/2] update readme --- benchmarks/README.md | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index dfde9e22d..9779a8946 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -21,22 +21,21 @@ The benchmarking pipeline has up to four stages: ## Configuration -Create a `.env` file based on the benchmark-specific template: +Benchmarks are deployed using `integration-tests/build-tests.sh`. Create a `.env` file in the `integration-tests/` directory: ```bash -cd benchmarks +cd integration-tests cp env.template .env # Edit .env with your bucket name, region, and benchmark preferences ``` -The benchmark deployment still uses `integration-tests/build-tests.sh`, so source the `.env` file before running: - -```bash -cd ../integration-tests -source ../benchmarks/.env -``` +The `benchmarks/env.template` documents the benchmark-specific variables (a subset of the full `integration-tests/env.template`). Key benchmark variables: -See `benchmarks/env.template` for all available benchmark configuration options. Note: benchmark variables are a subset of those in `integration-tests/env.template` — only variables relevant to benchmarking are included. +- `BENCHMARK_DATA_S3_URI` — S3 URI containing your benchmark datasets +- `BENCHMARK_DATA_DIR` — Local directory with benchmark data (alternative to S3) +- `BENCHMARK_QA_LIMIT` — Limit QA pairs for quick prototype runs +- `BENCHMARK_IS_PROTOTYPE` — Use prototype (small) datasets +- `EXISTING_VPC_ID` / `EXISTING_SUBNET_IDS` — Reuse an existing VPC to avoid quota limits ## S3 Data Layout