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

[joeyutong]

Linked issue: #941

Purpose of change

Flink may rematerialize a job's Python dependencies under a new python-dist-* directory after task failover, while Pemja's process-level main interpreter and its import state remain alive in the TaskManager JVM. Existing user modules can therefore keep __file__, __path__, or __spec__ entries that point to the removed dependency directory.

This change activates each job's current dependency generation before loading Python actions or resources. When the generation changes, it:

• clears the Flink Agents Python function cache;
• evicts modules attributed to the previous generation from sys.modules;
• removes stale sys.path and importer-cache entries;
• prepends the current generation's PYTHONPATH entries; and
• preserves modules and paths owned by other active jobs.

The generation is recorded only after the refresh succeeds, so a failed refresh is retried on the next initialization.

Tests

• uv run --project python pytest -q python/flink_agents/runtime/tests/test_python_dependency.py (3 passed)
• mvn -B --no-transfer-progress -pl runtime -am -Dtest=PythonDependencyGenerationManagerTest -Dsurefire.failIfNoSpecifiedTests=false test (1 passed)
• Maven Spotless formatting/checks passed as part of the Java verification

API

No public API changes.

Documentation

• doc-needed
• doc-not-needed
• doc-included

[weiqingy]

Thanks for taking this on. A few questions inline.

[weiqingy]

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 [Bug] Refresh Python import state when Flink replaces dependency directories #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?

[wenjin272]

Thanks for pointing this out. I think #941 and this PR should focus on dependency-generation changes during failover of the same Job. Reusing the TaskManager across different Job IDs is a separate multi-job lifecycle problem, relevant to applications containing multiple jobs and Session Mode.

Supporting that scenario requires additional generation cleanup and coordination of process-global import/cache state. I suggest tracking it in a follow-up issue/PR and clarifying that this PR only guarantees same-Job failover recovery.

[weiqingy]

_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?

[weiqingy]

_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"?

[weiqingy]

_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?

[weiqingy]

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?