abstract runner logic out + create stubs for new DB's

Signed-off-by: Michael Oviedo <mikeovi@amazon.com>

Author: OVI3D0

benchmarks/worker_coordinator/parsing_test.py

@@ -30,7 +30,7 @@
 import pytest
 import ujson
 
-from osbenchmark.worker_coordinator import runner
+from osbenchmark.worker_coordinator import runners as runner
 
 @pytest.mark.benchmark(
     group="parse",

benchmarks/worker_coordinator/runner_test.py

@@ -24,7 +24,7 @@
 
 import pytest
 
-from osbenchmark.worker_coordinator import runner
+from osbenchmark.worker_coordinator import runners as runner
 
 bulk_index = runner.BulkIndex()
 

osbenchmark/worker_coordinator/runners/__init__.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+# Modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+# Licensed to Elasticsearch B.V. under one or more contributor
+# license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright
+# ownership. Elasticsearch B.V. licenses this file to you under
+# the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#	http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Runner package for OpenSearch Benchmark.
+
+This package contains:
+- base.py: Database-agnostic runner infrastructure (base classes, utilities)
+- opensearch.py: OpenSearch-specific runner implementations
+- vespa.py: Vespa runner implementations (stub for future PR)
+
+The registry functions (register_runner, runner_for, etc.) live here in __init__.py
+and are the central coordination point for all runner registration.
+"""
+
+import logging
+import types
+
+from osbenchmark import exceptions, workload
+from osbenchmark.worker_coordinator.runners.base import (
+    Runner,
+    Delegator,
+    time_func,
+    request_context_holder,
+    mandatory,
+    escape,
+    remove_prefix,
+    unwrap,
+    _single_cluster_runner,
+    _multi_cluster_runner,
+    _with_assertions,
+    _with_completion,
+    MultiClientRunner,
+    AssertingRunner,
+    NoCompletion,
+    WithCompletion,
+)
+
+__RUNNERS = {}
+
+
+def register_runner(operation_type, runner, **kwargs):
+    logger = logging.getLogger(__name__)
+    async_runner = kwargs.get("async_runner", False)
+    if isinstance(operation_type, workload.OperationType):
+        operation_type = operation_type.to_hyphenated_string()
+
+    if not async_runner:
+        raise exceptions.BenchmarkAssertionError(
+            "Runner [{}] must be implemented as async runner and registered with async_runner=True.".format(str(runner)))
+
+    if getattr(runner, "multi_cluster", False):
+        if "__aenter__" in dir(runner) and "__aexit__" in dir(runner):
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug("Registering runner object [%s] for [%s].", str(runner), str(operation_type))
+            cluster_aware_runner = _multi_cluster_runner(runner, str(runner), context_manager_enabled=True)
+        else:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug("Registering context-manager capable runner object [%s] for [%s].", str(runner), str(operation_type))
+            cluster_aware_runner = _multi_cluster_runner(runner, str(runner))
+    # we'd rather use callable() but this will erroneously also classify a class as callable...
+    elif isinstance(runner, types.FunctionType):
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug("Registering runner function [%s] for [%s].", str(runner), str(operation_type))
+        cluster_aware_runner = _single_cluster_runner(runner, runner.__name__)
+    elif "__aenter__" in dir(runner) and "__aexit__" in dir(runner):
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug("Registering context-manager capable runner object [%s] for [%s].", str(runner), str(operation_type))
+        cluster_aware_runner = _single_cluster_runner(runner, str(runner), context_manager_enabled=True)
+    else:
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug("Registering runner object [%s] for [%s].", str(runner), str(operation_type))
+        cluster_aware_runner = _single_cluster_runner(runner, str(runner))
+
+    __RUNNERS[operation_type] = _with_completion(_with_assertions(cluster_aware_runner))
+
+
+def runner_for(operation_type):
+    try:
+        return __RUNNERS[operation_type]
+    except KeyError:
+        raise exceptions.BenchmarkError("No runner available for operation type [%s]" % operation_type)
+
+
+def enable_assertions(enabled):
+    """
+    Changes whether assertions are enabled. The status changes for all tasks that are executed after this call.
+
+    :param enabled: ``True`` to enable assertions, ``False`` to disable them.
+    """
+    AssertingRunner.assertions_enabled = enabled
+
+
+# Only intended for unit-testing!
+def remove_runner(operation_type):
+    del __RUNNERS[operation_type]
+
+
+# Re-export OpenSearch runners and register_default_runners for convenience
+from osbenchmark.worker_coordinator.runners.opensearch import *  # noqa: F401,F403,E402
+from osbenchmark.worker_coordinator.runners.opensearch import register_default_runners  # noqa: E402