weiqingy commented on code in PR #943: URL: https://github.com/apache/flink-agents/pull/943#discussion_r3726296144
########## python/flink_agents/runtime/_python_dependency.py: ########## @@ -0,0 +1,180 @@ +################################################################################ +# 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) -> 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) + + 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( + _paths_for_generation(sys.path, current_generation) + ) + return False + + if previous_generation is not None: + _deactivate_generation(previous_generation) + + _activate_generation(current_generation) + + _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() Review Comment: `_clear_python_function_cache()` reaches `flink_agents.plan.function.clear_python_function_cache`, which clears `_PYTHON_FUNCTION_CACHE` wholesale. The key is `(module, qualname)` (`plan/function.py:30,364`), with no job or generation component, so a generation change for job A also wipes every other job's entries in the same TaskManager JVM. The rebuild is harmless, but the clear seems to open a `KeyError` window at `plan/function.py:368-373`: the `not in` test and the `[cache_key]` read are separate operations, and a `clear()` landing between them would raise on the other job's action thread. `_GENERATION_LOCK` doesn't cover that path, since `call_python_function` never takes the lock. Before this PR `clear_python_function_cache()` had no production caller outside `plan/tests/test_function.py`, so the window looks new here. I'm confident about the mechanism. Hitting it needs two agent jobs on one TaskManager with one of them restarting, so rare rather than impossible. Is a process-global wipe the right lever for a job-scoped problem? Selective clearing isn't possible as written since the key carries no generation, but could evicting just the modules that were dropped from `sys.modules` keep this inside the generation being deactivated? ########## python/flink_agents/runtime/_python_dependency.py: ########## @@ -0,0 +1,180 @@ +################################################################################ +# 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) -> 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) + + 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( + _paths_for_generation(sys.path, current_generation) + ) + return False + + if previous_generation is not None: Review Comment: This branch is what decides whether anything gets evicted, and it only fires when this job's recorded generation changed. Nothing deactivates a generation when a job *ends*: `PythonBridgeManager.close()` (line 303) closes the interpreter and the environment manager, with no counterpart to `ensurePythonDependencyGeneration`. On a session-cluster TaskManager that looks like it leaves the #941 shape reachable across jobs: 1. Job A imports top-level package `my_agent` from `gen-A` on TM X. 2. Job A finishes, `PythonEnvResources.release` deletes `gen-A` at refcount 0, but `sys.modules['my_agent']` still points into it. 3. Job B, a different `JobID` shipping its own `my_agent`, opens on the same TM. `_JOB_GENERATIONS.get('<B>')` is `None`, so this branch is skipped and only `_activate_generation(gen-B)` runs. 4. `import my_agent` hands back job A's module from the deleted directory, and `files('my_agent').joinpath(...)` fails the way #941 describes. `test_python_dependency.py:145` asserts exactly this preservation, which is right for a job that is still live. The piece I can't see is what would tell a live job apart from a finished one. The PR doesn't break this, it just doesn't close it, so it may well be deliberate scope. Is the cross-job case out of scope here? If it's in scope, could keying on "the recorded generation directory no longer exists" rather than "this job's previous generation" reach both, zipimport and namespace packages aside? ########## python/flink_agents/runtime/_python_dependency.py: ########## @@ -0,0 +1,180 @@ +################################################################################ +# 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) -> 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) + + 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( + _paths_for_generation(sys.path, current_generation) + ) + return False + + if previous_generation is not None: + _deactivate_generation(previous_generation) + + _activate_generation(current_generation) + + _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) -> None: Review Comment: `_paths_for_generation` (line 93) filters entries that are already in `sys.path`; it never derives paths from the generation directory. So if the new generation's entries aren't on `sys.path` yet, `_activate_generation` becomes a no-op while `ensure_python_dependency_generation` still records the generation, returns `True`, and Java logs "Activated Python dependency generation" (`PythonBridgeManager.java:161-165`). I checked by activating a generation whose path was absent from `sys.path`: it returned `True` with the path still absent. It works today because `new PythonInterpreter(config)` runs `configSearchPaths`, which emits `sys.path.insert(0, r'%s')` for each configured path (pemja 0.5.7), and that happens at `env.getInterpreter()` on `PythonBridgeManager.java:155`, just above the guard. If the interpreter construction ever moved below the guard, everything would still report success, and the failure would surface later as a `ModuleNotFoundError` in user code. A job with no `add_python_file` / requirements / archives legitimately has zero generation paths, so failing hard probably isn't right. But the ordering is invisible to anyone editing `open()`, and `ensurePythonDependencyGeneration` has no Javadoc today. Does the constraint belong there, something like "must be called after the interpreter is constructed and before any user module import"? ########## python/flink_agents/runtime/_python_dependency.py: ########## @@ -0,0 +1,180 @@ +################################################################################ +# 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) -> 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) + + 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( + _paths_for_generation(sys.path, current_generation) + ) + return False + + if previous_generation is not None: + _deactivate_generation(previous_generation) + + _activate_generation(current_generation) + + _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) -> None: + _deduplicate_and_prepend_paths(_paths_for_generation(sys.path, generation)) + _clear_importer_cache(generation) + importlib.invalidate_caches() + + +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: + 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: Review Comment: `_evict_modules_from_generation` drops any module whose `__file__` / `__path__` / `__spec__` resolves under the old generation directory, with no exemption for `flink_agents` itself. The generation directory holds more than user files: `AbstractPythonEnvironmentManager` puts `python-files` on `PYTHONPATH`, pip-installs requirements under `python-requirements` and puts the resulting site-packages on `PYTHONPATH` too, and `PythonEnvironmentManager.createEnvironment()` supports running the interpreter out of a venv extracted under `python-archives`. So if `apache-flink-agents` is in the job's `requirements.txt` or ships inside that venv, `flink_agents.runtime._python_dependency` is itself under the generation and gets evicted along with the user code. The next `interpreter.exec("from flink_agents.runtime import _python_dependency")` then imports a fresh module with an empty `_JOB_GENERATIONS`, every later call sees `previous_generation is None`, and #941 is back while the guard still returns `True` and logs "Activated". I verified the eviction itself: a package living under the generation directory does get dropped. What I haven't verified is whether a real deployment ever resolves `flink_agents` from under the base directory rather than from the TaskManager image, where it would sit outside the generation, so this may well be unreachable in practice. If it isn't, would exempting `flink_agents.*` from eviction be the cheaper protection, or holding the generation record somewhere eviction can't reach? ########## runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonDependencyGenerationManagerTest.java: ########## @@ -0,0 +1,58 @@ +/* + * 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"; + + when(interpreter.invoke( Review Comment: Nothing exercises the `checkState` at `PythonDependencyGenerationManager.java:46-49`. The stub here already mocks the interpreter, so a second case returning a non-`Boolean` would cover it in a few lines. Worth adding? -- 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]
