fix(framework): Restart AgentApps after dependency installation

framework/py/flwr/supercore/task_process/agent/run_agentapp.py

@@ -18,6 +18,7 @@
 import hashlib
 import os
 import signal
+import sys
 import threading
 from collections.abc import Callable, Iterator
 from concurrent.futures import Future
@@ -26,6 +27,7 @@
 from logging import DEBUG, ERROR
 from pathlib import Path
 from queue import Queue
+from tempfile import NamedTemporaryFile
 from time import monotonic
 from types import FrameType
 from typing import Any, cast
@@ -107,6 +109,29 @@
 _RUNTIME_API_KEY_ENV = "FLWR_RUNTIME_API_KEY"
 _RUNTIME_BASE_URL_ENV = "FLWR_RUNTIME_BASE_URL"
 _SSL_CERT_FILE_ENV = "SSL_CERT_FILE"
+_RUNTIME_ENV = "FLWR_AGENTAPP_RUNTIME_ENV"
+
+
+def _restart_with_runtime_env(
+    runtime_env_dir: Path, task_input: PullTaskInputResponse, token: str
+) -> None:
+    """Restart with app dependencies available before framework imports."""
+    input_path = runtime_env_dir / "task-input.pb"
+    descriptor = os.open(input_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
+    with os.fdopen(descriptor, "wb") as input_file:
+        input_file.write(task_input.SerializeToString())
+    environ = os.environ.copy()
+    environ[_RUNTIME_ENV] = str(runtime_env_dir)
+    environ["PYTHONPATH"] = os.pathsep.join(sys.path)
+    args = [sys.executable, *sys.orig_argv[1:]]
+    if "--token-stdin" in args:
+        with NamedTemporaryFile(
+            mode="w", encoding="utf-8", dir=runtime_env_dir, delete=False
+        ) as token_file:
+            token_file.write(token)
+        args.remove("--token-stdin")
+        args.extend(["--token-file", token_file.name])
+    os.execve(sys.executable, args, environ)
 
 
 @dataclass(frozen=True)
@@ -262,9 +287,8 @@ def run(self) -> int:  # pylint: disable=too-many-locals,too-many-statements
                 self._heartbeat_sender.start()
 
                 log(DEBUG, "[flwr-agentapp] Pull task input")
-                res: PullTaskInputResponse = grid._runtime_client.PullTaskInput(
-                    PullTaskInputRequest()
-                )
+                res = self._pull_task_input()
+                runtime_ready = self._runtime_env_dir is not None
 
                 self._context = context_from_proto(res.context)
                 run = run_from_proto(res.run)
@@ -295,6 +319,10 @@ def run(self) -> int:  # pylint: disable=too-many-locals,too-many-statements
                 "agentapp.prepare",
                 attributes=self._trace_attributes,
             ):
+                app_path, agent_app, agent_app_attr = self._prepare_task_app(fab, run)
+                if self._runtime_env_dir is not None and not runtime_ready:
+                    _restart_with_runtime_env(self._runtime_env_dir, res, self._token)
+
                 # Initialize the AgentApp session
                 with trace_span("agentapp.session"):
                     prompt, instruction_metadata = pull_prompt(grid)
@@ -328,7 +356,6 @@ def run(self) -> int:  # pylint: disable=too-many-locals,too-many-statements
                         grid=agent_grid,
                     )
 
-                app_path, agent_app, agent_app_attr = self._prepare_task_app(fab, run)
                 self._context.run_config = get_fused_config_from_dir(
                     app_path, run.override_config
                 )
@@ -382,6 +409,17 @@ def run(self) -> int:  # pylint: disable=too-many-locals,too-many-statements
 
         return exit_code
 
+    def _pull_task_input(self) -> PullTaskInputResponse:
+        """Restore claimed input after restart, or claim the task once."""
+        runtime_env = os.environ.pop(_RUNTIME_ENV, None)
+        if runtime_env is None:
+            return self._grid._runtime_client.PullTaskInput(PullTaskInputRequest())
+        self._runtime_env_dir = Path(runtime_env)
+        input_path = self._runtime_env_dir / "task-input.pb"
+        response = PullTaskInputResponse.FromString(input_path.read_bytes())
+        input_path.unlink()
+        return response
+
     def _prepare_task_app(
         self, fab: Fab, run: Run
     ) -> tuple[Path, AgentApp | None, str | None]:
