datastore/ci/jobs/integration_test_job.py

import argparse
import os
import re
import shlex
import subprocess
import time
from pathlib import Path
from typing import List, Optional, Tuple

from ci.jobs.scripts.bugfix_validation import bugfix_build_types, find_master_builds
from ci.jobs.scripts.find_tests import Targeting
from ci.jobs.scripts.integration_tests_configs import (
    IMAGES_ENV,
    LLVM_COVERAGE_SKIP_PREFIXES,
    force_heavy_modules_sequential,
    get_optimal_test_batch,
)
from ci.jobs.scripts.workflow_hooks.pr_labels_and_category import Labels
from ci.praktika.info import Info
from ci.praktika.result import Result
from ci.praktika.utils import Shell, Utils

repo_dir = Utils.cwd()
temp_path = f"{repo_dir}/ci/tmp"


MAX_FAILS_BEFORE_DROP = 5
# Flaky-check best-effort scope cap: the maximum number of changed test modules a single
# flaky-check run will execute. A PR can mechanically touch a large number of integration
# test modules (e.g. a repo-wide lint/format change), and running every one of them
# repeatedly under `--dist=each` cannot fit the flaky-check time budget - the job would be
# hard-killed by the external CI timeout before producing any report. When the cap is
# exceeded the extra modules are skipped (best effort) and the selected ones get full
# flakiness coverage instead of a truncated run. See FLAKY_CHECK_TIME_LIMIT for the hard
# time guarantee that backstops this.
MAX_FLAKY_CHECK_MODULES = 10
OOM_IN_DMESG_TEST_NAME = "OOM in dmesg"
ncpu = Utils.cpu_count()
mem_gb = round(Utils.physical_memory() // (1024**3), 1)

MAX_CPUS_PER_WORKER = 5
MAX_MEM_PER_WORKER = 11
# Flaky/targeted checks run with --dist=each, so every worker runs the full set
# of changed modules concurrently (each with its own Docker cluster) instead of
# splitting modules across workers. A worker's peak footprint is therefore much
# larger, so it needs a bigger memory budget to avoid exhausting the container
# cgroup and tripping the kernel OOM killer (see OOM_IN_DMESG_TEST_NAME).
MAX_MEM_PER_WORKER_DIST_EACH = 20

INFRASTRUCTURE_ERROR_PATTERNS = [
    "timed out after",
    "TimeoutExpired",
    "Cannot connect to the Docker daemon",
    "Error response from daemon",
    "Name or service not known",
    "Temporary failure in name resolution",
    "Network is unreachable",
    "Connection reset by peer",
    "No space left on device",
    "Cannot allocate memory",
    "OCI runtime create failed",
    "toomanyrequests",
    "pull access denied",
    "Got exception pulling images:",  # docker pull failure during cluster.start()
]


def _is_infrastructure_error(result: Result) -> bool:
    """Returns True if the result is a failure caused by infrastructure issues."""
    if not result.info:
        return False
    if result.status == Result.Status.ERROR:
        return any(pattern in result.info for pattern in INFRASTRUCTURE_ERROR_PATTERNS)
    # Docker compose/pull infrastructure failures may appear with FAIL status
    # when pytest reports fixture (setup phase) errors as test failures.
    # Require both docker context and an infrastructure pattern to avoid
    # false positives on genuine test failures.
    if result.status == Result.Status.FAIL:
        has_docker_context = (
            "'docker'" in result.info or "images_pull_cmd" in result.info
        )
        return has_docker_context and any(
            p in result.info for p in INFRASTRUCTURE_ERROR_PATTERNS
        )
    return False


def _mark_infrastructure_errors(results: list) -> int:
    """Scan results, label infrastructure errors with INFRA and change their status to SKIPPED.

    Returns the number of results that were relabeled.
    """
    count = 0
    for r in results:
        if _is_infrastructure_error(r):
            r.set_label(Result.Label.INFRA)
            r.status = Result.Status.SKIPPED
            count += 1
    if count:
        print(f"Marked {count} test result(s) as infrastructure errors")
    return count


def quote_tests(tests: List[str]) -> str:
    """Join test node IDs into a shell-safe, space-separated string.

    A parametrized integration test node ID can contain spaces, parentheses and
    quotes when the test is parametrized with SQL (e.g.
    `test.py::test_simple_append[SELECT now() FROM numbers(2)]`). The pytest
    command is executed through a shell, so each node ID must be quoted to
    survive as a single argument instead of being split or mis-parsed.
    """
    return " ".join(shlex.quote(t) for t in tests)


def start_docker_in_docker():
    with open("./ci/tmp/docker-in-docker.log", "w") as log_file:
        dockerd_proc = subprocess.Popen(
            "./ci/jobs/scripts/docker_in_docker.sh",
            stdout=log_file,
            stderr=subprocess.STDOUT,
        )
    retries = 20
    for i in range(retries):
        # On last retry, show errors; otherwise suppress them
        cmd = "docker info > /dev/null" if i == retries - 1 else "docker info > /dev/null 2>&1"
        if Shell.check(cmd, verbose=True):
            break
        if i == retries - 1:
            raise RuntimeError(
                f"Docker daemon didn't responded after {retries} attempts"
            )
        time.sleep(2)
    print(f"Started docker-in-docker asynchronously with PID {dockerd_proc.pid}")


_COMPOSE_DIR = Path("./tests/integration/compose")

# Explicit mapping for with_* flags whose compose file name cannot be derived
# by simply prepending "docker_compose_" and appending ".yml".
_WITH_FLAG_TO_COMPOSE: dict[str, List[str]] = {
    "mysql57": ["docker_compose_mysql.yml"],
    "mysql8": ["docker_compose_mysql_8_0.yml"],
    "dremio26": ["docker_compose_dremio_26_0.yml"],
    "kerberos_kdc": ["docker_compose_kerberos_kdc.yml"],
    # with_iceberg_catalog can use any of the iceberg catalogs; include them all
    "iceberg_catalog": [
        "docker_compose_iceberg_rest_catalog.yml",
        "docker_compose_iceberg_hms_catalog.yml",
        "docker_compose_iceberg_lakekeeper_catalog.yml",
        "docker_compose_iceberg_nessie_catalog.yml",
    ],
    "hms_catalog": ["docker_compose_iceberg_hms_catalog.yml"],
    "glue_catalog": ["docker_compose_glue_catalog.yml"],
    "prometheus_writer": ["docker_compose_prometheus_writer.yml"],
    "prometheus_reader": ["docker_compose_prometheus_reader.yml"],
    "prometheus_receiver": ["docker_compose_prometheus_receiver.yml"],
    # with_odbc_drivers implicitly sets up mysql8 + postgres
    "odbc_drivers": ["docker_compose_mysql_8_0.yml", "docker_compose_postgres.yml"],
    # Flags with no separate compose file of their own
    "jdbc_bridge": [],
    "net_trics": [],
}


def get_compose_files_for_test_modules(test_modules: List[str]) -> List[Path]:
    """Return compose files needed by the given test modules.

    Grep every Python source file in each test suite directory for:
    - `with_X=True` patterns (mapped via `_WITH_FLAG_TO_COMPOSE` or the obvious
      `docker_compose_{X}.yml` naming convention), and
    - explicit `docker_compose_*.yml` file name strings (used e.g. via
      `extra_parameters={"docker_compose_file_name": "..."}` calls).
    """
    needed: set[Path] = set()
    suite_dirs = {m.split("/")[0] for m in test_modules}

    for suite_dir in suite_dirs:
        suite_path = Path("./tests/integration/") / suite_dir
        if not suite_path.is_dir():
            continue
        for py_file in suite_path.glob("**/*.py"):
            try:
                content = py_file.read_text(errors="replace")
            except OSError:
                continue

            # 1. with_X=True → compose file via mapping or naming convention
            for m in re.finditer(r"\bwith_(\w+)\s*=\s*True", content):
                flag = m.group(1)
                if flag in _WITH_FLAG_TO_COMPOSE:
                    for fname in _WITH_FLAG_TO_COMPOSE[flag]:
                        p = _COMPOSE_DIR / fname
                        if p.exists():
                            needed.add(p)
                else:
                    p = _COMPOSE_DIR / f"docker_compose_{flag}.yml"
                    if p.exists():
                        needed.add(p)

            # 2. Directly named compose files (e.g. in extra_parameters dicts)
            for m in re.finditer(r"(docker_compose_\w+\.yml)", content):
                p = _COMPOSE_DIR / m.group(1)
                if p.exists():
                    needed.add(p)

    return sorted(needed)


def get_images_from_compose_files(compose_files: List[Path]) -> List[str]:
    """Parse compose files and return a deduplicated list of image references.

    Environment variable placeholders like `${DOCKER_NGINX_DAV_TAG:-latest}` are
    resolved from `os.environ`.  For clickhouse images that appear without a tag
    (e.g. `clickhouse/integration-test`) the tag is looked up from `IMAGES_ENV`.
    Images with still-unresolvable variables are silently skipped.
    """
    known_image_tags: dict[str, str] = {}
    for image_name, env_var in IMAGES_ENV.items():
        tag = os.environ.get(env_var)
        if tag:
            known_image_tags[image_name] = tag

    def resolve_image(raw: str) -> Optional[str]:
        def replace_var(m: re.Match) -> str:
            var_name = m.group(1)
            default = m.group(2) if m.group(2) is not None else "latest"
            return os.environ.get(var_name, default)

        resolved = re.sub(r"\$\{(\w+)(?::-([^}]*))?\}", replace_var, raw)
        if "${" in resolved:
            return None  # Still-unresolvable variable — skip
        # Append the correct tag for tagless known clickhouse images
        if ":" not in resolved and resolved in known_image_tags:
            resolved = f"{resolved}:{known_image_tags[resolved]}"
        return resolved

    images: set[str] = set()
    for compose_file in compose_files:
        try:
            content = compose_file.read_text()
        except OSError:
            continue
        for m in re.finditer(r"^\s+image:\s+(.+)$", content, re.MULTILINE):
            # Strip inline YAML comments from unquoted values before resolving
            # (e.g. `coredns/coredns:1.9.3 # :latest broke this test`).
            raw = re.sub(r"\s+#.*$", "", m.group(1).strip())
            resolved = resolve_image(raw)
            if resolved:
                images.add(resolved)

    return sorted(images)


def prefetch_images(
    images: List[str], retries: int = 3, pull_timeout: int = 300
) -> bool:
    """Pull every image in parallel using `ci/prefetch-integration-test-images`.

    Images with no manifest for the current architecture (e.g. amd64-only images
    on arm64 runners) are silently skipped.  Returns True on success, False if any
    image fails to pull for a real reason.
    """
    if not images:
        print("No images to pre-fetch.")
        return True

    script = f"{repo_dir}/ci/jobs/scripts/prefetch-integration-test-images"
    env = {
        **os.environ,
        "PULL_RETRIES": str(retries),
        "PULL_TIMEOUT": str(pull_timeout),
    }
    return Shell.check(
        f"{script} {' '.join(images)}",
        verbose=True,
        env=env,
    )


def parse_args():
    parser = argparse.ArgumentParser(description="Datastore Build Job")
    parser.add_argument("--options", help="Job parameters: ...")
    parser.add_argument(
        "--test",
        help="Optional. Test name patterns (space-separated)",
        default=[],
        nargs="+",
        action="extend",
    )
    parser.add_argument(
        "--count",
        help="Optional. Number of times to repeat each test",
        default=None,
        type=int,
    )
    parser.add_argument(
        "--debug",
        help="Optional. Open python debug console on exception",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "--path",
        help="Optional. Path to custom clickhouse binary",
        type=str,
        default="",
    )
    parser.add_argument(
        "--path_1",
        help="Optional. Path to custom server config",
        type=str,
        default="",
    )
    parser.add_argument(
        "--workers",
        help="Optional. Number of parallel workers for pytest",
        default=None,
        type=int,
    )
    parser.add_argument(
        "--session-timeout",
        help="Optional. Session timeout in seconds",
        default=None,
        type=int,
    )
    parser.add_argument(
        "--param",
        help=(
            "Optional. Comma-separated KEY=VALUE pairs to inject as environment "
            "variables for pytest (e.g. --param PYTEST_ADDOPTS=-vv,CUSTOM_FLAG=1)"
        ),
        type=str,
        default="",
    )
    return parser.parse_args()


def merge_profraw_files(llvm_profdata_cmd: str, job_params: list):
    """Merge all profraw files into final profdata file.

    Args:
        llvm_profdata_cmd: Path to llvm-profdata tool
        job_params: List of job parameters for naming output file
    """
    import subprocess
    from pathlib import Path

    # Find all profraw files
    profraw_files = [str(p) for p in Path(".").rglob("*.profraw")]

    if not profraw_files:
        print("No profraw files found", flush=True)
        return

    joined_job_params = "_".join(job_params) if job_params else "all"
    joined_job_params = joined_job_params.replace(" ", "_").replace("/", "_")
    final_file = f"./it-{joined_job_params}.profdata"
    print(f"Merging {len(profraw_files)} profraw files into {final_file}", flush=True)

    result = subprocess.run(
        [llvm_profdata_cmd, "merge", "-sparse", "-failure-mode=warn"]
        + profraw_files
        + ["-o", final_file],
        capture_output=True,
        text=True,
    )

    # Check for corrupted files in stderr
    corrupted_count = result.stderr.count(
        "invalid instrumentation profile"
    ) + result.stderr.count("file header is corrupt")
    if corrupted_count > 0:
        print(f"  WARNING: Found {corrupted_count} corrupted profraw files", flush=True)
        # Extract and display corrupted filenames from stderr
        corrupted_files = set()
        for line in result.stderr.split("\n"):
            if (
                "invalid instrumentation profile" in line
                or "file header is corrupt" in line
                or "error:" in line.lower()
            ):
                print(f"    {line.strip()}", flush=True)
                # Extract filename from error message (format: "error: file.profraw: ..." or "warning: file.profraw: ...")
                parts = line.split(":")
                if len(parts) >= 2:
                    potential_file = parts[1].strip()
                    if potential_file.endswith(".profraw"):
                        corrupted_files.add(potential_file)
        if corrupted_files:
            print(f"  Corrupted files: {', '.join(corrupted_files)}", flush=True)

    if result.returncode == 0:
        print(f"Successfully created final coverage file: {final_file}", flush=True)

        # Delete merged profraw files to save disk space
        deleted_count = 0
        for profraw_file in profraw_files:
            try:
                Path(profraw_file).unlink()
                deleted_count += 1
            except Exception as e:
                print(f"  WARNING: Failed to delete {profraw_file}: {e}", flush=True)
        print(f"  Deleted {deleted_count} profraw files", flush=True)
        return final_file
    else:
        print("ERROR: Failed to create final coverage file", flush=True)
        if result.stderr:
            print(result.stderr, flush=True)
        return None




def get_parallel_sequential_tests_to_run(
    batch_num: int,
    total_batches: int,
    args_test: List[str],
    workers: int,
    job_options: str,
    info: Info,
    no_strict: bool = False,
) -> Tuple[List[str], List[str]]:
    if args_test:
        batch_num = 1
        total_batches = 1

    test_files = [
        str(p.relative_to("./tests/integration/"))
        for p in Path("./tests/integration/").glob("test_*/test*.py")
    ]

    if "amd_llvm_coverage" in (job_options or ""):
        before = len(test_files)
        test_files = [
            f
            for f in test_files
            if not any(f.startswith(prefix) for prefix in LLVM_COVERAGE_SKIP_PREFIXES)
        ]
        print(
            f"LLVM coverage: skipped {before - len(test_files)} test files matching LLVM_COVERAGE_SKIP_PREFIXES"
        )

    assert len(test_files) > 100

    parallel_test_modules, sequential_test_modules = get_optimal_test_batch(
        test_files, total_batches, batch_num, workers, job_options, info
    )

    if "excluded_from_llvm" in (job_options or ""):
        excluded_from_llvm_set = {
            f
            for f in (parallel_test_modules + sequential_test_modules)
            if any(f.startswith(prefix) for prefix in LLVM_COVERAGE_SKIP_PREFIXES)
            or "is_built_with_llvm_coverage" in Path(f"./tests/integration/{f}").read_text()
        }
        parallel_test_modules = [f for f in parallel_test_modules if f in excluded_from_llvm_set]
        sequential_test_modules = [f for f in sequential_test_modules if f in excluded_from_llvm_set]
        print(
            f"LLVM coverage disabled-only: kept {len(parallel_test_modules)} parallel and "
            f"{len(sequential_test_modules)} sequential test files "
            f"(from LLVM_COVERAGE_SKIP_PREFIXES or containing is_built_with_llvm_coverage)"
        )

    if not args_test:
        return parallel_test_modules, sequential_test_modules

    # there are following possible values for args.test:
    # 1) test suit (e.g. test_directory or test_directory/)
    # 2) test module (e.g. test_directory/test_module or test_directory/test_module.py)
    # 3) test case (e.g. test_directory/test_module.py::test_case or test_directory/test_module::test_case[test_param])
    def normalize_test_path(test_arg: str) -> str:
        """Normalize test path by removing integration test directory prefixes."""
        # Handle: tests/integration/, integration/, ./tests/integration/, or full paths
        if "tests/integration/" in test_arg:
            # Extract everything after tests/integration/
            test_arg = test_arg.split("tests/integration/", 1)[1]
        elif test_arg.startswith("integration/"):
            # Handle integration/ prefix
            test_arg = test_arg[len("integration/"):]
        return test_arg

    def test_match(test_file: str, test_arg: str) -> bool:
        if "/" not in test_arg:
            return f"{test_arg}/" in test_file
        if test_arg.endswith(".py"):
            return test_file == test_arg
        parts = test_arg.split("::", maxsplit=1)
        test_module = parts[0]
        if test_file.removesuffix(".py") != test_module.removesuffix(".py"):
            return False
        # When a specific test function is requested, verify it exists in the
        # file.  Targeted CI runs pull test names from CIDB, but the test may
        # have been moved or removed since the record was written.  Passing a
        # stale nodeID to pytest causes the entire collection to fail with
        # exit-code 5 ("no tests collected"), aborting all other tests too.
        if len(parts) > 1:
            test_func = parts[1].split("[")[0]  # strip parametrization
            file_path = Path("./tests/integration/") / test_file
            try:
                content = file_path.read_text()
                if f"def {test_func}(" not in content:
                    print(
                        f"WARNING: test function '{test_func}' not found in {test_file}, skipping stale target"
                    )
                    return False
            except OSError:
                return False
        return True

    parallel_tests = []
    sequential_tests = []
    for test_arg in args_test:
        # Normalize the test path first
        normalized_test_arg = normalize_test_path(test_arg)
        matched = False
        for test_file in parallel_test_modules:
            if test_match(test_file, normalized_test_arg):
                parallel_tests.append(normalized_test_arg)
                matched = True
        for test_file in sequential_test_modules:
            if test_match(test_file, normalized_test_arg):
                sequential_tests.append(normalized_test_arg)
                matched = True
        if not no_strict:
            assert matched, f"Test [{test_arg}] not found"

    return parallel_tests, sequential_tests


def tail(filepath: str, buff_len: int = 1024) -> List[str]:
    with open(filepath, "rb") as f:
        # Get file size to avoid seeking before start of file
        f.seek(0, os.SEEK_END)
        file_size = f.tell()

        if file_size <= buff_len:
            # File is smaller than buffer, read from beginning
            f.seek(0)
        else:
            # File is larger, seek from end
            f.seek(-buff_len, os.SEEK_END)
            f.readline()  # Skip partial line

        data = f.read()
        return data.decode(errors="replace")


def run_pytest_and_collect_results(
    command: str, env: str, report_name: str, timeout: int = None
) -> Tuple[Result, bool]:
    """
    Runs a pytest command and reports whether the run was cut short by a timeout.

    Returns `(test_result, timed_out)`. `timed_out` is True when the run did not finish
    on its own but was stopped by either:
      - the graceful xdist `--session-timeout` (pytest interrupts itself and writes the
        `xdist.dsession.Interrupted: session-timeout:` marker to the log), or
      - the hard subprocess backstop (`Shell.run` `SIGTERM`s/`SIGKILL`s pytest after
        `timeout` seconds, e.g. when a test hangs past the session-timeout).

    Callers use `timed_out` to tell an empty result set caused by expected time-budget
    exhaustion (best effort, may be downgraded to `SKIPPED`) apart from one caused by a
    real pytest/harness failure (no timeout - must stay `ERROR`).
    """

    run_sw = Utils.Stopwatch()

    test_result = Result.from_pytest_run(
        command=command,
        env=env,
        cwd="./tests/integration/",
        pytest_report_file=f"{temp_path}/pytest_{report_name}.jsonl",
        pytest_logfile=f"{temp_path}/pytest_{report_name}.log",
        logfile=f"{temp_path}/{report_name}.log",
        timeout=timeout,
    )

    timed_out = False
    if "!!!!!!! xdist.dsession.Interrupted: session-timeout:" in tail(
        f"{temp_path}/{report_name}.log"
    ):
        timed_out = True
        test_result.info = "ERROR: session-timeout occurred during test execution"
        assert test_result.status == Result.Status.ERROR
        test_result.results.append(
            Result(
                name="Timeout",
                status=Result.Status.FAIL,
                info=test_result.info,
            )
        )
    elif timeout is not None and run_sw.duration >= timeout:
        # The graceful session-timeout marker is absent but the run still reached the
        # hard subprocess `timeout`: `Shell.run` killed pytest before it could finish
        # (a normal run returns well under `timeout`). Treat this as a timeout so an
        # empty result is reported as best-effort rather than as a harness failure.
        timed_out = True

    return test_result, timed_out


def is_empty_best_effort_skip(
    is_flaky_check: bool,
    is_targeted_check: bool,
    has_results: bool,
    timed_out: bool,
) -> bool:
    """
    Decide whether a flaky/targeted check that produced no test results should be
    reported as `SKIPPED` (best effort) instead of falling through to `create_from`'s
    default `ERROR`.

    Only the expected timeout path is downgraded: a flaky/targeted run whose time budget
    was exhausted (graceful xdist session-timeout or the hard subprocess backstop) before
    any test produced a result is best-effort `SKIPPED`. When no timeout was observed, an
    empty result means pytest failed to produce any output for some other reason (it
    crashed before writing the jsonl report, or exited with a plugin/internal error and no
    test rows). That is a real harness failure and must stay `ERROR`, so this returns
    False.
    """
    return (is_flaky_check or is_targeted_check) and not has_results and timed_out


def main():
    sw = Utils.Stopwatch()
    info = Info()
    args = parse_args()
    job_params = args.options.split(",") if args.options else []
    job_params = [to.strip() for to in job_params]
    use_old_analyzer = False
    use_distributed_plan = False
    use_database_disk = False
    is_flaky_check = False
    is_bugfix_validation = False
    is_parallel = False
    is_sequential = False
    is_targeted_check = False
    is_llvm_coverage = False
    llvm_profdata_cmd = None

    # Set on_error_hook to collect logs on hard timeout
    Result.from_fs(info.job_name).set_on_error_hook(
        """
dmesg -T >./ci/tmp/dmesg.log
sudo chown -R $(id -u):$(id -g) ./tests/integration
tar -czf ./ci/tmp/logs.tar.gz \
  ./tests/integration/test_*/_instances*/ \
  ./ci/tmp/*.log \
  ./ci/tmp/*.jsonl || :
"""
    ).set_files(
        [
            "./ci/tmp/logs.tar.gz",
            "./ci/tmp/dmesg.log",
            "./ci/tmp/docker-in-docker.log",
        ],
        strict=False,
    )

    if args.param:
        for item in args.param.split(","):
            print(f"Setting env variable: {item}")
            key, _, value = item.partition("=")
            key = key.strip()
            if not key:
                continue
            os.environ[key] = value.strip()

    java_path = Shell.get_output(
        "update-alternatives --config java | sed -n 's/.*(providing \/usr\/bin\/java): //p'",
        verbose=True,
    )
    repeat_option = ""
    if "bugfix" in info.job_name.lower():
        is_bugfix_validation = True

    batch_num, total_batches = 1, 1
    for to in job_params:
        if "/" in to:
            batch_num, total_batches = map(int, to.split("/"))
        elif any(build in to for build in ("amd_", "arm_")):
            if "amd_llvm_coverage" in to:
                is_llvm_coverage = True
        elif to == "old analyzer":
            use_old_analyzer = True
        elif to == "distributed plan":
            use_distributed_plan = True
        elif to == "db disk":
            use_database_disk = True
        elif to == "flaky":
            is_flaky_check = True
        elif to == "parallel":
            is_parallel = True
        elif to == "sequential":
            is_sequential = True
        elif "bugfix" in to.lower() or "validation" in to.lower():
            is_bugfix_validation = True
        elif "targeted" in to:
            is_targeted_check = True
        else:
            assert False, f"Unknown job option [{to}]"

    if args.count:
        repeat_option = f"--count {args.count} --random-order"
    # For flaky/targeted checks, --count is not used. Instead, --dist=each runs N workers
    # each executing all modules independently with their own isolated Docker cluster
    # (ClickHouseCluster appends PYTEST_XDIST_WORKER to project_name for isolation).

    if args.workers:
        workers = args.workers
    else:
        # --dist=each (flaky/targeted checks) makes every worker run all modules,
        # so budget more memory per worker than the module-splitting loadfile runs.
        mem_per_worker = (
            MAX_MEM_PER_WORKER_DIST_EACH
            if is_flaky_check or is_targeted_check
            else MAX_MEM_PER_WORKER
        )
        print("ncpu:", ncpu)
        print("mem_gb:", mem_gb)
        workers = min(ncpu // MAX_CPUS_PER_WORKER, mem_gb // mem_per_worker) or 1

    clickhouse_path = f"{Utils.cwd()}/ci/tmp/clickhouse"
    clickhouse_server_config_dir = f"{Utils.cwd()}/programs/server"
    if info.is_local_run:
        if args.path:
            clickhouse_path = args.path
        else:
            paths_to_check = [
                clickhouse_path,  # it's set for CI runs, but we need to check it
                f"{Utils.cwd()}/build/programs/clickhouse",
                f"{Utils.cwd()}/clickhouse",
            ]
            for path in paths_to_check:
                if Path(path).is_file():
                    clickhouse_path = path
                    break
            else:
                raise FileNotFoundError(
                    "Datastore binary not found in any of the paths: "
                    + ", ".join(paths_to_check)
                    + ". You can also specify path to binary via --path argument"
                )
        if args.path_1:
            clickhouse_server_config_dir = args.path_1
    assert Path(
        clickhouse_server_config_dir
    ), f"Datastore config dir does not exist [{clickhouse_server_config_dir}]"
    print(f"Using Datastore binary at [{clickhouse_path}]")

    changed_test_modules = []
    if is_bugfix_validation or is_flaky_check or is_targeted_check:
        if info.is_local_run:
            assert (
                args.test
            ), "--test must be provided for flaky or bugfix job flavor with local run"
        else:
            if is_bugfix_validation and Labels.PR_BUGFIX not in info.pr_labels and Labels.PR_CRITICAL_BUGFIX not in info.pr_labels:
                # Not a bugfix PR - run a simple sanity test
                changed_test_modules = ["test_accept_invalid_certificate/test.py"]
            else:
                # TODO: reduce scope to modified test cases instead of entire modules
                changed_files = info.get_changed_files()
                for file in changed_files:
                    if Targeting.is_integration_test_file(file):
                        changed_test_modules.append(
                            file.removeprefix("tests/integration/")
                        )
                if not changed_test_modules and Labels.CI_FORCE_ALL in info.pr_labels:
                    print(
                        f"NOTE: No changed test modules found, but '{Labels.CI_FORCE_ALL}' label forces run - using sanity test"
                    )
                    changed_test_modules = ["test_accept_invalid_certificate/test.py"]

    # Best-effort scope cap for the flaky check (see MAX_FLAKY_CHECK_MODULES). When a PR
    # touches more changed test modules than can be repeatedly run within the time budget,
    # run a deterministic subset (sorted for reproducibility) and skip the rest rather than
    # truncating the whole run. The remaining ones are reported below as skipped so the
    # reduced coverage is explicit, not silent.
    skipped_flaky_modules = []
    if is_flaky_check and len(changed_test_modules) > MAX_FLAKY_CHECK_MODULES:
        changed_test_modules = sorted(changed_test_modules)
        skipped_flaky_modules = changed_test_modules[MAX_FLAKY_CHECK_MODULES:]
        changed_test_modules = changed_test_modules[:MAX_FLAKY_CHECK_MODULES]
        print(
            f"Flaky check: best-effort scope cap - running {len(changed_test_modules)} of "
            f"{len(changed_test_modules) + len(skipped_flaky_modules)} changed modules "
            f"to fit the time budget.\n"
            f"  Running: {changed_test_modules}\n"
            f"  Skipped (best effort): {skipped_flaky_modules}"
        )

    if is_bugfix_validation:
        # Download the master-HEAD binaries matching this job's runner arch:
        # the aarch64 job runs on an ARM runner and must use the ARM builds.
        build_types = bugfix_build_types(info.job_name)
        bt_paths = {bt: f"{temp_path}/clickhouse_{bt}" for bt in build_types}
        # In local runs, only reuse existing binaries; probing master commits in S3
        # depends on `master_commits` workflow data populated by CI workflow hooks
        # and is not available locally.
        if info.is_local_run:
            missing = [str(p) for p in bt_paths.values() if not Path(p).is_file()]
            assert not missing, (
                "Local bugfix validation requires all build-type binaries to be "
                f"present under {temp_path}; missing: {missing}"
            )
            build_urls = None
        else:
            build_urls = find_master_builds(build_types)
            assert build_urls, "Could not find master builds in S3"
        if build_urls:
            for bt, url in build_urls.items():
                bt_path = bt_paths[bt]
                if not info.is_local_run or not Path(bt_path).is_file():
                    print(f"NOTE: Downloading {bt} build to [{bt_path}]")
                    Shell.run(
                        f"wget -nv -O {bt_path} {url}", verbose=True, strict=True
                    )
                    Shell.run(f"chmod +x {bt_path}", verbose=True)
        clickhouse_path = f"{temp_path}/clickhouse_{build_types[0]}"

    if is_bugfix_validation or is_flaky_check:
        assert (
            changed_test_modules or (info.is_local_run and args.test)
        ), "No changed test modules found, either job must be skipped or bug in changed test search logic"

    Shell.check(f"chmod +x {clickhouse_path}", verbose=True, strict=True)
    Shell.check(f"{clickhouse_path} --version", verbose=True, strict=True)

    targeted_tests = []
    if is_targeted_check:
        assert not args.test, "--test not supposed to be used for targeted check ???"
        targeter = Targeting(info=info)
        tests, results_with_info = targeter.get_all_relevant_tests_with_info()
        # no subtask level for integration tests - cannot add this info to the report now
        # results.append(results_with_info)
        if not tests:
            # early exit
            Result.create_from(
                status=Result.Status.SKIPPED,
                info="No failed tests found from previous runs",
            ).complete_job()

        # Parse test names from the query result
        for test_ in tests:
            if test_.strip():
                test_name = test_.strip()
                targeted_tests.append(
                    test_name.split("[")[0]
                )  # remove parametrization - does not work with test repeat with --count
        print(f"Parsed {len(targeted_tests)} test names: {targeted_tests}")

    if not Shell.check("docker info > /dev/null 2>&1", verbose=True):
        start_docker_in_docker()
    Shell.check("docker info > /dev/null", verbose=True, strict=True)

    parallel_test_modules, sequential_test_modules = (
        get_parallel_sequential_tests_to_run(
            batch_num,
            total_batches,
            args.test or targeted_tests or changed_test_modules,
            workers,
            args.options,
            info,
            no_strict=is_targeted_check or is_flaky_check,  # targeted check might want to run test that was removed on a merge-commit; flaky check might pick up a changed test filtered out by SKIP_LIST in the private fork
        )
    )

    if is_flaky_check or is_targeted_check:
        # The flaky/targeted parallel bucket runs `--dist=each`: every worker runs
        # every parallel module at once. TEST_CONFIGS `dist_each_sequential` modules
        # would start one cluster per worker and OOM small runners, so move them to
        # the looped sequential phase. Normal `--dist=loadfile` runs do not call this.
        before = list(parallel_test_modules)
        parallel_test_modules, sequential_test_modules = force_heavy_modules_sequential(
            parallel_test_modules, sequential_test_modules
        )
        moved = [m for m in before if m not in parallel_test_modules]
        if moved:
            print(f"Forced heavy modules to the sequential phase (avoid concurrent --dist=each clusters): {moved}")

    if is_sequential:
        parallel_test_modules = []
        assert not is_parallel
    elif is_parallel:
        sequential_test_modules = []
        assert not is_sequential

    # If this PR only touches test files (no production/config code changed),
    # this batch only needs to run whichever of parallel_test_modules /
    # sequential_test_modules actually contains a changed module - the other
    # side would produce results identical to master and can be dropped
    # outright (saving the time to run it), and if neither side contains a
    # changed module the whole batch can be skipped. Placed after the
    # is_sequential/is_parallel handling above so it sees the modules this
    # job invocation will actually run, not the pre-flavor-filter set.
    if (
        total_batches > 1
        and not is_flaky_check
        and not is_targeted_check
        and not is_bugfix_validation
        and not args.test
    ):
        changed_files = info.get_changed_files()
        if changed_files and all(
            Targeting.is_functional_test_file(f)
            or Targeting.is_integration_test_file(f)
            or Targeting.is_ci_job_script(f)
            for f in changed_files
        ):
            changed_integration_modules = {
                f.removeprefix("tests/integration/")
                for f in changed_files
                if Targeting.is_integration_test_file(f)
            }
            if not changed_integration_modules:
                Result.create_from(
                    status=Result.Status.SKIPPED,
                    info="Only non-integration test files changed in this PR - nothing for this job to run",
                ).complete_job()
            if not (changed_integration_modules & set(parallel_test_modules)):
                parallel_test_modules = []
            if not (changed_integration_modules & set(sequential_test_modules)):
                sequential_test_modules = []
            if not parallel_test_modules and not sequential_test_modules:
                Result.create_from(
                    status=Result.Status.SKIPPED,
                    info="Only test files changed in this PR and none of the changed test modules fall into this batch",
                ).complete_job()

    if (is_targeted_check or is_flaky_check) and not parallel_test_modules and not sequential_test_modules:
        # Targeted check: all selected tests were stale (removed or renamed since the CIDB record).
        # Flaky check: all changed tests were filtered out (e.g. by SKIP_LIST in the private fork).
        # Either way, skip gracefully instead of producing a "no results" error.
        skip_info = (
            "All targeted tests are stale (removed or renamed)"
            if is_targeted_check
            else "All changed tests were filtered out (e.g. by SKIP_LIST)"
        )
        Result.create_from(
            status=Result.Status.SKIPPED,
            info=skip_info,
        ).complete_job()

    if is_flaky_check or is_targeted_check:
        # Sort by module file so all tests from the same file are consecutive.
        # With --dist=each, pytest preserves CLI argument order and uses it as the
        # collection order. If tests from different modules interleave (e.g. CIDB
        # returns them sorted by failure time), pytest finalizes and re-enters
        # module-scoped fixtures between them, breaking tests that call
        # cluster.add_instance() inside the fixture.
        # For regular jobs, preserve the duration-aware ordering from get_optimal_test_batch.
        parallel_test_modules = sorted(parallel_test_modules, key=lambda t: t.split("::")[0])
        sequential_test_modules = sorted(sequential_test_modules, key=lambda t: t.split("::")[0])

    # Setup environment variables for tests
    for image_name, env_name in IMAGES_ENV.items():
        tag = info.docker_tag(image_name)
        if tag:
            print(f"Setting environment variable [{env_name}] to [{tag}]")
            os.environ[env_name] = tag
        else:
            assert False, f"No tag found for image [{image_name}]"

    # Pre-fetch all Docker images needed by the selected test suites.
    # This is done after IMAGES_ENV vars are set so tag resolution works correctly.
    # Fail fast here rather than discovering missing images mid-test-run.
    all_test_modules = parallel_test_modules + sequential_test_modules
    compose_files = get_compose_files_for_test_modules(all_test_modules)
    print(
        f"Compose files detected for this batch ({len(compose_files)}): "
        + ", ".join(str(f.name) for f in compose_files)
    )
    images_to_prefetch = get_images_from_compose_files(compose_files)
    if not prefetch_images(images_to_prefetch):
        Result.create_from(
            status=Result.Status.ERROR,
            info="Failed to pre-pull Docker images needed by the test batch",
            labels=[Result.Label.INFRA],
        ).complete_job()

    test_env = {
        "CLICKHOUSE_TESTS_BASE_CONFIG_DIR": clickhouse_server_config_dir,
        "CLICKHOUSE_TESTS_SERVER_BIN_PATH": clickhouse_path,
        "CLICKHOUSE_BINARY": clickhouse_path,  # some test cases support alternative binary location
        "CLICKHOUSE_TESTS_CLIENT_BIN_PATH": clickhouse_path,
        "CLICKHOUSE_USE_OLD_ANALYZER": "1" if use_old_analyzer else "0",
        "CLICKHOUSE_USE_DISTRIBUTED_PLAN": "1" if use_distributed_plan else "0",
        "CLICKHOUSE_USE_DATABASE_DISK": "1" if use_database_disk else "0",
        "PYTEST_CLEANUP_CONTAINERS": "1",
        "JAVA_PATH": java_path,
        # PromQL compliance: deterministic JSON for post-hook (see promql_compliance_hook.py).
        "COMPLIANCE_RESULT_FILE": os.environ.get(
            "COMPLIANCE_RESULT_FILE", os.path.join(temp_path, "promql_compliance_result.json")
        ),
    }
    if is_llvm_coverage:
        test_env["LLVM_PROFILE_FILE"] = "it-%4m.profraw"
        print(
            f"NOTE: This is LLVM coverage run, setting LLVM_PROFILE_FILE to [{test_env['LLVM_PROFILE_FILE']}]"
        )
        # Auto-detect available LLVM profdata tool
        for ver in ["21", "20", "18", "19", "17", "16", ""]:
            cmd = f"llvm-profdata{'-' + ver if ver else ''}"
            if Shell.check(f"command -v {cmd}", verbose=False):
                llvm_profdata_cmd = cmd
                break

        if not llvm_profdata_cmd:
            print("ERROR: llvm-profdata not found in PATH")
        else:
            print(f"Using {llvm_profdata_cmd} to merge coverage files")

    test_results = []
    failed_tests_files = []

    has_error = False
    # Set when a pytest run was cut short by a timeout (graceful session-timeout or the
    # hard subprocess backstop). Used below to keep an empty flaky/targeted result a
    # best-effort SKIPPED only when a timeout actually exhausted the budget.
    timed_out = False
    session_timeout_parallel = 3600 * 2
    session_timeout_sequential = 3600

    if is_llvm_coverage:
        session_timeout_parallel = 7200
        session_timeout_sequential = 7200

    if args.session_timeout:
        session_timeout_parallel = args.session_timeout * 2
        session_timeout_sequential = args.session_timeout

    # Flaky-check soft timeout. Mirrors the pattern in `ci/jobs/functional_tests.py`:
    # bound the total time spent inside pytest so the job has headroom for cleanup,
    # log collection and reporting before the job is cancelled. Without this, a
    # flaky-check run over many modified test modules can be hard-killed, producing no
    # report at all instead of a best-effort partial one.
    #
    # The budget must stay well below the external ceiling at which a lone integration
    # job is cancelled (observed at ~80-90 min from job start in CI). The previous 90 min
    # was above that ceiling: the graceful xdist `--session-timeout` and the subprocess
    # hard-kill backstop never fired before the external cancellation, so the whole
    # process was killed and 0 results were reported (the job showed a red failure rather
    # than a best-effort report). 45 min matches the functional flaky check and leaves
    # ample room for the hard-kill backstop (+600s below), cleanup and reporting.
    FLAKY_CHECK_TIME_LIMIT = 45 * 60  # 45 min - kept below the external job-cancellation ceiling
    if is_flaky_check:
        elapsed_for_flaky = int(sw.duration)
        flaky_check_remaining_s = max(FLAKY_CHECK_TIME_LIMIT - elapsed_for_flaky, 60)
        print(
            f"Flaky-check time limit: {FLAKY_CHECK_TIME_LIMIT}s "
            f"(elapsed so far: {elapsed_for_flaky}s, remaining: {flaky_check_remaining_s}s)"
        )
        # Cap per-phase session timeouts so a single phase cannot consume the entire budget.
        session_timeout_parallel = min(session_timeout_parallel, flaky_check_remaining_s)
        session_timeout_sequential = min(session_timeout_sequential, flaky_check_remaining_s)

    error_info = []

    failed_test_cases = []

    # Clear dmesg to avoid false OOM detection from previous CI jobs on the same host.
    # Do this only in CI (non-local runs) and via a non-interactive privileged helper.
    if not info.is_local_run:
        try:
            Utils.clear_dmesg()
        except Exception as ex:
            print(f"Failed to clear dmesg before integration tests: {ex}")

    if is_flaky_check or is_targeted_check:
        # Each xdist worker runs all modules independently with its own isolated Docker cluster.
        # ClickHouseCluster appends PYTEST_XDIST_WORKER to the project name, so clusters
        # from different workers never interfere. --dist=each sends all tests to every worker.
        parallel_dist = "--dist=each"
        parallel_workers = workers
        # Sequential tests cannot run in parallel, so we loop over them instead.
        # Run at least 3 times to have meaningful flakiness signal, at most workers times.
        sequential_repeat_cnt = max(3, workers)
    else:
        parallel_dist = "--dist=loadfile"
        parallel_workers = workers
        sequential_repeat_cnt = 1

    if parallel_test_modules:
        log_file = f"{temp_path}/pytest_parallel.log"
        test_result_parallel, parallel_timed_out = run_pytest_and_collect_results(
            command=f"{quote_tests(parallel_test_modules)} --report-log-exclude-logs-on-passed-tests -n {parallel_workers} {parallel_dist} --tb=short {repeat_option} --session-timeout={session_timeout_parallel}",
            env=test_env,
            report_name="parallel",
            timeout=session_timeout_parallel + 600,
        )
        timed_out = timed_out or parallel_timed_out
        test_results.extend(test_result_parallel.results)
        _mark_infrastructure_errors(test_result_parallel.results)
        failed_test_cases.extend(
            [t.name for t in test_result_parallel.results if t.is_failure()]
        )
        if test_result_parallel.files:
            failed_tests_files.extend(test_result_parallel.files)
        if test_result_parallel.is_error():
            if not is_targeted_check and not is_flaky_check:
                # In targeted checks we may overload the run with many heavy tests running
                # in parallel; in flaky checks the soft FLAKY_CHECK_TIME_LIMIT may cap the
                # pytest session-timeout. In both cases a session-timeout is an expected
                # risk rather than an infrastructure problem, so we do not treat such
                # errors as job-level failures.
                has_error = True
                error_info.append(test_result_parallel.info)

    fail_num = len([r for r in test_results if not r.is_ok()])
    if sequential_test_modules and fail_num < MAX_FAILS_BEFORE_DROP and not has_error:
        for attempt in range(sequential_repeat_cnt):
            # Recompute remaining budget for flaky-check at every iteration and stop
            # scheduling new runs once it is exhausted (soft timeout).
            iter_session_timeout_sequential = session_timeout_sequential
            if is_flaky_check:
                elapsed_for_flaky = int(sw.duration)
                flaky_check_remaining_s = max(FLAKY_CHECK_TIME_LIMIT - elapsed_for_flaky, 0)
                if flaky_check_remaining_s < 60:
                    print(
                        f"Flaky-check time limit reached after [{attempt}/{sequential_repeat_cnt}] sequential attempts "
                        f"(elapsed: {elapsed_for_flaky}s, limit: {FLAKY_CHECK_TIME_LIMIT}s); stopping"
                    )
                    break
                iter_session_timeout_sequential = min(
                    session_timeout_sequential, flaky_check_remaining_s
                )
            test_result_sequential, sequential_timed_out = run_pytest_and_collect_results(
                command=f"{quote_tests(sequential_test_modules)} --report-log-exclude-logs-on-passed-tests --tb=short {repeat_option} -n 1 --dist=loadfile --session-timeout={iter_session_timeout_sequential}",
                env=test_env,
                report_name="sequential",
                timeout=iter_session_timeout_sequential + 600,
            )
            timed_out = timed_out or sequential_timed_out
            test_results.extend(test_result_sequential.results)
            _mark_infrastructure_errors(test_result_sequential.results)
            failed_test_cases.extend(
                [t.name for t in test_result_sequential.results if t.is_failure()]
            )
            if test_result_sequential.files:
                failed_tests_files.extend(test_result_sequential.files)
            if test_result_sequential.is_error():
                if not is_targeted_check and not is_flaky_check:
                    # In targeted checks we may overload the run with many heavy tests running
                    # sequentially; in flaky checks the per-iteration `iter_session_timeout_sequential`
                    # may be capped by FLAKY_CHECK_TIME_LIMIT. In both cases a session-timeout is
                    # an expected risk rather than an infrastructure problem, so we do not treat
                    # such errors as job-level failures.
                    has_error = True
                    error_info.append(test_result_sequential.info)
                break
            if (is_flaky_check or is_targeted_check) and not test_result_sequential.is_ok():
                print(
                    f"Flaky/targeted check: sequential test run fails after attempt [{attempt+1}/{sequential_repeat_cnt}] - break"
                )
                break

    # Run additional build types for bugfix validation.
    # Exit early on first failure to avoid duplicate test names and workspace pollution.
    if is_bugfix_validation:
        build_types = bugfix_build_types(info.job_name)
        for r in test_results:
            r.set_label(build_types[0])

        if all(r.is_ok() for r in test_results):
            for bugfix_bt in build_types[1:]:
                print(f"\n=== Bugfix validation with {bugfix_bt} ===")
                bt_clickhouse_path = f"{temp_path}/clickhouse_{bugfix_bt}"
                test_env["CLICKHOUSE_TESTS_SERVER_BIN_PATH"] = bt_clickhouse_path
                test_env["CLICKHOUSE_BINARY"] = bt_clickhouse_path
                test_env["CLICKHOUSE_TESTS_CLIENT_BIN_PATH"] = bt_clickhouse_path
                Shell.check(
                    f"{bt_clickhouse_path} --version", verbose=True, strict=True
                )

                bt_test_results = []

                if parallel_test_modules:
                    bt_result_parallel, _ = run_pytest_and_collect_results(
                        command=f"{quote_tests(parallel_test_modules)} --report-log-exclude-logs-on-passed-tests -n {workers} --dist=loadfile --tb=short {repeat_option} --session-timeout={session_timeout_parallel}",
                        env=test_env,
                        report_name=f"parallel_{bugfix_bt}",
                        timeout=session_timeout_parallel + 600,
                    )
                    bt_test_results.extend(bt_result_parallel.results)
                    _mark_infrastructure_errors(bt_result_parallel.results)
                    if bt_result_parallel.files:
                        failed_tests_files.extend(bt_result_parallel.files)
                    if bt_result_parallel.is_error():
                        has_error = True
                        error_info.append(bt_result_parallel.info)

                bt_fail_num = len([r for r in bt_test_results if not r.is_ok()])
                if sequential_test_modules and bt_fail_num < MAX_FAILS_BEFORE_DROP and not has_error:
                    bt_result_sequential, _ = run_pytest_and_collect_results(
                        command=f"{quote_tests(sequential_test_modules)} --report-log-exclude-logs-on-passed-tests --tb=short {repeat_option} -n 1 --dist=loadfile --session-timeout={session_timeout_sequential}",
                        env=test_env,
                        report_name=f"sequential_{bugfix_bt}",
                        timeout=session_timeout_sequential + 600,
                    )
                    bt_test_results.extend(bt_result_sequential.results)
                    _mark_infrastructure_errors(bt_result_sequential.results)
                    if bt_result_sequential.files:
                        failed_tests_files.extend(bt_result_sequential.files)
                    if bt_result_sequential.is_error():
                        has_error = True
                        error_info.append(bt_result_sequential.info)

                for r in bt_test_results:
                    r.set_label(bugfix_bt)
                test_results = bt_test_results

                if any(not r.is_ok() for r in bt_test_results):
                    break

    # Collect logs before re-run
    attached_files = []
    if not info.is_local_run:
        failed_suits = []
        # Collect docker compose configs used in tests
        config_files = [
            str(p)
            for p in Path("./tests/integration/").glob("test_*/_instances*/*/configs/")
        ]
        for test_result in test_results:
            if not test_result.is_ok() and ".py" in test_result.name:
                failed_suits.append(test_result.name.split("/")[0])
        failed_suits = list(set(failed_suits))
        for failed_suit in failed_suits:
            failed_tests_files.append(f"tests/integration/{failed_suit}")

        # Add all files matched ./ci/tmp/*.log ./ci/tmp/*.jsonl into failed_tests_files
        for pattern in ["*.log", "*.jsonl"]:
            for log_file in Path("./ci/tmp/").glob(pattern):
                if log_file.is_file():
                    failed_tests_files.append(str(log_file))

        if failed_suits:
            attached_files.append(
                Utils.compress_files_gz(failed_tests_files, f"{temp_path}/logs.tar.gz")
            )
            attached_files.append(
                Utils.compress_files_gz(config_files, f"{temp_path}/configs.tar.gz")
            )
            if Path("./ci/tmp/docker-in-docker.log").exists():
                attached_files.append("./ci/tmp/docker-in-docker.log")

    # Rerun failed tests if any to check if failure is reproducible
    if 0 < len(failed_test_cases) < 10 and not (
        is_flaky_check or is_bugfix_validation or is_targeted_check or info.is_local_run
    ):
        test_result_retries, _ = run_pytest_and_collect_results(
            command=f"{quote_tests(failed_test_cases)} --report-log-exclude-logs-on-passed-tests --tb=short -n 1 --dist=loadfile --session-timeout=1200",
            env=test_env,
            report_name="retries",
            timeout=1200 + 600,
        )
        successful_retries = [t.name for t in test_result_retries.results if t.is_ok()]
        failed_retries = [t.name for t in test_result_retries.results if t.is_failure()]
        if successful_retries or failed_retries:
            for test_case in test_results:
                if test_case.name in successful_retries:
                    test_case.set_label(Result.Label.OK_ON_RETRY)
                elif test_case.name in failed_retries:
                    test_case.set_label(Result.Label.FAILED_ON_RETRY)

    # Remove iptables rule added in tests
    Shell.check("sudo iptables -D DOCKER-USER 1 ||:", verbose=True)

    if not info.is_local_run:
        print("Dumping dmesg")
        Shell.check("dmesg -T > ./ci/tmp/dmesg.log", verbose=True, strict=True)
        with open("./ci/tmp/dmesg.log", "rb") as dmesg:
            dmesg = dmesg.read()
            if (
                b"Out of memory: Killed process" in dmesg
                or b"oom_reaper: reaped process" in dmesg
                or b"oom-kill:constraint=CONSTRAINT_NONE" in dmesg
            ):
                if is_bugfix_validation:
                    # A host OOM is an infrastructure/resource failure, not bug
                    # reproduction. Report it as `ERROR` and set `has_error` so
                    # the status inversion below is skipped (same as `Timeout`),
                    # otherwise the inverted `FAIL` would flip the job to green.
                    test_results.append(
                        Result(
                            name=OOM_IN_DMESG_TEST_NAME, status=Result.Status.ERROR
                        )
                    )
                    has_error = True
                    error_info.append(
                        "OOM in dmesg - infrastructure/resource failure, "
                        "not bug reproduction"
                    )
                else:
                    test_results.append(
                        Result(
                            name=OOM_IN_DMESG_TEST_NAME, status=Result.Status.FAIL
                        )
                    )
                attached_files.append("./ci/tmp/dmesg.log")

    # For targeted, flaky checks, and bugfix validation, the synthetic "Timeout"
    # result must not be propagated as a top-level `FAIL`: for targeted checks a
    # session-timeout is an expected risk (because of `--count N` overloading), for
    # flaky checks because of the soft `FLAKY_CHECK_TIME_LIMIT`, and for bugfix
    # validation an inverted `FAIL` would be mistakenly treated as successful bug
    # reproduction.
    if is_targeted_check or is_flaky_check or is_bugfix_validation:
        test_results = [r for r in test_results if r.name != "Timeout"]

    # Whether pytest produced any real test results, captured *before* the synthetic
    # `skipped_flaky_modules` entries are appended below. The empty-result status decision
    # must look only at real pytest output: in a scope-capped flaky run the synthetic
    # `SKIPPED` entries would otherwise make `test_results` non-empty even when the selected
    # modules produced nothing, masking a timeout-empty run (should be `SKIPPED`) or an empty
    # harness failure (should be `ERROR`) as a green top-level result.
    pytest_has_results = bool(test_results)

    # Make the best-effort scope cap explicit in the report: list the modules that were
    # skipped to fit the time budget as SKIPPED entries rather than dropping them silently.
    for skipped_module in skipped_flaky_modules:
        test_results.append(
            Result(
                name=skipped_module,
                status=Result.Status.SKIPPED,
                info="Skipped by flaky-check best-effort scope cap (MAX_FLAKY_CHECK_MODULES)",
            )
        )

    # If a timeout exhausted the time budget before any test produced a result (e.g. a
    # single very slow or hanging module consumed the whole budget), report SKIPPED rather
    # than letting `create_from` default an empty result set to ERROR, which would block
    # the PR. Crucially, this best-effort downgrade applies *only* when a timeout was
    # actually observed: an empty result without a timeout means pytest failed to produce
    # any output for some other reason (crashed before writing the jsonl report, plugin or
    # internal error, ...), which is a real harness failure and must stay ERROR.
    #
    # Both decisions use `pytest_has_results` (real pytest output) rather than the current
    # `test_results`, which may carry only the synthetic `skipped_flaky_modules` entries
    # appended above. With a scope cap in effect those synthetic `SKIPPED` rows make
    # `test_results` non-empty, so the empty-result status must be forced here: `create_from`
    # only defaults an empty result set to `ERROR`, and a list of `SKIPPED` rows would
    # otherwise collapse to a green top-level status.
    empty_best_effort = is_empty_best_effort_skip(
        is_flaky_check, is_targeted_check, pytest_has_results, timed_out
    )
    empty_harness_failure = (
        (is_flaky_check or is_targeted_check) and not pytest_has_results and not timed_out
    )
    R = Result.create_from(
        results=test_results,
        status=(
            Result.Status.SKIPPED
            if empty_best_effort
            else Result.Status.ERROR if empty_harness_failure else ""
        ),
        info=(
            "No test results collected within the flaky-check time budget (best effort)"
            if empty_best_effort
            else (
                "No test results collected and no timeout was observed - reporting the "
                "empty pytest run as a harness ERROR"
                if empty_harness_failure
                else ""
            )
        ),
        stopwatch=sw,
        files=attached_files,
    )

    if is_llvm_coverage:
        assert (
            is_bugfix_validation is False
        ), "LLVM coverage with bugfix validation is not supported"
        has_failure = False
        for r in R.results:
            if r.status == Result.Status.FAIL:
                if r.has_label(Result.Label.OK_ON_RETRY):
                    # Remove label and set to OK
                    r.remove_label(Result.Label.OK_ON_RETRY)
                    r.status = Result.Status.OK
                else:
                    has_failure = True
        if has_failure:
            R.set_failed()
            R.set_info("Some tests failed during LLVM coverage run")
        else:
            R.set_success()
            has_error = False

    # Capture whether this run saw any infrastructure problems BEFORE the
    # clearing block below resets `has_error`. If the answer is yes, the
    # bugfix-validation inversion path further down must be skipped: we have
    # no reliable signal about whether the bug reproduces on this arch, and
    # running the inversion would let an infra `FAIL` be flipped to `OK`
    # (counted as validation) or rewrite an `ERROR`-only outcome as
    # `SKIPPED` via the no-`has_failure` branch. See bot review on
    # ClickHouse/ClickHouse#103541 (2026-05-15).
    had_infra_or_error = has_error or any(
        r.has_label(Result.Label.INFRA) for r in test_results
    )

    # If all non-OK results are infrastructure errors, do not treat as a real failure
    if has_error:
        non_ok = [r for r in test_results if not r.is_ok()]
        if non_ok and all(r.has_label(Result.Label.INFRA) for r in non_ok):
            print(
                "All failures are infrastructure errors - clearing error flag"
            )
            has_error = False
            force_ok_exit = True

    if has_error:
        R.set_error().set_info("\n".join(error_info))

    if is_bugfix_validation and (Labels.PR_BUGFIX in info.pr_labels or Labels.PR_CRITICAL_BUGFIX in info.pr_labels):
        assert (
            is_llvm_coverage is False
        ), "Bugfix validation with LLVM coverage is not supported"
        if had_infra_or_error:
            # Infrastructure errors or session-level failures were observed
            # during this run. Skip the inversion path so the per-arch job
            # cannot be silently counted as a validation. The post-hook in
            # `new_tests_check.py` uses strict `is_success` (`OK` / `XFAIL`
            # only); leaving the result in a non-success state is enough to
            # prevent this arch from contributing a false validation.
            #
            # If, after all the upstream handling, the result is still in a
            # success-equivalent state (e.g. every surviving child is `OK`
            # because all infra failures were already relabeled to `SKIPPED`
            # by `_mark_infrastructure_errors`), force `ERROR` here so the
            # post-hook cannot accidentally treat this arch as validated.
            print(
                "Bugfix validation: infrastructure error or session-level "
                "failure detected - skipping status inversion to avoid "
                "leaking an infra outcome into validation success."
            )
            if R.is_success():
                R.set_error().set_info(
                    "Bugfix validation aborted: infrastructure error during "
                    "the run - no reliable signal about whether the bug "
                    "reproduces on this arch"
                )
        else:
            has_failure = False
            for r in R.results:
                # invert statuses: only `FAIL` is treated as a successful
                # reproduction signal. Generic `ERROR` is left untouched
                # because in integration tests `ERROR` is also used for
                # runner-level problems (for example session-timeout paths
                # in `run_pytest_and_collect_results`), and infrastructure
                # errors that escape `_mark_infrastructure_errors` could
                # otherwise flip the job to green.
                r.set_label(Result.Label.XFAIL)
                if r.status == Result.Status.FAIL:
                    r.status = Result.Status.OK
                    has_failure = True
                elif r.status == Result.Status.OK:
                    r.status = Result.Status.FAIL
            if not has_failure:
                # See the matching comment in `ci/jobs/functional_tests.py`. The
                # bug did not reproduce on this arch, so report SKIPPED instead
                # of FAIL: `Result.is_ok` includes SKIPPED so the job exits 0,
                # while `is_success` (used by the post-hook) excludes SKIPPED so
                # the per-arch job does not count as a validation. Contract:
                # at least one per-arch job must end up `OK`/`XFAIL` for the
                # post-hook to consider the bug validated.
                print("Bug does not reproduce on this arch - bugfix validation N/A")
                R.set_status(Result.Status.SKIPPED)
                R.set_info("Bug does not reproduce on this arch - bugfix validation N/A")
            else:
                R.set_success()

    force_ok_exit = False
    if R:
        failures_cnt = len([r for r in R.results if not r.is_ok()])
        if failures_cnt > 0 and failures_cnt < 4:
            print(
                f"NOTE: Failed {failures_cnt} tests - do not block pipeline, exit with 0"
            )
            force_ok_exit = True
        elif failures_cnt > 0 and "ci-non-blocking" in info.pr_labels:
            print(
                f"NOTE: Failed {failures_cnt} tests, label 'ci-non-blocking' is set - do not block pipeline - exit with 0"
            )
            force_ok_exit = True
    if is_bugfix_validation:
        # Per-arch bugfix-validation jobs are advisory: their pass/fail status
        # records "did the bug reproduce on this arch?", not whether the PR
        # should be blocked. Setting `do_not_block_pipeline_on_failure=True`
        # marks the job as non-blocking so downstream jobs are not dropped
        # when this job reports FAIL. The process itself still exits with
        # the natural status (`Result.complete_job` calls `sys.exit(1)` on
        # non-OK results); the non-blocking flag is metadata for the
        # pipeline scheduler. The PR-merge-blocking decision lives in the
        # `new_tests_check.py` workflow post-hook, which OR's the per-arch
        # bugfix-validation job statuses.
        print(
            "NOTE: Bugfix validation job - marking as non-blocking; "
            "failure here will not block downstream pipeline jobs "
            "(process exit code still reflects the actual job status)"
        )
        force_ok_exit = True
    if is_llvm_coverage and llvm_profdata_cmd:
        print("Collecting and merging LLVM coverage files...")

        # Merge all profraw files into final profdata file
        merged_profdata = merge_profraw_files(llvm_profdata_cmd, job_params)

        # Attach profdata file to the result report so it is uploaded
        # unconditionally (even when tests fail) and visible in the CI report.
        if merged_profdata and os.path.exists(merged_profdata):
            R.files.append(merged_profdata)

        force_ok_exit = True
        print("NOTE: LLVM coverage job - do not block pipeline - exit with 0")

    R.sort().complete_job(do_not_block_pipeline_on_failure=force_ok_exit)


if __name__ == "__main__":
    main()