diff --git a/python/flink_agents/runtime/_python_dependency.py b/python/flink_agents/runtime/_python_dependency.py new file mode 100644 index 000000000..3d39283f7 --- /dev/null +++ b/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 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 000000000..cd722ce49 --- /dev/null +++ b/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 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 cea9cce54..c9aafb89b 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 @@ -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, 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 000000000..4479bd80e --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManager.java @@ -0,0 +1,56 @@ +/* + * 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() {} + + 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 000000000..a6bc528d9 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManagerTest.java @@ -0,0 +1,61 @@ +/* + * 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.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 = generation + "/python-files/user-code"; + + 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); + } +}