[runtime][python] Refresh imports for new dependency generations

python/flink_agents/runtime/_python_dependency.py

@@ -0,0 +1,173 @@
+################################################################################
+#  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 typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from types import ModuleType
+    from typing import Any
+
+_GENERATION_LOCK = threading.RLock()
+_JOB_GENERATIONS: dict[str, str] = {}
+
+
+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.
+
+    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)
+
+    current_paths = _paths_for_generation(python_path, current_generation)
+
+    with _GENERATION_LOCK:
+        previous_generation = _JOB_GENERATIONS.get(job_id)
+        if previous_generation == current_generation:
+            return False
+
+        if previous_generation is not None:
+            _clear_python_function_cache()
+            _evict_modules_from_generation(previous_generation)
+            _remove_paths_from_generation(sys.path, previous_generation)
+
+        _prepend_current_paths(current_paths)
+        _clear_importer_cache(previous_generation, current_generation)
+        importlib.invalidate_caches()
+
+        _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 _paths_for_generation(python_path: str, generation: str) -> list[str]:
+    paths = (_normalize_path(path) for path in python_path.split(os.pathsep) if path)
+    return list(
+        dict.fromkeys(
+            path for path in paths if _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:
+    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 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 _prepend_current_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(
+    previous_generation: str | None, current_generation: str
+) -> None:
+    for path in list(sys.path_importer_cache):
+        if _path_belongs_to_generation(path, current_generation) or (
+            previous_generation is not None
+            and _path_belongs_to_generation(path, previous_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

python/flink_agents/runtime/tests/test_python_dependency.py

@@ -0,0 +1,196 @@
+################################################################################
+#  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 shutil
+import sys
+import uuid
+from importlib.resources import files
+from pathlib import Path
+
+import pytest
+
+from flink_agents.plan import function as plan_function
+from flink_agents.runtime import _python_dependency
+
+
+@pytest.fixture(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)
+    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)
+    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"
+    )
+
+    assert _python_dependency.ensure_python_dependency_generation(
+        "job-1", str(old_generation), str(old_python_path)
+    )
+    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"
+    )
+
+    assert _python_dependency.ensure_python_dependency_generation(
+        "job-1", str(new_generation), str(new_python_path)
+    )
+    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"
+
+    assert not _python_dependency.ensure_python_dependency_generation(
+        "job-1", str(new_generation), str(new_python_path)
+    )
+    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"
+    )
+
+    _python_dependency.ensure_python_dependency_generation(
+        "job-1", str(first_old_generation), str(first_old_python_path)
+    )
+    first_old_module = importlib.import_module(f"{first_package}.action")
+
+    _python_dependency.ensure_python_dependency_generation(
+        "job-2", str(second_generation), str(second_python_path)
+    )
+    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"
+    )
+    _python_dependency.ensure_python_dependency_generation(
+        "job-1", str(first_new_generation), str(first_new_python_path)
+    )
+
+    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"
+    )
+    _python_dependency.ensure_python_dependency_generation(
+        "job-retry", str(old_generation), str(old_python_path)
+    )
+    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"
+    )
+    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), str(new_python_path)
+        )
+
+    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), str(new_python_path)
+    )
+    assert importlib.import_module(f"{package_name}.action").VALUE == "new"
+
+
+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

runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java

@@ -153,6 +153,19 @@ void open(
             pythonEnvironmentManager.open();
             EmbeddedPythonEnvironment env = pythonEnvironmentManager.createEnvironment();
             pythonInterpreter = env.getInterpreter();
+            String dependencyGeneration = pythonEnvironmentManager.getBaseDirectory();
+            boolean dependencyGenerationChanged =
+                    PythonDependencyGenerationManager.ensurePythonDependencyGeneration(
+                            pythonInterpreter,
+                            jobId,
+                            dependencyGeneration,
+                            env.getEnv().getOrDefault("PYTHONPATH", ""));
+            if (dependencyGenerationChanged) {
+                LOG.info(
+                        "Activated Python dependency generation {} for job {}.",
+                        dependencyGeneration,
+                        jobId);
+            }
             pythonRunnerContext =
                     new PythonRunnerContextImpl(
                             metricGroup,