@@ -400,7 +438,7 @@ def _prepare_task_app(
             install_from_fab(fab.content, skip_prompt=True)
         fab_id, fab_version = get_fab_metadata(fab.content)
         app_path = get_project_dir(fab_id, fab_version, fab.hash_str)
-        if self._runtime_dependency_install:
+        if self._runtime_dependency_install and self._runtime_env_dir is None:
             log(DEBUG, "[flwr-agentapp] Installing app dependencies.")
             with trace_span("agentapp.dependencies"):
                 self._runtime_env_dir = install_app_dependencies(
@@ -418,7 +456,7 @@ def _prepare_task_app(
                     },
                 )
         else:
-            log(DEBUG, "[flwr-agentapp] Runtime dependency installation is disabled.")
+            log(DEBUG, "[flwr-agentapp] Using existing runtime dependencies.")
 
         config = get_project_config(app_path)
         agent_app_attr = config["tool"]["flwr"]["app"]["components"]["agentapp"]

framework/py/flwr/supercore/task_process/agent/run_agentapp_test.py

@@ -20,6 +20,8 @@
 import importlib
 import os
 import signal
+import subprocess
+import sys
 import threading
 from pathlib import Path
 from queue import Queue
@@ -29,7 +31,11 @@
 
 from flwr.agentapp import AgentApp, LoadAgentAppError
 from flwr.app import ConfigRecord, Message, RecordDict
-from flwr.common.constant import SubStatus
+from flwr.common.constant import (
+    FLWR_AGENTAPP_TOKEN_STDIN_ACKNOWLEDGEMENT,
+    FLWR_TASK_TOKEN_LENGTH,
+    SubStatus,
+)
 from flwr.supercore.constant import (
     AGENT_MESSAGE_CONTENT_RECORD_KEY,
     AGENT_MESSAGE_TEXT_KEY,
@@ -572,3 +578,196 @@ def test_run_agentapp_keeps_cold_process_behavior(
             "success": False,
         },
     )
+
+
+def test_restart_loads_app_dependencies_before_cached_modules(tmp_path: Path) -> None:
+    """Restart discards cached dependencies and restores input without reclaiming it."""
+    base = tmp_path / "base"
+    runtime = tmp_path / "runtime"
+    base.mkdir()
+    runtime.mkdir()
+    (base / "cached_dependency.py").write_text("version = 'old'\n")
+    (runtime / "cached_dependency.py").write_text("Converter = object\n")
+    (runtime / "app_dependency.py").write_text(
+        "from cached_dependency import Converter\n"
+    )
+    probe = tmp_path / "probe.py"
+    probe.write_text(
+        """
+import importlib
+import os
+import stat
+import sys
+from pathlib import Path
+from queue import Queue
+from unittest.mock import Mock, patch
+
+from flwr.proto.runtime_pb2 import PullTaskInputResponse
+module = importlib.import_module('flwr.supercore.task_process.agent.run_agentapp')
+base, runtime = map(Path, sys.argv[1:])
+if module._RUNTIME_ENV not in os.environ:
+    sys.path.insert(0, str(base))
+    import cached_dependency
+    sys.path.insert(0, str(runtime))
+    try:
+        import app_dependency
+    except ImportError as error:
+        assert 'Converter' in str(error)
+    else:
+        raise AssertionError('Expected the cached-module import failure')
+    module._restart_with_runtime_env(
+        runtime, PullTaskInputResponse(task_id=123), 'token'
+    )
+    raise AssertionError('Restart returned')
+
+assert stat.S_IMODE((runtime / 'task-input.pb').stat().st_mode) == 0o600
+with patch.object(module, 'HttpGrid') as grid:
+    lifecycle = module._AgentAppTaskLifecycle(
+        runtime_api_address='localhost:9092', log_queue=Queue(), token='token',
+        insecure=True, certificates=None, certificates_path=None,
+        runtime_dependency_install=True, preloaded=None,
+    )
+    response = lifecycle._pull_task_input()
+    assert response.task_id == 123
+    assert lifecycle._runtime_env_dir == runtime
+    grid.return_value._runtime_client.PullTaskInput.assert_not_called()
+assert module._RUNTIME_ENV not in os.environ
+assert not (runtime / 'task-input.pb').exists()
+import app_dependency
+assert app_dependency.Converter is object
+lifecycle.finalize()
+lifecycle.finalize()
+grid.return_value._runtime_client.PushTaskOutput.assert_called_once()
+assert not runtime.exists()
+print('Fresh dependency imports and one-time task input: PASS')
+"""
+    )
+    result = subprocess.run(
+        [sys.executable, str(probe), str(base), str(runtime)],
+        check=True,
+        capture_output=True,
+        text=True,
+        env={**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)},
+        timeout=60,
+    )
+    assert "Fresh dependency imports and one-time task input: PASS" in result.stdout
+
+
+def test_dependency_restart_precedes_initial_prompt(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    """Do not consume a chat prompt before restarting into installed dependencies."""
+    grid = Mock()
+    monkeypatch.setattr(run_agentapp_module, "HttpGrid", Mock(return_value=grid))
+    monkeypatch.setattr(run_agentapp_module, "HeartbeatSender", Mock())
+    monkeypatch.setattr(run_agentapp_module, "start_log_uploader", Mock())
+    monkeypatch.setattr(
+        run_agentapp_module, "context_from_proto", Mock(return_value=Mock(node_id=0))
+    )
+    monkeypatch.setattr(
+        run_agentapp_module, "run_from_proto", Mock(return_value=Run.create_empty(1))
+    )
+    monkeypatch.setattr(run_agentapp_module, "fab_from_proto", Mock())
+    monkeypatch.delenv(run_agentapp_module._RUNTIME_ENV, raising=False)
+    lifecycle = _AgentAppTaskLifecycle(
+        runtime_api_address="runtime.example:9092",
+        log_queue=Queue(),
+        token="task-token",
+        insecure=True,
+        certificates=None,
+        certificates_path=None,
+        runtime_dependency_install=True,
+        preloaded=None,
+    )
+
+    def prepare(_fab: Fab, _run: Run) -> tuple[Path, None, str]:
+        lifecycle._runtime_env_dir = tmp_path
+        return tmp_path, None, "app:agent"
+
+    monkeypatch.setattr(lifecycle, "_prepare_task_app", prepare)
+    restart = Mock(side_effect=SystemExit(0))
+    prompt = Mock()
+    monkeypatch.setattr(run_agentapp_module, "_restart_with_runtime_env", restart)
+    monkeypatch.setattr(run_agentapp_module, "pull_prompt", prompt)
+    with pytest.raises(SystemExit):
+        lifecycle.run()
+    grid._runtime_client.PullTaskInput.assert_called_once()
+    restart.assert_called_once_with(
+        tmp_path, grid._runtime_client.PullTaskInput.return_value, "task-token"
+    )
+    prompt.assert_not_called()
+
+
+def test_restart_preserves_stdin_cli_token(tmp_path: Path) -> None:
+    """Preserve stdin authentication privately and acknowledge startup only once."""
+    probe = tmp_path / "stdin_probe.py"
+    probe.write_text(
+        """
+import importlib
+import os
+import stat
+import sys
+from pathlib import Path
+from queue import Queue
+from unittest.mock import patch
+
+from flwr.proto.runtime_pb2 import PullTaskInputResponse
+cli = importlib.import_module('flwr.supercore.cli.flwr_agentapp')
+module = importlib.import_module('flwr.supercore.task_process.agent.run_agentapp')
+runtime = Path(__file__).parent / 'runtime'
+
+def run_agentapp(**kwargs):
+    assert kwargs['runtime_dependency_install'] is True
+    token = kwargs['token']
+    if module._RUNTIME_ENV not in os.environ:
+        assert sys.stdin.closed
+        runtime.mkdir()
+        module._restart_with_runtime_env(
+            runtime, PullTaskInputResponse(task_id=123), token
+        )
+        raise AssertionError('Restart returned')
+    assert '--token-stdin' not in sys.argv
+    assert token not in '\\0'.join(sys.orig_argv)
+    token_path = Path(sys.argv[sys.argv.index('--token-file') + 1])
+    assert stat.S_IMODE(token_path.stat().st_mode) == 0o600
+    assert token_path.read_text() == token
+    with patch.object(module, 'HttpGrid') as grid:
+        lifecycle = module._AgentAppTaskLifecycle(
+            runtime_api_address=kwargs['runtime_api_address'], log_queue=Queue(),
+            token=token, insecure=True, certificates=None, certificates_path=None,
+            runtime_dependency_install=True, preloaded=None,
+        )
+        assert lifecycle._pull_task_input().task_id == 123
+        grid.return_value._runtime_client.PullTaskInput.assert_not_called()
+        lifecycle.finalize()
+        grid.return_value._runtime_client.PushTaskOutput.assert_called_once()
+    assert not runtime.exists()
+    print('Stdin-authenticated CLI restart: PASS')
+
+with (
+    patch.object(cli, 'run_agentapp', run_agentapp),
+    patch.object(cli, 'mirror_output_to_queue'),
+    patch.object(cli, 'restore_output'),
+):
+    cli.flwr_agentapp()
+"""
+    )
+    token = "ab" * FLWR_TASK_TOKEN_LENGTH
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(probe),
+            "--token-stdin",
+            "--insecure",
+            "--allow-runtime-dependency-installation",
+        ],
+        input=f"{token}\n",
+        check=True,
+        capture_output=True,
+        text=True,
+        env={**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)},
+        timeout=60,
+    )
+    assert result.stdout.count(FLWR_AGENTAPP_TOKEN_STDIN_ACKNOWLEDGEMENT) == 1
+    assert "Stdin-authenticated CLI restart: PASS" in result.stdout
+    assert token not in result.stdout + result.stderr