Close owned agent clients on exit

 * Track AgentClient ownership for public and internal Polyaxon clients.
 * Close owned clients from sync and async agent exit hooks to avoid leaking async SDK sessions.

Author: md-mq

cli/polyaxon/_runner/agent/async_agent.py

@@ -61,7 +61,10 @@ async def __aenter__(self):
         return await self._enter()
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self._exit()
+        try:
+            await self._exit()
+        finally:
+            await self.client.aclose()
 
     async def refresh_executor(self):
         if (

cli/polyaxon/_runner/agent/base_agent.py

@@ -8,7 +8,7 @@
 from polyaxon._auxiliaries import V1PolyaxonInitContainer, V1PolyaxonSidecarContainer
 from polyaxon._connections import V1Connection
 from polyaxon._constants.globals import DEFAULT
-from polyaxon._runner.agent.client import AgentClient
+from polyaxon._runner.agent.client import AgentClient, AsyncAgentClient
 from polyaxon._runner.executor import BaseExecutor
 from polyaxon._schemas.checks import ChecksConfig
 from polyaxon._schemas.lifecycle import LiveState, V1Statuses
@@ -44,9 +44,8 @@ def __init__(
         self._graceful_shutdown = False
         self._last_data_collected_at = last_hour
         self._last_reconciled_at = last_hour
-        self.client = AgentClient(
-            owner=owner, agent_uuid=agent_uuid, is_async=self.IS_ASYNC
-        )
+        agent_client_cls = AsyncAgentClient if self.IS_ASYNC else AgentClient
+        self.client = agent_client_cls(owner=owner, agent_uuid=agent_uuid)
         self.executor = self.EXECUTOR()
         self.content = settings.AGENT_CONFIG.to_json()
 

cli/polyaxon/_runner/agent/client.py

@@ -3,39 +3,57 @@
 
 from polyaxon._schemas.lifecycle import V1StatusCondition, V1Statuses
 from polyaxon.client import PolyaxonClient, V1Agent, V1AgentStateResponse
+from polyaxon.exceptions import PolyaxonClientException
 from polyaxon.logger import logger
 
 
-class AgentClient:
+class _AgentClientBase:
+    _IS_ASYNC = False
+
     def __init__(
         self,
         owner: Optional[str] = None,
         agent_uuid: Optional[str] = None,
         client: Optional[PolyaxonClient] = None,
         internal_client: Optional[PolyaxonClient] = None,
-        is_async: bool = False,
     ):
         self.owner = owner
         self.agent_uuid = agent_uuid
-        self.is_async = is_async
+        self._validate_client_mode(client, "client")
+        self._validate_client_mode(internal_client, "internal_client")
         self._client = client
         self._internal_client = internal_client
+        self._created_client = None
+        self._created_internal_client = None
+
+    def _validate_client_mode(self, client: Optional[PolyaxonClient], name: str):
+        if client is None:
+            return
+        is_async = getattr(client, "is_async", None)
+        if isinstance(is_async, bool) and is_async != self._IS_ASYNC:
+            raise PolyaxonClientException(
+                "Injected `{}` transport mode does not match AgentClient mode.".format(
+                    name
+                )
+            )
 
     @property
     def client(self):
         if self._client:
             return self._client
-        self._client = PolyaxonClient(is_async=self.is_async)
+        self._client = PolyaxonClient(is_async=self._IS_ASYNC)
+        self._created_client = self._client
         return self._client
 
     @property
     def internal_client(self):
         if self._internal_client:
             return self._internal_client
         self._internal_client = PolyaxonClient(
-            is_async=self.is_async,
+            is_async=self._IS_ASYNC,
             is_internal=True,
         )
+        self._created_internal_client = self._internal_client
         return self._internal_client
 
     @property
@@ -181,3 +199,21 @@ def log_run_status(
             uuid=run_uuid,
             body={"condition": status_condition},
         )
+
+
+class AgentClient(_AgentClientBase):
+    def close(self):
+        if self._created_client is not None:
+            self._created_client.close()
+        if self._created_internal_client is not None:
+            self._created_internal_client.close()
+
+
+class AsyncAgentClient(_AgentClientBase):
+    _IS_ASYNC = True
+
+    async def aclose(self):
+        if self._created_client is not None:
+            await self._created_client.aclose()
+        if self._created_internal_client is not None:
+            await self._created_internal_client.aclose()

cli/polyaxon/_runner/agent/sync_agent.py

@@ -60,7 +60,10 @@ def __enter__(self):
         return self._enter()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self._exit()
+        try:
+            self._exit()
+        finally:
+            self.client.close()
 
     def refresh_executor(self):
         if (

cli/tests/test_k8s/test_async_agent.py

@@ -3,7 +3,7 @@
 
 from polyaxon._k8s.agent.async_agent import AsyncAgent
 from polyaxon._k8s.executor.async_executor import AsyncExecutor
-from polyaxon._runner.agent.client import AgentClient
+from polyaxon._runner.agent.client import AsyncAgentClient
 from polyaxon._utils.test_utils import AsyncMock, patch_settings
 
 
@@ -18,7 +18,7 @@ async def test_init_agent_component(register):
     agent = AsyncAgent(owner="foo", agent_uuid="uuid")
     assert agent.max_interval == 6
     assert isinstance(agent.executor, AsyncExecutor)
-    assert isinstance(agent.client, AgentClient)
+    assert isinstance(agent.client, AsyncAgentClient)
     assert register.call_count == 0
 
 
@@ -43,7 +43,7 @@ async def test_init_agent(
     agent.executor.manager.get_version.return_value = {}
     assert agent.max_interval == 6
     assert agent.executor is not None
-    assert isinstance(agent.client, AgentClient)
+    assert isinstance(agent.client, AsyncAgentClient)
     assert get_agent.call_count == 0
     assert get_agent_state.call_count == 0
     assert create_agent_status.call_count == 0
@@ -53,9 +53,24 @@ async def test_init_agent(
     await agent._enter()
     assert agent.max_interval == 6
     assert agent.executor is not None
-    assert isinstance(agent.client, AgentClient)
+    assert isinstance(agent.client, AsyncAgentClient)
     assert get_agent.call_count == 1
     assert get_agent_state.call_count == 0
     assert create_agent_status.call_count == 1
     assert sync_agent.call_count == 1
     assert agent.executor.manager.get_version.call_count == 1
+
+
+@pytest.mark.agent_mark
+@pytest.mark.asyncio
+async def test_async_agent_aexit_closes_client_in_finally():
+    patch_settings()
+    agent = AsyncAgent(owner="foo", agent_uuid="uuid")
+    agent.client = MagicMock()
+    agent.client.aclose = AsyncMock()
+    agent._exit = AsyncMock(side_effect=RuntimeError("exit failed"))
+
+    with pytest.raises(RuntimeError):
+        await agent.__aexit__(None, None, None)
+
+    agent.client.aclose.assert_called_once()

cli/tests/test_runner/test_agent_client.py

@@ -0,0 +1,131 @@
+from mock import MagicMock, patch
+import pytest
+
+from polyaxon._runner.agent.client import AgentClient, AsyncAgentClient
+from polyaxon.exceptions import PolyaxonClientException
+
+
+pytestmark = pytest.mark.agent_mark
+
+
+class ClientMock:
+    def __init__(self, is_async: bool):
+        self.is_async = is_async
+        self.agents_v1 = MagicMock()
+        self.runs_v1 = MagicMock()
+        self.close_calls = 0
+        self.aclose_calls = 0
+
+    def close(self):
+        self.close_calls += 1
+
+    async def aclose(self):
+        self.aclose_calls += 1
+
+
+@patch("polyaxon._runner.agent.client.PolyaxonClient")
+def test_agent_client_close_closes_owned_clients(client_cls):
+    public_client = ClientMock(is_async=False)
+    internal_client = ClientMock(is_async=False)
+    client_cls.side_effect = [public_client, internal_client]
+    client = AgentClient(owner="foo", agent_uuid="uuid")
+
+    assert client.client is public_client
+    assert client.internal_client is internal_client
+
+    client.close()
+
+    assert public_client.close_calls == 1
+    assert internal_client.close_calls == 1
+
+
+@pytest.mark.asyncio
+@patch("polyaxon._runner.agent.client.PolyaxonClient")
+async def test_agent_client_aclose_closes_owned_clients(client_cls):
+    public_client = ClientMock(is_async=True)
+    internal_client = ClientMock(is_async=True)
+    client_cls.side_effect = [public_client, internal_client]
+    client = AsyncAgentClient(owner="foo", agent_uuid="uuid")
+
+    assert client.client is public_client
+    assert client.internal_client is internal_client
+
+    await client.aclose()
+
+    assert public_client.aclose_calls == 1
+    assert internal_client.aclose_calls == 1
+
+
+def test_agent_client_close_does_not_close_injected_clients():
+    public_client = ClientMock(is_async=False)
+    internal_client = ClientMock(is_async=False)
+    client = AgentClient(
+        owner="foo",
+        agent_uuid="uuid",
+        client=public_client,
+        internal_client=internal_client,
+    )
+
+    client.close()
+
+    assert public_client.close_calls == 0
+    assert internal_client.close_calls == 0
+
+
+@pytest.mark.asyncio
+async def test_agent_client_aclose_does_not_close_injected_clients():
+    public_client = ClientMock(is_async=True)
+    internal_client = ClientMock(is_async=True)
+    client = AsyncAgentClient(
+        owner="foo",
+        agent_uuid="uuid",
+        client=public_client,
+        internal_client=internal_client,
+    )
+
+    await client.aclose()
+
+    assert public_client.aclose_calls == 0
+    assert internal_client.aclose_calls == 0
+
+
+def test_async_agent_client_does_not_expose_sync_close():
+    client = AsyncAgentClient(owner="foo", agent_uuid="uuid")
+
+    assert not hasattr(client, "close")
+
+
+def test_agent_client_does_not_expose_async_close():
+    client = AgentClient(owner="foo", agent_uuid="uuid")
+
+    assert not hasattr(client, "aclose")
+
+
+def test_agent_client_rejects_mode_mismatch_public_client():
+    with pytest.raises(PolyaxonClientException):
+        AsyncAgentClient(
+            owner="foo",
+            agent_uuid="uuid",
+            client=ClientMock(is_async=False),
+        )
+
+
+def test_agent_client_rejects_mode_mismatch_internal_client():
+    with pytest.raises(PolyaxonClientException):
+        AgentClient(
+            owner="foo",
+            agent_uuid="uuid",
+            internal_client=ClientMock(is_async=True),
+        )
+
+
+def test_agent_client_allows_loose_mocks_without_mode_attribute():
+    client = AgentClient(
+        owner="foo",
+        agent_uuid="uuid",
+        client=MagicMock(),
+        internal_client=MagicMock(),
+    )
+
+    assert client.client is not None
+    assert client.internal_client is not None

cli/tests/test_runner/test_base_agent.py

@@ -2,7 +2,7 @@
 import pytest
 
 from polyaxon._constants.globals import DEFAULT
-from polyaxon._runner.agent.client import AgentClient
+from polyaxon._runner.agent.client import AgentClient, AsyncAgentClient
 from polyaxon._runner.agent.sync_agent import BaseSyncAgent
 from polyaxon._utils.test_utils import BaseTestCase
 
@@ -148,8 +148,18 @@ def test_agent_client_uses_internal_client_for_collect_agent_data(self):
 
     @patch("polyaxon._runner.agent.client.PolyaxonClient")
     def test_agent_client_creates_internal_client_with_internal_mode(self, client_cls):
-        client = AgentClient(owner="foo", agent_uuid="uuid", is_async=True)
+        client = AsyncAgentClient(owner="foo", agent_uuid="uuid")
 
         _ = client.internal_client
 
         client_cls.assert_called_once_with(is_async=True, is_internal=True)
+
+    def test_sync_agent_exit_closes_client_in_finally(self):
+        agent = DummyAgent(owner="foo", agent_uuid="uuid")
+        agent.client = MagicMock()
+        agent._exit = MagicMock(side_effect=RuntimeError("exit failed"))
+
+        with pytest.raises(RuntimeError):
+            agent.__exit__(None, None, None)
+
+        agent.client.close.assert_called_once()