This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 12262f63 [runtime][python] Refresh imports for new dependency
generations (#943)
12262f63 is described below
commit 12262f63bd52a872dfca26ead964a98ce64274e8
Author: Yu Tang <[email protected]>
AuthorDate: Mon Aug 24 17:31:21 2026 +0800
[runtime][python] Refresh imports for new dependency generations (#943)
Co-authored-by: Codex <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Claude Code <[email protected]>
---
python/flink_agents/runtime/_python_dependency.py | 229 ++++++++++++++++
.../runtime/tests/test_python_dependency.py | 295 +++++++++++++++++++++
.../runtime/operator/PythonBridgeManager.java | 22 +-
.../PythonDependencyGenerationManager.java | 73 +++++
.../PythonDependencyGenerationManagerTest.java | 84 ++++++
5 files changed, 700 insertions(+), 3 deletions(-)
diff --git a/python/flink_agents/runtime/_python_dependency.py
b/python/flink_agents/runtime/_python_dependency.py
new file mode 100644
index 00000000..5fe65e00
--- /dev/null
+++ b/python/flink_agents/runtime/_python_dependency.py
@@ -0,0 +1,229 @@
+################################################################################
+# 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.
+################################################################################
+
+from __future__ import annotations
+
+import importlib
+import os
+import sys
+import threading
+from pathlib import Path
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+ from typing import Any
+
+# Interpreter-scoped generation records. This helper may be imported from a
+# job's python-dist directory and evicted when that directory is replaced; the
+# records must outlive any one generation-scoped module.
+_STATE_MODULE_NAME = "_flink_agents_python_dependency_state"
+
+
+def _get_state_module() -> ModuleType:
+ state = sys.modules.get(_STATE_MODULE_NAME)
+ if state is not None:
+ return state
+
+ candidate = ModuleType(_STATE_MODULE_NAME)
+ candidate.generation_lock = threading.RLock()
+ candidate.job_generations = {}
+ # setdefault prevents concurrent Pemja threads from installing different
+ # state modules.
+ return sys.modules.setdefault(_STATE_MODULE_NAME, candidate)
+
+
+_STATE = _get_state_module()
+_GENERATION_LOCK = _STATE.generation_lock
+_JOB_GENERATIONS = _STATE.job_generations
+
+
+def ensure_python_dependency_generation(
+ job_id: str, generation: str, python_path: str = ""
+) -> bool:
+ """Activate a Flink-managed dependency generation in the Pemja interpreter.
+
+ When Flink replaces a job's temporary dependency directory, remove imports
+ owned by the previous directory before user actions or resources are
loaded.
+
+ This tracks one generation per job id and only refreshes that job's
previous
+ directory. It does not deactivate generations when a job ends, and it does
+ not isolate import caches across jobs that share a TaskManager.
+
+ ``python_path`` is the generation's configured ``PYTHONPATH``. Activation
+ prepends those entries even if they are not already on ``sys.path``.
Callers
+ must invoke this after the interpreter is constructed and before any user
+ module is imported.
+
+ Returns:
+ ``True`` when a different generation was activated, otherwise
``False``.
+ """
+ if not job_id:
+ msg = "job_id must not be empty"
+ raise ValueError(msg)
+
+ current_generation = _normalize_path(generation)
+ if not Path(current_generation).is_dir():
+ msg = f"Python dependency generation does not exist:
{current_generation}"
+ raise RuntimeError(msg)
+
+ with _GENERATION_LOCK:
+ previous_generation = _JOB_GENERATIONS.get(job_id)
+ if previous_generation == current_generation:
+ # Pemja inserts configured paths for every interpreter sharing this
+ # generation.
+ _deduplicate_and_prepend_paths(
+ _configured_paths_for_generation(current_generation,
python_path)
+ )
+ return False
+
+ if previous_generation is not None:
+ _deactivate_generation(previous_generation)
+
+ _activate_generation(current_generation, python_path)
+
+ _JOB_GENERATIONS[job_id] = current_generation
+ return True
+
+
+def _normalize_path(path: str | os.PathLike[str]) -> str:
+ # Keep Flink's symlink path so imported modules remain attributable to
+ # their owning python-dist generation.
+ return os.path.normcase(str(Path(path).absolute()))
+
+
+def _deactivate_generation(generation: str) -> None:
+ _clear_python_function_cache()
+ _evict_modules_from_generation(generation)
+ _remove_paths_from_generation(sys.path, generation)
+ _clear_importer_cache(generation)
+
+
+def _activate_generation(generation: str, python_path: str = "") -> None:
+ _deduplicate_and_prepend_paths(
+ _configured_paths_for_generation(generation, python_path)
+ )
+ _clear_importer_cache(generation)
+ importlib.invalidate_caches()
+
+
+def _python_path_entries(python_path: str) -> list[str]:
+ if not python_path:
+ return []
+ return [entry for entry in python_path.split(os.pathsep) if entry]
+
+
+def _configured_paths_for_generation(generation: str, python_path: str) ->
list[str]:
+ return _paths_for_generation(
+ [*_python_path_entries(python_path), *sys.path], generation
+ )
+
+
+def _paths_for_generation(paths: list[str], generation: str) -> list[str]:
+ paths = (_try_normalize_path(path) for path in paths)
+ return list(
+ dict.fromkeys(
+ path
+ for path in paths
+ if path is not None and _path_belongs_to_generation(path,
generation)
+ )
+ )
+
+
+def _module_paths(module: ModuleType) -> Iterator[Any]:
+ spec = getattr(module, "__spec__", None)
+ path_values = (
+ getattr(module, "__file__", None),
+ getattr(module, "__path__", None),
+ getattr(spec, "origin", None),
+ getattr(spec, "submodule_search_locations", None),
+ )
+ for value in path_values:
+ if isinstance(value, str | bytes | os.PathLike):
+ yield value
+ elif value is not None:
+ try:
+ yield from value
+ except (TypeError, ValueError):
+ continue
+
+
+def _try_normalize_path(path: Any) -> str | None:
+ if not isinstance(path, str | bytes | os.PathLike):
+ return None
+ try:
+ return _normalize_path(os.fsdecode(path))
+ except (OSError, TypeError, ValueError):
+ return None
+
+
+def _clear_python_function_cache() -> None:
+ # Same-job failover runs this during operator open after the previous
+ # attempt has closed, so no concurrent call_python_function is expected.
+ function_module = sys.modules.get("flink_agents.plan.function")
+ if function_module is not None:
+ function_module.clear_python_function_cache()
+
+
+def _evict_modules_from_generation(generation: str) -> None:
+ modules_to_remove = []
+ for module_name, module in list(sys.modules.items()):
+ if module_name == _STATE_MODULE_NAME:
+ continue
+ if module is not None and any(
+ _path_belongs_to_generation(path, generation)
+ for path in _module_paths(module)
+ ):
+ modules_to_remove.append(module_name)
+
+ for module_name in sorted(
+ modules_to_remove, key=lambda name: name.count("."), reverse=True
+ ):
+ sys.modules.pop(module_name, None)
+
+
+def _remove_paths_from_generation(paths: list[str], generation: str) -> None:
+ paths[:] = [
+ path for path in paths if not _path_belongs_to_generation(path,
generation)
+ ]
+
+
+def _deduplicate_and_prepend_paths(current_paths: list[str]) -> None:
+ normalized_current_paths = set(current_paths)
+ sys.path[:] = [
+ path
+ for path in sys.path
+ if _try_normalize_path(path) not in normalized_current_paths
+ ]
+ sys.path[0:0] = current_paths
+
+
+def _clear_importer_cache(generation: str) -> None:
+ for path in list(sys.path_importer_cache):
+ if _path_belongs_to_generation(path, generation):
+ sys.path_importer_cache.pop(path, None)
+
+
+def _path_belongs_to_generation(path: Any, generation: str) -> bool:
+ normalized = _try_normalize_path(path)
+ if normalized is None:
+ return False
+ candidate = Path(normalized)
+ generation_path = Path(generation)
+ return candidate == generation_path or generation_path in candidate.parents
diff --git a/python/flink_agents/runtime/tests/test_python_dependency.py
b/python/flink_agents/runtime/tests/test_python_dependency.py
new file mode 100644
index 00000000..2b96d559
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_python_dependency.py
@@ -0,0 +1,295 @@
+################################################################################
+# 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.
+################################################################################
+
+import importlib
+import importlib.util
+import shutil
+import sys
+import uuid
+from importlib.resources import files
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+from flink_agents.plan import function as plan_function
+from flink_agents.runtime import _python_dependency
+
+_HELPER_MODULE = "flink_agents.runtime._python_dependency"
+
+
[email protected](autouse=True)
+def restore_import_state():
+ original_sys_path = list(sys.path)
+ original_importer_cache = dict(sys.path_importer_cache)
+ original_generations = dict(_python_dependency._JOB_GENERATIONS)
+ original_function_cache = dict(plan_function._PYTHON_FUNCTION_CACHE)
+ original_helper = sys.modules.get(_HELPER_MODULE)
+ imported_packages: set[str] = set()
+
+ yield imported_packages
+
+ sys.path[:] = original_sys_path
+ sys.path_importer_cache.clear()
+ sys.path_importer_cache.update(original_importer_cache)
+ _python_dependency._JOB_GENERATIONS.clear()
+ _python_dependency._JOB_GENERATIONS.update(original_generations)
+ plan_function._PYTHON_FUNCTION_CACHE.clear()
+ plan_function._PYTHON_FUNCTION_CACHE.update(original_function_cache)
+ if original_helper is not None:
+ sys.modules[_HELPER_MODULE] = original_helper
+ for package_name in imported_packages:
+ for module_name in list(sys.modules):
+ if module_name == package_name or module_name.startswith(
+ f"{package_name}."
+ ):
+ sys.modules.pop(module_name, None)
+ importlib.invalidate_caches()
+
+
+def test_generation_change_reloads_user_package_and_resources(
+ tmp_path: Path, restore_import_state
+):
+ package_name = f"generation_package_{uuid.uuid4().hex}"
+ restore_import_state.add(package_name)
+ old_generation, old_python_path = _create_generation(
+ tmp_path, "old", package_name, "old"
+ )
+ sys.path.insert(0, str(old_python_path))
+
+ assert _python_dependency.ensure_python_dependency_generation(
+ "job-1", str(old_generation)
+ )
+ old_module = importlib.import_module(f"{package_name}.action")
+ assert old_module.VALUE == "old"
+ plan_function._PYTHON_FUNCTION_CACHE[(f"{package_name}.action",
"handler")] = (
+ object()
+ )
+ assert files(package_name).joinpath("skills", "SKILL.md").read_text() ==
"old"
+
+ shutil.rmtree(old_generation)
+ new_generation, new_python_path = _create_generation(
+ tmp_path, "new", package_name, "new"
+ )
+ sys.path.insert(0, str(new_python_path))
+
+ assert _python_dependency.ensure_python_dependency_generation(
+ "job-1", str(new_generation)
+ )
+ assert package_name not in sys.modules
+ assert f"{package_name}.action" not in sys.modules
+ assert str(old_python_path) not in sys.path
+ assert str(old_python_path) not in sys.path_importer_cache
+ assert sys.path[0] == str(new_python_path)
+ assert plan_function.get_python_function_cache_size() == 0
+
+ new_module = importlib.import_module(f"{package_name}.action")
+ assert new_module.VALUE == "new"
+ assert files(package_name).joinpath("skills", "SKILL.md").read_text() ==
"new"
+
+ # Each Pemja interpreter inserts its configured paths before this guard
runs.
+ sys.path.insert(0, str(new_python_path))
+ assert not _python_dependency.ensure_python_dependency_generation(
+ "job-1", str(new_generation)
+ )
+ assert sys.path.count(str(new_python_path)) == 1
+ assert importlib.import_module(f"{package_name}.action") is new_module
+
+
+def test_generation_change_preserves_other_active_job(
+ tmp_path: Path, restore_import_state
+):
+ first_package = f"first_package_{uuid.uuid4().hex}"
+ second_package = f"second_package_{uuid.uuid4().hex}"
+ restore_import_state.update({first_package, second_package})
+
+ first_old_generation, first_old_python_path = _create_generation(
+ tmp_path, "first-old", first_package, "first-old"
+ )
+ second_generation, second_python_path = _create_generation(
+ tmp_path, "second", second_package, "second"
+ )
+ sys.path.insert(0, str(first_old_python_path))
+
+ _python_dependency.ensure_python_dependency_generation(
+ "job-1", str(first_old_generation)
+ )
+ first_old_module = importlib.import_module(f"{first_package}.action")
+
+ sys.path.insert(0, str(second_python_path))
+ _python_dependency.ensure_python_dependency_generation(
+ "job-2", str(second_generation)
+ )
+ second_module = importlib.import_module(f"{second_package}.action")
+
+ shutil.rmtree(first_old_generation)
+ first_new_generation, first_new_python_path = _create_generation(
+ tmp_path, "first-new", first_package, "first-new"
+ )
+ sys.path.insert(0, str(first_new_python_path))
+ _python_dependency.ensure_python_dependency_generation(
+ "job-1", str(first_new_generation)
+ )
+
+ assert first_old_module.VALUE == "first-old"
+ assert first_package not in sys.modules
+ assert importlib.import_module(f"{first_package}.action").VALUE ==
"first-new"
+ assert importlib.import_module(f"{second_package}.action") is second_module
+ assert str(second_python_path) in sys.path
+
+
+def test_failed_refresh_is_retried(tmp_path: Path, restore_import_state,
monkeypatch):
+ package_name = f"retry_package_{uuid.uuid4().hex}"
+ restore_import_state.add(package_name)
+ old_generation, old_python_path = _create_generation(
+ tmp_path, "retry-old", package_name, "old"
+ )
+ sys.path.insert(0, str(old_python_path))
+ _python_dependency.ensure_python_dependency_generation(
+ "job-retry", str(old_generation)
+ )
+ importlib.import_module(f"{package_name}.action")
+
+ shutil.rmtree(old_generation)
+ new_generation, new_python_path = _create_generation(
+ tmp_path, "retry-new", package_name, "new"
+ )
+ sys.path.insert(0, str(new_python_path))
+ original_invalidate_caches = importlib.invalidate_caches
+
+ def fail_invalidate_caches() -> None:
+ msg = "injected cache invalidation failure"
+ raise RuntimeError(msg)
+
+ monkeypatch.setattr(
+ _python_dependency.importlib,
+ "invalidate_caches",
+ fail_invalidate_caches,
+ )
+ with pytest.raises(RuntimeError, match="injected cache invalidation
failure"):
+ _python_dependency.ensure_python_dependency_generation(
+ "job-retry", str(new_generation)
+ )
+
+ assert _python_dependency._JOB_GENERATIONS["job-retry"] ==
str(old_generation)
+
+ monkeypatch.setattr(
+ _python_dependency.importlib,
+ "invalidate_caches",
+ original_invalidate_caches,
+ )
+ assert _python_dependency.ensure_python_dependency_generation(
+ "job-retry", str(new_generation)
+ )
+ assert importlib.import_module(f"{package_name}.action").VALUE == "new"
+
+
+def test_activate_uses_configured_python_path(tmp_path: Path,
restore_import_state):
+ package_name = f"configured_path_{uuid.uuid4().hex}"
+ restore_import_state.add(package_name)
+ generation, python_path = _create_generation(
+ tmp_path, "configured", package_name, "configured"
+ )
+ normalized = _python_dependency._normalize_path(python_path)
+ assert normalized not in sys.path
+
+ assert _python_dependency.ensure_python_dependency_generation(
+ "job-configured", str(generation), str(python_path)
+ )
+ assert sys.path[0] == normalized
+ assert importlib.import_module(f"{package_name}.action").VALUE ==
"configured"
+
+ sys.path.remove(normalized)
+ assert not _python_dependency.ensure_python_dependency_generation(
+ "job-configured", str(generation), str(python_path)
+ )
+ assert sys.path[0] == normalized
+
+
+def test_generation_state_survives_helper_reload_from_job_requirements(
+ tmp_path: Path, restore_import_state
+):
+ package_name = f"helper_reload_{uuid.uuid4().hex}"
+ restore_import_state.add(package_name)
+
+ gen_a, path_a = _create_generation(tmp_path, "helper-a", package_name, "a")
+ gen_b, path_b = _create_generation(tmp_path, "helper-b", package_name, "b")
+ gen_c, path_c = _create_generation(tmp_path, "helper-c", package_name, "c")
+ helper_a = _copy_helper_into_generation(path_a)
+ helper_b = _copy_helper_into_generation(path_b)
+ helper_c = _copy_helper_into_generation(path_c)
+
+ loaded_a = _load_helper_from(helper_a)
+ assert loaded_a.ensure_python_dependency_generation(
+ "job-fa", str(gen_a), str(path_a)
+ )
+ assert importlib.import_module(f"{package_name}.action").VALUE == "a"
+
+ shutil.rmtree(gen_a)
+ assert loaded_a.ensure_python_dependency_generation(
+ "job-fa", str(gen_b), str(path_b)
+ )
+ assert sys.modules.get(_HELPER_MODULE) is not loaded_a
+
+ loaded_b = _load_helper_from(helper_b)
+ state = sys.modules[_python_dependency._STATE_MODULE_NAME]
+ assert loaded_b._JOB_GENERATIONS is state.job_generations
+ assert loaded_a._JOB_GENERATIONS is state.job_generations
+ assert loaded_b._JOB_GENERATIONS["job-fa"] ==
loaded_b._normalize_path(gen_b)
+ assert importlib.import_module(f"{package_name}.action").VALUE == "b"
+
+ shutil.rmtree(gen_b)
+ assert loaded_b.ensure_python_dependency_generation(
+ "job-fa", str(gen_c), str(path_c)
+ )
+ loaded_c = _load_helper_from(helper_c)
+ assert loaded_c._JOB_GENERATIONS["job-fa"] ==
loaded_c._normalize_path(gen_c)
+ assert importlib.import_module(f"{package_name}.action").VALUE == "c"
+
+
+def _create_generation(
+ tmp_path: Path, generation_name: str, package_name: str, value: str
+) -> tuple[Path, Path]:
+ generation = tmp_path / f"python-dist-{generation_name}"
+ python_path = generation / "python-files" / "user-code"
+ package_path = python_path / package_name
+ skills_path = package_path / "skills"
+ skills_path.mkdir(parents=True)
+ (package_path / "__init__.py").write_text("")
+ (package_path / "action.py").write_text(f"VALUE = {value!r}\n")
+ (skills_path / "SKILL.md").write_text(value)
+ return generation, python_path
+
+
+def _copy_helper_into_generation(python_path: Path) -> Path:
+ dest = python_path / "flink_agents" / "runtime" / "_python_dependency.py"
+ dest.parent.mkdir(parents=True, exist_ok=True)
+ (python_path / "flink_agents" / "__init__.py").write_text("")
+ (dest.parent / "__init__.py").write_text("")
+ shutil.copy(Path(_python_dependency.__file__), dest)
+ return dest
+
+
+def _load_helper_from(helper_file: Path) -> ModuleType:
+ spec = importlib.util.spec_from_file_location(_HELPER_MODULE, helper_file)
+ assert spec is not None
+ assert spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[_HELPER_MODULE] = module
+ spec.loader.exec_module(module)
+ return module
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index 95e215d7..ae3bedb3 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -96,10 +96,12 @@ class PythonBridgeManager implements AutoCloseable {
* <p>Scans the agent plan for any {@link PythonFunction} action or {@link
* PythonResourceProvider}. If neither is present, this method is a no-op
and {@link
* #isInitialized()} stays {@code false}. Otherwise it builds the {@link
- * PythonEnvironmentManager}, opens an embedded {@link PythonInterpreter},
constructs the shared
- * {@link PythonRunnerContextImpl}, wires the Java/Python resource
adapters, and conditionally
+ * PythonEnvironmentManager}, opens an embedded {@link PythonInterpreter},
refreshes the shared
+ * import state for the current dependency generation, constructs the
shared {@link
+ * PythonRunnerContextImpl}, wires the Java/Python resource adapters, and
conditionally
* initializes the Python action executor and the Python resource adapter
(each only when the
- * corresponding component is present in the plan).
+ * corresponding component is present in the plan). The generation guard
runs immediately after
+ * interpreter construction and before any user module import.
*
* @param agentPlan the agent plan describing actions and resources.
* @param resourceCache the resource cache visible to both languages.
@@ -154,6 +156,20 @@ class PythonBridgeManager implements AutoCloseable {
pythonEnvironmentManager.open();
EmbeddedPythonEnvironment env =
pythonEnvironmentManager.createEnvironment();
pythonInterpreter = env.getInterpreter();
+ String dependencyGeneration =
pythonEnvironmentManager.getBaseDirectory();
+ String pythonPath = env.getEnv().get("PYTHONPATH");
+ boolean dependencyGenerationChanged =
+
PythonDependencyGenerationManager.ensurePythonDependencyGeneration(
+ pythonInterpreter,
+ jobId,
+ dependencyGeneration,
+ pythonPath == null ? "" : pythonPath);
+ if (dependencyGenerationChanged) {
+ LOG.info(
+ "Activated Python dependency generation {} for job
{}.",
+ dependencyGeneration,
+ jobId);
+ }
pythonRunnerContext =
new PythonRunnerContextImpl(
metricGroup,
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManager.java
new file mode 100644
index 00000000..9c5ad934
--- /dev/null
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManager.java
@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.runtime.operator;
+
+import org.apache.flink.api.common.JobID;
+import pemja.core.PythonInterpreter;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/** Coordinates Flink Python dependency generations with Pemja's shared import
state. */
+final class PythonDependencyGenerationManager {
+
+ private static final String PYTHON_IMPORT =
+ "from flink_agents.runtime import _python_dependency";
+ private static final String ENSURE_PYTHON_DEPENDENCY_GENERATION =
+ "_python_dependency.ensure_python_dependency_generation";
+
+ private PythonDependencyGenerationManager() {}
+
+ /**
+ * Refreshes Pemja's process-level import state for this job's current
dependency directory.
+ *
+ * <p>Must be called after {@link
+ *
org.apache.flink.agents.runtime.env.EmbeddedPythonEnvironment#getInterpreter()}
has
+ * constructed the interpreter and before any user action or resource is
imported.
+ *
+ * <p>This guard only tracks same-job failover, where Flink rematerializes
the same artifacts
+ * under a new {@code python-dist-*} directory. Cross-job reuse of a
TaskManager is out of
+ * scope.
+ *
+ * @param interpreter the Pemja interpreter that shares process-level
import state
+ * @param jobId the Flink job whose dependency generation should be
activated
+ * @param generation the current Flink-managed dependency directory
+ * @param pythonPath the environment {@code PYTHONPATH} for this
generation; used to activate
+ * generation entries even if they are not yet present on {@code
sys.path}
+ */
+ static boolean ensurePythonDependencyGeneration(
+ PythonInterpreter interpreter, JobID jobId, String generation,
String pythonPath) {
+ checkNotNull(interpreter);
+ checkNotNull(jobId);
+ checkNotNull(generation);
+ checkNotNull(pythonPath);
+
+ interpreter.exec(PYTHON_IMPORT);
+ Object result =
+ interpreter.invoke(
+ ENSURE_PYTHON_DEPENDENCY_GENERATION,
+ jobId.toHexString(),
+ generation,
+ pythonPath);
+ checkState(
+ result instanceof Boolean,
+ "Python dependency generation guard returned an invalid
result: %s",
+ result);
+ return (Boolean) result;
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManagerTest.java
new file mode 100644
index 00000000..7a954ee3
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManagerTest.java
@@ -0,0 +1,84 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.runtime.operator;
+
+import org.apache.flink.api.common.JobID;
+import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
+import pemja.core.PythonInterpreter;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Tests for {@link PythonDependencyGenerationManager}. */
+class PythonDependencyGenerationManagerTest {
+
+ @Test
+ void importsGuardModuleBeforeInvokingGenerationCheck() {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ JobID jobId = new JobID();
+ String generation = "/tmp/python-dist-current";
+ String pythonPath = "/tmp/python-dist-current/python-files";
+
+ when(interpreter.invoke(
+
"_python_dependency.ensure_python_dependency_generation",
+ jobId.toHexString(),
+ generation,
+ pythonPath))
+ .thenReturn(true);
+
+ assertThat(
+
PythonDependencyGenerationManager.ensurePythonDependencyGeneration(
+ interpreter, jobId, generation, pythonPath))
+ .isTrue();
+
+ InOrder calls = inOrder(interpreter);
+ calls.verify(interpreter).exec("from flink_agents.runtime import
_python_dependency");
+ calls.verify(interpreter)
+ .invoke(
+
"_python_dependency.ensure_python_dependency_generation",
+ jobId.toHexString(),
+ generation,
+ pythonPath);
+ }
+
+ @Test
+ void rejectsNonBooleanGenerationResult() {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ JobID jobId = new JobID();
+ String generation = "/tmp/python-dist-current";
+ String pythonPath = "/tmp/python-dist-current/python-files";
+
+ when(interpreter.invoke(
+
"_python_dependency.ensure_python_dependency_generation",
+ jobId.toHexString(),
+ generation,
+ pythonPath))
+ .thenReturn("yes");
+
+ assertThatThrownBy(
+ () ->
+
PythonDependencyGenerationManager.ensurePythonDependencyGeneration(
+ interpreter, jobId, generation,
pythonPath))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("invalid result");
+ }
+}
