joeyutong commented on code in PR #943: URL: https://github.com/apache/flink-agents/pull/943#discussion_r3796287385
########## 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: Agreed. The guard now takes the environment `PYTHONPATH` and prepends those generation entries even if they are not already on `sys.path`. Also documented that it must run after interpreter construction and before any user import. ########## 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: Added a test that a non-Boolean invoke result fails the `checkState`. ########## 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: For same-job failover this runs during operator open after the previous attempt has closed, so no concurrent `call_python_function` is expected on this job. Leaving the process-wide cache clear as-is for this PR. ########## 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: Agreed this PR is scoped to same-job failover. Cross-job reuse of a TaskManager is a separate lifecycle; we'll track that as a follow-up. The guard now documents that it only refreshes the previous generation for the same job id. ########## 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: You're right that the helper can live under the generation. Job generation records now live on an interpreter-scoped in-memory module so they survive helper eviction and reload. Added a regression for A→B→C with the helper loaded from `python-dist`. -- 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]
