wenjin272 commented on code in PR #992: URL: https://github.com/apache/flink-agents/pull/992#discussion_r3794819796
########## e2e-test/test-scripts/test_checkpoint_recovery.sh: ########## @@ -0,0 +1,1307 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF 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. +# +# +# Verifies that Python agent memory survives the loss of a TaskManager process. +# +# Submits checkpoint_recovery_job.py to a local Flink standalone cluster. The job +# parks itself inside a tool call that blocks on a file this script creates, so a +# completed checkpoint provably holds the agent's memory before anything is killed. +# The script then hard-kills the TaskManager, restarts it, waits for a real restore, +# releases the tool and reads the verdict the job publishes. +# +# Unlike its sibling test_submit_examples_to_flink.sh, a successful submission is +# NOT a pass: the only pass is a verdict file that says so. +# +# Env: FLINK_VERSION (default 2.3.0), FLINK_HOME (reuse existing install), +# VERBOSE=1 (set -x), plus the *_TIMEOUT overrides below. +# +# A FLINK_HOME passed in has to be that same FLINK_VERSION. The run copies +# opt/flink-python-<FLINK_VERSION>.jar into its lib/ and stops when that file is +# not there, so point FLINK_VERSION at whatever installation FLINK_HOME names. + +set -euo pipefail + +if [[ "${VERBOSE:-0}" == "1" ]]; then + set -x +fi + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[0;33m' +BLUE='\033[0;34m' +NC='\033[0m' + +log_info() { printf "${BLUE}[INFO]${NC} %s\n" "$*" >&2; } +log_ok() { printf "${GREEN}[OK]${NC} %s\n" "$*" >&2; } +log_warn() { printf "${YELLOW}[WARN]${NC} %s\n" "$*" >&2; } +log_error() { printf "${RED}[ERROR]${NC} %s\n" "$*" >&2; } +log_section() { + printf "\n${BLUE}==============================================================${NC}\n" >&2 + printf "${BLUE}>>> %s${NC}\n" "$*" >&2 + printf "${BLUE}==============================================================${NC}\n" >&2 +} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/../.."; pwd)" +log_info "Project root: $ROOT_DIR" + +FLINK_VERSION="${FLINK_VERSION:-2.3.0}" +FLINK_MAJOR_MINOR="${FLINK_VERSION%.*}" +REST_URL="${REST_URL:-http://localhost:8081}"; + +JOB_MODULE="flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_job.py" +EXPECTED_AGENTS_VERSION="${EXPECTED_AGENTS_VERSION:-0.3.dev0}" + +# Checkpoint interval is deliberately short: the run is parked while we wait for two +# checkpoints to complete, and that wait is charged against the tool's own deadline. +# Held in milliseconds because that is the unit the checkpoint-config endpoint reports, +# so the value written and the value asserted are the same number. +CHECKPOINT_INTERVAL_MS="${CHECKPOINT_INTERVAL_MS:-5000}" +RESTART_ATTEMPTS="${RESTART_ATTEMPTS:-3}" +# Unset, this falls back to slot.request.timeout (5 min), after which a pending slot +# request is failed and one restart attempt is burned while the TM is being replaced. +STANDALONE_STARTUP_TIME="${STANDALONE_STARTUP_TIME:-600s}" + +# Budgets. Setup waits are unconstrained; the four marked ones run after the tool has +# parked and are therefore charged against the tool's own release deadline. They are +# clamped at runtime to the time actually remaining (see charged_timeout), so these +# numbers are ceilings rather than guarantees. +CLUSTER_TIMEOUT="${CLUSTER_TIMEOUT:-120}" +SUBMIT_TIMEOUT="${SUBMIT_TIMEOUT:-300}" +JOB_RUNNING_TIMEOUT="${JOB_RUNNING_TIMEOUT:-180}" +IDENTITY_TIMEOUT="${IDENTITY_TIMEOUT:-180}" +TOOL_ENTERED_TIMEOUT="${TOOL_ENTERED_TIMEOUT:-300}" +CHECKPOINT_TIMEOUT="${CHECKPOINT_TIMEOUT:-30}" # charged against the tool deadline +TM_GONE_TIMEOUT="${TM_GONE_TIMEOUT:-45}" # charged against the tool deadline +TM_UP_TIMEOUT="${TM_UP_TIMEOUT:-30}" # charged against the tool deadline +RESTORE_TIMEOUT="${RESTORE_TIMEOUT:-30}" # charged against the tool deadline +VERDICT_TIMEOUT="${VERDICT_TIMEOUT:-120}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +# Every REST call is bounded. This is charged in the budget arithmetic, so keep the +# two in step: a healthy local JobManager answers in milliseconds, and a large value +# here buys nothing while inflating the worst-case overrun of every wait. +CURL_MAX_TIME="${CURL_MAX_TIME:-5}" + +# Wall-clock costs inside the charged window that belong to no single wait: the +# individual REST reads between the waits, two jps invocations, the SIGKILL grace and +# starting the replacement TaskManager. This is an estimate, and deliberately only +# feeds the pre-flight feasibility check — the runtime clamp is what actually holds +# the guarantee, so being wrong here cannot let the tool self-release. +HANDSHAKE_FIXED_COST_S="${HANDSHAKE_FIXED_COST_S:-45}" + +# Kept in hand so `release` is written comfortably before the tool gives up, covering +# the tool's own poll interval and the one-second granularity of SECONDS. +HANDSHAKE_SAFETY_S="${HANDSHAKE_SAFETY_S:-15}" + +# Bash 3 (default on macOS) lacks associative arrays. +RESULT_NAMES=() +RESULT_STATES=() +SUBMITTED_JOB_IDS=() +CONFIG_KEYS_SET=() +JOB_ID="" +WORK_DIR="" +HANDSHAKE_DIR="" +VERDICT_DIR="" +CHECKPOINT_DIR="" +INPUT_DIR="" +INPUT_FILE="" +FLINK_CONF="" +TM_PID_BEFORE="" +TM_RESOURCE_ID_BEFORE="" +RESTORED_BEFORE="" +RELEASE_DEADLINE_S="" +HANDSHAKE_DEADLINE_AT="" + +cleanup() { + local exit_code=$? + log_section "Cleanup" + + if [[ -n "${FLINK_HOME:-}" && -x "$FLINK_HOME/bin/flink" ]]; then + for jid in "${SUBMITTED_JOB_IDS[@]:-}"; do + [[ -n "$jid" ]] || continue + log_info "Cancelling job $jid" + "$FLINK_HOME/bin/flink" cancel "$jid" >/dev/null 2>&1 || true + done + + if [[ -x "$FLINK_HOME/bin/stop-cluster.sh" ]]; then + # If the TaskManager was killed and never replaced, this prints + # "No taskexecutor daemon (pid: N) is running anymore" into the archive. + # That line is expected here and is not the failure. + log_info "Stopping Flink cluster" + "$FLINK_HOME/bin/stop-cluster.sh" >/dev/null 2>&1 || true + fi + + # Archive the whole log directory, never named files: `kill -9` leaves the + # dead TaskManager's pid in the pid file, so the replacement writes to a + # higher log index and a named copy would miss the post-restore TaskManager. + if [[ -d "$FLINK_HOME/log" ]]; then + local log_archive="$ROOT_DIR/flink-logs-$(date +%Y%m%d-%H%M%S).tar.gz" + tar -czf "$log_archive" -C "$FLINK_HOME" log >/dev/null 2>&1 \ + && log_info "Flink logs archived to: $log_archive" \ + || log_warn "Failed to archive Flink logs" + fi + fi + + # Strip the keys we appended. A second run against a reused Flink home would + # otherwise append them twice, and a duplicate key is a hard YAML parse failure + # that stops the cluster from starting at all. + if [[ -n "$FLINK_CONF" && -f "$FLINK_CONF" ]]; then + local key + for key in "${CONFIG_KEYS_SET[@]:-}"; do + [[ -n "$key" ]] || continue + delete_config_key "$key" + done + log_info "Reverted ${#CONFIG_KEYS_SET[@]} appended config key(s)" + fi + + # The killed TaskManager's entry stays in /tmp/flink-*-taskexecutor.pid, which + # only nudges the next local run's log-file index upward. That file is shared by + # any Flink cluster this user runs, so it is left alone rather than risking a + # concurrent cluster's shutdown for a cosmetic gain. + + # The work directory holds the verdict file, the handshake markers and the + # checkpoints — everything needed to explain a failure. Remove it only when the + # run both recorded assertions and is exiting cleanly; otherwise keep it and say + # where it is. A non-zero exit code is its own reason to keep: the setup steps and + # several assertions abort under set -e without recording anything, so judging by + # the recorded results alone would delete the diagnostics for exactly those. + if [[ -n "$WORK_DIR" && -d "$WORK_DIR" ]]; then + local keep=0 state + for state in "${RESULT_STATES[@]:-}"; do + [[ "$state" == "FAIL" ]] && keep=1 + done + (( ${#RESULT_STATES[@]} == 0 )) && keep=1 + (( exit_code != 0 )) && keep=1 + if (( keep == 1 )); then + log_warn "Keeping diagnostics in $WORK_DIR (verdict, handshake markers, checkpoints)" + else + rm -rf "$WORK_DIR" + fi + fi + + print_summary + exit "$exit_code" +} +trap cleanup EXIT + +print_summary() { + log_section "Test summary" + local total=${#RESULT_NAMES[@]} + if (( total == 0 )); then + # Exit non-zero rather than return: a run that recorded no assertion verified + # nothing, and reporting that as success is the precise failure this whole + # script exists to make impossible. + log_error "No assertion was recorded, so nothing was verified" + exit 1 + fi + local passed=0 + local failed=0 + local i + for (( i = 0; i < total; i++ )); do + local name="${RESULT_NAMES[$i]}" + local state="${RESULT_STATES[$i]}" + if [[ "$state" == "PASS" ]]; then + printf " ${GREEN}PASS${NC} %s\n" "$name" + passed=$((passed + 1)) + else + printf " ${RED}FAIL${NC} %s\n" "$name" + failed=$((failed + 1)) + fi + done + printf "\nTotal: %d Passed: %d Failed: %d\n" "$total" "$passed" "$failed" + + if (( failed > 0 )); then + log_error "$failed assertion(s) failed" + exit 1 + fi +} + +record_result() { + RESULT_NAMES+=("$1") + RESULT_STATES+=("$2") +} + +# --------------------------------------------------------------------------- +# JSON extraction. python3 rather than jq, whose presence on the runner is not +# guaranteed. +# +# Exit 1 means "I could not read this": empty body, unparsable JSON, a path that +# does not exist, or a leaf of the wrong type. Callers MUST keep that distinct +# from "I read a value and did not like it yet" — reporting the first as the +# second turns a wrong field name into a phantom timeout. +# --------------------------------------------------------------------------- +json_query() { + local mode="$1" path="$2" + python3 -c ' +import json +import sys + +mode, path = sys.argv[1], sys.argv[2] +raw = sys.stdin.read() +if not raw.strip(): + sys.exit(1) +try: + doc = json.loads(raw) +except ValueError: + sys.exit(1) + +# The cluster config endpoint answers with a list of {"key":..,"value":..} pairs, +# and its keys contain dots, so it cannot go through the path walker below. +if mode == "conf": + if not isinstance(doc, list): + sys.exit(1) + for item in doc: + if isinstance(item, dict) and item.get("key") == path: + value = item.get("value") + print("" if value is None else value) + sys.exit(0) + sys.exit(1) + +node = doc +for part in path.split("."): + if isinstance(node, list) and part.isdigit() and int(part) < len(node): + node = node[int(part)] + elif isinstance(node, dict) and part in node: + node = node[part] + else: + sys.exit(1) + +if mode == "nullable": + # The key must be present; only its value may be null. A missing key is a + # read failure, so a misspelled field cannot masquerade as "not yet". + print("null" if node is None else "present") +elif mode == "len": + if not isinstance(node, list): + sys.exit(1) + print(len(node)) +elif mode == "int": + if isinstance(node, bool) or not isinstance(node, int): + sys.exit(1) + print(node) +elif mode == "str": + if not isinstance(node, str): + sys.exit(1) + print(node) +else: + sys.exit(2) +' "$mode" "$path" +} + +rest_get() { + curl -fsS --max-time "$CURL_MAX_TIME" "$REST_URL$1" 2>/dev/null +} + +rest_field() { + local path="$1" mode="$2" field="$3" + rest_get "$path" | json_query "$mode" "$field" +} + +# --------------------------------------------------------------------------- +# Compare an observed value against a spec: eq:V, ge:N, nonnull. +# +# Returns 0 on a match, 1 on a mismatch, and 2 when the comparison cannot be +# evaluated at all — an unknown operator, or a non-integer operand for ge. Callers +# must keep 2 apart from 1: the first means the spec is wrong, the second means the +# condition is not satisfied yet. +# --------------------------------------------------------------------------- +value_matches() { + local observed="$1" spec="$2" + local op="${spec%%:*}" want="${spec#*:}" + case "$op" in + eq) [[ "$observed" == "$want" ]] ;; + ge) + # (( )) evaluates its operands as shell arithmetic, so a non-numeric + # observation would be treated as a variable name and abort the script + # under set -u instead of simply not matching. Return 2 so the caller can + # tell "cannot compare" from "does not match". + if [[ ! "$observed" =~ ^-?[0-9]+$ ]] || [[ ! "$want" =~ ^-?[0-9]+$ ]]; then + log_error "value_matches: 'ge' needs two integers, got observed='$observed' want='$want'" + return 2 + fi + (( observed >= want )) + ;; + nonnull) [[ "$observed" == "present" ]] ;; + *) log_error "value_matches: unknown comparison '$spec'"; return 2 ;; + esac +} + +# --------------------------------------------------------------------------- +# Poll a REST field until it satisfies a spec. +# +# $1 label human name of the step, used in both failure messages +# $2 path REST path, e.g. /jobs/<id>/checkpoints +# $3 probe "mode:field" that must ALWAYS parse on a healthy endpoint +# $4 target "mode:field" being tested +# $5 spec comparison, see value_matches +# $6 timeout seconds +# +# The probe is what separates the two failure modes. If the probe never parses we +# are not talking to the endpoint we think we are; if the probe parses but the +# target never does, the target's field name is wrong (or this Flink build omits +# it) — neither is evidence that the condition was false. +# --------------------------------------------------------------------------- +wait_for_rest() { + local label="$1" path="$2" probe="$3" target="$4" spec="$5" timeout="$6" + local probe_mode="${probe%%:*}" probe_field="${probe#*:}" + local target_mode="${target%%:*}" target_field="${target#*:}" + + local probe_seen=0 target_seen=0 last="" + local body value rc + # Deadline, not accumulated sleeps: curl and the two python3 spawns per iteration + # are not free, so counting only the sleeps would let a nominal budget overrun by + # several times against a slow endpoint — and these budgets are what keep the + # recovery inside the tool's own release deadline. + local deadline=$((SECONDS + timeout)) + log_info "$label: polling GET $path for $target_field ($spec), budget ${timeout}s" + while (( SECONDS < deadline )); do + body=$(rest_get "$path") || body="" + if [[ -n "$body" ]]; then + if printf '%s' "$body" | json_query "$probe_mode" "$probe_field" >/dev/null; then + probe_seen=1 + fi + if value=$(printf '%s' "$body" | json_query "$target_mode" "$target_field"); then + target_seen=1 + last="$value" + rc=0 + value_matches "$value" "$spec" || rc=$? + if (( rc == 0 )); then + log_ok "$label: $target_field is $value" + return 0 + elif (( rc > 1 )); then + log_error "$label: comparison '$spec' cannot be evaluated against '$value'" + return 1 + fi + fi + fi + sleep "$POLL_INTERVAL" + done + + # Order matters: an observed target is reported even when the probe never parsed, + # because in that state the condition really was evaluated and really was false, + # and claiming otherwise would withhold the one value a reader needs. + if (( target_seen == 1 )); then + log_error "$label: '$target_field' from GET $REST_URL$path was last '$last' after ${timeout}s and never satisfied '$spec'." + if (( probe_seen == 0 )); then + log_warn "$label: the probe field '$probe_field' never parsed, so the response shape is only partly as expected." + fi + elif (( probe_seen == 1 )); then + log_error "$label: parsed '$probe_field' from GET $REST_URL$path, so the endpoint is right, but never parsed '$target_field' in ${timeout}s. Either the field name is wrong or this Flink build omits it. This is NOT evidence that the condition was false." + else + log_error "$label: never parsed '$probe_field' or '$target_field' from GET $REST_URL$path in ${timeout}s. The endpoint did not answer, or its response shape is not what this script expects. This is NOT evidence that the condition was false." + fi + return 1 +} + +# --------------------------------------------------------------------------- +# Wait for an exact filename. Never a glob: every file the job publishes is +# written to "<name>.tmp" and renamed, so a crash between the two steps leaves a +# permanent ".tmp" twin that a glob would happily read as the real thing. +# --------------------------------------------------------------------------- +wait_for_file() { + local label="$1" file="$2" timeout="$3" + local deadline=$((SECONDS + timeout)) + log_info "$label: waiting for $file, budget ${timeout}s" + while (( SECONDS < deadline )); do + if [[ -f "$file" ]]; then + log_ok "$label: $file appeared" + return 0 + fi + sleep "$POLL_INTERVAL" + done + local dir + dir="$(dirname "$file")" + log_error "$label: $file did not appear within ${timeout}s" + if [[ -f "$file.tmp" ]]; then + log_error "$label: found a leftover $file.tmp — the writer was interrupted between write and rename" + fi + if [[ -d "$dir" ]]; then + local listing + listing=$(find "$dir" -maxdepth 1 -mindepth 1 2>/dev/null | sed 's|.*/||' | sort | tr '\n' ' ') + log_info "$label: $dir currently holds: $listing" + else + log_error "$label: the directory $dir does not exist" + fi + return 1 +} + +install_flink_distribution() { + local install_dir="$1" + + # Reuse install.sh's download, archive validation and extraction logic, but + # stop before it installs a *released* Flink Agents JAR. That step derives its + # download URL from the Flink version, and no such artifact is published for + # every Flink version this test runs against — asking for one that does not + # exist would abort the install. This test must run the artifacts built from + # the current checkout anyway, which stage_dist_jars puts in place below. + FLINK_AGENTS_INSTALL_SH_NO_RUN=1 \ + FLINK_VERSION="$FLINK_VERSION" \ + INSTALL_FLINK=Yes \ + INSTALL_DIR="$install_dir" \ + NO_PROMPT=1 \ + bash -c ' + source "$1" + plan_flink + install_flink_if_needed + ' _ "$ROOT_DIR/tools/install.sh" +} + +# Replaces install.sh's setup_python_env, which would populate the venv with a +# released flink-agents. This creates a plain venv and nothing more; the packages +# arrive later, from tools/build.sh's `uv pip install dist/*.whl` and from +# install_built_python_package. +prepare_python_venv() { + if [[ -f "$VENV_DIR/pyvenv.cfg" && -x "$VENV_DIR/bin/python" ]]; then + # Reuse is unconditional on the interpreter: unlike install.sh, this accepts + # a venv built by a Python outside the range python/pyproject.toml declares. + log_info "Reusing Python venv: $VENV_DIR" + return + fi + if [[ -e "$VENV_DIR" ]] \ + && [[ ! -d "$VENV_DIR" \ + || -n "$(find "$VENV_DIR" -mindepth 1 -print -quit 2>/dev/null)" ]]; then + log_error "VENV_DIR is not an empty directory or a valid venv: $VENV_DIR" + return 1 + fi + + local python_bin="${PYTHON_BIN:-python3}" + log_info "Creating Python venv: $VENV_DIR" + "$python_bin" -m venv "$VENV_DIR" +} + +install_flink() { + log_section "Install Flink standalone (version $FLINK_VERSION)" + + # Anchor VENV_DIR to the repo so a fresh and a reused Flink installation share + # one Python environment, and so later steps can address it by path. + export VENV_DIR="${VENV_DIR:-$ROOT_DIR/.flink-agents-env}" + + if [[ -n "${FLINK_HOME:-}" && -x "$FLINK_HOME/bin/flink" ]]; then + log_info "Reusing existing FLINK_HOME: $FLINK_HOME" + export FLINK_HOME + else + local install_dir="${INSTALL_DIR:-$HOME/.local/flink}" + log_info "Installing the Flink distribution with tools/install.sh helpers" + install_flink_distribution "$install_dir" + export FLINK_HOME="${install_dir}/flink-${FLINK_VERSION}" + + if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then + log_error "Flink installation not found at expected path: $FLINK_HOME" + exit 1 + fi + log_ok "Flink installed at: $FLINK_HOME" + fi + + # The distribution ships the PyFlink jar under opt/, outside the lib/ directory + # the cluster loads. install.sh copies it across as part of --enable-pyflink, + # which is no longer reached from here, so the copy happens here — on both branches + # above, since a reused FLINK_HOME need not have been prepared by install.sh. + local pyflink_jar="$FLINK_HOME/opt/flink-python-${FLINK_VERSION}.jar" + if [[ ! -f "$pyflink_jar" ]]; then + log_error "PyFlink JAR not found in Flink distribution: $pyflink_jar" + exit 1 + fi + cp "$pyflink_jar" "$FLINK_HOME/lib/" + + prepare_python_venv + + # Both branches above fall through to here, so the venv is activated whether + # FLINK_HOME was reused or freshly installed. That is load-bearing: the + # TaskManager daemon inherits this shell's PATH and resolves a bare `python` + # from it, and this script restarts the TaskManager mid-test — so an + # unactivated shell would give the replacement TaskManager a different + # interpreter from the original. Activation settles which interpreter that is + # and nothing more — what is installed in it is install_built_python_package's + # business, and that runs later. + if [[ ! -x "$VENV_DIR/bin/python" ]]; then + log_error "Expected Python venv not found at: $VENV_DIR" + exit 1 + fi + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" + log_ok "Activated Python venv: $VENV_DIR" + + FLINK_CONF="$FLINK_HOME/conf/config.yaml" + if [[ ! -f "$FLINK_CONF" ]]; then + log_error "Flink config not found: $FLINK_CONF" + exit 1 + fi +} + +# The wheel this run must exercise is the one tools/build.sh just produced, so it is +# installed from python/dist rather than from PyPI. apache-flink is pinned to the +# Flink version the cluster runs, because the submitting client and the TaskManager +# both execute this same interpreter. +install_built_python_package() { + local wheel + wheel=$(find "$ROOT_DIR/python/dist" -maxdepth 1 -name '*.whl' | head -n 1) + if [[ -z "$wheel" ]]; then + log_error "Python wheel not found after build in: $ROOT_DIR/python/dist" + return 1 + fi + + log_info "Installing $(basename "$wheel") and apache-flink==$FLINK_VERSION into $VENV_DIR" + "$VENV_DIR/bin/python" -m pip install --quiet \ + "$wheel" "apache-flink==$FLINK_VERSION" + + if ! "$VENV_DIR/bin/python" -c 'import flink_agents, pyflink' >/dev/null 2>&1; then + log_error "The built Flink Agents wheel or PyFlink is not importable from: $VENV_DIR/bin/python" + return 1 + fi + + export PYFLINK_CLIENT_EXECUTABLE="$VENV_DIR/bin/python" + log_ok "Installed the built wheel and PyFlink into: $VENV_DIR" +} + +build_project() { + log_section "Build flink-agents (Java + Python)" + ( + cd "$ROOT_DIR" + SKIP_SPOTLESS_CHECK=true bash tools/build.sh + ) + install_built_python_package + log_ok "Build completed" +} + +stage_dist_jars() { + log_section "Stage dist uber jar into \$FLINK_HOME/lib" + + local project_version + project_version=$(sed -n 's/.*<version>\(.*\)<\/version>.*/\1/p' \ + "$ROOT_DIR/pom.xml" | head -n 2 | tail -n 1) + log_info "Detected project version: $project_version" + + # The flink-version uber jar already bundles the common deps. + local flink_jar="$ROOT_DIR/dist/flink-${FLINK_MAJOR_MINOR}/target/flink-agents-dist-flink-${FLINK_MAJOR_MINOR}-${project_version}.jar" + + if [[ ! -f "$flink_jar" ]]; then + log_error "Flink dist jar not found: $flink_jar" + exit 1 + fi + + # Drop any flink-agents-dist jar already in lib/ before copying. A reused + # FLINK_HOME may still carry one — from an earlier tools/install.sh run, or from + # an earlier run of this test at a different project version — and a jar of a + # different version has a different filename, so the copy below would not + # replace it and both would sit on the classpath. + rm -f "$FLINK_HOME/lib/"/flink-agents-dist-*.jar + cp "$flink_jar" "$FLINK_HOME/lib/" + log_ok "Staged: $(basename "$flink_jar")" +} + +prepare_work_dirs() { + log_section "Create input, handshake, verdict and checkpoint directories" + WORK_DIR="$(mktemp -d "${TMPDIR:-/tmp}/flink-agents-recovery.XXXXXX")" + HANDSHAKE_DIR="$WORK_DIR/handshake" + VERDICT_DIR="$WORK_DIR/verdict" + CHECKPOINT_DIR="$WORK_DIR/checkpoints" + INPUT_DIR="$WORK_DIR/input" + # The name is load-bearing. Flink's default file filter skips any name beginning + # with '.' or '_', and applies that even to a path the source was handed directly, + # which would leave the job running with nothing to read and nothing logged. + INPUT_FILE="$INPUT_DIR/trigger.txt" + # The job must never create these itself: a tool that provisions its own + # handshake directory cannot tell "the harness set me up" from "my path is wrong". + mkdir -p "$HANDSHAKE_DIR" "$VERDICT_DIR" "$CHECKPOINT_DIR" "$INPUT_DIR" + + # One line is all the source needs: the job builds the record from a constant it + # shares with the assertion and ignores this content. Written under a temporary + # name and renamed, so the file is whole the instant it appears under the name the + # source reads. Submission is several steps away, and that distance is what + # actually orders the two — the rename is redundancy, not the ordering mechanism. + printf 'trigger\n' > "$INPUT_FILE.tmp" + mv "$INPUT_FILE.tmp" "$INPUT_FILE" + # An empty file produces no record and no error. Nothing that depends on a record + # then runs, the agent's identity report included, so the run surfaces as the + # runtime-identity wait expiring with nothing pointing at the input. A wrong path + # needs no guard here: enumeration is eager, so the job fails within seconds with + # "Could not enumerate file splits". + if [[ ! -s "$INPUT_FILE" ]]; then + log_error "Trigger file is missing or empty: $INPUT_FILE" + return 1 + fi + + log_ok "Work directory: $WORK_DIR" +} + +# --------------------------------------------------------------------------- +# Config editing. Delete-then-append, so applying the same key twice is a no-op +# rather than a duplicate YAML key. +# +# A flat dotted key at column 0 is read correctly regardless of the nested blocks +# the shipped file uses, because Flink flattens the parsed document before +# building its Configuration. Appending (never prepending) also means our value +# wins any collision with a nested block, since the flattened map is populated in +# document order. +# --------------------------------------------------------------------------- +delete_config_key() { + local config_key="$1" + local tmp + tmp="$(mktemp)" + # grep -v exits 1 when it filters everything out, which set -e would treat as + # fatal. The dots in the key are regex "any char"; harmless for these keys. + grep -v "^${config_key}: " "$FLINK_CONF" > "$tmp" || true + mv "$tmp" "$FLINK_CONF" +} + +set_config_key() { + local config_key="$1" value="$2" + delete_config_key "$config_key" Review Comment: Could we preserve and restore the original `config.yaml`? When reusing `FLINK_HOME`, `set_config_key` removes pre-existing values, while cleanup only removes the test values, so the original configuration is lost. Backing up the file before modification and restoring it on `EXIT` would avoid mutating the reused installation. ########## python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py: ########## @@ -0,0 +1,447 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF 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. +################################################################################# +"""Agent that parks a built-in tool-call round so a checkpoint can capture it. + +A deterministic mock chat model emits one tool call, and the tool then blocks on a +filesystem handshake, so the harness rather than the clock decides when the agent +run may finish. While the run is parked, the built-in tool-call context and a +``bytes`` short-term-memory value sit in checkpointable state; the harness kills and +restarts the TaskManager, and the assertions below then run in a TaskManager process +that never performed any of the writes. + +Submitted to a real cluster by ``checkpoint_recovery_job``. The module name is +deliberately neither ``*_test.py`` nor ``*_example.py``: the first would make pytest +collect it, the second would make the example-submission nightly submit it. +""" + +import json +import sys +import time +import uuid +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any, Dict, List, Sequence + +from pydantic import BaseModel +from pyflink.datastream import KeySelector + +import flink_agents.api.memory_object as memory_object_module +from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.decorators import ( + action, + chat_model_connection, + chat_model_setup, + prompt, + tool, +) +from flink_agents.api.events.chat_event import ChatRequestEvent, ChatResponseEvent +from flink_agents.api.events.event import Event, InputEvent, OutputEvent +from flink_agents.api.events.event_type import EventType +from flink_agents.api.prompts.prompt import Prompt +from flink_agents.api.resource import ResourceDescriptor +from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.tools import InjectedArg +from flink_agents.api.tools.tool import Tool as BaseTool +from flink_agents.api.tools.tool import ToolType + +# Config keys the submitting job must set; a missing one fails the run loudly. +HANDSHAKE_DIR_CONFIG_KEY = "handshake_dir" +VERDICT_DIR_CONFIG_KEY = "verdict_dir" + +# Filenames exchanged with the harness. +TOOL_ENTERED_MARKER = "tool-entered" +RELEASE_MARKER = "release" +IDENTITY_MARKER = "runtime-identity.json" +VERDICT_MARKER = "verdict.json" + +RELEASE_TIMEOUT_S = 240.0 +POLL_INTERVAL_S = 0.2 + +# The user message content, also built into the input record by the job. +USER_CONTENT = "hold the tool call open across a taskmanager restart" + +BLOCKING_TOOL_NAME = "block_until_released" + +_ROUND_ONE_MARKER = "recovery-round-one-assistant" +_BLOB_MEMORY_KEY = "blob" + +# Returned by the tool only on the released path, and required in the transcript. +# The framework converts a tool exception into the response string +# "Tool `block_until_released` execute failed." rather than failing the action, so a +# timed-out handshake reaches the assertions only as the absence of this value. That +# makes the value load-bearing: it must not be a substring of that failure string, +# which rules out anything built from the tool's own name. +_TOOL_SENTINEL = "handshake-done-sentinel" + +# A NUL plus bytes that are not valid UTF-8, so a value handled as UTF-8 text anywhere +# on the path either raises or comes back corrupted. A single-byte codec such as +# latin-1 would round trip it intact, so this catches UTF-8 handling rather than every +# possible string conversion. +_KNOWN_BLOB = b"\x00\x01flink-agents\xff\xfe" + +# Process-local and deliberately outside Flink state. Anything checkpointed is +# restored along with the payload and therefore cannot tell a restored read apart +# from the reader's own write. +# +# Pemja runs as ExecType.MULTI_THREAD against a singleton main interpreter, so these +# globals outlive an in-place task restart and are reset only when the TaskManager +# process itself dies. A read that returns the known value while they are still zero +# therefore proves the write happened in a different TaskManager process, which is +# stronger than proving it happened in a different interpreter. +_PROCESS_EPOCH = uuid.uuid4().hex +_BLOB_WRITES_IN_THIS_PROCESS = 0 +_TOOL_CALLS_EMITTED_IN_THIS_PROCESS = 0 + + +def _atomic_write(target: Path, content: str) -> None: + """Publish a file by renaming it into place so no reader sees a partial write.""" + tmp = target.parent / f"{target.name}.tmp" + tmp.write_text(content, encoding="utf-8") + tmp.replace(target) + + +def await_release( + handshake_dir: str, + *, + timeout_s: float = RELEASE_TIMEOUT_S, + poll_interval_s: float = POLL_INTERVAL_S, +) -> None: + """Announce that the tool was entered, then block until the harness releases it. + + The marker is written before the wait begins so the harness can order its kill + strictly after the tool-call context reached checkpointable state. + + Raises ``TimeoutError`` at the deadline so the wait is bounded and the reason + reaches the TaskManager log. The raise has no effect on control flow: the caller + converts every tool exception into an ordinary tool response, so raising and + returning drive the run identically. What makes a timeout observable is the + sentinel the tool returns only on this function's success path. + + Parameters + ---------- + handshake_dir : str + Directory the harness created before submitting the job. + timeout_s : float + How long to wait for the release marker before failing the run. + poll_interval_s : float + Interval between existence checks. + """ + base = Path(handshake_dir) + _atomic_write(base / TOOL_ENTERED_MARKER, _PROCESS_EPOCH) + + release = base / RELEASE_MARKER + deadline = time.monotonic() + timeout_s + while not release.exists(): + if time.monotonic() > deadline: + msg = f"waited {timeout_s}s for the release marker at {release}" + raise TimeoutError(msg) + time.sleep(poll_interval_s) + + +def _flink_agents_version() -> str: + """Report the installed distribution version, or ``unknown`` if there is none.""" + try: + return version("flink-agents") + except PackageNotFoundError: + return "unknown" + + +def _runtime_identity() -> Dict[str, Any]: + """Describe the flink-agents installation this process is actually running. + + The probe is ``flink_agents.api.memory_object`` rather than the ``flink_agents`` + package. ``flink_agents/__init__.py`` is a single ``pkgutil.extend_path`` call, so + ``flink_agents.__file__`` names whichever ``__init__.py`` came first on the path + while ``flink_agents.api`` can be served from a different entry. ``api`` extends + nothing, so a module under it pins the installation that supplied the code under + test. + """ + return { + "flink_agents_api_file": memory_object_module.__file__, + "flink_agents_version": _flink_agents_version(), + "python_executable": sys.executable, + "python_version": sys.version, + "process_epoch": _PROCESS_EPOCH, + } + + +def _write_runtime_identity(handshake_dir: str) -> None: + """Publish the pre-kill identity so the harness can fail before it kills anything. + + This copy always describes the process that is about to be killed. The copy the + assertions rely on rides in the verdict record, which the surviving process + writes. + """ + _atomic_write( + Path(handshake_dir) / IDENTITY_MARKER, + json.dumps(_runtime_identity(), sort_keys=True), + ) + + +def _mark_blob_written() -> None: + global _BLOB_WRITES_IN_THIS_PROCESS + _BLOB_WRITES_IN_THIS_PROCESS += 1 + + +def _mark_tool_call_emitted() -> None: + global _TOOL_CALLS_EMITTED_IN_THIS_PROCESS + _TOOL_CALLS_EMITTED_IN_THIS_PROCESS += 1 + + +def _blob_matches(raw: Any) -> bool: + """Check the value is ``bytes`` or ``bytearray`` holding the known blob. + + The type gate is part of the assertion, not defensive coding. ``bytes()`` + accepts any iterable of ints, so without it a ``byte[]`` materialized as + ``[0, 1, 102, ...]`` would compare equal to the blob and pass — and this + predicate is the one whose bug produces a false pass. + + ``bytes`` is admitted because it is the only one of the two the memory value + validator accepts at write time. ``bytearray`` is admitted because the value + comes back across the bridge from a Java ``byte[]`` and may materialize as + either — a read-side allowance this test makes, not one the write-side contract + grants. Nothing else is admitted on the chance it might appear; anything else + fails here rather than aborting the action, so the run still publishes a verdict + and ``blob_observed_type`` records what did come back. + """ + if not isinstance(raw, bytes | bytearray): + return False + return bytes(raw) == _KNOWN_BLOB + + +def _required_config(ctx: RunnerContext, key: str) -> str: + value = ctx.config.get_str(key) + if not value: + msg = f"Missing config for the checkpoint recovery job: {key}" + raise ValueError(msg) + return value + + +class CheckpointRecoveryInput(BaseModel): + """Input record for the checkpoint recovery agent. + + Attributes: + ---------- + id : int + Unique identifier used as the partition key. + content : str + The user message content fed to the agent. + """ + + id: int + content: str + + +class CheckpointRecoveryKeySelector(KeySelector): + """KeySelector extracting the partition key from a CheckpointRecoveryInput.""" + + def get_key(self, value: CheckpointRecoveryInput) -> int: + """Extract key from CheckpointRecoveryInput.""" + return value.id + + +class RecoveryMockChatConnection(BaseChatModelConnection): + """Mock connection emitting one tool call, then joining the whole transcript.""" + + def chat( + self, + messages: Sequence[ChatMessage], + tools: List[BaseTool] | None = None, + output_schema: OutputSchema | None = None, + **kwargs: Any, + ) -> ChatMessage: + """Request the blocking tool, or join every message once the tool replied. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + structured-output translation. Declaring the parameter keeps a caller-supplied + schema out of ``**kwargs``. + """ + self._reject_unsupported_output_schema(output_schema) + if messages[-1].role == MessageRole.TOOL: + # Joining every message carries the rebuilt transcript out to the + # emitted content, which is where the assertion can reach it. + content = "\n".join(message.content for message in messages) + return ChatMessage(role=MessageRole.ASSISTANT, content=content) + + # Validate the tool was bound before the model was invoked. + assert tools[0].name == BLOCKING_TOOL_NAME + _mark_tool_call_emitted() + tool_call = { + "id": str(uuid.uuid4()), + "type": ToolType.FUNCTION, + "function": {"name": BLOCKING_TOOL_NAME, "arguments": {}}, + } + return ChatMessage( + role=MessageRole.ASSISTANT, + content=_ROUND_ONE_MARKER, + tool_calls=[tool_call], + ) + + +class RecoveryMockChatModel(BaseChatModelSetup): + """Mock chat model setup for the checkpoint recovery agent.""" + + @property + def model_kwargs(self) -> Dict[str, Any]: + """Return model kwargs.""" + return {} + + +class CheckpointRecoveryAgent(Agent): + """Agent held mid-tool-call so a checkpoint captures its memory.""" + + @prompt + @staticmethod + def recovery_prompt() -> Prompt: + """Prompt used by the mock chat model.""" + return Prompt.from_text( + text="Please call the appropriate tool to do the following task: {task}", + ) + + @chat_model_connection + @staticmethod + def recovery_connection() -> ResourceDescriptor: + """Chat model connection used by the mock chat model.""" + return ResourceDescriptor( + clazz=f"{RecoveryMockChatConnection.__module__}." + f"{RecoveryMockChatConnection.__name__}" + ) + + @chat_model_setup + @staticmethod + def recovery_chat_model() -> ResourceDescriptor: + """Chat model referenced by the ChatRequestEvent.""" + return ResourceDescriptor( + clazz=f"{RecoveryMockChatModel.__module__}." + f"{RecoveryMockChatModel.__name__}", + connection="recovery_connection", + model="mock-model", + prompt="recovery_prompt", + tools=[BLOCKING_TOOL_NAME], + ) + + @tool( + injected_args={ + "handshake_dir": InjectedArg.from_config(HANDSHAKE_DIR_CONFIG_KEY) + } + ) + @staticmethod + def block_until_released(handshake_dir: str) -> str: + """Hold the agent run open until the harness releases it. + + Takes no model-visible arguments, so the mock never has to fabricate one. + The post-restore re-execution finds the release marker already present and + returns immediately, which makes the handshake idempotent. + + Parameters + ---------- + handshake_dir : str + The handshake directory, injected by runtime. + + Returns: + ------- + str: + The sentinel, which reaches the transcript as the tool message and is + required by the assertion. A timed-out handshake never gets here, so the + transcript carries the framework's failure string instead. + """ + await_release(handshake_dir) + return _TOOL_SENTINEL + + @action(EventType.InputEvent) + @staticmethod + def process_input(event: Event, ctx: RunnerContext) -> None: + """Record the payload and start the tool-call round.""" + input_data = CheckpointRecoveryInput.model_validate( + InputEvent.from_event(event).input + ) + # Resolve both directories up front: a misspelled verdict_dir would otherwise + # surface only after the whole park/kill/restart cycle has been paid for. + handshake_dir = _required_config(ctx, HANDSHAKE_DIR_CONFIG_KEY) + _required_config(ctx, VERDICT_DIR_CONFIG_KEY) + _write_runtime_identity(handshake_dir) + + # Carry the record id across the chat round-trip via per-key memory, + # since ChatResponseEvent does not echo the original input. + ctx.short_term_memory.set("input_id", input_data.id) + ctx.short_term_memory.set(_BLOB_MEMORY_KEY, _KNOWN_BLOB) + _mark_blob_written() + + ctx.send_event( + ChatRequestEvent( + model="recovery_chat_model", + messages=[ + ChatMessage(role=MessageRole.USER, content=input_data.content) + ], + prompt_args={"task": input_data.content}, Review Comment: This request leaves `output_schema` unset, so the E2E only exercises the `None` branch. Since #828 specifically normalized `OutputSchema` for checkpoint safety, a regression in its serialization or reconstruction would remain undetected. Could we include a non-null `OutputSchema` in this recovery flow? ########## python/flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_agent.py: ########## @@ -0,0 +1,447 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF 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. +################################################################################# +"""Agent that parks a built-in tool-call round so a checkpoint can capture it. + +A deterministic mock chat model emits one tool call, and the tool then blocks on a +filesystem handshake, so the harness rather than the clock decides when the agent +run may finish. While the run is parked, the built-in tool-call context and a +``bytes`` short-term-memory value sit in checkpointable state; the harness kills and +restarts the TaskManager, and the assertions below then run in a TaskManager process +that never performed any of the writes. + +Submitted to a real cluster by ``checkpoint_recovery_job``. The module name is +deliberately neither ``*_test.py`` nor ``*_example.py``: the first would make pytest +collect it, the second would make the example-submission nightly submit it. +""" + +import json +import sys +import time +import uuid +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any, Dict, List, Sequence + +from pydantic import BaseModel +from pyflink.datastream import KeySelector + +import flink_agents.api.memory_object as memory_object_module +from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.decorators import ( + action, + chat_model_connection, + chat_model_setup, + prompt, + tool, +) +from flink_agents.api.events.chat_event import ChatRequestEvent, ChatResponseEvent +from flink_agents.api.events.event import Event, InputEvent, OutputEvent +from flink_agents.api.events.event_type import EventType +from flink_agents.api.prompts.prompt import Prompt +from flink_agents.api.resource import ResourceDescriptor +from flink_agents.api.runner_context import RunnerContext +from flink_agents.api.tools import InjectedArg +from flink_agents.api.tools.tool import Tool as BaseTool +from flink_agents.api.tools.tool import ToolType + +# Config keys the submitting job must set; a missing one fails the run loudly. +HANDSHAKE_DIR_CONFIG_KEY = "handshake_dir" +VERDICT_DIR_CONFIG_KEY = "verdict_dir" + +# Filenames exchanged with the harness. +TOOL_ENTERED_MARKER = "tool-entered" +RELEASE_MARKER = "release" +IDENTITY_MARKER = "runtime-identity.json" +VERDICT_MARKER = "verdict.json" + +RELEASE_TIMEOUT_S = 240.0 +POLL_INTERVAL_S = 0.2 + +# The user message content, also built into the input record by the job. +USER_CONTENT = "hold the tool call open across a taskmanager restart" + +BLOCKING_TOOL_NAME = "block_until_released" + +_ROUND_ONE_MARKER = "recovery-round-one-assistant" +_BLOB_MEMORY_KEY = "blob" + +# Returned by the tool only on the released path, and required in the transcript. +# The framework converts a tool exception into the response string +# "Tool `block_until_released` execute failed." rather than failing the action, so a +# timed-out handshake reaches the assertions only as the absence of this value. That +# makes the value load-bearing: it must not be a substring of that failure string, +# which rules out anything built from the tool's own name. +_TOOL_SENTINEL = "handshake-done-sentinel" + +# A NUL plus bytes that are not valid UTF-8, so a value handled as UTF-8 text anywhere +# on the path either raises or comes back corrupted. A single-byte codec such as +# latin-1 would round trip it intact, so this catches UTF-8 handling rather than every +# possible string conversion. +_KNOWN_BLOB = b"\x00\x01flink-agents\xff\xfe" + +# Process-local and deliberately outside Flink state. Anything checkpointed is +# restored along with the payload and therefore cannot tell a restored read apart +# from the reader's own write. +# +# Pemja runs as ExecType.MULTI_THREAD against a singleton main interpreter, so these +# globals outlive an in-place task restart and are reset only when the TaskManager +# process itself dies. A read that returns the known value while they are still zero +# therefore proves the write happened in a different TaskManager process, which is +# stronger than proving it happened in a different interpreter. +_PROCESS_EPOCH = uuid.uuid4().hex +_BLOB_WRITES_IN_THIS_PROCESS = 0 +_TOOL_CALLS_EMITTED_IN_THIS_PROCESS = 0 + + +def _atomic_write(target: Path, content: str) -> None: + """Publish a file by renaming it into place so no reader sees a partial write.""" + tmp = target.parent / f"{target.name}.tmp" + tmp.write_text(content, encoding="utf-8") + tmp.replace(target) + + +def await_release( + handshake_dir: str, + *, + timeout_s: float = RELEASE_TIMEOUT_S, + poll_interval_s: float = POLL_INTERVAL_S, +) -> None: + """Announce that the tool was entered, then block until the harness releases it. + + The marker is written before the wait begins so the harness can order its kill + strictly after the tool-call context reached checkpointable state. + + Raises ``TimeoutError`` at the deadline so the wait is bounded and the reason + reaches the TaskManager log. The raise has no effect on control flow: the caller + converts every tool exception into an ordinary tool response, so raising and + returning drive the run identically. What makes a timeout observable is the + sentinel the tool returns only on this function's success path. + + Parameters + ---------- + handshake_dir : str + Directory the harness created before submitting the job. + timeout_s : float + How long to wait for the release marker before failing the run. + poll_interval_s : float + Interval between existence checks. + """ + base = Path(handshake_dir) + _atomic_write(base / TOOL_ENTERED_MARKER, _PROCESS_EPOCH) + + release = base / RELEASE_MARKER + deadline = time.monotonic() + timeout_s + while not release.exists(): + if time.monotonic() > deadline: + msg = f"waited {timeout_s}s for the release marker at {release}" + raise TimeoutError(msg) + time.sleep(poll_interval_s) + + +def _flink_agents_version() -> str: + """Report the installed distribution version, or ``unknown`` if there is none.""" + try: + return version("flink-agents") + except PackageNotFoundError: + return "unknown" + + +def _runtime_identity() -> Dict[str, Any]: + """Describe the flink-agents installation this process is actually running. + + The probe is ``flink_agents.api.memory_object`` rather than the ``flink_agents`` + package. ``flink_agents/__init__.py`` is a single ``pkgutil.extend_path`` call, so + ``flink_agents.__file__`` names whichever ``__init__.py`` came first on the path + while ``flink_agents.api`` can be served from a different entry. ``api`` extends + nothing, so a module under it pins the installation that supplied the code under + test. + """ + return { + "flink_agents_api_file": memory_object_module.__file__, + "flink_agents_version": _flink_agents_version(), + "python_executable": sys.executable, + "python_version": sys.version, + "process_epoch": _PROCESS_EPOCH, + } + + +def _write_runtime_identity(handshake_dir: str) -> None: + """Publish the pre-kill identity so the harness can fail before it kills anything. + + This copy always describes the process that is about to be killed. The copy the + assertions rely on rides in the verdict record, which the surviving process + writes. + """ + _atomic_write( + Path(handshake_dir) / IDENTITY_MARKER, + json.dumps(_runtime_identity(), sort_keys=True), + ) + + +def _mark_blob_written() -> None: + global _BLOB_WRITES_IN_THIS_PROCESS + _BLOB_WRITES_IN_THIS_PROCESS += 1 + + +def _mark_tool_call_emitted() -> None: + global _TOOL_CALLS_EMITTED_IN_THIS_PROCESS + _TOOL_CALLS_EMITTED_IN_THIS_PROCESS += 1 + + +def _blob_matches(raw: Any) -> bool: + """Check the value is ``bytes`` or ``bytearray`` holding the known blob. + + The type gate is part of the assertion, not defensive coding. ``bytes()`` + accepts any iterable of ints, so without it a ``byte[]`` materialized as + ``[0, 1, 102, ...]`` would compare equal to the blob and pass — and this + predicate is the one whose bug produces a false pass. + + ``bytes`` is admitted because it is the only one of the two the memory value + validator accepts at write time. ``bytearray`` is admitted because the value + comes back across the bridge from a Java ``byte[]`` and may materialize as + either — a read-side allowance this test makes, not one the write-side contract + grants. Nothing else is admitted on the chance it might appear; anything else + fails here rather than aborting the action, so the run still publishes a verdict + and ``blob_observed_type`` records what did come back. + """ + if not isinstance(raw, bytes | bytearray): Review Comment: Could we require exact `bytes` here? The memory contract deliberately accepts `bytes` but rejects `bytearray`. With the current check, a restored `bytearray` with the same content would still pass. Using `type(raw) is bytes and raw == _KNOWN_BLOB` would make this a true `bytes` round-trip assertion. ########## e2e-test/test-scripts/test_checkpoint_recovery.sh: ########## @@ -0,0 +1,1307 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF 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. +# +# +# Verifies that Python agent memory survives the loss of a TaskManager process. +# +# Submits checkpoint_recovery_job.py to a local Flink standalone cluster. The job +# parks itself inside a tool call that blocks on a file this script creates, so a +# completed checkpoint provably holds the agent's memory before anything is killed. +# The script then hard-kills the TaskManager, restarts it, waits for a real restore, +# releases the tool and reads the verdict the job publishes. +# +# Unlike its sibling test_submit_examples_to_flink.sh, a successful submission is +# NOT a pass: the only pass is a verdict file that says so. +# +# Env: FLINK_VERSION (default 2.3.0), FLINK_HOME (reuse existing install), +# VERBOSE=1 (set -x), plus the *_TIMEOUT overrides below. +# +# A FLINK_HOME passed in has to be that same FLINK_VERSION. The run copies +# opt/flink-python-<FLINK_VERSION>.jar into its lib/ and stops when that file is +# not there, so point FLINK_VERSION at whatever installation FLINK_HOME names. + +set -euo pipefail + +if [[ "${VERBOSE:-0}" == "1" ]]; then + set -x +fi + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[0;33m' +BLUE='\033[0;34m' +NC='\033[0m' + +log_info() { printf "${BLUE}[INFO]${NC} %s\n" "$*" >&2; } +log_ok() { printf "${GREEN}[OK]${NC} %s\n" "$*" >&2; } +log_warn() { printf "${YELLOW}[WARN]${NC} %s\n" "$*" >&2; } +log_error() { printf "${RED}[ERROR]${NC} %s\n" "$*" >&2; } +log_section() { + printf "\n${BLUE}==============================================================${NC}\n" >&2 + printf "${BLUE}>>> %s${NC}\n" "$*" >&2 + printf "${BLUE}==============================================================${NC}\n" >&2 +} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/../.."; pwd)" +log_info "Project root: $ROOT_DIR" + +FLINK_VERSION="${FLINK_VERSION:-2.3.0}" +FLINK_MAJOR_MINOR="${FLINK_VERSION%.*}" +REST_URL="${REST_URL:-http://localhost:8081}"; + +JOB_MODULE="flink_agents/e2e_tests/e2e_tests_integration/checkpoint_recovery_job.py" +EXPECTED_AGENTS_VERSION="${EXPECTED_AGENTS_VERSION:-0.3.dev0}" + +# Checkpoint interval is deliberately short: the run is parked while we wait for two +# checkpoints to complete, and that wait is charged against the tool's own deadline. +# Held in milliseconds because that is the unit the checkpoint-config endpoint reports, +# so the value written and the value asserted are the same number. +CHECKPOINT_INTERVAL_MS="${CHECKPOINT_INTERVAL_MS:-5000}" +RESTART_ATTEMPTS="${RESTART_ATTEMPTS:-3}" +# Unset, this falls back to slot.request.timeout (5 min), after which a pending slot +# request is failed and one restart attempt is burned while the TM is being replaced. +STANDALONE_STARTUP_TIME="${STANDALONE_STARTUP_TIME:-600s}" + +# Budgets. Setup waits are unconstrained; the four marked ones run after the tool has +# parked and are therefore charged against the tool's own release deadline. They are +# clamped at runtime to the time actually remaining (see charged_timeout), so these +# numbers are ceilings rather than guarantees. +CLUSTER_TIMEOUT="${CLUSTER_TIMEOUT:-120}" +SUBMIT_TIMEOUT="${SUBMIT_TIMEOUT:-300}" +JOB_RUNNING_TIMEOUT="${JOB_RUNNING_TIMEOUT:-180}" +IDENTITY_TIMEOUT="${IDENTITY_TIMEOUT:-180}" +TOOL_ENTERED_TIMEOUT="${TOOL_ENTERED_TIMEOUT:-300}" +CHECKPOINT_TIMEOUT="${CHECKPOINT_TIMEOUT:-30}" # charged against the tool deadline +TM_GONE_TIMEOUT="${TM_GONE_TIMEOUT:-45}" # charged against the tool deadline +TM_UP_TIMEOUT="${TM_UP_TIMEOUT:-30}" # charged against the tool deadline +RESTORE_TIMEOUT="${RESTORE_TIMEOUT:-30}" # charged against the tool deadline +VERDICT_TIMEOUT="${VERDICT_TIMEOUT:-120}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +# Every REST call is bounded. This is charged in the budget arithmetic, so keep the +# two in step: a healthy local JobManager answers in milliseconds, and a large value +# here buys nothing while inflating the worst-case overrun of every wait. +CURL_MAX_TIME="${CURL_MAX_TIME:-5}" + +# Wall-clock costs inside the charged window that belong to no single wait: the +# individual REST reads between the waits, two jps invocations, the SIGKILL grace and +# starting the replacement TaskManager. This is an estimate, and deliberately only +# feeds the pre-flight feasibility check — the runtime clamp is what actually holds +# the guarantee, so being wrong here cannot let the tool self-release. +HANDSHAKE_FIXED_COST_S="${HANDSHAKE_FIXED_COST_S:-45}" + +# Kept in hand so `release` is written comfortably before the tool gives up, covering +# the tool's own poll interval and the one-second granularity of SECONDS. +HANDSHAKE_SAFETY_S="${HANDSHAKE_SAFETY_S:-15}" + +# Bash 3 (default on macOS) lacks associative arrays. +RESULT_NAMES=() +RESULT_STATES=() +SUBMITTED_JOB_IDS=() +CONFIG_KEYS_SET=() +JOB_ID="" +WORK_DIR="" +HANDSHAKE_DIR="" +VERDICT_DIR="" +CHECKPOINT_DIR="" +INPUT_DIR="" +INPUT_FILE="" +FLINK_CONF="" +TM_PID_BEFORE="" +TM_RESOURCE_ID_BEFORE="" +RESTORED_BEFORE="" +RELEASE_DEADLINE_S="" +HANDSHAKE_DEADLINE_AT="" + +cleanup() { + local exit_code=$? + log_section "Cleanup" + + if [[ -n "${FLINK_HOME:-}" && -x "$FLINK_HOME/bin/flink" ]]; then + for jid in "${SUBMITTED_JOB_IDS[@]:-}"; do + [[ -n "$jid" ]] || continue + log_info "Cancelling job $jid" + "$FLINK_HOME/bin/flink" cancel "$jid" >/dev/null 2>&1 || true + done + + if [[ -x "$FLINK_HOME/bin/stop-cluster.sh" ]]; then + # If the TaskManager was killed and never replaced, this prints + # "No taskexecutor daemon (pid: N) is running anymore" into the archive. + # That line is expected here and is not the failure. + log_info "Stopping Flink cluster" + "$FLINK_HOME/bin/stop-cluster.sh" >/dev/null 2>&1 || true + fi + + # Archive the whole log directory, never named files: `kill -9` leaves the + # dead TaskManager's pid in the pid file, so the replacement writes to a + # higher log index and a named copy would miss the post-restore TaskManager. + if [[ -d "$FLINK_HOME/log" ]]; then + local log_archive="$ROOT_DIR/flink-logs-$(date +%Y%m%d-%H%M%S).tar.gz" + tar -czf "$log_archive" -C "$FLINK_HOME" log >/dev/null 2>&1 \ + && log_info "Flink logs archived to: $log_archive" \ + || log_warn "Failed to archive Flink logs" + fi + fi + + # Strip the keys we appended. A second run against a reused Flink home would + # otherwise append them twice, and a duplicate key is a hard YAML parse failure + # that stops the cluster from starting at all. + if [[ -n "$FLINK_CONF" && -f "$FLINK_CONF" ]]; then + local key + for key in "${CONFIG_KEYS_SET[@]:-}"; do + [[ -n "$key" ]] || continue + delete_config_key "$key" + done + log_info "Reverted ${#CONFIG_KEYS_SET[@]} appended config key(s)" + fi + + # The killed TaskManager's entry stays in /tmp/flink-*-taskexecutor.pid, which + # only nudges the next local run's log-file index upward. That file is shared by + # any Flink cluster this user runs, so it is left alone rather than risking a + # concurrent cluster's shutdown for a cosmetic gain. + + # The work directory holds the verdict file, the handshake markers and the + # checkpoints — everything needed to explain a failure. Remove it only when the + # run both recorded assertions and is exiting cleanly; otherwise keep it and say + # where it is. A non-zero exit code is its own reason to keep: the setup steps and + # several assertions abort under set -e without recording anything, so judging by + # the recorded results alone would delete the diagnostics for exactly those. + if [[ -n "$WORK_DIR" && -d "$WORK_DIR" ]]; then + local keep=0 state + for state in "${RESULT_STATES[@]:-}"; do + [[ "$state" == "FAIL" ]] && keep=1 + done + (( ${#RESULT_STATES[@]} == 0 )) && keep=1 + (( exit_code != 0 )) && keep=1 + if (( keep == 1 )); then + log_warn "Keeping diagnostics in $WORK_DIR (verdict, handshake markers, checkpoints)" + else + rm -rf "$WORK_DIR" + fi + fi + + print_summary + exit "$exit_code" +} +trap cleanup EXIT + +print_summary() { + log_section "Test summary" + local total=${#RESULT_NAMES[@]} + if (( total == 0 )); then + # Exit non-zero rather than return: a run that recorded no assertion verified + # nothing, and reporting that as success is the precise failure this whole + # script exists to make impossible. + log_error "No assertion was recorded, so nothing was verified" + exit 1 + fi + local passed=0 + local failed=0 + local i + for (( i = 0; i < total; i++ )); do + local name="${RESULT_NAMES[$i]}" + local state="${RESULT_STATES[$i]}" + if [[ "$state" == "PASS" ]]; then + printf " ${GREEN}PASS${NC} %s\n" "$name" + passed=$((passed + 1)) + else + printf " ${RED}FAIL${NC} %s\n" "$name" + failed=$((failed + 1)) + fi + done + printf "\nTotal: %d Passed: %d Failed: %d\n" "$total" "$passed" "$failed" + + if (( failed > 0 )); then + log_error "$failed assertion(s) failed" + exit 1 + fi +} + +record_result() { + RESULT_NAMES+=("$1") + RESULT_STATES+=("$2") +} + +# --------------------------------------------------------------------------- +# JSON extraction. python3 rather than jq, whose presence on the runner is not +# guaranteed. +# +# Exit 1 means "I could not read this": empty body, unparsable JSON, a path that +# does not exist, or a leaf of the wrong type. Callers MUST keep that distinct +# from "I read a value and did not like it yet" — reporting the first as the +# second turns a wrong field name into a phantom timeout. +# --------------------------------------------------------------------------- +json_query() { + local mode="$1" path="$2" + python3 -c ' +import json +import sys + +mode, path = sys.argv[1], sys.argv[2] +raw = sys.stdin.read() +if not raw.strip(): + sys.exit(1) +try: + doc = json.loads(raw) +except ValueError: + sys.exit(1) + +# The cluster config endpoint answers with a list of {"key":..,"value":..} pairs, +# and its keys contain dots, so it cannot go through the path walker below. +if mode == "conf": + if not isinstance(doc, list): + sys.exit(1) + for item in doc: + if isinstance(item, dict) and item.get("key") == path: + value = item.get("value") + print("" if value is None else value) + sys.exit(0) + sys.exit(1) + +node = doc +for part in path.split("."): + if isinstance(node, list) and part.isdigit() and int(part) < len(node): + node = node[int(part)] + elif isinstance(node, dict) and part in node: + node = node[part] + else: + sys.exit(1) + +if mode == "nullable": + # The key must be present; only its value may be null. A missing key is a + # read failure, so a misspelled field cannot masquerade as "not yet". + print("null" if node is None else "present") +elif mode == "len": + if not isinstance(node, list): + sys.exit(1) + print(len(node)) +elif mode == "int": + if isinstance(node, bool) or not isinstance(node, int): + sys.exit(1) + print(node) +elif mode == "str": + if not isinstance(node, str): + sys.exit(1) + print(node) +else: + sys.exit(2) +' "$mode" "$path" +} + +rest_get() { + curl -fsS --max-time "$CURL_MAX_TIME" "$REST_URL$1" 2>/dev/null +} + +rest_field() { + local path="$1" mode="$2" field="$3" + rest_get "$path" | json_query "$mode" "$field" +} + +# --------------------------------------------------------------------------- +# Compare an observed value against a spec: eq:V, ge:N, nonnull. +# +# Returns 0 on a match, 1 on a mismatch, and 2 when the comparison cannot be +# evaluated at all — an unknown operator, or a non-integer operand for ge. Callers +# must keep 2 apart from 1: the first means the spec is wrong, the second means the +# condition is not satisfied yet. +# --------------------------------------------------------------------------- +value_matches() { + local observed="$1" spec="$2" + local op="${spec%%:*}" want="${spec#*:}" + case "$op" in + eq) [[ "$observed" == "$want" ]] ;; + ge) + # (( )) evaluates its operands as shell arithmetic, so a non-numeric + # observation would be treated as a variable name and abort the script + # under set -u instead of simply not matching. Return 2 so the caller can + # tell "cannot compare" from "does not match". + if [[ ! "$observed" =~ ^-?[0-9]+$ ]] || [[ ! "$want" =~ ^-?[0-9]+$ ]]; then + log_error "value_matches: 'ge' needs two integers, got observed='$observed' want='$want'" + return 2 + fi + (( observed >= want )) + ;; + nonnull) [[ "$observed" == "present" ]] ;; + *) log_error "value_matches: unknown comparison '$spec'"; return 2 ;; + esac +} + +# --------------------------------------------------------------------------- +# Poll a REST field until it satisfies a spec. +# +# $1 label human name of the step, used in both failure messages +# $2 path REST path, e.g. /jobs/<id>/checkpoints +# $3 probe "mode:field" that must ALWAYS parse on a healthy endpoint +# $4 target "mode:field" being tested +# $5 spec comparison, see value_matches +# $6 timeout seconds +# +# The probe is what separates the two failure modes. If the probe never parses we +# are not talking to the endpoint we think we are; if the probe parses but the +# target never does, the target's field name is wrong (or this Flink build omits +# it) — neither is evidence that the condition was false. +# --------------------------------------------------------------------------- +wait_for_rest() { + local label="$1" path="$2" probe="$3" target="$4" spec="$5" timeout="$6" + local probe_mode="${probe%%:*}" probe_field="${probe#*:}" + local target_mode="${target%%:*}" target_field="${target#*:}" + + local probe_seen=0 target_seen=0 last="" + local body value rc + # Deadline, not accumulated sleeps: curl and the two python3 spawns per iteration + # are not free, so counting only the sleeps would let a nominal budget overrun by + # several times against a slow endpoint — and these budgets are what keep the + # recovery inside the tool's own release deadline. + local deadline=$((SECONDS + timeout)) + log_info "$label: polling GET $path for $target_field ($spec), budget ${timeout}s" + while (( SECONDS < deadline )); do + body=$(rest_get "$path") || body="" + if [[ -n "$body" ]]; then + if printf '%s' "$body" | json_query "$probe_mode" "$probe_field" >/dev/null; then + probe_seen=1 + fi + if value=$(printf '%s' "$body" | json_query "$target_mode" "$target_field"); then + target_seen=1 + last="$value" + rc=0 + value_matches "$value" "$spec" || rc=$? + if (( rc == 0 )); then + log_ok "$label: $target_field is $value" + return 0 + elif (( rc > 1 )); then + log_error "$label: comparison '$spec' cannot be evaluated against '$value'" + return 1 + fi + fi + fi + sleep "$POLL_INTERVAL" + done + + # Order matters: an observed target is reported even when the probe never parsed, + # because in that state the condition really was evaluated and really was false, + # and claiming otherwise would withhold the one value a reader needs. + if (( target_seen == 1 )); then + log_error "$label: '$target_field' from GET $REST_URL$path was last '$last' after ${timeout}s and never satisfied '$spec'." + if (( probe_seen == 0 )); then + log_warn "$label: the probe field '$probe_field' never parsed, so the response shape is only partly as expected." + fi + elif (( probe_seen == 1 )); then + log_error "$label: parsed '$probe_field' from GET $REST_URL$path, so the endpoint is right, but never parsed '$target_field' in ${timeout}s. Either the field name is wrong or this Flink build omits it. This is NOT evidence that the condition was false." + else + log_error "$label: never parsed '$probe_field' or '$target_field' from GET $REST_URL$path in ${timeout}s. The endpoint did not answer, or its response shape is not what this script expects. This is NOT evidence that the condition was false." + fi + return 1 +} + +# --------------------------------------------------------------------------- +# Wait for an exact filename. Never a glob: every file the job publishes is +# written to "<name>.tmp" and renamed, so a crash between the two steps leaves a +# permanent ".tmp" twin that a glob would happily read as the real thing. +# --------------------------------------------------------------------------- +wait_for_file() { + local label="$1" file="$2" timeout="$3" + local deadline=$((SECONDS + timeout)) + log_info "$label: waiting for $file, budget ${timeout}s" + while (( SECONDS < deadline )); do + if [[ -f "$file" ]]; then + log_ok "$label: $file appeared" + return 0 + fi + sleep "$POLL_INTERVAL" + done + local dir + dir="$(dirname "$file")" + log_error "$label: $file did not appear within ${timeout}s" + if [[ -f "$file.tmp" ]]; then + log_error "$label: found a leftover $file.tmp — the writer was interrupted between write and rename" + fi + if [[ -d "$dir" ]]; then + local listing + listing=$(find "$dir" -maxdepth 1 -mindepth 1 2>/dev/null | sed 's|.*/||' | sort | tr '\n' ' ') + log_info "$label: $dir currently holds: $listing" + else + log_error "$label: the directory $dir does not exist" + fi + return 1 +} + +install_flink_distribution() { + local install_dir="$1" + + # Reuse install.sh's download, archive validation and extraction logic, but + # stop before it installs a *released* Flink Agents JAR. That step derives its + # download URL from the Flink version, and no such artifact is published for + # every Flink version this test runs against — asking for one that does not + # exist would abort the install. This test must run the artifacts built from + # the current checkout anyway, which stage_dist_jars puts in place below. + FLINK_AGENTS_INSTALL_SH_NO_RUN=1 \ + FLINK_VERSION="$FLINK_VERSION" \ + INSTALL_FLINK=Yes \ + INSTALL_DIR="$install_dir" \ + NO_PROMPT=1 \ + bash -c ' + source "$1" + plan_flink + install_flink_if_needed + ' _ "$ROOT_DIR/tools/install.sh" +} + +# Replaces install.sh's setup_python_env, which would populate the venv with a +# released flink-agents. This creates a plain venv and nothing more; the packages +# arrive later, from tools/build.sh's `uv pip install dist/*.whl` and from +# install_built_python_package. +prepare_python_venv() { + if [[ -f "$VENV_DIR/pyvenv.cfg" && -x "$VENV_DIR/bin/python" ]]; then + # Reuse is unconditional on the interpreter: unlike install.sh, this accepts + # a venv built by a Python outside the range python/pyproject.toml declares. + log_info "Reusing Python venv: $VENV_DIR" + return + fi + if [[ -e "$VENV_DIR" ]] \ + && [[ ! -d "$VENV_DIR" \ + || -n "$(find "$VENV_DIR" -mindepth 1 -print -quit 2>/dev/null)" ]]; then + log_error "VENV_DIR is not an empty directory or a valid venv: $VENV_DIR" + return 1 + fi + + local python_bin="${PYTHON_BIN:-python3}" + log_info "Creating Python venv: $VENV_DIR" + "$python_bin" -m venv "$VENV_DIR" +} + +install_flink() { + log_section "Install Flink standalone (version $FLINK_VERSION)" + + # Anchor VENV_DIR to the repo so a fresh and a reused Flink installation share + # one Python environment, and so later steps can address it by path. + export VENV_DIR="${VENV_DIR:-$ROOT_DIR/.flink-agents-env}" + + if [[ -n "${FLINK_HOME:-}" && -x "$FLINK_HOME/bin/flink" ]]; then + log_info "Reusing existing FLINK_HOME: $FLINK_HOME" + export FLINK_HOME + else + local install_dir="${INSTALL_DIR:-$HOME/.local/flink}" + log_info "Installing the Flink distribution with tools/install.sh helpers" + install_flink_distribution "$install_dir" + export FLINK_HOME="${install_dir}/flink-${FLINK_VERSION}" + + if [[ ! -x "$FLINK_HOME/bin/flink" ]]; then + log_error "Flink installation not found at expected path: $FLINK_HOME" + exit 1 + fi + log_ok "Flink installed at: $FLINK_HOME" + fi + + # The distribution ships the PyFlink jar under opt/, outside the lib/ directory + # the cluster loads. install.sh copies it across as part of --enable-pyflink, + # which is no longer reached from here, so the copy happens here — on both branches + # above, since a reused FLINK_HOME need not have been prepared by install.sh. + local pyflink_jar="$FLINK_HOME/opt/flink-python-${FLINK_VERSION}.jar" + if [[ ! -f "$pyflink_jar" ]]; then + log_error "PyFlink JAR not found in Flink distribution: $pyflink_jar" + exit 1 + fi + cp "$pyflink_jar" "$FLINK_HOME/lib/" + + prepare_python_venv + + # Both branches above fall through to here, so the venv is activated whether + # FLINK_HOME was reused or freshly installed. That is load-bearing: the + # TaskManager daemon inherits this shell's PATH and resolves a bare `python` + # from it, and this script restarts the TaskManager mid-test — so an + # unactivated shell would give the replacement TaskManager a different + # interpreter from the original. Activation settles which interpreter that is + # and nothing more — what is installed in it is install_built_python_package's + # business, and that runs later. + if [[ ! -x "$VENV_DIR/bin/python" ]]; then + log_error "Expected Python venv not found at: $VENV_DIR" + exit 1 + fi + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" + log_ok "Activated Python venv: $VENV_DIR" + + FLINK_CONF="$FLINK_HOME/conf/config.yaml" + if [[ ! -f "$FLINK_CONF" ]]; then + log_error "Flink config not found: $FLINK_CONF" + exit 1 + fi +} + +# The wheel this run must exercise is the one tools/build.sh just produced, so it is +# installed from python/dist rather than from PyPI. apache-flink is pinned to the +# Flink version the cluster runs, because the submitting client and the TaskManager +# both execute this same interpreter. +install_built_python_package() { + local wheel + wheel=$(find "$ROOT_DIR/python/dist" -maxdepth 1 -name '*.whl' | head -n 1) + if [[ -z "$wheel" ]]; then + log_error "Python wheel not found after build in: $ROOT_DIR/python/dist" + return 1 + fi + + log_info "Installing $(basename "$wheel") and apache-flink==$FLINK_VERSION into $VENV_DIR" + "$VENV_DIR/bin/python" -m pip install --quiet \ + "$wheel" "apache-flink==$FLINK_VERSION" + + if ! "$VENV_DIR/bin/python" -c 'import flink_agents, pyflink' >/dev/null 2>&1; then + log_error "The built Flink Agents wheel or PyFlink is not importable from: $VENV_DIR/bin/python" + return 1 + fi + + export PYFLINK_CLIENT_EXECUTABLE="$VENV_DIR/bin/python" + log_ok "Installed the built wheel and PyFlink into: $VENV_DIR" +} + +build_project() { + log_section "Build flink-agents (Java + Python)" + ( + cd "$ROOT_DIR" + SKIP_SPOTLESS_CHECK=true bash tools/build.sh + ) + install_built_python_package + log_ok "Build completed" +} + +stage_dist_jars() { + log_section "Stage dist uber jar into \$FLINK_HOME/lib" + + local project_version + project_version=$(sed -n 's/.*<version>\(.*\)<\/version>.*/\1/p' \ + "$ROOT_DIR/pom.xml" | head -n 2 | tail -n 1) + log_info "Detected project version: $project_version" + + # The flink-version uber jar already bundles the common deps. + local flink_jar="$ROOT_DIR/dist/flink-${FLINK_MAJOR_MINOR}/target/flink-agents-dist-flink-${FLINK_MAJOR_MINOR}-${project_version}.jar" + + if [[ ! -f "$flink_jar" ]]; then + log_error "Flink dist jar not found: $flink_jar" + exit 1 + fi + + # Drop any flink-agents-dist jar already in lib/ before copying. A reused + # FLINK_HOME may still carry one — from an earlier tools/install.sh run, or from + # an earlier run of this test at a different project version — and a jar of a + # different version has a different filename, so the copy below would not + # replace it and both would sit on the classpath. + rm -f "$FLINK_HOME/lib/"/flink-agents-dist-*.jar + cp "$flink_jar" "$FLINK_HOME/lib/" + log_ok "Staged: $(basename "$flink_jar")" +} + +prepare_work_dirs() { + log_section "Create input, handshake, verdict and checkpoint directories" + WORK_DIR="$(mktemp -d "${TMPDIR:-/tmp}/flink-agents-recovery.XXXXXX")" + HANDSHAKE_DIR="$WORK_DIR/handshake" + VERDICT_DIR="$WORK_DIR/verdict" + CHECKPOINT_DIR="$WORK_DIR/checkpoints" + INPUT_DIR="$WORK_DIR/input" + # The name is load-bearing. Flink's default file filter skips any name beginning + # with '.' or '_', and applies that even to a path the source was handed directly, + # which would leave the job running with nothing to read and nothing logged. + INPUT_FILE="$INPUT_DIR/trigger.txt" + # The job must never create these itself: a tool that provisions its own + # handshake directory cannot tell "the harness set me up" from "my path is wrong". + mkdir -p "$HANDSHAKE_DIR" "$VERDICT_DIR" "$CHECKPOINT_DIR" "$INPUT_DIR" + + # One line is all the source needs: the job builds the record from a constant it + # shares with the assertion and ignores this content. Written under a temporary + # name and renamed, so the file is whole the instant it appears under the name the + # source reads. Submission is several steps away, and that distance is what + # actually orders the two — the rename is redundancy, not the ordering mechanism. + printf 'trigger\n' > "$INPUT_FILE.tmp" + mv "$INPUT_FILE.tmp" "$INPUT_FILE" + # An empty file produces no record and no error. Nothing that depends on a record + # then runs, the agent's identity report included, so the run surfaces as the + # runtime-identity wait expiring with nothing pointing at the input. A wrong path + # needs no guard here: enumeration is eager, so the job fails within seconds with + # "Could not enumerate file splits". + if [[ ! -s "$INPUT_FILE" ]]; then + log_error "Trigger file is missing or empty: $INPUT_FILE" + return 1 + fi + + log_ok "Work directory: $WORK_DIR" +} + +# --------------------------------------------------------------------------- +# Config editing. Delete-then-append, so applying the same key twice is a no-op +# rather than a duplicate YAML key. +# +# A flat dotted key at column 0 is read correctly regardless of the nested blocks +# the shipped file uses, because Flink flattens the parsed document before +# building its Configuration. Appending (never prepending) also means our value +# wins any collision with a nested block, since the flattened map is populated in +# document order. +# --------------------------------------------------------------------------- +delete_config_key() { + local config_key="$1" + local tmp + tmp="$(mktemp)" + # grep -v exits 1 when it filters everything out, which set -e would treat as + # fatal. The dots in the key are regex "any char"; harmless for these keys. + grep -v "^${config_key}: " "$FLINK_CONF" > "$tmp" || true + mv "$tmp" "$FLINK_CONF" +} + +set_config_key() { + local config_key="$1" value="$2" + delete_config_key "$config_key" + printf '%s: %s\n' "$config_key" "$value" >> "$FLINK_CONF" + CONFIG_KEYS_SET+=("$config_key") +} + +configure_flink() { + log_section "Configure checkpointing and restart strategy" + + # Without an interval, checkpointing is simply off — the option has no default. + # That is the one misconfiguration that would let this whole test pass having + # verified nothing, which is why the effective value is read back below. + set_config_key "execution.checkpointing.interval" "${CHECKPOINT_INTERVAL_MS}ms" + set_config_key "execution.checkpointing.dir" "file://$CHECKPOINT_DIR" + # Every other Python e2e test in this repo disables restarts. This one must not: + # without a restart strategy the job dies on the TaskManager loss instead of + # recovering from its checkpoint. + set_config_key "restart-strategy.type" "fixed-delay" + set_config_key "restart-strategy.fixed-delay.attempts" "$RESTART_ATTEMPTS" + set_config_key "resourcemanager.standalone.start-up-time" "$STANDALONE_STARTUP_TIME" + + log_ok "Appended ${#CONFIG_KEYS_SET[@]} config key(s) to $FLINK_CONF" +} + +# --------------------------------------------------------------------------- +# Read the keys back from the running JobManager. Writing to a file only proves +# we wrote to a file. +# +# What this does and does not establish: /jobmanager/config echoes the loaded +# configuration including keys Flink does not recognize, so a value matching here +# proves the append reached the configuration the JobManager parsed — not that the +# option has any effect. The one key where that distinction matters is the +# checkpoint interval, and it is asserted separately and authoritatively by +# assert_checkpointing_enabled once a job exists. +# --------------------------------------------------------------------------- +assert_effective_config() { + log_section "Verify the JobManager loaded the config we wrote" + + local exact_keys=( + "execution.checkpointing.dir=file://$CHECKPOINT_DIR" + "restart-strategy.type=fixed-delay" + "restart-strategy.fixed-delay.attempts=$RESTART_ATTEMPTS" + "resourcemanager.standalone.start-up-time=$STANDALONE_STARTUP_TIME" + ) + + local body + body=$(rest_get "/jobmanager/config") || body="" + if [[ -z "$body" ]]; then + log_error "Could not read GET $REST_URL/jobmanager/config" + record_result "effective-config" "FAIL" + return 1 + fi + + local ok=1 entry key want got + for entry in "${exact_keys[@]}"; do + key="${entry%%=*}"; want="${entry#*=}" + if got=$(printf '%s' "$body" | json_query conf "$key"); then + if [[ "$got" == "$want" ]]; then + log_ok "config $key = $got" + else + log_error "config $key is '$got', expected '$want'" + ok=0 + fi + else + log_error "config $key is absent from the cluster configuration — the append did not take effect" + ok=0 + fi + done + + if (( ok == 0 )); then + record_result "effective-config" "FAIL" + return 1 + fi + record_result "effective-config" "PASS" +} + +# --------------------------------------------------------------------------- +# The authoritative check that checkpointing is actually on for this job. +# +# /jobs/{id}/checkpoints/config is served from the job's own CheckpointCoordinator +# configuration rather than from the raw config file, so unlike /jobmanager/config +# it cannot echo a key back that Flink ignored. It 404s outright when the job has +# no checkpointing, and reports the interval in exact milliseconds. +# +# This is the guard on the one misconfiguration that would let the whole test pass +# having verified nothing. +# --------------------------------------------------------------------------- +assert_checkpointing_enabled() { + log_section "Verify checkpointing is enabled for this job" + + local body + body=$(rest_get "/jobs/$JOB_ID/checkpoints/config") || body="" + if [[ -z "$body" ]]; then + log_error "Could not read GET $REST_URL/jobs/$JOB_ID/checkpoints/config. This endpoint answers only when the job has checkpointing configured, so an empty or 404 response means checkpointing is off and nothing would be captured to recover." + record_result "checkpointing-enabled" "FAIL" + return 1 + fi + + local interval + if ! interval=$(printf '%s' "$body" | json_query int "interval"); then + log_error "Could not read 'interval' from GET $REST_URL/jobs/$JOB_ID/checkpoints/config; the response was: $body" + record_result "checkpointing-enabled" "FAIL" + return 1 + fi + if [[ "$interval" != "$CHECKPOINT_INTERVAL_MS" ]]; then + log_error "The job's checkpoint interval is ${interval}ms, expected ${CHECKPOINT_INTERVAL_MS}ms. The value written to $FLINK_CONF is not the value in effect." + record_result "checkpointing-enabled" "FAIL" + return 1 + fi + log_ok "checkpointing-enabled: interval ${interval}ms" + + # The two-checkpoint containment gate assumes checkpoints do not overlap, which + # is Flink's default of one concurrent checkpoint. Read from the same + # authoritative endpoint rather than inferred. + local max_concurrent + if max_concurrent=$(printf '%s' "$body" | json_query int "max_concurrent"); then + log_info "checkpointing-enabled: max_concurrent = $max_concurrent" + (( max_concurrent == 1 )) || log_warn "More than one concurrent checkpoint is allowed, so the two-checkpoint containment gate is weaker than intended: the second completion no longer implies it started after the first finished." + else + log_warn "Could not read 'max_concurrent'; the containment gate assumes it is 1" + fi + + record_result "checkpointing-enabled" "PASS" +} + +start_cluster() { + log_section "Start Flink standalone cluster" + "$FLINK_HOME/bin/start-cluster.sh" + + log_info "Waiting for JobManager REST API at $REST_URL ..." + local deadline=$((SECONDS + CLUSTER_TIMEOUT)) + while (( SECONDS < deadline )); do + if curl -fsS "$REST_URL/overview" >/dev/null 2>&1; then Review Comment: Could this poll use `rest_get` or add `curl --max-time`? The current request is unbounded, so a single stalled connection can exceed `CLUSTER_TIMEOUT` and potentially hang until the workflow-level timeout. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
